PoissonDeterministic.ipynb

{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "# Example: 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 first 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_{a} J(a):=\\frac{1}{2}\\int_\\Omega (u-u_d)^2\\, dx   \\frac{\\gamma}{2}\\int_\\Omega|\\nabla a|^2\\,dx,\n",
      "$$\n",
      "\n",
      "where $u$ is the solution of\n",
      "\n",
      "$$\n",
      "\\begin{split}\n",
      "\\quad -\\nabla\\cdot(\\exp(a)\\nabla u) &= f \\text{ in }\\Omega,\\\\\n",
      "u &= 0 \\text{ on }\\partial\\Omega.\n",
      "\\end{split}\n",
      "$$\n",
      "\n",
      "Here $a\\in U_{ad}:=\\{a\\in L^{\\infty}(\\Omega)\\}$ the unknown coefficient field, $u_d$ denotes (possibly noisy) data, $f\\in H^{-1}(\\Omega)$ a given force, and $\\gamma\\ge 0$ the regularization parameter.\n",
      "\n",
      "### The variational (or weak) form of the state equation:\n",
      "\n",
      "Find $u\\in H_0^1(\\Omega)$ such that $(\\exp(a)\\nabla u,\\nabla v) - (f,v) = 0, \\text{ for all } v\\in H_0^1(\\Omega),$\n",
      "where $H_0^1(\\Omega)$ is the space of functions vanishing on $\\partial\\Omega$ with square integrable derivatives. Here, $(\\cdot\\,,\\cdot)$ denotes the $L^2$-inner product, i.e, for scalar functions $u,v$ defined on $\\Omega$ we denote $(u,v) := \\int_\\Omega u(x) v(x) \\,dx$.\n",
      "\n",
      "### Optimality System:\n",
      "\n",
      "The Lagrangian functional $\\mathscr{L}:L^\\infty(\\Omega)\\times H_0^1(\\Omega)\\times H_0^1(\\Omega)\\rightarrow \\mathbb{R}$, which we use as a tool to derive the optimality system, is given by\n",
      "\n",
      "$$\n",
      "\\mathscr{L}(a,u,p):= \\frac{1}{2}(u-u_d,u-u_d)  \n",
      "\\frac{\\gamma}{2}(\\nabla a, \\nabla a)    (\\exp(a)\\nabla u,\\nabla p) - (f,p).\n",
      "$$\n",
      "\n",
      "The Lagrange multiplier theory shows that, at a solution all variations of the Lagrangian functional with respect to all variables must vanish. These variations of $\\mathscr{L}$ with respect to $(p,u,a)$ in directions $(\\tilde{u}, \\tilde{p}, \\tilde{a})$ are given by\n",
      "\n",
      "$$\n",
      "  \\begin{alignat}{2}\n",
      "    \\mathscr{L}_p(a,u,p)(\\tilde{p})  &= (\\exp(a)\\nabla u, \\nabla \\tilde{p}) -\n",
      "    (f,\\tilde{p}) &&= 0,\\\\\n",
      "     \\mathscr{L}_u(a,u,p)(\\tilde{u}) &= (\\exp(a)\\nabla p, \\nabla \\tilde{u})  \n",
      "     (u-u_d,\\tilde{u}) && = 0,\\\\\n",
      "     \\mathscr{L}_a(a,u,p)(\\tilde{a})  &= \\gamma(\\nabla a, \\nabla \\tilde{a})  \n",
      "     (\\tilde{a}\\exp(a)\\nabla u, \\nabla p) &&= 0,\n",
      "  \\end{alignat}\n",
      "$$\n",
      "\n",
      "where the variations $(\\tilde{u}, \\tilde{p}, \\tilde{a})$ are taken from the same spaces as $(u,p,a)$. \n",
      "\n",
      "The gradient of the cost functional $\\mathcal{J}(a)$ therefore is\n",
      "\n",
      "$$\n",
      "    \\mathcal{G}(a)(\\tilde a) = \\gamma(\\nabla a, \\nabla \\tilde{a})  \n",
      "     (\\tilde{a}\\exp(a)\\nabla u, \\nabla \\tilde{p}).\n",
      "$$\n",
      "\n",
      "### Inexact Newton-CG:\n",
      "\n",
      "Newton's method requires second-order variational derivatives of the Lagrangian . Written in abstract form, it computes an update direction $(\\hat a_k, \\hat u_k,\\hat p_k)$ from the following Newton step for the Lagrangian functional:\n",
      "\n",
      "$$\n",
      "\\mathscr{L}''(a_k, u_k, p_k)\\left[(\\tilde\n",
      "  a, \\tilde u, \\tilde p),(\\hat a_k, \\hat u_k, \\hat p_k)\\right] =\n",
      "-\\mathscr{L}'(a_k,u_k,p_k)(\\tilde a, \\tilde u, \\tilde p),\n",
      "$$\n",
      "\n",
      "for all variations $(\\tilde a, \\tilde u, \\tilde p)$, where $\\mathscr{L}'$ and $\\mathscr{L}''$ denote the first and\n",
      "second variations of the Lagrangian. For the elliptic parameter inversion problem, this Newton step (written in variatonal form) is as follows: Find $(\\hat u_k, \\hat a_k,\\hat p_k)$ as the solution of the linear system\n",
      "\n",
      "$$\n",
      "  \\begin{array}{llll}\n",
      "    (\\hat{u}_k, \\tilde u) &  (\\hat{a}_k \\exp(a_k)\\nabla p_k, \\nabla\n",
      "    \\tilde u) &  (\\exp(a_k) \\nabla \\tilde u,\n",
      "    \\nabla \\hat p_k) &= (u_d - u_k, \\tilde u)- (\\exp(a_k) \\nabla\n",
      "    p_k, \\nabla \\tilde u)\\\\\n",
      "    (\\tilde a \\exp(a_k) \\nabla \\hat u_k, \\nabla p_k) &  \\gamma\n",
      "    (\\nabla \\hat a_k, \\nabla \\tilde a)   (\\tilde a \\hat a_k \\exp(a_k)\\nabla u, \\nabla p) &  (\\tilde a\n",
      "     \\exp(a_k) \\nabla u_k, \\nabla \\hat p_k) &= - \\gamma (\\nabla a_k, \\nabla\\tilde a) - (\\tilde\n",
      "      a  \\exp(a_k) \\nabla u_k, \\nabla p_k)\\\\\n",
      "    (\\exp(a_k) \\nabla \\hat u_k, \\nabla \\tilde p) &  (\\hat a_k \\exp(a_k) \\nabla u_k, \\nabla\n",
      "      \\tilde p) & &= - (\\exp(a_k) \\nabla u_k,\n",
      "    \\nabla \\tilde p)   (f, \\tilde p),\n",
      "  \\end{array}\n",
      "$$\n",
      "\n",
      "for all $(\\tilde u, \\tilde a, \\tilde p)$.\n",
      "\n",
      "### Discrete Newton system:\n",
      "$\n",
      "\\def\\tu{\\tilde u}\n",
      "\\def\\btu{\\bf \\tilde u}\n",
      "\\def\\ta{\\tilde a}\n",
      "\\def\\bta{\\bf \\tilde a}\n",
      "\\def\\tp{\\tilde p}\n",
      "\\def\\btp{\\bf \\tilde p}\n",
      "\\def\\hu{\\hat u}\n",
      "\\def\\bhu{\\bf \\hat u}\n",
      "\\def\\ha{\\hat a}\n",
      "\\def\\bha{\\bf \\hat a}\n",
      "\\def\\hp{\\hat p}\n",
      "\\def\\bhp{\\bf \\hat p}\n",
      "$\n",
      "The discretized Newton step: denote the vectors corresponding to the discretization of the functions $\\ha_k,\\hu_k, \\hp_k$ by $\\bf \\bha_k, \\bhu_k$ and $\\bhp_k$. Then, the discretization of the above system is given by the following symmetric linear system:\n",
      "\n",
      "$$\n",
      "  \\begin{bmatrix}\n",
      "    \\bf W_{\\scriptsize\\mbox{uu}} & \\bf W_{\\scriptsize\\mbox{ua}} & \\bf A^T \\\\\n",
      "    \\bf W_{\\scriptsize\\mbox{au}} & \\bf R   \\bf R_{\\scriptsize\\mbox{aa}}& \\bf C^T \\\\\n",
      "    \\bf A & \\bf C & 0\n",
      "\\end{bmatrix}\n",
      "\\left[\n",
      "  \\begin{array}{c}\n",
      "    \\bhu_k \\\\\n",
      "    \\bha_k \\\\\n",
      "    \\bhp_k\n",
      "  \\end{array} \\right] =\n",
      "-\\left[\n",
      "  \\begin{array}{ccc}\n",
      "    \\bf{g}_u\\\\\n",
      "    \\bf{g}_a\\\\\n",
      "    \\bf{g}_p\n",
      "\\end{array}\n",
      "  \\right],\n",
      "$$\n",
      "\n",
      "where $\\bf W_{\\scriptsize \\mbox{uu}}$, $\\bf W_{\\scriptsize\\mbox{ua}}$, $\\bf W_{\\scriptsize\\mbox{au}}$, and $\\bf R$ are the components of the Hessian matrix of the Lagrangian, $\\bf A$ and $\\bf C$ are the Jacobian of the state equation with respect to the state and the control variables, respectively and $\\bf g_u$, $\\bf g_a$, and $\\bf g_p$ are the discrete gradients of the Lagrangian with respect to $\\bf u $, $\\bf a$ and $\\bf p$, respectively.\n",
      "\n",
      "### Reduced Hessian apply:\n",
      "\n",
      "To eliminate the incremental state and adjoint variables, $\\bhu_k$ and $\\bhp_k$, from the first and last equations we use\n",
      "\n",
      "$$\n",
      "\\begin{align}\n",
      "\\bhu_k &= -\\bf A^{-1} \\bf C \\, \\bha_k,\\\\\n",
      "\\bhp_k &= -\\bf A^{-T} (\\bf W_{\\scriptsize\\mbox{uu}} \\bhu_k  \n",
      "\\bf W_{\\scriptsize\\mbox{ua}}\\,\\bha_k).\n",
      "\\end{align}\n",
      "$$\n",
      "\n",
      "This results in the following reduced linear system for the Newton step\n",
      "\n",
      "$$\n",
      "  \\bf H \\, \\bha_k = -\\bf{g}_a,\n",
      "$$\n",
      "\n",
      "with the reduced Hessian $\\bf H$ applied to a vector $\\bha$ given by\n",
      "\n",
      "$$\n",
      "  \\bf H \\bha = \\underbrace{(\\bf R   \\bf R_{\\scriptsize\\mbox{aa}})}_{\\text{Hessian of the regularization}} \\bha  \n",
      "    \\underbrace{(\\bf C^{T}\\bf A^{-T} (\\bf W_{\\scriptsize\\mbox{uu}}\n",
      "    \\bf A^{-1} \\bf C - \\bf W_{\\scriptsize\\mbox{ua}}) -\n",
      "    \\bf W_{\\scriptsize\\mbox{au}} \\bf A^{-1}\n",
      "    \\bf C)}_{\\text{Hessian of the data misfit}}\\;\\bha.\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\n",
      "\n",
      "### List of software used:\n",
      "\n",
      "- <a href=\"http://fenicsproject.org/\">FEniCS</a>, a parallel finite element element library for the discretization of partial differential equations\n",
      "- <a href=\"http://www.mcs.anl.gov/petsc/\">PETSc</a>, for scalable and efficient linear algebra operations and solvers\n",
      "- <a href=\"http://matplotlib.org/\">Matplotlib</a>, a python package used for plotting the results\n",
      "- <a href=\"http://www.numpy.org/\">Numpy</a>, a python package for linear algebra"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Set up\n",
      "\n",
      "### Import dependencies"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from dolfin import *\n",
      "\n",
      "import numpy as np\n",
      "import time\n",
      "import logging\n",
      "\n",
      "import matplotlib.pyplot as plt\n",
      "%matplotlib inline\n",
      "import nb\n",
      "\n",
      "start = time.clock()\n",
      "\n",
      "logging.getLogger('FFC').setLevel(logging.WARNING)\n",
      "logging.getLogger('UFL').setLevel(logging.WARNING)\n",
      "set_log_active(False)\n",
      "\n",
      "np.random.seed(seed=1)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "### The cost function evaluation:\n",
      "\n",
      "$$\n",
      "J(a):=\\underbrace{\\frac{1}{2}\\int_\\Omega (u-u_d)^2\\, dx}_{\\text misfit}   \\underbrace{\\frac{\\gamma}{2}\\int_\\Omega|\\nabla a|^2\\,dx}_{\\text reg}\n",
      "$$"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Define cost function\n",
      "def cost(u, ud, a, W, R):\n",
      "    diff = u.vector() - ud.vector()\n",
      "    reg = 0.5 * a.vector().inner(R*a.vector() ) \n",
      "    misfit = 0.5 * diff.inner(W * diff)\n",
      "    return [reg   misfit, misfit, reg]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "### The reduced Hessian apply to a vector v:\n",
      "\n",
      "$$\n",
      "\\begin{align}\n",
      "\\bhu &= -\\bf A^{-1} \\bf C \\, & \\text{linearized forward}\\\\\n",
      "\\bhp &= -\\bf A^{-T} (\\bf W_{\\scriptsize\\mbox{uu}} \\bhu  \n",
      "\\bf W_{\\scriptsize\\mbox{ua}}\\,\\bha) & \\text{adjoint}\\\\\n",
      "\\bf H \\bf v &= (\\bf R   \\bf R_{\\scriptsize\\mbox{aa}})\\bf v   \\bf C^T \\bhp   \\bf W_{\\scriptsize\\mbox{au}} \\bhu.\n",
      "\\end{align}\n",
      "$$"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# define (Gauss-Newton) Hessian apply H * v\n",
      "def Hess_GN (v, R, C, A, W):\n",
      "    solve (A, du, - (C * v))\n",
      "    solve (A, dp, - (W * du))\n",
      "    CT_dp = Vector()\n",
      "    C.init_vector(CT_dp, 1)\n",
      "    C.transpmult(dp, CT_dp)\n",
      "    H_V = R * v   CT_dp\n",
      "    return H_V\n",
      "\n",
      "# define (Newton) Hessian apply H * v\n",
      "def Hess_Newton (v, R, Raa, C, A, W, Wua):\n",
      "    RHS = -(C * v)\n",
      "    bc2.apply(RHS)\n",
      "    solve (A, du, RHS)\n",
      "    RHS = -(W * du) -  Wua * v\n",
      "    bc2.apply(RHS)\n",
      "    solve (A, dp, RHS)\n",
      "    CT_dp = Vector()\n",
      "    C.init_vector(CT_dp, 1)\n",
      "    C.transpmult(dp, CT_dp)\n",
      "    Wua_du = Vector()\n",
      "    Wua.init_vector(Wua_du, 1)\n",
      "    Wua.transpmult(du, Wua_du)\n",
      "    H_V = R*v   Raa*v   CT_dp   Wua_du\n",
      "    return H_V\n",
      "\n",
      "# Creat Class MyLinearOperator to perform Hessian function\n",
      "class MyLinearOperator(LinearOperator):\n",
      "    cgiter = 0\n",
      "    def __init__(self, R, Raa, C, A, W, Wua):\n",
      "        LinearOperator.__init__(self, a_delta, a_delta)\n",
      "        self.R = R\n",
      "        self.Raa = Raa\n",
      "        self.C = C\n",
      "        self.A = A\n",
      "        self.W = W\n",
      "        self.Wua = Wua\n",
      "\n",
      "    # Hessian performed on x, output as generic vector y\n",
      "    def mult(self, x, y):\n",
      "        self.cgiter  = 1\n",
      "        y.zero()\n",
      "        if iter <= 500:\n",
      "            y.axpy(1., Hess_GN (x, self.R, self.C, self.A, self.W) )\n",
      "        else:\n",
      "            y.axpy(1., Hess_Newton (x, self.R, self.Raa, self.C, self.A, self.W, self.Wua) )"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "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, p, g$ corresponding to state, adjoint and coefficient/gradient variables, and the corresponding test functions and the parameters for the optimization."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# create mesh and define function spaces\n",
      "nx = 64\n",
      "ny = 64\n",
      "mesh = UnitSquareMesh(nx, ny)\n",
      "V = FunctionSpace(mesh, 'Lagrange', 1)\n",
      "V2 = FunctionSpace(mesh, 'Lagrange', 2)\n",
      "\n",
      "# define Trial and Test Functions\n",
      "u, p, g, a_delta = TrialFunction(V2), TrialFunction(V2), TrialFunction(V), TrialFunction(V)\n",
      "u_test, p_test, ud_test, g_test = TestFunction(V2), TestFunction(V2), TestFunction(V2), TestFunction(V)\n",
      "p = Function(V2)\n",
      "\n",
      "# initialize input functions\n",
      "atrue = Expression('log(2   7*(pow(pow(x[0] - 0.5,2)   pow(x[1] - 0.5,2),0.5) > 0.2))')\n",
      "f = Constant(\"1.0\")\n",
      "u0 = Constant(\"0.0\")\n",
      "a = interpolate(Expression(\"log(2.0)\"),V)\n",
      "\n",
      "# plot\n",
      "atruef = interpolate(atrue, V)\n",
      "plt.figure(figsize=(15,5))\n",
      "nb.plot(mesh,subplot_loc=121, mytitle=\"Mesh\", show_axis='on')\n",
      "nb.plot(atruef,subplot_loc=122, mytitle=\"True parameter field\")\n",
      "plt.show()\n",
      "\n",
      "# noise level\n",
      "noise_level = 0.05\n",
      "\n",
      "# define parameters for the optimization\n",
      "tol = 1e-8\n",
      "gamma = 1e-8\n",
      "c = 1e-4\n",
      "maxiter = 100\n",
      "#eps = 1e-4\n",
      "plot_on = 0\n",
      "\n",
      "# initialize iter counters\n",
      "iter = 1\n",
      "total_cg_iter = 0\n",
      "solution = 0"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# set up dirichlet boundary conditions\n",
      "def u0_boundary(x,on_boundary):\n",
      "    return on_boundary\n",
      "bc2 = DirichletBC(V2, u0, u0_boundary)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "### Set up synthetic observations:\n",
      "\n",
      "- Propose a coefficient field $a_{\\text true}$ shown above\n",
      "- The weak form of the pde: \n",
      "    Find $u\\in H_0^1(\\Omega)$ such that $\\underbrace{(\\exp(a_{\\text true})\\nabla u,\\nabla v)}_{\\; := \\; a_{pde}} - \\underbrace{(f,v)}_{\\; := \\;L_{pde}} = 0, \\text{ for all } v\\in H_0^1(\\Omega)$.\n",
      "\n",
      "- Perturb the solution: $u = u   \\eta$, where $\\eta \\sim \\mathcal{N}(0, \\sigma)$"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# weak form for setting up the synthetic observations\n",
      "ud = TrialFunction(V2)\n",
      "a_goal = inner(exp(atrue) * nabla_grad(ud), nabla_grad(ud_test)) * dx\n",
      "L_goal = f * ud_test * dx\n",
      "\n",
      "# solve the forward/state problem to generate synthetic observations\n",
      "goal_A, goal_b = assemble_system(a_goal, L_goal, bc2)\n",
      "ud = Function(V2)\n",
      "solve(goal_A, ud.vector(), goal_b)\n",
      "\n",
      "utrue = Function(V2)\n",
      "utrue.assign(ud)\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 = Vector()\n",
      "goal_A.init_vector(noise,1)\n",
      "noise.set_local( noise_level * MAX * np.random.normal(0, 1, len(ud.vector().array())) )\n",
      "ud.vector().axpy(1., noise)\n",
      "bc2.apply(ud.vector())\n",
      "\n",
      "# plot\n",
      "nb.multi1_plot([utrue, ud], [\"State solution with atrue\", \"Synthetic observations\"])\n",
      "plt.show()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Setting up the state equations, right hand side for the adjoint and the neccessary matrices:\n"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# weak form for setting up the state equation\n",
      "a_state = inner(exp(a) * nabla_grad(u), nabla_grad(u_test)) * dx\n",
      "L_state = f * u_test * dx\n",
      "W_equ   = inner(u, u_test) * dx\n",
      "\n",
      "# # weak form for setting up the right hand side of the adjoint\n",
      "u = Function(V2)\n",
      "L_adjoint = -inner(u - ud, u_test) * dx\n",
      "\n",
      "# weak form for setting up matrices\n",
      "Wua_equ = inner(exp(a) * a_delta * nabla_grad(p_test), nabla_grad(p)) * dx\n",
      "C_equ   = inner(exp(a) * a_delta * nabla_grad(u), nabla_grad(u_test)) * dx\n",
      "M_equ   = inner(g, g_test) * dx\n",
      "R_equ   = gamma * inner(nabla_grad(g), nabla_grad(g_test)) * dx\n",
      "Raa_equ = inner(exp(a) * a_delta * g_test *  nabla_grad(u),  nabla_grad(p)) * dx\n",
      "\n",
      "# assemble matrices M, W, and R\n",
      "M = assemble(M_equ)\n",
      "W = assemble(W_equ)\n",
      "R = assemble(R_equ)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# solve state equation\n",
      "A, state_b = assemble_system (a_state, L_state, bc2)\n",
      "solve (A, u.vector(), state_b)\n",
      "\n",
      "# evaluate cost\n",
      "[cost_old, misfit_old, reg_old] = cost(u, ud, a, W, R)\n",
      "\n",
      "# plot\n",
      "plt.figure(figsize=(15,5))\n",
      "nb.plot(a,subplot_loc=121, mytitle=\"a_ini\", vmin=atruef.vector().min(), vmax=atruef.vector().max())\n",
      "nb.plot(u,subplot_loc=122, mytitle=\"u(a_ini)\")\n",
      "plt.show()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## The inexact Newton-CG optimization with Armijo line search:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# initializations\n",
      "g, a_delta = Vector(), Vector()\n",
      "R.init_vector(a_delta,0)\n",
      "R.init_vector(g,0)\n",
      "\n",
      "du, dp = Vector(), Vector()\n",
      "W.init_vector(du,1)\n",
      "W.init_vector(dp,0)\n",
      "\n",
      "a_prev, a_diff = Function(V), Function(V)\n",
      "\n",
      "print \"Nit   CGit   cost          misfit        reg           sqrt(-G*D)    ||grad||       alpha  tolcg\"\n",
      "\n",
      "while iter <  maxiter and solution == 0:\n",
      "\n",
      "    # assemble matrix C\n",
      "    C =  assemble(C_equ)\n",
      "\n",
      "    # solve the adoint problem\n",
      "    A, adjoint_RHS = assemble_system(a_state, L_adjoint, bc2)\n",
      "    solve(A, p.vector(), adjoint_RHS)\n",
      "\n",
      "    # assemble W_ua and R\n",
      "    Wua = assemble (Wua_equ)\n",
      "    Raa = assemble (Raa_equ)\n",
      "\n",
      "    # evaluate the  gradient\n",
      "    CT_p = Vector()\n",
      "    C.init_vector(CT_p,1)\n",
      "    C.transpmult(p.vector(), CT_p)\n",
      "    MG = CT_p   R * a.vector()\n",
      "    solve(M, g, MG)\n",
      "\n",
      "    # calculate the norm of the gradient\n",
      "    grad2 = g.inner(MG)\n",
      "    gradnorm = sqrt(grad2)\n",
      "\n",
      "    # set the CG tolerance (use Eisenstat\u2013Walker termination criterion)\n",
      "    if iter == 1:\n",
      "        gradnorm_ini = gradnorm\n",
      "    tolcg = min(0.5, sqrt(gradnorm/gradnorm_ini))\n",
      "\n",
      "    # define the Hessian apply operator (with preconditioner)\n",
      "    Hess_Apply = MyLinearOperator(R, Raa, C, A, W, Wua )\n",
      "    P = R   gamma * M\n",
      "    solver = PETScKrylovSolver(\"cg\", \"amg\")\n",
      "    solver.set_operators(Hess_Apply, P)\n",
      "    solver.parameters[\"relative_tolerance\"] = tolcg\n",
      "\n",
      "    # solve the Newton system H a_delta = - MG\n",
      "    solver.solve(a_delta, -MG)\n",
      "    total_cg_iter  = Hess_Apply.cgiter\n",
      "\n",
      "    # linesearch\n",
      "    alpha = 1\n",
      "    descent = 0\n",
      "    no_backtrack = 0\n",
      "    a_prev.assign(a)\n",
      "    while descent == 0 and no_backtrack < 10:\n",
      "        a.vector().axpy(alpha, a_delta )\n",
      "\n",
      "        # solve the state/forward problem\n",
      "        state_A, state_b = assemble_system(a_state, L_state, bc2)\n",
      "        solve(state_A, u.vector(), state_b)\n",
      "\n",
      "        # evaluate cost\n",
      "        [cost_new, misfit_new, reg_new] = cost(u, ud, a, W, R)\n",
      "\n",
      "        # check if Armijo conditions are satisfied\n",
      "        if cost_new < cost_old - alpha * c * grad2:\n",
      "            cost_old = cost_new\n",
      "            descent = 1\n",
      "        else:\n",
      "            no_backtrack  = 1\n",
      "            alpha *= 0.5\n",
      "            a.assign(a_prev)  # reset a\n",
      "\n",
      "    # calculate sqrt(-G * D)\n",
      "    graddir = sqrt(- MG.inner(a_delta) )\n",
      "\n",
      "    sp = \"\"\n",
      "    print \"- %2s - %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 == 1:\n",
      "        nb.multi1_plot([a,u,p], [\"a\",\"u\",\"p\"], same_colorbar=False)\n",
      "        plt.show()\n",
      "    \n",
      "    # check for convergence\n",
      "    if sqrt(grad2) < tol and iter > 1:\n",
      "        solution = 1\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 solution == 0:\n",
      "    print \"Newton's method did not converge in \", maxiter, \" iterations\"\n",
      "\n",
      "print \"Time elapsed: \", time.clock()-start"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "nb.multi1_plot([atruef, a], [\"atrue\", \"a\"])\n",
      "nb.multi1_plot([u,p], [\"u\",\"p\"], same_colorbar=False)\n",
      "plt.show()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}