{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7.3 PDE-Constrained Shape Optimization (semi-automated)"
   ]
  },
  {
   "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|^q \\; dx, \\quad q\\ge 2\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)$. \n",
    "\n",
    "Here, we want to compute the shape derivative by differentiation of a suitably defined perturbed Lagrangian using automated differentiation. For details we refer to \n",
    "\n",
    "P. Gangl, K. Sturm, M. Neunteufel, J. Schöberl.\n",
    "Fully and Semi-Automated Shape Differentiation in NGSolve,\n",
    "Struct. Multidisc. Optim., 63, pp.1579-1607, 2021."
   ]
  },
  {
   "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.AddCircle(c=(0.5,0.5), r=0.5, bc = 'circle')\n",
    "mesh = Mesh(geo.GenerateMesh(maxh = 0.08))"
   ]
  },
  {
   "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",
    "\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|^q \\mbox dx$, we get\n",
    "$$\n",
    "    \\partial_u J(u)(w) = q \\int_\\Omega (u-u_d)^{d-1}w \\,\\mbox dx.\n",
    "$$\n",
    "However, we can also use the Diff(...) command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q=4\n",
    "def Cost(u): \n",
    "    return (u-ud)**q*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": [
    "### Automatic Shape Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The formula for the shape derivative was derived as the partial derivative of the perturbed Lagrangian (brought back to the original domain):\n",
    "$$\n",
    "    d\\mathcal J(\\Omega; X) = \\frac{\\partial}{\\partial t} \\left. \\left( \\int_\\Omega \\xi(t)|u - u_d^t|^q \\mbox{d} x \n",
    "         +  \\int_{\\Omega} (F_t^{-\\top}\\nabla u) \\cdot (F_t^{-\\top} \\nabla p) \\xi(t) \\, dx - \\int_{\\Omega} \\xi(t) f^t p \\,dx \\right) \\right\\rvert_{t=0} \n",
    "$$\n",
    "where \n",
    "<ul>\n",
    "    <li>   $T_t(x)=x+tX(x)=y$ \n",
    "    <li> $F_t = DT_t = I+t DX$\n",
    "    <li> $\\xi(t) = \\mbox{det}(F_t)$\n",
    "    <li>  $u_d^t = u_d \\circ T_t$\n",
    "    <li> $f^t = f \\circ T_t$\n",
    "</ul>\n",
    "\n",
    "The integrand depends on the parameter $t$ only via $\\xi(t)$, $F_t$ and $T_t$. Denoting the integrand by $G^{u,p}$, the derivative is given by\n",
    "$$\n",
    "\\begin{array}{rl}\n",
    "     d\\mathcal J(\\Omega; X) =& \\frac{d G^{u,p}}{d \\xi} \\frac{d \\xi}{d t} + \\frac{d  G^{u,p}}{d F} \\frac{d F}{d t} + \\frac{d  G^{u,p}}{dy} \\cdot \\frac{d T_t}{dt} \\\\\n",
    "     =& \\frac{d G^{u,p}}{d \\xi} \\mbox{div}(X) + \\frac{d  G^{u,p}}{d F} DX + \\frac{d  G^{u,p}}{dy} \\cdot X\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the perturbed PDE and cost function involving parameters $F$ and $J$ representing the Jacobian of a shape transformation and its determinant, respectively. While their values are set to unity and thus does not influence the solution of the state or adjoint equation, this procedure allows us to differentiate with respect to them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J = Parameter(1)\n",
    "F = Id(2)\n",
    "#Finv = 2*Id(2) - F    # consistent linearization at F = I (avoids calling Inv() for speedup)\n",
    "\n",
    "def Equation(u,v):\n",
    "    return ( (Inv(F.trans)*grad(u))*(Inv(F.trans)*grad(v))-f*v)*J*dx\n",
    "#    return ( (Finv.trans*grad(u))*(Finv.trans*grad(v))-f*v)*J*dx\n",
    "\n",
    "def CostAuto(u): \n",
    "    return J*(u-ud)**q*dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VEC = H1(mesh, order=2, dim=2)\n",
    "PHI, X = VEC.TnT()\n",
    "\n",
    "Lagrangian = CostAuto(gfu) + Equation(gfu,gfp)\n",
    "dJOmegaAuto = LinearForm(VEC)\n",
    "dJOmegaAuto += Lagrangian.Diff(J, div(X))\n",
    "dJOmegaAuto += Lagrangian.Diff(F, grad(X).trans)   # grad(X).trans is Jacobian\n",
    "dJOmegaAuto += Lagrangian.Diff(x, X[0]) + Lagrangian.Diff(y, X[1])"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveDeformationEquationAuto():\n",
    "    rhs = gfX.vec.CreateVector()\n",
    "    rhs.data = dJOmegaAuto.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",
    "dJOmegaAuto.Assemble()\n",
    "SolveDeformationEquationAuto()\n",
    "Draw(-gfX, mesh, \"gfX\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "mesh.SetDeformation(gfset)\n",
    "sceneSet = Draw(gfset,mesh,\"gfset\")\n",
    "SetVisualization (deformation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfset.Set((0,0))\n",
    "mesh.SetDeformation(gfset)\n",
    "print('Cost at initial design', Integrate (CostAuto(gfu), mesh))\n",
    "\n",
    "scale = 0.5 / Norm(gfX.vec)\n",
    "gfset.vec.data -= scale * gfX.vec\n",
    "mesh.SetDeformation(gfset)\n",
    "sceneSet.Redraw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.Assemble()\n",
    "fstate.Assemble()\n",
    "SolveStateEquation()\n",
    "print('Cost at new design', Integrate (CostAuto(gfu), mesh))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Equation can also be used to define the bilinear form. The following defines the same bilinear form as above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aAuto = BilinearForm(fes, symmetric=True)\n",
    "aAuto += Equation(u,v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, the user has to enter the PDE (in its transformed form) only once."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let us again run the full algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#reset to and solve for initial configuration\n",
    "gfset.Set((0,0))\n",
    "mesh.SetDeformation(gfset)\n",
    "scene = Draw(gfset,mesh,\"gfset\")\n",
    "SetVisualization (deformation=True)\n",
    "\n",
    "a.Assemble()\n",
    "fstate.Assemble()\n",
    "SolveStateEquation()\n",
    "\n",
    "LineSearch = False\n",
    "\n",
    "iter_max = 600\n",
    "Jold = Integrate(CostAuto(gfu), mesh)\n",
    "converged = False\n",
    "for k in range(iter_max):\n",
    "    scene.Redraw()\n",
    "    print('cost at iteration', k, ': ', Jold)\n",
    "    mesh.SetDeformation(gfset)\n",
    "    \n",
    "    a.Assemble()\n",
    "    fstate.Assemble()\n",
    "    SolveStateEquation()\n",
    "    \n",
    "    fadjoint.Assemble()\n",
    "    SolveAdjointEquation()\n",
    "    \n",
    "    b.Assemble()\n",
    "    dJOmegaAuto.Assemble()\n",
    "    SolveDeformationEquationAuto()\n",
    "\n",
    "    scale = 0.01 / Norm(gfX.vec)\n",
    "    gfsetOld = gfset\n",
    "    gfset.vec.data -= scale * gfX.vec\n",
    "    \n",
    "    Jnew = Integrate(CostAuto(gfu), mesh)\n",
    "    \n",
    "    if LineSearch:\n",
    "        while Jnew > Jold and scale > 1e-12:\n",
    "            scale = scale / 2\n",
    "            \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(CostAuto(gfu), mesh)\n",
    "    \n",
    "    if converged==True:\n",
    "        break\n",
    "    Jold = Jnew\n",
    "\n"
   ]
  },
  {
   "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
}