How do I break up a squared term in sympy

user_14415333 0
Viewed 51

I am using python (3.7.3) with sympy (1.6.2) to store a function with squared terms and non-squared terms, with each term being the product of exactly two variables.

For example,

>> import sympy as sy
>> x = sy.Symbol('x')
>> y = sy.Symbol('y')
>> F = x*x+x*y
>> print(F)
x**2+x*y

I want to be able to iterate through the terms and get each operand. For example,

terms = F.expand(basic=True).args
for term in terms
    (t0,t1) = term.args
    print('t0:{}, t1:{}'.format(t0,t1))
    # do some stuff using t0, t1

This works for the x*y term, but not the x**2 term.

>> print((x*y).args)
(x,y)
>> print((x**2).args) # I want this to be (x,x)
(x,2)

I tried running (x**2).expand(), but this appears to be the fully expanded version of the expression.

My question is twofold:

  1. is there a way to expand x**2 so that it is stored as x*x?
  2. is there a better way to go about getting each operand in each term than the for loop I show above?
sympy python-3.x python
1 Answers

You could define a custom function that defactors in the way you want:

def get_factors(expr):
    if expr.func == sy.Mul:
        return expr.args
    elif expr.func == sy.Pow:
        return tuple(expr.args[0] for _ in range(expr.args[1]))
    else:
        raise NotImplementedError()

Usage:

>>> a, b = terms
>>> get_factors(a)
(x, x)
>>> get_factors(b)
(x, y)
user_10162160
Related