audio_lab/signal/data/
complex_data.rs

1use ndarray::prelude::*;
2use num::complex::Complex64;
3
4#[cfg(feature = "serde")]
5use super::check_channel_labels;
6use super::{DataError, XAxis};
7
8#[derive(Clone, Debug)]
9#[cfg_attr(
10    feature = "serde",
11    derive(serde::Serialize, serde::Deserialize),
12    serde(try_from = "ComplexDataFields")
13)]
14pub struct ComplexData {
15    x_axis: XAxis,
16    y_data: Array2<Complex64>,
17    comment: Option<String>,
18    channel_labels: Vec<Option<String>>,
19}
20
21/// Deserialized fields, checked by `TryFrom` before they become a `ComplexData`.
22#[cfg(feature = "serde")]
23#[derive(serde::Deserialize)]
24struct ComplexDataFields {
25    x_axis: XAxis,
26    y_data: Array2<Complex64>,
27    comment: Option<String>,
28    channel_labels: Vec<Option<String>>,
29}
30
31#[cfg(feature = "serde")]
32impl TryFrom<ComplexDataFields> for ComplexData {
33    type Error = DataError;
34
35    fn try_from(fields: ComplexDataFields) -> Result<Self, Self::Error> {
36        if fields.y_data.ncols() != fields.x_axis.len() {
37            return Err(DataError::NotMatching);
38        }
39        check_channel_labels(&fields.channel_labels, fields.y_data.nrows())?;
40        // Struct literal on purpose: a new field fails to compile here instead
41        // of being silently dropped when reading.
42        Ok(Self {
43            x_axis: fields.x_axis,
44            y_data: fields.y_data,
45            comment: fields.comment,
46            channel_labels: fields.channel_labels,
47        })
48    }
49}
50
51impl ComplexData {
52    pub fn new(x_data: Array1<f64>, y_data: Array2<Complex64>) -> Result<Self, DataError> {
53        if y_data.ncols() != x_data.len() {
54            return Err(DataError::NotMatching);
55        }
56        Self::from_axis(XAxis::explicit(x_data)?, y_data)
57    }
58
59    /// Create data from an existing axis. Only the length is checked, since
60    /// an [`XAxis`] is always valid.
61    pub fn from_axis(x_axis: XAxis, y_data: Array2<Complex64>) -> Result<Self, DataError> {
62        if y_data.ncols() != x_axis.len() {
63            return Err(DataError::NotMatching);
64        }
65        let num_channels = y_data.nrows();
66        Ok(Self {
67            x_axis,
68            y_data,
69            comment: None,
70            channel_labels: vec![None; num_channels],
71        })
72    }
73
74    pub fn zeros(num_channels: usize, num_data_points: usize) -> Self {
75        Self {
76            x_axis: XAxis::indices(num_data_points),
77            y_data: Array2::zeros((num_channels, num_data_points)),
78            comment: None,
79            channel_labels: vec![None; num_channels],
80        }
81    }
82
83    pub fn from_x_data(x_data: Array1<f64>, num_channels: usize) -> Result<Self, DataError> {
84        let x_axis = XAxis::explicit(x_data)?;
85        Ok(Self {
86            y_data: Array2::zeros((num_channels, x_axis.len())),
87            x_axis,
88            comment: None,
89            channel_labels: vec![None; num_channels],
90        })
91    }
92
93    pub fn from_y_data(y_data: Array2<Complex64>) -> Self {
94        let num_channels = y_data.nrows();
95        Self {
96            x_axis: XAxis::indices(y_data.ncols()),
97            y_data,
98            comment: None,
99            channel_labels: vec![None; num_channels],
100        }
101    }
102
103    pub fn comment(&self) -> Option<&str> {
104        self.comment.as_deref()
105    }
106
107    pub fn set_comment(&mut self, comment: Option<&str>) {
108        self.comment = comment.map(Into::into);
109    }
110
111    pub fn channel_label(&self, channel: usize) -> Option<&str> {
112        self.channel_labels.get(channel).and_then(|s| s.as_deref())
113    }
114
115    pub fn set_channel_label(&mut self, channel: usize, label: Option<&str>) {
116        if let Some(slot) = self.channel_labels.get_mut(channel) {
117            *slot = label.map(Into::into);
118        }
119    }
120
121    /// All channel labels, one per channel.
122    #[cfg(feature = "serde")]
123    pub(crate) fn channel_labels(&self) -> &[Option<String>] {
124        &self.channel_labels
125    }
126
127    /// Replace all channel labels. There must be one per channel.
128    #[cfg(feature = "serde")]
129    pub(crate) fn set_channel_labels(
130        &mut self,
131        channel_labels: Vec<Option<String>>,
132    ) -> Result<(), DataError> {
133        check_channel_labels(&channel_labels, self.num_channels())?;
134        self.channel_labels = channel_labels;
135        Ok(())
136    }
137
138    pub fn num_channels(&self) -> usize {
139        self.y_data.nrows()
140    }
141
142    pub fn num_data_points(&self) -> usize {
143        self.y_data.ncols()
144    }
145
146    /// The x axis. Prefer this over [`x_data`](Self::x_data) for single
147    /// values ([`XAxis::get`]) or iteration, because it never allocates.
148    pub fn x_axis(&self) -> &XAxis {
149        &self.x_axis
150    }
151
152    /// The x values as an array. Borrowed for an explicit axis, computed
153    /// (one allocation) for a uniform one.
154    pub fn x_data(&self) -> CowArray<'_, f64, Ix1> {
155        self.x_axis.to_cow_array()
156    }
157
158    /// Replace the x axis. It must have one value per data point.
159    pub fn set_x_axis(&mut self, x_axis: XAxis) -> Result<(), DataError> {
160        if x_axis.len() != self.num_data_points() {
161            return Err(DataError::NotMatching);
162        }
163        self.x_axis = x_axis;
164        Ok(())
165    }
166
167    pub fn y_data(&self) -> ArrayView2<'_, Complex64> {
168        self.y_data.view()
169    }
170
171    pub fn y_data_mut(&mut self) -> ArrayViewMut2<'_, Complex64> {
172        self.y_data.view_mut()
173    }
174
175    /// Consume the data and return the owned y values without copying.
176    pub fn into_y_data(self) -> Array2<Complex64> {
177        self.y_data
178    }
179
180    pub fn channel(&self, ch: usize) -> ArrayView1<'_, Complex64> {
181        self.y_data.row(ch)
182    }
183
184    pub fn channel_mut(&mut self, ch: usize) -> ArrayViewMut1<'_, Complex64> {
185        self.y_data.row_mut(ch)
186    }
187
188    pub fn channel_iter(&self) -> ndarray::iter::AxisIter<'_, Complex64, Ix1> {
189        self.y_data.outer_iter()
190    }
191
192    pub fn channel_iter_mut(&mut self) -> ndarray::iter::AxisIterMut<'_, Complex64, Ix1> {
193        self.y_data.outer_iter_mut()
194    }
195
196    pub fn iter(&self) -> ndarray::iter::Iter<'_, Complex64, Ix2> {
197        self.y_data.iter()
198    }
199
200    pub fn iter_mut(&mut self) -> ndarray::iter::IterMut<'_, Complex64, Ix2> {
201        self.y_data.iter_mut()
202    }
203
204    /// Index of the x value closest to `x`. On a tie the lower index wins.
205    ///
206    /// O(1) for a uniform axis, a binary search otherwise.
207    pub fn nearest_x_index(&self, x: f64) -> usize {
208        self.x_axis.nearest_index(x)
209    }
210}
211
212impl std::fmt::Display for ComplexData {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        let base_info = format!(
215            "Complex data with {} channels and {} data points.",
216            self.num_channels(),
217            self.num_data_points()
218        );
219
220        if let Some(comment) = self.comment() {
221            write!(f, "{}\nComment: {:?}", base_info, comment)
222        } else {
223            write!(f, "{}", base_info)
224        }
225    }
226}
227
228impl PartialEq for ComplexData {
229    fn eq(&self, other: &Self) -> bool {
230        self.x_axis == other.x_axis && self.y_data == other.y_data
231    }
232}
233
234impl approx::AbsDiffEq for ComplexData {
235    type Epsilon = f64;
236
237    fn default_epsilon() -> Self::Epsilon {
238        f64::EPSILON
239    }
240
241    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
242        if !self.x_axis.abs_diff_eq(&other.x_axis, epsilon) {
243            return false;
244        }
245
246        // Compare y_data
247        if !self.y_data().abs_diff_eq(&other.y_data(), epsilon) {
248            return false;
249        }
250
251        true
252    }
253}
254
255// Immutable IntoIterator
256impl<'a> IntoIterator for &'a ComplexData {
257    type Item = &'a Complex64;
258    type IntoIter = ndarray::iter::Iter<'a, Complex64, ndarray::Ix2>;
259
260    fn into_iter(self) -> Self::IntoIter {
261        self.iter()
262    }
263}
264
265// Mutable IntoIterator
266impl<'a> IntoIterator for &'a mut ComplexData {
267    type Item = &'a mut Complex64;
268    type IntoIter = ndarray::iter::IterMut<'a, Complex64, ndarray::Ix2>;
269
270    fn into_iter(self) -> Self::IntoIter {
271        self.iter_mut()
272    }
273}
274
275// Owned IntoIterator
276impl IntoIterator for ComplexData {
277    type Item = Complex64;
278    type IntoIter = ndarray::iter::IntoIter<Complex64, Ix2>;
279
280    fn into_iter(self) -> Self::IntoIter {
281        self.y_data.into_iter()
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    #[test]
290    fn complex_data_errors() {
291        let error = ComplexData::new(
292            array![0.0, 1.0, 2.0, 5.0],
293            arr2(&[[Complex64::from(0.0); 3]; 2]),
294        )
295        .err();
296        assert!(matches!(error, Some(DataError::NotMatching)));
297
298        let error =
299            ComplexData::new(array![0.0, 0.0, 0.0], arr2(&[[Complex64::from(0.0); 3]; 2])).err();
300        assert!(matches!(error, Some(DataError::NotIncreasing)));
301
302        let error = ComplexData::from_x_data(array![0.0, 0.0, 0.0], 3).err();
303        assert!(matches!(error, Some(DataError::NotIncreasing)));
304    }
305
306    #[test]
307    fn nearest_x_index_matches_linear_scan() {
308        let data = ComplexData::from_x_data(array![0.0, 1.0, 2.5, 4.0, 10.0], 1).unwrap();
309        let linear_scan = |x: f64| {
310            data.x_data()
311                .iter()
312                .enumerate()
313                .min_by(|(_, a), (_, b)| {
314                    (*a - x)
315                        .abs()
316                        .partial_cmp(&(*b - x).abs())
317                        .unwrap_or(std::cmp::Ordering::Equal)
318                })
319                .map(|(index, _)| index)
320                .unwrap()
321        };
322
323        // Exact hits, midpoints (ties go to the lower index), values outside
324        // the range, and NaN.
325        let queries = [
326            -5.0,
327            0.0,
328            0.5,
329            0.6,
330            1.75,
331            2.5,
332            3.25,
333            3.3,
334            7.0,
335            7.1,
336            10.0,
337            42.0,
338            f64::NAN,
339        ];
340        for x in queries {
341            assert_eq!(data.nearest_x_index(x), linear_scan(x), "x = {x}");
342        }
343
344        assert_eq!(ComplexData::zeros(1, 0).nearest_x_index(1.0), 0);
345    }
346
347    #[test]
348    fn test_complex_data() -> Result<(), DataError> {
349        let data = ComplexData::zeros(2, 4);
350        assert_eq!(data.x_data(), arr1(&[0.0, 1.0, 2.0, 3.0]));
351        assert_eq!(data.y_data(), arr2(&[[Complex64::from(0.0); 4]; 2]),);
352
353        let data = ComplexData::from_x_data(arr1(&[0.0, 1.0, 2.0, 5.0]), 2).unwrap();
354        assert_eq!(data.x_data(), arr1(&[0.0, 1.0, 2.0, 5.0]));
355        assert_eq!(data.y_data(), arr2(&[[Complex64::from(0.0); 4]; 2]),);
356
357        let mut i = 0;
358        let mut data = ComplexData::new(
359            Array1::from_iter((0..10).map(|v| v as f64 / 10.0)),
360            Array2::from_shape_fn((2, 10), |_| {
361                let value = Complex64::from(i as f64);
362                i += 1;
363                value
364            }),
365        )?;
366
367        assert_eq!(data.comment(), None);
368
369        data.set_comment(Some("This is just test data"));
370
371        assert_eq!(data.comment(), Some("This is just test data"));
372
373        assert_eq!(data.num_channels(), 2);
374        assert_eq!(data.num_data_points(), 10);
375
376        assert_eq!(
377            data.x_data(),
378            arr1(&[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
379        );
380
381        assert_eq!(
382            data.y_data(),
383            arr2(&[
384                [
385                    Complex64::from(0.0),
386                    Complex64::from(1.0),
387                    Complex64::from(2.0),
388                    Complex64::from(3.0),
389                    Complex64::from(4.0),
390                    Complex64::from(5.0),
391                    Complex64::from(6.0),
392                    Complex64::from(7.0),
393                    Complex64::from(8.0),
394                    Complex64::from(9.0)
395                ],
396                [
397                    Complex64::from(10.0),
398                    Complex64::from(11.0),
399                    Complex64::from(12.0),
400                    Complex64::from(13.0),
401                    Complex64::from(14.0),
402                    Complex64::from(15.0),
403                    Complex64::from(16.0),
404                    Complex64::from(17.0),
405                    Complex64::from(18.0),
406                    Complex64::from(19.0),
407                ]
408            ])
409        );
410
411        assert_eq!(
412            data.y_data_mut(),
413            arr2(&[
414                [
415                    Complex64::from(0.0),
416                    Complex64::from(1.0),
417                    Complex64::from(2.0),
418                    Complex64::from(3.0),
419                    Complex64::from(4.0),
420                    Complex64::from(5.0),
421                    Complex64::from(6.0),
422                    Complex64::from(7.0),
423                    Complex64::from(8.0),
424                    Complex64::from(9.0)
425                ],
426                [
427                    Complex64::from(10.0),
428                    Complex64::from(11.0),
429                    Complex64::from(12.0),
430                    Complex64::from(13.0),
431                    Complex64::from(14.0),
432                    Complex64::from(15.0),
433                    Complex64::from(16.0),
434                    Complex64::from(17.0),
435                    Complex64::from(18.0),
436                    Complex64::from(19.0),
437                ]
438            ])
439        );
440
441        let a: Vec<Complex64> = data.clone().into_iter().collect();
442        assert_eq!(
443            a,
444            (0..20)
445                .map(|v| Complex64::from(v as f64))
446                .collect::<Vec<_>>()
447        );
448
449        assert_eq!(
450            data.channel(0),
451            arr1(&[
452                Complex64::from(0.0),
453                Complex64::from(1.0),
454                Complex64::from(2.0),
455                Complex64::from(3.0),
456                Complex64::from(4.0),
457                Complex64::from(5.0),
458                Complex64::from(6.0),
459                Complex64::from(7.0),
460                Complex64::from(8.0),
461                Complex64::from(9.0)
462            ])
463        );
464
465        assert_eq!(
466            data.channel_mut(0),
467            arr1(&[
468                Complex64::from(0.0),
469                Complex64::from(1.0),
470                Complex64::from(2.0),
471                Complex64::from(3.0),
472                Complex64::from(4.0),
473                Complex64::from(5.0),
474                Complex64::from(6.0),
475                Complex64::from(7.0),
476                Complex64::from(8.0),
477                Complex64::from(9.0)
478            ])
479        );
480
481        assert_eq!(
482            data.channel(1),
483            arr1(&[
484                Complex64::from(10.0),
485                Complex64::from(11.0),
486                Complex64::from(12.0),
487                Complex64::from(13.0),
488                Complex64::from(14.0),
489                Complex64::from(15.0),
490                Complex64::from(16.0),
491                Complex64::from(17.0),
492                Complex64::from(18.0),
493                Complex64::from(19.0)
494            ])
495        );
496
497        assert_eq!(
498            data.channel_mut(1),
499            arr1(&[
500                Complex64::from(10.0),
501                Complex64::from(11.0),
502                Complex64::from(12.0),
503                Complex64::from(13.0),
504                Complex64::from(14.0),
505                Complex64::from(15.0),
506                Complex64::from(16.0),
507                Complex64::from(17.0),
508                Complex64::from(18.0),
509                Complex64::from(19.0)
510            ])
511        );
512
513        for (i, ch) in data.channel_iter().enumerate() {
514            match i {
515                0 => assert_eq!(
516                    ch,
517                    arr1(&[
518                        Complex64::from(0.0),
519                        Complex64::from(1.0),
520                        Complex64::from(2.0),
521                        Complex64::from(3.0),
522                        Complex64::from(4.0),
523                        Complex64::from(5.0),
524                        Complex64::from(6.0),
525                        Complex64::from(7.0),
526                        Complex64::from(8.0),
527                        Complex64::from(9.0)
528                    ])
529                ),
530                1 => assert_eq!(
531                    ch,
532                    arr1(&[
533                        Complex64::from(10.0),
534                        Complex64::from(11.0),
535                        Complex64::from(12.0),
536                        Complex64::from(13.0),
537                        Complex64::from(14.0),
538                        Complex64::from(15.0),
539                        Complex64::from(16.0),
540                        Complex64::from(17.0),
541                        Complex64::from(18.0),
542                        Complex64::from(19.0)
543                    ])
544                ),
545                _ => panic!("Should only have two channels"),
546            }
547        }
548
549        for (i, ch) in data.channel_iter_mut().enumerate() {
550            match i {
551                0 => assert_eq!(
552                    ch,
553                    arr1(&[
554                        Complex64::from(0.0),
555                        Complex64::from(1.0),
556                        Complex64::from(2.0),
557                        Complex64::from(3.0),
558                        Complex64::from(4.0),
559                        Complex64::from(5.0),
560                        Complex64::from(6.0),
561                        Complex64::from(7.0),
562                        Complex64::from(8.0),
563                        Complex64::from(9.0)
564                    ])
565                ),
566                1 => assert_eq!(
567                    ch,
568                    arr1(&[
569                        Complex64::from(10.0),
570                        Complex64::from(11.0),
571                        Complex64::from(12.0),
572                        Complex64::from(13.0),
573                        Complex64::from(14.0),
574                        Complex64::from(15.0),
575                        Complex64::from(16.0),
576                        Complex64::from(17.0),
577                        Complex64::from(18.0),
578                        Complex64::from(19.0)
579                    ])
580                ),
581                _ => panic!("Should only have two channels"),
582            }
583        }
584
585        Ok(())
586    }
587}
audio_lab/signal/data/complex_data.rs

audio_lab/signal/data/
complex_data.rs
