keep array_annotations in the output of signal_processing functions (N…

…euralEnsemble#258)

* pass array_annotations on to output of signal_processing functions

* test keep array_annotations for zscore, butter, hilbert

elephant/signal_processing.py

@@ -126,6 +126,9 @@ def zscore(signal, inplace=True):
             sig_normalized = sig
         else:
             sig_normalized = sig.duplicate_with_new_data(sig_normalized)
+            # todo use flag once is fixed
+            #      https://github.com/NeuralEnsemble/python-neo/issues/752
+            sig_normalized.array_annotate(**sig.array_annotations)
         sig_dimless = sig_normalized / sig.units
         result.append(sig_dimless)
 
@@ -142,16 +145,16 @@ def cross_correlation_function(signal, ch_pairs, env=False, nlags=None):
     Computes unbiased estimator of the cross-correlation function.
 
     Calculates the unbiased estimator of the cross-correlation function [1]_
-    
+
     .. math::
              R(\\tau) = \\frac{1}{N-|k|} R'(\\tau) \\ ,
-    
-    where :math:`R'(\\tau) = \\left<x(t)y(t+\\tau)\\right>` in a pairwise 
+
+    where :math:`R'(\\tau) = \\left<x(t)y(t+\\tau)\\right>` in a pairwise
     manner, i.e. `signal[ch_pairs[0,0]]` vs `signal2[ch_pairs[0,1]]`,
     `signal[ch_pairs[1,0]]` vs `signal2[ch_pairs[1,1]]`, and so on. The
     cross-correlation function is obtained by `scipy.signal.fftconvolve`.
     Time series in signal are zscored beforehand. Alternatively returns the
-    Hilbert envelope of :math:`R(\\tau)`, which is useful to determine the 
+    Hilbert envelope of :math:`R(\\tau)`, which is useful to determine the
     correlation length of oscillatory signals.
 
     Parameters
@@ -394,7 +397,11 @@ def butter(signal, highpass_freq=None, lowpass_freq=None, order=4,
 
     if isinstance(signal, neo.AnalogSignal):
         filtered_data = np.rollaxis(filtered_data, -1, 0)
-        return signal.duplicate_with_new_data(filtered_data)
+        signal_out = signal.duplicate_with_new_data(filtered_data)
+        # todo use flag once is fixed
+        #      https://github.com/NeuralEnsemble/python-neo/issues/752
+        signal_out.array_annotate(**signal.array_annotations)
+        return signal_out
     elif isinstance(signal, pq.quantity.Quantity):
         return filtered_data * signal.units
     else:
@@ -628,6 +635,9 @@ def hilbert(signal, N='nextpow'):
 
     output = signal.duplicate_with_new_data(
         scipy.signal.hilbert(signal.magnitude, N=n, axis=0)[:n_org])
+    # todo use flag once is fixed
+    #      https://github.com/NeuralEnsemble/python-neo/issues/752
+    output.array_annotate(**signal.array_annotations)
     return output / output.units
 
 

elephant/test/test_signal_processing.py

@@ -195,6 +195,15 @@ def test_zscore_single_multidim_dup(self):
         # Assert original signal is untouched
         self.assertEqual(signal[0, 0].magnitude, self.test_seq1[0])
 
+    def test_zscore_array_annotations(self):
+        signal = neo.AnalogSignal(
+            self.test_seq1, units='mV',
+            t_start=0. * pq.ms, sampling_rate=1000. * pq.Hz,
+            array_annotations=dict(valid=True, my_list=[0]))
+        zscored = elephant.signal_processing.zscore(signal, inplace=False)
+        self.assertDictEqual(signal.array_annotations,
+                             zscored.array_annotations)
+
     def test_zscore_single_multidim_inplace(self):
         """
         Test z-score on a single AnalogSignal with multiple dimensions, asking
@@ -391,7 +400,8 @@ def test_butter_filter_function(self):
         # generate white noise AnalogSignal
         noise = neo.AnalogSignal(
             np.random.normal(size=5000),
-            sampling_rate=1000 * pq.Hz, units='mV')
+            sampling_rate=1000 * pq.Hz, units='mV',
+            array_annotations=dict(valid=True, my_list=[0]))
 
         kwds = {'signal': noise, 'highpass_freq': 250.0 * pq.Hz,
                 'lowpass_freq': None, 'filter_function': 'filtfilt'}
@@ -412,6 +422,10 @@ def test_butter_filter_function(self):
         self.assertAlmostEqual(psd_filtfilt[0, 0], psd_lfilter[0, 0])
         self.assertAlmostEqual(psd_filtfilt[0, 0], psd_sosfiltfilt[0, 0])
 
+        # Test if array_annotations are preserved
+        self.assertDictEqual(noise.array_annotations,
+                             filtered_noise.array_annotations)
+
     def test_butter_invalid_filter_function(self):
         # generate a dummy AnalogSignal
         anasig_dummy = neo.AnalogSignal(
@@ -521,11 +535,13 @@ def setUp(self):
             self.amplitude[:, 2] * np.cos(self.phase[:, 2]),
             self.amplitude[:, 3] * np.cos(self.phase[:, 3])])
 
+        array_annotations = dict(my_list=np.arange(sigs.shape[0]))
         self.long_signals = neo.AnalogSignal(
             sigs.T, units='mV',
             t_start=0. * pq.ms,
             sampling_rate=(len(time) / (time[-1] - time[0])).rescale(pq.Hz),
-            dtype=float)
+            dtype=float,
+            array_annotations=array_annotations)
 
         # Generate test data covering a single oscillation cycle in 1s only
         phases = np.arange(0, 2 * np.pi, np.pi / 256)
@@ -564,6 +580,14 @@ def test_hilbert_output_shape(self):
         self.assertEqual(np.shape(output), true_shape)
         self.assertEqual(output.units, pq.dimensionless)
 
+    def test_hilbert_array_annotations(self):
+        output = elephant.signal_processing.hilbert(self.long_signals,
+                                                    N='nextpow')
+        # Test if array_annotations are preserved
+        self.assertSetEqual(set(output.array_annotations.keys()), {"my_list"})
+        assert_array_equal(output.array_annotations['my_list'],
+                           self.long_signals.array_annotations['my_list'])
+
     def test_hilbert_theoretical_long_signals(self):
         """
         Tests the output of the hilbert function with regard to amplitude and