Multiple Vector Summation using Einsum

This is an example of how to properly use the einsum function to compute the summation of multiple vectors.

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


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

        a = np.arange(4)
        vec = self.declare_input('a', val=a)

        # Special operation: sum all the entries of the first and second
        # vector to a single scalar
        self.register_output(
            'einsum_special2',
            ot.einsum_new_api(vec, vec, operation=[(1, ), (2, )]))


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

print('a', prob['a'].shape)
print(prob['a'])
print('einsum_special2', prob['einsum_special2'].shape)
print(prob['einsum_special2'])

a (4,)
[0. 1. 2. 3.]
einsum_special2 (1,)
[36.]