{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pystencils as ps\n",
    "import sympy as sp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Demo: Working with derivatives\n",
    "\n",
    "\n",
    "## Overview\n",
    "This notebook demonstrates how to formulate continuous differential operators in *pystencils* and automatically derive finite difference stencils from them.\n",
    "\n",
    "Instead of using the built-in derivatives in *sympy*, *pystencils* comes with its own derivative objects. They represent spatial derivatives of pystencils fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} {f}_{(0,0)}}$"
      ],
      "text/plain": [
       "Diff(f_C, 0, -1)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f = ps.fields(\"f: [2D]\")\n",
    "first_derivative_x = ps.fd.diff(f, 0)\n",
    "first_derivative_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This object is the derivative of the field $f$ with respect to the first spatial coordinate $x$. To get a finite difference approximation a discretization strategy is required:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{{f}_{(1,0)} - {f}_{(-1,0)}}{2 h}$"
      ],
      "text/plain": [
       "(f_E - f_W)/(2*h)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "discretize_2nd_order = ps.fd.Discretization2ndOrder(dx=sp.symbols(\"h\"))\n",
    "discretize_2nd_order(first_derivative_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Strictly speaking, derivative objects act on *field accesses*, not *fields*, that why there is a $(0,0)$ index at the field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} {f}_{(0,0)}}$"
      ],
      "text/plain": [
       "Diff(f_C, 0, -1)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "first_derivative_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes it might be useful to specify derivatives at an offset e.g."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Diff(f_N, 0, -1), (f_NE - f_NW)/(2*h))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "derivative_offset = ps.fd.diff(f[0, 1], 0)\n",
    "derivative_offset, discretize_2nd_order(derivative_offset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another example with second order derivatives:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} {\\partial_{0} {f}_{(0,0)}}} + {\\partial_{1} {\\partial_{1} {f}_{(0,0)}}}$"
      ],
      "text/plain": [
       "Diff(Diff(f_C, 0, -1), 0, -1) + Diff(Diff(f_C, 1, -1), 1, -1)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "laplacian = ps.fd.diff(f, 0, 0) + ps.fd.diff(f, 1, 1)\n",
    "laplacian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{- 2 {f}_{(0,0)} + {f}_{(1,0)} + {f}_{(-1,0)}}{h^{2}} + \\frac{- 2 {f}_{(0,0)} + {f}_{(0,1)} + {f}_{(0,-1)}}{h^{2}}$"
      ],
      "text/plain": [
       "(-2*f_C + f_E + f_W)/h**2 + (-2*f_C + f_N + f_S)/h**2"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "discretize_2nd_order(laplacian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with derivative terms\n",
    "\n",
    "No automatic simplifications are done on derivative terms i.e. linearity relations or product rule are not applied automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} (c + {\\partial_{0} {f}_{(0,0)}} + {\\partial_{0} {g}_{(0,0)}} + 5) }$"
      ],
      "text/plain": [
       "Diff(c + Diff(f_C, 0, -1) + Diff(g_C, 0, -1) + 5, 0, -1)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f, g = ps.fields(\"f, g :[2D]\")\n",
    "c = sp.symbols(\"c\")\n",
    "δ = ps.fd.diff\n",
    "\n",
    "expr = δ( δ(f, 0) +  δ(g, 0) + c + 5 , 0)\n",
    "expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This nested term can not be discretized automatically. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Only derivatives with field or field accesses as arguments can be discretized\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    discretize_2nd_order(expr)\n",
    "except ValueError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Linearity\n",
    "The following function expands all derivatives exploiting linearity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} c} + {\\partial_{0} {\\partial_{0} {f}_{(0,0)}}} + {\\partial_{0} {\\partial_{0} {g}_{(0,0)}}}$"
      ],
      "text/plain": [
       "Diff(c, 0, -1) + Diff(Diff(f_C, 0, -1), 0, -1) + Diff(Diff(g_C, 0, -1), 0, -1)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps.fd.expand_diff_linear(expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The symbol $c$ that was included is interpreted as a function by default. \n",
    "We can control the simplification behaviour by specifying all functions or all constants:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} {\\partial_{0} {f}_{(0,0)}}} + {\\partial_{0} {\\partial_{0} {g}_{(0,0)}}}$"
      ],
      "text/plain": [
       "Diff(Diff(f_C, 0, -1), 0, -1) + Diff(Diff(g_C, 0, -1), 0, -1)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps.fd.expand_diff_linear(expr, functions=(f[0,0], g[0, 0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} {\\partial_{0} {f}_{(0,0)}}} + {\\partial_{0} {\\partial_{0} {g}_{(0,0)}}}$"
      ],
      "text/plain": [
       "Diff(Diff(f_C, 0, -1), 0, -1) + Diff(Diff(g_C, 0, -1), 0, -1)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps.fd.expand_diff_linear(expr, constants=[c])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The expanded term can then be discretized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{- 2 {f}_{(0,0)} + {f}_{(1,0)} + {f}_{(-1,0)}}{h^{2}} + \\frac{- 2 {g}_{(0,0)} + {g}_{(1,0)} + {g}_{(-1,0)}}{h^{2}}$"
      ],
      "text/plain": [
       "(-2*f_C + f_E + f_W)/h**2 + (-2*g_C + g_E + g_W)/h**2"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "discretize_2nd_order(ps.fd.expand_diff_linear(expr, constants=[c]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Product rule\n",
    "\n",
    "The next cells show how to apply product rule and its reverse:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} ({f}_{(0,0)} {g}_{(0,0)}) }$"
      ],
      "text/plain": [
       "Diff(f_C*g_C, 0, -1)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expr = δ(f[0, 0] * g[0, 0], 0 )\n",
    "expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {f}_{(0,0)} {\\partial_{0} {g}_{(0,0)}} + {g}_{(0,0)} {\\partial_{0} {f}_{(0,0)}}$"
      ],
      "text/plain": [
       "f_C*Diff(g_C, 0, -1) + g_C*Diff(f_C, 0, -1)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expanded_expr = ps.fd.expand_diff_products(expr)\n",
    "expanded_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} ({f}_{(0,0)} {g}_{(0,0)}) }$"
      ],
      "text/plain": [
       "Diff(f_C*g_C, 0, -1)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recombined_expr = ps.fd.combine_diff_products(expanded_expr)\n",
    "recombined_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert recombined_expr == expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate derivatives\n",
    "\n",
    "Arguments of derivative need not be to be fields, only when trying to discretize them. \n",
    "The next cells show how to transform them to *sympy* derivatives and evaluate them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle {\\partial_{0} (k^{3} + 2 k) }$"
      ],
      "text/plain": [
       "Diff(k**3 + 2*k, 0, -1)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k = sp.symbols(\"k\")\n",
    "expr = δ(k**3 + 2 * k, 0 )\n",
    "expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 3 k^{2} + 2$"
      ],
      "text/plain": [
       "3*k**2 + 2"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps.fd.evaluate_diffs(expr, var=k)"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}