src/rand_distr.rs

Thu, 08 Dec 2022 14:10:07 +0200

author
Tuomo Valkonen <tuomov@iki.fi>
date
Thu, 08 Dec 2022 14:10:07 +0200
changeset 21
0771706f472f
parent 0
eb3c7813b67a
child 23
9869fa1e0ccd
permissions
-rw-r--r--

Save more CSV files when iteration-wise plotting is enabled.

This helps to generate TikZ illustrations for presentations.

0
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
1 //! Random distribution wrappers and implementations
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
2
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
3 use numeric_literals::replace_float_literals;
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
4 use rand::Rng;
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
5 use rand_distr::{Distribution, Normal, StandardNormal, NormalError};
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
6 use serde::{Serialize, Deserialize};
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
7 use serde::ser::{Serializer, SerializeStruct};
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
8 use alg_tools::types::*;
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
9
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
10 /// Wrapper for [`Normal`] that can be serialized by serde.
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
11 pub struct SerializableNormal<T : Float>(Normal<T>)
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
12 where StandardNormal : Distribution<T>;
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
13
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
14 impl<T : Float> Distribution<T> for SerializableNormal<T>
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
15 where StandardNormal : Distribution<T> {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
16 fn sample<R>(&self, rng: &mut R) -> T
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
17 where
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
18 R : Rng + ?Sized
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
19 { self.0.sample(rng) }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
20 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
21
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
22 impl<T : Float> SerializableNormal<T>
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
23 where StandardNormal : Distribution<T> {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
24 pub fn new(mean : T, std_dev : T) -> Result<SerializableNormal<T>, NormalError> {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
25 Ok(SerializableNormal(Normal::new(mean, std_dev)?))
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
26 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
27 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
28
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
29 impl<F> Serialize for SerializableNormal<F>
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
30 where
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
31 StandardNormal : Distribution<F>,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
32 F: Float + Serialize,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
33 {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
34 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
35 where
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
36 S: Serializer,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
37 {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
38 let mut s = serializer.serialize_struct("Normal", 2)?;
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
39 s.serialize_field("mean", &self.0.mean())?;
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
40 s.serialize_field("std_dev", &self.0.std_dev())?;
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
41 s.end()
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
42 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
43 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
44
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
45 /// Salt-and-pepper noise distribution
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
46 ///
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
47 /// This is the distribution that outputs each $\\{-m,0,m\\}$ with the corresponding
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
48 /// probabilities $\\{1-p, p/2, p/2\\}$.
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
49 #[derive(Copy, Clone, Debug, Serialize, Deserialize)]
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
50 pub struct SaltAndPepper<T : Float>{
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
51 /// The magnitude parameter $m$
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
52 magnitude : T,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
53 /// The probability parameter $p$
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
54 probability : T
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
55 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
56
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
57 /// Error for [`SaltAndPepper`].
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
58 #[derive(Copy, Clone, Debug, Serialize, Deserialize)]
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
59 pub enum SaltAndPepperError {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
60 /// The probability parameter $p$ is not in the range [0, 1].
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
61 InvalidProbability,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
62 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
63 impl std::error::Error for SaltAndPepperError {}
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
64
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
65 impl std::fmt::Display for SaltAndPepperError {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
66 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
67 f.write_str(match self {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
68 SaltAndPepperError::InvalidProbability =>
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
69 " The probability parameter is not in the range [0, 1].",
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
70 })
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
71 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
72 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
73
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
74 #[replace_float_literals(T::cast_from(literal))]
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
75 impl<T : Float> SaltAndPepper<T> {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
76 pub fn new(magnitude : T, probability : T) -> Result<SaltAndPepper<T>, SaltAndPepperError> {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
77 if probability > 1.0 || probability < 0.0 {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
78 Err(SaltAndPepperError::InvalidProbability)
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
79 } else {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
80 Ok(SaltAndPepper { magnitude, probability })
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
81 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
82 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
83 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
84
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
85 #[replace_float_literals(T::cast_from(literal))]
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
86 impl<T : Float> Distribution<T> for SaltAndPepper<T> {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
87 fn sample<R>(&self, rng: &mut R) -> T
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
88 where
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
89 R : Rng + ?Sized
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
90 {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
91 let (p, sign) : (float, bool) = rng.gen();
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
92 match (p < self.probability.as_(), sign) {
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
93 (false, _) => 0.0,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
94 (true, true) => self.magnitude,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
95 (true, false) => -self.magnitude,
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
96 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
97 }
eb3c7813b67a Initial version
Tuomo Valkonen <tuomov@iki.fi>
parents:
diff changeset
98 }

mercurial