Fix tests

Author: robertkleef

dedisperse/dedisperse.py

@@ -4,7 +4,7 @@
 # pylint: disable-msg=C0103
 import numpy as np
 
-def dedisperse(samples, highest_x=None, max_delay=None, dm=None):
+def dedisperse(samples, highest_x=None, max_delay=0, dm=None):
     '''
     This method performs dedispersion on the filterbank data
     The maximum_delay specifies between the currently considered pulsar signal and the next pulsar
@@ -19,22 +19,23 @@ def dedisperse(samples, highest_x=None, max_delay=None, dm=None):
         pulsar_intensity = find_estimation_intensity(samples, highest_x)
         dm = find_dm(samples, pulsar_intensity, max_delay)
 
-    # Distribute the DM over the amount of samples
-    delays_per_sample = np.round(np.linspace(dm, 0, samples.shape[1])).astype(int)
+    if dm:
+        # Distribute the DM over the amount of samples
+        delays_per_sample = (np.round(np.linspace(dm, 0, samples.shape[1]))).astype(int)
 
-    # Loop over the frequencies
-    for i, _ in enumerate(delays_per_sample):
+        # Loop over the frequencies
+        for i, _ in enumerate(delays_per_sample):
 
-        # Temporary array that is used to later delay the frequency
-        temporary_samples = []
+            # Temporary array that is used to later delay the frequency
+            temporary_samples = []
 
-        # Select frequency/column 'i' of all samples
-        temporary_samples = samples[:, i]
+            # Select frequency/column 'i' of all samples
+            temporary_samples = samples[:, i]
 
-        # Write back the frequency/column 'i' to the samples, using numpy's roll function
-        samples[:, i] = np.roll(temporary_samples, delays_per_sample[i])
+            # Write back the frequency/column 'i' to the samples, using numpy's roll function
+            samples[:, i] = np.roll(temporary_samples, delays_per_sample[i])
 
-    return samples
+        return samples
 
 def find_dm(samples, pulsar_intensity, max_delay):
     '''

examples/dedisperse_stream.py

@@ -14,8 +14,6 @@
 import dedisperse as dedisperse
 from plot.static_waterfall import waterfall_plot
 
-from timeseries.timeseries import Timeseries
-
 from clipping import clipping
 
 # Read filterbank data,

tests/test_dedisperse.py

@@ -3,6 +3,7 @@
 """
 
 import unittest
+import numpy as np 
 
 from .context import dedisperse # pylint: disable-msg=E0611
 
@@ -25,5 +26,5 @@ def test_dedisperse(self):
             expect moved frequencies per sample
         """
         disp_measure = 6
-        results = dedisperse.dedisperse(self.samples, disp_measure)
+        results = dedisperse.dedisperse(np.array(self.samples), None, None, disp_measure)
         self.assertListEqual(list(results[len(self.samples)-1]), [10]*7)