From 3ed6fc540b78e90a49984bc43a47f59b206b0994 Mon Sep 17 00:00:00 2001
From: nilay <n.y.ozge@gmail.com>
Date: Wed, 6 Mar 2019 21:23:30 -0800
Subject: [PATCH] Adding Quat Calculations

---
 pymo/preprocessing.py  | 14 +++++++++++---
 pymo/rotation_tools.py | 28 +++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/pymo/preprocessing.py b/pymo/preprocessing.py
index 5ec88f0..992e064 100644
--- a/pymo/preprocessing.py
+++ b/pymo/preprocessing.py
@@ -24,12 +24,15 @@ def fit(self, X, y=None):
 
     def transform(self, X, y=None):
         if self.param_type == 'euler':
+            print('euler')
             return X
         elif self.param_type == 'expmap':
             return self._to_expmap(X)
         elif self.param_type == 'quat':
-            return X
+            print('Quats')
+            return self._to_pos(X)
         elif self.param_type == 'position':
+            print('pos')
             return self._to_pos(X)
         else:
             raise 'param types: euler, quat, expmap, position'
@@ -108,8 +111,13 @@ def _to_pos(self, X):
                 #euler_values = [[0,0,0] for f in rc.iterrows()] #for deugging
                 #pos_values = [[0,0,0] for f in pc.iterrows()] #for deugging
                 
-                # Convert the eulers to rotation matrices
-                rotmats = np.asarray([Rotation([f[0], f[1], f[2]], 'euler', from_deg=True).rotmat for f in euler_values])                  
+                # Convert the eulers or quats to rotation matrices
+                if self.param_type == 'euler':
+                    rotmats = np.asarray([Rotation([f[0], f[1], f[2]], 'euler', from_deg=True).rotmat for f in euler_values])
+                elif self.param_type == 'quat':
+                    rotmats = np.asarray([Rotation([f[0], f[1], f[2]], 'quat', from_deg=False).rotmat for f in euler_values])
+                else:
+                    raise 'Type not supported'
                 tree_data[joint]=[
                                     [], # to store the rotation matrix
                                     []  # to store the calculated position
diff --git a/pymo/rotation_tools.py b/pymo/rotation_tools.py
index 1fe6072..3c88742 100644
--- a/pymo/rotation_tools.py
+++ b/pymo/rotation_tools.py
@@ -24,6 +24,8 @@ def __init__(self,rot, param_type, **params):
             self._from_euler(rot[0],rot[1],rot[2], params)
         elif param_type == 'expmap':
             self._from_expmap(rot[0], rot[1], rot[2], params)
+        elif param_type == 'quat':
+            self._from_quat(rot[0], rot[1], rot[2], params)
 
     def _from_euler(self, alpha, beta, gamma, params):
         '''Expecting degress'''
@@ -84,7 +86,31 @@ def _from_expmap(self, alpha, beta, gamma, params):
             [2*x*z*s**2-2*y*c*s, 2*y*z*s**2+2*x*c*s , 2*(z**2-1)*s**2+1]
         ])
         
-
+    def _from_quat(self, alpha, beta, gamma, params):
+        if params['from_deg']==True:
+            alpha = deg2rad(alpha)
+            beta = deg2rad(beta)
+            gamma = deg2rad(gamma)
+        
+        axis = [alpha, beta, gamma]
+        theta = np.linalg.norm(axis)
+        
+        if not np.isclose(np.linalg.norm(axis), 1.):
+            axis = axis/np.linalg.norm(axis)
+        
+        sinThetaOver2 = np.sin(theta * 0.5)
+		
+        q1 = sinThetaOver2 * axis[0]
+        q2 = sinThetaOver2 * axis[1]
+        q3 = sinThetaOver2 * axis[2]
+        q4 = np.cos(theta * 0.5)
+		
+        self.rotmat = np.asarray([
+            [1-2*q2*q2-2*q3*q3,  2*(q1*q2-q3*q4),    2*(q1*q3+q2*q4)],
+            [2*(q1*q2+q3*q4),    1-2*q1*q1-2*q3*q3,  2*(q2*q3-q1*q4)],
+            [2*(q1*q3-q2*q4),    2*(q1*q4+q2*q3),    1-2*q1*q1-2*q2*q2]
+        ])	
+		
 
     def get_euler_axis(self):
         R = self.rotmat