vectorized conversion from ax(is angle)

python/damask/_rotation.py

@@ -653,27 +653,50 @@ class Rotation:
     @staticmethod
     def ax2qu(ax):
         """Axis angle pair to quaternion."""
-        if np.abs(ax[3])<1.e-6:
+        if len(ax.shape) == 1:
-            qu = np.array([ 1.0, 0.0, 0.0, 0.0 ])
+            if np.abs(ax[3])<1.e-6:
+                qu = np.array([ 1.0, 0.0, 0.0, 0.0 ])
+            else:
+                c = np.cos(ax[3]*0.5)
+                s = np.sin(ax[3]*0.5)
+                qu = np.array([ c, ax[0]*s, ax[1]*s, ax[2]*s ])
+            return qu
         else:
-            c = np.cos(ax[3]*0.5)
+            c = np.cos(ax[...,3:4]*.5)
-            s = np.sin(ax[3]*0.5)
+            s = np.sin(ax[...,3:4]*.5)
-            qu = np.array([ c, ax[0]*s, ax[1]*s, ax[2]*s ])
+            qu = np.where(np.abs(ax[...,3:4])<1.e-12,
-        return qu
+                          [1.0, 0.0, 0.0, 0.0],
+                          np.block([c, ax[...,:3]*s]))
+            return qu
 
     @staticmethod
     def ax2om(ax):
         """Axis angle pair to rotation matrix."""
-        c = np.cos(ax[3])
+        if len(ax.shape) == 1:
-        s = np.sin(ax[3])
+            c = np.cos(ax[3])
-        omc = 1.0-c
+            s = np.sin(ax[3])
-        om=np.diag(ax[0:3]**2*omc + c)
+            omc = 1.0-c
+            om=np.diag(ax[0:3]**2*omc + c)
 
-        for idx in [[0,1,2],[1,2,0],[2,0,1]]:
+            for idx in [[0,1,2],[1,2,0],[2,0,1]]:
-            q = omc*ax[idx[0]] * ax[idx[1]]
+                q = omc*ax[idx[0]] * ax[idx[1]]
-            om[idx[0],idx[1]] = q + s*ax[idx[2]]
+                om[idx[0],idx[1]] = q + s*ax[idx[2]]
-            om[idx[1],idx[0]] = q - s*ax[idx[2]]
+                om[idx[1],idx[0]] = q - s*ax[idx[2]]
-        return om if P < 0.0 else om.T
+            return om if P < 0.0 else om.T
+        else:
+            c = np.cos(ax[...,3:4])
+            s = np.sin(ax[...,3:4])
+            omc = 1. -c
+            ax = np.block([c+omc*ax[...,0:1]**2,
+                           omc*ax[...,0:1]*ax[...,1:2] + s*ax[...,2:3],
+                           omc*ax[...,0:1]*ax[...,2:3] - s*ax[...,1:2],
+                           omc*ax[...,0:1]*ax[...,1:2] - s*ax[...,2:3],
+                           c+omc*ax[...,1:2]**2,
+                           omc*ax[...,1:2]*ax[...,2:3] + s*ax[...,0:1],
+                           omc*ax[...,0:1]*ax[...,2:3] + s*ax[...,1:2],
+                           omc*ax[...,1:2]*ax[...,2:3] - s*ax[...,0:1],
+                           c+omc*ax[...,2:3]**2]).reshape(ax.shape[:-1]+(3,3))
+            return ax
 
     @staticmethod
     def ax2eu(ax):
@@ -683,21 +706,35 @@ class Rotation:
     @staticmethod
     def ax2ro(ax):
         """Axis angle pair to Rodrigues-Frank vector."""
-        if np.abs(ax[3])<1.e-6:
+        if len(ax.shape) == 1:
-            ro = [ 0.0, 0.0, P, 0.0 ]
+            if np.abs(ax[3])<1.e-6:
+                ro = [ 0.0, 0.0, P, 0.0 ]
+            else:
+                ro = [ax[0], ax[1], ax[2]]
+                # 180 degree case
+                ro += [np.inf] if np.isclose(ax[3],np.pi,atol=1.0e-15,rtol=0.0) else \
+                      [np.tan(ax[3]*0.5)]
+            return np.array(ro)
         else:
-            ro = [ax[0], ax[1], ax[2]]
+            ro = np.block([ax[...,:3],
-            # 180 degree case
+                           np.where(np.isclose(ax[...,3:4],np.pi,atol=1.e-15,rtol=.0),
-            ro += [np.inf] if np.isclose(ax[3],np.pi,atol=1.0e-15,rtol=0.0) else \
+                                    np.inf,
-                  [np.tan(ax[3]*0.5)]
+                                    np.tan(ax[...,3:4]*0.5))
-        return np.array(ro)
+                          ])
+            ro[np.abs(ax[...,3])<1.e-6] = [.0,.0,P,.0]
+            return ro
 
     @staticmethod
     def ax2ho(ax):
         """Axis angle pair to homochoric vector."""
-        f = (0.75 * ( ax[3] - np.sin(ax[3]) ))**(1.0/3.0)
+        if len(ax.shape) == 1:
-        ho = ax[0:3] * f
+            f = (0.75 * ( ax[3] - np.sin(ax[3]) ))**(1.0/3.0)
-        return ho
+            ho = ax[0:3] * f
+            return ho
+        else:
+            f = (0.75 * ( ax[...,3:4] - np.sin(ax[...,3:4]) ))**(1.0/3.0)
+            ho = ax[...,:3] * f
+            return ho
 
     @staticmethod
     def ax2cu(ax):

python/tests/test_Rotation.py

@@ -78,7 +78,7 @@ def default():
     specials /= np.linalg.norm(specials,axis=1).reshape(-1,1)
     specials[specials[:,0]<0]*=-1
     return [Rotation.fromQuaternion(s) for s in specials] + \
-           [Rotation.fromRandom() for r in range(n-len(specials))]
+           [Rotation.fromRandom() for _ in range(n-len(specials))]
 
 @pytest.fixture
 def reference_dir(reference_dir_base):
@@ -149,3 +149,15 @@ class TestRotation:
             o = Rotation.fromQuaternion(rot.asQuaternion()).asCubochoric()
             print(m,o,rot.asQuaternion())
             assert np.allclose(m,o,atol=atol)
+
+    @pytest.mark.parametrize('conversion',[Rotation.ax2qu,
+                                           Rotation.ax2om,
+                                           Rotation.ax2ro,
+                                           Rotation.ax2ho,
+                                          ])
+    def test_axisAngle_vectorization(self,default,conversion):
+        ax = np.array([rot.asAxisAngle() for rot in default])
+        dev_null = conversion(ax.reshape(ax.shape[0]//2,-1,4))
+        co = conversion(ax)
+        for a,c in zip(ax,co):
+             assert np.allclose(conversion(a),c)