TravelSignal: extract time_offset into method

simulations/lib/travelsignal.py

@@ -11,7 +11,7 @@ class TravelSignal:
     Model an arbitrary digitised signal that can be translated to another position and time.
     """
 
-    def __init__(self, signal, sample_rate, t_0 = 0, x_0 = 0, periodic=True, interp1d_kw = None):
+    def __init__(self, signal, sample_rate, t_0 = 0, x_0 = 0, periodic=True, interp1d_kw = None, velocity=None):
         """
         Initialise by saving the raw signal
 
@@ -32,17 +32,21 @@ class TravelSignal:
             Use scipy.interpolate's interp1d_kw for interpolation.
             Set to True, or a dictionary to enable.
             Dictionary will be entered in as **kwargs.
+        velocity : float, optional
+            Defaults to the speed of light in m/s.
         """
 
-        self.raw = signal
+        self.raw = np.asarray(signal)
         self.periodic = periodic
 
         self.sample_rate = sample_rate # Hz
         self.sample_length = len(self.raw)
         self.time_length = self.sample_length*sample_rate # s
 
-        self.x_0 = x_0
+        self.velocity = 299792458 if velocity is None else velocity # m / s
-        self.t_0 = t_0
+
+        self.x_0 = x_0 # m
+        self.t_0 = t_0 # s
 
         # choose interpolation method
         if not interp1d_kw:
@@ -87,36 +91,7 @@ class TravelSignal:
         Returns the signal at (t_f, x_f)
         """
 
-        if t_0 is None:
+        total_time_offset = self.total_time_offset(t_f=t_f, x_f=x_f, t_0=t_0, x_0=x_0, velocity=velocity)
-            t_0 = self.t_0
-
-        if velocity is None:
-            velocity = 1
-
-
-        ## spatial offset
-        if x_f is None:
-            spatial_time_offset = 0
-        else:
-            x_f = np.asarray(x_f)
-            if x_0 is None:
-                x_0 = self.x_0
-
-            spatial_time_offset = np.sum(np.sqrt( (x_f - x_0)**2 )/velocity)
-
-        ## temporal offset
-        if t_f is None:
-            temporal_time_offset = 0
-        else:
-            t_f = np.asarray(t_f)
-
-            if t_0 is None:
-                t_0 = self.t_0
-
-            temporal_time_offset = t_f - t_0
-
-        # total offset
-        total_time_offset = temporal_time_offset - spatial_time_offset
         n_offset = (total_time_offset * self.sample_rate )
 
         # periodic signal
@@ -169,6 +144,52 @@ class TravelSignal:
 
         return amplitude, total_time_offset
 
+    def spatial_time(self, x_f, velocity=None, x_0=None):
+        """
+        Calculate the time offset caused by a spatial difference.
+        """
+        if velocity is None:
+            velocity = self.velocity
+
+        if x_0 is None:
+            x_0 = self.x_0
+
+        return np.sum(np.sqrt( (x_f - x_0)**2 )/velocity)
+
+    def total_time_offset(self, t_f = None, x_f = None, t_0 = None, x_0 = None, velocity = None):
+        """
+        Calculate how much time shifting is needed to go from (t_0, x_0) to (t_f, x_f).
+
+        Convention:
+            (t_0, x_0) < (t_f, x_0) gives a positive time shift,
+            (t_0, x_0) != (t_0, x_f) gives a negative time shift
+
+        Returns:
+         the time shift
+        """
+
+        ## spatial offset
+        if x_f is None:
+            spatial_time_offset = 0
+        else:
+            x_f = np.asarray(x_f)
+            if x_0 is None:
+                x_0 = self.x_0
+
+            spatial_time_offset = self.spatial_time(x_f, x_0=x_0, velocity=velocity)
+
+        ## temporal offset
+        if t_f is None:
+            temporal_time_offset = 0
+        else:
+            t_f = np.asarray(t_f)
+
+            if t_0 is None:
+                t_0 = self.t_0
+
+            temporal_time_offset = t_f - t_0
+
+        return temporal_time_offset - spatial_time_offset
 
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
@@ -192,9 +213,9 @@ if __name__ == "__main__":
     ax.set_ylabel("Amplitude")
     ax.set_xlabel("Time")
 
-    ax.plot(time,   signal,                    label='Raw signal')
+    ax.plot(time,   signal,                     label='Raw signal')
-    ax.plot(time2,  mysignal(time2)+0.5,     '.-',     label='TravelSignal(periodic)+0.5')
+    ax.plot(time2,  mysignal(time2) +0.5, '.-', label='TravelSignal(periodic)+0.5')
-    ax.plot(time2,  mysignal2(time2)-0.5,     '.-',     label='TravelSignal-0.5')
+    ax.plot(time2,  mysignal2(time2)-0.5, '.-', label='TravelSignal-0.5')
 
     ax.legend()