From 8346e3cc27480307abd64a95e755e79fd413b704 Mon Sep 17 00:00:00 2001
From: Liam Gray <lgray@phas.ubc.ca>
Date: Fri, 12 Apr 2024 12:26:03 -0700
Subject: [PATCH] feat(delay): add option to low-pass with null delay filter

---
 draco/analysis/delay.py | 36 ++++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/draco/analysis/delay.py b/draco/analysis/delay.py
index 669257454..a3908c3d8 100644
--- a/draco/analysis/delay.py
+++ b/draco/analysis/delay.py
@@ -1861,15 +1861,24 @@ def delay_spectrum_wiener_filter(
     return y_spec
 
 
-def null_delay_filter(freq, max_delay, mask, num_delay=200, tol=1e-8, window=True):
+def null_delay_filter(
+    freq,
+    delay_cut,
+    mask,
+    num_delay=200,
+    tol=1e-8,
+    window=True,
+    type_="high",
+    lapack_driver="gesvd",
+):
     """Take frequency data and null out any delays below some value.
 
     Parameters
     ----------
     freq : np.ndarray[freq]
         Frequencies we have data at.
-    max_delay : float
-        Maximum delay to keep.
+    delay_cut : float
+        Delay cut to apply.
     mask : np.ndarray[freq]
         Frequencies to mask out.
     num_delay : int, optional
@@ -1878,17 +1887,27 @@ def null_delay_filter(freq, max_delay, mask, num_delay=200, tol=1e-8, window=Tru
         Cut off value for singular values.
     window : bool, optional
         Apply a window function to the data while filtering.
+    type_ : str, optional
+        Whether to apply a high-pass or low-pass filter. Options are
+        `high` or `low`. Default is `high`.
+    lapack_driver : str, optional
+        Which lapack driver to use in the SVD. Options are 'gesvd' or 'gesdd'.
+        'gesdd' is generally faster, but seems to experience convergence issues.
+        Default is 'gesvd'.
 
     Returns
     -------
     filter : np.ndarray[freq, freq]
         The filter as a 2D matrix.
     """
+    if type_ not in {"high", "low"}:
+        raise ValueError(f"Filter type must be one of [high, low]. Got {type_}")
+
     # Construct the window function
     x = (freq - freq.min()) / freq.ptp()
     w = tools.window_generalised(x, window="nuttall")
 
-    delay = np.linspace(-max_delay, max_delay, num_delay)
+    delay = np.linspace(-delay_cut, delay_cut, num_delay)
 
     # Construct the Fourier matrix
     F = mask[:, np.newaxis] * np.exp(
@@ -1904,9 +1923,14 @@ def null_delay_filter(freq, max_delay, mask, num_delay=200, tol=1e-8, window=Tru
     # to be the fault of MKL (see https://github.com/scipy/scipy/issues/10032 and links
     # therein). This seems to be limited to the `gesdd` LAPACK routine, so we can get
     # around it by switching to `gesvd`.
-    u, sig, vh = la.svd(F, lapack_driver="gesvd")
+    u, sig, vh = la.svd(F, full_matrices=False, lapack_driver=lapack_driver)
     nmodes = np.sum(sig > tol * sig.max())
-    p = u[:, :nmodes]
+
+    # Select the modes to null out based on the filter type
+    if type_ == "high":
+        p = u[:, :nmodes]
+    elif type_ == "low":
+        p = u[:, nmodes:]
 
     # Construct a projection matrix for the filter
     proj = np.identity(len(freq)) - np.dot(p, p.T.conj())