{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "#### Sympy code for FDM expansion of LBM diffusive equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "# -----------------------------------------------------------------------------\n",
    "# 1) Symbols (parameters and step sizes)\n",
    "# -----------------------------------------------------------------------------\n",
    "dx, dt = sp.symbols('Delta_x Delta_t', positive=True, real=True)\n",
    "omega, a = sp.symbols('omega a', real=True)\n",
    "\n",
    "Omega = 1 - omega\n",
    "\n",
    "alpha1 = sp.simplify(Omega + a*omega)\n",
    "alpha2 = sp.simplify(Omega + (1 - 2*a)*omega)\n",
    "beta1  = sp.simplify(-Omega*(Omega + (1 - a)*omega))\n",
    "beta2  = sp.simplify(-Omega*(Omega + 2*a*omega))\n",
    "\n",
    "# -----------------------------------------------------------------------------\n",
    "# 2) Continuous variables and field\n",
    "# -----------------------------------------------------------------------------\n",
    "x, t = sp.symbols('x t', real=True)\n",
    "phi = sp.Function('phi')\n",
    "# Shorthand for the base (continuous) field value at (x,t)\n",
    "PHI = phi(x, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def taylor_phi(sx=0, st=0, order_x=4, order_t=4):\n",
    "    expr = sp.S(0)\n",
    "    # Keep all mixed terms up to the specified individual orders.\n",
    "    for m in range(order_x + 1):\n",
    "        for n in range(order_t + 1):\n",
    "            # skip the (0,0) term? no, keep it.\n",
    "            term = ( (sx*dx)**m * (st*dt)**n\n",
    "                    / (sp.factorial(m)*sp.factorial(n))\n",
    "                    * sp.diff(PHI, x, m, t, n) )\n",
    "            expr += term\n",
    "    return sp.expand(expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\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)}$"
      ],
      "text/plain": [
       "Delta_t**4*Derivative(phi(x, t), (t, 4))/24 + Delta_t**3*Derivative(phi(x, t), (t, 3))/6 + Delta_t**2*Derivative(phi(x, t), (t, 2))/2 + Delta_t*Derivative(phi(x, t), t) + phi(x, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\phi{\\left(x,t \\right)}$"
      ],
      "text/plain": [
       "phi(x, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\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)}$"
      ],
      "text/plain": [
       "Delta_t**4*Derivative(phi(x, t), (t, 4))/24 - Delta_t**3*Derivative(phi(x, t), (t, 3))/6 + Delta_t**2*Derivative(phi(x, t), (t, 2))/2 - Delta_t*Derivative(phi(x, t), t) + phi(x, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\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)}$"
      ],
      "text/plain": [
       "2*Delta_t**4*Derivative(phi(x, t), (t, 4))/3 - 4*Delta_t**3*Derivative(phi(x, t), (t, 3))/3 + 2*Delta_t**2*Derivative(phi(x, t), (t, 2)) - 2*Delta_t*Derivative(phi(x, t), t) + phi(x, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\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)}$"
      ],
      "text/plain": [
       "Delta_x**4*Derivative(phi(x, t), (x, 4))/24 - Delta_x**3*Derivative(phi(x, t), (x, 3))/6 + Delta_x**2*Derivative(phi(x, t), (x, 2))/2 - Delta_x*Derivative(phi(x, t), x) + phi(x, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\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)}$"
      ],
      "text/plain": [
       "Delta_x**4*Derivative(phi(x, t), (x, 4))/24 + Delta_x**3*Derivative(phi(x, t), (x, 3))/6 + Delta_x**2*Derivative(phi(x, t), (x, 2))/2 + Delta_x*Derivative(phi(x, t), x) + phi(x, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# -----------------------------------------------------------------------------\n",
    "# 4) Discrete stencil terms (expanded)\n",
    "# -----------------------------------------------------------------------------\n",
    "# phi_x^{t+1}, phi_x^{t}, phi_x^{t-1}, phi_x^{t-2}\n",
    "phi_t_p1 = taylor_phi(sx=0,  st=+1, order_x=0, order_t=4)\n",
    "display(phi_t_p1)\n",
    "phi_t_0  = taylor_phi(sx=0,  st=0,  order_x=0, order_t=0)  # exactly phi(x,t)\n",
    "display(phi_t_0)\n",
    "phi_t_m1 = taylor_phi(sx=0,  st=-1, order_x=0, order_t=4)\n",
    "display(phi_t_m1)\n",
    "phi_t_m2 = taylor_phi(sx=0,  st=-2, order_x=0, order_t=4)\n",
    "display(phi_t_m2)\n",
    "\n",
    "# phi_{x±1}^{t}\n",
    "phi_xm1_t0 = taylor_phi(sx=-1, st=0,  order_x=4, order_t=0)\n",
    "display(phi_xm1_t0)\n",
    "phi_xp1_t0 = taylor_phi(sx=+1, st=0,  order_x=4, order_t=0)\n",
    "display(phi_xp1_t0)\n",
    "\n",
    "# phi_{x±1}^{t-1}\n",
    "phi_xm1_tm1 = taylor_phi(sx=-1, st=-1, order_x=4, order_t=4)\n",
    "phi_xp1_tm1 = taylor_phi(sx=+1, st=-1, order_x=4, order_t=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------\n",
    "# 5) Build the discrete equation (expanded) and move everything to one side\n",
    "# -----------------------------------------------------------------------------\n",
    "lhs = phi_t_p1\n",
    "\n",
    "rhs = (alpha1*phi_xm1_t0 + alpha2*phi_t_0 + alpha1*phi_xp1_t0\n",
    "       + beta1*phi_xm1_tm1 + beta2*phi_t_m1 + beta1*phi_xp1_tm1\n",
    "       + Omega**2 * phi_t_m2)\n",
    "\n",
    "residual = sp.simplify(sp.expand(lhs - rhs))  # = 0 is the modified equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_dx=4\n",
    "max_dt=2\n",
    "max_total=4\n",
    "expr = sp.expand(residual)\n",
    "poly = sp.Poly(expr, dx, dt, domain='EX')  # keep symbolic coeffs\n",
    "kept = sp.S(0)\n",
    "\n",
    "for (px, pt), coeff in poly.terms():\n",
    "    if (px <= max_dx) and (pt <= max_dt) and ((px + pt) <= max_total):\n",
    "        kept += coeff * dx**px * dt**pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# display(sp.simplify(residual))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\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)}$"
      ],
      "text/plain": [
       "Delta_t**2*omega*(4 - 3*omega)*Derivative(phi(x, t), (t, 2))/2 - Delta_t*Delta_x**2*(a*omega**2 - a*omega - omega + 1)*Derivative(phi(x, t), t, (x, 2)) + Delta_t*omega**2*Derivative(phi(x, t), t) + Delta_x**4*a*omega*(omega - 2)*Derivative(phi(x, t), (x, 4))/12 + Delta_x**2*a*omega*(omega - 2)*Derivative(phi(x, t), (x, 2))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(sp.simplify(kept.subs(dt**2*dx**2,0)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}