add pyfftw_sdp, tests, and relevant fixes

Author: NimaSarajpoor

docstring.py

@@ -23,7 +23,7 @@ def get_docstring_args(fd, file_name, func_name, class_name=None):
         msg += f"function/method: {func_name}\n"
         raise RuntimeError(msg)
 
-    if class_name is None:
+    if len(re.findall(r"Returns", docstring)) > 0:
         params_section = re.findall(
             r"(?<=Parameters)(.*)(?=Returns)", docstring, re.DOTALL
         )[0]
@@ -43,7 +43,7 @@ def get_signature_args(fd):
     return set([a.arg for a in fd.args.args if a.arg != "self"])
 
 
-def check_args(doc_args, sig_args, file_name, func_name, class_name=None):
+def check_args(docstring_args, signature_args, file_name, func_name, class_name=None):
     """
     Compare docstring arguments and signature argments
     """

environment.yml

@@ -27,3 +27,5 @@ dependencies:
   - jupyterlab-myst>=2.0.0
   - myst-nb>=1.0.0
   - polars>=1.14.0
+  - fftw>=3.3
+  - pyfftw>=0.15.0

stumpy/sdp.py

@@ -6,6 +6,13 @@
 
 from . import config
 
+try:
+    import pyfftw
+
+    FFTW_IS_AVAILABLE = True
+except ImportError:  # pragma: no cover
+    FFTW_IS_AVAILABLE = False
+
 
 @njit(fastmath=config.STUMPY_FASTMATH_TRUE)
 def _njit_sliding_dot_product(Q, T):
@@ -99,3 +106,180 @@ def _pocketfft_sliding_dot_product(Q, T):
     fft_2d = r2c(True, tmp, axis=-1)
 
     return c2r(False, np.multiply(fft_2d[0], fft_2d[1]), n=next_fast_n)[m - 1 : n]
+
+
+class _PYFFTW_SLIDING_DOT_PRODUCT:
+    """
+    A class to compute the sliding dot product using FFTW via pyfftw.
+
+    This class uses FFTW (via pyfftw) to efficiently compute the sliding dot product
+    between a query sequence Q and a time series T. It preallocates arrays and caches
+    FFTW objects to optimize repeated computations with similar-sized inputs.
+
+    Parameters
+    ----------
+    max_n : int, default=2**20
+        Maximum length to preallocate arrays for. This will be the size of the
+        real-valued array. A complex-valued array of size `1 + (max_n // 2)`
+        will also be preallocated. If inputs exceed this size, arrays will be
+        reallocated to accommodate larger sizes.
+
+    Attributes
+    ----------
+    real_arr : pyfftw.empty_aligned
+        Preallocated real-valued array for FFTW computations.
+
+    complex_arr : pyfftw.empty_aligned
+        Preallocated complex-valued array for FFTW computations.
+
+    rfft_objects : dict
+        Cache of FFTW forward transform objects, keyed by
+        (next_fast_n, n_threads, planning_flag).
+
+    irfft_objects : dict
+        Cache of FFTW inverse transform objects, keyed by
+        (next_fast_n, n_threads, planning_flag).
+
+    Notes
+    -----
+    The class maintains internal caches of FFTW objects to avoid redundant planning
+    operations when called multiple times with similar-sized inputs and parameters.
+
+    Examples
+    --------
+    >>> sdp_obj = _PYFFTW_SLIDING_DOT_PRODUCT(max_n=1000)
+    >>> Q = np.array([1, 2, 3])
+    >>> T = np.array([4, 5, 6, 7, 8])
+    >>> result = sdp_obj(Q, T)
+
+    References
+    ----------
+    `FFTW documentation <http://www.fftw.org/>`__
+
+    `pyfftw documentation <https://pyfftw.readthedocs.io/>`__
+    """
+
+    def __init__(self, max_n=2**20):
+        """
+        Initialize the `_PYFFTW_SLIDING_DOT_PRODUCT` object, which can be called
+        to compute the sliding dot product using FFTW via pyfftw.
+
+        Parameters
+        ----------
+        max_n : int, default=2**20
+            Maximum length to preallocate arrays for. This will be the size of the
+            real-valued array. A complex-valued array of size `1 + (max_n // 2)`
+            will also be preallocated.
+        """
+        # Preallocate arrays
+        self.real_arr = pyfftw.empty_aligned(max_n, dtype="float64")
+        self.complex_arr = pyfftw.empty_aligned(1 + (max_n // 2), dtype="complex128")
+
+        # Store FFTW objects, keyed by (next_fast_n, n_threads, planning_flag)
+        self.rfft_objects = {}
+        self.irfft_objects = {}
+
+    def __call__(self, Q, T, n_threads=1, planning_flag="FFTW_ESTIMATE"):
+        """
+        Compute the sliding dot product between `Q` and `T` using FFTW via pyfftw,
+        and cache FFTW objects if not already cached.
+
+        Parameters
+        ----------
+        Q : numpy.ndarray
+            Query array or subsequence.
+
+        T : numpy.ndarray
+            Time series or sequence.
+
+        n_threads : int, default=1
+            Number of threads to use for FFTW computations.
+
+        planning_flag : str, default="FFTW_ESTIMATE"
+            The planning flag that will be used in FFTW for planning.
+            See pyfftw documentation for details. Current options, ordered
+            ascendingly by the level of aggressiveness in planning, are:
+            "FFTW_ESTIMATE", "FFTW_MEASURE", "FFTW_PATIENT", and "FFTW_EXHAUSTIVE".
+            The more aggressive the planning, the longer the planning time, but
+            the faster the execution time.
+
+        Returns
+        -------
+        out : numpy.ndarray
+            Sliding dot product between `Q` and `T`.
+
+        Notes
+        -----
+        The planning_flag is defaulted to "FFTW_ESTIMATE" to be aligned with
+        MATLAB's FFTW usage (as of version R2025b)
+        See: https://www.mathworks.com/help/matlab/ref/fftw.html
+
+        This implementation is inspired by the answer on StackOverflow:
+        https://stackoverflow.com/a/30615425/2955541
+        """
+        m = Q.shape[0]
+        n = T.shape[0]
+        next_fast_n = pyfftw.next_fast_len(n)
+
+        # Update preallocated arrays if needed
+        if next_fast_n > len(self.real_arr):
+            self.real_arr = pyfftw.empty_aligned(next_fast_n, dtype="float64")
+            self.complex_arr = pyfftw.empty_aligned(
+                1 + (next_fast_n // 2), dtype="complex128"
+            )
+
+        real_arr = self.real_arr[:next_fast_n]
+        complex_arr = self.complex_arr[: 1 + (next_fast_n // 2)]
+
+        # Get or create FFTW objects
+        key = (next_fast_n, n_threads, planning_flag)
+
+        rfft_obj = self.rfft_objects.get(key, None)
+        if rfft_obj is None:
+            rfft_obj = pyfftw.FFTW(
+                input_array=real_arr,
+                output_array=complex_arr,
+                direction="FFTW_FORWARD",
+                flags=(planning_flag,),
+                threads=n_threads,
+            )
+            self.rfft_objects[key] = rfft_obj
+        else:
+            rfft_obj.update_arrays(real_arr, complex_arr)
+
+        irfft_obj = self.irfft_objects.get(key, None)
+        if irfft_obj is None:
+            irfft_obj = pyfftw.FFTW(
+                input_array=complex_arr,
+                output_array=real_arr,
+                direction="FFTW_BACKWARD",
+                flags=(planning_flag, "FFTW_DESTROY_INPUT"),
+                threads=n_threads,
+            )
+            self.irfft_objects[key] = irfft_obj
+        else:
+            irfft_obj.update_arrays(complex_arr, real_arr)
+
+        # RFFT(T)
+        real_arr[:n] = T
+        real_arr[n:] = 0.0
+        rfft_obj.execute()  # output is in complex_arr
+        complex_arr_T = complex_arr.copy()
+
+        # RFFT(Q)
+        # Scale by 1/next_fast_n to account for
+        # FFTW's unnormalized inverse FFT via execute()
+        np.multiply(Q[::-1], 1.0 / next_fast_n, out=real_arr[:m])
+        real_arr[m:] = 0.0
+        rfft_obj.execute()  # output is in complex_arr
+
+        # RFFT(T) * RFFT(Q)
+        np.multiply(complex_arr, complex_arr_T, out=complex_arr)
+
+        # IRFFT (input is in complex_arr)
+        irfft_obj.execute()  # output is in real_arr
+
+        return real_arr[m - 1 : n]
+
+
+_pyfftw_sliding_dot_product = _PYFFTW_SLIDING_DOT_PRODUCT()

tests/test_sdp.py

@@ -172,3 +172,79 @@ def test_sdp_power2():
             raise e
 
     return
+
+
+def test_pyfftw_sdp_max_n():
+    # When `len(T)` larger than `max_n` in pyfftw_sdp,
+    # the internal preallocated arrays should be resized.
+    # This test checks that functionality.
+    sliding_dot_product = sdp._PYFFTW_SLIDING_DOT_PRODUCT(max_n=2**10)
+
+    # len(T) > max_n to trigger array resizing
+    T = np.random.rand(2**11)
+    Q = np.random.rand(2**8)
+
+    comp = sliding_dot_product(Q, T)
+    ref = naive_sliding_dot_product(Q, T)
+
+    np.testing.assert_allclose(comp, ref)
+
+    return
+
+
+def test_pyfftw_sdp_cache():
+    # To ensure that the caching mechanism in
+    # pyfftw_sdp is working as intended
+    sliding_dot_product = sdp._PYFFTW_SLIDING_DOT_PRODUCT(max_n=2**10)
+    assert sliding_dot_product.rfft_objects == {}
+    assert sliding_dot_product.irfft_objects == {}
+
+    T = np.random.rand(2**5)
+    Q = np.random.rand(2**2)
+
+    n_threads = 1
+    planning_flag = "FFTW_ESTIMATE"
+    sliding_dot_product(Q, T, n_threads=n_threads, planning_flag=planning_flag)
+
+    # Check that the FFTW objects are cached
+    expected_key = (len(T), n_threads, planning_flag)
+    assert expected_key in sliding_dot_product.rfft_objects
+    assert expected_key in sliding_dot_product.irfft_objects
+
+    return
+
+
+def test_pyfftw_sdp_update_arrays():
+    # To ensure that the cached FFTW objects
+    # can be reused when preallocated arrays
+    # are updated.
+    sliding_dot_product = sdp._PYFFTW_SLIDING_DOT_PRODUCT(max_n=2**10)
+
+    n_threads = 1
+    planning_flag = "FFTW_ESTIMATE"
+
+    T1 = np.random.rand(2**5)
+    Q1 = np.random.rand(2**2)
+    sliding_dot_product(Q1, T1, n_threads=n_threads, planning_flag=planning_flag)
+
+    # len(T2) > max_n to trigger array resizing
+    T2 = np.random.rand(2**11)
+    Q2 = np.random.rand(2**3)
+    sliding_dot_product(Q2, T2, n_threads=n_threads, planning_flag=planning_flag)
+
+    # Check if the FFTW objects cached for inputs (Q1, T1)
+    # can be reused when preallocated arrays are resized
+    # after calling with (Q2, T2)
+    key1 = (len(T1), n_threads, planning_flag)
+    rfft_obj_before = sliding_dot_product.rfft_objects[key1]
+    irfft_obj_before = sliding_dot_product.irfft_objects[key1]
+
+    comp = sliding_dot_product(Q1, T1, n_threads=n_threads, planning_flag=planning_flag)
+    ref = naive_sliding_dot_product(Q1, T1)
+
+    # test for correctness
+    np.testing.assert_allclose(comp, ref)
+
+    # Check that the same FFTW objects are reused
+    assert sliding_dot_product.rfft_objects[key1] is rfft_obj_before
+    assert sliding_dot_product.irfft_objects[key1] is irfft_obj_before