Reordering a Matrix using Einsum with Sparse Partials

This is an example of how to properly use the einsum function to reorder a matrix while using sparse partial derivatives.

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


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

        shape2 = (5, 4)
        b = np.arange(20).reshape(shape2)
        mat = self.declare_input('b', val=b)

        self.register_output(
            'einsum_reorder1_sparse_derivs',
            ot.einsum(
                mat,
                subscripts='ij->ji',
                partial_format='sparse',
            ))


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

print('b', prob['b'].shape)
print(prob['b'])
print('einsum_reorder1_sparse_derivs', prob['einsum_reorder1_sparse_derivs'].shape)
print(prob['einsum_reorder1_sparse_derivs'])

b (5, 4)
[[ 0.  1.  2.  3.]
 [ 4.  5.  6.  7.]
 [ 8.  9. 10. 11.]
 [12. 13. 14. 15.]
 [16. 17. 18. 19.]]
einsum_reorder1_sparse_derivs (4, 5)
[[ 0.  4.  8. 12. 16.]
 [ 1.  5.  9. 13. 17.]
 [ 2.  6. 10. 14. 18.]
 [ 3.  7. 11. 15. 19.]]