{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7.7 PDE-Constrained Topology Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are again interested in minimizing a shape function of the type\n",
    "$$\n",
    "     \\mathcal J(\\Omega) = \\int_{\\mathsf{D}} (u-u_d)^2 \\; dx\n",
    "$$\n",
    "subject to $u:\\mathsf{D} \\to \\mathbb{R}$ solves\n",
    "$$ \n",
    "\\int_{\\mathsf{D}} \\beta_\\Omega\\nabla u\\cdot \\nabla \\varphi\\;dx = \\int_{\\mathsf{D}} f_\\Omega\\varphi \\quad \\text{ for all } \\varphi \\in H^1_0(\\mathsf{D}), \n",
    "$$\n",
    "where\n",
    "<ul>\n",
    "    <li> $\\beta_\\Omega = \\beta_1\\chi_\\Omega + \\beta_2 \\chi_{\\mathsf{D}\\setminus \\Omega}$, $\\qquad \\beta_1,\\beta_2>0$, \n",
    "    <li> $f_\\Omega = f_1\\chi_\\Omega + f_2 \\chi_{\\mathsf{D}\\setminus \\Omega}$, $\\qquad f_1,f_2\\in \\mathbb{R}$, \n",
    "    <li> $u_d:\\mathsf{D} \\to \\mathbb{R}$. \n",
    "</ul>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netgen.geom2d import SplineGeometry  # SplieGeometry to define a 2D mesh\n",
    "\n",
    "from ngsolve import * # import everything from the ngsolve module \n",
    "from ngsolve.webgui import Draw\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from interpolations import InterpolateLevelSetToElems # function which interpolates a levelset function\n",
    "\n",
    "myMAXH = 0.05\n",
    "EPS = myMAXH * 1e-6      #variable defining when a point is on the interface and when not"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The file interpolations.py can contains a routine to compute the volume ratio of each element with respect to a given level set function. The file can be downloaded from [interpolations.py](interpolations.py) . \n",
    "\n",
    "We begin with defining the domain $\\mathsf{D}$.  For this we define a rectangular mesh with side length $R$ using the NGSolve function SplineGeometry() with maximum triangle size maxh. The boundary name of $\\partial\\mathsf{D}$ is \"rectangle\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo = SplineGeometry()   \n",
    "\n",
    "R = 2\n",
    "\n",
    "## add a rectangle\n",
    "geo.AddRectangle(p1=(-R,-R),\n",
    "                 p2=(R,R),\n",
    "                 bc=\"rectangle\",\n",
    "                 leftdomain=1,\n",
    "                 rightdomain=0)\n",
    "\n",
    "\n",
    "geo.SetMaterial (1, \"outer\") # give the domain the name \"outer\"\n",
    "\n",
    "mesh = Mesh(geo.GenerateMesh(maxh=myMAXH)) # generate ngsolve mesh\n",
    "Draw(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<ul>\n",
    "    <li> fes_state,fes_adj - $H^1$ conforming finite element space of order 1 </li>\n",
    "    <li> pwc - $L^2$ subspace of functions constant on each element of mesh </li>\n",
    "    <li> $u,v$ are our trial and test functions </li>\n",
    "    <li> gfu, gfp are Gridfunctions storing the state and adjoint variable, respectively </li>\n",
    "    <li> gfud stores $u_d$ - the target function </li>\n",
    "</ul>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# H1-conforming finite element space\n",
    "\n",
    "fes_state = H1(mesh, order=1, dirichlet=\"rectangle\")\n",
    "fes_adj = H1(mesh, order=1, dirichlet=\"rectangle\")\n",
    "fes_level = H1(mesh, order=1)\n",
    "\n",
    "pwc = L2(mesh)   #piecewise constant space\n",
    "\n",
    "\n",
    "## test and trial functions\n",
    "u, v = fes_state.TnT()\n",
    "\n",
    "p, q = fes_adj.TnT()\n",
    "\n",
    "gfu = GridFunction(fes_state)\n",
    "gfp = GridFunction(fes_adj)\n",
    "\n",
    "gfud = GridFunction(fes_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We represent the design $\\Omega \\subset \\mathsf{D}$ by means of a level set function $\\psi: \\mathsf{D} \\rightarrow \\mathbb R$ in the following way:\n",
    "\n",
    "$$\n",
    "\\begin{array}{rl}\n",
    "    \\psi(x) < 0 &\\quad \\Longleftrightarrow\\quad  x \\in \\Omega \\\\\n",
    "    \\psi(x) = 0 &\\quad \\Longleftrightarrow\\quad  x \\in \\partial \\Omega \\\\\n",
    "    \\psi(x) > 0 &\\quad \\Longleftrightarrow\\quad  x \\in \\mathsf{D} \\setminus \\overline \\Omega.\n",
    "\\end{array}\n",
    "$$\n",
    "<ul>\n",
    "    <li> psi - gridfunction in fes_state defining $\\psi$ \n",
    "    <li> psides - gridfunction describing the optimal shape\n",
    "    <li> psinew - dummy function for line search \n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi = GridFunction(fes_level)\n",
    "psi.Set(1) \n",
    "\n",
    "psides = GridFunction(fes_level)\n",
    "psinew = GridFunction(fes_level)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we solve the state equation\n",
    "\\begin{equation}\n",
    "\\int_{\\mathsf{D}} \\beta_\\Omega \\nabla u \\cdot \\nabla \\varphi \\;dx = \\int_{\\mathsf{D}} f_\\Omega\\varphi \\;dx \\quad \\text{ for all } \\varphi \\in H^1_0(\\mathsf{D}).\n",
    "\\end{equation}\n",
    "<ul>\n",
    "    <li> beta, f_rhs - piecewise constant gridfunction </li> \n",
    "    <li> B - bilinear form defining $\\int_{\\mathsf{D}} \\beta_\\Omega\\nabla \\psi \\cdot \\nabla \\varphi \\;dx$ </li>\n",
    "    <li> B_adj - bilinear form defining $\\int_{\\mathsf{D}} \\beta_\\Omega\\nabla \\psi \\cdot \\nabla \\varphi \\;dx$ </li>\n",
    "    <li> L - right hand side $\\int_{\\mathsf{D}} f_\\Omega \\varphi \\;dx$ </li>\n",
    "</ul>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# constants for f_rhs and beta\n",
    "f1 = 10\n",
    "f2 = 1\n",
    "\n",
    "beta1 = 10\n",
    "beta2 = 1\n",
    "\n",
    "# piecewise constant coefficient functions beta and f_rhs\n",
    "\n",
    "beta = GridFunction(pwc)\n",
    "beta.Set(beta1)\n",
    "\n",
    "f_rhs = GridFunction(pwc)\n",
    "f_rhs.Set(f1)\n",
    "\n",
    "# bilinear form for state equation\n",
    "B = BilinearForm(fes_state)\n",
    "B += beta*grad(u) * grad(v) * dx\n",
    "\n",
    "B_adj = BilinearForm(fes_adj)\n",
    "B_adj += beta*grad(p) * grad(q) * dx\n",
    "\n",
    "L = LinearForm(fes_state)\n",
    "L += f_rhs * v *dx\n",
    "\n",
    "duCost = LinearForm(fes_adj)\n",
    "\n",
    "# solving\n",
    "psides.Set(1)\n",
    "InterpolateLevelSetToElems(psides, beta1, beta2, beta, mesh, EPS)\n",
    "InterpolateLevelSetToElems(psides, f1, f2, f_rhs, mesh, EPS)\n",
    "\n",
    "B.Assemble()\n",
    "L.Assemble()\n",
    "\n",
    "inv = B.mat.Inverse(fes_state.FreeDofs(), inverse=\"sparsecholesky\") # inverse of bilinear form\n",
    "gfu.vec.data = inv*L.vec\n",
    "\n",
    "scene_u = Draw(gfu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an optimal shape we define $\\Omega_{opt} := \\{f<0\\}$, where $f$ is the function from the levelset example describing a clover shape. We construct $u_d$ by solving \n",
    "$$\n",
    "\\int_{\\mathsf{D}} \\beta_{\\Omega_{opt}} \\nabla u_d \\cdot \\nabla \\varphi \\;dx = \\int_{\\mathsf{D}} f_{\\Omega_{opt}}\\varphi \\;dx \\quad \\text{ for all } \\varphi \\in H^1_0(\\mathsf{D}).\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = 4.0/5.0\n",
    "b = 2\n",
    "f = CoefficientFunction( 0.1*( (sqrt((x - a)**2 + b * y**2) - 1) \\\n",
    "                * (sqrt((x + a)**2 + b * y**2) - 1) \\\n",
    "                * (sqrt(b * x**2 + (y - a)**2) - 1) \\\n",
    "                * (sqrt(b * x**2 + (y + a)**2) - 1) - 0.001) )\n",
    "\n",
    "psides.Set(f)\n",
    "InterpolateLevelSetToElems(psides, beta1, beta2, beta, mesh, EPS)\n",
    "InterpolateLevelSetToElems(psides, f1, f2, f_rhs, mesh, EPS)\n",
    "\n",
    "B.Assemble()\n",
    "L.Assemble()\n",
    "inv = B.mat.Inverse(fes_state.FreeDofs(), inverse=\"sparsecholesky\") # inverse of bilinear form\n",
    "\n",
    "gfud.vec.data = inv*L.vec\n",
    "\n",
    "Draw(gfud, mesh, 'gfud')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function SolvePDE solves the state and adjoint state each time it is called. Note that since the functions \n",
    "beta and f_rhs are GridFunctions we do not have to re-define the bilinear form and linear form when beta and f_rhs change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolvePDE(adjoint=False):\n",
    "    #Newton(a, gfu, printing = False, maxerr = 3e-9)\n",
    "\n",
    "    B.Assemble()\n",
    "    L.Assemble()\n",
    "\n",
    "    inv_state = B.mat.Inverse(fes_state.FreeDofs(), inverse=\"sparsecholesky\")\n",
    "    \n",
    "    # solve state equation\n",
    "    gfu.vec.data = inv_state*L.vec\n",
    "\n",
    "    if adjoint == True: \n",
    "        # solve adjoint state equatoin\n",
    "        duCost.Assemble()\n",
    "        B_adj.Assemble()\n",
    "        inv_adj = B_adj.mat.Inverse(fes_adj.FreeDofs(), inverse=\"sparsecholesky\")\n",
    "        gfp.vec.data = -inv_adj * duCost.vec\n",
    "    scene_u.Redraw()\n",
    "\n",
    "InterpolateLevelSetToElems(psi, beta1, beta2, beta, mesh, EPS)\n",
    "InterpolateLevelSetToElems(psi, f1, f2, f_rhs, mesh, EPS)\n",
    "\n",
    "SolvePDE()\n",
    "\n",
    "# define the cost function\n",
    "def Cost(u):\n",
    "    return (u - gfud)**2*dx\n",
    "\n",
    "# derivative of cost function\n",
    "duCost += 2*(gfu-gfud) * q * dx\n",
    "    \n",
    "print(\"initial cost = \", Integrate(Cost(gfu), mesh))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It remains to implement the topological derivative:\n",
    "In define $NegPos$ and $PosNeg$ as the topological derivative in $\\Omega$ respectively $D \\setminus \\overline{\\Omega}$:\n",
    "$$\n",
    "\\begin{array}{rl}\n",
    "    \\mathsf{NegPos}(x) :=  -DJ(\\Omega,\\omega)(x) &= (-1) * \\left(   2 \\beta_2 \\left(\\frac{\\beta_2-\\beta_1}{\\beta_1+\\beta_2}\\right)\\nabla u(x)\\cdot \\nabla p(x) - (f_2-f_1)p(x) \\right) \\\\\n",
    "    &= 2 \\beta_2 \\left(\\frac{\\beta_1-\\beta_2}{\\beta_1+\\beta_2}\\right)\\nabla u(x)\\cdot \\nabla p(x) - (f_1-f_2)p(x) \\quad \\text{ for }x\\in \\Omega \\\\\n",
    "    \\mathsf{PosNeg}(x) := DJ(\\Omega,\\omega)(x) &= 2 \\beta_1 \\left(\\frac{\\beta_1-\\beta_2}{\\beta_1+\\beta_2}\\right)\\nabla u(x)\\cdot \\nabla p(x) - (f_1-f_2)p(x) \\quad \\text{ for }x\\in D \\setminus \\overline{\\Omega}\n",
    "\\end{array}\n",
    "$$\n",
    "and the update function (also called generalised topological derivative)\n",
    "$$\n",
    "g_\\Omega(x):= -\\chi_\\Omega (x)DJ(\\Omega,\\omega)(x) + (1-\\chi_\\Omega)*DJ(\\Omega,\\omega)(x) . \n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BetaPosNeg = 2 * beta2 * (beta1-beta2)/(beta1+beta2)    # factors in TD in positive part {x: \\psi(x)>0} ={x: \\beta(x) = \\beta_2}\n",
    "BetaNegPos = 2 * beta1 * (beta1-beta2)/(beta1+beta2)    # factors in TD in positive part {x: \\psi(x)>0} ={x: \\beta(x) = \\beta_2}\n",
    "FPosNeg = -(f1-f2)\n",
    "\n",
    "psinew.vec.data = psi.vec\n",
    "\n",
    "## normalise psi in L2\n",
    "normPsi = sqrt(Integrate(psi**2*dx, mesh)) \n",
    "psi.vec.data = 1/normPsi * psi.vec\n",
    "\n",
    "kappa = 0.05\n",
    "\n",
    "# set level set function to data\n",
    "InterpolateLevelSetToElems(psi, beta1, beta2, beta, mesh, EPS)\n",
    "InterpolateLevelSetToElems(psi, f1, f2, f_rhs, mesh, EPS)\n",
    "\n",
    "Redraw()\n",
    "\n",
    "TD_node = GridFunction(fes_level)\n",
    "#scene1 = Draw(TD_node, mesh, \"TD_node\")\n",
    "\n",
    "TD_pwc = GridFunction(pwc)\n",
    "#scene2 = Draw(TD_pwc, mesh, \"TD_pwc\")\n",
    "\n",
    "TDPosNeg_pwc = GridFunction(pwc)\n",
    "TDNegPos_pwc = GridFunction(pwc)\n",
    "\n",
    "cutRatio = GridFunction(pwc)\n",
    "\n",
    "# solve for current configuration\n",
    "SolvePDE(adjoint=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we can define a loop defining the levelset update:\n",
    "\\begin{align*}\n",
    "    \\psi_{k+1} = (1-s_k) \\frac{\\psi_k}{|\\psi_k|} + s_k \\frac{g_{\\psi_k}}{|g_{\\psi_k}|}\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene = Draw(psi)\n",
    "scene_cr = Draw(cutRatio, mesh, \"cutRatio\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iter_max = 20\n",
    "converged = False\n",
    "\n",
    "xm=0.\n",
    "ym=0.\n",
    "psi.Set( (x-xm)**2+(y-ym)**2-0.25**2)\n",
    "psinew.vec.data= psi.vec\n",
    "scene.Redraw()\n",
    "\n",
    "J = Integrate(Cost(gfu),mesh)\n",
    "\n",
    "with TaskManager():\n",
    "\n",
    "    for k in range(iter_max):\n",
    "        print(\"================ iteration \", k, \"===================\")\n",
    "\n",
    "        # copy new levelset data from psinew into psi\n",
    "        psi.vec.data = psinew.vec\n",
    "        scene.Redraw()\n",
    "        \n",
    "        \n",
    "        SolvePDE(adjoint=True)\n",
    "        \n",
    "        J_current = Integrate(Cost(gfu),mesh)\n",
    "\n",
    "        print( Integrate( (gfu-gfud)**2*dx, mesh) )\n",
    "\n",
    "        print(\"cost in beginning of iteration\", k, \": Cost = \", J_current)        \n",
    "        \n",
    "        # compute the piecewise constant topological derivative in each domain\n",
    "        TDPosNeg_pwc.Set( BetaPosNeg * grad(gfu) * grad(gfp)  + FPosNeg*gfp )\n",
    "        TDNegPos_pwc.Set( BetaNegPos * grad(gfu) * grad(gfp)  + FPosNeg*gfp )\n",
    "\n",
    "        # compute the cut ratio of the interface elements\n",
    "        InterpolateLevelSetToElems(psi, 1, 0, cutRatio, mesh, EPS)\n",
    "        scene_cr.Redraw()\n",
    "        \n",
    "        # compute the combined topological derivative using the cut ratio information\n",
    "        for j in range(len(TD_pwc.vec)):\n",
    "            TD_pwc.vec[j] = cutRatio.vec[j] * TDNegPos_pwc.vec[j] + (1-cutRatio.vec[j])*TDPosNeg_pwc.vec[j]\n",
    "        \n",
    "        TD_node.Set(TD_pwc)\n",
    "                \n",
    "        normTD = sqrt(Integrate(TD_node**2*dx, mesh)) # L2 norm of TD_node\n",
    "        TD_node.vec.data = 1/normTD * TD_node.vec # normalised TD_node\n",
    "        \n",
    "        normPsi = sqrt(Integrate(psi**2*dx, mesh)) # L2 norm of psi\n",
    "        psi.vec.data = 1/normPsi * psi.vec  # normalised psi\n",
    "        \n",
    "        linesearch = True\n",
    "       \n",
    "        for j in range(10): \n",
    "\n",
    "            # update the level set function\n",
    "            psinew.vec.data = (1-kappa)*psi.vec + kappa*TD_node.vec\n",
    "            \n",
    "            # update beta and f_rhs\n",
    "            InterpolateLevelSetToElems(psinew, beta1, beta2, beta, mesh, EPS)\n",
    "            InterpolateLevelSetToElems(psinew, f1, f2, f_rhs, mesh, EPS)\n",
    "\n",
    "            # solve PDE without adjoint\n",
    "            SolvePDE()\n",
    "            \n",
    "            Redraw(blocking=True)\n",
    " \n",
    "            Jnew = Integrate(Cost(gfu), mesh)\n",
    "            \n",
    "            if Jnew > J_current:\n",
    "                print(\"--------------------\")\n",
    "                print(\"-----line search ---\")\n",
    "                print(\"--------------------\")\n",
    "                kappa = kappa*0.8\n",
    "                print(\"kappa\", kappa)\n",
    "            else:\n",
    "                break\n",
    "        \n",
    "        Redraw(blocking=True)\n",
    "        print(\"----------- Jnew in  iteration \", k, \" = \", Jnew, \" (kappa = \", kappa, \")\")\n",
    "        print('')\n",
    "        print(\"iter\" + str(k) + \", Jnew = \" + str(Jnew) + \" (kappa = \", kappa, \")\")\n",
    "        kappa = min(1, kappa*1.2)\n",
    "\n",
    "        print(\"end of iter \" + str(k))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After 999 iterations: $J = $\n",
    "<center>\n",
    "   <img src=\"psi_iter999.png\" width=\"800\" aligh=\"center\"/>\n",
    "</center>"
   ]
  },
  {
   "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
}
