nmr_schedule/generators/
poisson_gap.rs

1use core::cmp::Ordering;
2use core::fmt::Display;
3
4use alloc::vec;
5use ndarray::{Array, Ix1};
6use rand::{Rng, SeedableRng};
7use rand_chacha::ChaCha12Rng;
8
9use crate::{Monotonicity, Schedule};
10
11use super::{xor_iteration, Generator, Trace};
12
13fn sample_poisson(λ: f64, fail_after: usize, mut rng: impl Rng) -> usize {
14    let threshold = (-λ).exp();
15    let mut i = 0;
16    let mut u = 1.;
17
18    loop {
19        u *= rng.random::<f64>();
20
21        if u < threshold || i > fail_after {
22            return i;
23        }
24
25        i += 1;
26    }
27}
28
29/// The original Sine Weighted Poisson-gap algorithm.
30///
31/// The iteration parameter alters the random seed.
32///
33/// S. G. Hyberts, K. Takeuchi, G. Wagner, J Am Chem Soc. 2010, 132(7), 2145-2147. <https://doi.org/10.1021/ja908004w>
34#[derive(Clone, Copy, Debug)]
35pub struct SinWeightedPoissonGap {
36    seed: [u8; 32],
37}
38
39impl SinWeightedPoissonGap {
40    /// Create a new `SinWeightedPoissonGap` from a random seed.
41    pub const fn new(seed: [u8; 32]) -> SinWeightedPoissonGap {
42        SinWeightedPoissonGap { seed }
43    }
44}
45
46/// The trace for `SinWeightedPoissonGap`
47#[derive(Clone, Copy, Debug)]
48pub struct SinWeightedPoissonGapTrace {
49    /// How many iterations were required to find a λ that gave the correct count
50    pub iterations: usize,
51    /// Which values of λ and `iteration` gave the correct count
52    pub final_params: (f64, u64),
53}
54
55impl Display for SinWeightedPoissonGapTrace {
56    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
57        write!(
58            f,
59            "Generation tries needed: {}, Final parameters: ⟨λ {}, Seed XOR = {}⟩",
60            self.iterations, self.final_params.0, self.final_params.1
61        )
62    }
63}
64
65impl Generator<Ix1> for SinWeightedPoissonGap {
66    fn _generate(&self, count: usize, dims: Ix1, iteration: u64) -> Trace<Ix1> {
67        if count == 0 {
68            return Trace::new(
69                Schedule::new(Array::from_vec(vec![false; dims[0]])),
70                SinWeightedPoissonGapTrace {
71                    iterations: 0,
72                    final_params: (0., 0),
73                },
74            );
75        }
76
77        let (sched, λ, iterations) = generate_by_fuzz(
78            count,
79            iteration,
80            Monotonicity::Decreasing,
81            dims[0] as f64 / count as f64,
82            |arbitrary, iteration| {
83                let mut rng = ChaCha12Rng::from_seed(xor_iteration(self.seed, iteration));
84                let mut schedule = alloc::vec![false; dims[0]];
85
86                let mut i = 0;
87
88                while i < dims[0] {
89                    schedule[i] = true;
90
91                    i += 1 + sample_poisson(
92                        arbitrary
93                            * ((i as f64 + 0.5) / (dims[0] as f64 + 1.)
94                                * core::f64::consts::FRAC_PI_2)
95                                .sin(),
96                        dims[0],
97                        &mut rng,
98                    );
99                }
100
101                Schedule::new(Array::from_vec(schedule))
102            },
103        );
104
105        Trace::new(
106            sched,
107            SinWeightedPoissonGapTrace {
108                iterations,
109                final_params: (λ),
110            },
111        )
112    }
113}
114
115/// Generate schedules by randomly varying the `iteration` parameter while at the same time nudging an arbitrary parameter until the algorithm gives the correct number of points.
116///
117/// Binary search cannot be used for Poisson Gap because it does not converge for all random seeds (small variations tend to snowball and prevent convergence)
118///
119/// This is also the method used by the code in the supplement of the original paper.
120fn generate_by_fuzz(
121    count: usize,
122    original_iteration: u64,
123    monotonicity: Monotonicity,
124    start_value: f64,
125    generate_one: impl Fn(f64, u64) -> Schedule<Ix1>,
126) -> (Schedule<Ix1>, (f64, u64), usize) {
127    let mut rng = ChaCha12Rng::from_seed([0; 32]);
128    let mut arbitrary = start_value;
129    let mut iterations = 0;
130
131    loop {
132        iterations += 1;
133
134        let iteration = if iterations == 1 {
135            0
136        } else {
137            rng.random::<u64>()
138        } ^ original_iteration;
139
140        let sched = generate_one(arbitrary, iteration);
141
142        let sched_count = sched.count();
143        let adjust = 1. + (2.5 / (sched.len() as f64) * (count as f64 - sched_count as f64).abs());
144
145        match (sched_count.cmp(&count), monotonicity) {
146            (Ordering::Less, Monotonicity::Increasing)
147            | (Ordering::Greater, Monotonicity::Decreasing) => arbitrary *= adjust,
148            (Ordering::Equal, _) => return (sched, (arbitrary, iteration), iterations),
149            (Ordering::Greater, Monotonicity::Increasing)
150            | (Ordering::Less, Monotonicity::Decreasing) => arbitrary /= adjust,
151        }
152    }
153}
nmr_schedule/generators/poisson_gap.rs

nmr_schedule/generators/
poisson_gap.rs
