multi_poly.py

from functools import reduce
from operator import mul
import sympy as sp
import numpy as np

# Both of these function will be vectorized via np.vectorize.
# This will not speed up the computation at all (behind np.vectorize, there still are for loops)
# but this avoid nested for loop, making that easier to read (and deal with arbitrary number of dims).

def build_X(a, xs):
    """
    take the indices of a and apply it to the x_i.
    reduce (from the built-in functools module) combined with mul (from the built-in operator module)
    is able to reproduce the Gamma behavior : like sum for product.
    """
    orders = a.args[1:]
    return reduce(mul, [xs[i]**order for i, order in enumerate(orders)])
build_X = np.vectorize(build_X, excluded=[1])

def build_A(a, *idxs):
    """
    simple function that fill the IndexedBase a from the indices generated via np.indices.
    """
    return a[idxs]
build_A = np.vectorize(build_A)

D = [2, 3]

# sympy.symbols is able to generate a range of symbols like symbols("x3:5") => (x3, x4)
xs = np.array(sp.symbols(f"x_1:{len(D)+1}"))

# sympy provide an IndexedBase to represent indexed vector / matrix / more.
a = sp.IndexedBase("a", shape=D)

# I use your idea : build A and X then using the Frobenius inner product.
idxs = np.indices(D)
A = build_A(a, *idxs)
X = build_X(A, xs)

# Then, you should have the proper result.
polynomial = np.sum(np.multiply(A, X))
print(polynomial)

forked from this

other solutions here and here