{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Non-orthogonal curvilinear coordinates\n",
    "\n",
    "In this notebook we solve Poisson's equation on a 2D wavy domain, using non-orthogonal basis vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import Math\n",
    "from shenfun import *\n",
    "from shenfun.la import Solver2D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the non-orthogonal curvilinear coordinates. Here `psi=(u, v)` are the new coordinates and `rv` is the position vector\n",
    "\n",
    "$$\n",
    "\\vec{r} = u \\mathbf{i} + v\\left(1- \\frac{\\sin 2u}{10} \\right) \\mathbf{j}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{i}$ and $\\mathbf{j}$ are the Cartesian unit vectors in $x$- and $y$-directions, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = sp.Symbol('x', real=True, positive=True)\n",
    "v = sp.Symbol('y', real=True)\n",
    "psi = (u, v)\n",
    "rv = (u, v*(1-sp.sin(2*u)/10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now choose basis functions and create tensor product space. Notice that one has to use complex Fourier space and not the real, because the integral measure is a function of u."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 20\n",
    "#B0 = FunctionSpace(N, 'C', bc=(0, 0), domain=(0, 2*np.pi))\n",
    "B0 = FunctionSpace(N, 'F', dtype='D', domain=(0, 2*np.pi))\n",
    "B1 = FunctionSpace(N, 'L', bc=(0, 0), domain=(-1, 1))\n",
    "\n",
    "T = TensorProductSpace(comm, (B0, B1), dtype='D', coordinates=(psi, rv, sp.Q.negative(sp.sin(2*u)-10) & sp.Q.negative(sp.sin(2*u)/10-1)))\n",
    "p = TrialFunction(T)\n",
    "q = TestFunction(T)\n",
    "b = T.coors.get_covariant_basis()\n",
    "T.coors.sg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.Matrix(T.coors.get_covariant_metric_tensor())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Math(T.coors.latex_basis_vectors(covariant=True, symbol_names={u: 'u', v: 'v'}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the mesh to see the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = T.local_cartesian_mesh()\n",
    "x, y = mesh\n",
    "plt.figure(figsize=(10, 4))\n",
    "for i, (xi, yi) in enumerate(zip(x, y)):\n",
    "    plt.plot(xi, yi, 'b')\n",
    "    plt.plot(x[:, i], y[:, i], 'r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print the Laplace operator in curvilinear coordinates. We use `replace` to simplify the expression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dp = div(grad(p))\n",
    "g = sp.Symbol('g', real=True, positive=True)\n",
    "replace = [(1-sp.sin(2*u)/10, sp.sqrt(g)), (sp.sin(2*u)-10, -10*sp.sqrt(g)), (5*sp.sin(2*u)-50, -50*sp.sqrt(g))]\n",
    "Math((dp*T.coors.sg**2).tolatex(funcname='p', symbol_names={u: 'u', v: 'v'}, replace=replace))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve Poisson's equation. First define a manufactured solution and assemble the right hand side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ue = sp.sin(2*u)*(1-v**2)\n",
    "f = (div(grad(p))).tosympy(basis=ue, psi=psi)\n",
    "fj = Array(T, buffer=f*T.coors.sg)\n",
    "f_hat = Function(T)\n",
    "f_hat = inner(q, fj, output_array=f_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then assemble the left hand side and solve using a generic 2D solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M = inner(q, div(grad(p))*T.coors.sg)\n",
    "#M = inner(grad(q*T.coors.sg), -grad(p))\n",
    "u_hat = Function(T)\n",
    "Sol1 = Solver2D(M)\n",
    "u_hat = Sol1(f_hat, u_hat)\n",
    "uj = u_hat.backward()\n",
    "uq = Array(T, buffer=ue)\n",
    "print('Error =', np.linalg.norm(uj-uq))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(M)):\n",
    "    print(len(M[i].mats[0].keys()), len(M[i].mats[1].keys()), M[i].mats[0].measure, M[i].mats[1].measure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the solution in the wavy Cartesian domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xx, yy = T.local_cartesian_mesh()\n",
    "plt.figure(figsize=(12, 4))\n",
    "plt.contourf(xx, yy, uj.real)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect the sparsity pattern of the generated matrix on the left hand side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import spy\n",
    "plt.figure()\n",
    "spy(Sol1.M, markersize=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse.linalg import eigs\n",
    "mats = inner(q*T.coors.sg, -div(grad(p)))\n",
    "Sol1 = Solver2D(mats)\n",
    "BB = inner(p, q*T.coors.sg)\n",
    "Sol2 = Solver2D(BB)\n",
    "f = eigs(Sol1.M, k=20, M=Sol2.M, which='LM', sigma=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l = 10\n",
    "u_hat = Function(T)\n",
    "u_hat[:, :-2] = f[1][:, l].reshape(T.dims())\n",
    "plt.contourf(xx, yy, u_hat.backward().real, 100)\n",
    "plt.colorbar()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "shenfun38",
   "language": "python",
   "name": "shenfun38"
  },
  "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.6"
  },
  "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
}