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