From 8f01f8d84f0e23093fc5907fc51d38c78086d602 Mon Sep 17 00:00:00 2001 From: Camilo Julian Berutti <104410591+camiloberutti@users.noreply.github.com> Date: Sun, 12 Apr 2026 23:01:06 +0200 Subject: [PATCH] Enable IRASA multi-method support and add regression tests --- neurodsp/aperiodic/irasa.py | 17 ++- neurodsp/spectral/power.py | 1 + .../tests/aperiodic/test_irasa_multimethod.py | 108 ++++++++++++++++++ 3 files changed, 120 insertions(+), 6 deletions(-) create mode 100644 neurodsp/tests/aperiodic/test_irasa_multimethod.py diff --git a/neurodsp/aperiodic/irasa.py b/neurodsp/aperiodic/irasa.py index 60b6dbcbd..5176ec11f 100644 --- a/neurodsp/aperiodic/irasa.py +++ b/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 diff --git a/neurodsp/spectral/power.py b/neurodsp/spectral/power.py index b9a11b813..392c842d1 100644 --- a/neurodsp/spectral/power.py +++ b/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'], } diff --git a/neurodsp/tests/aperiodic/test_irasa_multimethod.py b/neurodsp/tests/aperiodic/test_irasa_multimethod.py new file mode 100644 index 000000000..2115141c2 --- /dev/null +++ b/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'])