Reordering a Tensor

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


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

        # Declare tens as an input tensor with shape = (4, 3, 2, 5)
        tens = self.declare_input(
            'T1',
            val=np.arange(4 * 3 * 5 * 2).reshape((4, 3, 5, 2)),
        )

        # Compute a new tensor by reordering axes of tens; reordering is
        # given by 'ijkl->ljki'
        self.register_output('axes_reordered_tensor',
                             ot.reorder_axes(tens, 'ijkl->ljki'))


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

print('T1', prob['T1'].shape)
print(prob['T1'])
print('axes_reordered_tensor', prob['axes_reordered_tensor'].shape)
print(prob['axes_reordered_tensor'])

T1 (4, 3, 5, 2)
[[[[  0.   1.]
   [  2.   3.]
   [  4.   5.]
   [  6.   7.]
   [  8.   9.]]

  [[ 10.  11.]
   [ 12.  13.]
   [ 14.  15.]
   [ 16.  17.]
   [ 18.  19.]]

  [[ 20.  21.]
   [ 22.  23.]
   [ 24.  25.]
   [ 26.  27.]
   [ 28.  29.]]]


 [[[ 30.  31.]
   [ 32.  33.]
   [ 34.  35.]
   [ 36.  37.]
   [ 38.  39.]]

  [[ 40.  41.]
   [ 42.  43.]
   [ 44.  45.]
   [ 46.  47.]
   [ 48.  49.]]

  [[ 50.  51.]
   [ 52.  53.]
   [ 54.  55.]
   [ 56.  57.]
   [ 58.  59.]]]


 [[[ 60.  61.]
   [ 62.  63.]
   [ 64.  65.]
   [ 66.  67.]
   [ 68.  69.]]

  [[ 70.  71.]
   [ 72.  73.]
   [ 74.  75.]
   [ 76.  77.]
   [ 78.  79.]]

  [[ 80.  81.]
   [ 82.  83.]
   [ 84.  85.]
   [ 86.  87.]
   [ 88.  89.]]]


 [[[ 90.  91.]
   [ 92.  93.]
   [ 94.  95.]
   [ 96.  97.]
   [ 98.  99.]]

  [[100. 101.]
   [102. 103.]
   [104. 105.]
   [106. 107.]
   [108. 109.]]

  [[110. 111.]
   [112. 113.]
   [114. 115.]
   [116. 117.]
   [118. 119.]]]]
axes_reordered_tensor (2, 3, 5, 4)
[[[[  0.  30.  60.  90.]
   [  2.  32.  62.  92.]
   [  4.  34.  64.  94.]
   [  6.  36.  66.  96.]
   [  8.  38.  68.  98.]]

  [[ 10.  40.  70. 100.]
   [ 12.  42.  72. 102.]
   [ 14.  44.  74. 104.]
   [ 16.  46.  76. 106.]
   [ 18.  48.  78. 108.]]

  [[ 20.  50.  80. 110.]
   [ 22.  52.  82. 112.]
   [ 24.  54.  84. 114.]
   [ 26.  56.  86. 116.]
   [ 28.  58.  88. 118.]]]


 [[[  1.  31.  61.  91.]
   [  3.  33.  63.  93.]
   [  5.  35.  65.  95.]
   [  7.  37.  67.  97.]
   [  9.  39.  69.  99.]]

  [[ 11.  41.  71. 101.]
   [ 13.  43.  73. 103.]
   [ 15.  45.  75. 105.]
   [ 17.  47.  77. 107.]
   [ 19.  49.  79. 109.]]

  [[ 21.  51.  81. 111.]
   [ 23.  53.  83. 113.]
   [ 25.  55.  85. 115.]
   [ 27.  57.  87. 117.]
   [ 29.  59.  89. 119.]]]]