add column names for compute_wavelet_transform

pynapple-org · Jan 29, 2025 · 71a1f54 · 71a1f54
1 parent c961c66
commit 71a1f54
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/pynapple/process/wavelets.py b/pynapple/process/wavelets.py
@@ -127,7 +127,10 @@ def compute_wavelet_transform(
 
     if len(output_shape) == 2:
         return nap.TsdFrame(
-            t=sig.index, d=np.squeeze(cwt, axis=1), time_support=sig.time_support
+            t=sig.index,
+            d=np.squeeze(cwt, axis=1),
+            time_support=sig.time_support,
+            columns=freqs,
         )
     else:
         return nap.TsdTensor(

diff --git a/tests/test_signal_processing.py b/tests/test_signal_processing.py
@@ -397,6 +397,9 @@ def test_compute_wavelet_transform(
     np.testing.assert_array_almost_equal(
         mwt.time_support.values, sig.time_support.values
     )
+    if isinstance(mwt, nap.TsdFrame):
+        # test column names if TsdFrame
+        np.testing.assert_array_almost_equal(mwt.columns, freqs)
 
 
 @pytest.mark.parametrize(