{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The HFM library - A fast marching solver with adaptive stencils \n",
    "\n",
    "## Part : Algorithmic enhancements to the fast marching method\n",
    "## Chapter : Sensitivity in semi-Lagrangian schemes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!---The HFM library provides a variety of metric classes : isotropic, Riemannian, Rander, asymmetric quadratic, defined by a Hooke tensor, and more.--->\n",
    "\n",
    "This notebook illustrates automatic differentiation with anisotropic metrics, and requires specific computational techniques in comparison to the isotropic case. We illustrate here their foundations and usage.\n",
    "For that purpose, we are dependent on an implementation detail - semi-Lagrangian implementation - that is usually transparent to the user. \n",
    "\n",
    "The methods presented in this notebook apply to the following HFM models : `Rander2`, `AsymmetricQuadratic2`, `Seismic_`, `SeismicTopographic_` as well as `Isotropic_`, `Diagonal_`.\n",
    "The underscore stands for $2$ or $3$. For two-dimensional Riemannian metrics, use `Rander2`.\n",
    "\n",
    "**Eulerian vs Semi-Lagrangian schemes.**\n",
    "The HFM (Hamiltonian Fast Marching) library takes its name from a family of Eulerian schemes based on a specific representation of the Hamiltonian, which are particularly efficient for Riemannian metrics and curvature penalized metrics, see:\n",
    "* Jean-Marie Mirebeau, Jorg Portegies, \"Hamiltonian Fast Marching: A numerical solver for anisotropic and non-holonomic eikonal PDEs\", 2019 [(link)](https://hal.archives-ouvertes.fr/hal-01778322)\n",
    "\n",
    "However, other classes of metrics (such as Rander, asymmetric quadratic, or hooke tensor defined), require a different approach. In their case, we rely on an adaptive semi-Lagrangian scheme, in the spirit of \n",
    "* Mirebeau, J.-M. (2014). Efficient fast marching with Finsler metrics. Numerische Mathematik, 126(3), 515–557. [link](https://hal.archives-ouvertes.fr/hal-00736431)\n",
    "\n",
    "**Semi-Lagrangian schemes** are discretizations of the eikonal equation \n",
    "$$\n",
    "    F_x^*(\\nabla u(x)) = 1\n",
    "$$\n",
    "in the following form which is closely related with Bellman's optimality principle\n",
    "$$\n",
    "    u(x) = \\min_{y \\in S(x)} F_x(x-y) + I_{S(x)} u(y),\n",
    "$$\n",
    "where $S(x)$ the *stencil* is a polygonal neighborhood of $x$ whose vertices are grid-points, and $I_S$ denotes piecewise linear interpolation.\n",
    "The equality is imposed at all points of the discrete or continuous domain. In addition, adequate boundary conditions are imposed on the domain boundary, usually of Dirichlet and of outflow type.\n",
    "\n",
    "**The geodesic flow.**\n",
    "The minimizer $y_*(x) \\in S(x)$ in the semi-lagrangian discretization of the eikonal equation has a geometrical interpretation: $v_*(x) = x-y_*(x)$ approximates the direction of the minimal geodesic passing through $x$, referred to as the geodesic flow direction. This data it is optionally exported by the HFM library, after normalization, under the name `geodesicFlow`.\n",
    "\n",
    "<!---\n",
    "Said otherwise, one has $V(x) \\approx v_*$, where \n",
    "$$\n",
    "    \\gamma'(t) = V(\\gamma(t)) \\quad V(x) = \\nabla F_x^*(\\nabla u(x))\n",
    "$$\n",
    "--->\n",
    "\n",
    "**The envelope theorem.**\n",
    "Under suitable assumptions, one can differentiate the value of an optimization problem w.r.t. parameters, using the so-called envelope theorem: the extremizer variation can be ignored at first order.\n",
    "\n",
    "Said otherwise, if the metric depends on a small parameter $\\epsilon$, then\n",
    "$$\n",
    "    u(x;\\epsilon) = F_x(v_*(x);\\epsilon) + I_{S(x)}u(y_*(x);\\epsilon) + o(\\epsilon).\n",
    "$$\n",
    "\n",
    "As a result, if two parameter dependent metrics $F$ and $G$ obey\n",
    "$$\n",
    "    F_x(v_*(x); \\epsilon) = G_x(v_*(x); \\epsilon) + o(\\epsilon),\n",
    "$$\n",
    "at all discretization points $x$, then they produce identical first order perturbations of $u$.\n",
    "The HFM library takes advantage of this fact to simplify computations, but for that purposes needs the geodesic flow direction $v_*(x)$ to be computed in a preliminary step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[**Summary**](Summary.ipynb) of volume Fast Marching Methods, this series of notebooks.\n",
    "\n",
    "[**Main summary**](../Summary.ipynb) of the Adaptive Grid Discretizations \n",
    "\tbook of notebooks, including the other volumes.\n",
    "\n",
    "# Table of contents\n",
    "  * [1. Forward differentiation](#1.-Forward-differentiation)\n",
    "    * [1.1 Isotropic example](#1.1-Isotropic-example)\n",
    "    * [1.2 Rander example](#1.2-Rander-example)\n",
    "    * [1.3 Seismic example](#1.3-Seismic-example)\n",
    "  * [2. Reverse differentiation](#2.-Reverse-differentiation)\n",
    "    * [2.1 Isotropic example](#2.1-Isotropic-example)\n",
    "    * [2.2 Rander example](#2.2-Rander-example)\n",
    "    * [2.3 Seismic example](#2.3-Seismic-example)\n",
    "  * [3. Geodesic flow and gradient of the value function](#3.-Geodesic-flow-and-gradient-of-the-value-function)\n",
    "    * [3.1 Computing the gradient](#3.1-Computing-the-gradient)\n",
    "    * [3.2 Differentiating the gradient](#3.2-Differentiating-the-gradient)\n",
    "    * [3.3 Differentiating the geodesic flow](#3.3-Differentiating-the-geodesic-flow)\n",
    "\n",
    "\n",
    "\n",
    "This Python&reg; notebook is intended as documentation and testing for the [HamiltonFastMarching (HFM) library](https://github.com/mirebeau/HamiltonFastMarching), which also has interfaces to the Matlab&reg; and Mathematica&reg; languages. \n",
    "More information on the HFM library in the manuscript:\n",
    "* Jean-Marie Mirebeau, Jorg Portegies, \"Hamiltonian Fast Marching: A numerical solver for anisotropic and non-holonomic eikonal PDEs\", 2019 [(link)](https://hal.archives-ouvertes.fr/hal-01778322)\n",
    "\n",
    "Copyright Jean-Marie Mirebeau, University Paris-Sud, CNRS, University Paris-Saclay"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Importing the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T08:58:38.921802Z",
     "iopub.status.busy": "2024-04-30T08:58:38.921510Z",
     "iopub.status.idle": "2024-04-30T08:58:38.930724Z",
     "shell.execute_reply": "2024-04-30T08:58:38.930265Z"
    }
   },
   "outputs": [],
   "source": [
    "import sys; sys.path.insert(0,\"..\") # Allow import of agd from parent directory (useless if conda package installed)\n",
    "#from Miscellaneous import TocTools; print(TocTools.displayTOC('SensitivitySL','FMM'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T08:58:38.933344Z",
     "iopub.status.busy": "2024-04-30T08:58:38.933128Z",
     "iopub.status.idle": "2024-04-30T08:58:39.264189Z",
     "shell.execute_reply": "2024-04-30T08:58:39.263673Z"
    }
   },
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01magd\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Eikonal\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01magd\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutomaticDifferentiation \u001b[38;5;28;01mas\u001b[39;00m ad\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01magd\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FiniteDifferences \u001b[38;5;28;01mas\u001b[39;00m fd\n",
      "File \u001b[0;32m~/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/agd/Eikonal/__init__.py:22\u001b[0m\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mfunctools\u001b[39;00m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mLibraryCall\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m GetBinaryDir\n\u001b[0;32m---> 22\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrun_detail\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Cache\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mDictIn\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m dictIn,dictOut,CenteredLinspace\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_VoronoiDecomposition_noAD\u001b[39m(arr,mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mBoth\u001b[39m\u001b[38;5;124m'\u001b[39m,\u001b[38;5;241m*\u001b[39margs,\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n",
      "File \u001b[0;32m~/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/agd/Eikonal/run_detail.py:6\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mLibraryCall\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m RunDispatch,GetBinaryDir\n\u001b[0;32m----> 6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mDictIn_detail\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m factoringPointChoice\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Metrics\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutomaticDifferentiation \u001b[38;5;28;01mas\u001b[39;00m ad\n",
      "File \u001b[0;32m~/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/agd/Eikonal/DictIn_detail.py:13\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FiniteDifferences \u001b[38;5;28;01mas\u001b[39;00m fd\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LinearParallel \u001b[38;5;28;01mas\u001b[39;00m lp\n\u001b[0;32m---> 13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Metrics\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfactoringPointChoice\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m     16\u001b[0m \t\u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mfactoringPointChoice\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m: \u001b[38;5;66;03m# 'Key' and 'Seed' are synonyms here\u001b[39;00m\n",
      "File \u001b[0;32m~/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/agd/Metrics/__init__.py:27\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Copyright 2020 Jean-Marie Mirebeau, University Paris-Sud, CNRS, University Paris-Saclay\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Distributed WITHOUT ANY WARRANTY. Licensed under the Apache License, Version 2.0, see http://www.apache.org/licenses/LICENSE-2.0\u001b[39;00m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;124;03mThe Metrics package defines a variety of norms on R^d, defined by appropriate parameters,\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;124;03mpossibly non-symmetric. The data of a such a norm, at each point of a domain, is referred \u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[38;5;124;03mAdditional norm/metric classes are defined in the Seismic subpackage.\u001b[39;00m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m---> 27\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase\u001b[39;00m \t    \u001b[38;5;28;01mimport\u001b[39;00m Base\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01misotropic\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Isotropic\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagonal\u001b[39;00m  \u001b[38;5;28;01mimport\u001b[39;00m Diagonal \n",
      "File \u001b[0;32m~/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/agd/Metrics/base.py:9\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LinearParallel \u001b[38;5;28;01mas\u001b[39;00m lp\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FiniteDifferences \u001b[38;5;28;01mas\u001b[39;00m fd\n\u001b[0;32m----> 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Interpolation\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mBase\u001b[39;00m:\n\u001b[1;32m     12\u001b[0m \t\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;124;03m\tBase class for a metric, in other words a family of norms.\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;124;03m\t\"\"\"\u001b[39;00m\n",
      "File \u001b[0;32m~/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/agd/Interpolation.py:16\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mitertools\u001b[39;00m\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mscipy\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlinalg\u001b[39;00m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutomaticDifferentiation \u001b[38;5;28;01mas\u001b[39;00m ad\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21morigin_scale_shape\u001b[39m(grid):\n",
      "File \u001b[0;32m~/opt/miniconda3/envs/agd-hfm_dev_310/lib/python3.10/site-packages/scipy/__init__.py:94\u001b[0m\n\u001b[1;32m     92\u001b[0m         \u001b[38;5;28;01mcontinue\u001b[39;00m\n\u001b[1;32m     93\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m callable(_fun) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(_fun, \u001b[38;5;28mtype\u001b[39m):\n\u001b[0;32m---> 94\u001b[0m         _fun \u001b[38;5;241m=\u001b[39m \u001b[43m_deprecated\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_msg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mformat\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_key\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_fun\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     95\u001b[0m     \u001b[38;5;28mglobals\u001b[39m()[_key] \u001b[38;5;241m=\u001b[39m _fun\n\u001b[1;32m     96\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m np, _types\n",
      "File \u001b[0;32m~/opt/miniconda3/envs/agd-hfm_dev_310/lib/python3.10/site-packages/scipy/_lib/deprecation.py:17\u001b[0m, in \u001b[0;36m_deprecated.<locals>.wrap\u001b[0;34m(fun)\u001b[0m\n\u001b[1;32m     11\u001b[0m     warnings\u001b[38;5;241m.\u001b[39mwarn(\n\u001b[1;32m     12\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTrying to deprecate class \u001b[39m\u001b[38;5;132;01m{!r}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(fun),\n\u001b[1;32m     13\u001b[0m         category\u001b[38;5;241m=\u001b[39m\u001b[38;5;167;01mRuntimeWarning\u001b[39;00m, stacklevel\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m)\n\u001b[1;32m     14\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m fun\n\u001b[1;32m     16\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;129;43m@functools\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwraps\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m---> 17\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mdef\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43mcall\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwarnings\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwarn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmsg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcategory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;167;43;01mDeprecationWarning\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m                  \u001b[49m\u001b[43mstacklevel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstacklevel\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     20\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mreturn\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/miniconda3/envs/agd-hfm_dev_310/lib/python3.10/functools.py:35\u001b[0m, in \u001b[0;36mupdate_wrapper\u001b[0;34m(wrapper, wrapped, assigned, updated)\u001b[0m\n\u001b[1;32m     32\u001b[0m WRAPPER_ASSIGNMENTS \u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__module__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__name__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__qualname__\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__doc__\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     33\u001b[0m                        \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__annotations__\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     34\u001b[0m WRAPPER_UPDATES \u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__dict__\u001b[39m\u001b[38;5;124m'\u001b[39m,)\n\u001b[0;32m---> 35\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mupdate_wrapper\u001b[39m(wrapper,\n\u001b[1;32m     36\u001b[0m                    wrapped,\n\u001b[1;32m     37\u001b[0m                    assigned \u001b[38;5;241m=\u001b[39m WRAPPER_ASSIGNMENTS,\n\u001b[1;32m     38\u001b[0m                    updated \u001b[38;5;241m=\u001b[39m WRAPPER_UPDATES):\n\u001b[1;32m     39\u001b[0m     \u001b[38;5;124;03m\"\"\"Update a wrapper function to look like the wrapped function\u001b[39;00m\n\u001b[1;32m     40\u001b[0m \n\u001b[1;32m     41\u001b[0m \u001b[38;5;124;03m       wrapper is the function to be updated\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[38;5;124;03m       function (defaults to functools.WRAPPER_UPDATES)\u001b[39;00m\n\u001b[1;32m     49\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m     50\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m attr \u001b[38;5;129;01min\u001b[39;00m assigned:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "from agd import Eikonal\n",
    "from agd import AutomaticDifferentiation as ad\n",
    "from agd import FiniteDifferences as fd\n",
    "from agd import Metrics\n",
    "from agd import LinearParallel as lp\n",
    "from agd.Metrics.Seismic import Hooke\n",
    "from agd.Interpolation import UniformGridInterpolation\n",
    "from agd.Plotting import savefig; #savefig.dirName = 'Figures/Sensitivity'\n",
    "norm_infinity = ad.Optimization.norm_infinity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from copy import copy\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reload_packages():\n",
    "    from Miscellaneous.rreload import rreload\n",
    "    global ad,fd,Metrics,lp,Hooke,UniformGridInterpolation\n",
    "    [ad,fd,Metrics,lp,Hooke,UniformGridInterpolation] = rreload([ad,fd,Metrics,lp,Hooke,UniformGridInterpolation],rootdir='..')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Forward differentiation\n",
    "\n",
    "We illustrate forward differentiation of anisotropic fast marching, with Rander metrics.\n",
    "Note that isotropic and Riemannian metrics are special cases of Rander metrics, hence the method presented below extends to those using appropriate casts. And this is indeed what we do in the next subsection.\n",
    "\n",
    "<!---**Note** The default numerical schemes for the models 'Isotropic2' and 'Riemann2' do not --->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Isotropic example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use a Rander metric to encode an isotropic metric, so as to illustrate the semi-Lagrangian discretization and its automatic differentiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "hfmIn = Eikonal.dictIn({\n",
    "    'model':'Rander2',\n",
    "    'arrayOrdering':'RowMajor',\n",
    "    'seeds':[[0,0]],\n",
    "    'extractValues':True,\n",
    "})\n",
    "hfmIn.SetRect(sides=[[-1,1],[-1,1]],dimx=101)\n",
    "X = hfmIn.Grid()\n",
    "\n",
    "delta = ad.Dense.identity((1,))\n",
    "metric =  Metrics.Isotropic(1.+delta[0],vdim=2)\n",
    "#metric = Metrics.Riemann.from_diagonal(1.+delta[0],1.+delta[0]) #Alternatively"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As described above, automatic differentiation in the semi-Lagrangian setting is a two-step process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dictIn({'arrayOrdering': 'RowMajor', 'model': 'Rander2', 'seeds': array([[0., 0.]]), 'extractValues': True, 'gridScale': array(0.01980198), 'dims': array([101., 101.]), 'origin': array([-1., -1.])})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hfmIn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "hfmIn.factoringPointChoice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.004497 s.\n",
      "--- HFM call triggered above to compute geodesic flow ---\n",
      "Filling cache data\n",
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "hfmIn['metric'] = metric\n",
    "hfmOut_ad,values_ad = hfmIn.Run()\n",
    "plt.axis('equal')\n",
    "plt.contourf(*X,values_ad.gradient(0)); "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The caching of the geodesicflow could be done manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- First run for caching geodesic flow information ---\n",
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.004363 s.\n",
      "Filling cache data\n",
      "--- Second run using cached data :  dict_keys(['values', 'activeNeighs', 'geodesicFlow'])  ---\n",
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "print(\"--- First run for caching geodesic flow information ---\")\n",
    "hfmIn_noad = hfmIn.copy()\n",
    "hfmIn_noad['metric'] = ad.remove_ad(metric,iterables=(Metrics.Base,))\n",
    "cache = Eikonal.Cache(needsflow=True)\n",
    "hfmIn_noad.Run(cache=cache)\n",
    "\n",
    "print(\"--- Second run using cached data : \", cache.contents.keys(),\" ---\")\n",
    "hfmOut_ad,values_ad = hfmIn.Run(cache=cache)\n",
    "plt.axis('equal')\n",
    "plt.contourf(*X,values_ad.gradient(0));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Rander example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rander metrics are a generalization of Riemannian metrics, involving an additional asymmetric part. \n",
    "In particular, any Riemannian metric can be cast into a Rander metric, which allows to use automatic differentiation in this context. We reproduce an instance of Zermelo's navigation problem under a drift, which involves a Rander metric, see notebook [Rander](Rander.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "R = np.linalg.norm(X,axis=0)\n",
    "driftMult = 0.9*np.sin(2*np.pi*X[0])*np.sin(2.*np.pi*X[1])\n",
    "drift = (driftMult/R) * X\n",
    "zermelo = Metrics.Rander.from_Zermelo(np.eye(2),drift)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Independent perturbations are applied to the symmetric and antisymmetric part of the metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta = ad.Dense.identity(shape=(3,))\n",
    "zermelo_ad = Metrics.Rander(zermelo.m *( 1 + delta[0]*(X[0]>2*X[1]) + delta[1]*(X[0]<=2*X[1]) ), \n",
    "                            zermelo.w*(1+delta[2]) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.006755 s.\n",
      "--- HFM call triggered above to compute geodesic flow ---\n",
      "Filling cache data\n",
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "hfmIn['metric']=zermelo_ad\n",
    "hfmOut,values = hfmIn.Run()\n",
    "values_zermelo_forward = values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We rely on Euler's identity for validation, which implies here that a weighted sum of the chosen perturbations equals the distance map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,7))\n",
    "plt.subplot(2,3,1)\n",
    "plt.title('Distance map')\n",
    "plt.contourf(*X,values.value)\n",
    "plt.subplot(2,3,2)\n",
    "plt.title('Weighted sum of perturbations')\n",
    "plt.contourf(*X,2*values.gradient(0)+2*values.gradient(1)+values.gradient(2));\n",
    "plt.subplot(2,3,4)\n",
    "plt.title('First perturbation')\n",
    "plt.contourf(*X,values.gradient(0))\n",
    "plt.subplot(2,3,5)\n",
    "plt.title('Second perturbation')\n",
    "plt.contourf(*X,values.gradient(1))\n",
    "plt.subplot(2,3,6)\n",
    "plt.title('Third perturbation')\n",
    "plt.contourf(*X,values.gradient(2));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Seismic example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "seismicIn = Eikonal.dictIn({\n",
    "    'model':'Seismic2',\n",
    "    'arrayOrdering':'RowMajor',\n",
    "    'extractValues':1,\n",
    "    'seeds':[[0.,0.]],\n",
    "    'exportGeodesicFlow':1.,\n",
    "})\n",
    "\n",
    "# Define the domain, get the coordinate grid\n",
    "seismicIn.SetRect(sides=[[-2,2],[-1,1]],dimx=201)\n",
    "X,Z = seismicIn.Grid()\n",
    "\n",
    "# Make small oscillations in the tilting angle\n",
    "delta = ad.Dense.identity(shape=(3,))\n",
    "\n",
    "seismic_ad = Hooke.mica[0].extract_xz()\n",
    "seismic_ad.hooke = fd.as_field(seismic_ad.hooke,X.shape) * (1+delta[0]*(X>=2*Z)+delta[1]*(X<2*Z))\n",
    "\n",
    "tiltAngle = (np.pi/6.)*np.sin(np.pi*X+1.\n",
    "    + delta[0]*(X>=0) +delta[1]*(X<0) - delta[2])\n",
    "seismic_ad.rotate_by(tiltAngle)\n",
    "\n",
    "seismicIn['metric']=seismic_ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.071829 s.\n",
      "--- HFM call triggered above to compute geodesic flow ---\n",
      "Filling cache data\n",
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "hfmOut,values = seismicIn.Run()\n",
    "values_seismic_forward=values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,5))\n",
    "plt.subplot(2,3,1)\n",
    "plt.title('Distance map')\n",
    "plt.contourf(X,Z,values.value)\n",
    "plt.subplot(2,3,2)\n",
    "plt.title('Weighted sum of perturbations')\n",
    "values_wsum = -2*(values.gradient(0)+values.gradient(1)+values.gradient(2))\n",
    "plt.contourf(X,Z,values_wsum);\n",
    "plt.subplot(2,3,4)\n",
    "plt.title('First perturbation')\n",
    "plt.contourf(X,Z,values.gradient(0))\n",
    "plt.subplot(2,3,5)\n",
    "plt.title('Second perturbation')\n",
    "plt.contourf(X,Z,values.gradient(1))\n",
    "plt.subplot(2,3,6)\n",
    "plt.title('Third perturbation')\n",
    "plt.contourf(X,Z,values.gradient(2));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert norm_infinity(values.value-values_wsum)<1e-8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Reverse differentiation \n",
    "\n",
    "We reproduce and validate the previous examples using reverse automatic differentiation, which allows the efficient computation of the gradient of an objective function defined in terms of the output values of fast marching method, w.r.t. each component of the input metric."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Isotropic example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost = np.ones(hfmIn['dims'].astype(int))\n",
    "rev,cost_ad = ad.Reverse.empty(inputs=cost,input_iterables = (Eikonal.dictIn,Metrics.Base))\n",
    "metric_ad =  Metrics.Isotropic(cost_ad)\n",
    "hfmIn['metric'] = metric_ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.00449 s.\n",
      "Filling cache data\n"
     ]
    }
   ],
   "source": [
    "hfmOut,values_ad = rev.apply(Eikonal.dictIn.Run,hfmIn,cache=Eikonal.Cache())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "X=hfmIn.Grid()\n",
    "interp = UniformGridInterpolation(X,values_ad)\n",
    "x0 = (0.2,0.8); x1 = (-0.5,-0.9); alpha0 = 1.2; alpha1 = 0.7\n",
    "objective = alpha0*interp(x0)+alpha1*interp(x1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "spAD(array(1.72040888), array([0.648  , 0.432  , 0.072  , 0.048  , 0.07875, 0.09625, 0.23625,\n",
       "       0.28875]), array([-6151, -6152, -6252, -6253, -2429, -2430, -2530, -2531]))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "objective"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "grad, = rev.to_inputshapes(rev.gradient(objective))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The support of the gradient approximates the geodesics, which in the case of a constant metric are straight lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title(\"The support of the gradient\")\n",
    "plt.contourf(*X,grad);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Rander example\n",
    "\n",
    "We reproduce Zermelo's navigation problem presented in the first section, and cross validate the forward and reverse automatic differentiation methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "rev,zermelo_rev = ad.Reverse.empty(inputs=zermelo,input_iterables = (Eikonal.dictIn,Metrics.Base))\n",
    "hfmIn['metric'] = zermelo_rev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.006672 s.\n",
      "Filling cache data\n"
     ]
    }
   ],
   "source": [
    "hfmOut,values_ad = rev.apply(Eikonal.dictIn.Run,hfmIn,cache=Eikonal.Cache())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "X=hfmIn.Grid()\n",
    "interp = UniformGridInterpolation(X,values_ad)\n",
    "objective = alpha0*interp(x0)+alpha1*interp(x1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "grad_m,grad_w = rev.to_inputshapes(rev.gradient(objective))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We validate the results by comparing forward and reverse automatic differentiation, of the above objective function, with the above variations in the Rander metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient of the objective function w.r.t. three variations in the metric : [ 0.61552158  0.6992123  -0.82402924]\n"
     ]
    }
   ],
   "source": [
    "sensitivity_reverse = np.array(\n",
    "    [ (grad_m*zermelo_ad.m.gradient(i)).sum() + (grad_w*zermelo_ad.w.gradient(i)).sum() for i in range(3)])\n",
    "print(f\"Gradient of the objective function w.r.t. three variations in the metric : {sensitivity_reverse}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "interp = UniformGridInterpolation(X,values_zermelo_forward)\n",
    "sensitivity_forward = (alpha0*interp(x0)+alpha1*interp(x1)).coef"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert ad.Optimization.norm_infinity(sensitivity_reverse-sensitivity_forward) < 1e-15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Seismic example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "seismic_rev = copy(seismic_ad)\n",
    "rev,seismic_rev.hooke = ad.Reverse.empty(inputs=seismic_ad.hooke.value, input_iterables = (Eikonal.dictIn,Metrics.Base))\n",
    "seismicIn['metric'] = seismic_rev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.071248 s.\n",
      "Filling cache data\n"
     ]
    }
   ],
   "source": [
    "hfmOut,values_rev = rev.apply(Eikonal.dictIn.Run,seismicIn,cache=Eikonal.Cache())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "X,Z = seismicIn.Grid()\n",
    "interp = UniformGridInterpolation(np.array((X,Z)),values_rev)\n",
    "objective = alpha0*interp(x0)+alpha1*interp(x1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "grad_hooke, = rev.to_inputshapes(rev.gradient(objective))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we validate the results by comparing forward and reverse automatic differentiation, of the above objective function, with the above variations in the Rander metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient of the objective function w.r.t. three variations in the metric [-0.04829966 -0.06817261  0.        ]\n"
     ]
    }
   ],
   "source": [
    "sensitivity_reverse = np.array([ (grad_hooke*seismic_ad.hooke.gradient(i)).sum() for i in range(3)])\n",
    "print(f\"Gradient of the objective function w.r.t. three variations in the metric {sensitivity_reverse}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "interp = UniformGridInterpolation(np.array([X,Z]),values_seismic_forward)\n",
    "sensitivity_forward = (alpha0*interp(x0)+alpha1*interp(x1)).coef"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert norm_infinity(sensitivity_reverse-sensitivity_forward) < 1e-16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Geodesic flow and gradient of the value function\n",
    "\n",
    "Semi-Lagrangian schemes provide an intrinsic, upwind estimate of the geodesic flow and gradient of the value function. We illustrate here their computation, and discuss their differentiation, with respect to perturbations of the metric, which is subject to substantial limitations due to the nature of the numerical scheme.\n",
    "\n",
    "Note that there are alternative approaches to computing the gradient of the value function, which are simpler and may be adequate in applications, such as centered finite differences along the coordinate axes. The interest of the upwind gradient should be more stable is that it should be the most stable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Computing the gradient\n",
    "\n",
    "The HFM library provides a numerical approximation of geodesic flow, which can be used to reconstruct a gradient. For that purpose, one relies on norm duality, thanks to the formula:\n",
    "$$\n",
    "    \\nabla u(x) = \\nabla F_x( V(x) )\n",
    "$$\n",
    "Surprisingly, when the approximate geodesic flow is used in place of $V(x)$, then the resulting approximation of $\\nabla u(x)$ can be interpreted in terms of finite differences, in the causal stencil used by the fast marching method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "hfmIn = Eikonal.dictIn({\n",
    "    'model':'Rander2',\n",
    "    'metric':zermelo,\n",
    "    'seeds':[[0,0]],\n",
    "    'exportGeodesicFlow':True,\n",
    "    'extractValues':True,\n",
    "})\n",
    "hfmIn.SetRect(sides=[[-1,1],[-1,1]],dimx=101)\n",
    "X = hfmIn.Grid()\n",
    "h = hfmIn['gridScale']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.00658 s.\n"
     ]
    }
   ],
   "source": [
    "hfmOut,values = hfmIn.Run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the gradient using finite differences and norm duality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_fd = fd.DiffGradient(values,gridScale=h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jean-mariemirebeau/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/Notebooks_FMM/../agd/AutomaticDifferentiation/Base.py:27: RuntimeWarning: divide by zero encountered in power\n",
      "  def pow(x,n):\treturn (x**n,n*x**(n-1))\n",
      "/Users/jean-mariemirebeau/Dropbox/Programmes/GithubM1/AdaptiveGridDiscretizations/Notebooks_FMM/../agd/AutomaticDifferentiation/Dense.py:141: RuntimeWarning: invalid value encountered in multiply\n",
      "  return self.new(a,_add_dim(b)*self.coef)\n"
     ]
    }
   ],
   "source": [
    "grad_hfm = hfmIn['metric'].gradient(hfmOut['flow'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The two results are quite similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = 5\n",
    "plt.figure(figsize=(12,5))\n",
    "plt.subplot(1,2,1)\n",
    "plt.title(\"Gradient computed using centered finite differences\")\n",
    "plt.quiver(*X[:,::s,::s],*grad_fd[:,::s,::s])\n",
    "plt.subplot(1,2,2)\n",
    "plt.title(\"Gradient computed using HFM upwind finite differences\")\n",
    "plt.quiver(*X[:,::s,::s],*grad_hfm[:,::s,::s]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the centered finite differences gradient is not computed along the boundary, whereas the upwind gradient is not computed at the seed point(s)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ill defined centered gradients 400. \n",
      "Ill defined upwind gradients 1.\n"
     ]
    }
   ],
   "source": [
    "nan_fd = np.any(np.isnan(grad_fd),axis=0)\n",
    "nan_hfm = np.any(np.isnan(grad_hfm),axis=0)\n",
    "\n",
    "print(f\"Ill defined centered gradients {nan_fd.sum()}. \\n\"\n",
    "      f\"Ill defined upwind gradients {nan_hfm.sum()}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main differences between the centered and upwind gradients arise near the cut locus, and where the metric is strongly anisotropic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_diff = ad.Optimization.norm(grad_fd-grad_hfm,axis=0)\n",
    "plt.title(\"Difference of centered and upwind gradients\")\n",
    "plt.contourf(*X,grad_diff);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!---\n",
    "s = 5\n",
    "plt.figure(figsize=(12,5))\n",
    "plt.subplot(1,2,1)\n",
    "plt.title(\"Gradient computed using centered finite differences\")\n",
    "plt.contourf(*X,nan_fd)\n",
    "plt.subplot(1,2,2)\n",
    "plt.title(\"Gradient computed using HFM upwind finite differences\")\n",
    "plt.contourf(*X,nan_hfm);\n",
    "--->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Differentiating the gradient\n",
    "\n",
    "Some applications, such as travel-time tomography, require to differentiate the gradient of the value function w.r.t. variations in the cost function. \n",
    "The centered finite differences approach extends in a straightforward manner to this setting. \n",
    "\n",
    "For a technical reason, the situation is more complex with upwind finite differences, because they involve offsets defining an angular sector which often degenerates. More precisely, the semi-Lagrangian scheme solves a minimization problem over a polygonal neighborhood of each point, whose extremum may be attained in a lower dimensional facet. See introduction.\n",
    "\n",
    "In that case, the upwind gradient is well defined, but it cannot be differentiated in an upwind manner.\n",
    "\n",
    "<!---\n",
    "Semi-Lagrangian schemes can be used to compute the gradient \n",
    "\n",
    "\n",
    "The value function produced by semi-Lagrangian schemes is not differentiable at every point.\n",
    "Indeed, the optimal direction found may be within a lower dimensional simplex, in which case it a locking effect arises.\n",
    "\n",
    "There are several approaches to compute the gradient of the value function:\n",
    "- using finite differences on the value function, e.g. centered and along the principal axes.\n",
    "- using norm duality, and the fact that it is dual to the geodesic flow.\n",
    "$$\n",
    "    \\nabla u(x) = \\nabla F_x( V(x) )\n",
    "$$\n",
    "\n",
    "\n",
    "Both approaches are compatible with automatic differentiation, forward and reverse. Regarding the dual approach, the following comments are in order:\n",
    "- it is equivalent to computing an upwind gradient using the same stencil as the fast marching method.\n",
    "- compatibility with AD follows, once again, from the envelope theorem. Considering a perturbation parametrized by a small $\\epsilon$ one has \n",
    "$$\n",
    "    \\nabla u(x; \\epsilon) = \\nabla F_x( V(x; 0); \\epsilon) + o(\\epsilon)\n",
    "$$\n",
    "Note, crucially, that we can set $\\epsilon=0$ in $V$, by the envelope theorem.--->\n",
    "\n",
    "<!---\n",
    "Conversely, the derivative of the geodesic flow can be recovered from the gradient:\n",
    "$$\n",
    "V(x;\\epsilon) = \\nabla F_x^*(\\nabla u(x; \\epsilon); \\epsilon)\n",
    "$$\n",
    "--->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "hfmIn.update({'exportActiveOffsets':True,'metric':zermelo_ad})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requesting cacheable data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Fast marching solver completed in 0.006335 s.\n",
      "--- HFM call triggered above to compute geodesic flow ---\n",
      "Filling cache data\n",
      "Providing cached data\n",
      "Field verbosity defaults to 1\n",
      "Field cosAngleMin defaults to 0.5\n",
      "Field refineStencilAtWallBoundary defaults to 0\n",
      "Field order defaults to 1\n",
      "Field seedRadius defaults to 0\n",
      "Bypassing fast marching solver based on cached data.\n"
     ]
    }
   ],
   "source": [
    "hfmOut,values_ad = hfmIn.Run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_fd_ad = fd.DiffGradient(values_ad,gridScale=h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We exported the active stencils used in the fast-marching algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "offsets = hfmOut['offsets']\n",
    "du_ad = fd.DiffUpwind(values_ad,offsets,h)\n",
    "\n",
    "valid = lp.det(offsets)!=0\n",
    "inv_offsets = np.full(offsets.shape,np.nan)\n",
    "inv_offsets[:,:,valid] = lp.inverse(offsets[:,:,valid].astype(float))\n",
    "\n",
    "grad_hfm_ad = lp.dot_VA(du_ad,inv_offsets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This lets us reconstruct an upwind gradient. As discussed above, this does not apply in places where the semi-Lagrangian scheme selects an extremal point of the stencil."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Proportion of valid points 0.9322615429859817\n"
     ]
    }
   ],
   "source": [
    "print(f\"Proportion of valid points {valid.sum()/valid.size}\")\n",
    "plt.figure(figsize=(13,4))\n",
    "plt.subplot(1,3,1)\n",
    "plt.title(\"Centered gradient\")\n",
    "s=5; plt.quiver(*X[:,::s,::s],*grad_fd_ad.value[:,::s,::s]);\n",
    "plt.subplot(1,3,2)\n",
    "plt.title(\"Upwind gradient\")\n",
    "s=5; plt.quiver(*X[:,::s,::s],*grad_hfm_ad.value[:,::s,::s]);\n",
    "plt.subplot(1,3,3)\n",
    "plt.title(\"Differentiable points of the upwind gradient\")\n",
    "plt.contourf(*X,valid);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As announced, the upwind finite differences and the method based on norm duality produce the same values, up to machine precision, where they are defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_diff = grad_hfm_ad.value - grad_hfm\n",
    "grad_diff_norm = np.where(valid,ad.Optimization.norm(grad_diff,axis=0),0.)\n",
    "assert norm_infinity(grad_diff_norm)<1e-11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The value function is non-differentiable along the second diagonal, which is part of the cut locus (a.k.a. the points where the value function is not differentiable). For that reason, the centered gradient derivatives yield huge and meaningless values in this along this diagonal. In contrast, the upwind gradient is better behaver in the neighborhood of the cut locus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,8))\n",
    "for i,grad in enumerate((grad_fd_ad,grad_hfm_ad)):\n",
    "    s=5; sl = slice(None,None,s) if i==1 else slice(1,-1,s)\n",
    "    for j in range(3):\n",
    "        plt.subplot(2,3,1+3*i+j)\n",
    "        plt.title(f\"Derivative {j} of {'upwind' if i==1 else 'centered'} gradient \")\n",
    "        plt.quiver(*X[:,1:-1:s,1:-1:s],*grad.gradient(j)[:,1:-1:s,1:-1:s]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Differentiating the geodesic flow\n",
    "\n",
    "The geodesic flow can be recovered from the gradient using norm duality.\n",
    "$$\n",
    "    V(x) = \\nabla F^*_x(\\nabla u(x)),\n",
    "$$\n",
    "where $F^*$ denotes the dual metric. This approach is compatible with automatic differentiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "flow_hfm_ad = zermelo_ad.dual().gradient(grad_hfm_ad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The constant part matches the output of the HFM library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "flow_diff = flow_hfm_ad.value - hfmOut['flow']\n",
    "flow_diff_norm = np.where(valid,ad.Optimization.norm(flow_diff,axis=0),0.)\n",
    "assert norm_infinity(flow_diff_norm)<1e-11"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For validating the first order part, we dualize once more, to recover the upwind gradient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_hfm_ad2 = zermelo_ad.gradient(flow_hfm_ad)\n",
    "grad_diff_ad = grad_hfm_ad2 - grad_hfm_ad\n",
    "grad_diff_ad_norm = np.where(valid,ad.Optimization.norm(grad_diff_ad.gradient(),axis=(0,1)),0.)\n",
    "assert norm_infinity(grad_diff_ad_norm) < 1e-11"
   ]
  },
  {
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}