Shenfun-introduction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "<center>\n",
    "\n",
    "# Shenfun - High-Performance Computing platform for the Spectral Galerkin method\n",
    "\n",
    "<div><img src=\"https://rawcdn.githack.com/spectralDNS/spectralutilities/f3419a3e6c40dad55be5dcca51f6e0e21713dd90/figures/Chebyshev_Polynomials_of_the_First_Kind.svg\" width=\"300\"></div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Shenfun - facts\n",
    "\n",
    "1. Shenfun is named in honour of <strong>Professor Jie Shen</strong> for his seminal work on the spectral Galerkin method:-) \n",
    "2. Shenfun is a high performance computing platform for solving partial differential equations (PDEs) with the spectral Galerkin method.\n",
    "3. Shenfun has been run with 65,000 processors on a Cray XC40.\n",
    "4. Shenfun is a high-level <strong>Python</strong> package originally developed for pseudo-spectral turbulence simulations.\n",
    "\n",
    "\n",
    "<img src=\"https://rawcdn.githack.com/spectralDNS/spectralutilities/473129742f0b5f8d57e8c647809272c0ced99a45/movies/RB_200k_small.png\" style=\"float:left\" width=\"300\"> <img src=\"https://rawcdn.githack.com/spectralDNS/spectralutilities/473129742f0b5f8d57e8c647809272c0ced99a45/movies/isotropic_cropped.gif\" style=\"float:right\" width=\"200\"> \n",
    "<p style=\"clear: both;\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# The Spectral Galerkin method\n",
    "\n",
    "## Advantages\n",
    "\n",
    "- Accuracy (spectral)\n",
    "- Sparsity - very low memory use for coefficient matrices \n",
    "- Efficiency - Very fast linear algebra solvers due to sparsity\n",
    "- Robustness (round-off errors) - condition numbers much lower than for collocation methods\n",
    "\n",
    "## Disadvantages\n",
    "\n",
    "- No complex grids. Structured tensor product grids only\n",
    "- Nonlinear terms must be implemented explicitly (pseudospectral)\n",
    "- Generally considered more difficult to implement since we solve equations in spectral space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# The Spectral Galerkin method (in a nutshell)\n",
    "\n",
    "approximates solutions $u(x)$ using global <strong>trial</strong> functions $\\phi_k(x)$ and unknown expansion coefficients $\\hat{u}_k$\n",
    "\n",
    "$$\n",
    "u(x) = \\sum_{k=0}^{N-1}\\hat{u}_k \\phi_k(x)\n",
    "$$\n",
    "\n",
    "Multidimensional solutions are formed from outer (tensor) products of 1D bases\n",
    "\n",
    "$$\n",
    "u(x, y) = \\sum_{k=0}^{N_0-1}\\sum_{l=0}^{N_1-1}\\hat{u}_{kl} \\phi_{kl}(x, y)\\quad \\text{ or }\\quad\n",
    "u(x, y, z) = \\sum_{k=0}^{N_0-1}\\sum_{l=0}^{N_1-1} \\sum_{m=0}^{N_2-1}\\hat{u}_{klm} \\phi_{klm}(x, y, z)\n",
    "$$\n",
    "\n",
    "where, for example\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\phi_{kl}(x, y) &= T_k(x) L_l(y)\\\\\n",
    "\\phi_{klm}(x, y, z) &= T_k(x) L_l(y) \\exp(\\text{i}mz)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "$T_k$ and $L_k$ are Chebyshev and Legendre polynomials."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# The Spectral Galerkin method\n",
    "\n",
    "solves PDEs, like Poisson's equation\n",
    "\n",
    "\\begin{align}\n",
    "\\nabla^2 u(x) &= f(x), \\quad x \\in [-1, 1] \\\\\n",
    "u(\\pm 1) &= 0\n",
    "\\end{align}\n",
    "\n",
    "using variational forms by the <strong>method of weighted residuals</strong>. I.e., multiply PDE by a test function $v$ and integrate over the domain. For Poisson this leads to the problem:\n",
    "\n",
    "Find $u \\in H^1_0$ such that \n",
    "\n",
    "$$(\\nabla u, \\nabla v)_w^N = -(f, v)_w^N \\quad \\forall v \\in H^1_0$$\n",
    "\n",
    "Here $(u, v)_w^{N}$ is a weighted inner product and $v(=\\phi_j)$ is a <strong>test</strong> function. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Weighted inner products\n",
    "\n",
    "The weighted inner product is defined as\n",
    "\n",
    "$$\n",
    " (u, v)_w = \\int_{\\Omega} u \\overline{v} w \\, d\\Omega,\n",
    "$$\n",
    "\n",
    "where $w(\\mathbf{x})$ is a weight associated with the chosen basis (different bases have different weights). The overline represents a complex conjugate (for Fourier).\n",
    "\n",
    "$\\Omega$ is a Cartesian product domain spanned by the chosen 1D bases.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# In Shenfun quadrature is used for the integrals\n",
    "\n",
    "1D with Chebyshev basis:\n",
    "\n",
    "$$\n",
    "(u, v)_w ^N = \\sum_{i=0}^{N-1} u(x_i) v(x_i) \\omega_i \\approx \\int_{-1}^1 \\frac{u v}{\\sqrt{1-x^2}}  \\, {dx},\n",
    "$$\n",
    "\n",
    "where $\\{\\omega_i\\}_{i=0}^{N-1}$ are the quadrature weights associated with the chosen basis and quadrature rule. The associated quadrature points are denoted as $\\{x_i\\}_{i=0}^{N-1}$. \n",
    "\n",
    "2D with mixed Chebyshev-Fourier:\n",
    "\n",
    "$$\n",
    "(u, v)_w^N = \\int_{-1}^1\\int_{0}^{2\\pi} \\frac{u \\overline{v}}{2\\pi\\sqrt{1-x^2}} \\, {dxdy} \\approx \\sum_{i=0}^{N_0-1}\\sum_{j=0}^{N_1-1} u(x_i, y_j) \\overline{v}(x_i, y_j) \\omega^{(x)}_i \\omega_j^{(y)} ,\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Spectral Galerkin solution procedure\n",
    "\n",
    "1. Choose function space(s) satisfying the correct boundary conditions\n",
    "2. Transform PDEs to variational forms using inner products\n",
    "3. Assemble variational forms and solve resulting linear algebra systems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Orthogonal bases\n",
    "\n",
    "<p style=\"margin-bottom:1cm;\">\n",
    "\n",
    "| Family    | Basis                                    | Weight   | Domain    |\n",
    "|  :---:    |         :---:                            | :---:    |  :---:   |\n",
    "| Chebyshev | $$\\{T_k\\}_{k=0}^{N-1}$$                  | $$1/\\sqrt{1-x^2}$$ | $$[-1, 1]$$ |\n",
    "| Legendre  | $$\\{L_k\\}_{k=0}^{N-1}$$                  |  1       |$$[-1, 1]$$ |\n",
    "| Jacobi    | $$\\{J_k^{(\\alpha,\\beta)}\\}_{k=0}^{N-1}$$ | $$(1-x)^{\\alpha}(1+x)^{\\beta}$$ |  $$[-1, 1]$$ |\n",
    "| Fourier   | $$\\{\\exp(\\text{i}kx)\\}_{k=-N/2}^{N/2-1}$$| $$1/(2\\pi)$$  |$$[0, 2\\pi]$$ |\n",
    "| Hermite   | $$\\{H_k\\}_{k=0}^{N-1}$$                  | $$e^{-x^2/2}$$        | $$[-\\infty, \\infty]$$|\n",
    "| Laguerre  | $$\\{La_k\\}_{k=0}^{N-1}$$                 | $e^{-x/2}$        | $$[0, \\infty]$$ |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from shenfun import *\n",
    "N = 8\n",
    "C = FunctionSpace(N, 'Chebyshev', quad='GC', domain=[-2, 2])\n",
    "L = FunctionSpace(N, 'Legendre')\n",
    "x, w = C.points_and_weights()\n",
    "print(C.points_and_weights())\n",
    "C.sympy_basis(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Jie Shen's bases with Dirichlet bcs\n",
    "<p style=\"margin-bottom:1cm;\">\n",
    "\n",
    "\n",
    "| family    | Basis                 | Boundary condition |\n",
    "|-----------|-----------------------|----------|\n",
    "| Chebyshev | $$\\{T_k-T_{k+2}\\}_{k=0}^{N-3}$$ | $$u(\\pm 1) = 0$$ |\n",
    "| Legendre  | $$\\{L_k-L_{k+2}\\}_{k=0}^{N-3}$$ | $$u(\\pm 1) = 0$$ |\n",
    "| Hermite   | $$\\exp(-x^2)\\{H_k\\}_{k=0}^{N-1}$$ | $$u(\\pm \\infty) = 0$$ |\n",
    "| Laguerre  | $$\\exp(-x/2)\\{La_k-La_{k+1}\\}_{k=0}^{N-2}$$| $$u(0) = u(\\infty) = 0$$ |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Composite bases are implemented using a stencil matrix, such that\n",
    "\n",
    "$$\n",
    "\\phi_i = \\sum_{j=0}^{N-1} K_{ij}Q_j\n",
    "$$\n",
    "\n",
    "where $K$ is the stencil matrix and $\\{Q_j\\}_{j=0}^{N-1}$ are the orthogonal polynomials, like $T_j$ or $L_j$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "N = 8\n",
    "C0 = FunctionSpace(N, 'Chebyshev', bc=(0, 0))\n",
    "L0 = FunctionSpace(N, 'Legendre', bc=(0, 0))\n",
    "H0 = FunctionSpace(N, 'Hermite')\n",
    "La = FunctionSpace(N, 'Laguerre', bc=(0, 0))\n",
    "L0.stencil_matrix().diags().toarray()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Shen's bases with Neumann $u'(\\pm 1) = 0$\n",
    "\n",
    "<p style=\"margin-bottom:1cm;\">\n",
    "\n",
    "| family    | Basis                 |\n",
    "|-----------|-----------------------|\n",
    "| Chebyshev | $$\\left\\{T_k-\\frac{k^2}{(k+2)^2}T_{k+2}\\right\\}_{k=0}^{N-3}$$ | \n",
    "| Legendre  | $$\\left\\{L_k-\\frac{k(k+1)}{(k+2)(k+3)}L_{k+2}\\right\\}_{k=0}^{N-3}$$ |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "CN = FunctionSpace(N, 'Chebyshev', bc={'left': ('N', 0), 'right': ('N', 0)})\n",
    "LN = FunctionSpace(N, 'Legendre', bc={'left': ('N', 0), 'right': ('N', 0)})\n",
    "dict(CN.stencil_matrix())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "# Shen's biharmonic bases $u(\\pm 1) = u'(\\pm 1) = 0$\n",
    "<p style=\"margin-bottom:1cm;\">\n",
    "\n",
    "| family    | Basis                 |\n",
    "|-----------| :-----------------:   |\n",
    "| Chebyshev | $$\\left\\{T_k-\\frac{2(k+2)}{k+3}T_{k+2}+\\frac{k+1}{k+3} T_{k+4}\\right\\}_{k=0}^{N-5}$$ | \n",
    "| Legendre  | $$\\left\\{L_k-\\frac{2(2k+5)}{(2k+7)}L_{k+2}+\\frac{2k+3}{2k+7}L_{k+4}\\right\\}_{k=0}^{N-5}$$ |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "CB = FunctionSpace(N, 'Chebyshev', bc=(0, 0, 0, 0))\n",
    "LB = FunctionSpace(N, 'Legendre', bc=(0, 0, 0, 0))\n",
    "dict(CB.stencil_matrix())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Multidimensional tensor product spaces\n",
    "<p style=\"margin-bottom:0.5cm;\">\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "L_0 &= \\{L_k(x)-L_{k+2}(x)\\}_{k=0}^{N-3} \\\\\n",
    "C_0 &= \\{T_k(x)-T_{k+2}(x)\\}_{k=0}^{N-3} \\\\\n",
    "L_1 &= \\{L_l(y)\\}_{l=0}^{N-1} \\\\\n",
    "LL(x, y) &= L_0(x) \\otimes L_1(y) \\\\\n",
    "CL(x, y) &= C_0(x) \\otimes L_1(y)\n",
    "\\end{align}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "L0 = FunctionSpace(N, 'Legendre', bc=(0, 0))\n",
    "C0 = FunctionSpace(N, 'Chebyshev', bc=(0, 0))\n",
    "L1 = FunctionSpace(N, 'Legendre')\n",
    "LL = TensorProductSpace(comm, (L0, L1)) # comm is MPI.COMM_WORLD\n",
    "CL = TensorProductSpace(comm, (C0, L1))\n",
    "V = VectorSpace(LL) # For vector valued functions\n",
    "f = Array(LL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Operators in shenfun\n",
    "\n",
    "Act on instances of a `TestFunction`, `TrialFunction` or `Function`\n",
    "\n",
    "- div\n",
    "- grad\n",
    "- curl\n",
    "- Dx   (for a partial derivative)\n",
    "\n",
    "# Assembly\n",
    "- project\n",
    "- inner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "L0 = FunctionSpace(N, 'Legendre', bc=(0, 0))\n",
    "L1 = FunctionSpace(N, 'Legendre')\n",
    "u = TrialFunction(L0)\n",
    "v = TestFunction(L0)\n",
    "uh = Function(L0)\n",
    "g = Array(L0)\n",
    "du = grad(u)  # vector valued expression\n",
    "h = div(du)   # scalar valued expression\n",
    "\n",
    "A = inner(Dx(u, 0, 2), v)\n",
    "dict(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# The shenfun `Function` represents the solution\n",
    "\n",
    "`uh = Function(L0)`\n",
    "\n",
    "$$\n",
    "u_h(x) = \\sum_{k=0}^{N-1} \\hat{u}_k \\phi_{k}(x)\n",
    "$$\n",
    "\n",
    "The function evaluated for all quadrature points, $\\{x_j\\}_{j=0}^{N-1}$, is an `Array`\n",
    "\n",
    "`uj = Array(L0)`\n",
    "\n",
    "There is a (fast) `backward` transform for moving from `Function` to `Array`, and a `forward` transform to go the other way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "uj = Array(L0)\n",
    "uj = uh.backward(uj)\n",
    "uh = uj.forward(uh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Projections\n",
    "\n",
    "Project $g(\\mathbf{x})$ to $V$:\n",
    "\n",
    "Find $u$ in $V$ such that:\n",
    "\n",
    "$$(u, v)_w = (I^Ng, v)_w \\quad \\text{for} \\, v \\in V $$\n",
    "\n",
    "where $I^Ng$ is $\\{g(x_j)\\}_{j=0}^{N-1}$, i.e., $g(x)$ evaluated on the quadrature mesh.\n",
    "\n",
    "Works if $g(x)$ is \n",
    "\n",
    "- an `Array`, which is exactly a `Function` evaluated on the mesh\n",
    "- an expression involving a `Function`, like `div(grad(uh))`\n",
    "- a `sympy` expression, like `sin(x)`\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "dudx = project(Dx(uh, 0, 1), L1) # Compute du/dx\n",
    "wh = project(uj, L1)\n",
    "import sympy as sp\n",
    "x, y = sp.symbols('x,y')\n",
    "ws = project(sp.sin(4*x), L1)\n",
    "ws.eval(np.array([0.51]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Implementation matches mathematics\n",
    "<p style=\"margin-bottom:1cm;\">\n",
    "\n",
    "$$\n",
    "A = (\\nabla u, \\nabla v)_w^N\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "A = inner(grad(u), grad(v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "dict(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(A.diags().todense())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "A diagonal stiffness matrix!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Complete Poisson solver with error verification in 1D\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Solve Poisson's equation\n",
    "import matplotlib.pyplot as plt\n",
    "from sympy import symbols, sin, cos, lambdify\n",
    "from shenfun import * \n",
    "\n",
    "# Use sympy to compute manufactured solution\n",
    "x, y = symbols(\"x,y\")\n",
    "ue = sin(6*sp.pi*x)*(1-x**2) # `ue` is the manufactured solution\n",
    "fe = ue.diff(x, 2) # `fe` is Poisson's right hand side for `ue`\n",
    "\n",
    "SD = FunctionSpace(50, 'L', bc=(0, 0))\n",
    "u = TrialFunction(SD)\n",
    "v = TestFunction(SD)\n",
    "\n",
    "b = inner(v, Array(SD, buffer=fe)) # Array is initialized with `fe`\n",
    "A = inner(v, div(grad(u)))\n",
    "\n",
    "uh = Function(SD)\n",
    "uh = A.solve(b, uh)\n",
    "ue = Array(SD, buffer=ue)\n",
    "print(\"L2-error = \", np.sqrt(inner(1, (uh.backward()-ue)**2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "plt.plot(SD.mesh(), uh.backward(), SD.mesh(), ue)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# 2D - still closely matching mathematics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "L0 = FunctionSpace(N, 'Legendre', bc=(0, 0))\n",
    "F1 = FunctionSpace(N, 'Fourier', dtype='d')\n",
    "TP = TensorProductSpace(comm, (L0, F1))\n",
    "u = TrialFunction(TP)\n",
    "v = TestFunction(TP)\n",
    "A = inner(grad(u), grad(v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "# ? \n",
    "\n",
    "A is a list of two TPMatrix objects???\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# `TPMatrix` is a Tensor Product matrix\n",
    "\n",
    "A `TPMatrix` is the outer product of smaller matrices (2 in 2D, 3 in 3D etc).  \n",
    "\n",
    "Consider the inner product:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "(\\nabla u, \\nabla v)_w &= \\frac{1}{2\\pi}\\int_{-1}^{1}\\int_{0}^{2\\pi} \\left(\\frac{\\partial u}{\\partial x}, \\frac{\\partial u}{\\partial y}\\right) \\cdot \\left(\\frac{\\partial \\overline{v}}{\\partial x}, \\frac{\\partial \\overline{v}}{\\partial y}\\right) {dxdy} \\\\\n",
    "(\\nabla u, \\nabla v)_w &= \\frac{1}{2\\pi}\\int_{-1}^1 \\int_{0}^{2\\pi} \\frac{\\partial u}{\\partial x}\\frac{\\partial \\overline{v}}{\\partial x} {dxdy} + \\int_{-1}^1 \\int_{0}^{2\\pi} \\frac{\\partial u}{\\partial y}\\frac{\\partial \\overline{v}}{\\partial y} {dxdy}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "which, like `A`, is a sum of two terms. These two terms are the two `TPMatrix`es returned by `inner` above.\n",
    "\n",
    "Now each one of these two terms can be written as the outer product of two smaller matrices. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Consider the first, inserting for test and trial functions\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "v &= \\phi_{kl} = (L_k(x)-L_{k+2}(x))\\exp(\\text{i}ly) \\\\\n",
    "u &= \\phi_{mn}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "The first term becomes\n",
    "\n",
    "$$\n",
    "\\small\n",
    "\\begin{align}\n",
    "\\int_{-1}^1 \\int_{0}^{2\\pi} \\frac{\\partial u}{\\partial x}\\frac{\\partial \\overline{v}}{\\partial x} \\frac{dxdy}{2\\pi} &= \\underbrace{\\int_{-1}^1 \\frac{\\partial (L_m-L_{m+2})}{\\partial x}\\frac{\\partial (L_k-L_{k+2})}{\\partial x} {dx}}_{a_{km}} \\underbrace{\\int_{0}^{2\\pi}  \\exp(iny) \\exp(-ily) \\frac{dy}{2\\pi}}_{\\delta_{ln}} \\\\\n",
    "   &= a_{km} \\delta_{ln}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "and the second\n",
    "\n",
    "$$\n",
    "\\small\n",
    "\\begin{align}\n",
    "\\int_{-1}^1 \\int_{0}^{2\\pi} \\frac{\\partial u}{\\partial y}\\frac{\\partial \\overline{v}}{\\partial y} \\frac{dxdy}{2\\pi} &= \\underbrace{\\int_{-1}^1 (L_m-L_{m+2})(L_k-L_{k+2}) {dx}}_{b_{km}} \\underbrace{\\int_{0}^{2\\pi} ln \\exp(iny) \\exp(-ily)\\frac{dy}{2\\pi}}_{l^2\\delta_{ln}} \\\\\n",
    "   &= l^2 b_{km} \\delta_{ln}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "All in all:\n",
    "\n",
    "$$\n",
    "(\\nabla u, \\nabla v)_w = \\left(a_{km} \\delta_{ln} + l^2 b_{km} \\delta_{ln}\\right)\n",
    "$$\n",
    "\n",
    "The sum of two tensor product matrices!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "A = inner(grad(u), grad(v)) # <- list of two TPMatrices\n",
    "for mat in A[1].mats:  # a_{km} and \\delta_{ln}\n",
    "    print(mat.diags().todense())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "plt.spy(A[1].diags(), markersize=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# 3D Poisson (with MPI and Fourier x 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from sympy import symbols, sin, cos, lambdify\n",
    "from shenfun import *\n",
    "\n",
    "# Use sympy to compute manufactured solution\n",
    "x, y, z = symbols(\"x,y,z\")\n",
    "ue = (cos(4*x) + sin(2*y) + sin(4*z))*(1-x**2)\n",
    "fe = ue.diff(x, 2) + ue.diff(y, 2) + ue.diff(z, 2)\n",
    "\n",
    "C0 = FunctionSpace(32, 'Chebyshev', bc=(0, 0))\n",
    "F1 = FunctionSpace(32, 'Fourier', dtype='D')\n",
    "F2 = FunctionSpace(32, 'Fourier', dtype='d')\n",
    "T = TensorProductSpace(comm, (C0, F1, F2))\n",
    "u = TrialFunction(T)\n",
    "v = TestFunction(T)\n",
    "\n",
    "# Assemble left and right hand\n",
    "f_hat = inner(v, Array(T, buffer=fe))\n",
    "A = inner(v, div(grad(u)))\n",
    "\n",
    "# Solve\n",
    "solver = chebyshev.la.Helmholtz(*A) # Very fast solver due to Jie Shen\n",
    "u_hat = Function(T)\n",
    "u_hat = solver(f_hat, u_hat)\n",
    "assert np.linalg.norm(u_hat.backward()-Array(T, buffer=ue)) < 1e-12\n",
    "print(u_hat.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Contour plot of slice with constant y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "X = T.local_mesh()\n",
    "ua = u_hat.backward()\n",
    "plt.contourf(X[2][0, 0, :], X[0][:, 0, 0], ua[:, 2], 100)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Run with MPI distribution of arrays\n",
    "\n",
    "Here we would normally run from a bash shell\n",
    "<p style=\"margin-bottom:0.5cm;\">\n",
    "\n",
    "<div style=\"color:black\"> <strong>[bash shell] mpirun -np 4 python poisson3D.py </strong> </div>\n",
    "\n",
    "Since we are in a Jupyter notebook, lets actually do this from python in a live cell:-) The exclamation mark '!' is a magic command to run bash scripts in Jupyter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "!mpirun -np 8 python poisson3D.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Note that Fourier bases are especially attractive because of features easily handled with MPI:\n",
    "\n",
    "    - diagonal matrices\n",
    "    - fast transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Nonlinearities (convolution)\n",
    "All treated with pseudo-spectral techniques\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\hat{w}_k &= \\widehat{u^2}_k\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "That is, transform `Function`s to real space `Array`s, perform the nonlinear operation there and transform the nonlinear product back to spectral space (to a `Function`).\n",
    "\n",
    "3/2-rule or 2/3-rule is possible for dealiasing with Fourier. Not for the remaining bases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "uj = Array(SD)\n",
    "#uj[:] = np.random.random(uj.shape)\n",
    "uj = uh.backward(uj)\n",
    "wh = Function(SD)\n",
    "wh = SD.forward(uj*uj, wh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Mixed tensor product spaces\n",
    "\n",
    "Solve several equations simultaneously\n",
    "\n",
    "- Coupled equations\n",
    "- Block matrices and vectors\n",
    "- Tensor spaces of vectors, like velocity $u \\in [\\mathbb{R}^3]^3$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Stokes equations\n",
    "### lid-driven cavity - coupled solver\n",
    "<p style=\"margin-bottom:0.25cm;\">\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\nabla^2 \\mathbf{u} - \\nabla p &= \\mathbf{f} \\quad \\text{in }  \\Omega, \\quad \\quad \\Omega = [-1, 1]\\times[-1, 1]\\\\ \n",
    "\\nabla \\cdot \\mathbf{u} &= h \\quad \\text{in } \\Omega  \\\\ \n",
    "\\int_{\\Omega} p dx &= 0 \\\\\n",
    "\\mathbf{u}(\\pm 1, y) = \\mathbf{u}(x, -1) = (0, 0) &\\text{ and }\\mathbf{u}(x, 1) = (1, 0) \\text{ or } ((1-x^2)(1+x^2), 0)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Given appropriate spaces $V$ and $Q$ a variational form reads: find $(\\mathbf{u}, p) \\in V \\times Q$ such that \n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "a((\\mathbf{u}, p), (\\mathbf{v}, q)) = L((\\mathbf{v}, q)) \\quad \\forall (\\mathbf{v}, q) \\in V \\times Q\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "where bilinear and linear forms are, respectively\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "    a((\\mathbf{u}, p), (\\mathbf{v}, q)) = \\int_{\\Omega} (\\nabla^2 \\mathbf{u} - \\nabla p) \\cdot {\\mathbf{v}} \\, dx_w + \\int_{\\Omega} \\nabla \\cdot \\mathbf{u} \\, {q} \\, dx_w,\n",
    "\\end{equation}\n",
    "$$\n",
    "$$\n",
    "\\begin{equation}  \n",
    "    L((\\mathbf{v}, q)) = \\int_{\\Omega} \\mathbf{f} \\cdot {\\mathbf{v}}\\, dx_w + \\int_{\\Omega} h {q} \\, dx_w\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Using integration by parts for Legendre\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "    a((\\mathbf{u}, p), (\\mathbf{v}, q)) = -\\int_{\\Omega} \\nabla \\mathbf{u} \\cdot \\nabla{\\mathbf{v}} \\, dx_w + \\int_{\\Omega} \\nabla \\cdot \\mathbf{v} \\, {p} \\, dx_w + \\int_{\\Omega} \\nabla \\cdot \\mathbf{u} \\, {q} \\, dx_w,\n",
    "\\end{equation}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Implementation of spaces, basis functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "N = (40, 40)\n",
    "family = 'Legendre'\n",
    "D0X = FunctionSpace(N[0], 'Legendre', bc=(0, 0))\n",
    "#D1Y = FunctionSpace(N[1], 'Legendre', bc=(1, 0)) # Regular lid\n",
    "D1Y = FunctionSpace(N[1], 'Legendre', bc=(0, (1-x)**2*(1+x)**2)) # Regularized lid\n",
    "D0Y = FunctionSpace(N[1], 'Legendre', bc=(0, 0))\n",
    "PX = FunctionSpace(N[0], 'Legendre')\n",
    "PY = FunctionSpace(N[1], 'Legendre')\n",
    "\n",
    "# All required spaces\n",
    "V1 = TensorProductSpace(comm, (D0X, D1Y))\n",
    "V0 = TensorProductSpace(comm, (D0X, D0Y))\n",
    "Q = TensorProductSpace(comm, (PX, PY), modify_spaces_inplace=True)\n",
    "V = VectorSpace([V1, V0])\n",
    "W = CompositeSpace([V0, V0])\n",
    "VQ = CompositeSpace([V, Q])\n",
    "\n",
    "# For inf-sup use P_N - P_{N-2} for velocity-pressure\n",
    "PX.slice = lambda: slice(0, PX.N-2)\n",
    "PY.slice = lambda: slice(0, PY.N-2)\n",
    "\n",
    "# All required test and trial functions\n",
    "up = TrialFunction(VQ)\n",
    "vq = TestFunction(VQ)\n",
    "u, p = up\n",
    "v, q = vq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Implementation Stokes - matrices and solve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Assemble matrices\n",
    "A = inner(grad(v), -grad(u))\n",
    "G = inner(div(v), p)\n",
    "D = inner(q, div(u))\n",
    "\n",
    "# Create Block matrix\n",
    "sol = la.BlockMatrixSolver(A+G+D)\n",
    "\n",
    "# Functions to hold solution and rhs\n",
    "up_hat = Function(VQ)\n",
    "fh_hat = Function(VQ)\n",
    "\n",
    "# Solve Stokes problem. Note constraint for pressure\n",
    "up_hat = sol(fh_hat, u=up_hat, constraints=((2, 0, 0),)) \n",
    "\n",
    "# Move solution to Array in real space\n",
    "up = up_hat.backward()\n",
    "u_, p_ = up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "X = Q.local_mesh(True)\n",
    "plt.quiver(X[0], X[1], u_[0], u_[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Sparsity pattern of block matrix\n",
    "\n",
    "$$\n",
    "M =\n",
    "  \\begin{bmatrix}\n",
    "      A[0]+A[1] & 0 & G[0] \\\\\n",
    "      0 & A[2]+A[3] & G[1] \\\\ \n",
    "      D[0] & D[1] & 0\n",
    "  \\end{bmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "plt.figure(figsize=(6,4))\n",
    "plt.spy(sol.mat.diags(), markersize=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "# Block matrix\n",
    "\n",
    "$$\n",
    "M =\n",
    "  \\begin{bmatrix}\n",
    "      A[0]+A[1] & 0 & G[0] \\\\\n",
    "      0 & A[2]+A[3] & G[1] \\\\ \n",
    "      D[0] & D[1] & 0\n",
    "  \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "where $D = G^T$ for the Legendre basis, making $M$ symmetric. For Chebyshev $M$ will not be symmetric.\n",
    "\n",
    "Solver through [scipy.sparse.linalg](https://docs.scipy.org/doc/scipy/reference/sparse.linalg.html)\n",
    "\n",
    "For Navier-Stokes of the lid-driven cavity, see https://github.com/spectralDNS/shenfun/blob/master/demo/NavierStokesDrivenCavity.py"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "shenfun39",
   "language": "python",
   "name": "shenfun39"
  },
  "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.9.7"
  },
  "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": 2
}