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)

def d(name, nx=0, nt=0):
    n = nx + nt
    if n == 1:
        num = rf"\partial{{{name}}}"
    else:
        num = rf"\partial^{{{n}}}{{{name}}}"
    den = ""
    if nx:
        den += rf"\partial{{x^{{{nx}}}}}" if nx > 1 else r"\partial{x}"
    if nt:
        den += rf"\partial{{t^{{{nt}}}}}" if nt > 1 else r"\partial{t}"
    return sp.Symbol(rf"\frac{{{num}}}{{{den}}}")
# -------------------1st-Order-Terms----------------------
dphit = d(r"\phi", nt=1); dphix   = d(r"\phi", nx=1)
# -------------------2nd-Order-Terms----------------------
dphix2 = d(r"\phi", nx=2); dphixt = d(r"\phi", nx=1, nt=1); dphit2 = d(r"\phi", nt=2)
# -------------------3rd-Order-Terms----------------------
dphix3 = d(r"\phi", nx=3); dphix2t = d(r"\phi", nx=2, nt=1); dphixt2 = d(r"\phi", nx=1, nt=2); dphit3  = d(r"\phi", nt=3)
# -------------------4th-Order-Terms----------------------
dphix4 = d(r"\phi", nx=4); dphix3t = d(r"\phi", nx=3, nt=1); dphix2t2 = d(r"\phi", nx=2, nt=2); dphixt3 = d(r"\phi", nx=1, nt=3); dphit4 = d(r"\phi", nt=4)

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)
from IPython.display import display, Math
# 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( Math(r"\phi_x^{t+1} = " + sp.latex(phi_t_p1)))
phi_t_0  = taylor_phi(sx=0,  st=0,  order_x=0, order_t=0)  # exactly phi(x,t)
display( Math(r"\phi_x^{t} = " + sp.latex(phi_t_0)))
phi_t_m1 = taylor_phi(sx=0,  st=-1, order_x=0, order_t=4)
display( Math(r"\phi_x^{t-1} = " + sp.latex(phi_t_m1)))
phi_t_m2 = taylor_phi(sx=0,  st=-2, order_x=0, order_t=4)
display( Math(r"\phi_x^{t-2} = " + sp.latex(phi_t_m2)))

# phi_{x±1}^{t}
phi_xm1_t0 = taylor_phi(sx=-1, st=0,  order_x=4, order_t=0)
display( Math(r"\phi_{x-1}^{t} = " + sp.latex(phi_xm1_t0)))
phi_xp1_t0 = taylor_phi(sx=+1, st=0,  order_x=4, order_t=0)
display( Math(r"\phi_{x+1}^{t} = " + sp.latex(phi_xp1_t0)))

# phi_{x±1}^{t-1}
phi_xm1_tm1 = taylor_phi(sx=-1, st=-1, order_x=4, order_t=4)
display( Math(r"\phi_{x-1}^{t-1} = " + sp.latex(phi_xm1_tm1)))
phi_xp1_tm1 = taylor_phi(sx=+1, st=-1, order_x=4, order_t=4)
display( Math(r"\phi_{x+1}^{t-1} = " + sp.latex(phi_xp1_tm1)))
\[\displaystyle \phi_x^{t+1} = \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_x^{t} = \phi{\left(x,t \right)}\]
\[\displaystyle \phi_x^{t-1} = \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_x^{t-2} = \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 \phi_{x-1}^{t} = \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 \phi_{x+1}^{t} = \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 \phi_{x-1}^{t-1} = \frac{\Delta_{t}^{4} \Delta_{x}^{4} \frac{\partial^{8}}{\partial x^{4}\partial t^{4}} \phi{\left(x,t \right)}}{576} - \frac{\Delta_{t}^{4} \Delta_{x}^{3} \frac{\partial^{7}}{\partial x^{3}\partial t^{4}} \phi{\left(x,t \right)}}{144} + \frac{\Delta_{t}^{4} \Delta_{x}^{2} \frac{\partial^{6}}{\partial x^{2}\partial t^{4}} \phi{\left(x,t \right)}}{48} - \frac{\Delta_{t}^{4} \Delta_{x} \frac{\partial^{5}}{\partial x\partial t^{4}} \phi{\left(x,t \right)}}{24} + \frac{\Delta_{t}^{4} \frac{\partial^{4}}{\partial t^{4}} \phi{\left(x,t \right)}}{24} - \frac{\Delta_{t}^{3} \Delta_{x}^{4} \frac{\partial^{7}}{\partial x^{4}\partial t^{3}} \phi{\left(x,t \right)}}{144} + \frac{\Delta_{t}^{3} \Delta_{x}^{3} \frac{\partial^{6}}{\partial x^{3}\partial t^{3}} \phi{\left(x,t \right)}}{36} - \frac{\Delta_{t}^{3} \Delta_{x}^{2} \frac{\partial^{5}}{\partial x^{2}\partial t^{3}} \phi{\left(x,t \right)}}{12} + \frac{\Delta_{t}^{3} \Delta_{x} \frac{\partial^{4}}{\partial x\partial t^{3}} \phi{\left(x,t \right)}}{6} - \frac{\Delta_{t}^{3} \frac{\partial^{3}}{\partial t^{3}} \phi{\left(x,t \right)}}{6} + \frac{\Delta_{t}^{2} \Delta_{x}^{4} \frac{\partial^{6}}{\partial x^{4}\partial t^{2}} \phi{\left(x,t \right)}}{48} - \frac{\Delta_{t}^{2} \Delta_{x}^{3} \frac{\partial^{5}}{\partial x^{3}\partial t^{2}} \phi{\left(x,t \right)}}{12} + \frac{\Delta_{t}^{2} \Delta_{x}^{2} \frac{\partial^{4}}{\partial x^{2}\partial t^{2}} \phi{\left(x,t \right)}}{4} - \frac{\Delta_{t}^{2} \Delta_{x} \frac{\partial^{3}}{\partial x\partial t^{2}} \phi{\left(x,t \right)}}{2} + \frac{\Delta_{t}^{2} \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)}}{2} - \frac{\Delta_{t} \Delta_{x}^{4} \frac{\partial^{5}}{\partial x^{4}\partial t} \phi{\left(x,t \right)}}{24} + \frac{\Delta_{t} \Delta_{x}^{3} \frac{\partial^{4}}{\partial x^{3}\partial t} \phi{\left(x,t \right)}}{6} - \frac{\Delta_{t} \Delta_{x}^{2} \frac{\partial^{3}}{\partial x^{2}\partial t} \phi{\left(x,t \right)}}{2} + \Delta_{t} \Delta_{x} \frac{\partial^{2}}{\partial x\partial t} \phi{\left(x,t \right)} - \Delta_{t} \frac{\partial}{\partial t} \phi{\left(x,t \right)} + \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 \phi_{x+1}^{t-1} = \frac{\Delta_{t}^{4} \Delta_{x}^{4} \frac{\partial^{8}}{\partial x^{4}\partial t^{4}} \phi{\left(x,t \right)}}{576} + \frac{\Delta_{t}^{4} \Delta_{x}^{3} \frac{\partial^{7}}{\partial x^{3}\partial t^{4}} \phi{\left(x,t \right)}}{144} + \frac{\Delta_{t}^{4} \Delta_{x}^{2} \frac{\partial^{6}}{\partial x^{2}\partial t^{4}} \phi{\left(x,t \right)}}{48} + \frac{\Delta_{t}^{4} \Delta_{x} \frac{\partial^{5}}{\partial x\partial t^{4}} \phi{\left(x,t \right)}}{24} + \frac{\Delta_{t}^{4} \frac{\partial^{4}}{\partial t^{4}} \phi{\left(x,t \right)}}{24} - \frac{\Delta_{t}^{3} \Delta_{x}^{4} \frac{\partial^{7}}{\partial x^{4}\partial t^{3}} \phi{\left(x,t \right)}}{144} - \frac{\Delta_{t}^{3} \Delta_{x}^{3} \frac{\partial^{6}}{\partial x^{3}\partial t^{3}} \phi{\left(x,t \right)}}{36} - \frac{\Delta_{t}^{3} \Delta_{x}^{2} \frac{\partial^{5}}{\partial x^{2}\partial t^{3}} \phi{\left(x,t \right)}}{12} - \frac{\Delta_{t}^{3} \Delta_{x} \frac{\partial^{4}}{\partial x\partial t^{3}} \phi{\left(x,t \right)}}{6} - \frac{\Delta_{t}^{3} \frac{\partial^{3}}{\partial t^{3}} \phi{\left(x,t \right)}}{6} + \frac{\Delta_{t}^{2} \Delta_{x}^{4} \frac{\partial^{6}}{\partial x^{4}\partial t^{2}} \phi{\left(x,t \right)}}{48} + \frac{\Delta_{t}^{2} \Delta_{x}^{3} \frac{\partial^{5}}{\partial x^{3}\partial t^{2}} \phi{\left(x,t \right)}}{12} + \frac{\Delta_{t}^{2} \Delta_{x}^{2} \frac{\partial^{4}}{\partial x^{2}\partial t^{2}} \phi{\left(x,t \right)}}{4} + \frac{\Delta_{t}^{2} \Delta_{x} \frac{\partial^{3}}{\partial x\partial t^{2}} \phi{\left(x,t \right)}}{2} + \frac{\Delta_{t}^{2} \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)}}{2} - \frac{\Delta_{t} \Delta_{x}^{4} \frac{\partial^{5}}{\partial x^{4}\partial t} \phi{\left(x,t \right)}}{24} - \frac{\Delta_{t} \Delta_{x}^{3} \frac{\partial^{4}}{\partial x^{3}\partial t} \phi{\left(x,t \right)}}{6} - \frac{\Delta_{t} \Delta_{x}^{2} \frac{\partial^{3}}{\partial x^{2}\partial t} \phi{\left(x,t \right)}}{2} - \Delta_{t} \Delta_{x} \frac{\partial^{2}}{\partial x\partial t} \phi{\left(x,t \right)} - \Delta_{t} \frac{\partial}{\partial t} \phi{\left(x,t \right)} + \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
kept
\[\displaystyle \Delta_{t}^{2} \Delta_{x}^{2} \left(\frac{a \omega^{2} \frac{\partial^{4}}{\partial x^{2}\partial t^{2}} \phi{\left(x,t \right)}}{2} - \frac{a \omega \frac{\partial^{4}}{\partial x^{2}\partial t^{2}} \phi{\left(x,t \right)}}{2} - \frac{\omega \frac{\partial^{4}}{\partial x^{2}\partial t^{2}} \phi{\left(x,t \right)}}{2} + \frac{\frac{\partial^{4}}{\partial x^{2}\partial t^{2}} \phi{\left(x,t \right)}}{2}\right) + \Delta_{t}^{2} \left(- \frac{3 \omega^{2} \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)}}{2} + 2 \omega \frac{\partial^{2}}{\partial t^{2}} \phi{\left(x,t \right)}\right) + \Delta_{t} \Delta_{x}^{2} \left(- a \omega^{2} \frac{\partial^{3}}{\partial x^{2}\partial t} \phi{\left(x,t \right)} + a \omega \frac{\partial^{3}}{\partial x^{2}\partial t} \phi{\left(x,t \right)} + \omega \frac{\partial^{3}}{\partial x^{2}\partial t} \phi{\left(x,t \right)} - \frac{\partial^{3}}{\partial x^{2}\partial t} \phi{\left(x,t \right)}\right) + \Delta_{t} \omega^{2} \frac{\partial}{\partial t} \phi{\left(x,t \right)} + \Delta_{x}^{4} \left(\frac{a \omega^{2} \frac{\partial^{4}}{\partial x^{4}} \phi{\left(x,t \right)}}{12} - \frac{a \omega \frac{\partial^{4}}{\partial x^{4}} \phi{\left(x,t \right)}}{6}\right) + \Delta_{x}^{2} \left(a \omega^{2} \frac{\partial^{2}}{\partial x^{2}} \phi{\left(x,t \right)} - 2 a \omega \frac{\partial^{2}}{\partial x^{2}} \phi{\left(x,t \right)}\right)\]
term1 = sp.expand(kept).subs({
    # ------------------------------------------------4th order--------------------------------------------------------------------------------------------
    PHI.diff(x, 4):dphix4,PHI.diff(x, 3, t):dphix3t,PHI.diff(x, 2, t, 2):dphix2t2,PHI.diff(x, t, 3):dphixt3,PHI.diff(t, 4):dphit4,
    # ------------------------------------------------3rd order--------------------------------------------------------------------------------------------
    PHI.diff(x, 3):dphix3, PHI.diff(x, 2, t):dphix2t, PHI.diff(x, t, 2):dphixt2, PHI.diff(t, 3):dphit3,
    # ------------------------------------------------2nd order--------------------------------------------------------------------------------------------
    PHI.diff(x, 2):dphix2, PHI.diff(x, t):dphixt, PHI.diff(t, 2):dphit2,
    # ------------------------------------------------1st order--------------------------------------------------------------------------------------------
    PHI.diff(x):dphix, PHI.diff(t):dphit,
})
display(term1)
\[\displaystyle \frac{\Delta_{t}^{2} \Delta_{x}^{2} \frac{\partial^{4}{\phi}}{\partial{x^{2}}\partial{t^{2}}} a \omega^{2}}{2} - \frac{\Delta_{t}^{2} \Delta_{x}^{2} \frac{\partial^{4}{\phi}}{\partial{x^{2}}\partial{t^{2}}} a \omega}{2} - \frac{\Delta_{t}^{2} \Delta_{x}^{2} \frac{\partial^{4}{\phi}}{\partial{x^{2}}\partial{t^{2}}} \omega}{2} + \frac{\Delta_{t}^{2} \Delta_{x}^{2} \frac{\partial^{4}{\phi}}{\partial{x^{2}}\partial{t^{2}}}}{2} - \frac{3 \Delta_{t}^{2} \frac{\partial^{2}{\phi}}{\partial{t^{2}}} \omega^{2}}{2} + 2 \Delta_{t}^{2} \frac{\partial^{2}{\phi}}{\partial{t^{2}}} \omega - \Delta_{t} \Delta_{x}^{2} \frac{\partial^{3}{\phi}}{\partial{x^{2}}\partial{t}} a \omega^{2} + \Delta_{t} \Delta_{x}^{2} \frac{\partial^{3}{\phi}}{\partial{x^{2}}\partial{t}} a \omega + \Delta_{t} \Delta_{x}^{2} \frac{\partial^{3}{\phi}}{\partial{x^{2}}\partial{t}} \omega - \Delta_{t} \Delta_{x}^{2} \frac{\partial^{3}{\phi}}{\partial{x^{2}}\partial{t}} + \Delta_{t} \frac{\partial{\phi}}{\partial{t}} \omega^{2} + \frac{\Delta_{x}^{4} \frac{\partial^{4}{\phi}}{\partial{x^{4}}} a \omega^{2}}{12} - \frac{\Delta_{x}^{4} \frac{\partial^{4}{\phi}}{\partial{x^{4}}} a \omega}{6} + \Delta_{x}^{2} \frac{\partial^{2}{\phi}}{\partial{x^{2}}} a \omega^{2} - 2 \Delta_{x}^{2} \frac{\partial^{2}{\phi}}{\partial{x^{2}}} a \omega\]
term2 = sp.collect(sp.expand(-term1/(dt*omega**2))+ dphit,[dphit,dphix,
                                                         dphix2*dx*dx/dt,dphixt,dphit2*dt,
                                                         dphix2t*dx*dx,dphixt2,dphix3,dphit3*dt**2,
                                                         dphix2t2*dt*dx**2,dphix4*dx**4/dt,
                                                        ]).subs(dphix2t2,0)
display( Math(r"\frac{\partial\phi}{\partial t}= " + sp.latex(term2)))
\[\displaystyle \frac{\partial\phi}{\partial t}= \Delta_{t} \frac{\partial^{2}{\phi}}{\partial{t^{2}}} \left(\frac{3}{2} - \frac{2}{\omega}\right) + \Delta_{x}^{2} \frac{\partial^{3}{\phi}}{\partial{x^{2}}\partial{t}} \left(a - \frac{a}{\omega} - \frac{1}{\omega} + \frac{1}{\omega^{2}}\right) + \frac{\Delta_{x}^{4} \frac{\partial^{4}{\phi}}{\partial{x^{4}}} \left(- \frac{a}{12} + \frac{a}{6 \omega}\right)}{\Delta_{t}} + \frac{\Delta_{x}^{2} \frac{\partial^{2}{\phi}}{\partial{x^{2}}} \left(- a + \frac{2 a}{\omega}\right)}{\Delta_{t}}\]
nus = sp.symbols('\\nu')
nut=dx**2/dt*(-a+(2*a/omega))
term3 = sp.collect(sp.expand(term2.subs(nut,nus)
                       .subs(dphit2,nut*nut*dphix4)
                       .subs(dphix2t,nut*dphix4)), [dphix4*dx**4/dt])
dphix4_term3=sp.simplify(sp.factor(term3.coeff(dphix4)).subs(dt*omega,a*(2-omega)*dx**2/nus))
# display(dphix4_term3)
display( Math(r"\frac{\partial\phi}{\partial t}= " + sp.latex(term3.subs(term3.coeff(dphix4),dphix4_term3))))
\[\displaystyle \frac{\partial\phi}{\partial t}= \frac{\Delta_{x}^{2} \frac{\partial^{4}{\phi}}{\partial{x^{4}}} \nu \left(- 6 a \omega^{2} + 48 a \omega - 48 a + \omega^{2} - 12 \omega + 12\right)}{12 \omega^{2}} + \frac{\partial^{2}{\phi}}{\partial{x^{2}}} \nu\]