{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TriaMesh Geodesics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lapy import TriaMesh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load a triangle mesh of a flat square (OFF file)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--> OFF format         ... \n",
      " --> DONE ( V: 415 , T: 768 )\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'TriaMesh'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "T = TriaMesh.read_off(\"../data/square-mesh.off\")\n",
    "type(T).__name__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "        <script type=\"text/javascript\">\n",
       "        window.PlotlyConfig = {MathJaxConfig: 'local'};\n",
       "        if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}\n",
       "        if (typeof require !== 'undefined') {\n",
       "        require.undef(\"plotly\");\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': ['https://cdn.plot.ly/plotly-2.18.0.min']\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(Plotly) {\n",
       "            window._Plotly = Plotly;\n",
       "        });\n",
       "        }\n",
       "        </script>\n",
       "        "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# import plotting functions\n",
    "# import plotly\n",
    "# plotly.offline.init_notebook_mode(connected=True)\n",
    "import plotly.io as pio\n",
    "\n",
    "from lapy.plot import plot_tria_mesh\n",
    "\n",
    "pio.renderers.default = \"sphinx_gallery\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now plot the triangle mesh with a function overlay of the triangle quality. Note, that this is a function attached to the triangles, not the vertices, so it is piecewise flat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "q = T.tria_qualities()\n",
    "plot_tria_mesh(T,plot_edges=True, tfunc=q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Laplace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we check a few properties of eigenfunctions. For this we get the first few and also the stiffness matrix (A) and the lumped mass matrix (B, reduced to a diagonal), of which we can easily compute the inverse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TriaMesh with regular Laplace-Beltrami\n",
      "Solver: spsolve (LU decomposition) ...\n"
     ]
    }
   ],
   "source": [
    "# compute first eigenfunction\n",
    "from lapy import Solver\n",
    "\n",
    "fem = Solver(T,lump=True)\n",
    "eval, evec = fem.eigs()\n",
    "vfunc = evec[:,1]\n",
    "\n",
    "# also get A,B (lumped), and inverse of B (easy as it is diagonal due to lumping)\n",
    "A, B = fem.stiffness, fem.mass\n",
    "Bi = B.copy()\n",
    "Bi.data **= -1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The mass matrix B represents the inner product so that the integral of the product of two functions x and y over the whole mesh is x B y'. The lumped mass matrix that we use here is a simplified version, as all off-diagonal elements are added to the diagonal. The entries on the diagonal represent the vertex areas and their sum is the total area of the mesh. For our unit square it will be 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what happens when we apply the Laplace operator to an eigenfunction. Eigenfunctions are solutions to Delta f = - lambda f, so we should obtain a scaled version of f."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_tria_mesh(T,Bi*(A*vfunc),plot_edges=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the same as the corresponding eigenvalue times the eigenfunction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_tria_mesh(T,eval[1]*vfunc,plot_edges=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Laplace is also defined as the -div(grad(f)). So first applying the gradient and then the divergence to an eigenfunction and then multiplying with inv(B) should yield the same result as above again. Note, that multiplying with inv(B) is necessary to get back from the integrated divergence to the original function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lapy.diffgeo import compute_divergence, compute_gradient\n",
    "\n",
    "grad = compute_gradient(T,vfunc)\n",
    "divx = -compute_divergence(T,grad)\n",
    "plot_tria_mesh(T,Bi*divx,plot_edges=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Geodesics\n",
    "\n",
    "Now we will replicate the idea of geodesics in heat, where first a heat diffusion is solved and massaged in the right way to yield an approximation of geodesics on the mesh. This also works on curved meshes, but for simplicity we keep using the square here. So let's start with computing the heat diffusion from the boundary (with default time factor m=1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TriaMesh with regular Laplace-Beltrami\n",
      "Matrix Format now:  csc\n",
      "Solver: spsolve (LU decomposition) ...\n"
     ]
    }
   ],
   "source": [
    "from lapy import heat\n",
    "\n",
    "bvert = T.boundary_loops()\n",
    "u = heat.diffusion(T,bvert,m=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We show some of the level sets. Note, that they are not evenly spaced and get steeper closer to the boundary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_tria_mesh(T,u,plot_edges=True,plot_levels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next step is to compute the gradient (vector field) of the heat diffusion function and normalize all vectors to unit length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# compute gradient of heat diffusion\n",
    "tfunc = compute_gradient(T,u)\n",
    "\n",
    "# normalize gradient\n",
    "X = -tfunc / np.sqrt((tfunc**2).sum(1))[:,np.newaxis]\n",
    "X = np.nan_to_num(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we get the integrated divergence of the normalized gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "divx = compute_divergence(T,X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, to obtain the distance function, we need to solve a Poisson equation. The solution can be shifted arbitrary, so we need to subtract the minimum, which should be along the boundary of the mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TriaMesh with regular Laplace-Beltrami\n",
      "Matrix Format now: csc\n",
      "Solver: cholesky decomp - performance optimal ...\n"
     ]
    }
   ],
   "source": [
    "# compute distance\n",
    "from scipy.sparse.linalg import splu\n",
    "\n",
    "useCholmod = True\n",
    "try:\n",
    "    from sksparse.cholmod import cholesky\n",
    "except ImportError:\n",
    "    useCholmod = False\n",
    "\n",
    "fem = Solver(T,lump=True)\n",
    "A, B = fem.stiffness, fem.mass\n",
    "\n",
    "H=-A\n",
    "b0=divx\n",
    "        \n",
    "# solve H x = b0\n",
    "# we don't need the B matrix here, as divx is the integrated divergence \n",
    "print(\"Matrix Format now: \"+H.getformat())\n",
    "if useCholmod:\n",
    "    print(\"Solver: cholesky decomp - performance optimal ...\")\n",
    "    chol = cholesky(H)\n",
    "    x = chol(b0)\n",
    "else:\n",
    "    print(\"Solver: spsolve (LU decomp) - performance not optimal ...\")\n",
    "    lu = splu(H)\n",
    "    x = lu.solve(b0)\n",
    "\n",
    "# remove shift\n",
    "x = x-min(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In short, the idea is to first get a function (heat diffusion) that flows from a point or from a set of points (boundary) through the mesh, then normalize its gradient, compute the divergence and finally step backward through the Laplace operator to find a function that has this normalized gradient. If we look at it, we actually notice that level sets are equally spaced now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_tria_mesh(T,x,plot_edges=True,plot_levels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice, but only an approximation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.6049783117351546, 0.7071067811865476)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# max distance (smoothed)\n",
    "(max(x), np.sqrt(2)/2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of computing the solution ourselves, we could simply employ the standard Poisson solver on inv(B) times divx."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matrix Format now: csc\n",
      "Solver: spsolve (LU decomposition) ...\n"
     ]
    }
   ],
   "source": [
    "vf = fem.poisson(-Bi*divx)\n",
    "vf = vf - min(vf)\n",
    "plot_tria_mesh(T,vf,plot_edges=True,plot_levels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This should give us the same result as what we had earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5.30470728232757e-07"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "max(abs(vf-x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or we can just call compute_geodesic_f which does all the work for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TriaMesh with regular Laplace-Beltrami\n",
      "Matrix Format now: csc\n",
      "Solver: spsolve (LU decomposition) ...\n"
     ]
    }
   ],
   "source": [
    "from lapy.diffgeo import compute_geodesic_f\n",
    "\n",
    "gf = compute_geodesic_f(T,u)\n",
    "plot_tria_mesh(T,gf,plot_edges=True,plot_levels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And verify it is still the same result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.2569845126451114e-07"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# again should be the same\n",
    "max(abs(gf-x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to normalizing the gradients of a function, we can see if we can rotate it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TriaMesh with regular Laplace-Beltrami\n",
      "Matrix Format now: csc\n",
      "Solver: spsolve (LU decomposition) ...\n"
     ]
    }
   ],
   "source": [
    "# testing if we can rotate the function\n",
    "from lapy.diffgeo import compute_rotated_f\n",
    "\n",
    "gf = compute_rotated_f(T,vf)\n",
    "plot_tria_mesh(T,gf,plot_edges=True,plot_levels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mean Curvature Mapping\n",
    "\n",
    "To demonstrate the mean curvature mapping to a sphere, we need to have a closed mesh. It should not have too many flat regions (so not the cube) as there is no cuvature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load your mesh here and uncomment. The mesh should have not too many flat regions (not a cube)\n",
    "#from lapy.diffgeo import tria_mean_curvature_flow\n",
    "#from lapy.plot import plot_tria_mesh\n",
    "#T = TriaIO.import_off(\"../data/???\")\n",
    "#T2 = tria_mean_curvature_flow(T)\n",
    "#plot_tria_mesh(T2,plot_edges=True,plot_levels=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3",
   "language": "python",
   "name": "python3"
  },
  "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"
  },
  "nbsphinx": {
   "execute": "always"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}