Labs/Lab3_Thursday/Poisson_INCG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Coefficient field inversion in an elliptic partial differential equation\n",
    "\n",
    "We consider the estimation of a coefficient in an elliptic partial\n",
    "differential equation as a model problem. Depending on the\n",
    "interpretation of the unknowns and the type of measurements, this\n",
    "model problem arises, for instance, in inversion for groundwater flow\n",
    "or heat conductivity.  It can also be interpreted as finding a\n",
    "membrane with a certain spatially varying stiffness. Let\n",
    "$\\Omega\\subset\\mathbb{R}^n$, $n\\in\\{1,2,3\\}$ be an open, bounded\n",
    "domain and consider the following problem:\n",
    "\n",
    "$$\n",
    "\\min_{m} J(m):=\\frac{1}{2}\\int_\\Omega (u-u_d)^2\\, dx + \\frac{\\gamma}{2}\\int_\\Omega |\\nabla m|^2\\,dx,\n",
    "$$\n",
    "\n",
    "where $u$ is the solution of\n",
    "\n",
    "$$\n",
    "\\begin{split}\n",
    "\\quad -\\nabla\\cdot(e^m\\nabla u) &= f \\text{ in }\\Omega,\\\\\n",
    "u &= 0 \\text{ on }\\partial\\Omega.\n",
    "\\end{split}\n",
    "$$\n",
    "\n",
    "Here $m\\in \\mathcal{M}:=\\{m\\in H^1(\\Omega) \\bigcap L^{\\infty}(\\Omega)\\}$ denotes the unknown coefficient field, $u \\in \\mathcal{V}:= H^1_0(\\Omega)$ the state variable, $u_d$ the (possibly noisy) data, $f\\in H^{-1}(\\Omega)$ a given volume force, and $\\gamma\\ge 0$ the regularization parameter.\n",
    "\n",
    "### The variational (or weak) form of the state equation:\n",
    "\n",
    "Find $u\\in \\mathcal{V}$ such that \n",
    "\n",
    "$$(e^m\\nabla u,\\nabla v) - (f,v) = 0, \\text{ for all } v\\in \\mathcal{V},$$\n",
    "\n",
    "where $\\mathcal{V}$ is the space of functions vanishing on $\\partial\\Omega$ with square integrable derivatives.\n",
    "\n",
    "Above, $(\\cdot\\,\\cdot)$ denotes the $L^2$-inner product, i.e, for scalar functions $u,v$ defined on $\\Omega$ we write \n",
    "\n",
    "$$(u,v) := \\int_\\Omega u(x) v(x) \\,dx, $$\n",
    "\n",
    "and similarly for vector functions $\\boldsymbol{u}, \\boldsymbol{v}$ defined on $\\Omega$ we write\n",
    "\n",
    "$$(\\boldsymbol{u},\\boldsymbol{v}) := \\int_\\Omega \\boldsymbol{u}(x) \\cdot \\boldsymbol{v}(x) \\,dx.$$\n",
    "\n",
    "### Gradient evaluation:\n",
    "\n",
    "The Lagrangian functional $\\mathscr{L}: \\mathcal{V} \\times \\mathcal{M} \\times \\mathcal{V} \\rightarrow \\mathbb{R}$ is given by\n",
    "\n",
    "$$\n",
    "\\mathscr{L}(u,m,p):= \\frac{1}{2}(u-u_d,u-u_d) +\n",
    "\\frac{\\gamma}{2}(\\nabla m, \\nabla m) +  (e^m\\nabla u,\\nabla p) - (f,p).\n",
    "$$\n",
    "\n",
    "Then the gradient of the cost functional $\\mathcal{J}(m)$ with respect to the parameter $m$ is\n",
    "\n",
    "$$\n",
    "    \\mathcal{G}(m)(\\tilde m) := \\gamma(\\nabla m, \\nabla \\tilde{m}) +\n",
    "     (\\tilde{m} e^m\\nabla u, \\nabla p) \\quad \\forall \\tilde{m} \\in \\mathcal{M},\n",
    "$$\n",
    "\n",
    "where $u \\in \\mathcal{V}$ is the solution of the forward problem,\n",
    "\n",
    "$$ \\mathscr{L}_p(u,m,p)(\\tilde{p})  := (e^m\\nabla u, \\nabla \\tilde{p}) - (f,\\tilde{p}) = 0\n",
    "\\quad \\forall \\tilde{p} \\in \\mathcal{V}, $$\n",
    "\n",
    "and $p \\in \\mathcal{V}$ is the solution of the adjoint problem,\n",
    "\n",
    "$$ \\mathscr{L}_u(u,m,p)(\\tilde{u}) := (e^m\\nabla p, \\nabla \\tilde{u}) + (u-u_d,\\tilde{u}) = 0\n",
    "\\quad \\forall \\tilde{u} \\in \\mathcal{V}.$$\n",
    "\n",
    "### Hessian action:\n",
    "\n",
    "To evaluate the action $\\mathcal{H}(m)(\\hat{m})$ of the Hessian is a given direction $\\hat{m}$ , we consider variations of the meta-Lagrangian functional\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathscr{L}^H(u,m,p; \\hat{u}, \\hat{m}, \\hat{p}) := & {} & {} \\\\\n",
    "{} & \\gamma(\\nabla m, \\nabla \\tilde{m}) + (\\tilde{m}e^m\\nabla u, \\nabla p) & \\text{gradient}\\\\\n",
    "{} & + (e^m\\nabla u, \\nabla \\hat{p}) - (f,\\hat{p}) & \\text{forward eq}\\\\\n",
    "{} & + (e^m\\nabla p, \\nabla \\hat{u}) + (u-u_d,\\hat{u}) & \\text{adjoint eq}.\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Then action of the Hessian is a given direction $\\hat{m}$ is\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "(\\tilde{m}, \\mathcal{H}(m)(\\hat{m}) ) & := \\mathscr{L}^H_m(u,m,p; \\hat{u}, \\hat{m}, \\hat{p})(\\tilde{m}) \\\\\n",
    "{} & =\n",
    "(\\tilde{m} e^m \\nabla \\hat{u}, \\nabla{p}) + \\gamma (\\nabla \\hat{m}, \\nabla \\tilde{m}) + (\\tilde{m} \\hat{m} e^m \\nabla u, \\nabla p) + (\\tilde{m} e^m \\nabla u, \\nabla \\hat{p}) \\quad \\forall \\tilde{m} \\in \\mathcal{M},\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where \n",
    "\n",
    "- $u\\in \\mathcal{V}$ and $p \\in \\mathcal{V}$ are the solution of the forward and adjoint problem, respectively;\n",
    "\n",
    "- $\\hat{u} \\in \\mathcal{V}$ is the solution of the incremental forward problem,\n",
    "\n",
    "$$\n",
    "\\mathscr{L}^H_p(u,m,p; \\hat{u}, \\hat{m}, \\hat{p})(\\tilde{p}) := (e^m \\nabla \\hat{u}, \\nabla \\tilde{p}) + (\\hat{m} e^m \\, \\nabla u, \\nabla \\tilde p) = 0 \\quad \\forall \\tilde{p} \\in \\mathcal{V};\n",
    "$$\n",
    "\n",
    "\n",
    "- and $\\hat{p} \\in \\mathcal{V}$ is the solution of the incremental adjoint problem,\n",
    "$$\n",
    "\\mathscr{L}^H_u(u,m,p; \\hat{u}, \\hat{m}, \\hat{p})(\\tilde{u}) := (\\hat{u}, \\tilde{u}) + (\\hat{m} e^m\\nabla p, \\nabla \\tilde{u}) + (e^m \\nabla \\tilde u, \\nabla \\hat{p}) = 0 \\quad \\forall \\tilde{u} \\in \\mathcal{V}.\n",
    "$$\n",
    "\n",
    "### Inexact Newton-CG:\n",
    "\n",
    "Written in abstract form, the Newton Method computes an update direction $\\hat{m}_k$ by solving the linear system \n",
    "\n",
    "$$\n",
    "(\\tilde{m}, \\mathcal{H}(m_k)(\\hat{m}_k) ) = -\\mathcal{G}(m_k)(\\tilde m) \\quad \\forall \\tilde{m} \\in H^1(\\Omega),\n",
    "$$\n",
    "\n",
    "where the evaluation of the gradient $\\mathcal{G}(m_k)$ involve the solution $u_k$ and $p_k$ of the forward and adjoint problem (respectively) for $m = m_k$.\n",
    "Similarly, the Hessian action $\\mathcal{H}(m_k)(\\hat{m}_k)$ requires to additional solve the incremental forward and adjoint problems.\n",
    "\n",
    "### Discrete Newton system:\n",
    "$\n",
    "\\def\\tu{\\tilde u}\n",
    "\\def\\tm{\\tilde m}\n",
    "\\def\\tp{\\tilde p}\n",
    "\\def\\hu{\\hat u}\n",
    "\\def\\hp{\\hat p}\n",
    "\\def\\hm{\\hat m}\n",
    "$\n",
    "$\n",
    "\\def\\bu{{\\bf u}}\n",
    "\\def\\bm{{\\bf m}}\n",
    "\\def\\bp{{\\bf p}}\n",
    "\\def\\btu{{\\bf \\tilde u}}\n",
    "\\def\\btm{{\\bf \\tilde m}}\n",
    "\\def\\btp{{\\bf \\tilde p}}\n",
    "\\def\\bhu{{\\bf \\hat u}}\n",
    "\\def\\bhm{{\\bf \\hat m}}\n",
    "\\def\\bhp{{\\bf \\hat p}}\n",
    "\\def\\bg{{\\bf g}}\n",
    "$\n",
    "$\n",
    "\\def\\bA{{\\bf A}}\n",
    "\\def\\bC{{\\bf C}}\n",
    "\\def\\bH{{\\bf H}}\n",
    "\\def\\bR{{\\bf R}}\n",
    "\\def\\bW{{\\bf W}}\n",
    "$\n",
    "\n",
    "Let us denote the vectors corresponding to the discretization of the functions $u_k, m_k, p_k$ by $\\bu_k, \\bm_k, \\bp_k$ and of the functions $\\hu_k, \\hm_k, \\hp_k$ by $\\bhu_k, \\bhm_k,\\bhp_k$.\n",
    "\n",
    "Then, the discretization of the above system is given by the following symmetric linear system:\n",
    "\n",
    "$$\n",
    "  \\bH_k \\, \\bhm_k = -\\bg_k.\n",
    "$$\n",
    "\n",
    "The gradient $\\bg_k$ is computed using the following three steps\n",
    "\n",
    "- Given $\\bm_k$ we solve the forward problem\n",
    "\n",
    "$$ \\bA_k \\bu_k = {\\bf f}, $$\n",
    "\n",
    "where $\\bA_k \\bu_k$ stems from the discretization $(e^{m_k}\\nabla u_k, \\nabla \\tilde{p})$, and ${\\bf f}$ stands for the discretization of the right hand side $f$.\n",
    "\n",
    "- Given $\\bm_k$ and $\\bu_k$ solve the adjoint problem\n",
    "\n",
    "$$ \\bA_k^T \\bp_k = - \\bW_{\\scriptsize\\mbox{uu}}\\,(\\bu_k-\\bu_d) $$\n",
    "\n",
    "where $\\bA_k^T \\bp_k$ stems from the discretization of $(e^{m_k}\\nabla \\tilde{u}, \\nabla p_k)$, $\\bW_{\\scriptsize\\mbox{uu}}$ is the mass matrix corresponding to the $L^2$ inner product in the state space, and $\\bu_d$ stems from the data.\n",
    "\n",
    "- Define the gradient \n",
    "\n",
    "$$ \\bg_k = \\bR \\bm_k + \\bC_k^T \\bp_k, $$\n",
    "\n",
    "where $\\bR$ is the matrix stemming from discretization of the regularization operator $\\gamma ( \\nabla \\hat{m}, \\nabla \\tilde{m})$, and $\\bC_k$ stems from discretization of the term $(\\tilde{m} e^{m_k} \\, \\nabla u_k, \\nabla p_k)$.\n",
    "\n",
    "Similarly the action of the Hessian $\\bH_k \\, \\bhm_k$ in a direction $\\bhm_k$ (by using the CG algorithm we only need the action of $\\bH_k$ to solve the Newton step) is given by\n",
    "\n",
    "- Solve the incremental forward problem\n",
    "\n",
    "$$ \\bA_k \\bhu_k = -\\bC_k \\bhm_k, $$\n",
    "\n",
    "where $\\bC_k \\bm_k$ stems from discretization of $(\\hat{m} e^{m_k} \\nabla u_k, \\nabla \\tilde p)$.\n",
    "\n",
    "- Solve the incremental adjoint problem\n",
    "\n",
    "$$ \\bA_k^T \\bhp_k = -(\\bW_{\\scriptsize\\mbox{uu}} \\bhu_k + \\bW_{\\scriptsize\\mbox{um}}\\,\\bhm_k),$$\n",
    "\n",
    "where $\\bW_{\\scriptsize\\mbox{um}}\\,\\bhm_k$ stems for the discretization of $(\\hat{m}_k e^{m_k}\\nabla p_k, \\nabla \\tilde{u})$.\n",
    "\n",
    "- Define the Hessian action\n",
    "\n",
    "$$\n",
    "  \\bH_k \\, \\bhm = \\underbrace{(\\bR + \\bW_{\\scriptsize\\mbox{mm}})}_{\\text{Hessian of the regularization}} \\bhm +\n",
    "    \\underbrace{(\\bC_k^{T}\\bA_k^{-T} (\\bW_{\\scriptsize\\mbox{uu}}\n",
    "    \\bA_k^{-1} \\bC_k - \\bW_{\\scriptsize\\mbox{um}}) -\n",
    "    \\bW_{\\scriptsize\\mbox{mu}} \\bA_k^{-1}\n",
    "    \\bC_k)}_{\\text{Hessian of the data misfit}}\\;\\bhm.\n",
    "$$\n",
    "\n",
    "### Goals:\n",
    "\n",
    "By the end of this notebook, you should be able to:\n",
    "\n",
    "- solve the forward and adjoint Poisson equations\n",
    "- understand the inverse method framework\n",
    "- visualise and understand the results\n",
    "- modify the problem and code\n",
    "\n",
    "### Mathematical tools used:\n",
    "\n",
    "- Finite element method\n",
    "- Derivation of gradiant and Hessian via the adjoint method\n",
    "- inexact Newton-CG\n",
    "- Armijo line search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfin as dl\n",
    "import numpy as np\n",
    "\n",
    "from hippylib import *\n",
    "\n",
    "import logging\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "logging.getLogger('FFC').setLevel(logging.WARNING)\n",
    "logging.getLogger('UFL').setLevel(logging.WARNING)\n",
    "dl.set_log_active(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model set up:\n",
    "\n",
    "As in the introduction, the first thing we need to do is set up the numerical model.  In this cell, we set the mesh, the finite element functions $u, m, p$ corresponding to state, parameter and adjoint variables, and the corresponding test functions and the parameters for the optimization.\n",
    "\n",
    "The true parameter ``mtrue`` is the finite element interpolant of the function\n",
    "\n",
    "$$ m_{\\rm true} = \\left\\{ \\begin{array}{l} \\ln 4 \\; \\forall \\,(x,y) \\, {\\rm s.t.}\\, \\sqrt{ (x-.5)^2 + (y-.5)^2} \\leq 0.2 \\\\ \\ln 8 \\; {\\rm otherwise}. \\end{array}\\right. $$\n",
    "\n",
    "The forcing term ``f`` and the boundary conditions ``u0`` for the forward problem are\n",
    "\n",
    "$$ f = 1 \\; \\forall {\\bf x} \\in \\Omega, \\quad u = 0 \\; \\forall {\\bf x} \\in \\partial \\Omega. $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create mesh and define function spaces\n",
    "nx = 64\n",
    "ny = 64\n",
    "mesh = dl.UnitSquareMesh(nx, ny)\n",
    "Vm = dl.FunctionSpace(mesh, 'Lagrange', 1)\n",
    "Vu = dl.FunctionSpace(mesh, 'Lagrange', 2)\n",
    "\n",
    "# The true and initial guess inverted parameter\n",
    "mtrue = dl.interpolate(dl.Expression('std::log( 8. - 4.*(pow(x[0] - 0.5,2) + pow(x[1] - 0.5,2) < pow(0.2,2) ) )', degree=5), Vm)\n",
    "\n",
    "# define function for state and adjoint\n",
    "u = dl.Function(Vu)\n",
    "m = dl.Function(Vm)\n",
    "p = dl.Function(Vu)\n",
    "\n",
    "# define Trial and Test Functions\n",
    "u_trial, m_trial, p_trial = dl.TrialFunction(Vu), dl.TrialFunction(Vm), dl.TrialFunction(Vu)\n",
    "u_test,  m_test,  p_test  = dl.TestFunction(Vu),  dl.TestFunction(Vm),  dl.TestFunction(Vu)\n",
    "\n",
    "# initialize input functions\n",
    "f = dl.Constant(1.0)\n",
    "u0 = dl.Constant(0.0)\n",
    "\n",
    "# plot\n",
    "plt.figure(figsize=(15,5))\n",
    "nb.plot(mesh,subplot_loc=121, mytitle=\"Mesh\", show_axis='on')\n",
    "nb.plot(mtrue,subplot_loc=122, mytitle=\"True parameter field\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up dirichlet boundary conditions\n",
    "def boundary(x,on_boundary):\n",
    "    return on_boundary\n",
    "\n",
    "bc_state = dl.DirichletBC(Vu, u0, boundary)\n",
    "bc_adj = dl.DirichletBC(Vu, dl.Constant(0.), boundary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up synthetic observations:\n",
    "\n",
    "- Propose a coefficient field $m_{\\rm true}$ shown above\n",
    "- The weak form of the PDE:\n",
    "\n",
    "    Find $u\\in \\mathcal{V}$ such that \n",
    "    $$\\underbrace{(e^{m_{\\rm true}}\\nabla u,\\nabla v)}_{\\; := \\; a_{\\rm true}} - \\underbrace{(f,v)}_{\\; := \\;L_{\\rm true}} = 0, \\text{ for all } v\\in \\mathcal{V}.$$\n",
    "\n",
    "- Perturb the solution: $u = u + \\eta$, where $\\eta \\sim \\mathcal{N}(0, \\sigma^2)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# noise level\n",
    "noise_level = 0.05\n",
    "\n",
    "# weak form for setting up the synthetic observations\n",
    "a_true = dl.inner(dl.exp(mtrue) * dl.grad(u_trial), dl.grad(u_test)) * dl.dx\n",
    "L_true = f * u_test * dl.dx\n",
    "\n",
    "# solve the forward/state problem to generate synthetic observations\n",
    "A_true, b_true = dl.assemble_system(a_true, L_true, bc_state)\n",
    "\n",
    "utrue = dl.Function(Vu)\n",
    "dl.solve(A_true, utrue.vector(), b_true)\n",
    "\n",
    "ud = dl.Function(Vu)\n",
    "ud.assign(utrue)\n",
    "\n",
    "# perturb state solution and create synthetic measurements ud\n",
    "# ud = u + ||u||/SNR * random.normal\n",
    "MAX = ud.vector().norm(\"linf\")\n",
    "noise = dl.Vector()\n",
    "A_true.init_vector(noise,1)\n",
    "noise.set_local( noise_level * MAX * np.random.normal(0, 1, len(ud.vector().get_local())) )\n",
    "bc_adj.apply(noise)\n",
    "\n",
    "ud.vector().axpy(1., noise)\n",
    "\n",
    "# plot\n",
    "nb.multi1_plot([utrue, ud], [\"State solution with mtrue\", \"Synthetic observations\"])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The cost function evaluation:\n",
    "\n",
    "$$\n",
    "J(m):=\\underbrace{\\frac{1}{2}\\int_\\Omega (u-u_d)^2\\, dx}_{\\text{misfit} } + \\underbrace{\\frac{\\gamma}{2}\\int_\\Omega|\\nabla m|^2\\,dx}_{\\text{reg}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# regularization parameter\n",
    "gamma = 1e-8\n",
    "\n",
    "# define cost function\n",
    "def cost(u, ud, m,gamma):\n",
    "    reg = 0.5*gamma * dl.assemble( dl.inner(dl.grad(m), dl.grad(m))*dl.dx ) \n",
    "    misfit = 0.5 * dl.assemble( (u-ud)**2*dl.dx)\n",
    "    return [reg + misfit, misfit, reg]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up the variational form for the state/adjoint equations and gradient evaluation\n",
    "\n",
    "Below we define the variational forms that appears in the the state/adjoint equations and gradient evaluations.\n",
    "\n",
    "Specifically,\n",
    "\n",
    "- `a_state`, `L_state` stand for the bilinear and linear form of the state equation, repectively;\n",
    "- `a_adj`, `L_adj` stand for the bilinear and linear form of the adjoint equation, repectively;\n",
    "- `CTvarf`, `gradRvarf` stand for the contributions to the gradient coming from the PDE and the regularization, respectively.\n",
    "\n",
    "We also build the *mass* matrix $M$ that is used to discretize the $L^2(\\Omega)$ inner product."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# weak form for setting up the state equation\n",
    "a_state = dl.inner(dl.exp(m) * dl.grad(u_trial), dl.grad(u_test)) * dl.dx\n",
    "L_state = f * u_test * dl.dx\n",
    "\n",
    "# weak form for setting up the adjoint equation\n",
    "a_adj = dl.inner(dl.exp(m) * dl.grad(p_trial), dl.grad(p_test)) * dl.dx\n",
    "L_adj = -dl.inner(u - ud, p_test) * dl.dx\n",
    "\n",
    "# weak form for gradient\n",
    "CTvarf    = dl.inner(dl.exp(m)*m_test*dl.grad(u), dl.grad(p)) * dl.dx\n",
    "gradRvarf = gamma*dl.inner(dl.grad(m), dl.grad(m_test))*dl.dx\n",
    "\n",
    "# L^2 weighted inner product\n",
    "M_varf   = dl.inner(m_trial, m_test) * dl.dx\n",
    "M = dl.assemble(M_varf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initial guess\n",
    "We solve the state equation and compute the cost functional for the initial guess of the parameter ``m0``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m0 = dl.interpolate(dl.Constant(np.log(4.) ), Vm )\n",
    "\n",
    "m.assign(m0)\n",
    "\n",
    "# solve state equation\n",
    "state_A, state_b = dl.assemble_system (a_state, L_state, bc_state)\n",
    "dl.solve (state_A, u.vector(), state_b)\n",
    "\n",
    "# evaluate cost\n",
    "[cost_old, misfit_old, reg_old] = cost(u, ud, m, gamma)\n",
    "\n",
    "# plot\n",
    "plt.figure(figsize=(15,5))\n",
    "nb.plot(m,subplot_loc=121, mytitle=\"m0\", vmin=mtrue.vector().min(), vmax=mtrue.vector().max())\n",
    "nb.plot(u,subplot_loc=122, mytitle=\"u(m0)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Variational forms for Hessian action\n",
    "\n",
    "We define the following variational forms that are needed for the Hessian evaluation\n",
    "\n",
    "- `W_varf`, `R_varf` are the second variation of the data-misfit and regularization component of the cost functional respectively (note since `W_varf`, `R_varf` are independent of $u$, $m$, $p$ they can be preassembled);\n",
    "\n",
    "- `C_varf` is the second variation of the PDE with respect to $p$ and $m$;\n",
    "\n",
    "- `Wum_varf` is the second variation of the PDE with respect to $u$ and $m$;\n",
    "\n",
    "- `Wmm_varf` is the second variation of the PDE with respect to $m$.\n",
    "\n",
    "> **Note**: Since the forward problem is linear, the bilinear forms for the incremental state and adjoint equations are the same as the bilinear forms for the state and adjoint equations, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W_varf   = dl.inner(u_trial, u_test) * dl.dx\n",
    "R_varf   = dl.Constant(gamma) * dl.inner(dl.grad(m_trial), dl.grad(m_test)) * dl.dx\n",
    "\n",
    "C_varf   = dl.inner(dl.exp(m) * m_trial * dl.grad(u), dl.grad(u_test)) * dl.dx\n",
    "Wum_varf = dl.inner(dl.exp(m) * m_trial * dl.grad(p_test), dl.grad(p)) * dl.dx\n",
    "Wmm_varf = dl.inner(dl.exp(m) * m_trial * m_test *  dl.grad(u),  dl.grad(p)) * dl.dx\n",
    "\n",
    "# Assemble constant matrices\n",
    "W = dl.assemble(W_varf)\n",
    "R = dl.assemble(R_varf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hessian action on a vector $\\bhm$:\n",
    "\n",
    "Here we describe how to apply the Hessian operator to a vector $\\bhm$. For an opportune choice of the regularization, the Hessian operator evaluated in a neighborhood of the solution is positive define, whereas far from the solution the reduced Hessian may be indefinite. On the constrary, the Gauss-Newton approximation of the Hessian is always positive defined.\n",
    "\n",
    "For this reason, it is beneficial to perform a few initial Gauss-Newton steps (5 in this particular example) to accelerate the convergence of the inexact Newton-CG algorithm.\n",
    "\n",
    "The Hessian action reads:\n",
    "$$\n",
    "\\begin{align}\n",
    "\\bhu &= -\\bA^{-1} \\bC \\bhm\\, & \\text{incremental forward}\\\\\n",
    "\\bhp &= -\\bA^{-T} (\\bW_{\\scriptsize\\mbox{uu}} \\bhu +\n",
    "\\bW_{\\scriptsize\\mbox{um}}\\,\\bhm) & \\text{incremental adjoint}\\\\\n",
    "\\bH \\bhm &= (\\bR + \\bW_{\\scriptsize\\mbox{mm}})\\bhm + \\bC^T \\bhp + \\bW_{\\scriptsize\\mbox{mu}} \\bhu.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "The Gauss-Newton Hessian action is obtained by dropping the second derivatives operators $\\bW_{\\scriptsize\\mbox{um}}\\,\\bhm$, $\\bW_{\\scriptsize\\mbox{mm}}\\bf \\bhm$, and $\\bW_{\\scriptsize\\mbox{mu}} \\bhu$:\n",
    "$$\n",
    "\\begin{align}\n",
    "\\bhu &= -\\bA^{-1} \\bC \\bf \\bhm\\, & \\text{incremental forward}\\\\\n",
    "\\bhp &= -\\bA^{-T} \\bW_{\\scriptsize\\mbox{uu}} \\bhu & \\text{incremental adjoint}\\\\\n",
    "\\bH_{\\rm GN} \\bhm &= \\bR \\bhm + \\bC^T \\bhp.\n",
    "\\end{align}\n",
    "$$\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Class HessianOperator to perform Hessian apply to a vector\n",
    "class HessianOperator():\n",
    "    cgiter = 0\n",
    "    def __init__(self, R, Wmm, C, A, adj_A, W, Wum, bc0, use_gaussnewton=False):\n",
    "        self.R = R\n",
    "        self.Wmm = Wmm\n",
    "        self.C = C\n",
    "        self.A = A\n",
    "        self.adj_A = adj_A\n",
    "        self.W = W\n",
    "        self.Wum = Wum\n",
    "        self.bc0 = bc0\n",
    "        self.use_gaussnewton = use_gaussnewton\n",
    "        \n",
    "        # incremental state\n",
    "        self.du = dl.Vector()\n",
    "        self.A.init_vector(self.du,0)\n",
    "        \n",
    "        #incremental adjoint\n",
    "        self.dp = dl.Vector()\n",
    "        self.adj_A.init_vector(self.dp,0)\n",
    "        \n",
    "        # auxiliary vector\n",
    "        self.Wum_du = dl.Vector()\n",
    "        self.Wum.init_vector(self.Wum_du, 1)\n",
    "        \n",
    "    def init_vector(self, v, dim):\n",
    "        self.R.init_vector(v,dim)\n",
    "\n",
    "    # Hessian performed on v, output as generic vector y\n",
    "    def mult(self, v, y):\n",
    "        self.cgiter += 1\n",
    "        y.zero()\n",
    "        if self.use_gaussnewton:\n",
    "            self.mult_GaussNewton(v,y)\n",
    "        else:\n",
    "            self.mult_Newton(v,y)\n",
    "            \n",
    "    # define (Gauss-Newton) Hessian apply H * v\n",
    "    def mult_GaussNewton(self, v, y):\n",
    "        \n",
    "        #incremental forward\n",
    "        rhs = -(self.C * v)\n",
    "        self.bc0.apply(rhs)\n",
    "        dl.solve (self.A, self.du, rhs)\n",
    "        \n",
    "        #incremental adjoint\n",
    "        rhs = - (self.W * self.du)\n",
    "        self.bc0.apply(rhs)\n",
    "        dl.solve (self.adj_A, self.dp, rhs)\n",
    "                \n",
    "        # Misfit term\n",
    "        self.C.transpmult(self.dp, y)\n",
    "        \n",
    "        if self.R:\n",
    "            Rv = self.R*v\n",
    "            y.axpy(1, Rv)\n",
    "        \n",
    "    # define (Newton) Hessian apply H * v\n",
    "    def mult_Newton(self, v, y):\n",
    "        \n",
    "        #incremental forward\n",
    "        rhs = -(self.C * v)\n",
    "        self.bc0.apply(rhs)\n",
    "        dl.solve (self.A, self.du, rhs)\n",
    "        \n",
    "        #incremental adjoint\n",
    "        rhs = -(self.W * self.du) -  self.Wum * v\n",
    "        self.bc0.apply(rhs)\n",
    "        dl.solve (self.adj_A, self.dp, rhs)\n",
    "                \n",
    "        #Misfit term\n",
    "        self.C.transpmult(self.dp, y)\n",
    "        \n",
    "        self.Wum.transpmult(self.du, self.Wum_du)\n",
    "        y.axpy(1., self.Wum_du)\n",
    "        \n",
    "        y.axpy(1., self.Wmm*v)\n",
    "        \n",
    "        #Reg/Prior term\n",
    "        if self.R:\n",
    "            y.axpy(1., self.R*v)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The inexact Newton-CG optimization with Armijo line search:\n",
    "\n",
    "We solve the constrained optimization problem using the inexact Newton-CG method with Armijo line search.\n",
    "\n",
    "The stopping criterion is based on a relative reduction of the norm of the gradient (i.e. $\\frac{\\|g_{n}\\|}{\\|g_{0}\\|} \\leq \\tau$).\n",
    "\n",
    "First, we compute the gradient by solving the state and adjoint equation for the current parameter $m$, and then substituing the current state $u$, parameter $m$ and adjoint $p$ variables in the weak form expression of the gradient:\n",
    "$$ (g, \\tilde{m}) = \\gamma(\\nabla m, \\nabla \\tilde{m}) +(\\tilde{m}\\nabla u, \\nabla p).$$\n",
    "\n",
    "Then, we compute the Newton direction $\\hat m$ by iteratively solving $\\mathcal{H} {\\hat m} = -g$.\n",
    "The Newton system is solved inexactly by early termination of conjugate gradient iterations via Eisenstat–Walker (to prevent oversolving) and Steihaug  (to avoid negative curvature) criteria. \n",
    "\n",
    "> Usually, one uses the regularization matrix $R$ as preconditioner for the Hessian system, however since $R$ is singular (the constant vector is in the null space of $R$), here we use $P = R + \\frac{\\gamma}{10} M$, where $M$ is the mass matrix in parameter space.\n",
    "\n",
    "Finally, the Armijo line search uses backtracking to find $\\alpha$ such that a sufficient reduction in the cost functional is achieved.\n",
    "More specifically, we use backtracking to find $\\alpha$ such that:\n",
    "$$J( m + \\alpha \\hat m ) \\leq J(m) + \\alpha c_{\\rm armijo} (\\hat m,g). $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define parameters for the optimization\n",
    "tol = 1e-8\n",
    "c = 1e-4\n",
    "maxiter = 12\n",
    "plot_on = False\n",
    "\n",
    "# initialize iter counters\n",
    "iter = 1\n",
    "total_cg_iter = 0\n",
    "converged = False\n",
    "\n",
    "# initializations\n",
    "g, m_delta = dl.Vector(), dl.Vector()\n",
    "R.init_vector(m_delta,0)\n",
    "R.init_vector(g,0)\n",
    "\n",
    "m_prev = dl.Function(Vm)\n",
    "\n",
    "print( \"Nit   CGit   cost          misfit        reg           sqrt(-G*D)    ||grad||       alpha  tolcg\" )\n",
    "\n",
    "while iter <  maxiter and not converged:\n",
    "\n",
    "    # solve the adoint problem\n",
    "    adjoint_A, adjoint_RHS = dl.assemble_system(a_adj, L_adj, bc_adj)\n",
    "    dl.solve(adjoint_A, p.vector(), adjoint_RHS)\n",
    "\n",
    "    # evaluate the  gradient\n",
    "    MG = dl.assemble(CTvarf + gradRvarf)\n",
    "\n",
    "    # calculate the L^2 norm of the gradient\n",
    "    dl.solve(M, g, MG)\n",
    "    grad2 = g.inner(MG)\n",
    "    gradnorm = np.sqrt(grad2)\n",
    "\n",
    "    # set the CG tolerance (use Eisenstat–Walker termination criterion)\n",
    "    if iter == 1:\n",
    "        gradnorm_ini = gradnorm\n",
    "    tolcg = min(0.5, np.sqrt(gradnorm/gradnorm_ini))\n",
    "    \n",
    "    # assemble W_um and W_mm\n",
    "    C   = dl.assemble(C_varf)\n",
    "    Wum = dl.assemble(Wum_varf)\n",
    "    Wmm = dl.assemble(Wmm_varf)\n",
    "\n",
    "    # define the Hessian apply operator (with preconditioner)\n",
    "    Hess_Apply = HessianOperator(R, Wmm, C, state_A, adjoint_A, W, Wum, bc_adj, use_gaussnewton=(iter<6) )\n",
    "    P = R + 0.1*gamma * M\n",
    "    Psolver = dl.PETScKrylovSolver(\"cg\", amg_method())\n",
    "    Psolver.set_operator(P)\n",
    "    \n",
    "    solver = CGSolverSteihaug()\n",
    "    solver.set_operator(Hess_Apply)\n",
    "    solver.set_preconditioner(Psolver)\n",
    "    solver.parameters[\"rel_tolerance\"] = tolcg\n",
    "    solver.parameters[\"zero_initial_guess\"] = True\n",
    "    solver.parameters[\"print_level\"] = -1\n",
    "\n",
    "    # solve the Newton system H a_delta = - MG\n",
    "    solver.solve(m_delta, -MG)\n",
    "    total_cg_iter += Hess_Apply.cgiter\n",
    "    \n",
    "    # linesearch\n",
    "    alpha = 1\n",
    "    descent = 0\n",
    "    no_backtrack = 0\n",
    "    m_prev.assign(m)\n",
    "    while descent == 0 and no_backtrack < 10:\n",
    "        m.vector().axpy(alpha, m_delta )\n",
    "\n",
    "        # solve the state/forward problem\n",
    "        state_A, state_b = dl.assemble_system(a_state, L_state, bc_state)\n",
    "        dl.solve(state_A, u.vector(), state_b)\n",
    "\n",
    "        # evaluate cost\n",
    "        [cost_new, misfit_new, reg_new] = cost(u, ud, m, gamma)\n",
    "\n",
    "        # check if Armijo conditions are satisfied\n",
    "        if cost_new < cost_old + alpha * c * MG.inner(m_delta):\n",
    "            cost_old = cost_new\n",
    "            descent = 1\n",
    "        else:\n",
    "            no_backtrack += 1\n",
    "            alpha *= 0.5\n",
    "            m.assign(m_prev)  # reset a\n",
    "\n",
    "    # calculate sqrt(-G * D)\n",
    "    graddir = np.sqrt(- MG.inner(m_delta) )\n",
    "\n",
    "    sp = \"\"\n",
    "    print( \"%2d %2s %2d %3s %8.5e %1s %8.5e %1s %8.5e %1s %8.5e %1s %8.5e %1s %5.2f %1s %5.3e\" % \\\n",
    "        (iter, sp, Hess_Apply.cgiter, sp, cost_new, sp, misfit_new, sp, reg_new, sp, \\\n",
    "         graddir, sp, gradnorm, sp, alpha, sp, tolcg) )\n",
    "\n",
    "    if plot_on:\n",
    "        nb.multi1_plot([m,u,p], [\"m\",\"u\",\"p\"], same_colorbar=False)\n",
    "        plt.show()\n",
    "    \n",
    "    # check for convergence\n",
    "    if gradnorm < tol and iter > 1:\n",
    "        converged = True\n",
    "        print( \"Newton's method converged in \",iter,\"  iterations\" )\n",
    "        print( \"Total number of CG iterations: \", total_cg_iter )\n",
    "        \n",
    "    iter += 1\n",
    "    \n",
    "if not converged:\n",
    "    print( \"Newton's method did not converge in \", maxiter, \" iterations\" )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nb.multi1_plot([mtrue, m], [\"mtrue\", \"m\"])\n",
    "nb.multi1_plot([u,p], [\"u\",\"p\"], same_colorbar=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The generalized eigenvalues and eigenvectors of the Hessian misfit  \n",
    "\n",
    "We used the *double pass* randomized algorithm to compute the generalized eigenvalues and eigenvectors of the Hessian misfit.\n",
    "In particular, we solve\n",
    "\n",
    "$$ H_{\\rm misfit} \\hat{\\bf m}_i = \\lambda_i R \\hat{\\bf v}_i. $$\n",
    "\n",
    "The Figure shows the largest *k* generalized eigenvectors of the Hessian misfit.\n",
    "The effective rank of the Hessian misfit is the number of eigenvalues above the red line ($y=1$).\n",
    "The effective rank is independent of the mesh size.\n",
    "\n",
    "> **Note**: Since $R$ is singular (the constant are in the null space of $R$), we will add a small mass matrix $M$ to $R$ and use $P = R + \\frac{\\gamma}{10}M$ instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Hmisfit = HessianOperator(None, Wmm, C, state_A, adjoint_A, W, Wum, bc_adj, use_gaussnewton=False)\n",
    "k = 50\n",
    "p = 20\n",
    "print( \"Double Pass Algorithm. Requested eigenvectors: {0}; Oversampling {1}.\".format(k,p) )\n",
    "\n",
    "Omega = MultiVector(m.vector(), k+p)\n",
    "parRandom.normal(1., Omega)\n",
    "lmbda, evecs = doublePassG(Hmisfit, P, Psolver, Omega, k)\n",
    "\n",
    "plt.plot(range(0,k), lmbda, 'b*', range(0,k+1), np.ones(k+1), '-r')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('number')\n",
    "plt.ylabel('eigenvalue')\n",
    "\n",
    "nb.plot_eigenvectors(Vm, evecs, mytitle=\"Eigenvector\", which=[0,1,2,5,10,15])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hands on\n",
    "\n",
    "### Question 1\n",
    "\n",
    "> Report the number of inexact Newton and of total CG iterations for a discretization of the domain with $8 \\times 8$, $16 \\times 16$, $32 \\times 32$, $64 \\times 64$ finite elements and give the number of unknowns used to discretize the log diffusivity field m for each of these meshes. Discuss how the number of iterations changes as the inversion parameter mesh is refined, i.e., as the parameter dimension increases. Is inexact Newton-CG method scalable with respect to the parameter dimension?\n",
    "\n",
    "### Question 2\n",
    "\n",
    "> Compute the generalized eigenvalues and eigenvectors of the Hessian misfit at the solution of the inverse problem \n",
    "for a discretization of the domain with $8 \\times 8$, $16 \\times 16$, $32 \\times 32$, $64 \\times 64$ finite elements.\n",
    "What do you observe?\n",
    "\n",
    "### Question 3\n",
    "\n",
    "> Add the advective term $\\mathbf{v} = [30,0]^t$ to the inverse problem and its hIPPYlib/FEniCS implementation and plot the resulting reconstruction of $m$ for a noise level of 0.01 and for a reasonably chosen regularization parameter. The *optimal* regularization parameter can be found manually (i.e., by experimenting with a few different values and finding the one that results in a reconstruction that best matches the *true* log diffusivity field), or else by the discrepancy principle, if you are so inclined."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright &copy; 2016-2018, The University of Texas at Austin & University of California, Merced.\n",
    "All Rights reserved.\n",
    "See file COPYRIGHT for details.\n",
    "\n",
    "This file is part of the hIPPYlib library. For more information and source code\n",
    "availability see https://hippylib.github.io.\n",
    "\n",
    "hIPPYlib is free software; you can redistribute it and/or modify it under the terms of the GNU General Public License (as published by the Free Software Foundation) version 2.0 dated June 1991."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}