Add the ability to calculate power spectra directly from binned spike train objects

elephant/conversion.py

@@ -999,18 +999,52 @@ def to_bool_array(self):
         """
         return self.to_array(dtype=bool)
 
-    def to_array(self, dtype=None):
+    def to_bool_analogsignal(self):
+        """
+        Converts the binned spike train to a neo.AnalogSignal object with binary values.
+
+        Returns
+        -------
+        neo.AnalogSignal
+            A boolean analog signal where each time bin contains either 0 (no spike)
+            or 1 (one or more spikes).
+
+        See Also
+        --------
+        to_bool_array : Returns the binary representation as a NumPy array.
+        """
+        import neo
+        bool_arr = self.to_bool_array()
+        # Create an AnalogSignal with the same time base as the binned spike train
+        signal = neo.AnalogSignal(
+            bool_arr.astype(float).T,  # Transpose to match AnalogSignal convention
+            units=pq.dimensionless,  # Binary signal has no physical units
+            sampling_period=self.bin_size,
+            t_start=self.t_start
+        )
+        return signal
+
+    def to_array(self, dtype=None, scaling="counts"):
         """
         Returns a dense matrix, calculated from the sparse matrix, with counted
         time points of spikes. The rows correspond to spike trains and the
         columns correspond to bins in a `BinnedSpikeTrain`.
         Entries contain the count of spikes that occurred in the given bin of
         the given spike train.
 
+        Parameters
+        ----------
+        dtype : type, optional
+            The desired data-type for the array. Default is None.
+        scaling : {"counts", "normalized"}, optional
+            Determines the scaling of the returned array:
+            * "counts" : raw spike counts (default)
+            * "normalized" : spike counts divided by the bin size in seconds
+
         Returns
         -------
         matrix : np.ndarray
-            Matrix with spike counts. Columns represent the index positions of
+            Matrix with spike counts or rates. Columns represent the index positions of
             the binned spikes and rows represent the spike trains.
 
         Examples
@@ -1022,20 +1056,106 @@ def to_array(self, dtype=None):
         ...                  t_stop=10.0 * pq.s)
         >>> x = conv.BinnedSpikeTrain(a, n_bins=10, bin_size=1 * pq.s,
         ...                           t_start=0 * pq.s)
-        >>> print(x.to_array())
+        >>> print(x.to_array())  # counts
         [[2 1 0 1 1 1 1 0 0 0]]
+        >>> print(x.to_array(scaling="normalized"))  # counts/second
+        [[2. 1. 0. 1. 1. 1. 1. 0. 0. 0.]]
 
         See also
         --------
         scipy.sparse.csr_matrix
         scipy.sparse.csr_matrix.toarray
 
+        Raises
+        ------
+        ValueError
+            If the scaling parameter is not one of {"counts", "normalized"}.
         """
+        if scaling not in ("counts", "normalized"):
+            raise ValueError(
+                f"Invalid scaling parameter: {scaling}. "
+                "Available options: 'counts', 'normalized'"
+            )
+
         array = self.sparse_matrix.toarray()
         if dtype is not None:
             array = array.astype(dtype)
+
+        if scaling == "normalized":
+            # Convert bin_size to seconds for normalization
+            bin_size_seconds = self.bin_size.rescale('s').magnitude
+            array = array / bin_size_seconds
+
         return array
 
+    def to_analog_signal(self, dtype=None, scaling="counts"):
+        """
+        Returns the binned spike train as a neo.AnalogSignal. The signal values
+        represent either spike counts or normalized spike rates depending on the
+        scaling parameter.
+
+        Parameters
+        ----------
+        dtype : type, optional
+            The desired data-type for the signal values. Default is None.
+        scaling : {"counts", "normalized"}, optional
+            Determines the scaling of the returned signal:
+            * "counts" : raw spike counts (default)
+            * "normalized" : spike counts divided by the bin size in seconds,
+              resulting in firing rates in Hz
+
+        Returns
+        -------
+        neo.AnalogSignal
+            Signal containing spike counts or rates. Rows represent time bins
+            and columns represent different spike trains.
+            If scaling="counts", the signal has units of 1/b in Hz, where b is
+            the bin size.
+            If scaling="normalized", the signal has units of Hz (spikes/second).
+
+        Examples
+        --------
+        >>> import elephant.conversion as conv
+        >>> import neo as n
+        >>> import quantities as pq
+        >>> a = n.SpikeTrain([0.5, 0.7, 1.2, 3.1, 4.3, 5.5, 6.7] * pq.s,
+        ...                  t_stop=10.0 * pq.s)
+        >>> x = conv.BinnedSpikeTrain(a, n_bins=10, bin_size=1 * pq.s,
+        ...                           t_start=0 * pq.s)
+        >>> signal = x.to_analog_signal()  # counts (1/bin_size) Hz
+        >>> signal = x.to_analog_signal(scaling="normalized")  # rates (Hz)
+
+        See also
+        --------
+        to_array : Returns the binned spike train as a plain NumPy array
+        neo.AnalogSignal : The neo analog signal object
+
+        Raises
+        ------
+        ValueError
+            If the scaling parameter is not one of {"counts", "normalized"}.
+        """
+        # Get the array representation with the desired scaling
+        array = self.to_array(dtype=dtype, scaling=scaling)
+
+        # Determine the appropriate units based on scaling
+        if scaling == "normalized":
+            units = pq.Hz
+        else:  # counts
+            units = pq.CompoundUnit(f"1.0 / {self.bin_size.rescale('s').magnitude} * Hz")
+
+        # Create the AnalogSignal
+        # Note: AnalogSignal expects shape (time_steps, channels),
+        # so we transpose the array which has shape (channels, time_steps)
+        signal = neo.AnalogSignal(
+            array.T,
+            units=units,
+            sampling_period=self.bin_size,
+            t_start=self.t_start
+        )
+
+        return signal
+
     def binarize(self, copy=True):
         """
         Clip the internal array (no. of spikes in a bin) to `0` (no spikes) or

elephant/spectral.py

@@ -20,15 +20,17 @@
 
 import warnings
 
-import neo
 import numpy as np
 import quantities as pq
 import scipy.signal
+import neo
+import elephant.conversion
 
 __all__ = [
     "welch_psd",
     "welch_coherence",
     "multitaper_psd",
+    "segmented_multitaper_psd",
     "multitaper_cross_spectrum",
     "segmented_multitaper_cross_spectrum",
     "multitaper_coherence"
@@ -62,7 +64,7 @@ def welch_psd(signal, n_segments=8, len_segment=None,
 
     Parameters
     ----------
-    signal : neo.AnalogSignal or pq.Quantity or np.ndarray
+    signal : neo.AnalogSignal or pq.Quantity or np.ndarray or elephant.conversion.BinnedSpikeTrain
         Time series data, of which PSD is estimated. When `signal` is
         `pq.Quantity` or `np.ndarray`, sampling frequency should be given
         through the keyword argument `fs`. Otherwise, the default value is
@@ -201,6 +203,10 @@ def welch_psd(signal, n_segments=8, len_segment=None,
 
 
     """
+    if isinstance(signal,elephant.conversion.BinnedSpikeTrain):
+        # Convert spike train to an analog signal with units of Hz
+        signal = signal.to_analog_signal(scaling="normalized")
+
     # 'hanning' window was removed with release of scipy 1.9.0, it was
     # deprecated since 1.1.0.
     if window == 'hanning':
@@ -299,7 +305,7 @@ def multitaper_psd(signal, fs=1, nw=4, num_tapers=None, peak_resolution=None,
 
     Parameters
     ----------
-    signal : neo.AnalogSignal or pq.Quantity or np.ndarray
+    signal : neo.AnalogSignal or pq.Quantity or np.ndarray or elephant.conversion.BinnedSpikeTrain
         Time series data of which PSD is estimated. When `signal` is np.ndarray
         sampling frequency should be given through keyword argument `fs`.
         Signal should be passed as (n_channels, n_samples)
@@ -351,6 +357,13 @@ def multitaper_psd(signal, fs=1, nw=4, num_tapers=None, peak_resolution=None,
 
     # When the input is AnalogSignal, the data is added after rolling the axis
     # for time index to the last
+    if isinstance(signal,elephant.conversion.BinnedSpikeTrain):
+        # signal = neo.AnalogSignal(
+        #     signal.to_array().transpose()/signal.bin_size.rescale(pq.s).magnitude*pq.dimensionless,
+        #     t_start=signal.t_start,
+        #     sampling_period=signal.bin_size)
+        signal = signal.to_analog_signal(scaling="normalized")
+
     data = np.asarray(signal)
     if isinstance(signal, neo.AnalogSignal):
         data = np.moveaxis(data, 0, 1)
@@ -521,6 +534,13 @@ def segmented_multitaper_psd(signal, n_segments=1, len_segment=None,
     TypeError
         If `peak_resolution` is None and `num_tapers` is not an int.
     """
+    if isinstance(signal,elephant.conversion.BinnedSpikeTrain):
+        # signal = neo.AnalogSignal(
+        #     signal.to_array().transpose()/signal.bin_size.rescale(pq.s).magnitude*pq.dimensionless,
+        #     t_start=signal.t_start,
+        #     sampling_period=signal.bin_size)
+        signal = signal.to_analog_signal(scaling="normalized")
+
 
     # When the input is AnalogSignal, the data is added after rolling the axis
     # for time index to the last

elephant/test/test_spectral.py

@@ -149,6 +149,40 @@ def test_welch_psd_input_types(self):
             (freqs_neo == freqs_np).all() and (
                     psd_neo == psd_np).all())
 
+    def test_welch_psd_binned_spiketrain(self):
+        rate = 50 * pq.Hz
+        t_stop = 10 * pq.s
+        spiketrain = elephant.spike_train_generation.StationaryPoissonProcess(
+            rate, t_stop=t_stop).generate_spiketrain()
+        bin_size = 2 * pq.ms
+        binned_st = elephant.conversion.BinnedSpikeTrain(
+            spiketrain, bin_size=bin_size)
+
+        freqs, psd = elephant.spectral.welch_psd(binned_st)
+
+        # Check units: should be Hz (from Hz^2/Hz)
+        self.assertEqual(freqs.units, pq.Hz)
+        self.assertEqual(psd.units, pq.Hz)
+
+        # Check scaling: for large frequencies, PSD should approach 2*rate
+        # (factor 2 because it is a one-sided PSD)
+        avg_psd = np.mean(psd[0, freqs > 100 * pq.Hz])
+        self.assertAlmostEqual(avg_psd.rescale('Hz').magnitude,
+                               2*rate.rescale('Hz').magnitude,
+                               delta=0.2 * 2 * rate.rescale('Hz').magnitude)
+
+    def test_welch_psd_binned_spiketrain_empty(self):
+        """
+        Test Welch's PSD estimation for an empty binned spiketrain.
+        """
+        # Check behavior for an empty spiketrain
+        bin_size = 2 * pq.ms
+        empty_st = neo.SpikeTrain([] * pq.s, t_start=0 * pq.s, t_stop=10 * pq.s)
+        binned_empty = elephant.conversion.BinnedSpikeTrain(
+            empty_st, bin_size=bin_size)
+        freqs_empty, psd_empty = elephant.spectral.welch_psd(binned_empty)
+        self.assertEqual(np.max(psd_empty), 0 * pq.Hz)
+
     def test_welch_psd_multidim_input(self):
         # generate multidimensional data
         num_channel = 4
@@ -341,6 +375,37 @@ def test_multitaper_psd_input_types(self):
             (freqs_neo == freqs_np).all() and (
                     psd_neo == psd_np).all())
 
+    def test_multitaper_psd_binned_spiketrain(self):
+        rate = 50 * pq.Hz
+        t_stop = 10 * pq.s
+        spiketrain = elephant.spike_train_generation.StationaryPoissonProcess(
+            rate, t_stop=t_stop).generate_spiketrain()
+        bin_size = 2 * pq.ms
+        binned_st = elephant.conversion.BinnedSpikeTrain(
+            spiketrain, bin_size=bin_size)
+
+        freqs, psd = elephant.spectral.multitaper_psd(binned_st)
+
+        self.assertEqual(freqs.units, pq.Hz)
+        self.assertEqual(psd.units, pq.Hz)
+
+        avg_psd = np.mean(psd[0, freqs > 100 * pq.Hz])
+        self.assertAlmostEqual(avg_psd.rescale('Hz').magnitude,
+                               2 * rate.rescale('Hz').magnitude,
+                               delta=0.2 * 2 * rate.rescale('Hz').magnitude)
+
+    def test_multitaper_psd_binned_spiketrain_empty(self):
+        """
+        Test Multitaper PSD estimation for an empty binned spiketrain.
+        """
+        # Check behavior for an empty spiketrain
+        bin_size = 2 * pq.ms
+        empty_st = neo.SpikeTrain([] * pq.s, t_start=0 * pq.s, t_stop=10 * pq.s)
+        binned_empty = elephant.conversion.BinnedSpikeTrain(
+            empty_st, bin_size=bin_size)
+        freqs_empty, psd_empty = elephant.spectral.multitaper_psd(binned_empty)
+        self.assertEqual(np.max(psd_empty), 0 * pq.Hz)
+
 
 class SegmentedMultitaperPSDTestCase(unittest.TestCase):
     # The following assertions test _segmented_apply_func in the context
@@ -506,6 +571,40 @@ def test_segmented_multitaper_psd_input_types(self):
             (freqs_neo == freqs_np).all() and (
                     psd_neo == psd_np).all())
 
+    def test_segmented_multitaper_psd_binned_spiketrain(self):
+        rate = 50 * pq.Hz
+        t_stop = 10 * pq.s
+        spiketrain = elephant.spike_train_generation.StationaryPoissonProcess(
+            rate, t_stop=t_stop).generate_spiketrain()
+        bin_size = 2 * pq.ms
+        binned_st = elephant.conversion.BinnedSpikeTrain(
+            spiketrain, bin_size=bin_size)
+
+        freqs, psd = elephant.spectral.segmented_multitaper_psd(binned_st)
+
+        # Check units: should be Hz
+        self.assertEqual(freqs.units, pq.Hz)
+        self.assertEqual(psd.units, pq.Hz)
+
+        # Check scaling: for large frequencies, PSD should approach 2*rate
+        # Flattening because segmented_multitaper_psd might have extra dimension
+        avg_psd = np.mean(psd.flatten()[freqs > 100 * pq.Hz])
+        self.assertAlmostEqual(avg_psd.rescale('Hz').magnitude,
+                               2 * rate.rescale('Hz').magnitude,
+                               delta=0.2 * 2 * rate.rescale('Hz').magnitude)
+
+    def test_segmented_multitaper_psd_binned_spiketrain_empty(self):
+        """
+        Test Segmented Multitaper PSD estimation for an empty binned spiketrain.
+        """
+        # Check behavior for an empty spiketrain
+        bin_size = 2 * pq.ms
+        empty_st = neo.SpikeTrain([] * pq.s, t_start=0 * pq.s, t_stop=10 * pq.s)
+        binned_empty = elephant.conversion.BinnedSpikeTrain(
+            empty_st, bin_size=bin_size)
+        freqs_empty, psd_empty = elephant.spectral.multitaper_psd(binned_empty)
+        self.assertEqual(np.max(psd_empty), 0 * pq.Hz)
+
 
 class MultitaperCrossSpectrumTestCase(unittest.TestCase):
     def test_multitaper_cross_spectrum_errors(self):