nmr_schedule/modifiers/
tm_filter.rs

1use core::fmt::Display;
2
3use crate::{
4    Schedule,
5    generators::{Generator, Trace},
6    modifier,
7};
8
9pub(crate) mod generators {
10    /// The generator after applying [`super::TMFilter`]
11    #[derive(Debug)]
12    pub struct TMFilterGenerator<T> {
13        pub(crate) generator: T,
14    }
15}
16
17use alloc::vec::Vec;
18use generators::*;
19use ndarray::Ix1;
20
21use super::{Filter, Modifier};
22
23/// The trace for `TMFilter`
24#[derive(Debug)]
25pub struct TMFilterTrace {
26    /// Gives the indices of each swap performed in the schedule
27    pub swaps_done: Vec<usize>,
28}
29
30impl Display for TMFilterTrace {
31    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
32        write!(f, "Performed {} swaps", self.swaps_done.len())
33    }
34}
35
36impl<T: Generator<Ix1>> Generator<Ix1> for TMFilterGenerator<T> {
37    fn _generate(&self, count: usize, dims: Ix1, iteration: u64) -> Trace<Ix1> {
38        let trace = self
39            .generator
40            .generate_with_iter_and_trace(count, dims, iteration);
41
42        TMFilter::filter_from_iter_and_trace(&TMFilter::new(), trace, iteration)
43    }
44}
45
46modifier!(
47    TMFilter<Ix1>,
48    TMFilterBuilder,
49    r#"Filter the schedule to reduce patterns by swapping samples to make the schedule "look like" the Thue-Morse sequence.
50    
51The Thue-Morse sequence is known to avoid consecutive patterns and this effectively transfers that property to the schedule.
52    
53The iteration parameter alters the slice of the Thue-Morse sequence that is used for filtering."#,
54    tm_filter,
55);
56
57impl Modifier<Ix1> for TMFilter {
58    type Output<T: Generator<Ix1>> = TMFilterGenerator<T>;
59
60    fn modify<T: Generator<Ix1>>(self, generator: T) -> Self::Output<T> {
61        TMFilterGenerator { generator }
62    }
63}
64
65impl Filter<Ix1> for TMFilter {
66    fn filter_with_iter_and_trace(&self, sched: Schedule<Ix1>, mut iteration: u64) -> Trace<Ix1> {
67        if sched.len() < 2 {
68            return Trace::new(
69                sched,
70                TMFilterTrace {
71                    swaps_done: Vec::new(),
72                },
73            );
74        }
75
76        if iteration > u64::MAX - (sched.len() as u64) {
77            iteration -= sched.len() as u64;
78        }
79
80        let mut sched = sched.into_inner();
81
82        const fn thue_morse(i: u64) -> bool {
83            i.count_ones() % 2 == 0
84        }
85
86        let mut swaps_done = Vec::new();
87
88        for i in 1..sched.len() - 2 {
89            let v64 = i as u64;
90            if sched[i] == thue_morse(v64 + iteration + 1)
91                && sched[i + 1] == thue_morse(v64 + iteration)
92                && sched[i] != sched[i + 1]
93            {
94                swaps_done.push(i);
95                sched[i] = thue_morse(v64 + iteration);
96                sched[i + 1] = thue_morse(v64 + iteration + 1);
97            }
98        }
99
100        Trace::new(Schedule::new(sched), TMFilterTrace { swaps_done })
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use ndarray::Ix1;
107
108    use crate::{
109        generators::{Generator, Quantiles},
110        pdf::unweighted,
111    };
112
113    use super::TMFilterBuilder;
114
115    #[test]
116    fn tm_overflow() {
117        Quantiles::new(unweighted)
118            .tm_filter()
119            .generate_with_iter(64, Ix1(256), u64::MAX);
120
121        Quantiles::new(unweighted)
122            .tm_filter()
123            .generate_with_iter(64, Ix1(256), u64::MAX - 253);
124    }
125}