1use nalgebra::RealField;
2use ndarray::{Array, ArrayBase, Data, Dimension, RawData};
3
4pub 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 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}