Add general cosine window

SpookyYomo · SpookyYomo · commit ed969ebaddcb · 2024-11-19T21:18:32.000+08:00
diff --git a/sci-rs/src/signal/windows/general_cosine.rs b/sci-rs/src/signal/windows/general_cosine.rs
@@ -0,0 +1,217 @@
+use super::{extend, len_guard, truncate};
+use nalgebra::RealField;
+use num_traits::{real::Real, Float};
+
+#[cfg(feature = "alloc")]
+use super::GetWindow;
+#[cfg(feature = "alloc")]
+use alloc::{vec, vec::Vec};
+
+/// Collection of arguments for window `GeneralCosine` for use in [GetWindow].
+#[cfg(feature = "alloc")]
+#[derive(Debug, Clone, PartialEq)]
+pub struct GeneralCosine<A>
+where
+    A: Real + Float,
+{
+    /// Number of points in the output window. If zero, an empty array is returned in [GetWindow].
+    pub m: usize,
+    /// Sequence of weighting coefficients.
+    pub a: Vec<A>, // Is there a better type, such as impl Into?
+    /// Whether the window is symmetric.
+    ///
+    /// When true, generates a symmetric window, for use in filter design.  
+    /// When false, generates a periodic window, for use in spectral analysis.
+    pub sym: bool,
+}
+
+impl<A> GeneralCosine<A>
+where
+    A: Real + Float,
+{
+    /// Returns a GeneralCosine struct.  
+    ///
+    /// # Parameters
+    /// * `m`:  
+    ///     Number of points in the output window. If zero, an empty array is returned.  
+    /// * `a`:  
+    ///     Sequence of weighting coefficients. This uses the convention of being centered on the
+    ///     origin, so these will typically all be positive numbers, not alternating sign.
+    /// * `sym`:  
+    ///     When true, generates a symmetric window, for use in filter design.  
+    ///     When false, generates a periodic window, for use in spectral analysis.
+    pub fn new(m: usize, a: Vec<A>, sym: bool) -> Self {
+        GeneralCosine { m, a, sym }
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl<A> GetWindow for GeneralCosine<A>
+where
+    A: Real + Float,
+{
+    /// Return a window of type: GeneralCosine.
+    ///
+    /// The general cosine window is a generic weighted sum of cosine terms.
+    ///
+    /// # Parameters
+    /// `self`: [GeneralCosine]
+    ///
+    /// # Returns
+    /// `w`: `vec<F>`  
+    ///     The window, with the maximum value normalized to 1 (though the value 1 does not appear
+    ///     if `M` is even and `sym` is True).
+    ///
+    /// # Example
+    /// We can create a `flat-top` window named "HFT90D" as following:
+    /// ```
+    /// use sci_rs::signal::windows::{GeneralCosine, GetWindow};
+    ///
+    /// let hfd90 = [1., 1.942604, 1.340318, 0.440811, 0.043097].into();
+    /// let window: Vec<f64> = GeneralCosine::new(30, hfd90, true).get_window();
+    /// ```
+    ///
+    /// This is equivalent to the Python code:
+    /// ```custom,{class=language-python}
+    /// from scipy.signal.windows import general_cosine
+    /// HFT90D = [1, 1.942604, 1.340318, 0.440811, 0.043097]
+    /// window = general_cosine(30, HFT90D, sym=False)
+    /// ```
+    ///
+    /// # References
+    /// <https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.windows.general_cosine.html>
+    #[cfg(feature = "alloc")]
+    fn get_window<F>(&self) -> Vec<F>
+    where
+        F: Real + Float + RealField,
+    {
+        if len_guard(self.m) {
+            return Vec::<F>::new();
+        }
+        let (m, needs_trunc) = extend(self.m, self.sym);
+
+        let linspace = (0..m).map(|i| F::from(i).unwrap());
+        let fac: _ = linspace.map(|i| F::two_pi() * i / F::from(m - 1).unwrap() - F::pi());
+        let w: Vec<_> = self
+            .a
+            .iter()
+            .enumerate()
+            .map(|(k, a)| {
+                fac.clone()
+                    .map(move |f| Float::cos(f * F::from(k).unwrap()) * F::from(*a).unwrap())
+            })
+            .fold(
+                vec![F::from(0).unwrap(); fac.clone().count()],
+                // Would this have made more sense if using ndarray::Array1? No, Array1 has no pop.
+                |acc, x| acc.iter().zip(x).map(|(&a, b)| a + b).collect(),
+            )
+            .into_iter()
+            .collect();
+
+        return truncate(w, needs_trunc);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use approx::assert_abs_diff_eq;
+
+    #[test]
+    fn general_cosine_scipy_eg() {
+        // Created with
+        // >>> from scipy.signal.windows import general_cosine
+        // >>> HFT90D = [1, 1.942604, 1.340318, 0.440811, 0.043097]
+        // >>> window = general_cosine(30, HFT90D, sym=False)
+        let expected = vec![
+            -1.04083409e-16,
+            -3.49808964e-03,
+            -1.85322176e-02,
+            -5.60667246e-02,
+            -1.25488810e-01,
+            -2.22198500e-01,
+            -3.14696393e-01,
+            -3.38497088e-01,
+            -2.04818447e-01,
+            1.72651725e-01,
+            8.38783500e-01,
+            1.76097559e+00,
+            2.81469639e+00,
+            3.80321808e+00,
+            4.51005597e+00,
+            4.76683000e+00,
+            4.51005597e+00,
+            3.80321808e+00,
+            2.81469639e+00,
+            1.76097559e+00,
+            8.38783500e-01,
+            1.72651725e-01,
+            -2.04818447e-01,
+            -3.38497088e-01,
+            -3.14696393e-01,
+            -2.22198500e-01,
+            -1.25488810e-01,
+            -5.60667246e-02,
+            -1.85322176e-02,
+            -3.49808964e-03,
+        ];
+
+        let hfd90 = [1., 1.942604, 1.340318, 0.440811, 0.043097].into();
+        let gc = GeneralCosine::new(30, hfd90, false);
+        assert_vec_eq(expected, gc.get_window());
+    }
+
+    #[test]
+    fn constant() {
+        let x = 0.42;
+        let n = 6;
+        let expected = vec![x; n];
+
+        assert_eq!(
+            expected,
+            GeneralCosine::new(n, vec![x], true).get_window::<f64>()
+        )
+    }
+
+    #[test]
+    fn case_b() {
+        // Created with
+        // >>> from scipy.signal.windows import general_cosine
+        // >>> n = 20
+        // >>> a = [0.42, 0.50]
+        // >>> window = general_cosine(30, a, sym=False)
+        let n = 20;
+        let a = vec![0.42, 0.50];
+        let expected = vec![
+            -0.08,
+            -0.05552826,
+            0.0154915,
+            0.12610737,
+            0.2654915,
+            0.42,
+            0.5745085,
+            0.71389263,
+            0.8245085,
+            0.89552826,
+            0.92,
+            0.89552826,
+            0.8245085,
+            0.71389263,
+            0.5745085,
+            0.42,
+            0.2654915,
+            0.12610737,
+            0.0154915,
+            -0.05552826,
+        ];
+
+        assert_vec_eq(expected, GeneralCosine::new(n, a, false).get_window());
+    }
+
+    #[track_caller]
+    fn assert_vec_eq(a: Vec<f64>, b: Vec<f64>) {
+        for (a, b) in a.into_iter().zip(b) {
+            assert_abs_diff_eq!(a, b, epsilon = 1e-6);
+        }
+    }
+}
diff --git a/sci-rs/src/signal/windows/mod.rs b/sci-rs/src/signal/windows/mod.rs
@@ -1,5 +1,6 @@
 #[cfg(feature = "alloc")]
 use alloc::vec::Vec;
+use nalgebra::RealField;
 use num_traits::{real::Real, Float};
 
 /// Corresponding window representation for tuple-structs of [Window] variants.
@@ -25,7 +26,7 @@ pub trait GetWindow {
     #[cfg(feature = "alloc")]
     fn get_window<W>(&self) -> Vec<W>
     where
-        W: Real + Float;
+        W: Real + Float + RealField;
 }
 
 /// Private function for windows implementing [GetWindow]
@@ -57,14 +58,23 @@ fn truncate<'a, W>(mut w: Vec<W>, needed: bool) -> Vec<W> {
 }
 
 mod boxcar;
+mod general_cosine;
 mod triangle;
 pub use boxcar::Boxcar;
+pub use general_cosine::GeneralCosine;
 pub use triangle::Triangle;
 
-/// todo
+/// This collects all structs that implement the [GetWindow] trait.  
+/// This allows for running .get_window() on the struct, which can then be, for example, used in
+/// Firwin.
 // Ordering is as in accordance with
 // https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html.
-pub enum Window {
+// Is it possible for the enums which wraps the structs to only require the generic that the struct
+// implements GetWindow?
+pub enum Window<A>
+where
+    A: Float,
+{
     /// Boxcar window, also known as a rectangular window or Dirichlet window; This is equivalent
     /// to no window at all.
     Boxcar(Boxcar),
@@ -89,8 +99,9 @@ pub enum Window {
     // Kaiser, // Needs Beta
     // KaiserBesselDerived, // Needs Beta
     // Gaussian, // Needs Standard Deviation
-    // Generic weighted sum of cosine term windows.
-    // GenericCosine(GenericCosine<T>), // Needs Weighting Coefficients
+    /// Generic weighted sum of cosine term windows.
+    // Needs Weighting Coefficients
+    GeneralCosine(GeneralCosine<A>),
     // GeneralGaussian, // Needs Power, Width
     // GeneralHamming, // Needs Window Coefficients.
     // Dpss, // Needs Normalized Half-Bandwidth.
