Merge pull request #398 from pynapple-org/update_psd

Update the definition of `compute_power_spectral_density` to compute the true PSD

Author: gviejo

doc/external.md

@@ -8,11 +8,11 @@ As such, it can function as a foundational element for other analysis packages h
 
 ![image](https://raw.githubusercontent.com/flatironinstitute/nemos/main/docs/assets/glm_features_scheme.svg)
 
-[NeMoS](https://nemos.readthedocs.io/en/stable/) is a statistical modeling framework optimized for systems neuroscience and powered by JAX. It streamlines the process of defining and selecting models, through a collection of easy-to-use methods for feature design.
+[NeMoS](https://nemos.readthedocs.io/en/latest/index.html) is a statistical modeling framework optimized for systems neuroscience and powered by JAX. It streamlines the process of defining and selecting models, through a collection of easy-to-use methods for feature design.
 
 The core of nemos includes GPU-accelerated, well-tested implementations of standard statistical models, currently focusing on the Generalized Linear Model (GLM). 
 
-Check out this [page](https://nemos.readthedocs.io/en/stable/generated/neural_modeling/) for many examples of neural modelling using nemos and pynapple.
+Check out this [page](https://nemos.readthedocs.io/en/latest/tutorials/README.html) for many examples of neural modelling using nemos and pynapple.
 
 ```{eval-rst}
 .. Note::

pynapple/process/__init__.py

@@ -22,10 +22,7 @@
     shift_timestamps,
     shuffle_ts_intervals,
 )
-from .spectrum import (
-    compute_mean_power_spectral_density,
-    compute_power_spectral_density,
-)
+from .spectrum import compute_fft, compute_mean_fft, compute_power_spectral_density
 from .tuning_curves import (
     compute_1d_mutual_info,
     compute_1d_tuning_curves,

pynapple/process/spectrum.py

@@ -46,11 +46,9 @@ def wrapper(*args, **kwargs):
 
 
 @_validate_spectrum_inputs
-def compute_power_spectral_density(
-    sig, fs=None, ep=None, full_range=False, norm=False, n=None
-):
+def compute_fft(sig, fs=None, ep=None, full_range=False, norm=False, n=None):
     """
-    Compute Power Spectral Density over a single epoch.
+    Compute Fast Fourier Transform over a single epoch.
     Perform numpy fft on sig, returns output assuming a constant sampling rate for the signal.
 
     Parameters
@@ -74,12 +72,14 @@ def compute_power_spectral_density(
     -------
     pandas.DataFrame
         Time frequency representation of the input signal, indexes are frequencies, values
-        are powers.
+        are the FFT.
 
     Notes
     -----
     This function computes fft on only a single epoch of data. This epoch be given with the ep
     parameter otherwise will be sig.time_support, but it must only be a single epoch.
+
+    If `full_range` is False, the nyquist frequency is excluded in the output due to how computes the FFT frequencies in `numpy.fft.fftfreq`.
     """
     if ep is None:
         ep = sig.time_support
@@ -105,7 +105,73 @@ def compute_power_spectral_density(
 
 
 @_validate_spectrum_inputs
-def compute_mean_power_spectral_density(
+def compute_power_spectral_density(sig, fs=None, ep=None, full_range=False, n=None):
+    """
+    Compute Power Spectral Density over a single epoch.
+    Perform numpy fft on sig and obtain the periodogram, returns output assuming a constant sampling rate for the signal.
+
+    Parameters
+    ----------
+    sig : pynapple.Tsd or pynapple.TsdFrame
+        Time series.
+    fs : float, optional
+        Sampling rate, in Hz. If None, will be calculated from the given signal
+    ep : None or pynapple.IntervalSet, optional
+        The epoch to calculate the fft on. Must be length 1.
+    full_range : bool, optional
+        If true, will return full fft frequency range, otherwise will return only positive values
+    n: int, optional
+        Length of the transformed axis of the output. If n is smaller than the length of the input,
+        the input is cropped. If it is larger, the input is padded with zeros. If n is not given,
+        the length of the input along the axis specified by axis is used.
+
+    Returns
+    -------
+    pandas.DataFrame
+        Power spectral density of the input signal, indexes are frequencies, values
+        are powers/frequency.
+
+    Notes
+    -----
+    This function computes fft on only a single epoch of data. This epoch be given with the ep
+    parameter otherwise will be sig.time_support, but it must only be a single epoch.
+
+    The power spectral density is calculated as the square of the absolute value of the FFT, scaled by the sampling rate and length of the signal.
+    See [this tutorial](https://www.mathworks.com/help/signal/ug/power-spectral-density-estimates-using-fft.html) for more information.
+
+    If `full_range` is False, the nyquist frequency is excluded in the output due to how computes the FFT frequencies in `numpy.fft.fftfreq`.
+    """
+
+    if ep is None:
+        ep = sig.time_support
+    if len(ep) != 1:
+        raise ValueError("Given epoch (or signal time_support) must have length 1")
+    if fs is None:
+        fs = sig.rate
+    if n is None:
+        n = len(sig.restrict(ep))
+
+    fft = compute_fft(sig, fs=fs, ep=ep, n=n, full_range=full_range)
+
+    # transform to power spectral density, power/Hz
+    psd = (1 / (fs * n)) * np.abs(fft) ** 2
+
+    if full_range is False:
+        # frequencies not at 0 and not at the nyquist frequency occur twice
+        # subtract from the nyquist frequency to adjust for floating point error in np.fft.fftfreq
+        # nyquist freq may occur at negative end of frequencies if N is even
+        doubled_freqs = (
+            (fft.index != 0)  # not 0
+            & (fft.index < (fs / 2 - 1e-6))  # less than positive nyquist freq
+            & (fft.index > (-fs / 2 + 1e-6))  # greater than negative nyquist freq
+        )
+        psd[doubled_freqs] *= 2
+
+    return psd
+
+
+@_validate_spectrum_inputs
+def compute_mean_fft(
     sig,
     interval_size,
     fs=None,

tests/test_power_spectral_density.py

@@ -22,23 +22,40 @@ def get_sorted_fft(data, fs):
     return fft_freq[order], fft[order]
 
 
-def test_compute_power_spectral_density():
+def get_periodogram(data, fs, full_range=False):
+    return_onesided = not full_range
+    f, p = signal.periodogram(
+        data, fs, return_onesided=return_onesided, detrend=False, axis=0
+    )
+    if p.ndim == 1:
+        p = p[:, np.newaxis]
+    if full_range:
+        order = np.argsort(f)
+        f = f[order]
+        p = p[order]
+    else:
+        f = f[:-1]
+        p = p[:-1]
+    return f, p
+
+
+def test_compute_fft():
 
     t = np.linspace(0, 1, 1000)
     sig = nap.Tsd(d=np.random.random(1000), t=t)
-    r = nap.compute_power_spectral_density(sig)
+    r = nap.compute_fft(sig)
     assert isinstance(r, pd.DataFrame)
     assert r.shape[0] == 500
 
     a, b = get_sorted_fft(sig.values, sig.rate)
     np.testing.assert_array_almost_equal(r.index.values, a[a >= 0])
     np.testing.assert_array_almost_equal(r.values, b[a >= 0])
 
-    r = nap.compute_power_spectral_density(sig, norm=True)
+    r = nap.compute_fft(sig, norm=True)
     np.testing.assert_array_almost_equal(r.values, b[a >= 0] / len(sig))
 
     sig = nap.TsdFrame(d=np.random.random((1000, 4)), t=t)
-    r = nap.compute_power_spectral_density(sig)
+    r = nap.compute_fft(sig)
     assert isinstance(r, pd.DataFrame)
     assert r.shape == (500, 4)
 
@@ -47,14 +64,57 @@ def test_compute_power_spectral_density():
     np.testing.assert_array_almost_equal(r.values, b[a >= 0])
 
     sig = nap.TsdFrame(d=np.random.random((1000, 4)), t=t)
-    r = nap.compute_power_spectral_density(sig, full_range=True)
+    r = nap.compute_fft(sig, full_range=True)
     assert isinstance(r, pd.DataFrame)
     assert r.shape == (1000, 4)
 
     a, b = get_sorted_fft(sig.values, sig.rate)
     np.testing.assert_array_almost_equal(r.index.values, a)
     np.testing.assert_array_almost_equal(r.values, b)
 
+    t = np.linspace(0, 1, 1000)
+    sig = nap.Tsd(d=np.random.random(1000), t=t)
+    r = nap.compute_fft(sig, ep=sig.time_support)
+    assert isinstance(r, pd.DataFrame)
+    assert r.shape[0] == 500
+
+    t = np.linspace(0, 1, 1000)
+    sig = nap.Tsd(d=np.random.random(1000), t=t)
+    r = nap.compute_fft(sig, fs=1000)
+    assert isinstance(r, pd.DataFrame)
+    assert r.shape[0] == 500
+
+
+def test_compute_power_spectral_density():
+
+    t = np.linspace(0, 1, 1000)
+    sig = nap.Tsd(d=np.random.random(1000), t=t)
+    r = nap.compute_power_spectral_density(sig)
+    assert isinstance(r, pd.DataFrame)
+    assert r.shape[0] == 500
+
+    a, b = get_periodogram(sig.values, sig.rate)
+    np.testing.assert_array_almost_equal(r.index.values, a[a >= 0])
+    np.testing.assert_array_almost_equal(r.values, b[a >= 0])
+
+    sig = nap.TsdFrame(d=np.random.random((1000, 4)), t=t)
+    r = nap.compute_power_spectral_density(sig)
+    assert isinstance(r, pd.DataFrame)
+    assert r.shape == (500, 4)
+
+    a, b = get_periodogram(sig.values, sig.rate)
+    np.testing.assert_array_almost_equal(r.index.values, a[a >= 0])
+    np.testing.assert_array_almost_equal(r.values, b[a >= 0])
+
+    sig = nap.TsdFrame(d=np.random.random((1000, 4)), t=t)
+    r = nap.compute_power_spectral_density(sig, full_range=True)
+    assert isinstance(r, pd.DataFrame)
+    assert r.shape == (1000, 4)
+
+    a, b = get_periodogram(sig.values, sig.rate, full_range=True)
+    np.testing.assert_array_almost_equal(r.index.values, a)
+    np.testing.assert_array_almost_equal(r.values, b)
+
     t = np.linspace(0, 1, 1000)
     sig = nap.Tsd(d=np.random.random(1000), t=t)
     r = nap.compute_power_spectral_density(sig, ep=sig.time_support)
@@ -135,9 +195,12 @@ def test_compute_power_spectral_density():
         ),
     ],
 )
-def test_compute_power_spectral_density_raise_errors(sig, kwargs, expectation):
+def test_compute_fft_raise_errors(sig, kwargs, expectation):
     with expectation:
-        nap.compute_power_spectral_density(sig, **kwargs)
+        nap.compute_fft(sig, **kwargs)
+    if "norm" not in kwargs:
+        with expectation:
+            nap.compute_power_spectral_density(sig, **kwargs)
 
 
 ############################################################
@@ -172,23 +235,23 @@ def get_signal_and_output(f=2, fs=1000, duration=100, interval_size=10, overlap=
     return (sig, out, freq)
 
 
-def test_compute_mean_power_spectral_density():
+def test_compute_mean_fft():
     sig, out, freq = get_signal_and_output()
-    psd = nap.compute_mean_power_spectral_density(sig, 10)
+    psd = nap.compute_mean_fft(sig, 10)
     assert isinstance(psd, pd.DataFrame)
     assert psd.shape[0] > 0  # Check that the psd DataFrame is not empty
     np.testing.assert_array_almost_equal(psd.values.flatten(), out[freq >= 0])
     np.testing.assert_array_almost_equal(psd.index.values, freq[freq >= 0])
 
     # Full range
-    psd = nap.compute_mean_power_spectral_density(sig, 10, full_range=True)
+    psd = nap.compute_mean_fft(sig, 10, full_range=True)
     assert isinstance(psd, pd.DataFrame)
     assert psd.shape[0] > 0  # Check that the psd DataFrame is not empty
     np.testing.assert_array_almost_equal(psd.values.flatten(), out)
     np.testing.assert_array_almost_equal(psd.index.values, freq)
 
     # Norm
-    psd = nap.compute_mean_power_spectral_density(sig, 10, norm=True)
+    psd = nap.compute_mean_fft(sig, 10, norm=True)
     assert isinstance(psd, pd.DataFrame)
     assert psd.shape[0] > 0  # Check that the psd DataFrame is not empty
     np.testing.assert_array_almost_equal(
@@ -200,7 +263,7 @@ def test_compute_mean_power_spectral_density():
     sig2 = nap.TsdFrame(
         t=sig.t, d=np.repeat(sig.values[:, None], 2, 1), time_support=sig.time_support
     )
-    psd = nap.compute_mean_power_spectral_density(sig2, 10, full_range=True)
+    psd = nap.compute_mean_fft(sig2, 10, full_range=True)
     assert isinstance(psd, pd.DataFrame)
     assert psd.shape[0] > 0  # Check that the psd DataFrame is not empty
     np.testing.assert_array_almost_equal(psd.values, np.repeat(out[:, None], 2, 1))
@@ -210,7 +273,7 @@ def test_compute_mean_power_spectral_density():
     sig2 = nap.TsdFrame(
         t=sig.t, d=np.repeat(sig.values[:, None], 2, 1), time_support=sig.time_support
     )
-    psd = nap.compute_mean_power_spectral_density(sig2, 10, full_range=True, fs=1000)
+    psd = nap.compute_mean_fft(sig2, 10, full_range=True, fs=1000)
     assert isinstance(psd, pd.DataFrame)
     assert psd.shape[0] > 0  # Check that the psd DataFrame is not empty
     np.testing.assert_array_almost_equal(psd.values, np.repeat(out[:, None], 2, 1))
@@ -329,8 +392,6 @@ def test_compute_mean_power_spectral_density():
         ),
     ],
 )
-def test_compute_mean_power_spectral_density_raise_errors(
-    sig, interval_size, kwargs, expectation
-):
+def test_compute_mean_fft_raise_errors(sig, interval_size, kwargs, expectation):
     with expectation:
-        nap.compute_mean_power_spectral_density(sig, interval_size, **kwargs)
+        nap.compute_mean_fft(sig, interval_size, **kwargs)