From 0704f1122e76b68eb2350c1605cc1ca61609a107 Mon Sep 17 00:00:00 2001
From: Philip Eisenlohr <eisenlohr@egr.msu.edu>
Date: Tue, 19 Sep 2023 17:50:28 -0400
Subject: [PATCH] result (and object interpretation) "flows" from rotation and
 data shapes

---
 python/damask/_rotation.py    | 48 ++++++++++++++++++++---------------
 python/tests/test_Rotation.py |  2 +-
 2 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/python/damask/_rotation.py b/python/damask/_rotation.py
index cd3164db4..3353cce6d 100644
--- a/python/damask/_rotation.py
+++ b/python/damask/_rotation.py
@@ -375,6 +375,11 @@ class Rotation:
         Return self@other.
 
         Rotate vector, second-order tensor, or fourth-order tensor.
+        `other` is interpreted as an array of tensor quantities with the highest-possible order
+        considering the shape of `self`.
+        For instance, shapes of (2,) and (3,3) for `self` and `other` prompt interpretation of
+        `other` as a second-rank tensor and result in (2,) rotated tensors, whereas
+        shapes of (2,1) and (3,3) for `self` and `other` result in (2,3) rotated vectors.
 
         Parameters
         ----------
@@ -388,27 +393,30 @@ class Rotation:
 
         """
         if isinstance(other, np.ndarray):
-            if self.shape + (3,) == other.shape:
-                q_m = self.quaternion[...,0]
-                p_m = self.quaternion[...,1:]
-                A = q_m**2 - np.einsum('...i,...i',p_m,p_m)
-                B = 2. * np.einsum('...i,...i',p_m,other)
-                C = 2. * _P * q_m
-                return np.block([(A * other[...,i]).reshape(self.shape+(1,)) +
-                                 (B * p_m[...,i]).reshape(self.shape+(1,)) +
-                                 (C * (  p_m[...,(i+1)%3]*other[...,(i+2)%3]\
-                                       - p_m[...,(i+2)%3]*other[...,(i+1)%3])).reshape(self.shape+(1,))
-                                 for i in [0,1,2]])
-            if self.shape + (3,3) == other.shape:
-                R = self.as_matrix()
-                return np.einsum('...im,...jn,...mn',R,R,other)
-            if self.shape + (3,3,3,3) == other.shape:
-                R = self.as_matrix()
-                return np.einsum('...im,...jn,...ko,...lp,...mnop',R,R,R,R,other)
-            else:
-                raise ValueError('can only rotate vectors, second-order tensors, and fourth-order tensors')
+            obs = util.shapeblender(self.shape,other.shape,keep_ones=False)[len(self.shape):]
+            for l in [4,2,1]:
+                if obs[-l:] == l*(3,):
+                    bs = util.shapeblender(self.shape,other.shape[:-l],False)
+                    self_ = self.broadcast_to(bs) if self.shape != bs else self
+                    if l==1:
+                        q_m = self_.quaternion[...,0]
+                        p_m = self_.quaternion[...,1:]
+                        A = q_m**2 - np.einsum('...i,...i',p_m,p_m)
+                        B = 2. * np.einsum('...i,...i',p_m,other)
+                        C = 2. * _P * q_m
+                        return np.block([(A * other[...,i]) +
+                                         (B *   p_m[...,i]) +
+                                         (C * ( p_m[...,(i+1)%3]*other[...,(i+2)%3]
+                                              - p_m[...,(i+2)%3]*other[...,(i+1)%3]))
+                                        for i in [0,1,2]]).reshape(bs+(3,),order='F')
+                    else:
+                        return np.einsum({2: '...im,...jn,...mn',
+                                          4: '...im,...jn,...ko,...lp,...mnop'}[l],
+                                         *l*[self_.as_matrix()],
+                                         other)
+            raise ValueError('can only rotate vectors, second-order tensors, and fourth-order tensors')
         elif isinstance(other, Rotation):
-            raise TypeError('use "R1*R2", i.e. multiplication, to compose rotations "R1" and "R2"')
+            raise TypeError('use "R2*R1", i.e. multiplication, to compose rotations "R1" and "R2"')
         else:
             raise TypeError(f'cannot rotate "{type(other)}"')
 
diff --git a/python/tests/test_Rotation.py b/python/tests/test_Rotation.py
index 6b9cd2b1a..5a2d91a6e 100644
--- a/python/tests/test_Rotation.py
+++ b/python/tests/test_Rotation.py
@@ -1065,7 +1065,7 @@ class TestRotation:
 
     @pytest.mark.parametrize('data',[np.random.rand(4),
                                      np.random.rand(3,2),
-                                     np.random.rand(3,2,3,3)])
+                                     np.random.rand(3,3,3,1)])
     def test_rotate_invalid_shape(self,data):
         R = Rotation.from_random()
         with pytest.raises(ValueError):