docs/notebooks/Torus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Helmholtz equation on the Torus\n",
    "\n",
    "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 (without boundary conditions) on the surface of a torus, using curvilinear coordinates. The surface of the torus is parametrized by    \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "x(\\theta, \\phi) &= (R + r \\cos \\theta) \\cos \\phi \\\\\n",
    "y(\\theta, \\phi) &= (R + r \\cos \\theta) \\sin \\phi \\\\\n",
    "z(\\theta, \\phi) &= r \\sin \\theta\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where $\\theta, \\phi$ are angles which make a full circle, so that their values start and end at the same point, $R$ is the distance from the center of the tube to the center of the torus,\n",
    "$r$ is the radius of the tube. Note that $\\theta$ is the angle in the small circle (around its center), whereas $\\phi$ is the angle of the large circle, around origo.\n",
    "\n",
    "We start the implementation by importing necessary functionality from shenfun and sympy and then defining the coordinates of the surface of the torus. Note that `rv` represents the position vector spanning the surface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from shenfun import *\n",
    "from shenfun.la import SolverGeneric1ND\n",
    "import sympy as sp\n",
    "from IPython.display import Math\n",
    "\n",
    "N = 96\n",
    "R = 3\n",
    "r = 1\n",
    "theta, phi = psi = sp.symbols('x,y', real=True, positive=True)\n",
    "rv = ((R + r*sp.cos(theta))*sp.cos(phi), (R + r*sp.cos(theta))*sp.sin(phi), r*sp.sin(theta)) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now create necessary bases and function spaces. Due to the geometry of the problem, the solution will be periodic in both $\\theta$ and $\\phi$ directions, and we choose Fourier basis functions. The basis for the $\\phi$-direction can be either real to complex or complex to complex, depending on the type of the solution. Here we assume a real solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B1 = FunctionSpace(N, 'F', dtype='D', domain=(0, 2*np.pi))\n",
    "B2 = FunctionSpace(N, 'F', dtype='d', domain=(0, 2*np.pi))\n",
    "T = TensorProductSpace(comm, (B1, B2), coordinates=(psi, rv, sp.Q.positive(r*sp.cos(theta)+R)))\n",
    "V = VectorSpace(T)\n",
    "u = TrialFunction(T)\n",
    "v = TestFunction(T)\n",
    "T.coors.sg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the assumption `sp.Q.positive(r*sp.cos(theta)+R))`, which is there to help sympy when computing basis vectors and scaling factors. It is not completely necessary, but try to omit it and look at what happens to the expanded Helmholtz equation below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 0\n",
    "du = div(grad(u))+alpha*u\n",
    "g = sp.Symbol('g', real=True, positive=True) # The Jacobian**2 (T.coors.sg**2)\n",
    "#replace = [(3*sp.cos(theta)+5, g/2), (sp.cos(theta)**2+6*sp.cos(theta)+10, sp.cos(theta)**2+g), (3*sp.cos(theta)+10, g/2+5)] # to simplify the look\n",
    "replace = [(sp.cos(theta)+3, g)]\n",
    "Math(du.tolatex(symbol_names={r: 'r', theta: '\\\\theta', phi: '\\\\phi'}, replace=replace))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now create a manufactured solution that satisfies periodicity and compute the right hand side $f$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ue = sp.sin(sp.cos(theta*4))*sp.cos(4*phi)\n",
    "f = (div(grad(u))+alpha*u).tosympy(basis=ue, psi=psi)\n",
    "fj = Array(T, buffer=f*T.sg)\n",
    "f_hat = Function(T)\n",
    "f_hat = inner(v, fj, output_array=f_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assemble coefficient matrix and solve problem. Note that the tensorproduct matrices along axis 0 can be non-diagonal due to the measure $\\cos \\theta + 3$. The matrices along the second axis will all be diagonal, so we can choose to use `SolverGeneric1ND`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mats = inner(v*T.sg, (div(grad(u))+alpha*u))\n",
    "#mats = inner(grad(v*T.sg), -grad(u)) # + inner(v, alpha*u)\n",
    "u_hat = Function(T)\n",
    "sol = SolverGeneric1ND(mats)\n",
    "u_hat = sol(f_hat, u_hat, constraints=((0, 0),))\n",
    "uj = u_hat.backward()\n",
    "uq = Array(T, buffer=ue)\n",
    "print('Error =', np.sqrt(inner(1, (uj-uq)**2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mats = inner(v*T.sg, u)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, just plot the solution using plotly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "surf3D(u_hat, wrapaxes=[0, 1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect sparsity pattern at a given Fourier wavenumber "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib notebook\n",
    "plt.figure(figsize=(6,4))\n",
    "plt.spy(sol.solvers1D[6].mat, markersize=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect the vector Laplacian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uv = TrialFunction(V)\n",
    "vv = TestFunction(V)\n",
    "du = div(grad(uv))\n",
    "Math(du.tolatex(symbol_names={r: 'r', theta: '\\\\theta', phi: '\\\\phi'}, replace=replace))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Math(T.coors.latex_basis_vectors(symbol_names={r: 'r', theta: '\\\\theta', phi: '\\\\phi'}, replace=replace))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "M = inner(vv*T.sg, du)\n",
    "B = BlockMatrix(M)\n",
    "plt.figure(figsize=(6,4))\n",
    "plt.spy(B.diags(it=(4,)), markersize=0.2)"
   ]
  }
 ],
 "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
}