From cfa74f5ffc7358de1c550a680f7bcfd0f1ce508a Mon Sep 17 00:00:00 2001 From: Anton Soloviev Date: Tue, 13 Oct 2026 13:36:49 -0700 Subject: [PATCH] Align matrix periodograms with the selected sample axis --- src/spectrum/periodogram.py | 21 ++++++++++++--------- test/test_periodogram_axis.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 9 deletions(-) create mode 100644 test/test_periodogram_axis.py diff --git a/src/spectrum/periodogram.py b/src/spectrum/periodogram.py index df25122..20ffa2b 100644 --- a/src/spectrum/periodogram.py +++ b/src/spectrum/periodogram.py @@ -54,6 +54,7 @@ def speriodogram(x, NFFT=None, detrend=True, sampling=1., :param x: an array or matrix of data samples. :param NFFT: length of the data before FFT is computed (zero padding) + :param int axis: sample axis for matrix input (default 0). :param bool detrend: detrend the data before co,puteing the FFT :param float sampling: sampling frequency of the input :attr:`data`. @@ -62,9 +63,9 @@ def speriodogram(x, NFFT=None, detrend=True, sampling=1., :return: 2-sided PSD if complex data, 1-sided if real. - if a matrix is provided (using numpy.matrix), then a periodogram - is computed for each row. The returned matrix has the same shape as the input - matrix. + For matrix input, ``axis`` selects the sample dimension (default 0). + A separate periodogram is computed for each column with ``axis=0`` or + each row with ``axis=1``. That dimension is replaced by the frequency bins. The mean of the input data is also removed from the data before computing the psd. @@ -110,27 +111,29 @@ def speriodogram(x, NFFT=None, detrend=True, sampling=1., # matrix case elif x.ndim == 2: logging.debug('2D array. each row is a 1D array') - [r, c] = x.shape - w = np.array([Window(r, window).data for this in range(c)]).reshape(r,c) + r = x.shape[axis] + shape = [1, 1] + shape[axis] = r + w = Window(r, window).data.reshape(shape) if NFFT is None: - NFFT = len(x) + NFFT = r isreal = np.isrealobj(x) if detrend == True: - m = np.mean(x, axis=axis) + m = np.mean(x, axis=axis, keepdims=True) else: m = 0 if isreal == True: if x.ndim == 2: - res = (abs (rfft (x*w - m, NFFT, axis=0))) ** 2. / r + res = (abs (rfft (x*w - m, NFFT, axis=axis))) ** 2. / r else: res = (abs (rfft (x*w - m, NFFT, axis=-1))) ** 2. / r else: if x.ndim == 2: - res = (abs (fft (x*w - m, NFFT, axis=0))) ** 2. / r + res = (abs (fft (x*w - m, NFFT, axis=axis))) ** 2. / r else: res = (abs (fft (x*w - m, NFFT, axis=-1))) ** 2. / r diff --git a/test/test_periodogram_axis.py b/test/test_periodogram_axis.py new file mode 100644 index 0000000..463ac26 --- /dev/null +++ b/test/test_periodogram_axis.py @@ -0,0 +1,30 @@ +import numpy as np +import pytest + +from spectrum import speriodogram + + +@pytest.mark.parametrize("axis", [0, 1, -1, -2]) +@pytest.mark.parametrize("window", ["hamming", "hann", "rectangle"]) +@pytest.mark.parametrize("nfft", [None, 11, 32]) +@pytest.mark.parametrize("detrend", [False, True]) +@pytest.mark.parametrize("complex_data", [False, True]) +def test_matrix_periodogram_matches_separate_channels( + axis, window, nfft, detrend, complex_data +): + data = np.random.default_rng(5).normal(size=(13, 3)) + if complex_data: + data = data + 1j * np.random.default_rng(8).normal(size=data.shape) + columns = [ + speriodogram(data[:, i], NFFT=nfft, window=window, detrend=detrend) + for i in range(data.shape[1]) + ] + expected = np.column_stack(columns) + if axis in (1, -1): + data = data.T + expected = expected.T + before = data.copy() + actual = speriodogram(data, NFFT=nfft, window=window, axis=axis, detrend=detrend) + assert actual.shape == expected.shape + np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12) + np.testing.assert_array_equal(data, before)