sci_rs/signal/wave/
mod.rs

1use nalgebra::RealField;
2use ndarray::{Array, ArrayBase, Data, Dimension, RawData};
3
4/// """
5/// Return a periodic square-wave waveform.
6///
7/// The square wave has a period ``2*pi``, has value +1 from 0 to
8/// ``2*pi*duty`` and -1 from ``2*pi*duty`` to ``2*pi``. `duty` must be in
9/// the interval [0,1].
10///
11/// Note that this is not band-limited.  It produces an infinite number
12/// of harmonics, which are aliased back and forth across the frequency
13/// spectrum.
14///
15/// Parameters
16/// ----------
17/// t : array_like
18///     The input time array.
19/// duty : array_like, optional
20///     Duty cycle.  Default is 0.5 (50% duty cycle).
21///     If an array, causes wave shape to change over time, and must be the
22///     same length as t.
23///
24/// Returns
25/// -------
26/// y : ndarray
27///     Output array containing the square waveform.
28///
29/// Examples
30/// --------
31/// A 5 Hz waveform sampled at 500 Hz for 1 second:
32///
33/// >>> import numpy as np
34/// >>> from scipy import signal
35/// >>> import matplotlib.pyplot as plt
36/// >>> t = np.linspace(0, 1, 500, endpoint=False)
37/// >>> plt.plot(t, signal.square(2 * np.pi * 5 * t))
38/// >>> plt.ylim(-2, 2)
39///
40/// A pulse-width modulated sine wave:
41///
42/// >>> plt.figure()
43/// >>> sig = np.sin(2 * np.pi * t)
44/// >>> pwm = signal.square(2 * np.pi * 30 * t, duty=(sig + 1)/2)
45/// >>> plt.subplot(2, 1, 1)
46/// >>> plt.plot(t, sig)
47/// >>> plt.subplot(2, 1, 2)
48/// >>> plt.plot(t, pwm)
49/// >>> plt.ylim(-1.5, 1.5)
50///
51/// """
52pub fn square<F, S, D>(t: &ArrayBase<S, D>, duty: F) -> Array<F, D>
53where
54    F: RealField,
55    S: Data<Elem = F>,
56    D: Dimension,
57{
58    assert!(F::zero() <= duty && duty <= F::one());
59    let duty_threshold = F::two_pi() * duty;
60    t.mapv(|t| {
61        let x = t % F::two_pi();
62        // Because % is the reminder and not the modulo operator, x can be negative.
63        let x = if x < F::zero() { x + F::two_pi() } else { x };
64        debug_assert!(F::zero() <= x && x <= F::two_pi());
65        if x < duty_threshold {
66            F::one()
67        } else {
68            -F::one()
69        }
70    })
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use approx::assert_abs_diff_eq;
77    use ndarray::{arr1, arr3};
78
79    #[test]
80    fn test_square_zero_duty() {
81        let t = arr1(&[
82            -4.821, -4.15, -3.394, -3.386, -2.966, -2.735, -2.464, -2.277, -2.094, -0.8963,
83            0.03853, 1.432, 2.384, 2.522, 2.732, 3.125, 3.297, 3.517, 3.602, 4.908,
84        ]);
85        let expected = arr1(&[
86            -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0,
87            -1.0, -1.0, -1.0, -1.0, -1.0, -1.0,
88        ]);
89        let result = square(&t, 0.0);
90        assert_vec_eq(result, expected);
91    }
92
93    #[test]
94    fn test_square_one_duty() {
95        let t = arr1(&[
96            -3.521, -3.284, -3.257, -2.367, -1.965, -1.933, -0.5399, -0.4277, -0.3761, 0.3024,
97            0.3624, 0.6161, 0.784, 1.42, 1.869, 1.934, 4.0, 4.298, 4.389, 4.807,
98        ]);
99        let expected = arr1(&[
100            1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
101            1.0, 1.0, 1.0,
102        ]);
103        let result = square(&t, 1.0);
104        assert_vec_eq(result, expected);
105    }
106
107    #[test]
108    fn test_square_duty_03_05_07() {
109        let t = arr1(&[
110            -4.991, -4.973, -3.988, -3.084, -2.562, -2.378, -1.618, -1.449, -0.8056, -0.6883,
111            -0.5677, -0.5353, -0.1377, -0.1142, 0.8072, 0.821, 1.836, 2.722, 4.189, 4.384,
112        ]);
113
114        let expected_03 = arr1(&[
115            1.0, 1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, 1.0,
116            1.0, 1.0, -1.0, -1.0, -1.0,
117        ]);
118        let expected_05 = arr1(&[
119            1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, 1.0,
120            1.0, 1.0, 1.0, -1.0, -1.0,
121        ]);
122        let expected_07 = arr1(&[
123            1.0, 1.0, 1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, 1.0, 1.0,
124            1.0, 1.0, 1.0, 1.0,
125        ]);
126        let result_03 = square(&t, 0.3);
127        let result_05 = square(&t, 0.5);
128        let result_07 = square(&t, 0.7);
129        assert_vec_eq(result_03, expected_03);
130        assert_vec_eq(result_05, expected_05);
131        assert_vec_eq(result_07, expected_07);
132    }
133
134    #[test]
135    fn test_square_3d() {
136        let t = arr3(&[
137            [
138                [-4.452, -4.182, -3.663, -3.307, -2.995],
139                [-2.482, -2.46, -1.929, -1.823, -1.44],
140            ],
141            [
142                [-0.8743, 0.5359, 0.9073, 2.101, 2.161],
143                [2.582, 2.977, 3.966, 4.298, 4.659],
144            ],
145        ]);
146        let expected = arr3(&[
147            [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, -1.0, -1.0, -1.0]],
148            [[-1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, -1.0, -1.0]],
149        ]);
150        let result = square(&t, 0.67);
151        assert_vec_eq(result, expected);
152    }
153
154    #[track_caller]
155    fn assert_vec_eq<D: Dimension>(a: Array<f32, D>, b: Array<f32, D>) {
156        assert_eq!(a.shape(), b.shape());
157        for (a, b) in a.into_iter().zip(b) {
158            assert_abs_diff_eq!(a, b, epsilon = 1e-6);
159        }
160    }
161}