result (and object interpretation) "flows" from rotation and data shapes

This commit is contained in:
Philip Eisenlohr 2023-09-19 17:50:28 -04:00
parent b3f98ab877
commit 0704f1122e
2 changed files with 29 additions and 21 deletions

View File

@ -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)}"')

View File

@ -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):