Enable IRASA multi-method support and add regression tests

neurodsp/aperiodic/irasa.py

@@ -69,15 +69,15 @@ def compute_irasa(sig, fs, f_range=None, hset=None, thresh=None, **spectrum_kwar
     hset = np.arange(1.1, 1.95, 0.05) if hset is None else hset
     hset = np.round(hset, 4)
 
-    # The `nperseg` input needs to be set to lock in the size of the FFT's
-    if 'nperseg' not in spectrum_kwargs:
+    # Only Welch uses `nperseg`; avoid injecting it for non-Welch methods.
+    if spectrum_kwargs.get('method', 'welch') == 'welch' and 'nperseg' not in spectrum_kwargs:
         spectrum_kwargs['nperseg'] = int(4 * fs)
 
     # Calculate the original spectrum across the whole signal
     freqs, psd = compute_spectrum(sig, fs, **spectrum_kwargs)
 
     # Do the IRASA resampling procedure
-    psds = np.zeros((len(hset), *psd.shape))
+    psds = np.full((len(hset), *psd.shape), np.nan, dtype=float)
     for ind, h_val in enumerate(hset):
 
         # Get the up-sampling / down-sampling (h, 1/h) factors as integers
@@ -92,11 +92,16 @@ def compute_irasa(sig, fs, f_range=None, hset=None, thresh=None, **spectrum_kwar
         freqs_up, psd_up = compute_spectrum(sig_up, h_val * fs, **spectrum_kwargs)
         freqs_dn, psd_dn = compute_spectrum(sig_dn, fs / h_val, **spectrum_kwargs)
 
-        # Calculate the geometric mean of h and 1/h
-        psds[ind, :] = np.sqrt(psd_up * psd_dn)
+        # Align spectra to the original frequency grid for methods whose output length
+        # changes with signal length (for example medfilt and multitaper).
+        psd_up_i = np.interp(freqs, freqs_up, psd_up, left=np.nan, right=np.nan)
+        psd_dn_i = np.interp(freqs, freqs_dn, psd_dn, left=np.nan, right=np.nan)
+
+        # Calculate the geometric mean of h and 1/h on a shared frequency grid.
+        psds[ind, :] = np.sqrt(psd_up_i * psd_dn_i)
 
     # Take the median resampled spectra, as an estimate of the aperiodic component
-    psd_aperiodic = np.median(psds, axis=0)
+    psd_aperiodic = np.nanmedian(psds, axis=0)
 
     # Subtract aperiodic from original, to get the periodic component
     psd_periodic = psd - psd_aperiodic

neurodsp/spectral/power.py

@@ -78,6 +78,7 @@ def compute_spectrum(sig, fs, method='welch', **kwargs):
     'welch' : ['avg_type', 'window', 'nperseg', 'noverlap', \
                'nfft', 'fast_len', 'f_range'],
     'medfilt' : ['filt_len', 'f_range'],
+    'multitaper' : ['bandwidth', 'n_tapers', 'low_bias', 'eigenvalue_weighting'],
 }
 
 

neurodsp/tests/aperiodic/test_irasa_multimethod.py

@@ -0,0 +1,108 @@
+"""
+Test suite for IRASA multi-method support fix.
+
+This test file validates that compute_irasa now supports all spectral
+estimation methods: welch, medfilt, multitaper, and wavelet.
+
+Before patch: Only welch worked; others raised AssertionError or KeyError.
+After patch: All methods work with appropriate parameters.
+"""
+
+import numpy as np
+import pytest
+from neurodsp.sim import sim_combined
+from neurodsp.aperiodic import compute_irasa
+
+
+class TestIRASAMultiMethods:
+    """Test compute_irasa with all supported spectral methods."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """Create test signal and frequency range."""
+        self.fs = 50.0
+        self.sig = sim_combined(
+            n_seconds=60,
+            fs=self.fs,
+            components={
+                'sim_powerlaw': {},
+                'sim_oscillation': {'freq': 10.0}
+            }
+        )
+        self.f_range = [1, 20]
+
+    def test_irasa_welch(self):
+        """Test IRASA with Welch method (original working case)."""
+        freqs, ap, pe = compute_irasa(
+            self.sig, self.fs,
+            f_range=self.f_range,
+            method='welch',
+            nperseg=1000,
+            noverlap=500
+        )
+        assert len(freqs) > 0, "Welch returned empty frequency array"
+        assert len(ap) == len(freqs), "Aperiodic component length mismatch"
+        assert len(pe) == len(freqs), "Periodic component length mismatch"
+
+    def test_irasa_medfilt(self):
+        """Test IRASA with medfilt method (patched)."""
+        freqs, ap, pe = compute_irasa(
+            self.sig, self.fs,
+            f_range=self.f_range,
+            method='medfilt',
+            filt_len=1.0
+        )
+        assert len(freqs) > 0, "Medfilt returned empty frequency array"
+        assert len(ap) == len(freqs), "Aperiodic component length mismatch"
+        assert len(pe) == len(freqs), "Periodic component length mismatch"
+        assert np.isfinite(ap).all(), "Aperiodic component has non-finite values"
+        assert np.isfinite(pe).all(), "Periodic component has non-finite values"
+
+    def test_irasa_multitaper(self):
+        """Test IRASA with multitaper method (patched)."""
+        freqs, ap, pe = compute_irasa(
+            self.sig, self.fs,
+            f_range=self.f_range,
+            method='multitaper'
+        )
+        assert len(freqs) > 0, "Multitaper returned empty frequency array"
+        assert len(ap) == len(freqs), "Aperiodic component length mismatch"
+        assert len(pe) == len(freqs), "Periodic component length mismatch"
+
+    def test_irasa_wavelet(self):
+        """Test IRASA with wavelet method (patched)."""
+        freqs_wavelet = np.logspace(
+            np.log10(self.f_range[0]),
+            np.log10(self.f_range[1]),
+            30
+        )
+        freqs, ap, pe = compute_irasa(
+            self.sig, self.fs,
+            f_range=self.f_range,
+            method='wavelet',
+            freqs=freqs_wavelet,
+            n_cycles=3.0
+        )
+        assert len(freqs) > 0, "Wavelet returned empty frequency array"
+        assert len(ap) == len(freqs), "Aperiodic component length mismatch"
+        assert len(pe) == len(freqs), "Periodic component length mismatch"
+
+    def test_irasa_default_method_is_welch(self):
+        """Verify backward compatibility: default method is welch."""
+        freqs_default, ap_default, pe_default = compute_irasa(
+            self.sig, self.fs,
+            f_range=self.f_range,
+            nperseg=1000
+        )
+        freqs_explicit, ap_explicit, pe_explicit = compute_irasa(
+            self.sig, self.fs,
+            f_range=self.f_range,
+            method='welch',
+            nperseg=1000
+        )
+        assert len(freqs_default) == len(freqs_explicit)
+        np.testing.assert_array_almost_equal(freqs_default, freqs_explicit)
+
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-v'])