X-Axis Rotation Tensor For Different Angles in a Tensor

This example generates a tensor of rotation matrices that rotate a vector about the x-axis by angles given in the tensor.

from openmdao.api import Problem
from omtools.api import Group
import omtools.api as ot
import numpy as np


class ExampleDiffRadianTensorRotX(Group):
    def setup(self):

        # Shape of a random tensor rotation matrix
        shape = (2, 3, 4)

        num_elements = np.prod(shape)

        # Vector of angles in radians
        angle_val2 = np.repeat(
            np.pi / 3, num_elements) + 2 * np.pi * np.arange(num_elements)

        angle_val2 = angle_val2.reshape(shape)

        # Adding the vector as an input
        angle_tensor = self.declare_input('tensor', val=angle_val2)

        # Rotation in the x-axis for tensor2
        self.register_output('tensor_Rot_x', ot.rotmat(angle_tensor, axis='x'))


prob = Problem()
prob.model = ExampleDiffRadianTensorRotX()
prob.setup(force_alloc_complex=True)
prob.run_model()

print('tensor', prob['tensor'].shape)
print(prob['tensor'])
print('tensor_Rot_x', prob['tensor_Rot_x'].shape)
print(prob['tensor_Rot_x'])

tensor (2, 3, 4)
[[[  1.04719755   7.33038286  13.61356817  19.89675347]
  [ 26.17993878  32.46312409  38.74630939  45.0294947 ]
  [ 51.31268001  57.59586532  63.87905062  70.16223593]]

 [[ 76.44542124  82.72860654  89.01179185  95.29497716]
  [101.57816247 107.86134777 114.14453308 120.42771839]
  [126.71090369 132.994089   139.27727431 145.56045962]]]
tensor_Rot_x (2, 3, 4, 3, 3)
[[[[[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]]


  [[[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]]


  [[[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]]]



 [[[[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]]


  [[[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]]


  [[[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]

   [[ 1.         0.         0.       ]
    [ 0.         0.5       -0.8660254]
    [ 0.         0.8660254  0.5      ]]]]]