{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Description\n",
    "\n",
    "This notebook intends to gather all the functionalities you'll have to implement for assignment 2.3.\n",
    "\n",
    "# Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import igl\n",
    "import meshplot as mp\n",
    "import time\n",
    "\n",
    "import sys as _sys\n",
    "_sys.path.append(\"../src\")\n",
    "from elasticsolid import *\n",
    "from elasticenergy import *\n",
    "from matplotlib import gridspec\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "shadingOptions = {\n",
    "    \"flat\":True,\n",
    "    \"wireframe\":False,   \n",
    "}\n",
    "\n",
    "rot = np.array(\n",
    "    [[1, 0, 0 ],\n",
    "     [0, 0, 1],\n",
    "     [0, -1, 0 ]]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load mesh\n",
    "\n",
    "Several meshes are available for you to play with under `data/`: `ball.obj`, `dinosaur.obj`, and `beam.obj`. You can also uncomment the few commented lines below to manipulate a simple mesh made out of 2 tetrahedra."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc63cdd8cfe14603b4e54d871b4206d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(-1.987469…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "v, _, _, t, _, _ = igl.read_obj(\"../data/dinosaur.obj\")\n",
    "# v, _, _, t, _, _ = igl.read_obj(\"../data/beam.obj\")\n",
    "\n",
    "# t = np.array([\n",
    "#         [0, 1, 2, 3],\n",
    "#         [1, 2, 3, 4]\n",
    "#     ])\n",
    "# v = np.array([\n",
    "#     [0., 0., 0.],\n",
    "#     [1., 0., 0.],\n",
    "#     [0., 1., 0.],\n",
    "#     [0., 0., 1.],\n",
    "#     [2/3, 2/3, 2/3]\n",
    "# ])\n",
    "\n",
    "be = igl.edges(igl.boundary_facets(t))\n",
    "e = igl.edges(t)\n",
    "\n",
    "aabb = np.max(v, axis=0) - np.min(v, axis=0)\n",
    "length_scale = np.mean(aabb)\n",
    "\n",
    "p = mp.plot(v @ rot.T, t, shading=shadingOptions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linear/Non-Linear Elastic Solid\n",
    "\n",
    "## Instantiation\n",
    "\n",
    "We first specify the elasticity model to use for the elastic solid, as well as pinned vertices, and volumetric forces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho     = 131  # [kg.m-3]\n",
    "young   = 3e8 # [Pa] \n",
    "poisson = 0.2\n",
    "force_mass = np.zeros(shape = (3,))\n",
    "force_mass[2] = - rho * 9.81\n",
    "\n",
    "# minX    = np.min(v[:, 0])\n",
    "# pin_idx = np.arange(v.shape[0])[v[:, 0] < minX + 0.2*aabb[0]]\n",
    "minZ    = np.min(v[:, 2])\n",
    "pin_idx = np.arange(v.shape[0])[v[:, 2] < minZ + 0.1*aabb[2]]\n",
    "\n",
    "# ee    = LinearElasticEnergy(young, poisson)\n",
    "ee    = NeoHookeanElasticEnergy(young, poisson)\n",
    "solid = ElasticSolid(v, t, ee, rho=rho, pin_idx=pin_idx, f_mass=force_mass)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deform the mesh\n",
    "\n",
    "This should now involve elastic forces computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6c6920e14be4332b4a96e814c9236da",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(-3.974938…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v_def = v.copy()\n",
    "v_def[:, 0] *= 2.\n",
    "solid.update_def_shape(v_def)\n",
    "\n",
    "p = mp.plot(solid.v_def @ rot.T, solid.t, shading=shadingOptions)\n",
    "p.add_points(solid.v_def[solid.pin_idx, :] @ rot.T, shading={\"point_color\":\"black\", \"point_size\": 0.1 * length_scale})\n",
    "forcesScale = 2 * np.max(np.linalg.norm(solid.f_ext, axis=1))\n",
    "p.add_lines(solid.v_def @ rot.T, (solid.v_def + solid.f_ext / forcesScale) @ rot.T)\n",
    "p.add_edges(v @ rot.T, be, shading={\"line_color\": \"blue\"})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find equilibrium\n",
    "\n",
    "We compare different methods: number of steps, computation time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps = 1000\n",
    "thresh  = 1.\n",
    "v_init = v.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gradient descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fe18307f9c24439db1ed8fc58928c284",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(-1.987469…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%run ../src/Utils.py\n",
    "\n",
    "report_GD = equilibrium_convergence_report_GD(solid, v_init, n_steps, 1e-8, thresh=thresh)\n",
    "energies_el_GD  = report_GD['energies_el']\n",
    "energies_ext_GD = report_GD['energies_ext']\n",
    "energy_GD       = energies_el_GD + energies_ext_GD\n",
    "residuals_GD    = report_GD['residuals']\n",
    "times_GD        = report_GD['times']\n",
    "idx_stop_GD     = report_GD['idx_stop']\n",
    "v_def_GD        = solid.v_def.copy()\n",
    "\n",
    "# Lastly, plot the resulting shape\n",
    "p = mp.plot(v_def_GD @ rot.T, solid.t, shading=shadingOptions)\n",
    "forcesScale = 2 * np.max(np.linalg.norm(solid.f_ext, axis=1))\n",
    "p.add_lines(v_def_GD @ rot.T, (solid.v_def + solid.f_ext / forcesScale) @ rot.T)\n",
    "p.add_points(v_def_GD[pin_idx, :] @ rot.T, shading={\"point_color\":\"black\", \"point_size\": 0.1 * length_scale})\n",
    "p.add_edges(v @ rot.T, be, shading={\"line_color\": \"blue\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BFGS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0621481088744a34af8662b30613ceac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(children=(DirectionalLight(color='white', intensity=0.6, position=(-0.297981…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "report_BFGS = equilibrium_convergence_report_BFGS(solid, v_init, n_steps, 1e-8, thresh=thresh)\n",
    "energies_el_BFGS  = report_BFGS['energies_el']\n",
    "energies_ext_BFGS = report_BFGS['energies_ext']\n",
    "energy_BFGS       = energies_el_BFGS + energies_ext_BFGS\n",
    "residuals_BFGS    = report_BFGS['residuals']\n",
    "times_BFGS        = report_BFGS['times']\n",
    "idx_stop_BFGS     = report_BFGS['idx_stop']\n",
    "v_def_BFGS        = solid.v_def.copy()\n",
    "\n",
    "# Lastly, plot the resulting shape\n",
    "p = mp.plot(v_def_BFGS @ rot.T, solid.t, shading=shadingOptions)\n",
    "forcesScale = 2 * np.max(np.linalg.norm(solid.f_ext, axis=1))\n",
    "p.add_lines(v_def_BFGS @ rot.T, (solid.v_def + solid.f_ext / forcesScale) @ rot.T)\n",
    "p.add_points(v_def_BFGS[pin_idx, :] @ rot.T, shading={\"point_color\":\"black\", \"point_size\": 0.1 * length_scale})\n",
    "p.add_edges(v @ rot.T, be, shading={\"line_color\": \"blue\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Newton-CG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Energy: -127985.65419881282 Force norm: 287742.8969935862 Line search Iters: 0\n",
      "Energy: -178892.14400109363 Force norm: 182952.00531737506 Line search Iters: 0\n",
      "Energy: -206901.42653635694 Force norm: 151950.01953001413 Line search Iters: 0\n",
      "Energy: -227792.85782151608 Force norm: 135158.66283929092 Line search Iters: 0\n",
      "Energy: -244838.85692066932 Force norm: 123433.77675577633 Line search Iters: 0\n",
      "Energy: -259318.48388755036 Force norm: 114389.35466936302 Line search Iters: 0\n",
      "Energy: -271937.0843840749 Force norm: 107190.15322070086 Line search Iters: 0\n",
      "Energy: -283160.0740997393 Force norm: 101411.77988596291 Line search Iters: 0\n"
     ]
    }
   ],
   "source": [
    "report_NCG = equilibrium_convergence_report_NCG(solid, v_init, n_steps, thresh=thresh)\n",
    "energies_el_NCG  = report_NCG['energies_el']\n",
    "energies_ext_NCG = report_NCG['energies_ext']\n",
    "energy_NCG       = energies_el_NCG + energies_ext_NCG\n",
    "residuals_NCG    = report_NCG['residuals']\n",
    "times_NCG        = report_NCG['times']\n",
    "idx_stop_NCG     = report_NCG['idx_stop']\n",
    "v_def_NCG        = solid.v_def.copy()\n",
    "\n",
    "# Lastly, plot the resulting shape\n",
    "p = mp.plot(v_def_NCG @ rot.T, solid.t, shading=shadingOptions)\n",
    "forcesScale = 2 * np.max(np.linalg.norm(solid.f_ext, axis=1))\n",
    "p.add_lines(v_def_NCG @ rot.T, (solid.v_def + solid.f_ext / forcesScale) @ rot.T)\n",
    "p.add_points(v_def_NCG[pin_idx, :] @ rot.T, shading={\"point_color\":\"black\", \"point_size\": 0.1 * length_scale})\n",
    "p.add_edges(v @ rot.T, be, shading={\"line_color\": \"blue\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare the different algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap   = plt.get_cmap('viridis')\n",
    "colors = cmap(np.linspace(0., 1., 4))\n",
    "\n",
    "gs = gridspec.GridSpec(nrows=1, ncols=1, width_ratios=[1], height_ratios=[1])\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "\n",
    "axTmp = plt.subplot(gs[0, 0])\n",
    "axTmp.plot(times_GD[:idx_stop_GD+1], residuals_GD[:idx_stop_GD+1], c=colors[0], \n",
    "           label=\"Gradient Descent ({:}its)\".format(idx_stop_GD))\n",
    "axTmp.plot(times_BFGS[:idx_stop_BFGS+1], residuals_BFGS[:idx_stop_BFGS+1], c=colors[1], \n",
    "           label=\"BFGS ({:}its)\".format(idx_stop_BFGS))\n",
    "axTmp.plot(times_NCG[:idx_stop_NCG+1], residuals_NCG[:idx_stop_NCG+1], c=colors[2], \n",
    "           label=\"Newton-CG ({:}its)\".format(idx_stop_NCG))\n",
    "y_lim = axTmp.get_ylim()\n",
    "axTmp.set_title(\"Residuals as optimization goes\", fontsize=14)\n",
    "axTmp.set_xlabel(\"Computation time [s]\", fontsize=12)\n",
    "axTmp.set_ylabel(\"Force residuals [N]\", fontsize=12)\n",
    "axTmp.set_ylim(y_lim)\n",
    "plt.legend(fontsize=12)\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gs = gridspec.GridSpec(nrows=1, ncols=1, width_ratios=[1], height_ratios=[1])\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "\n",
    "axTmp = plt.subplot(gs[0, 0])\n",
    "axTmp.plot(times_GD[:idx_stop_GD+1], energy_GD[:idx_stop_GD+1], c=colors[0], \n",
    "           label=\"Gradient Descent ({:}its)\".format(idx_stop_GD))\n",
    "axTmp.plot(times_BFGS[:idx_stop_BFGS+1], energy_BFGS[:idx_stop_BFGS+1], c=colors[1], \n",
    "           label=\"BFGS ({:}its)\".format(idx_stop_BFGS))\n",
    "axTmp.plot(times_NCG[:idx_stop_NCG+1], energy_NCG[:idx_stop_NCG+1], c=colors[2], \n",
    "           label=\"Newton-CG ({:}its)\".format(idx_stop_NCG))\n",
    "y_lim = axTmp.get_ylim()\n",
    "axTmp.set_title(\"Total energy as optimization goes\", fontsize=14)\n",
    "axTmp.set_xlabel(\"Computation time [s]\", fontsize=12)\n",
    "axTmp.set_ylabel(\"Energy [J]\", fontsize=12)\n",
    "axTmp.set_ylim(y_lim)\n",
    "plt.legend(fontsize=12)\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "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
}