nmr_schedule/modifiers/
tm_filter.rs1use core::fmt::Display;
2
3use crate::{
4 Schedule,
5 generators::{Generator, Trace},
6 modifier,
7};
8
9pub(crate) mod generators {
10 #[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#[derive(Debug)]
25pub struct TMFilterTrace {
26 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}