{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Topology optimization\n",
    "\n",
    "**Keywords**: *Topology optimization*, *weak and strong form*, *non-linear problem*, *Laplacian*, *grid topology*\n",
    "\n",
    "![](bending_animation.gif)\n",
    "\n",
    "*Figure 1*: Optimization of the bending beam. Evolution of the density for fixed total mass."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Introduction\n",
    "\n",
    "Topology optimization is the task of finding structures that are mechanically ideal.\n",
    "In this example we cover the bending beam, where we specify a load, boundary conditions and the total mass. Then, our\n",
    "objective is to find the most suitable geometry within the design space minimizing the compliance (i.e. the inverse stiffness) of the structure.\n",
    "We shortly introduce our simplified model for regular meshes. A detailed derivation of the method and advanced techniques\n",
    "can be found in [JanHacJun2019regularizedthermotopopt](@cite) and\n",
    "[BlaJanJun2022taylorwlsthermotopopt](@cite).\n",
    "\n",
    "We start by introducing the local, elementwise density $\\chi \\in [\\chi_{\\text{min}}, 1]$ of the material, where we choose\n",
    "$\\chi_{\\text{min}}$ slightly above zero to prevent numerical instabilities. Here, $\\chi = \\chi_{\\text{min}}$ means void and $\\chi=1$\n",
    "means bulk material. Then, we use a SIMP ansatz (solid isotropic material with penalization) for the stiffness tensor\n",
    "$C(\\chi) = \\chi^p C_0$, where $C_0$ is the stiffness of the bulk material. The SIMP exponent $p>1$ ensures that the\n",
    "model prefers the density values void and bulk before the intermediate values. The variational formulation then yields\n",
    "the modified Gibbs energy\n",
    "$$\n",
    "G = \\int_{\\Omega} \\frac{1}{2} \\chi^p \\varepsilon : C : \\varepsilon \\; \\text{d}V - \\int_{\\Omega} \\boldsymbol{f} \\cdot \\boldsymbol{u} \\; \\text{d}V - \\int_{\\partial\\Omega} \\boldsymbol{t} \\cdot \\boldsymbol{u} \\; \\text{d}A.\n",
    "$$\n",
    "Furthermore, we receive the evolution equation of the density\n",
    "and the additional Neumann boundary condition in the strong form\n",
    "$$\n",
    "p_\\chi + \\eta \\dot{\\chi} + \\lambda + \\gamma - \\beta \\nabla^2 \\chi \\ni 0 \\quad \\forall \\textbf{x} \\in \\Omega,\n",
    "$$\n",
    "$$\n",
    "\\beta \\nabla \\chi \\cdot \\textbf{n} = 0 \\quad \\forall \\textbf{x} \\in \\partial \\Omega,\n",
    "$$\n",
    "with the thermodynamic driving force\n",
    "$$\n",
    "p_\\chi = \\frac{1}{2} p \\chi^{p-1} \\varepsilon : C : \\varepsilon.\n",
    "$$\n",
    "We obtain the mechanical displacement field by applying the Finite Element Method to the weak form\n",
    "of the Gibbs energy using Ferrite. In contrast, we use the evolution equation (i.e. the strong form) to calculate\n",
    "the value of the density field $\\chi$. The advantage of this \"split\" approach is the very high computation speed.\n",
    "The evolution equation consists of the driving force, the damping parameter $\\eta$, the regularization parameter $\\beta$ times the Laplacian,\n",
    "which is necessary to avoid numerical issues like mesh dependence or checkerboarding, and the constraint parameters $\\lambda$, to keep the mass constant,\n",
    "and $\\gamma$, to avoid leaving the set $[\\chi_{\\text{min}}, 1]$. By including gradient regularization, it becomes necessary to calculate the Laplacian.\n",
    "The Finite Difference Method for square meshes with the edge length $\\Delta h$ approximates the Laplacian as follows:\n",
    "$$\n",
    "\\nabla^2 \\chi_p = \\frac{1}{(\\Delta h)^2} (\\chi_n + \\chi_s + \\chi_w + \\chi_e - 4 \\chi_p)\n",
    "$$\n",
    "Here, the indices refer to the different cardinal directions. Boundary element do not have neighbors in each direction. However, we can calculate\n",
    "the central difference to fulfill Neumann boundary condition. For example, if the element is on the left boundary, we have to fulfill\n",
    "$$\n",
    "\\nabla \\chi_p \\cdot \\textbf{n} = \\frac{1}{\\Delta h} (\\chi_w - \\chi_e) = 0\n",
    "$$\n",
    "from which follows $\\chi_w = \\chi_e$. Thus for boundary elements we can replace the value for the missing neighbor by the value of the opposite neighbor.\n",
    "In order to find the corresponding neighbor elements, we will make use of Ferrites grid topology funcionalities.\n",
    "\n",
    "## Commented Program\n",
    "We now solve the problem in Ferrite. What follows is a program spliced with comments.\n",
    "\n",
    "First we load all necessary packages."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Ferrite, SparseArrays, LinearAlgebra, Tensors, Printf"
   ],
   "metadata": {},
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we create a simple square grid of the size 2x1. We apply a fixed Dirichlet boundary condition\n",
    "to the left facet set, called `clamped`. On the right facet, we create a small set `traction`, where we\n",
    "will later apply a force in negative y-direction."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "create_grid (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "cell_type": "code",
   "source": [
    "function create_grid(n)\n",
    "    corners = [\n",
    "        Vec{2}((0.0, 0.0)),\n",
    "        Vec{2}((2.0, 0.0)),\n",
    "        Vec{2}((2.0, 1.0)),\n",
    "        Vec{2}((0.0, 1.0)),\n",
    "    ]\n",
    "    grid = generate_grid(Quadrilateral, (2 * n, n), corners)\n",
    "\n",
    "    # node-/facesets for boundary conditions\n",
    "    addnodeset!(grid, \"clamped\", x -> x[1] ≈ 0.0)\n",
    "    addfacetset!(grid, \"traction\", x -> x[1] ≈ 2.0 && norm(x[2] - 0.5) <= 0.05)\n",
    "    return grid\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we create the FE values, the DofHandler and the Dirichlet boundary condition."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "create_bc (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 3
    }
   ],
   "cell_type": "code",
   "source": [
    "function create_values()\n",
    "    # quadrature rules\n",
    "    qr = QuadratureRule{RefQuadrilateral}(2)\n",
    "    facet_qr = FacetQuadratureRule{RefQuadrilateral}(2)\n",
    "\n",
    "    # cell and facetvalues for u\n",
    "    ip = Lagrange{RefQuadrilateral, 1}()^2\n",
    "    cellvalues = CellValues(qr, ip)\n",
    "    facetvalues = FacetValues(facet_qr, ip)\n",
    "\n",
    "    return cellvalues, facetvalues\n",
    "end\n",
    "\n",
    "function create_dofhandler(grid)\n",
    "    dh = DofHandler(grid)\n",
    "    add!(dh, :u, Lagrange{RefQuadrilateral, 1}()^2) # displacement\n",
    "    close!(dh)\n",
    "    return dh\n",
    "end\n",
    "\n",
    "function create_bc(dh)\n",
    "    dbc = ConstraintHandler(dh)\n",
    "    add!(dbc, Dirichlet(:u, getnodeset(dh.grid, \"clamped\"), (x, t) -> zero(Vec{2}), [1, 2]))\n",
    "    close!(dbc)\n",
    "    t = 0.0\n",
    "    update!(dbc, t)\n",
    "    return dbc\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we define a struct to store all necessary material parameters (stiffness tensor of the bulk material\n",
    "and the parameters for topology optimization) and add a constructor to the struct to\n",
    "initialize it by using the common material parameters Young's modulus and Poisson number."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "Main.var\"##351\".MaterialParameters"
     },
     "metadata": {},
     "execution_count": 4
    }
   ],
   "cell_type": "code",
   "source": [
    "struct MaterialParameters{T, S <: SymmetricTensor{4, 2, T}}\n",
    "    C::S\n",
    "    χ_min::T\n",
    "    p::T\n",
    "    β::T\n",
    "    η::T\n",
    "end\n",
    "\n",
    "function MaterialParameters(E, ν, χ_min, p, β, η)\n",
    "    δ(i, j) = i == j ? 1.0 : 0.0 # helper function\n",
    "\n",
    "    G = E / 2(1 + ν) # =μ\n",
    "    λ = E * ν / (1 - ν^2) # correction for plane stress included\n",
    "\n",
    "    C = SymmetricTensor{4, 2}((i, j, k, l) -> λ * δ(i, j) * δ(k, l) + G * (δ(i, k) * δ(j, l) + δ(i, l) * δ(j, k)))\n",
    "    return MaterialParameters(C, χ_min, p, β, η)\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "To store the density and the strain required to calculate the driving forces, we create the struct\n",
    "`MaterialState`. We add a constructor to initialize the struct. The function `update_material_states!`\n",
    "updates the density values once we calculated the new values."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "update_material_states! (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 5
    }
   ],
   "cell_type": "code",
   "source": [
    "mutable struct MaterialState{T, S <: AbstractArray{SymmetricTensor{2, 2, T, 3}, 1}}\n",
    "    χ::T # density\n",
    "    ε::S # strain in each quadrature point\n",
    "end\n",
    "\n",
    "function MaterialState(ρ, n_qp)\n",
    "    return MaterialState(ρ, Array{SymmetricTensor{2, 2, Float64, 3}, 1}(undef, n_qp))\n",
    "end\n",
    "\n",
    "function update_material_states!(χn1, states, dh)\n",
    "    for (element, state) in zip(CellIterator(dh), states)\n",
    "        state.χ = χn1[cellid(element)]\n",
    "    end\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "source": [
    "Next, we define a function to calculate the driving forces for all elements.\n",
    "For this purpose, we iterate through all elements and calculate the average strain in each\n",
    "element. Then, we compute the driving force from the formula introduced at the beginning.\n",
    "We create a second function to collect the density in each element."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "compute_densities (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 6
    }
   ],
   "cell_type": "code",
   "source": [
    "function compute_driving_forces(states, mp, dh, χn)\n",
    "    pΨ = zeros(length(states))\n",
    "    for (element, state) in zip(CellIterator(dh), states)\n",
    "        i = cellid(element)\n",
    "        ε = sum(state.ε) / length(state.ε) # average element strain\n",
    "        pΨ[i] = 1 / 2 * mp.p * χn[i]^(mp.p - 1) * (ε ⊡ mp.C ⊡ ε)\n",
    "    end\n",
    "    return pΨ\n",
    "end\n",
    "\n",
    "function compute_densities(states, dh)\n",
    "    χn = zeros(length(states))\n",
    "    for (element, state) in zip(CellIterator(dh), states)\n",
    "        i = cellid(element)\n",
    "        χn[i] = state.χ\n",
    "    end\n",
    "    return χn\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "source": [
    "For the Laplacian we need some neighboorhood information which is constant throughout the analysis so we compute it once and cache it.\n",
    "We iterate through each facet of each element,\n",
    "obtaining the neighboring element by using the `getneighborhood` function. For boundary facets,\n",
    "the function call will return an empty object. In that case we use the dictionary to instead find the opposite\n",
    "facet, as discussed in the introduction."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "cache_neighborhood (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 7
    }
   ],
   "cell_type": "code",
   "source": [
    "function cache_neighborhood(dh, topology)\n",
    "    nbgs = Vector{Vector{Int}}(undef, getncells(dh.grid))\n",
    "    _nfacets = nfacets(dh.grid.cells[1])\n",
    "    opp = Dict(1 => 3, 2 => 4, 3 => 1, 4 => 2)\n",
    "\n",
    "    for element in CellIterator(dh)\n",
    "        nbg = zeros(Int, _nfacets)\n",
    "        i = cellid(element)\n",
    "        for j in 1:_nfacets\n",
    "            nbg_cellid = getneighborhood(topology, dh.grid, FacetIndex(i, j))\n",
    "            if !isempty(nbg_cellid)\n",
    "                nbg[j] = first(nbg_cellid)[1] # assuming only one facet neighbor per cell\n",
    "            else # boundary facet\n",
    "                nbg[j] = first(getneighborhood(topology, dh.grid, FacetIndex(i, opp[j])))[1]\n",
    "            end\n",
    "        end\n",
    "\n",
    "        nbgs[i] = nbg\n",
    "    end\n",
    "\n",
    "    return nbgs\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we calculate the Laplacian using the previously cached neighboorhood information."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "approximate_laplacian (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 8
    }
   ],
   "cell_type": "code",
   "source": [
    "function approximate_laplacian(nbgs, χn, Δh)\n",
    "    ∇²χ = zeros(length(nbgs))\n",
    "    for i in 1:length(nbgs)\n",
    "        nbg = nbgs[i]\n",
    "        ∇²χ[i] = (χn[nbg[1]] + χn[nbg[2]] + χn[nbg[3]] + χn[nbg[4]] - 4 * χn[i]) / (Δh^2)\n",
    "    end\n",
    "\n",
    "    return ∇²χ\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "source": [
    "For the iterative computation of the solution, a function is needed to update the densities in each element.\n",
    "To ensure that the mass is kept constant, we have to calculate the constraint\n",
    "parameter $\\lambda$, which we do via the bisection method. We repeat the calculation\n",
    "until the difference between the average density (calculated from the element-wise trial densities) and the target density nearly vanishes.\n",
    "By using the extremal values of $\\Delta \\chi$ as the starting interval, we guarantee that the method converges eventually."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "compute_χn1 (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 9
    }
   ],
   "cell_type": "code",
   "source": [
    "function compute_χn1(χn, Δχ, ρ, ηs, χ_min)\n",
    "    n_el = length(χn)\n",
    "\n",
    "    χ_trial = zeros(n_el)\n",
    "    ρ_trial = 0.0\n",
    "\n",
    "    λ_lower = minimum(Δχ) - ηs\n",
    "    λ_upper = maximum(Δχ) + ηs\n",
    "    λ_trial = 0.0\n",
    "\n",
    "    while abs(ρ - ρ_trial) > 1.0e-7\n",
    "        for i in 1:n_el\n",
    "            Δχt = 1 / ηs * (Δχ[i] - λ_trial)\n",
    "            χ_trial[i] = max(χ_min, min(1.0, χn[i] + Δχt))\n",
    "        end\n",
    "\n",
    "        ρ_trial = 0.0\n",
    "        for i in 1:n_el\n",
    "            ρ_trial += χ_trial[i] / n_el\n",
    "        end\n",
    "\n",
    "        if ρ_trial > ρ\n",
    "            λ_lower = λ_trial\n",
    "        elseif ρ_trial < ρ\n",
    "            λ_upper = λ_trial\n",
    "        end\n",
    "        λ_trial = 1 / 2 * (λ_upper + λ_lower)\n",
    "    end\n",
    "\n",
    "    return χ_trial\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "source": [
    "Lastly, we use the following helper function to compute the average driving force, which is later\n",
    "used to normalize the driving forces. This makes the used material parameters and numerical parameters independent\n",
    "of the problem."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "compute_average_driving_force (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 10
    }
   ],
   "cell_type": "code",
   "source": [
    "function compute_average_driving_force(mp, pΨ, χn)\n",
    "    n = length(pΨ)\n",
    "    w = zeros(n)\n",
    "\n",
    "    for i in 1:n\n",
    "        w[i] = (χn[i] - mp.χ_min) * (1 - χn[i])\n",
    "    end\n",
    "\n",
    "    p_Ω = sum(w .* pΨ) / sum(w) # average driving force\n",
    "\n",
    "    return p_Ω\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "source": [
    "Finally, we put everything together to update the density. The loop ensures the stability of the\n",
    "updated solution."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "update_density (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 11
    }
   ],
   "cell_type": "code",
   "source": [
    "function update_density(dh, states, mp, ρ, neighboorhoods, Δh)\n",
    "    n_j = Int(ceil(6 * mp.β / (mp.η * Δh^2))) # iterations needed for stability\n",
    "    χn = compute_densities(states, dh) # old density field\n",
    "    χn1 = zeros(length(χn))\n",
    "\n",
    "    for j in 1:n_j\n",
    "        ∇²χ = approximate_laplacian(neighboorhoods, χn, Δh) # Laplacian\n",
    "        pΨ = compute_driving_forces(states, mp, dh, χn) # driving forces\n",
    "        p_Ω = compute_average_driving_force(mp, pΨ, χn) # average driving force\n",
    "\n",
    "        Δχ = pΨ / p_Ω + mp.β * ∇²χ\n",
    "\n",
    "        χn1 = compute_χn1(χn, Δχ, ρ, mp.η, mp.χ_min)\n",
    "\n",
    "        if j < n_j\n",
    "            χn[:] = χn1[:]\n",
    "        end\n",
    "    end\n",
    "\n",
    "    return χn1\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we move on to the Finite Element part of the program. We use the following function to assemble our linear system."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "doassemble! (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 12
    }
   ],
   "cell_type": "code",
   "source": [
    "function doassemble!(cellvalues::CellValues, facetvalues::FacetValues, K::SparseMatrixCSC, grid::Grid, dh::DofHandler, mp::MaterialParameters, u, states)\n",
    "    r = zeros(ndofs(dh))\n",
    "    assembler = start_assemble(K, r)\n",
    "    nu = getnbasefunctions(cellvalues)\n",
    "\n",
    "    re = zeros(nu) # local residual vector\n",
    "    Ke = zeros(nu, nu) # local stiffness matrix\n",
    "\n",
    "    for (element, state) in zip(CellIterator(dh), states)\n",
    "        fill!(Ke, 0)\n",
    "        fill!(re, 0)\n",
    "\n",
    "        eldofs = celldofs(element)\n",
    "        ue = u[eldofs]\n",
    "\n",
    "        elmt!(Ke, re, element, cellvalues, facetvalues, grid, mp, ue, state)\n",
    "        assemble!(assembler, celldofs(element), Ke, re)\n",
    "    end\n",
    "\n",
    "    return K, r\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "source": [
    "The element routine is used to calculate the elementwise stiffness matrix and the residual. In contrast to a purely\n",
    "elastomechanic problem, for topology optimization we additionally use our material state to receive the density value of\n",
    "the element and to store the strain at each quadrature point."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "symmetrize_lower! (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 13
    }
   ],
   "cell_type": "code",
   "source": [
    "function elmt!(Ke, re, element, cellvalues, facetvalues, grid, mp, ue, state)\n",
    "    n_basefuncs = getnbasefunctions(cellvalues)\n",
    "    reinit!(cellvalues, element)\n",
    "    χ = state.χ\n",
    "\n",
    "    # We only assemble lower half triangle of the stiffness matrix and then symmetrize it.\n",
    "    @inbounds for q_point in 1:getnquadpoints(cellvalues)\n",
    "        dΩ = getdetJdV(cellvalues, q_point)\n",
    "        state.ε[q_point] = function_symmetric_gradient(cellvalues, q_point, ue)\n",
    "\n",
    "        for i in 1:n_basefuncs\n",
    "            δεi = shape_symmetric_gradient(cellvalues, q_point, i)\n",
    "            for j in 1:i\n",
    "                δεj = shape_symmetric_gradient(cellvalues, q_point, j)\n",
    "                Ke[i, j] += (χ)^(mp.p) * (δεi ⊡ mp.C ⊡ δεj) * dΩ\n",
    "            end\n",
    "            re[i] += (-δεi ⊡ ((χ)^(mp.p) * mp.C ⊡ state.ε[q_point])) * dΩ\n",
    "        end\n",
    "    end\n",
    "\n",
    "    symmetrize_lower!(Ke)\n",
    "\n",
    "    @inbounds for facet in 1:nfacets(getcells(grid, cellid(element)))\n",
    "        if (cellid(element), facet) ∈ getfacetset(grid, \"traction\")\n",
    "            reinit!(facetvalues, element, facet)\n",
    "            t = Vec((0.0, -1.0)) # force pointing downwards\n",
    "            for q_point in 1:getnquadpoints(facetvalues)\n",
    "                dΓ = getdetJdV(facetvalues, q_point)\n",
    "                for i in 1:n_basefuncs\n",
    "                    δu = shape_value(facetvalues, q_point, i)\n",
    "                    re[i] += (δu ⋅ t) * dΓ\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    return\n",
    "end\n",
    "\n",
    "function symmetrize_lower!(K)\n",
    "    for i in 1:size(K, 1)\n",
    "        for j in (i + 1):size(K, 1)\n",
    "            K[i, j] = K[j, i]\n",
    "        end\n",
    "    end\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 13
  },
  {
   "cell_type": "markdown",
   "source": [
    "We put everything together in the main function. Here the user may choose the radius parameter, which\n",
    "is related to the regularization parameter as $\\beta = ra^2$, the starting density, the number of elements in vertical direction and finally the\n",
    "name of the output. Additionally, the user may choose whether only the final design (default)\n",
    "or every iteration step is saved.\n",
    "\n",
    "First, we compute the material parameters and create the grid, DofHandler, boundary condition and FE values.\n",
    "Then we prepare the iterative Newton-Raphson method by pre-allocating all important vectors. Furthermore,\n",
    "we create material states for each element and construct the topology of the grid.\n",
    "\n",
    "During each iteration step, first we solve our FE problem in the Newton-Raphson loop. With the solution of the\n",
    "elastomechanic problem, we check for convergence of our topology design. The criteria has to be fulfilled twice in\n",
    "a row to avoid oscillations. If no convergence is reached yet, we update our design and prepare the next iteration step.\n",
    "Finally, we output the results in paraview and calculate the relative stiffness of the final design, i.e. how much how\n",
    "the stiffness increased compared to the starting point."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "topopt (generic function with 1 method)"
     },
     "metadata": {},
     "execution_count": 14
    }
   ],
   "cell_type": "code",
   "source": [
    "function topopt(ra, ρ, n, filename; output = :false)\n",
    "    # material\n",
    "    mp = MaterialParameters(210.0e3, 0.3, 1.0e-3, 3.0, ra^2, 15.0)\n",
    "\n",
    "    # grid, dofhandler, boundary condition\n",
    "    grid = create_grid(n)\n",
    "    dh = create_dofhandler(grid)\n",
    "    Δh = 1 / n # element edge length\n",
    "    dbc = create_bc(dh)\n",
    "\n",
    "    # cellvalues\n",
    "    cellvalues, facetvalues = create_values()\n",
    "\n",
    "    # Pre-allocate solution vectors, etc.\n",
    "    n_dofs = ndofs(dh) # total number of dofs\n",
    "    u = zeros(n_dofs) # solution vector\n",
    "    un = zeros(n_dofs) # previous solution vector\n",
    "\n",
    "    Δu = zeros(n_dofs)  # previous displacement correction\n",
    "    ΔΔu = zeros(n_dofs) # new displacement correction\n",
    "\n",
    "    # create material states\n",
    "    states = [MaterialState(ρ, getnquadpoints(cellvalues)) for _ in 1:getncells(dh.grid)]\n",
    "\n",
    "    χ = zeros(getncells(dh.grid))\n",
    "\n",
    "    r = zeros(n_dofs) # residual\n",
    "    K = allocate_matrix(dh) # stiffness matrix\n",
    "\n",
    "    i_max = 300 ## maximum number of iteration steps\n",
    "    tol = 1.0e-4\n",
    "    compliance = 0.0\n",
    "    compliance_0 = 0.0\n",
    "    compliance_n = 0.0\n",
    "    conv = :false\n",
    "\n",
    "    topology = ExclusiveTopology(grid)\n",
    "    neighboorhoods = cache_neighborhood(dh, topology)\n",
    "\n",
    "    # Newton-Raphson loop\n",
    "    NEWTON_TOL = 1.0e-8\n",
    "    print(\"\\n Starting Newton iterations\\n\")\n",
    "\n",
    "    for it in 1:i_max\n",
    "        apply_zero!(u, dbc)\n",
    "        newton_itr = -1\n",
    "\n",
    "        while true\n",
    "            newton_itr += 1\n",
    "\n",
    "            if newton_itr > 10\n",
    "                error(\"Reached maximum Newton iterations, aborting\")\n",
    "                break\n",
    "            end\n",
    "\n",
    "            # current guess\n",
    "            u .= un .+ Δu\n",
    "            K, r = doassemble!(cellvalues, facetvalues, K, grid, dh, mp, u, states)\n",
    "            norm_r = norm(r[Ferrite.free_dofs(dbc)])\n",
    "\n",
    "            if (norm_r) < NEWTON_TOL\n",
    "                break\n",
    "            end\n",
    "\n",
    "            apply_zero!(K, r, dbc)\n",
    "            ΔΔu = Symmetric(K) \\ r\n",
    "\n",
    "            apply_zero!(ΔΔu, dbc)\n",
    "            Δu .+= ΔΔu\n",
    "        end # of loop while NR-Iteration\n",
    "\n",
    "        # calculate compliance\n",
    "        compliance = 1 / 2 * u' * K * u\n",
    "\n",
    "        if it == 1\n",
    "            compliance_0 = compliance\n",
    "        end\n",
    "\n",
    "        # check convergence criterium (twice!)\n",
    "        if abs(compliance - compliance_n) / compliance < tol\n",
    "            if conv\n",
    "                println(\"Converged at iteration number: \", it)\n",
    "                break\n",
    "            else\n",
    "                conv = :true\n",
    "            end\n",
    "        else\n",
    "            conv = :false\n",
    "        end\n",
    "\n",
    "        # update density\n",
    "        χ = update_density(dh, states, mp, ρ, neighboorhoods, Δh)\n",
    "\n",
    "        # update old displacement, density and compliance\n",
    "        un .= u\n",
    "        Δu .= 0.0\n",
    "        update_material_states!(χ, states, dh)\n",
    "        compliance_n = compliance\n",
    "\n",
    "        # output during calculation\n",
    "        if output\n",
    "            i = @sprintf(\"%3.3i\", it)\n",
    "            filename_it = string(filename, \"_\", i)\n",
    "\n",
    "            VTKGridFile(filename_it, grid) do vtk\n",
    "                write_cell_data(vtk, χ, \"density\")\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "\n",
    "    # export converged results\n",
    "    if !output\n",
    "        VTKGridFile(filename, grid) do vtk\n",
    "            write_cell_data(vtk, χ, \"density\")\n",
    "        end\n",
    "    end\n",
    "    @printf \"Rel. stiffness: %.4f \\n\" compliance^(-1) / compliance_0^(-1)\n",
    "\n",
    "    return\n",
    "end"
   ],
   "metadata": {},
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "source": [
    "Lastly, we call our main function and compare the results. To create the\n",
    "complete output with all iteration steps, it is possible to set the output\n",
    "parameter to `true`."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "grid, χ =topopt(0.02, 0.5, 60, \"small_radius\"; output=:false);"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Starting Newton iterations\n",
      "Converged at iteration number: 65\n",
      "Rel. stiffness: 4.8466 \n",
      "  5.068164 seconds (2.43 M allocations: 2.016 GiB, 8.89% gc time, 6.18% compilation time)\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "@time topopt(0.03, 0.5, 60, \"large_radius\"; output = :false);\n",
    "#topopt(0.02, 0.5, 60, \"topopt_animation\"; output=:true); # can be used to create animations"
   ],
   "metadata": {},
   "execution_count": 15
  },
  {
   "cell_type": "markdown",
   "source": [
    "We observe, that the stiffness for the lower value of $ra$ is higher,\n",
    "but also requires more iterations until convergence and finer structures to be manufactured, as can be seen in Figure 2:\n",
    "\n",
    "![](bending.png)\n",
    "\n",
    "*Figure 2*: Optimization results of the bending beam for smaller (left) and larger (right) value of the regularization parameter $\\beta$.\n",
    "\n",
    "To prove mesh independence, the user could vary the mesh resolution and compare the results."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "---\n",
    "\n",
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ],
   "metadata": {}
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.11.6"
  },
  "kernelspec": {
   "name": "julia-1.11",
   "display_name": "Julia 1.11.6",
   "language": "julia"
  }
 },
 "nbformat": 4
}