docs/notebooks/sphere-helmholtz.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spherical coordinates in shenfun\n",
    "\n",
    "The Helmholtz equation is given as\n",
    "\n",
    "$$\n",
    "-\\nabla^2 u + \\alpha u = f.\n",
    "$$\n",
    "\n",
    "In this notebook we will solve this equation on a unitsphere, using spherical coordinates. To verify the implementation we use a spherical harmonics function as manufactured solution. \n",
    "\n",
    "We start the implementation by importing necessary functionality from shenfun and sympy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from shenfun import *\n",
    "from shenfun.la import SolverGeneric1ND\n",
    "import sympy as sp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define spherical coordinates $(r, \\theta, \\phi)$\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "x &= r \\sin \\theta \\cos \\phi \\\\\n",
    "y &= r \\sin \\theta \\sin \\phi \\\\\n",
    "z &= r \\cos \\theta\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "using sympy. The radius `r` will be constant `r=1`. We create the three-dimensional position vector `rv` as a function of the two new coordinates $(\\theta, \\phi)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = 1\n",
    "theta, phi = psi = sp.symbols('x,y', real=True, positive=True)\n",
    "rv = (r*sp.sin(theta)*sp.cos(phi), r*sp.sin(theta)*sp.sin(phi), r*sp.cos(theta))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define bases with the domains $\\theta \\in [0, \\pi]$ and $\\phi \\in [0, 2\\pi]$. Also define a tensorproductspace, test- and trialfunction. Note that the new coordinates and the position vector are fed to the `TensorProductSpace` and not the individual spaces:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N, M = 64, 64\n",
    "L0 = FunctionSpace(N, 'C', domain=(0, np.pi))\n",
    "F1 = FunctionSpace(M, 'F', dtype='d')\n",
    "T = TensorProductSpace(comm, (L0, F1), coordinates=(psi, rv, sp.Q.positive(sp.sin(theta))))\n",
    "v = TestFunction(T)\n",
    "u = TrialFunction(T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T.coors.sg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use one spherical harmonic function as a manufactured solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sph = sp.functions.special.spherical_harmonics.Ynm\n",
    "#ue = sph(6, 3, theta, phi)\n",
    "ue = sp.cos(8*(sp.sin(theta)*sp.cos(phi) + sp.sin(theta)*sp.sin(phi) + sp.cos(theta)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the right hand side on the quadrature mesh and take the scalar product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1000\n",
    "h = T.sg**2\n",
    "g = (-div(grad(u))+alpha*u).tosympy(basis=ue, psi=psi)\n",
    "gj = Array(T, buffer=g*h)\n",
    "g_hat = Function(T)\n",
    "g_hat = inner(v, gj, output_array=g_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we can use the `shenfun` operators `div` and `grad` on a trialfunction `u`, and then switch the trialfunction for a sympy function `ue`. The operators will then make use of sympy's [derivative method](https://docs.sympy.org/latest/tutorial/calculus.html#derivatives) on the function `ue`. Here `(-div(grad(u))+alpha*u)` corresponds to the equation we are trying to solve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Math\n",
    "Math((-div(grad(u))*h+alpha*u).tolatex(funcname='u', symbol_names={theta: '\\\\theta', phi: '\\\\phi'}))\n",
    "#Math((grad(u)).tolatex(funcname='u', symbol_names={theta: '\\\\theta', phi: '\\\\phi'}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluated with `u=ue` and you get the exact right hand side `f`.\n",
    "\n",
    "Tensor product matrices that make up the Helmholtz equation are then assembled as "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mats = inner(v, (-div(grad(u))+alpha*u)*h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mats[3].mats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the linear system of equations can be solved using the generic `SolverGeneric1ND`, that can be used for any problem that only has non-periodic boundary conditions in one dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_hat = Function(T)\n",
    "Sol1 = SolverGeneric1ND(mats)\n",
    "u_hat = Sol1(g_hat, u_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transform back to real space and compute the error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uj = u_hat.backward()\n",
    "uq = Array(T, buffer=ue)\n",
    "print('Error =', np.sqrt(inner(1, (uj-uq)**2)), np.linalg.norm(uj-uq))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linalg.norm(u_hat - u_hat.backward().forward())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.spy(Sol1.solvers1D[1].mat, markersize=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#raise RuntimeError\n",
    "surf3D(u_hat, wrapaxes=[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Postprocessing\n",
    "Since we used quite few quadrature points in solving this problem, we refine the solution for a nicer plot. Note that `refine` simply pads Functions with zeros, which gives exactly the same accuracy, but more quadrature points in real space. `u_hat` has `NxM` quadrature points, here we refine using 3 times as many points along both dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_hat2 = u_hat.refine([N*3, M*3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The periodic solution does not contain the periodic points twice, i.e., the computational mesh contains $0$, but not $2\\pi$. It looks better if we wrap the periodic dimension all around to $2\\pi$, and this is achieved with "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "surf3D(u_hat2, wrapaxes=[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Biharmonic equation\n",
    "\n",
    "A biharmonic equation is given as\n",
    "\n",
    "$$\n",
    "\\nabla^4 u + \\alpha u = f.\n",
    "$$\n",
    "\n",
    "This equation is extremely messy in spherical coordinates. I cannot even find it posted anywhere. Nevertheless, we can solve it trivially with shenfun, and we can also see what it looks like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Math((div(grad(div(grad(u))))+alpha*u).tolatex(funcname='u', symbol_names={theta: '\\\\theta', phi: '\\\\phi'}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember that this equation uses constant radius `r=1`. We now solve the equation using the same manufactured solution as for the Helmholtz equation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = (div(grad(div(grad(u))))+alpha*u).tosympy(basis=ue, psi=psi)\n",
    "gj = Array(T, buffer=g)\n",
    "\n",
    "# Take scalar product\n",
    "g_hat = Function(T)\n",
    "g_hat = inner(v, gj, output_array=g_hat)\n",
    "\n",
    "mats = inner(v, div(grad(div(grad(u)))) + alpha*u)\n",
    "\n",
    "# Solve\n",
    "u_hat = Function(T)\n",
    "Sol1 = SolverGeneric1ND(mats)\n",
    "u_hat = Sol1(g_hat, u_hat)\n",
    "\n",
    "# Transform back to real space.\n",
    "uj = u_hat.backward()\n",
    "uq = Array(T, buffer=ue)\n",
    "print('Error =', np.sqrt(dx((uj-uq)**2)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Want to see what the regular 3-dimensional biharmonic equation looks like in spherical coordinates? This is extremely tedious to derive by hand, but in shenfun you can get there with the following few lines of code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r, theta, phi = psi = sp.symbols('x,y,z', real=True, positive=True)\n",
    "rv = (r*sp.sin(theta)*sp.cos(phi), r*sp.sin(theta)*sp.sin(phi), r*sp.cos(theta))\n",
    "L0 = FunctionSpace(20, 'L', domain=(0, 1))\n",
    "F1 = FunctionSpace(20, 'L', domain=(0, np.pi))\n",
    "F2 = FunctionSpace(20, 'F', dtype='D')\n",
    "T = TensorProductSpace(comm, (L0, F1, F2), coordinates=(psi, rv, sp.Q.positive(sp.sin(theta))))\n",
    "p = TrialFunction(T)\n",
    "Math((div(grad(div(grad(p))))).tolatex(funcname='u', symbol_names={r: 'r', theta: '\\\\theta', phi: '\\\\phi'}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q = TestFunction(T)\n",
    "A = inner(div(grad(q)), div(grad(p)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I don't know if this is actually correct, because I haven't derived it by hand and I haven't seen it printed anywhere, but at least I know the Cartesian equation is correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L0 = FunctionSpace(8, 'C', domain=(0, np.pi))\n",
    "F1 = FunctionSpace(8, 'F', dtype='D')\n",
    "F2 = FunctionSpace(8, 'F', dtype='D')\n",
    "T = TensorProductSpace(comm, (L0, F1, F2))\n",
    "p = TrialFunction(T)\n",
    "Math((div(grad(div(grad(p))))).tolatex(funcname='u'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "shenfun",
   "language": "python",
   "name": "shenfun"
  },
  "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.2"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}