modewise_spectral/
lib.rs

1use modewise_method::{MethodDescriptor, MethodError};
2use modewise_signals::{MultichannelSeries, SamplingRate};
3use ndarray::{Array2, Array3, ArrayView2, ArrayView3};
4
5/// Frequency bins `[edges[i], edges[i + 1])`, spaced linearly or logarithmically.
6#[derive(Debug, Clone, PartialEq)]
7pub struct FrequencyGrid {
8    edges: Vec<f64>,
9    scale: GridScale,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum GridScale {
14    Linear,
15    Log,
16}
17
18impl FrequencyGrid {
19    pub fn new(
20        f_min: f64,
21        f_max: f64,
22        n_bins: usize,
23        scale: GridScale,
24    ) -> Result<Self, MethodError> {
25        if n_bins == 0 {
26            return Err(MethodError::invalid("n_bins", "must be at least 1"));
27        }
28        let lowest = if scale == GridScale::Log {
29            0.0
30        } else {
31            f64::NEG_INFINITY
32        };
33        if !(f_min.is_finite() && f_max.is_finite() && f_min > lowest && f_max > f_min) {
34            return Err(MethodError::invalid(
35                "frequency range",
36                format!(
37                    "need finite f_min < f_max{}, got [{f_min}, {f_max}]",
38                    if scale == GridScale::Log {
39                        " and f_min > 0"
40                    } else {
41                        ""
42                    }
43                ),
44            ));
45        }
46        let edges = (0..=n_bins)
47            .map(|i| {
48                let t = i as f64 / n_bins as f64;
49                match scale {
50                    GridScale::Linear => f_min + t * (f_max - f_min),
51                    GridScale::Log => f_min * (f_max / f_min).powf(t),
52                }
53            })
54            .collect();
55        Ok(Self { edges, scale })
56    }
57
58    pub fn edges(&self) -> &[f64] {
59        &self.edges
60    }
61
62    pub fn n_bins(&self) -> usize {
63        self.edges.len() - 1
64    }
65
66    /// Bin centres: arithmetic for linear grids, geometric for log grids.
67    pub fn centres(&self) -> Vec<f64> {
68        self.edges
69            .windows(2)
70            .map(|w| match self.scale {
71                GridScale::Linear => (w[0] + w[1]) / 2.0,
72                GridScale::Log => (w[0] * w[1]).sqrt(),
73            })
74            .collect()
75    }
76
77    pub fn scale(&self) -> GridScale {
78        self.scale
79    }
80
81    /// The bin holding `f`, or `None` when `f` is outside `[first edge, last edge)` or not finite.
82    pub fn bin(&self, f: f64) -> Option<usize> {
83        let last = *self.edges.last()?;
84        if !(f >= self.edges[0] && f < last) {
85            return None;
86        }
87        Some(self.edges.partition_point(|&e| e <= f) - 1)
88    }
89}
90
91/// Maps a signal directly to time-frequency power.
92pub trait TimeFrequencyTransform {
93    fn descriptor(&self) -> MethodDescriptor;
94    fn transform(&self, signal: &MultichannelSeries) -> Result<TimeFrequencyMap, MethodError>;
95}
96
97/// Power laid out as channel × frequency × time.
98#[derive(Debug, Clone, PartialEq)]
99pub struct TimeFrequencyMap {
100    power: Array3<f64>,
101    frequencies_hz: Vec<f64>,
102    times_s: Vec<f64>,
103    labels: Vec<String>,
104    sampling_rate: Option<SamplingRate>,
105    reliable: Option<Array2<bool>>,
106}
107
108impl TimeFrequencyMap {
109    pub fn new(
110        power: Array3<f64>,
111        frequencies_hz: Vec<f64>,
112        times_s: Vec<f64>,
113        labels: Vec<String>,
114    ) -> Result<Self, MethodError> {
115        let expected = (labels.len(), frequencies_hz.len(), times_s.len());
116        if power.dim() != expected {
117            return Err(MethodError::ShapeMismatch(format!(
118                "power has shape {:?}, axes imply {expected:?}",
119                power.dim()
120            )));
121        }
122        Ok(Self {
123            power,
124            frequencies_hz,
125            times_s,
126            labels,
127            sampling_rate: None,
128            reliable: None,
129        })
130    }
131
132    pub fn power(&self) -> ArrayView3<'_, f64> {
133        self.power.view()
134    }
135
136    pub fn frequencies_hz(&self) -> &[f64] {
137        &self.frequencies_hz
138    }
139
140    pub fn times_s(&self) -> &[f64] {
141        &self.times_s
142    }
143
144    pub fn labels(&self) -> &[String] {
145        &self.labels
146    }
147
148    /// Set when there is one time column per sample at `i / fs` (CWT, Hilbert spectrum), which
149    /// makes the map epochable. STFT frames are not per-sample and leave this unset.
150    pub fn with_sampling_rate(mut self, sampling_rate: SamplingRate) -> Result<Self, MethodError> {
151        let fs = sampling_rate.hz();
152        let aligned = self
153            .times_s
154            .iter()
155            .enumerate()
156            .all(|(i, t)| (t - i as f64 / fs).abs() <= 1e-9 * t.abs().max(1.0));
157        if !aligned {
158            return Err(MethodError::ShapeMismatch(
159                "times are not one per sample at i / fs".into(),
160            ));
161        }
162        self.sampling_rate = Some(sampling_rate);
163        Ok(self)
164    }
165
166    /// Frequency × time mask: false where edge effects make the value unreliable (for example
167    /// inside a wavelet's cone of influence).
168    pub fn with_reliable(mut self, reliable: Array2<bool>) -> Result<Self, MethodError> {
169        let expected = (self.frequencies_hz.len(), self.times_s.len());
170        if reliable.dim() != expected {
171            return Err(MethodError::ShapeMismatch(format!(
172                "reliable mask {:?}, expected {expected:?}",
173                reliable.dim()
174            )));
175        }
176        self.reliable = Some(reliable);
177        Ok(self)
178    }
179
180    pub fn sampling_rate(&self) -> Option<SamplingRate> {
181        self.sampling_rate
182    }
183
184    pub fn reliable(&self) -> Option<ArrayView2<'_, bool>> {
185        self.reliable.as_ref().map(|r| r.view())
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn grid_bins_are_half_open() {
195        let g = FrequencyGrid::new(0.0, 1.0, 4, GridScale::Linear).unwrap();
196        assert_eq!(g.edges(), [0.0, 0.25, 0.5, 0.75, 1.0]);
197        assert_eq!(g.centres(), [0.125, 0.375, 0.625, 0.875]);
198        assert_eq!(
199            (
200                g.bin(0.0),
201                g.bin(0.25),
202                g.bin(0.999),
203                g.bin(1.0),
204                g.bin(-0.1),
205                g.bin(f64::NAN)
206            ),
207            (Some(0), Some(1), Some(3), None, None, None)
208        );
209        assert_eq!(g.n_bins(), 4);
210        assert_eq!(g.scale(), GridScale::Linear);
211        let uneven = FrequencyGrid::new(0.1, 0.7, 3, GridScale::Linear).unwrap();
212        assert!(
213            uneven
214                .centres()
215                .iter()
216                .zip(uneven.edges().windows(2))
217                .all(|(c, w)| (c - (w[0] + w[1]) / 2.0).abs() < 1e-15)
218        );
219
220        let log = FrequencyGrid::new(0.01, 1.0, 2, GridScale::Log).unwrap();
221        assert!((log.edges()[1] - 0.1).abs() < 1e-12);
222        assert!((log.centres()[0] - (0.01f64 * 0.1).sqrt()).abs() < 1e-12);
223
224        assert!(FrequencyGrid::new(0.0, 1.0, 0, GridScale::Linear).is_err());
225        assert!(FrequencyGrid::new(1.0, 1.0, 3, GridScale::Linear).is_err());
226        assert!(
227            FrequencyGrid::new(0.0, 1.0, 3, GridScale::Log).is_err(),
228            "log grids need f_min > 0"
229        );
230        assert!(FrequencyGrid::new(f64::NAN, 1.0, 3, GridScale::Linear).is_err());
231    }
232
233    #[test]
234    fn rejects_axis_mismatch() {
235        let power = Array3::zeros((1, 3, 2));
236        let err = TimeFrequencyMap::new(power, vec![0.0, 1.0], vec![0.0, 1.0], vec!["a".into()]);
237        assert!(matches!(err, Err(MethodError::ShapeMismatch(_))));
238    }
239
240    #[test]
241    fn optional_sampling_rate_and_reliability_mask_are_checked() {
242        let map = TimeFrequencyMap::new(
243            Array3::zeros((1, 2, 3)),
244            vec![1.0, 2.0],
245            vec![0.0, 0.5, 1.0],
246            vec!["a".into()],
247        )
248        .unwrap();
249        assert!(map.sampling_rate().is_none() && map.reliable().is_none());
250        let fs = SamplingRate::new(2.0).unwrap();
251        let with = map
252            .clone()
253            .with_sampling_rate(fs)
254            .unwrap()
255            .with_reliable(Array2::from_elem((2, 3), true))
256            .unwrap();
257        assert_eq!(with.sampling_rate(), Some(fs));
258        assert_eq!(with.reliable().unwrap().dim(), (2, 3));
259        assert!(
260            map.clone()
261                .with_sampling_rate(SamplingRate::new(4.0).unwrap())
262                .is_err(),
263            "frames are not at i / 4 Hz"
264        );
265        assert!(map.with_reliable(Array2::from_elem((3, 2), true)).is_err());
266    }
267
268    #[test]
269    fn accumulated_rounding_in_long_time_axes_is_accepted() {
270        // Times built by repeated addition drift by ~1e-7 s after 1000 s; that is still per-sample.
271        let n = 1250;
272        let mut t = 0.0;
273        let times: Vec<f64> = (0..n)
274            .map(|_| {
275                let now = t;
276                t += 0.8;
277                now
278            })
279            .collect();
280        let mut times = times;
281        times[n - 1] += 1e-7;
282        let map =
283            TimeFrequencyMap::new(Array3::zeros((1, 1, n)), vec![0.1], times, vec!["a".into()])
284                .unwrap();
285        assert!(
286            map.with_sampling_rate(SamplingRate::new(1.25).unwrap())
287                .is_ok()
288        );
289    }
290
291    #[test]
292    fn accepts_matching_axes() {
293        let power = Array3::zeros((1, 2, 3));
294        let map =
295            TimeFrequencyMap::new(power, vec![0.0, 1.0], vec![0.0, 1.0, 2.0], vec!["a".into()])
296                .unwrap();
297        assert_eq!(map.frequencies_hz(), [0.0, 1.0]);
298        assert_eq!(map.times_s(), [0.0, 1.0, 2.0]);
299        assert_eq!(map.labels(), ["a"]);
300        assert_eq!(map.power().dim(), (1, 2, 3));
301    }
302}
modewise_spectral/lib.rs

modewise_spectral/
lib.rs
