nmr_schedule/pdf/
presets.rs

1use alloc::{borrow::ToOwned, vec::Vec};
2
3use crate::{find_zero, InitialState, Monotonicity};
4
5use super::Pdf;
6
7/// An unweighted PDF.
8///
9/// ```
10/// # use nmr_schedule::pdf::*;
11/// assert!(
12///     unweighted(256)
13///         .get_distribution()
14///         .iter()
15///         .all(|v| *v == 1. / 256.)
16/// );
17/// ```
18pub fn unweighted(len: usize) -> Pdf {
19    Pdf::from_integral(move |pos| pos / len as f64, len)
20}
21
22/// A linear PDF
23///
24/// ```
25/// # use nmr_schedule::pdf::*;
26/// assert!(
27///     linear(256)
28///         .get_distribution()
29///         .windows(2)
30///         .all(|v| v[0] - v[1] == 1. / 256. / 128.)
31/// );
32/// ```
33pub fn linear(len: usize) -> Pdf {
34    let inv_len = 1. / len as f64;
35    Pdf::from_integral(move |pos| inv_len * pos * (2. - inv_len * pos), len)
36}
37
38/// TODO: Citation
39/// An exponential PDF. `scale` represents how quickly the PDF decays.
40///
41/// For QSched compatibility, `scale` can be set to `bias·time`
42pub fn exponential(len: usize, scale: f64) -> Pdf {
43    let c = -scale / len as f64;
44    let normalize = 1. / ((-scale).exp() - 1.);
45
46    Pdf::from_integral(move |pos| normalize * (c * pos).exp(), len)
47}
48
49/// Options for the `bias` parameter used in QSched
50#[derive(Clone, Copy, Debug)]
51pub enum QSinBias {
52    /// Equivalent to setting 'bias' to '1'.
53    /// PDF is { 1 - sin(x) | 0 < x < time/2 }
54    Low,
55    /// Equivalent to setting 'bias' to '1.5'.
56    /// PDF is { (1 - sin(x))^1.8 | 0 < x < time/2 }
57    Med,
58    /// Equivalent to setting 'bias' to '2'.
59    /// PDF is { (0.85 - sin(x))^2 | 0 < x < time/2 }
60    High,
61}
62
63impl core::fmt::Display for QSinBias {
64    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
65        f.write_str(match self {
66            QSinBias::Low => "Low",
67            QSinBias::Med => "Med",
68            QSinBias::High => "High",
69        })
70    }
71}
72
73/// TODO: Citation
74/// Equivalent to the `quant-sin` option in QSched
75pub fn qsin(len: usize, bias: QSinBias, time: f64) -> Pdf {
76    let half_time = time / 2.;
77    let len_f = len as f64;
78
79    let s = half_time / len_f;
80
81    match bias {
82        QSinBias::Low => {
83            let f = |x: f64| x + x.cos();
84
85            let normalize = (f(half_time) - f(0.)).recip();
86
87            Pdf::from_integral(move |pos| f(pos * s) * normalize, len)
88        }
89        QSinBias::Med => {
90            let bars = 16;
91            let bars_f = bars as f64;
92
93            let mut discrete = alloc::vec![0.; len];
94
95            let mut sum = 0.;
96
97            for (i, v) in discrete.iter_mut().enumerate() {
98                let mut approximation = 0.;
99
100                for bar in 0..=bars {
101                    let x = i as f64 + bar as f64 / bars_f;
102                    approximation += (1. - (x * s).sin()).powf(1.8)
103                        * if bar == 0 || bar == bars { 0.5 } else { 1. };
104                }
105
106                *v = approximation;
107                sum += *v;
108            }
109
110            let scale = sum.recip();
111
112            discrete.iter_mut().for_each(|v| *v *= scale);
113
114            Pdf::from_discrete(discrete)
115        }
116        QSinBias::High => {
117            let f = |x: f64| 1.2225 * x + 1.7 * x.cos() - 0.25 * (2. * x).sin();
118
119            let normalize = (f(half_time) - f(0.)).recip();
120
121            Pdf::from_integral(move |pos| f(pos * s) * normalize, len)
122        }
123    }
124}
125
126/// Estimate the effective PDF of Sin Weighted Poisson Gap, if you would like to try it in other generators.
127///
128/// This has not been published and may be incorrect.
129#[allow(clippy::missing_panics_doc)] // Panicking is a bug
130pub fn sin_weighted_poisson_gap(len: usize, count: usize) -> Pdf {
131    // An expected gap length of 2 corresponds to a probability of (1/(2+1)) or 1/3:
132    // ▩__▩__▩__
133    // We want the sum of probabilities over the length to equal the count, so we can do a binary search for the correct value. Contributions welcome if you have a closed form for this.
134
135    let (_, mk_probabilities) = find_zero(
136        Monotonicity::Decreasing,
137        InitialState::new(len as f64 / count as f64, 0., f64::INFINITY),
138        crate::Precision::Image {
139            amount: 0.001,
140            debug: (&|_, _| panic!("Convergence failed")) as &_,
141        },
142        |arbitrary| {
143            let probabilities = (0..len).map(move |i| {
144                1. / (arbitrary
145                    * ((i as f64 + 0.5) / (len as f64 + 1.) * core::f64::consts::FRAC_PI_2).sin()
146                    + 1.)
147            });
148
149            (
150                probabilities.to_owned().sum::<f64>() - count as f64,
151                move || {
152                    probabilities.collect() // Return a function to collect instead of allocating a vector every time (which would be slow)
153                },
154            )
155        },
156    );
157
158    let probabilities: Vec<f64> = mk_probabilities();
159    let probabilities_ref = &probabilities;
160
161    // Binary search the count parameter to turn the probabilities into a PDF
162    // Again, contributions welcome if anyone knows a closed form for this.
163    let (_, mk_pdf) = find_zero(
164        Monotonicity::Increasing,
165        InitialState::new(count as f64, 0., f64::INFINITY),
166        crate::Precision::Image {
167            amount: 0.00001,
168            debug: (&|_, _| panic!("Convergence failed")) as &_,
169        },
170        |inverse_count| {
171            // This is actually the inverse operation of 1 - (1 - v)^(count)
172            // Not bothering to take 1/count, the binary search will do that
173            let pdf_gen = move |i| {
174                let complement: f64 = 1. - probabilities_ref[i];
175                1. - complement.powf(inverse_count)
176            };
177
178            ((0..len).map(pdf_gen).sum::<f64>() - 1., move || {
179                Pdf::from_discrete((0..len).map(pdf_gen).collect())
180            })
181        },
182    );
183
184    mk_pdf()
185}