{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7.2 PDE-Constrained Shape Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to solve the PDE-constrained shape optimization problem\n",
    "$$\n",
    "            \\underset{\\Omega\\subset \\mathsf{D}}{\\mbox{min}} \\; J(u) := \\int_\\Omega |u-u_d|^2 \\; dx\n",
    "$$\n",
    " subject to that $(\\Omega,u)$ satisfy\n",
    " $$\n",
    "           \\int_\\Omega \\nabla u \\cdot \\nabla v \\; dx = \\int_\\Omega f v \\; dx \\; \\quad \\text{ for all } v \\in H_0^1(\\Omega),\n",
    "$$\n",
    "where $\\Omega \\subset \\mathbb R^2$ for given $u_d, f \\in C^1(\\mathbb R^2)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "from netgen.geom2d import SplineGeometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo = SplineGeometry()\n",
    "geo.AddRectangle( (0.2, 0.2), (0.8, 0.8), bcs = ('b','r','t','l'))\n",
    "mesh = Mesh(geo.GenerateMesh(maxh = 0.1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#given data of our problem (chosen such that \\Omega^* = [0,1]^2 is the optimal shape)\n",
    "f = CoefficientFunction(2*y*(1-y)+2*x*(1-x))\n",
    "ud = x*(1-x)*y*(1-y)\n",
    "\n",
    "grad_f = CoefficientFunction( (f.Diff(x), f.Diff(y) ) )\n",
    "grad_ud = CoefficientFunction( (ud.Diff(x), ud.Diff(y) ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### State equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fes = H1(mesh, order=2, dirichlet=\".*\")\n",
    "u, v = fes.TnT()\n",
    "gfu = GridFunction(fes)\n",
    "scene = Draw (gfu, mesh, \"state\")\n",
    "\n",
    "a = BilinearForm(fes, symmetric=True)\n",
    "a += grad(u)*grad(v)*dx\n",
    "\n",
    "fstate = LinearForm(fes)\n",
    "fstate += f*v*dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveStateEquation():\n",
    "    rhs = gfu.vec.CreateVector()\n",
    "    rhs.data = fstate.vec - a.mat * gfu.vec\n",
    "    update = gfu.vec.CreateVector()\n",
    "    update.data = a.mat.Inverse(fes.FreeDofs()) * rhs\n",
    "    gfu.vec.data += update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.Assemble()\n",
    "fstate.Assemble()\n",
    "SolveStateEquation()\n",
    "scene.Redraw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adjoint equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set up the adjoint equation\n",
    "$$\n",
    "    \\mbox{Find } p \\in H_0^1(\\Omega): \\int_\\Omega \\nabla w \\cdot \\nabla p \\, \\mbox dx = - \\partial_u J(u)(w) \\quad \\text{ for all } w \\in H_0^1(\\Omega)\n",
    "$$\n",
    "where $u$ is the solution to the state equation. For $J(u) = \\int_\\Omega |u-u_d|^2 \\mbox dx$, we get\n",
    "$$\n",
    "    \\partial_u J(u)(w) = -2 \\int_\\Omega (u-u_d)w \\,\\mbox dx.\n",
    "$$\n",
    "However, we can also use the Diff(...) command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Cost(u): \n",
    "    return (u-ud)**2*dx\n",
    "\n",
    "p, w = fes.TnT()\n",
    "gfp = GridFunction(fes)\n",
    "scene = Draw (gfp, mesh, \"adjoint\")\n",
    "\n",
    "fadjoint = LinearForm(fes)\n",
    "fadjoint += -1*Cost(gfu).Diff(gfu,w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveAdjointEquation():\n",
    "    rhs = gfp.vec.CreateVector()\n",
    "    rhs.data = fadjoint.vec - a.mat.T * gfp.vec\n",
    "    update = gfp.vec.CreateVector()\n",
    "    update.data = a.mat.Inverse(fes.FreeDofs()).T * rhs\n",
    "    gfp.vec.data += update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fadjoint.Assemble()\n",
    "SolveAdjointEquation()\n",
    "scene.Redraw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that (for linear problems) the operator on the left hand side of the adjoint equation is just the transpose of the state operator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Shape derivative"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The shape derivative of the shape function \n",
    "$$\n",
    "\\mathcal{J}(\\Omega):= \\int_\\Omega |u-u_d|^2\\;dx\n",
    "$$\n",
    "at $\\Omega$ in direction $X$ is given by the formula\n",
    "$$\n",
    "    \\begin{array}{rl}\n",
    "    d\\mathcal J(\\Omega; X) =&\\int_{\\Omega}  \\mbox{div}(X) |u-u_d|^2  - 2(u-u_d)\\nabla u_d \\cdot X \\, dx \\\\\n",
    "\t\t& + \\int_{\\Omega} (\\mbox{div}(X) I - D X - D X^\\top )\\nabla u \\cdot \\nabla p \\, dx \\\\\n",
    "\t\t&- \\int_\\Omega (\\nabla f\\cdot X + f \\mbox{div}(X)) p\\, dx.\n",
    "    \\end{array}\n",
    "$$\n",
    "We now assemble this shape derivaitve in NGSolve as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VEC = H1(mesh, order=2, dim=2)\n",
    "PHI, X = VEC.TnT()\n",
    "\n",
    "dJOmega = LinearForm(VEC)\n",
    "dJOmega += SymbolicLFI(div(X)*( (gfu-ud)**2 - f*gfp + InnerProduct(grad(gfu),grad(gfp))))\n",
    "dJOmega += SymbolicLFI(-2*(gfu-ud)*InnerProduct(grad_ud,X) - InnerProduct(grad_f,X)*gfp)\n",
    "dJOmega += SymbolicLFI(-InnerProduct(X.Deriv()*grad(gfu),grad(gfp))-InnerProduct(grad(gfu),X.Deriv()*grad(gfp)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Descent Direction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<ul>\n",
    " <li> Next, we want to find a vector field $X$ which yields a decrease of the objective function, i.e. we want to find a vector field $X$ such that $$d \\mathcal J(\\Omega;X)<0.$$ </li>\n",
    "<li> This can be achieved by solving an auxiliary boundary value problem of the type \n",
    "$$\n",
    "    \\mbox{Find } X \\in H: \\qquad b(X, \\Phi) = d\\mathcal J(\\Omega; \\Phi) \\; \\text{ for all } \\Phi \\in H\n",
    "$$\n",
    "for a suitable Hilbert space $H$ (e.g. $H=H^1(\\Omega)^2$). Here, $b(\\cdot, \\cdot): H \\times H \\rightarrow \\mathbb R$ is a positive definite bilinear form which we are free to choose. Then, $-X$ is a descent direction since $$ d\\mathcal J(\\Omega; -X) = -d\\mathcal J(\\Omega; X) = - b(X, X) < 0.$$ </li>\n",
    "</ul>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = BilinearForm(VEC)\n",
    "b += InnerProduct(grad(X),grad(PHI))*dx + InnerProduct(X,PHI)*dx\n",
    "\n",
    "gfX = GridFunction(VEC)\n",
    "\n",
    "# gfset denotes the deformation of the original domain and will be updated during the shape optimization\n",
    "gfset = GridFunction(VEC)\n",
    "gfset.Set((0,0))\n",
    "scene = Draw(gfset,mesh,\"gfset\")\n",
    "SetVisualization (deformation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveDeformationEquation():\n",
    "    rhs = gfX.vec.CreateVector()\n",
    "    rhs.data = dJOmega.vec - b.mat * gfX.vec\n",
    "    update = gfX.vec.CreateVector()\n",
    "    update.data = b.mat.Inverse(VEC.FreeDofs()) * rhs\n",
    "    gfX.vec.data += update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b.Assemble()\n",
    "dJOmega.Assemble()\n",
    "SolveDeformationEquation()\n",
    "Draw(-gfX, mesh, \"gfX\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us update the geometry in the direction of $-X$. Here, we need to choose the step size accordingly (a descent is only guaranteed for this step size being sufficiently small)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Cost at initial design', Integrate (Cost(gfu), mesh))\n",
    "gfset.Set((0,0))\n",
    "scale = 0.5 / Norm(gfX.vec)\n",
    "gfset.vec.data -= scale * gfX.vec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now update the design and compute the new value of the cost function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene = Draw(gfset)\n",
    "mesh.SetDeformation(gfset)\n",
    "scene.Redraw()\n",
    "\n",
    "a.Assemble()\n",
    "fstate.Assemble()\n",
    "SolveStateEquation()\n",
    "print('Cost at new design', Integrate (Cost(gfu), mesh))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#reset the design\n",
    "gfset.Set((0,0))\n",
    "mesh.SetDeformation(gfset)\n",
    "\n",
    "#check if initial value of cost function 0.000486578350214552 is recovered\n",
    "a.Assemble()\n",
    "fstate.Assemble()\n",
    "SolveStateEquation()\n",
    "print('Cost at new design', Integrate (Cost(gfu), mesh))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, this can be applied in an iterative algorithm with a line search for choosing the step size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene = Draw(gfset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfset.Set((0,0))\n",
    "mesh.SetDeformation(gfset)\n",
    "scene.Redraw()\n",
    "a.Assemble()\n",
    "fstate.Assemble()\n",
    "SolveStateEquation()\n",
    "\n",
    "LineSearch = True\n",
    "\n",
    "iter_max = 600\n",
    "Jold = Integrate(Cost(gfu), mesh)\n",
    "converged = False\n",
    "\n",
    "# input(\"Press enter to start optimization\")\n",
    "for k in range(iter_max):\n",
    "    mesh.SetDeformation(gfset)\n",
    "    scene.Redraw()\n",
    "    \n",
    "    print('cost at iteration', k, ': ', Jold)\n",
    "    \n",
    "    a.Assemble()\n",
    "    fstate.Assemble()\n",
    "    SolveStateEquation()\n",
    "    \n",
    "    fadjoint.Assemble()\n",
    "    SolveAdjointEquation()\n",
    "    \n",
    "    b.Assemble()\n",
    "    dJOmega.Assemble()\n",
    "    SolveDeformationEquation()\n",
    "\n",
    "    scale = 0.01 / Norm(gfX.vec)\n",
    "    gfsetOld = gfset\n",
    "    gfset.vec.data -= scale * gfX.vec\n",
    "    \n",
    "    Jnew = Integrate(Cost(gfu), mesh)\n",
    "    \n",
    "    if LineSearch:\n",
    "        while Jnew > Jold and scale > 1e-12:\n",
    "            #input('a')\n",
    "            scale = scale / 2\n",
    "            print(\"scale = \", scale)\n",
    "            if scale <= 1e-12:\n",
    "                converged = True\n",
    "                break\n",
    "\n",
    "            gfset.vec.data = gfsetOld.vec - scale * gfX.vec\n",
    "            mesh.SetDeformation(gfset)\n",
    "            \n",
    "            a.Assemble()\n",
    "            fstate.Assemble()\n",
    "            SolveStateEquation()\n",
    "            Jnew = Integrate(Cost(gfu), mesh)\n",
    "    \n",
    "    if converged==True:\n",
    "        print(\"No more descent can be found\")\n",
    "        break\n",
    "    Jold = Jnew\n",
    "\n",
    "    Redraw(blocking=True)"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}