Adapt multitaper power spectrum engines to changing window lengths

python/mspasspy/algorithms/MTPowerSpectrumEngine.py

@@ -26,7 +26,8 @@ class MTPowerSpectrumEngine:
 
     :param winsize:   Number of samples expected for data window.  Note this
         argument defines the size of the eigentapers so it defines the size of
-        input time series the object expect to process.
+        input time series the object initially expects to process.  apply adapts
+        the operator and recomputes its tapers when the input length changes.
 
     :param tbp:  multitaper time-bandwidth product
 
@@ -35,8 +36,9 @@ class MTPowerSpectrumEngine:
 
     :param nfft:  optional size of the fft work arrays to use.  nfft should be
         greater than or equal winsize.  When greater zero padding is used in the
-        usual way.  Default is 0 which sets nfft to 2*winsize.  If nfft<winsize
-        a diagnostic message will be printed and nfft will be set equal to winsize.
+        usual way.  Default is 0, which lets MTSpec choose the FFT size for each input
+        length.  An explicit size is retained across length changes while it
+        is large enough; otherwise MTSpec chooses a suitable size automatically.
 
     :param iadapt:  integer argument passed to Prieto's MTspec that sets the
         eigentaper weighting scheme.  See Prieto's MTspect class documentation
@@ -56,14 +58,21 @@ def __init__(
         self.number_tapers = number_tapers  # kspec in multitaper
         self.vn = None  # constructor for MTSpec will set this when passed None
         self.nfft = nfft
+        self._requested_nfft = nfft
         self.idapt = iadapt
         self.lamb = None  # constructor for MTSpec set this
         self.MTSpec_instance = None
 
     def apply(self, d, dt=1.0):
         """
-        needs to support vector input or a TimeSeries - returns a PowerSpectrum
-        Cannot use mspass decorator because it supports raw data vectors
+        Estimate a PowerSpectrum from a TimeSeries or one-dimensional vector.
+
+        A changed input length resizes the operator and recomputes its tapers.
+        An explicit FFT size is retained when large enough; otherwise MTSpec
+        chooses the FFT size automatically.  Failed construction preserves the
+        previous operator state.  Empty inputs are invalid.
+
+        Cannot use mspass decorator because it supports raw data vectors.
         """
         if isinstance(d, TimeSeries):
             input_length = d.npts
@@ -80,11 +89,16 @@ def apply(self, d, dt=1.0):
                 "MTPowerSpectrumEngine.apply: arg0 must be a TimeSeries, "
                 "DoubleVector, or one-dimensional numpy array"
             )
-        if input_length != self.winsize:
-            raise ValueError(
-                "MTPowerSpectrumEngine.apply: input length={} does not match "
-                "winsize={}".format(input_length, self.winsize)
-            )
+        if input_length == 0:
+            raise ValueError("MTPowerSpectrumEngine.apply: input must not be empty")
+        resized = input_length != self.winsize
+        # Tapers and eigenvalues are specific to a window length.  Do not
+        # replace the cached state until MTSpec has constructed successfully.
+        vn = None if resized else self.vn
+        lamb = None if resized else self.lamb
+        nfft = self._requested_nfft if resized else self.nfft
+        if nfft < input_length:
+            nfft = 0
 
         # All inputs end up filling a numpy array passed to MTSpec below
         if isinstance(d, TimeSeries):
@@ -108,20 +122,19 @@ def apply(self, d, dt=1.0):
             # a future maintenance issue
             md["npts"] = len(d)
             md["delta"] = dt
-        self.MTSpec_instance = MTSpec(
+        instance = MTSpec(
             y,
             nw=self.tbp,
             kspec=self.number_tapers,
             dt=dt,
-            nfft=self.nfft,
+            nfft=nfft,
             iadapt=self.idapt,
-            vn=self.vn,
-            lamb=self.lamb,
+            vn=vn,
+            lamb=lamb,
         )
-        # Set these so they will be automatically reused on a
-        # successive call.   May need to trap size mismatch above
-        # if these were previously set - this needs some testing
-        #  REMOVE ME WHEN VERIFIED
+        self.MTSpec_instance = instance
+        self.winsize = input_length
+        # Reuse the length-specific tapers on subsequent equal-length calls.
         self.vn = self.MTSpec_instance.vn
         self.lamb = self.MTSpec_instance.lamb
         self.nfft = self.MTSpec_instance.nfft

python/tests/algorithms/test_mt_power_spectrum_cache_contract.py

@@ -81,16 +81,46 @@ def test_exact_length_is_repeatable_and_reuses_cached_tapers(kind):
 
 
 @pytest.mark.parametrize("kind", ["timeseries", "doublevector", "ndarray"])
-@pytest.mark.parametrize("length", [WINSIZE - 1, WINSIZE + 1])
-def test_length_mismatch_preserves_cache_and_later_valid_result(kind, length):
+@pytest.mark.parametrize("nfft", [0, 32, 64])
+def test_length_changes_match_fresh_engine(kind, nfft):
+    engine = MTPowerSpectrumEngine(WINSIZE, 3.0, 5, nfft=nfft)
+    for length in [WINSIZE, WINSIZE - 1, WINSIZE + 1, WINSIZE - 1, WINSIZE]:
+        datum = _input(kind, length)
+        actual = engine.apply(datum, dt=0.1)
+        expected = MTPowerSpectrumEngine(length, 3.0, 5, nfft=nfft).apply(datum, dt=0.1)
+        assert engine.winsize == length
+        assert engine.vn.shape[0] == length
+        assert actual["npts"] == length
+        assert actual["MTSpec_npts"] == length
+        assert actual["MTSpec_nfft"] == expected["MTSpec_nfft"]
+        assert actual.df() == expected.df()
+        np.testing.assert_allclose(_spectrum_values(actual), _spectrum_values(expected))
+        assert np.all(np.isfinite(_spectrum_values(actual)))
+
+
+@pytest.mark.parametrize("kind", ["timeseries", "doublevector", "ndarray"])
+def test_empty_input_preserves_cache(kind):
     engine = MTPowerSpectrumEngine(WINSIZE, 3.0, 5, nfft=64)
-    valid = _input(kind)
+    engine.apply(_input(kind), dt=0.1)
+    before = _cache_snapshot(engine)
+    with pytest.raises(ValueError, match="empty"):
+        engine.apply(_input(kind, 0), dt=0.1)
+    _assert_cache_unchanged(engine, before)
+    assert engine.winsize == WINSIZE
+
+
+def test_failed_numerical_resize_preserves_cache_and_later_valid_result():
+    engine = MTPowerSpectrumEngine(WINSIZE, 3.0, 5, nfft=64)
+    valid = _input("ndarray")
     expected = engine.apply(valid, dt=0.1)
     before = _cache_snapshot(engine)
 
-    with pytest.raises(ValueError, match="does not match winsize"):
-        engine.apply(_input(kind, length), dt=0.1)
+    # scipy's DPSS implementation requires NW < M/2.  This reaches the
+    # numerical backend with a nonempty window too short for the configured NW.
+    with pytest.raises(ValueError, match="NW must be less than M/2"):
+        engine.apply(_samples(6), dt=0.1)
 
+    assert engine.winsize == WINSIZE
     _assert_cache_unchanged(engine, before)
     actual = engine.apply(valid, dt=0.1)
     np.testing.assert_allclose(_spectrum_values(actual), _spectrum_values(expected))
@@ -122,3 +152,67 @@ def test_multidimensional_array_preserves_initialized_cache():
     _assert_cache_unchanged(engine, before)
     actual = engine.apply(_input("ndarray"), dt=0.1)
     np.testing.assert_allclose(_spectrum_values(actual), _spectrum_values(expected))
+
+
+@pytest.mark.parametrize("length", [WINSIZE - 1, WINSIZE + 1])
+def test_native_timeseries_mismatch_pads_or_truncates(length):
+    from mspasspy.ccore.algorithms.deconvolution import (
+        MTPowerSpectrumEngine as NativeEngine,
+    )
+    from mspasspy.ccore.utility import MsPASSError
+
+    engine = NativeEngine(WINSIZE, 3.0, 5, 64, 0.1)
+    actual = engine.apply(_input("timeseries", length))
+    assert actual.live
+    padded = np.zeros(WINSIZE)
+    padded[: min(length, WINSIZE)] = _samples(length)[:WINSIZE]
+    expected = engine.apply(DoubleVector(padded))
+    np.testing.assert_allclose(_spectrum_values(actual), expected)
+    with pytest.raises(MsPASSError, match="does not match operator"):
+        engine.apply(DoubleVector(_samples(length)))
+
+
+def test_broadband_snr_qc_with_reused_resized_engines():
+    import pickle
+    from pathlib import Path
+    from mspasspy.algorithms.snr import broadband_snr_QC
+    from mspasspy.ccore.algorithms.basic import TimeWindow
+    from mspasspy.algorithms.window import WindowData
+
+    path = Path(__file__).parents[1] / "data" / "snrtestdata"
+    with path.open("rb") as stream:
+        datum = pickle.load(stream)
+    datum["Ptime"] = 0.0
+    noise_window = TimeWindow(-200.0, -50.0)
+    signal_window = TimeWindow(-10.0, 50.0)
+    noise_npts = WindowData(datum, noise_window.start, noise_window.end).npts
+    signal_npts = WindowData(datum, signal_window.start, signal_window.end).npts
+    noise_engine = MTPowerSpectrumEngine(noise_npts - 1, 4.0, 6)
+    signal_engine = MTPowerSpectrumEngine(signal_npts + 1, 4.0, 6)
+    options = dict(
+        noise_window=noise_window,
+        signal_window=signal_window,
+        optional_metrics=[],
+        kill_null_signals=False,
+        save_spectra=True,
+    )
+    for _ in range(2):
+        actual = broadband_snr_QC(
+            TimeSeries(datum),
+            noise_spectrum_engine=noise_engine,
+            signal_spectrum_engine=signal_engine,
+            **options,
+        )
+        expected = broadband_snr_QC(
+            TimeSeries(datum),
+            noise_spectrum_engine=MTPowerSpectrumEngine(noise_npts, 4.0, 6),
+            signal_spectrum_engine=MTPowerSpectrumEngine(signal_npts, 4.0, 6),
+            **options,
+        )
+        assert actual.live
+        assert actual.is_defined("Parrival")
+        assert actual["Parrival"].keys() == expected["Parrival"].keys()
+        for key in ["low_f_band_edge", "high_f_band_edge", "bandwidth"]:
+            assert actual["Parrival"][key] == expected["Parrival"][key]
+    assert noise_engine.winsize == noise_npts
+    assert signal_engine.winsize == signal_npts