CODEWITHSUNDEEP
pythonnumpysympy
qAp
qAp

Reputation: 1189

how to write symbol for sum over a variable's subscript in sympy

I want to write a sympy symbol for a summation, but the index summed over also appears as the subscript of a variable name in the summand. For example, 

import numpy as np
import sympy
sympy.init_printing()

r = sympy.Symbol('r')

a = sympy.Matrix(sympy.symbols('a:4'))
rpowers = sympy.Matrix([r**i for i in range(len(a))])
long_expr = a.dot(rpowers)

n = sympy.Symbol('n')
a_n = sympy.Symbol('a_n')
short_expr = sympy.Sum(a_n * r**n, (n, 0, 3))

long_expr and short_expr denote the same thing mathematically. But with long_expr, I can substitute in the values for the a's and then lambdify that expression into a numpy function:

coeffed_long_expr = long_expr.subs(zip(a, [-1, 3, 23, 8]))
func_long_expr = sympy.lambdify([r], coeffed_long_expr, 'numpy')

How can I do the same with short_expr? Or is short_expr only useful for displaying the expression with a summation sign in this case? I would like to be able to display using the summation sign, especially for large ns.

Upvotes: 2

Views: 2534

Answers (1)

llf
llf

Reputation: 630

You can accomplish this by using sympy.Function:

import sympy
a_seq = [-1, 3, 23, 8]
n, r = sympy.symbols('n, r')
a_n = sympy.Function('a')(n)
terms = 4
short_expr = sympy.Sum(a_n * r**n, (n, 0, terms - 1))
coeffed_short_expr = short_expr.doit().subs(
    (a_n.subs(n, i), a_seq[i]) for i in range(terms)) # 8*r**3 + 23*r**2 + 3*r - 1
func_short_expr = sympy.lambdify(r, coeffed_short_expr, 'numpy')

If you wish for a cleaner, more efficient implementation, I suspect you may be able to define a subclass of sympy.Symbol that implements subs() properly for summations.

Upvotes: 6

Related Questions