{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4e52e00c-f620-481f-b7d4-50212d7e7c57",
   "metadata": {},
   "source": [
    "# Variables, equations and forcing terms\n",
    "\n",
    "This demo notebook summarizes the variables, equations and forcings that can be used in **PlesioGeostroPy**.\n",
    "\n",
    "It is assumed that readers are already familiar with the collection interface (see [Demo_Collections](Demo_Collections.ipynb))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c1f963e1-870f-4fe9-9753-5a378a68c496",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "root_dir = \".\"\n",
    "\n",
    "# The following 2 lines are a hack to import from parent directory\n",
    "# If the notebook is run in the root directory, comment out these 2 lines\n",
    "sys.path.append(os.path.dirname(os.getcwd()))\n",
    "root_dir = \"..\"\n",
    "\n",
    "from pg_utils.pg_model import core, base\n",
    "from sympy import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97b0a658-586f-49ff-8792-e179abbdbde9",
   "metadata": {},
   "source": [
    "---\n",
    "## Core variables\n",
    "\n",
    "The core variables that are concerned in the model are constructed in `pg_model.core` module.\n",
    "These first include the coordinates $s$, $\\phi$, $z$; Cartesian coordinates $x$, $y$; as well as spherical coordinates $r$, $\\theta$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "613729c3-1e30-4ab7-b901-7a3a33cc2a7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cylindrical oordinates:\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle s$"
      ],
      "text/plain": [
       "s"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\phi$"
      ],
      "text/plain": [
       "\\phi"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle z$"
      ],
      "text/plain": [
       "z"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(\"Cylindrical oordinates:\")\n",
    "display(*[core.cyl[i_comp] for i_comp in range(3)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59ae0d88-1f64-4d2b-894a-b2e4878fd696",
   "metadata": {},
   "source": [
    "Next, we have the 3-D vector fields. These include\n",
    "- magnetic field: in cylindrical coordinates: `B_vec`, background field `B0_vec`; in spherical coordinates `B_sph`, background field `B0_sph`;\n",
    "- velocity field: in cylindrical coordinates: `U_vec`, background field `U0_vec`; in spherical coordinates `U_sph`, background field `U0_sph`;\n",
    "\n",
    "we can take a look at one of these variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2b8451f7-95d5-4fc2-9ebf-217683ec8585",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{s}{\\left(s,\\phi,z,t \\right)}$"
      ],
      "text/plain": [
       "B_s(s, \\phi, z, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{\\phi}{\\left(s,\\phi,z,t \\right)}$"
      ],
      "text/plain": [
       "B_\\phi(s, \\phi, z, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{z}{\\left(s,\\phi,z,t \\right)}$"
      ],
      "text/plain": [
       "B_z(s, \\phi, z, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(*[core.B_vec[i_comp] for i_comp in range(3)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd9f0894-eaae-4774-8292-d3e95b6d13cf",
   "metadata": {},
   "source": [
    "The variables that represent quantities under PG formulation can be divided into two classes.\n",
    "The first class contains the PG variables; the second class contains the transformed variables. Please refer to [formulation PDF](https://gentlemin.github.io/assets/pdf/Ingredients.pdf) for details.\n",
    "Each of these two sets of variables have their own total field, background field and perturbation field:\n",
    "\n",
    "| Type       | PG                |  Transformed      |\n",
    "|------------|-------------------|-------------------|\n",
    "| Total      | `core.pgvar`      | `core.cgvar`      |\n",
    "| Background | `core.pgvar_bg`   | `core.cgvar_bg`   |\n",
    "| Perturbation   | `core.pgvar_ptb`  | `core.cgvar_ptb`  |\n",
    "\n",
    "Within the collection of all PG fields, we have in all 21 variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "37949c43-2196-4341-bb78-c788a66c19eb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Psi',\n",
       " 'Mss',\n",
       " 'Mpp',\n",
       " 'Msp',\n",
       " 'Msz',\n",
       " 'Mpz',\n",
       " 'zMss',\n",
       " 'zMpp',\n",
       " 'zMsp',\n",
       " 'Bs_e',\n",
       " 'Bp_e',\n",
       " 'Bz_e',\n",
       " 'dBs_dz_e',\n",
       " 'dBp_dz_e',\n",
       " 'Br_b',\n",
       " 'Bs_p',\n",
       " 'Bp_p',\n",
       " 'Bz_p',\n",
       " 'Bs_m',\n",
       " 'Bp_m',\n",
       " 'Bz_m']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "core.pgvar._field_names"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95f77962-2756-4c2c-bd3d-77054890123a",
   "metadata": {},
   "source": [
    "The first is the stream function; the second to the eight are magnetic moments; then we have the fields in the equatorial plane (ones with `_e` suffix). Finally, we have the boundary terms.\n",
    "\n",
    "Note that in standard PG, only `Br_b` = \"radial magnetic field at the boundary\" is used; however, for many problems such as the eigenvalue problem, we can use `Bs_p` (s-magnetic field at $z=+H$) - `Bz_m` (z-magnetic field at $z=-H$) instead of `Br_b`. Here is an overview of all available fields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7dd2fc98-76fd-4cd6-8a4d-ba637174ab2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\Psi{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\Psi(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{M_{ss}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\overline{M_{ss}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{M_{\\phi\\phi}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\overline{M_{\\phi\\phi}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{M_{s\\phi}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\overline{M_{s\\phi}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\widetilde{M_{sz}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\widetilde{M_{sz}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\widetilde{M_{\\phi z}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\widetilde{M_{\\phi z}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\widetilde{zM_{ss}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\widetilde{zM_{ss}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\widetilde{zM_{\\phi\\phi}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\widetilde{zM_{\\phi\\phi}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\widetilde{zM_{s\\phi}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\widetilde{zM_{s\\phi}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{s}^e{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_{s}^e(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{\\phi}^e{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_{\\phi}^e(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{z}^e{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_{z}^e(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{s, z}^e{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_{s, z}^e(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{\\phi, z}^e{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_{\\phi, z}^e(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B_{r1}{\\left(\\theta,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_{r1}(\\theta, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B^{+}_{s}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_s^+(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B^{+}_{\\phi}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_\\phi^+(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B^{+}_{z}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_z^+(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B^{-}_{s}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_s^-(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B^{-}_{\\phi}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_\\phi^-(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle B^{-}_{z}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "B_z^-(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(*[field_symb for field_symb in core.pgvar])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "955343cb-98f0-4030-8ae3-1315bacd7624",
   "metadata": {},
   "source": [
    "The same thing holds for `pgvar_bg` and `pgvar_ptb`. The background fields are denoted with an additional $0$ superscript, and the perturbed fields are denoted with lower-case letters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8715d37b-dd01-4faf-8116-f841d5a710d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{M_{ss}}^0{\\left(s,\\phi \\right)}$"
      ],
      "text/plain": [
       "\\overline{M_{ss}}^0(s, \\phi)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{m_{ss}}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\overline{m_{ss}}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(core.pgvar_bg.Mss, core.pgvar_ptb.Mss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8af6f685-e358-4e36-8176-a802827d8be3",
   "metadata": {},
   "source": [
    "We can also inspect the fields for transformed quantities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8c69d81a-7a4d-4722-bfcd-0414ac5afa94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{M_1}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\overline{M_1}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{M_1}^0{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\overline{M_1}^0(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\overline{m_1}{\\left(s,\\phi,t \\right)}$"
      ],
      "text/plain": [
       "\\overline{m_1}(s, \\phi, t)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(core.cgvar.M_1, core.cgvar_bg.M_1, core.cgvar_ptb.M_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fbe4372-f986-41d5-83c0-c2e774f8a3c9",
   "metadata": {},
   "source": [
    "Last but not least, the `core` module also stores the PG ansatz of the velocity, which can be invoked whenever necessary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2949e063-2ca4-4294-9f00-66cb55475c00",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\frac{\\partial}{\\partial \\phi} \\Psi{\\left(s,\\phi,t \\right)}}{s H{\\left(s \\right)}}$"
      ],
      "text/plain": [
       "Derivative(\\Psi(s, \\phi, t), \\phi)/(s*H(s))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - \\frac{\\frac{\\partial}{\\partial s} \\Psi{\\left(s,\\phi,t \\right)}}{H{\\left(s \\right)}}$"
      ],
      "text/plain": [
       "-Derivative(\\Psi(s, \\phi, t), s)/H(s)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{z \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\Psi{\\left(s,\\phi,t \\right)}}{s H^{2}{\\left(s \\right)}}$"
      ],
      "text/plain": [
       "z*Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), \\phi)/(s*H(s)**2)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(*[core.U_pg[i_comp] for i_comp in range(3)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc2441e4-ce13-4af2-a3d9-7f5dd70fca10",
   "metadata": {},
   "source": [
    "---\n",
    "## Equations\n",
    "\n",
    "With all the variables defined, we can now introduce the equations.\n",
    "\n",
    "The 15 original PG equations (see Jackson and Maffei [2020](https://royalsocietypublishing.org/doi/10.1098/rspa.2020.0513) or Holdenried-Chernoff [2021](https://www.research-collection.ethz.ch/handle/20.500.11850/509840)) have been typed in manually in the `equations` module.\n",
    "There these equations undergo two operations.\n",
    "- First, they are converted to their counterparts in transformed variables using `core.conjugate_to_PG` and `core.PG_to_conjugate` (see [Demo_Collections](Demo_Collections.ipynb)).\n",
    "- Then, these equations are linearized around a background field.\n",
    "\n",
    "The linearization procedure, borrowed from Daria's Mathematica notebook, is as follows: \n",
    "- each fields are expanded as the background state + $\\epsilon\\cdot$ a perturbation state. \n",
    "- all the terms that are linear in $\\epsilon$ are collected to yields the linearized form.\n",
    "  \n",
    "The fields in `equations` module are summarized as follows\n",
    "\n",
    "| Type       | PG                     |  Transformed          |\n",
    "|------------|------------------------|-----------------------|\n",
    "| Total      | `equations.eqs_pg`     | `equations.eqs_cg`    |\n",
    "| Linearized | `equations.eqs_pg_lin` | `equations.eqs_cg_lin`|\n",
    "\n",
    "These equations are similarly constructed as `base.CollectionPG` or `base.CollectionConjugate` objects.\n",
    "\n",
    "Since these equations are constructed when the module is imported, the first import of `equations` is slow.\n",
    "Loading the equations now may take 20-30s (tested on my laptop, AMD R5-4500U).\n",
    "The reason is that it takes a while to compile the equations, as they are not all written out in explicit forms in the code.\n",
    "In linearization especially, the program needs to simplify the expressions to a form where the perturbation $\\epsilon$ is taken out of all brackets.\n",
    "Another problem is that so far, compilation and derivation of equations is done sequentially, not in parallel.\n",
    "\n",
    "While the compilation of equations is not expected to be the bottleneck of the program in the end, it would still be desirable if it can be loaded faster.\n",
    "A method is to compute the equations and store them in a string that can be quickly loaded.\n",
    "This has already been done, and loading the precomputed equations are now the recommended method.\n",
    "\n",
    "The precomputed equations are stored as follows\n",
    "\n",
    "| Type       | PG                     |  Transformed          |\n",
    "|------------|------------------------|-----------------------|\n",
    "| Total      | /out/symbolic/eqs_pg.json     | /out/symbolic/eqs_cg.json     |\n",
    "| Linearized | /out/symbolic/eqs_pg_lin.json | /out/symbolic/eqs_cg_lin.json |\n",
    "\n",
    "These can be loaded using the `deserialize` methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9c25d3e8-a480-4b98-a835-4f363f08a6d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(root_dir, 'out/symbolic/eqs_pg.json'), 'r') as fread:\n",
    "    eqs_pg = base.CollectionPG.load_json(fread, parser=parse_expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbdf1769-2d03-42b8-a514-99b3afce9c6e",
   "metadata": {},
   "source": [
    "We can take a look at the vorticity equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "24697412-dafe-4a6a-98e6-b90fa4ab3b1d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{s \\frac{\\partial^{3}}{\\partial t\\partial s^{2}} \\Psi{\\left(s,\\phi,t \\right)}}{H{\\left(s \\right)}} - \\frac{s \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial^{2}}{\\partial t\\partial s} \\Psi{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}} + \\frac{\\frac{\\partial^{2}}{\\partial t\\partial s} \\Psi{\\left(s,\\phi,t \\right)}}{H{\\left(s \\right)}} - \\frac{\\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial^{3}}{\\partial t\\partial \\phi^{2}} \\Psi{\\left(s,\\phi,t \\right)}}{2 H^{2}{\\left(s \\right)}} + \\frac{\\frac{\\partial^{3}}{\\partial t\\partial \\phi^{2}} \\Psi{\\left(s,\\phi,t \\right)}}{s H{\\left(s \\right)}} = \\frac{s f_{\\phi}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} - \\frac{s \\frac{\\partial}{\\partial s} \\overline{f_\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\overline{f_\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\widetilde{f_z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\frac{\\partial}{\\partial \\phi} \\overline{f_s}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{2 \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\Psi{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}}$"
      ],
      "text/plain": [
       "Eq(s*Derivative(\\Psi(s, \\phi, t), (s, 2), t)/H(s) - s*Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), s, t)/H(s)**2 + Derivative(\\Psi(s, \\phi, t), s, t)/H(s) - Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), (\\phi, 2), t)/(2*H(s)**2) + Derivative(\\Psi(s, \\phi, t), (\\phi, 2), t)/(s*H(s)), s*f_{\\phi}^e(s, \\phi, t)*Derivative(H(s), s)/H(s) - s*Derivative(\\overline{f_\\phi}(s, \\phi, t), s)/(2*H(s)) - \\overline{f_\\phi}(s, \\phi, t)/(2*H(s)) + Derivative(H(s), s)*Derivative(\\widetilde{f_z}(s, \\phi, t), \\phi)/(2*H(s)) + Derivative(\\overline{f_s}(s, \\phi, t), \\phi)/(2*H(s)) - 2*Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), \\phi)/H(s)**2)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eqs_pg.Psi.doit().expand()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2de8125-db78-460b-b886-c3f28bdb554c",
   "metadata": {},
   "source": [
    "We can similarly load the vorticity equation for transformed variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d0628d12-42ba-4fb4-b784-24dcbe93e24f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{s \\frac{\\partial^{3}}{\\partial t\\partial s^{2}} \\Psi{\\left(s,\\phi,t \\right)}}{H{\\left(s \\right)}} - \\frac{s \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial^{2}}{\\partial t\\partial s} \\Psi{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}} + \\frac{\\frac{\\partial^{2}}{\\partial t\\partial s} \\Psi{\\left(s,\\phi,t \\right)}}{H{\\left(s \\right)}} - \\frac{\\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial^{3}}{\\partial t\\partial \\phi^{2}} \\Psi{\\left(s,\\phi,t \\right)}}{2 H^{2}{\\left(s \\right)}} + \\frac{\\frac{\\partial^{3}}{\\partial t\\partial \\phi^{2}} \\Psi{\\left(s,\\phi,t \\right)}}{s H{\\left(s \\right)}} = \\frac{s f_{\\phi}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} - \\frac{s \\frac{\\partial}{\\partial s} \\overline{f_\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\overline{f_\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\widetilde{f_z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\frac{\\partial}{\\partial \\phi} \\overline{f_s}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{2 \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\Psi{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}}$"
      ],
      "text/plain": [
       "Eq(s*Derivative(\\Psi(s, \\phi, t), (s, 2), t)/H(s) - s*Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), s, t)/H(s)**2 + Derivative(\\Psi(s, \\phi, t), s, t)/H(s) - Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), (\\phi, 2), t)/(2*H(s)**2) + Derivative(\\Psi(s, \\phi, t), (\\phi, 2), t)/(s*H(s)), s*f_{\\phi}^e(s, \\phi, t)*Derivative(H(s), s)/H(s) - s*Derivative(\\overline{f_\\phi}(s, \\phi, t), s)/(2*H(s)) - \\overline{f_\\phi}(s, \\phi, t)/(2*H(s)) + Derivative(H(s), s)*Derivative(\\widetilde{f_z}(s, \\phi, t), \\phi)/(2*H(s)) + Derivative(\\overline{f_s}(s, \\phi, t), \\phi)/(2*H(s)) - 2*Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), \\phi)/H(s)**2)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(os.path.join(root_dir, 'out/symbolic/eqs_cg.json'), 'r') as fread:\n",
    "    eqs_cg = base.CollectionConjugate.load_json(fread, parser=parse_expr)\n",
    "\n",
    "eqs_cg.Psi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a1b253b-40ca-480f-a108-bd593538e5fb",
   "metadata": {},
   "source": [
    "or the magnetic induction equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5704a292-6630-478e-8475-d86a4571cb78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\partial}{\\partial t} \\overline{M_{ss}}{\\left(s,\\phi,t \\right)} = - U_{s}{\\left(s,\\phi,z,t \\right)} \\frac{\\partial}{\\partial s} \\overline{M_{ss}}{\\left(s,\\phi,t \\right)} + 2 \\overline{M_{ss}}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} U_{s}{\\left(s,\\phi,z,t \\right)} + \\frac{U_{s}{\\left(s,\\phi,z,t \\right)} \\overline{M_{ss}}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} - \\frac{U_{\\phi}{\\left(s,\\phi,z,t \\right)} \\frac{\\partial}{\\partial \\phi} \\overline{M_{ss}}{\\left(s,\\phi,t \\right)}}{s} + \\frac{2 \\overline{M_{s\\phi}}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} U_{s}{\\left(s,\\phi,z,t \\right)}}{s}$"
      ],
      "text/plain": [
       "Eq(Derivative(\\overline{M_{ss}}(s, \\phi, t), t), -U_s(s, \\phi, z, t)*Derivative(\\overline{M_{ss}}(s, \\phi, t), s) + 2*\\overline{M_{ss}}(s, \\phi, t)*Derivative(U_s(s, \\phi, z, t), s) + U_s(s, \\phi, z, t)*\\overline{M_{ss}}(s, \\phi, t)*Derivative(H(s), s)/H(s) - U_\\phi(s, \\phi, z, t)*Derivative(\\overline{M_{ss}}(s, \\phi, t), \\phi)/s + 2*\\overline{M_{s\\phi}}(s, \\phi, t)*Derivative(U_s(s, \\phi, z, t), \\phi)/s)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eqs_pg.Mss.doit().expand()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c0c3732-a720-4d13-ad09-579d81ac85b3",
   "metadata": {},
   "source": [
    "Note that the induction equations are already closed.\n",
    "However, vorticity equations contain $\\overline{f_s}$, $\\overline{f_\\phi}$, $\\widetilde{f_z}$ and $f_{e\\phi}$, which are the vertically evenly integrated $s$-component, $\\phi$-component, vertically oddly integrated $z$-component and the equatorial $\\phi$-component of body force $\\mathbf{f}$, respectively.\n",
    "\n",
    "We do not yet know what constitutes these body force.\n",
    "This allows some flexibility regarding which forces are included in the system.\n",
    "Currently, only the forms of the Lorentz force are compiled in the package, but adding viscous diffusion as well as buoyancy force will be straightforward."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f4bac58-02da-4a25-bacf-ba020e200eaa",
   "metadata": {},
   "source": [
    "---\n",
    "## Forcing\n",
    "\n",
    "As mentioned, the forcing in the vorticity equation is only given by placeholder functions. It remains to be defined what forces are involved in the system. This is the task for the `forcing` module.\n",
    "\n",
    "The forcing in `PlesioGeostroPy` are arranged into three layers:\n",
    "1. The placeholder functions for general body forces. We have seen these variables above, they are `core.fs_sym`, `core.fp_sym`, `core.fz_asym` and `core.fe_p`;\n",
    "2. The placeholder functions for specific forces. For Lorentz force, we have `forcing.Ls_sym`, `forcing.Lp_sym`, `forcing.Lz_asym` and `forcing.Le_p`\n",
    "3. Expressions for specific forces in terms of PG/transformed variables. This include\n",
    "    - expressions of the forcing terms in full PG equation, e.g. `forcing.Ls_sym_expr`, `forcing.Lp_sym_expr`, `forcing.Lz_asym_expr` and `forcing.Le_p_expr`\n",
    "    - expressions of the forcing terms in linearized PG equation, e.g. `forcing.Ls_sym_lin`, `forcing.Lp_sym_lin`, `forcing.Lz_asym_lin` and `forcing.Le_p_lin`\n",
    "    - expressions of the forcing terms in transformed equation, e.g. `forcing.Ls_sym_cg`, `forcing.Lp_sym_cg`, `forcing.Lz_asym_cg` and `forcing.Le_p_cg`\n",
    "    - expressions of the forcing terms in linearized transformed equation, e.g. `forcing.Ls_sym_lin_cg`, `forcing.Lp_sym_lin_cg`, `forcing.Lz_asym_lin_cg` and `forcing.Le_p_lin_cg`\n",
    "\n",
    "To assemble the desired forcing in the system, one only needs to replace the placeholder functions in the first layer with a combination of placeholder functions of the second layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "19ea7f93-6c13-4f4b-9c1e-178c5aaccf24",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\mathrm{Le}^{2} s L_{\\phi}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s \\frac{\\partial}{\\partial s} \\overline{L_\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} \\overline{L_\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\widetilde{L_z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{\\partial}{\\partial \\phi} \\overline{L_s}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{2 \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\Psi{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}}$"
      ],
      "text/plain": [
       "\\mathrm{Le}**2*s*L_{\\phi}^e(s, \\phi, t)*Derivative(H(s), s)/H(s) - \\mathrm{Le}**2*s*Derivative(\\overline{L_\\phi}(s, \\phi, t), s)/(2*H(s)) - \\mathrm{Le}**2*\\overline{L_\\phi}(s, \\phi, t)/(2*H(s)) + \\mathrm{Le}**2*Derivative(H(s), s)*Derivative(\\widetilde{L_z}(s, \\phi, t), \\phi)/(2*H(s)) + \\mathrm{Le}**2*Derivative(\\overline{L_s}(s, \\phi, t), \\phi)/(2*H(s)) - 2*Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), \\phi)/H(s)**2"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pg_utils.pg_model import forcing, params\n",
    "\n",
    "cf_lorentz = params.Le**2\n",
    "\n",
    "forcing_map = {\n",
    "    core.fs_sym: cf_lorentz*forcing.Ls_sym,\n",
    "    core.fp_sym: cf_lorentz*forcing.Lp_sym,\n",
    "    core.fz_asym: cf_lorentz*forcing.Lz_asym,\n",
    "    core.fe_p: cf_lorentz*forcing.Le_p,\n",
    "}\n",
    "\n",
    "eqs_pg_lorentz = eqs_pg.subs(forcing_map)\n",
    "eqs_pg_lorentz.Psi.rhs.doit().expand()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2953568a-7b62-4095-a40b-0089bbf52fff",
   "metadata": {},
   "source": [
    "The expression above is the right hand side of the vorticity equation. Pay attention that we used the rotation period as timescale, and ended up with $Le^2$ as a prefactor for the Lorentz force (see [formulation PDF](https://gentlemin.github.io/assets/pdf/Ingredients.pdf) for details). Next, we simply invoke a predefined dictionary to substitute the explicit forcing expressions (3rd layer) for the placeholder functions (2nd layer). These dictionaries are summarized as follows\n",
    "\n",
    "| Type       | PG                     |  Transformed          |\n",
    "|------------|-----------------------------|-----------------------|\n",
    "| Total      | `forcing.force_explicit`    | `forcing.force_explicit_cg`   |\n",
    "| Linearized | `forcing.force_explicit_lin`| `forcing.force_explicit_lin_cg` |\n",
    "\n",
    "As we are currently working with total field PG equation, we just do"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2e0aa8f2-28fe-402a-96e8-7d75420bb442",
   "metadata": {},
   "outputs": [],
   "source": [
    "eqs_pg_explicit = eqs_pg_lorentz.subs(forcing.force_explicit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b2bde9-5e1c-463e-8141-2cea561c6ebd",
   "metadata": {},
   "source": [
    "The resulting formula is massively complicated already"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9f7bff9b-0cfc-4163-b028-6a05c5f86285",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\mathrm{Le}^{2} s^{2} B^{+}_{\\phi}{\\left(s,\\phi,t \\right)} B^{+}_{s}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{2 H^{3}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s^{2} B^{+}_{\\phi}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{+}_{s}{\\left(s,\\phi,t \\right)}}{2 H^{2}{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s^{2} B^{-}_{\\phi}{\\left(s,\\phi,t \\right)} B^{-}_{s}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{2 H^{3}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s^{2} B^{-}_{\\phi}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{-}_{s}{\\left(s,\\phi,t \\right)}}{2 H^{2}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s^{2} B^{+}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{+}_{\\phi}{\\left(s,\\phi,t \\right)}}{2 H^{2}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s^{2} B^{-}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{-}_{\\phi}{\\left(s,\\phi,t \\right)}}{2 H^{2}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s B^{+}_{\\phi}{\\left(s,\\phi,t \\right)} B^{+}_{s}{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s B^{+}_{\\phi}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{+}_{z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s B^{-}_{\\phi}{\\left(s,\\phi,t \\right)} B^{-}_{s}{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B^{-}_{\\phi}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{-}_{z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B^{+}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{+}_{s}{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B^{+}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{+}_{z}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{2 H^{2}{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B^{-}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{-}_{s}{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s B^{-}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{-}_{z}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{2 H^{2}{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s B^{+}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{+}_{\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B^{+}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{+}_{s}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{2 H^{2}{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B^{-}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B^{-}_{\\phi}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s B^{-}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{-}_{s}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{2 H^{2}{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B_{\\phi, z}^e{\\left(s,\\phi,t \\right)} B_{z}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} s B_{s}^e{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial s} B_{\\phi}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} s \\frac{\\partial^{2}}{\\partial s^{2}} \\overline{M_{s\\phi}}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} B^{+}_{\\phi}{\\left(s,\\phi,t \\right)} B^{+}_{z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} B^{-}_{\\phi}{\\left(s,\\phi,t \\right)} B^{-}_{z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} B^{+}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{+}_{z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} B^{-}_{s}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{-}_{z}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} B^{+}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{+}_{s}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} B^{+}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{+}_{z}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} B^{-}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{-}_{s}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} B^{-}_{z}{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B^{-}_{z}{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} B_{\\phi}^e{\\left(s,\\phi,t \\right)} B_{s}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} B_{\\phi}^e{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B_{\\phi}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} - \\frac{2 \\mathrm{Le}^{2} B_{z}^e{\\left(s,\\phi,t \\right)} \\frac{\\partial}{\\partial \\phi} B_{z}^e{\\left(s,\\phi,t \\right)} \\frac{d}{d s} H{\\left(s \\right)}}{H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial^{2}}{\\partial s\\partial \\phi} \\widetilde{M_{sz}}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{3 \\mathrm{Le}^{2} \\frac{\\partial}{\\partial s} \\overline{M_{s\\phi}}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} \\frac{\\partial^{2}}{\\partial s\\partial \\phi} \\overline{M_{\\phi\\phi}}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{\\partial^{2}}{\\partial s\\partial \\phi} \\overline{M_{ss}}{\\left(s,\\phi,t \\right)}}{2 H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial^{2}}{\\partial \\phi^{2}} \\widetilde{M_{\\phi z}}{\\left(s,\\phi,t \\right)}}{2 s H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\widetilde{M_{sz}}{\\left(s,\\phi,t \\right)}}{2 s H{\\left(s \\right)}} - \\frac{\\mathrm{Le}^{2} \\frac{\\partial}{\\partial \\phi} \\overline{M_{\\phi\\phi}}{\\left(s,\\phi,t \\right)}}{2 s H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{\\partial^{2}}{\\partial \\phi^{2}} \\overline{M_{s\\phi}}{\\left(s,\\phi,t \\right)}}{2 s H{\\left(s \\right)}} + \\frac{\\mathrm{Le}^{2} \\frac{\\partial}{\\partial \\phi} \\overline{M_{ss}}{\\left(s,\\phi,t \\right)}}{2 s H{\\left(s \\right)}} - \\frac{2 \\frac{d}{d s} H{\\left(s \\right)} \\frac{\\partial}{\\partial \\phi} \\Psi{\\left(s,\\phi,t \\right)}}{H^{2}{\\left(s \\right)}}$"
      ],
      "text/plain": [
       "\\mathrm{Le}**2*s**2*B_\\phi^+(s, \\phi, t)*B_s^+(s, \\phi, t)*Derivative(H(s), s)/(2*H(s)**3) - \\mathrm{Le}**2*s**2*B_\\phi^+(s, \\phi, t)*Derivative(B_s^+(s, \\phi, t), s)/(2*H(s)**2) + \\mathrm{Le}**2*s**2*B_\\phi^-(s, \\phi, t)*B_s^-(s, \\phi, t)*Derivative(H(s), s)/(2*H(s)**3) - \\mathrm{Le}**2*s**2*B_\\phi^-(s, \\phi, t)*Derivative(B_s^-(s, \\phi, t), s)/(2*H(s)**2) - \\mathrm{Le}**2*s**2*B_s^+(s, \\phi, t)*Derivative(B_\\phi^+(s, \\phi, t), s)/(2*H(s)**2) - \\mathrm{Le}**2*s**2*B_s^-(s, \\phi, t)*Derivative(B_\\phi^-(s, \\phi, t), s)/(2*H(s)**2) - \\mathrm{Le}**2*s*B_\\phi^+(s, \\phi, t)*B_s^+(s, \\phi, t)/H(s)**2 - \\mathrm{Le}**2*s*B_\\phi^+(s, \\phi, t)*Derivative(B_z^+(s, \\phi, t), s)/(2*H(s)) - \\mathrm{Le}**2*s*B_\\phi^-(s, \\phi, t)*B_s^-(s, \\phi, t)/H(s)**2 + \\mathrm{Le}**2*s*B_\\phi^-(s, \\phi, t)*Derivative(B_z^-(s, \\phi, t), s)/(2*H(s)) + \\mathrm{Le}**2*s*B_s^+(s, \\phi, t)*Derivative(B_s^+(s, \\phi, t), \\phi)/H(s)**2 + \\mathrm{Le}**2*s*B_s^+(s, \\phi, t)*Derivative(B_z^+(s, \\phi, t), \\phi)*Derivative(H(s), s)/(2*H(s)**2) + \\mathrm{Le}**2*s*B_s^-(s, \\phi, t)*Derivative(B_s^-(s, \\phi, t), \\phi)/H(s)**2 - \\mathrm{Le}**2*s*B_s^-(s, \\phi, t)*Derivative(B_z^-(s, \\phi, t), \\phi)*Derivative(H(s), s)/(2*H(s)**2) - \\mathrm{Le}**2*s*B_z^+(s, \\phi, t)*Derivative(B_\\phi^+(s, \\phi, t), s)/(2*H(s)) + \\mathrm{Le}**2*s*B_z^+(s, \\phi, t)*Derivative(B_s^+(s, \\phi, t), \\phi)*Derivative(H(s), s)/(2*H(s)**2) + \\mathrm{Le}**2*s*B_z^-(s, \\phi, t)*Derivative(B_\\phi^-(s, \\phi, t), s)/(2*H(s)) - \\mathrm{Le}**2*s*B_z^-(s, \\phi, t)*Derivative(B_s^-(s, \\phi, t), \\phi)*Derivative(H(s), s)/(2*H(s)**2) + \\mathrm{Le}**2*s*B_{\\phi, z}^e(s, \\phi, t)*B_{z}^e(s, \\phi, t)*Derivative(H(s), s)/H(s) + \\mathrm{Le}**2*s*B_{s}^e(s, \\phi, t)*Derivative(B_{\\phi}^e(s, \\phi, t), s)*Derivative(H(s), s)/H(s) - \\mathrm{Le}**2*s*Derivative(\\overline{M_{s\\phi}}(s, \\phi, t), (s, 2))/(2*H(s)) - \\mathrm{Le}**2*B_\\phi^+(s, \\phi, t)*B_z^+(s, \\phi, t)/(2*H(s)) + \\mathrm{Le}**2*B_\\phi^-(s, \\phi, t)*B_z^-(s, \\phi, t)/(2*H(s)) + \\mathrm{Le}**2*B_s^+(s, \\phi, t)*Derivative(B_z^+(s, \\phi, t), \\phi)/(2*H(s)) - \\mathrm{Le}**2*B_s^-(s, \\phi, t)*Derivative(B_z^-(s, \\phi, t), \\phi)/(2*H(s)) + \\mathrm{Le}**2*B_z^+(s, \\phi, t)*Derivative(B_s^+(s, \\phi, t), \\phi)/(2*H(s)) + \\mathrm{Le}**2*B_z^+(s, \\phi, t)*Derivative(B_z^+(s, \\phi, t), \\phi)*Derivative(H(s), s)/H(s) - \\mathrm{Le}**2*B_z^-(s, \\phi, t)*Derivative(B_s^-(s, \\phi, t), \\phi)/(2*H(s)) + \\mathrm{Le}**2*B_z^-(s, \\phi, t)*Derivative(B_z^-(s, \\phi, t), \\phi)*Derivative(H(s), s)/H(s) + \\mathrm{Le}**2*B_{\\phi}^e(s, \\phi, t)*B_{s}^e(s, \\phi, t)*Derivative(H(s), s)/H(s) + \\mathrm{Le}**2*B_{\\phi}^e(s, \\phi, t)*Derivative(B_{\\phi}^e(s, \\phi, t), \\phi)*Derivative(H(s), s)/H(s) - 2*\\mathrm{Le}**2*B_{z}^e(s, \\phi, t)*Derivative(B_{z}^e(s, \\phi, t), \\phi)*Derivative(H(s), s)/H(s) + \\mathrm{Le}**2*Derivative(H(s), s)*Derivative(\\widetilde{M_{sz}}(s, \\phi, t), \\phi, s)/(2*H(s)) - 3*\\mathrm{Le}**2*Derivative(\\overline{M_{s\\phi}}(s, \\phi, t), s)/(2*H(s)) - \\mathrm{Le}**2*Derivative(\\overline{M_{\\phi\\phi}}(s, \\phi, t), \\phi, s)/(2*H(s)) + \\mathrm{Le}**2*Derivative(\\overline{M_{ss}}(s, \\phi, t), \\phi, s)/(2*H(s)) + \\mathrm{Le}**2*Derivative(H(s), s)*Derivative(\\widetilde{M_{\\phi z}}(s, \\phi, t), (\\phi, 2))/(2*s*H(s)) + \\mathrm{Le}**2*Derivative(H(s), s)*Derivative(\\widetilde{M_{sz}}(s, \\phi, t), \\phi)/(2*s*H(s)) - \\mathrm{Le}**2*Derivative(\\overline{M_{\\phi\\phi}}(s, \\phi, t), \\phi)/(2*s*H(s)) + \\mathrm{Le}**2*Derivative(\\overline{M_{s\\phi}}(s, \\phi, t), (\\phi, 2))/(2*s*H(s)) + \\mathrm{Le}**2*Derivative(\\overline{M_{ss}}(s, \\phi, t), \\phi)/(2*s*H(s)) - 2*Derivative(H(s), s)*Derivative(\\Psi(s, \\phi, t), \\phi)/H(s)**2"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eqs_pg_explicit.Psi.rhs.doit().expand()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a99b68e3-dd17-43f1-ac8e-9cfb8ffb7103",
   "metadata": {},
   "source": [
    "---\n",
    "## Carry on!\n",
    "\n",
    "We have now finished with the fundamentals of the symbolic part in **PlesioGeostroPy**.\n",
    "There are many modules we haven't talked about yet, such as the entire spectral expansion module, and the entire numerics module.\n",
    "However, those will require a concrete problem to demonstrate how they actually work.\n",
    "Therefore, we are now ready to (and have to) dive into actually solving problems!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "020b4db0-e625-44be-9d2e-792c204d5b23",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3 (geomag)",
   "language": "python",
   "name": "geomag"
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}