diff --git a/navipy/trajectories/__init__.py b/navipy/trajectories/__init__.py index 924b6d87a4aeb13247b23594613149c6deacdf81..aff0b8d4b508eb79dc123c770a42db9527a61e1f 100644 --- a/navipy/trajectories/__init__.py +++ b/navipy/trajectories/__init__.py @@ -17,6 +17,7 @@ from functools import partial import time from scipy import signal from scipy.interpolate import CubicSpline +import warnings def _markers2position(x, kwargs): @@ -51,6 +52,18 @@ def _markerstransform(index_i, trajectory, return tmarker.transpose().flatten() +def _invmarkerstransform(index_i, trajectory, + homogeneous_markers, rotation_mode): + row = trajectory.loc[index_i] + angles = row.loc[rotation_mode].values + translate = row.loc['location'].values + trans_mat = htf.compose_matrix(angles=angles, + translate=translate, + axes=rotation_mode) + tmarker = np.linalg.inv(trans_mat).dot(homogeneous_markers)[:3] + return tmarker.transpose().flatten() + + class Trajectory(pd.DataFrame): def __init__(self, rotconv='rzyx', indeces=np.arange(1)): columns = self.__build_columns(rotconv) @@ -479,9 +492,8 @@ class Trajectory(pd.DataFrame): # ----------------------------------------------- # ---------------- TRANSFORM -------------------- # ----------------------------------------------- - - def world2body(self, markers, indeces=None): - """ Transform markers in world coordinate to body coordinate + def body2world(self, markers, indeces=None): + """ Transform markers in body coordinate to world coordinate """ if not isinstance(markers, pd.Series): msg = 'markers should be of type pd.Series and not' @@ -516,16 +528,44 @@ class Trajectory(pd.DataFrame): dtype=float) return transformed_markers - for index_i in indeces: - row = self.loc[index_i] - angles = row.loc[self.rotation_mode].values - translate = row.loc['location'].values - trans_mat = htf.compose_matrix(angles=angles, - translate=translate, - axes=self.rotation_mode) - tmarker = trans_mat.dot(homogeneous_markers.transpose())[:3] - transformed_markers.loc[index_i, :] = tmarker.transpose().flatten() - + def world2body(self, markers, indeces=None): + """ Transform markers in world coordinate to body coordinate + """ + msg = 'Prior to 12/09/2018:\n' + msg += 'world2body was doing a reverse transformed\n' + msg += 'Please use body2world instead' + warnings.warn(msg) + if not isinstance(markers, pd.Series): + msg = 'markers should be of type pd.Series and not' + msg += ' {}'.format(type(markers)) + raise TypeError(msg) + if not isinstance(markers.index, pd.MultiIndex): + msg = 'markers should have a multiIndex index \n' + msg += ' (i,"x"), (i,"y"),(i,"z")\n' + msg += 'here i is the index of the marker' + raise TypeError(msg) + if indeces is None: + # Looping through each time point along the trajectory + indeces = self.index + # More than one marker may be transformed + # The marker are assume to be a multiIndex dataframe + homogeneous_markers = markers.unstack() + homogeneous_markers['w'] = 1 + homogeneous_markers = homogeneous_markers.transpose() + # Looping throught the indeces + # to get the homogeneous transformation from the position orientation + # and then apply the transformed to the marker position + with Pool() as p: + result = p.map( + partial(_invmarkerstransform, + trajectory=self, + homogeneous_markers=homogeneous_markers, + rotation_mode=self.rotation_mode), + indeces) + transformed_markers = pd.DataFrame(data=result, + index=indeces, + columns=markers.index, + dtype=float) return transformed_markers def velocity(self): @@ -538,8 +578,7 @@ class Trajectory(pd.DataFrame): 'dalpha_2'], dtype=float) diffrow = self.diff() - velocity.loc[:, ['dx', 'dy', 'dz']] = \ - diffrow.loc[:, 'location'].values + velocity.loc[:, ['dx', 'dy', 'dz']] = diffrow.loc[:, 'location'].values for index_i, row in self.iterrows(): if self.rotation_mode == 'quaternion': raise NameError('Not implemented') @@ -742,7 +781,7 @@ series. tailmarker.loc[0, 'x'] = lollipop_tail_length else: tailmarker.loc[0, 'x'] = -lollipop_tail_length - tail = self.world2body(tailmarker, indeces=indeces) + tail = self.body2world(tailmarker, indeces=indeces) tail = tail.loc[:, 0] # Plot the agent trajectory # - loop through consecutive point diff --git a/navipy/trajectories/test_trajectory.py b/navipy/trajectories/test_trajectory.py index c5c4cc269842a6289579baa634fcea1f23524aa5..4e023e89665db87e2e9f4c3d026c06c50dd770a4 100644 --- a/navipy/trajectories/test_trajectory.py +++ b/navipy/trajectories/test_trajectory.py @@ -40,7 +40,7 @@ class TestTrajectoryTransform(unittest.TestCase): for euler_axes in list(htconst._AXES2TUPLE.keys()): trajectory.rotation_mode = euler_axes traj_test.rotation_mode = euler_axes - transformed_markers = trajectory.world2body(markers) + transformed_markers = trajectory.body2world(markers) # reverse for triangle_mode in ['x-axis=median-from-0', 'y-axis=1-2']: