How to substitute partial powers?

Viewed 44

I have an expression which contains terms like expr = a**m * b**n where a,b are symbols and m,n are integers. I want to use expr.subs(a**i * b**j, 0) to set the expression to zero if a**i * b**j is a factor of a**m * b**n (here i,j are integers too). This works in some cases, but not when only one of the symbols is used for the substitution:

>>> from sympy import symbols
>>> a, b = symbols('a b')
>>> (a**2 * b**3).subs(a**2 * b**2, 0)  # (1) works
0
>>> (a**2 * b**3).subs(a * b**2, 0)  # (2) works
0
>>> (a**2 * b**3).subs(b**2, 0)  # (3) does not work
a**2*b**3
>>> (a**3).subs(a**2, 0)  # (4) does not work
a**3
>>> (a**4).subs(a**2, 0)  # (5) works
0

I would like this substitution to work in all of the above cases and logically it should (since for example a**2 * b**3 is equivalent to a**2 * b * (b**2) and thus the substitution (b**2, 0) should yield 0). However, if using a single symbol in the substitution, it seems to work only for those cases where the expression contains the subs-term raised to a power; for example:

>>> (a**2 * b**4).subs(b**2, 0)  # (6) works
0
>>> (a**2 * b**4).subs(b**3, 0)  # (7) does not work
a**2*b**4

This behavior seems a little odd since it works for b**1 and also when artificially augmenting both, the actual expression and the subs-term, with some other symbol:

>>> (a**2 * b**4).subs(b, 0)  # (8) works
0
>>> c = symbols('c')
>>> (c * (a**2 * b**4)).subs(c * (b**3), 0)  # (9) works
0

While I could use this last observation (augmenting expressions) as a workaround, it doesn't seem neither clear nor efficient. Hence, is there another way to make the substitution work for all cases where the subs-term is contained in the actual expression?


>>> sympy.__version__
'1.10.1'
2 Answers

To check if something is a factor of a product you can use extract_multiplicatively:

>>> expr = a**3*b**3*c
>>> 0 if expr.extract_multiplicatively(a*b**2) else expr
0

You could also use the rem function to see if the remainder is 0:

>>> rem(expr, a*b**2)
0

If your expressions occur within a larger expression, then these could be incorporated into a replace call, expr.replace(...). If you are trying to ignore high-order terms in an expression you can just make a function to add the order of all symbol-based powers in a term:

>>> expr = expand((x + 1)**3*(y + 1)**2)
>>> def order(m):
...     return sum([degree(m, x) for x in m.free_symbols])
...
>>> expr.replace(lambda x: not x.is_Add,
...     lambda x: 0 if order(x)>2 else x)
3*x**2 + 6*x*y + 3*x + y**2 + 2*y + 1
  1. works, because a2 = 0 is 0 * b3 what leads to 0
  2. a = 0 means a**2 = 0 -> 0
  3. b2 = 0 but b2 is not in b**3, so it does not work
  4. same issue as 3
  5. a4 is in a2 so 0
  6. same as 5 but with b´s
  7. same issue as 3
  8. b is in b**4
  9. c = c so it comes to 0 because only symbols get replaced

Everything works like intended

Related