Add interpolation functions and extract consecutive non nan blocks

navipy/trajectories/__init__.py

@@ -15,6 +15,8 @@ from mpl_toolkits.mplot3d import Axes3D  # noqa F401
 from multiprocessing import Pool
 from functools import partial
 import time
+from scipy import signal
+from scipy.interpolate import CubicSpline
 
 
 def _markers2position(x, kwargs):
@@ -250,7 +252,9 @@ class Trajectory(pd.DataFrame):
     def q_3(self, q_3):
         self.__set_q_i(3, q_3)
 
-    # overload of save/load function
+    # -------------------------------------------
+    # ---------------- IO -----------------------
+    # -------------------------------------------
     def read_csv(self, filename, sep=',', header=[0, 1], index_col=0):
         """ Load from a hdf file
         """
@@ -272,8 +276,10 @@ class Trajectory(pd.DataFrame):
         df = pd.DataFrame(self)
         df.to_csv(filename)
 
-    # initialisation from variables
-    def from_array(self, nparray, rotconv):
+    # -------------------------------------------
+    # ---------------- INITS FROM VAR------------
+    # -------------------------------------------
+    def from_array(self, nparray, rotconv, indeces=None):
         """ Assign trajectory from a numpy array
             N x 6 (rotconv = Euler angles)
             N x 7 (rotconv = quaternion)
@@ -283,7 +289,17 @@ class Trajectory(pd.DataFrame):
             msg = 'nparray should be a np.ndarray and not {}'
             msg = msg.format(type(nparray))
             raise TypeError(msg)
-        indeces = np.arange(0, nparray.shape[0])
+        if indeces is None:
+            indeces = np.arange(0, nparray.shape[0])
+        if not isinstance(indeces, np.ndarray):
+            msg = 'indeces should be a np.ndarray and not {}'
+            msg = msg.format(type(indeces))
+            raise TypeError(msg)
+        if indeces.shape[0] != nparray.shape[0]:
+            msg = 'indeces and nparray should have same number of rows'
+            msg += '{}!={}'
+            msg = msg.format(indeces.shape[0], nparray.shape[0])
+            raise TypeError(msg)
         if rotconv == 'quaternion':
             if nparray.shape[1] != 7:
                 msg = 'nparray should have size Nx7 and not {}'
@@ -463,6 +479,10 @@ class Trajectory(pd.DataFrame):
             self.loc[index_i, self.rotation_mode] = orientation
         return self
 
+    # -----------------------------------------------
+    # ---------------- TRANSFORM --------------------
+    # -----------------------------------------------
+
     def world2body(self, markers, indeces=None):
         """ Transform markers in world coordinate to body coordinate
         """
@@ -543,6 +563,108 @@ class Trajectory(pd.DataFrame):
                                    'dalpha_2']] = rot.squeeze()
         return velocity
 
+    # --------------------------------------------
+    # ---------------- FILTER --------------------
+    # --------------------------------------------
+    def filtfilt(self, order, cutoff, padlen=None):
+        """
+        Filter the trajectory with order and cutoff by
+        using a lowpass filter twice (forward and backward)
+        to correct for phase shift
+
+        :param order: the order of the lowpass filter. Either a number \
+or a pandas series. The series should be multiindexed as the columns of \
+the trajectory.
+        :param cutoff: cut off frequency in Hz if sampling rate is known\
+otherwise relative to the Nyquist frequency. Either a number or a pandas \
+series.
+        """
+        if isinstance(order, [int, float]):
+            order = pd.Series(data=order, index=self.columns)
+        if isinstance(cutoff, [int, float]):
+            cutoff = pd.Series(data=cutoff, index=self.columns)
+        if not np.isnan(self.__sampling_rate):
+            nyquist = self.__sampling_rate/2
+            cutoff /= nyquist
+
+        subtraj = self.consecutive_blocks()
+        for trajno_nan in enumerate(subtraj):
+            indeces = trajno_nan.index
+            for col in self.columns:
+                b, a = signal.butter(order.loc[col], cutoff.loc[col])
+                if padlen is None:
+                    padlen = 3*max(len(a), len(b))
+                if trajno_nan.shape[0] < padlen:
+                    self.loc[indeces, col] *= np.nan
+                else:
+                    self.loc[indeces, col] = signal.filtfilt(
+                        b, a,
+                        trajno_nan.loc[:, col],
+                        padlen=padlen)
+
+    def fillna(self, method='Cubic'):
+        """ fillna with a given method
+        """
+        customs_method = ['Cubic']
+        if not (method in customs_method):
+            # fall back to pandas fillna function
+            return self.fillna(method)
+        # Start implementing customs_method
+        if method == 'Cubic':
+            for col in self.loc[:, 'location'].columns:
+                values = self.loc[:, ('location', col)]
+                validtime = values.dropna().index
+                validvalues = values.dropna().values
+                cs = CubicSpline(validtime, validvalues)
+                time = self.index
+                self.loc[:, ('location', col)] = cs(time)
+            # for the angles we first do a ffill and then
+            # unwrap and interpolate on the unwrap angles
+            rotconv = self.rotation_mode
+            for col in self.loc[:, rotconv].columns:
+                values = self.loc[:, (rotconv, col)]
+                validtime = values.dropna().index
+                unwrapvalues = np.unwrap(values.fillna(method='ffill'))
+                validvalues = unwrapvalues[validtime]
+                cs = CubicSpline(validtime, validvalues)
+                time = self.index
+                self.loc[:, (rotconv, col)] = cs(time)
+            return self
+
+        else:
+            msg = 'Method {} is not supported.'
+            msg += 'please use method supported by pd.fillna'
+            msg += ' or one of the following methods {}'
+            msg = msg.format(method, customs_method)
+            raise NameError(msg)
+
+    # --------------------------------------------
+    # ---------------- EXTRACT -------------------
+    # --------------------------------------------
+
+    def consecutive_blocks(self):
+        """ Return a list of subtrajectory withtout nans
+        """
+        # get a numpy array from the trajectory,
+        # because we are using numpy arrays later
+        np_traj = self.values
+        np_traj = np.hstack([self.index[:, np.newaxis], np_traj])
+        # Look for row containing at least one nan
+        nonans = np.any(np.isnan(np_traj), axis=1)
+        # spliting the trajectory according to nan location
+        events = np.split(np_traj, np.where(nonans)[0])
+
+        # removing NaN entries
+        events = [ev[~np.any(np.isnan(ev), axis=1)]
+                  for ev in events if isinstance(ev, np.ndarray)]
+        # removing empty DataFrames
+        subtraj = [Trajectory().from_dataframe(self.loc[ev[:, 0]])
+                   for ev in events if ev.size > 0]
+        return subtraj
+    # -------------------------------------------
+    # ---------------- PLOTS --------------------
+    # -------------------------------------------
+
     def lollipops(self, ax=None,
                   colors=None, step_lollipop=1,
                   offset_lollipop=0, lollipop_marker='o',