1use modewise_method::{MethodDescriptor, MethodError};
2use modewise_signals::{MultichannelSeries, SamplingRate};
3use ndarray::{Array2, Array3, ArrayView2, ArrayView3, Axis};
4
5mod loadings;
6
7pub use loadings::Loadings;
8
9#[derive(Debug, Clone, PartialEq)]
12pub struct InstantaneousModes {
13 amplitude: Array3<f64>,
14 frequency_hz: Array3<f64>,
15 centre_frequencies_hz: Vec<f64>,
16 sampling_rate: SamplingRate,
17 labels: Vec<String>,
18}
19
20impl InstantaneousModes {
21 pub fn new(
22 amplitude: Array3<f64>,
23 frequency_hz: Array3<f64>,
24 centre_frequencies_hz: Vec<f64>,
25 sampling_rate: SamplingRate,
26 labels: Vec<String>,
27 ) -> Result<Self, MethodError> {
28 let (n_modes, n_channels, _) = amplitude.dim();
29 if frequency_hz.dim() != amplitude.dim() {
30 return Err(MethodError::ShapeMismatch(format!(
31 "amplitude {:?} and frequency {:?} differ",
32 amplitude.dim(),
33 frequency_hz.dim()
34 )));
35 }
36 if centre_frequencies_hz.len() != n_modes || labels.len() != n_channels {
37 return Err(MethodError::ShapeMismatch(format!(
38 "{n_modes} modes and {n_channels} channels, but {} centre frequencies and {} labels",
39 centre_frequencies_hz.len(),
40 labels.len()
41 )));
42 }
43 Ok(Self {
44 amplitude,
45 frequency_hz,
46 centre_frequencies_hz,
47 sampling_rate,
48 labels,
49 })
50 }
51
52 pub fn amplitude(&self) -> ArrayView3<'_, f64> {
53 self.amplitude.view()
54 }
55
56 pub fn frequency_hz(&self) -> ArrayView3<'_, f64> {
57 self.frequency_hz.view()
58 }
59
60 pub fn centre_frequencies_hz(&self) -> &[f64] {
61 &self.centre_frequencies_hz
62 }
63
64 pub fn sampling_rate(&self) -> SamplingRate {
65 self.sampling_rate
66 }
67
68 pub fn labels(&self) -> &[String] {
69 &self.labels
70 }
71
72 pub fn n_samples(&self) -> usize {
73 self.amplitude.dim().2
74 }
75}
76
77fn median(mut values: Vec<f64>) -> f64 {
78 if values.is_empty() {
79 return f64::NAN;
80 }
81 values.sort_by(f64::total_cmp);
82 let mid = values.len() / 2;
83 if values.len() % 2 == 1 {
84 values[mid]
85 } else {
86 (values[mid - 1] + values[mid]) / 2.0
87 }
88}
89
90pub trait ModeDecomposer {
92 fn descriptor(&self) -> MethodDescriptor;
93 fn decompose(&self, signal: &MultichannelSeries) -> Result<Decomposition, MethodError>;
94}
95
96#[derive(Debug, Clone, PartialEq)]
98pub struct Decomposition {
99 modes: Array3<f64>,
100 centre_frequencies_hz: Vec<f64>,
101 sampling_rate: SamplingRate,
102 labels: Vec<String>,
103 aligned_across_channels: bool,
104 channel_centre_frequencies_hz: Option<Array2<f64>>,
106}
107
108impl Decomposition {
109 pub fn new(
110 modes: Array3<f64>,
111 centre_frequencies_hz: Vec<f64>,
112 sampling_rate: SamplingRate,
113 labels: Vec<String>,
114 aligned_across_channels: bool,
115 ) -> Result<Self, MethodError> {
116 let (n_modes, n_channels, _) = modes.dim();
117 if centre_frequencies_hz.len() != n_modes {
118 return Err(MethodError::ShapeMismatch(format!(
119 "{n_modes} modes but {} centre frequencies",
120 centre_frequencies_hz.len()
121 )));
122 }
123 if labels.len() != n_channels {
124 return Err(MethodError::ShapeMismatch(format!(
125 "{n_channels} channels but {} labels",
126 labels.len()
127 )));
128 }
129 Ok(Self {
130 modes,
131 centre_frequencies_hz,
132 sampling_rate,
133 labels,
134 aligned_across_channels,
135 channel_centre_frequencies_hz: None,
136 })
137 }
138
139 pub fn per_channel(
144 modes: Array3<f64>,
145 channel_centre_frequencies_hz: Array2<f64>,
146 sampling_rate: SamplingRate,
147 labels: Vec<String>,
148 ) -> Result<Self, MethodError> {
149 let centres = channel_centre_frequencies_hz
150 .rows()
151 .into_iter()
152 .map(|row| median(row.iter().copied().filter(|v| v.is_finite()).collect()))
153 .collect();
154 Self::new(modes, centres, sampling_rate, labels, false)?
155 .with_channel_centre_frequencies(channel_centre_frequencies_hz)
156 }
157
158 pub fn with_channel_centre_frequencies(
160 mut self,
161 frequencies_hz: Array2<f64>,
162 ) -> Result<Self, MethodError> {
163 let (n_modes, n_channels, _) = self.modes.dim();
164 if frequencies_hz.dim() != (n_modes, n_channels) {
165 return Err(MethodError::ShapeMismatch(format!(
166 "{n_modes} modes and {n_channels} channels, but channel centre frequencies are {:?}",
167 frequencies_hz.dim()
168 )));
169 }
170 self.channel_centre_frequencies_hz = Some(frequencies_hz);
171 Ok(self)
172 }
173
174 pub fn channel_centre_frequencies_hz(&self) -> Option<ArrayView2<'_, f64>> {
176 self.channel_centre_frequencies_hz
177 .as_ref()
178 .map(|f| f.view())
179 }
180
181 pub fn modes(&self) -> ArrayView3<'_, f64> {
182 self.modes.view()
183 }
184
185 pub fn mode(&self, index: usize) -> ArrayView2<'_, f64> {
186 self.modes.index_axis(Axis(0), index)
187 }
188
189 pub fn centre_frequencies_hz(&self) -> &[f64] {
190 &self.centre_frequencies_hz
191 }
192
193 pub fn sampling_rate(&self) -> SamplingRate {
194 self.sampling_rate
195 }
196
197 pub fn labels(&self) -> &[String] {
198 &self.labels
199 }
200
201 pub fn aligned_across_channels(&self) -> bool {
203 self.aligned_across_channels
204 }
205
206 pub fn n_modes(&self) -> usize {
207 self.modes.len_of(Axis(0))
208 }
209
210 pub fn reconstruct(&self, modes: &[usize]) -> Result<MultichannelSeries, MethodError> {
212 if modes.is_empty() {
213 return Err(MethodError::invalid("modes", "choose at least one mode"));
214 }
215 let (_, n_channels, n_samples) = self.modes.dim();
216 let mut sum = Array2::zeros((n_channels, n_samples));
217 for (i, &k) in modes.iter().enumerate() {
218 if k >= self.n_modes() {
219 return Err(MethodError::invalid(
220 "modes",
221 format!(
222 "mode {k} does not exist; there are {} modes",
223 self.n_modes()
224 ),
225 ));
226 }
227 if modes[..i].contains(&k) {
228 return Err(MethodError::invalid(
229 "modes",
230 format!("mode {k} is listed twice"),
231 ));
232 }
233 sum += &self.mode(k);
234 }
235 MultichannelSeries::new(sum, self.sampling_rate, self.labels.clone())
236 .map_err(|e| MethodError::ShapeMismatch(e.to_string()))
237 }
238}
239
240#[cfg(test)]
241mod tests {
242 use super::*;
243
244 fn fs() -> SamplingRate {
245 SamplingRate::new(1.0).unwrap()
246 }
247
248 #[test]
249 fn rejects_frequency_count_mismatch() {
250 let modes = Array3::zeros((2, 1, 4));
251 let err = Decomposition::new(modes, vec![0.1], fs(), vec!["a".into()], true);
252 assert!(matches!(err, Err(MethodError::ShapeMismatch(_))));
253 }
254
255 #[test]
256 fn rejects_label_count_mismatch() {
257 let modes = Array3::zeros((1, 2, 4));
258 let err = Decomposition::new(modes, vec![0.1], fs(), vec!["a".into()], true);
259 assert!(matches!(err, Err(MethodError::ShapeMismatch(_))));
260 }
261
262 #[test]
263 fn reconstruct_sums_chosen_modes() {
264 let mut modes = Array3::zeros((3, 1, 2));
265 modes[[0, 0, 0]] = 1.0;
266 modes[[1, 0, 0]] = 2.0;
267 modes[[2, 0, 1]] = 4.0;
268 let d =
269 Decomposition::new(modes, vec![0.1, 0.2, 0.3], fs(), vec!["a".into()], true).unwrap();
270 let series = d.reconstruct(&[0, 2]).unwrap();
271 assert_eq!(series.data().row(0).to_vec(), [1.0, 4.0]);
272 assert_eq!(series.labels(), ["a"]);
273 }
274
275 #[test]
276 fn reconstruct_rejects_bad_selection() {
277 let d = Decomposition::new(
278 Array3::zeros((2, 1, 2)),
279 vec![0.1, 0.2],
280 fs(),
281 vec!["a".into()],
282 true,
283 )
284 .unwrap();
285 for selection in [&[][..], &[2], &[1, 1]] {
286 assert!(d.reconstruct(selection).is_err(), "accepted {selection:?}");
287 }
288 }
289
290 #[test]
291 fn getters_return_what_was_stored() {
292 let d = Decomposition::new(
293 Array3::zeros((2, 1, 3)),
294 vec![0.1, 0.2],
295 fs(),
296 vec!["amy".into()],
297 false,
298 )
299 .unwrap();
300 assert_eq!(d.centre_frequencies_hz(), [0.1, 0.2]);
301 assert_eq!(d.labels(), ["amy"]);
302 assert!(!d.aligned_across_channels());
303 assert_eq!(d.sampling_rate(), fs());
304 let aligned = Decomposition::new(
305 Array3::zeros((1, 1, 3)),
306 vec![0.1],
307 fs(),
308 vec!["a".into()],
309 true,
310 )
311 .unwrap();
312 assert!(aligned.aligned_across_channels());
313 }
314
315 #[test]
316 fn instantaneous_modes_check_shapes() {
317 let ok = InstantaneousModes::new(
318 Array3::zeros((2, 1, 4)),
319 Array3::zeros((2, 1, 4)),
320 vec![0.1, 0.2],
321 fs(),
322 vec!["a".into()],
323 )
324 .unwrap();
325 assert_eq!(
326 (ok.n_samples(), ok.centre_frequencies_hz(), ok.labels()),
327 (4, [0.1, 0.2].as_slice(), ["a".to_string()].as_slice())
328 );
329 assert_eq!(ok.sampling_rate(), fs());
330 assert_eq!(ok.frequency_hz().dim(), ok.amplitude().dim());
331 assert!(
332 InstantaneousModes::new(
333 Array3::zeros((2, 1, 4)),
334 Array3::zeros((2, 1, 5)),
335 vec![0.1, 0.2],
336 fs(),
337 vec!["a".into()]
338 )
339 .is_err()
340 );
341 assert!(
342 InstantaneousModes::new(
343 Array3::zeros((2, 1, 4)),
344 Array3::zeros((2, 1, 4)),
345 vec![0.1],
346 fs(),
347 vec!["a".into()]
348 )
349 .is_err()
350 );
351 assert!(
352 InstantaneousModes::new(
353 Array3::zeros((2, 1, 4)),
354 Array3::zeros((2, 1, 4)),
355 vec![0.1, 0.2],
356 fs(),
357 vec![]
358 )
359 .is_err()
360 );
361 }
362
363 #[test]
364 fn exposes_single_mode() {
365 let mut modes = Array3::zeros((2, 1, 3));
366 modes[[1, 0, 2]] = 5.0;
367 let d = Decomposition::new(modes, vec![0.1, 0.2], fs(), vec!["a".into()], true).unwrap();
368 assert_eq!(d.n_modes(), 2);
369 assert_eq!(d.mode(1)[[0, 2]], 5.0);
370 }
371}
372
373#[cfg(test)]
374mod per_channel {
375 use super::*;
376
377 #[test]
378 fn mode_frequencies_are_channel_medians_ignoring_missing_modes() {
379 let fs = SamplingRate::new(1.0).unwrap();
380 let channel = ndarray::array![
381 [0.1, 0.3, 0.2],
382 [0.05, f64::NAN, 0.25],
383 [f64::NAN, f64::NAN, f64::NAN]
384 ];
385 let labels = vec!["a".into(), "b".into(), "c".into()];
386 let d = Decomposition::per_channel(
387 Array3::zeros((3, 3, 4)),
388 channel.clone(),
389 fs,
390 labels.clone(),
391 )
392 .unwrap();
393 assert_eq!(d.centre_frequencies_hz()[..2], [0.2, 0.15]);
394 assert!(d.centre_frequencies_hz()[2].is_nan());
395 assert!(!d.aligned_across_channels());
396 let stored = d.channel_centre_frequencies_hz().unwrap();
397 assert!(
398 stored
399 .iter()
400 .zip(&channel)
401 .all(|(a, b)| a.to_bits() == b.to_bits())
402 );
403 assert!(
404 Decomposition::per_channel(Array3::zeros((3, 2, 4)), channel, fs, labels[..2].to_vec())
405 .is_err()
406 );
407 let plain = Decomposition::new(
408 Array3::zeros((1, 1, 4)),
409 vec![0.1],
410 fs,
411 vec!["a".into()],
412 true,
413 )
414 .unwrap();
415 assert!(plain.channel_centre_frequencies_hz().is_none());
416 }
417}