nmr_schedule/modifiers/
mod.rs

1//! Contains various modifiers and filters that can be applied to generators.
2
3mod basic;
4mod iterate;
5mod psf_polisher;
6mod tm_filter;
7
8pub use basic::*;
9pub use iterate::*;
10use ndarray::Dimension;
11pub use psf_polisher::*;
12pub use tm_filter::*;
13
14use alloc::borrow::ToOwned;
15
16use crate::{
17    generators::{Generator, Trace},
18    Schedule,
19};
20
21/// Generators returned by modifiers
22pub mod generators {
23    pub use super::basic::generators::*;
24    pub use super::iterate::generators::*;
25    pub use super::psf_polisher::generators::*;
26    pub use super::tm_filter::generators::*;
27}
28
29/// Post-processes existing generators.
30///
31/// Modifiers get access to the original generator, so they may call the underlying generator repeatedly to get schedules they would like.
32pub trait Modifier<Dim: Dimension> {
33    /// The new generator after applying the modifier
34    type Output<T: Generator<Dim>>: Generator<Dim>;
35
36    /// Apply the modifier to a schedule generator
37    fn modify<T: Generator<Dim>>(self, generator: T) -> Self::Output<T>;
38}
39
40/// Post-process an existing schedule without requiring access to the generator.
41///
42/// Filters do not get access to the underlying generator and operate only on a schedule. This allows filters to be applied to user-supplied schedules unlike modifiers.
43pub trait Filter<Dim: Dimension> {
44    /// Apply the filter to a schedule
45    fn filter(&self, sched: Schedule<Dim>) -> Schedule<Dim> {
46        self.filter_with_iter_and_trace(sched, 0).into_sched()
47    }
48
49    /// Apply the filter to a schedule while providing a trace
50    fn filter_with_iter_and_trace(&self, sched: Schedule<Dim>, iteration: u64) -> Trace<Dim>;
51
52    /// Filter a schedule from an existing trace while pushing a new trace onto it
53    fn filter_from_iter_and_trace(&self, mut trace: Trace<Dim>, iteration: u64) -> Trace<Dim> {
54        let mut trace_after = self.filter_with_iter_and_trace(trace.sched().to_owned(), iteration);
55
56        trace.stack.append(&mut trace_after.stack);
57
58        trace
59    }
60}
61
62/// Creates boilerplate (including a builder extension trait) for modifiers
63///
64/// Example:
65/// ```
66/// # use nmr_schedule::{*, generators::*, modifiers::*, pdf::*};
67/// # use ndarray::*;
68/// # pub struct MagicifyGenerator<T, F>(T, F);
69/// # impl<T, F> Generator<Ix1> for MagicifyGenerator<T, F> {
70/// #     fn _generate(&self, count: usize, dims: Ix1, iteration: u64) -> Trace<Ix1> {
71/// #         todo!()
72/// #     }
73/// # }
74/// # fn typecheck() {
75/// #     Quantiles::new(linear).magicify(12, |s| 5.);
76/// # }
77///
78/// modifier!(<F: Fn(&Schedule<Ix1>) -> f64> Magicify<Ix1>,
79///     MagicifyBuilder,
80///     r"Apply magic to the schedule",
81///     magicify,
82///     magic_amt: usize,
83///     magic_fn: F
84/// );
85///
86/// impl<F: Fn(&Schedule<Ix1>) -> f64> Modifier<Ix1> for Magicify<F> {
87///     type Output<T: Generator<Ix1>> = MagicifyGenerator<T, F>;
88///
89///     fn modify<T: Generator<Ix1>>(self, generator: T) -> Self::Output<T> {
90///         // ...
91///         # todo!()
92///     }
93/// }
94/// ```
95///
96/// - The parameter in angle brackets is an optional list generic parameters with bounds
97/// - The second parameter is the name of the modifier. The macro will create a tuple-struct with that name and all the parameters.
98/// - The third supplies documentation to the modifier and to the builder method.
99/// - The fourth is the name of the builder trait that will be automatically implemented on all generators
100/// - The fifth is the name of the builder method that will be able to be called on generators (for example `.tm_filter()`).
101/// - The rest are arguments to be passed into the modifier's `new` function, the builder method, and added as fields to the modifier.
102#[macro_export]
103macro_rules! modifier {
104    ($(<$($name:ident : $bounds:path),+>)? $modifier:ident <$dim:ident>, $trait_name:ident, $doc:expr, $builder_name: ident, $($arg:ident : $ty:ty),*) => {
105        #[doc=$doc]
106        #[doc=concat!("\n\nCan be used as the `.", stringify!($builder_name), "()` method.")]
107        #[derive(Debug)]
108        #[allow(missing_copy_implementations)]
109        pub struct $modifier $(<$($name: $bounds),+>)? ($($ty),*);
110
111        impl $(<$($name: $bounds),+>)? $modifier $(<$($name),+>)? {
112            #[doc=concat!("Create a new instance of the modifier.")]
113            #[allow(clippy::new_without_default)]
114            pub const fn new($($arg : $ty),*) -> $modifier $(<$($name),+>)? {
115                $modifier($($arg),*)
116            }
117        }
118
119        #[doc=concat!("Builder trait for the [`", stringify!($modifier), "`] modifier.")]
120        pub trait $trait_name: Generator<$dim> {
121            #[doc=$doc]
122            fn $builder_name<$($($name: $bounds),+)?>(self, $($arg : $ty),*) -> <$modifier $(<$($name),+>)? as Modifier<$dim>>::Output<Self> where Self: Sized +  {
123                self.then($modifier::new($($arg),*))
124            }
125        }
126
127        impl<T: Generator<$dim>> $trait_name for T {}
128    };
129    ($dim:ident $(<$($name:ident : $bounds:path),+>)? $modifier:ident, $trait_name:ident, $doc:expr, $builder_name:ident, $($arg:ident : $ty:ty),*) => {
130        #[doc=$doc]
131        #[doc=concat!("\n\nCan be used as the `.", stringify!($builder_name), "()` method.")]
132        #[derive(Debug)]
133        #[allow(missing_copy_implementations)]
134        pub struct $modifier <$dim: Dimension, $($($name: $bounds),+)?> ($($ty),*, PhantomData<$dim>);
135
136        impl <$dim: Dimension, $($($name: $bounds),+)?> $modifier <$dim $(, $($name),+)?> {
137            #[doc=concat!("Create a new instance of the modifier.")]
138            #[allow(clippy::new_without_default)]
139            pub const fn new($($arg : $ty),*) -> $modifier <$dim $(, $($name),+)?> {
140                $modifier($($arg),*, PhantomData)
141            }
142        }
143
144        #[doc=concat!("Builder trait for the [`", stringify!($modifier), "`] modifier.")]
145        pub trait $trait_name<$dim: Dimension>: Generator<$dim> {
146            #[doc=$doc]
147            fn $builder_name<$($($name: $bounds),+)?>(self, $($arg : $ty),*) -> <$modifier <$dim $(, $($name),+)?> as Modifier<$dim>>::Output<Self> where Self: Sized +  {
148                self.then($modifier::new($($arg),*))
149            }
150        }
151
152        impl<Dim: Dimension, T: Generator<Dim>> $trait_name<Dim> for T {}
153    };
154}