Remove mixed-variable terms in SymPy series expansion

Question

Consider two functions of SymPy symbols e and i:

from sympy import Symbol, expand, Order
i = Symbol('i')
e = Symbol('e')
f = (i**3 + i**2 + i + 1)
g = (e**3 + e**2 + e + 1)
z = expand(f*g)

This will produce

z = e**3*i**3 + e**3*i**2 + e**3*i + e**3 + e**2*i**3 + e**2*i**2 + e**2*i + e**2 + e*i**3 + e*i**2 + e*i + e + i**3 + i**2 + i + 1

However, assume that e and i are both small and we can neglect both terms that are order three or higher. Using Sympy’s series tool or simply adding an O-notation Order class can handle this:

In : z = expand(f*g + Order(i**3) + Order(e**3))
Out: 1 + i + i**2 + e + e*i + e*i**2 + e**2 + e**2*i + e**2*i**2 + O(i**3) + O(e**3)

Looks great. However, I am still left with mixed terms e**2 * i**2. Individual variables in these terms are less than the desired cut-off so SymPy keeps them. However, mathematically small²·small² = small⁴. Likewise, e·i² = small·small² = small³.

At least for my purposes, I want these mixed terms dropped. Adding a mixed Order does not produce the desired result (it seems to ignore the first two orders).

In : expand(f*g + Order(i**3) + Order(e**3) + Order((i**2)*(e**2)))
Out: 1 + i + i**2 + i**3 + e + e*i + e*i**2 + e*i**3 + e**2 + e**2*i + e**3 + e**3*i + O(e**2*i**2, e, i)

Question: Does SymPy have an easy system to quickly remove the n-th order terms, as well as terms that are (e^a)·(i^b) where a+b > n?

Messy Solution: I have found a way to solve this, but it is messy and potentially not general.

z = expand(f*g + Order((e**2)*i) + Order(e*(i**2)))
zz = expand(z.removeO() + Order(e**3) + Order(i**3))

produces

zz = 1 + i + i**2 + e + e*i + e**2 + O(i**3) + O(e**3)

which is exactly what I want. So to specify my question: Is there a way to do this in one step that can be generalized to any n? Also, my solution loses the big-O notation that indicates mixed-terms were lost. This is not needed but would be nice.

Answer 1

As you have a dual limit, you must specify both infinitesimal variables (e and i) in all Order objects, even if they don’t appear in the first argument.

The reason for this is that Order(expr) only automatically chooses those symbols as infinitesimal that actually appear in the expr and thus, e.g., O(e) is only for the limit e→0. Now, Order objects with different limits don’t mix well, e.g.:

O(e*i)+O(e) == O(e*i) != O(e)+O(e*i) == O(e) # True

This leads to a mess where results depend on the order of addition, which is a good indicator that this is something to avoid. This can be avoided by explicitly specifying the infinitesimal symbols (as addition arguments of Order), e.g.:

O(e*i)+O(e,e,i) == O(e,e,i)+O(e*i) == O(e,e,i) # True

I haven’t found a way to avoid going through all combinations of e and i manually, but this can be done by a simple iteration:

orders = sum( Order(e**a*i**(n-a),e,i) for a in range(n+1) )
expand(f*g+orders)
# 1 + i + i**2 + e + e*i + e**2 + O(e**2*i, e, i) + O(e*i**2, e, i) + O(i**3, e, i) + O(e**3, e, i)

Answer 2

Without using Order you might try something simple like this:

>>> eq = expand(f*g)  # as you defined
>>> def total_degree(e):
...     x = Dummy()
...     free = e.free_symbols
...     if not free: return S.Zero
...     for f in free:
...         e = e.subs(f, x)
...     return degree(e)
>>> eq.replace(lambda x: total_degree(x) > 2, lambda x: S.Zero)
e**2 + e*i + e + i**2 + i + 1

Answer 3

There is a way about it using Poly. I have made a function that keeps the O(...) term and another that does not (faster).

from sympy import Symbol, expand, Order, Poly
i = Symbol('i')
e = Symbol('e')
f = (i**3 + i**2 + i + 1)
g = (e**3 + e**2 + e + 1)
z = expand(f*g)

def neglect(expr, order=3):
    z = Poly(expr)
    
    # extract all terms and keep the lower order ones
    d = z.as_dict()
    d = {t: c for t,c in d.items() if sum(t) < order}
    
    # Build resulting polynomial
    return Poly(d, z.gens).as_expr()
        
def neglectO(expr, order=3):
    # This one keeps O terms
    z = Poly(expr)
    
    # extract terms of higher "order"
    d = z.as_dict()
    large = {t: c for t,c in d.items() if sum(t) >= order}
    for t in large:  # Add each O(large monomial) to the expression
        expr += Order(Poly({t:1},z.gens).as_expr(), *z.gens)
    
    return expr

print(neglect(z))
print(neglectO(z))

This code prints the following:

e**2 + e*i + e + i**2 + i + 1
1 + i + i**2 + e + e*i + e**2 + O(e**2*i, e, i) + O(e*i**2, e, i) + O(i**3, e, i) + O(e**3, e, i)