Sympy code for FDM expansion of LBM diffusive equation

import sympy as sp
# -----------------------------------------------------------------------------
# 1) Symbols (parameters and step sizes)
# -----------------------------------------------------------------------------
dx, dt = sp.symbols('Delta_x Delta_t', positive=True, real=True)
omega, a = sp.symbols('omega a', real=True)

Omega = 1 - omega

alpha1 = sp.simplify(Omega + a*omega)
alpha2 = sp.simplify(Omega + (1 - 2*a)*omega)
beta1  = sp.simplify(-Omega*(Omega + (1 - a)*omega))
beta2  = sp.simplify(-Omega*(Omega + 2*a*omega))

# -----------------------------------------------------------------------------
# 2) Continuous variables and field
# -----------------------------------------------------------------------------
x, t = sp.symbols('x t', real=True)
phi = sp.Function('phi')
# Shorthand for the base (continuous) field value at (x,t)
PHI = phi(x, t)
def taylor_phi(sx=0, st=0, order_x=4, order_t=4):
    expr = sp.S(0)
    # Keep all mixed terms up to the specified individual orders.
    for m in range(order_x + 1):
        for n in range(order_t + 1):
            # skip the (0,0) term? no, keep it.
            term = ( (sx*dx)**m * (st*dt)**n
                    / (sp.factorial(m)*sp.factorial(n))
                    * sp.diff(PHI, x, m, t, n) )
            expr += term
    return sp.expand(expr)
# -----------------------------------------------------------------------------
# 4) Discrete stencil terms (expanded)
# -----------------------------------------------------------------------------
# phi_x^{t+1}, phi_x^{t}, phi_x^{t-1}, phi_x^{t-2}
phi_t_p1 = taylor_phi(sx=0,  st=+1, order_x=0, order_t=4)
display(phi_t_p1)
phi_t_0  = taylor_phi(sx=0,  st=0,  order_x=0, order_t=0)  # exactly phi(x,t)
display(phi_t_0)
phi_t_m1 = taylor_phi(sx=0,  st=-1, order_x=0, order_t=4)
display(phi_t_m1)
phi_t_m2 = taylor_phi(sx=0,  st=-2, order_x=0, order_t=4)
display(phi_t_m2)

# phi_{x±1}^{t}
phi_xm1_t0 = taylor_phi(sx=-1, st=0,  order_x=4, order_t=0)
display(phi_xm1_t0)
phi_xp1_t0 = taylor_phi(sx=+1, st=0,  order_x=4, order_t=0)
display(phi_xp1_t0)

# phi_{x±1}^{t-1}
phi_xm1_tm1 = taylor_phi(sx=-1, st=-1, order_x=4, order_t=4)
phi_xp1_tm1 = taylor_phi(sx=+1, st=-1, order_x=4, order_t=4)
\[\displaystyle \frac{\Delta_{t}^{4} \frac{\partial^{4}}{\partial t^{4}} \phi{\left(x,t \right)}}{24} + \frac{\Delta_{t}^{3} \frac{\partial^{3}}{\partial t^{3}} \phi{\left(x,t \right)}}{6} + \frac{\Delta_{t}^{2} \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)}}{2} + \Delta_{t} \frac{\partial}{\partial t} \phi{\left(x,t \right)} + \phi{\left(x,t \right)}\]
\[\displaystyle \phi{\left(x,t \right)}\]
\[\displaystyle \frac{\Delta_{t}^{4} \frac{\partial^{4}}{\partial t^{4}} \phi{\left(x,t \right)}}{24} - \frac{\Delta_{t}^{3} \frac{\partial^{3}}{\partial t^{3}} \phi{\left(x,t \right)}}{6} + \frac{\Delta_{t}^{2} \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)}}{2} - \Delta_{t} \frac{\partial}{\partial t} \phi{\left(x,t \right)} + \phi{\left(x,t \right)}\]
\[\displaystyle \frac{2 \Delta_{t}^{4} \frac{\partial^{4}}{\partial t^{4}} \phi{\left(x,t \right)}}{3} - \frac{4 \Delta_{t}^{3} \frac{\partial^{3}}{\partial t^{3}} \phi{\left(x,t \right)}}{3} + 2 \Delta_{t}^{2} \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)} - 2 \Delta_{t} \frac{\partial}{\partial t} \phi{\left(x,t \right)} + \phi{\left(x,t \right)}\]
\[\displaystyle \frac{\Delta_{x}^{4} \frac{\partial^{4}}{\partial x^{4}} \phi{\left(x,t \right)}}{24} - \frac{\Delta_{x}^{3} \frac{\partial^{3}}{\partial x^{3}} \phi{\left(x,t \right)}}{6} + \frac{\Delta_{x}^{2} \frac{\partial^{2}}{\partial x^{2}} \phi{\left(x,t \right)}}{2} - \Delta_{x} \frac{\partial}{\partial x} \phi{\left(x,t \right)} + \phi{\left(x,t \right)}\]
\[\displaystyle \frac{\Delta_{x}^{4} \frac{\partial^{4}}{\partial x^{4}} \phi{\left(x,t \right)}}{24} + \frac{\Delta_{x}^{3} \frac{\partial^{3}}{\partial x^{3}} \phi{\left(x,t \right)}}{6} + \frac{\Delta_{x}^{2} \frac{\partial^{2}}{\partial x^{2}} \phi{\left(x,t \right)}}{2} + \Delta_{x} \frac{\partial}{\partial x} \phi{\left(x,t \right)} + \phi{\left(x,t \right)}\]
# -----------------------------------------------------------------------------
# 5) Build the discrete equation (expanded) and move everything to one side
# -----------------------------------------------------------------------------
lhs = phi_t_p1

rhs = (alpha1*phi_xm1_t0 + alpha2*phi_t_0 + alpha1*phi_xp1_t0
       + beta1*phi_xm1_tm1 + beta2*phi_t_m1 + beta1*phi_xp1_tm1
       + Omega**2 * phi_t_m2)

residual = sp.simplify(sp.expand(lhs - rhs))  # = 0 is the modified equation
max_dx=4
max_dt=2
max_total=4
expr = sp.expand(residual)
poly = sp.Poly(expr, dx, dt, domain='EX')  # keep symbolic coeffs
kept = sp.S(0)

for (px, pt), coeff in poly.terms():
    if (px <= max_dx) and (pt <= max_dt) and ((px + pt) <= max_total):
        kept += coeff * dx**px * dt**pt
# display(sp.simplify(residual))
display(sp.simplify(kept.subs(dt**2*dx**2,0)))
\[\displaystyle \frac{\Delta_{t}^{2} \omega \left(4 - 3 \omega\right) \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)}}{2} - \Delta_{t} \Delta_{x}^{2} \left(a \omega^{2} - a \omega - \omega + 1\right) \frac{\partial^{3}}{\partial x^{2}\partial t} \phi{\left(x,t \right)} + \Delta_{t} \omega^{2} \frac{\partial}{\partial t} \phi{\left(x,t \right)} + \frac{\Delta_{x}^{4} a \omega \left(\omega - 2\right) \frac{\partial^{4}}{\partial x^{4}} \phi{\left(x,t \right)}}{12} + \Delta_{x}^{2} a \omega \left(\omega - 2\right) \frac{\partial^{2}}{\partial x^{2}} \phi{\left(x,t \right)}\]