{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true,
    "tags": [
     "active-ipynb",
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimizing the Thickness Distribution of a Cantilever Beam Using the Adjoint Method\n",
    "\n",
    "In this example, we optimize the thickness (height) distribution of a cantilever beam\n",
    "using the adjoint method to compute the gradient. We use Euler--Bernoulli beam theory and assume a rectangular section.\n",
    "\n",
    "## Background\n",
    "\n",
    "The optimization problem is:\n",
    "\n",
    "$$\n",
    "  \\begin{array}{r c l}\n",
    "    \\text{minimize} & & f^T d \\\\\n",
    "    \\text{with respect to} & & h \\\\\n",
    "    \\text{subject to} & & \\text{sum}(h) b L_0 = \\text{volume} \\\\\n",
    "  \\end{array}\n",
    "$$\n",
    "\n",
    "where $f$ is the vector of forces, $h$ is the vector of beam heights, and $L_0$ is the length of a single beam element.\n",
    "\n",
    "The displacements vector $d$ is given by\n",
    "\n",
    "$$\n",
    "  K d = f \n",
    "$$\n",
    "\n",
    "where $K$ is the stiffness matrix. However, in practice, we augment the linear system with Lagrange multipliers to apply the boundary constraints at the first node.\n",
    "\n",
    "Since our model contains a system of equations, we use the adjoint method to compute the gradient of the objective with respect to the beam height vector. The model is shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.test_suite.test_examples.beam_optimization.beam_group import BeamGroup\n",
    "\n",
    "if __name__ == '__main__':\n",
    "\n",
    "    E = 1.\n",
    "    L = 1.\n",
    "    b = 0.1\n",
    "    volume = 0.01\n",
    "\n",
    "    num_elements = 50\n",
    "\n",
    "    prob = om.Problem(model=BeamGroup(E=E, L=L, b=b, volume=volume, num_elements=num_elements),\n",
    "                      driver=om.ScipyOptimizeDriver(optimizer='SLSQP', tol=1e-9, disp=True))\n",
    "    prob.setup()\n",
    "    prob.run_driver()\n",
    "\n",
    "    om.n2(prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert np.linalg.norm(prob['h'] - [\n",
    "    0.14915754,  0.14764328,  0.14611321,  0.14456715,  0.14300421,  0.14142417,\n",
    "    0.13982611,  0.13820976,  0.13657406,  0.13491866,  0.13324268,  0.13154528,\n",
    "    0.12982575,  0.12808305,  0.12631658,  0.12452477,  0.12270701,  0.12086183,\n",
    "    0.11898809,  0.11708424,  0.11514904,  0.11318072,  0.11117762,  0.10913764,\n",
    "    0.10705891,  0.10493903,  0.10277539,  0.10056526,  0.09830546,  0.09599246,\n",
    "    0.09362243,  0.09119084,  0.08869265,  0.08612198,  0.08347229,  0.08073573,\n",
    "    0.07790323,  0.07496382,  0.07190453,  0.06870925,  0.0653583,   0.06182632,\n",
    "    0.05808044,  0.05407658,  0.04975295,  0.0450185,   0.03972912,  0.03363155,\n",
    "    0.02620192,  0.01610863\n",
    "]) < 1e-4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation: list of components\n",
    "\n",
    "There are 5 components that compute:\n",
    "\n",
    "1. Moment of inertia for each element\n",
    "2. Local stiffness matrix for each element\n",
    "3. Displacements from solution of the $Kd=f$ linear system augmented with the Lagrange multipliers\n",
    "4. Compliance\n",
    "5. Volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MomentOfInertiaComp(om.ExplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('num_elements', types=int)\n",
    "        self.options.declare('b')\n",
    "\n",
    "    def setup(self):\n",
    "        num_elements = self.options['num_elements']\n",
    "\n",
    "        self.add_input('h', shape=num_elements)\n",
    "        self.add_output('I', shape=num_elements)\n",
    "\n",
    "    def setup_partials(self):\n",
    "        rows = cols = np.arange(self.options['num_elements'])\n",
    "        self.declare_partials('I', 'h', rows=rows, cols=cols)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['I'] = 1./12. * self.options['b'] * inputs['h'] ** 3\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        partials['I', 'h'] = 1./4. * self.options['b'] * inputs['h'] ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LocalStiffnessMatrixComp(om.ExplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('num_elements', types=int)\n",
    "        self.options.declare('E')\n",
    "        self.options.declare('L')\n",
    "\n",
    "    def setup(self):\n",
    "        num_elements = self.options['num_elements']\n",
    "        E = self.options['E']\n",
    "        L = self.options['L']\n",
    "\n",
    "        self.add_input('I', shape=num_elements)\n",
    "        self.add_output('K_local', shape=(num_elements, 4, 4))\n",
    "\n",
    "        L0 = L / num_elements\n",
    "        coeffs = np.empty((4, 4))\n",
    "        coeffs[0, :] = [12, 6 * L0, -12, 6 * L0]\n",
    "        coeffs[1, :] = [6 * L0, 4 * L0 ** 2, -6 * L0, 2 * L0 ** 2]\n",
    "        coeffs[2, :] = [-12, -6 * L0, 12, -6 * L0]\n",
    "        coeffs[3, :] = [6 * L0, 2 * L0 ** 2, -6 * L0, 4 * L0 ** 2]\n",
    "        coeffs *= E / L0 ** 3\n",
    "\n",
    "        self.mtx = np.zeros((num_elements, 4, 4, num_elements))\n",
    "        for ind in range(num_elements):\n",
    "            self.mtx[ind, :, :, ind] = coeffs\n",
    "\n",
    "        self.declare_partials('K_local', 'I',\n",
    "            val=self.mtx.reshape(16 * num_elements, num_elements))\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['K_local'] = 0\n",
    "        for ind in range(self.options['num_elements']):\n",
    "            outputs['K_local'][ind, :, :] = self.mtx[ind, :, :, ind] * inputs['I'][ind]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse import coo_matrix\n",
    "from scipy.sparse.linalg import splu\n",
    "\n",
    "class StatesComp(om.ImplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('num_elements', types=int)\n",
    "        self.options.declare('force_vector', types=np.ndarray)\n",
    "\n",
    "    def setup(self):\n",
    "        num_elements = self.options['num_elements']\n",
    "        num_nodes = num_elements + 1\n",
    "        size = 2 * num_nodes + 2\n",
    "\n",
    "        self.add_input('K_local', shape=(num_elements, 4, 4))\n",
    "        self.add_output('d', shape=size)\n",
    "\n",
    "        cols = np.arange(16*num_elements)\n",
    "        rows = np.repeat(np.arange(4), 4)\n",
    "        rows = np.tile(rows, num_elements) + np.repeat(np.arange(num_elements), 16) * 2\n",
    "\n",
    "        self.declare_partials('d', 'K_local', rows=rows, cols=cols)\n",
    "        self.declare_partials('d', 'd')\n",
    "\n",
    "    def apply_nonlinear(self, inputs, outputs, residuals):\n",
    "        force_vector = np.concatenate([self.options['force_vector'], np.zeros(2)])\n",
    "\n",
    "        self.K = self.assemble_CSC_K(inputs)\n",
    "        residuals['d'] = self.K.dot(outputs['d'])  - force_vector\n",
    "\n",
    "    def solve_nonlinear(self, inputs, outputs):\n",
    "        force_vector = np.concatenate([self.options['force_vector'], np.zeros(2)])\n",
    "\n",
    "        self.K = self.assemble_CSC_K(inputs)\n",
    "        self.lu = splu(self.K)\n",
    "\n",
    "        outputs['d'] = self.lu.solve(force_vector)\n",
    "\n",
    "    def linearize(self, inputs, outputs, jacobian):\n",
    "        num_elements = self.options['num_elements']\n",
    "\n",
    "        self.K = self.assemble_CSC_K(inputs)\n",
    "        self.lu = splu(self.K)\n",
    "\n",
    "        i_elem = np.tile(np.arange(4), 4)\n",
    "        i_d = np.tile(i_elem, num_elements) + np.repeat(np.arange(num_elements), 16) * 2\n",
    "\n",
    "        jacobian['d', 'K_local'] = outputs['d'][i_d]\n",
    "\n",
    "        jacobian['d', 'd'] = self.K.toarray()\n",
    "\n",
    "    def solve_linear(self, d_outputs, d_residuals, mode):\n",
    "        if mode == 'fwd':\n",
    "            d_outputs['d'] = self.lu.solve(d_residuals['d'])\n",
    "        else:\n",
    "            d_residuals['d'] = self.lu.solve(d_outputs['d'])\n",
    "\n",
    "    def assemble_CSC_K(self, inputs):\n",
    "        \"\"\"\n",
    "        Assemble the stiffness matrix in sparse CSC format.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        ndarray\n",
    "            Stiffness matrix as dense ndarray.\n",
    "        \"\"\"\n",
    "        num_elements = self.options['num_elements']\n",
    "        num_nodes = num_elements + 1\n",
    "        num_entry = num_elements * 12 + 4\n",
    "        ndim = num_entry + 4\n",
    "\n",
    "        data = np.zeros((ndim, ), dtype=inputs._get_data().dtype)\n",
    "        cols = np.empty((ndim, ))\n",
    "        rows = np.empty((ndim, ))\n",
    "\n",
    "        # First element.\n",
    "        data[:16] = inputs['K_local'][0, :, :].flat\n",
    "        cols[:16] = np.tile(np.arange(4), 4)\n",
    "        rows[:16] = np.repeat(np.arange(4), 4)\n",
    "\n",
    "        j = 16\n",
    "        for ind in range(1, num_elements):\n",
    "            ind1 = 2 * ind\n",
    "            K = inputs['K_local'][ind, :, :]\n",
    "\n",
    "            # NW quadrant gets summed with previous connected element.\n",
    "            data[j-6:j-4] += K[0, :2]\n",
    "            data[j-2:j] += K[1, :2]\n",
    "\n",
    "            # NE quadrant\n",
    "            data[j:j+4] = K[:2, 2:].flat\n",
    "            rows[j:j+4] = np.array([ind1, ind1, ind1 + 1, ind1 + 1])\n",
    "            cols[j:j+4] = np.array([ind1 + 2, ind1 + 3, ind1 + 2, ind1 + 3])\n",
    "\n",
    "            # SE and SW quadrants together\n",
    "            data[j+4:j+12] = K[2:, :].flat\n",
    "            rows[j+4:j+12] = np.repeat(np.arange(ind1 + 2, ind1 + 4), 4)\n",
    "            cols[j+4:j+12] = np.tile(np.arange(ind1, ind1 + 4), 2)\n",
    "\n",
    "            j += 12\n",
    "\n",
    "        data[-4:] = 1.0\n",
    "        rows[-4] = 2 * num_nodes\n",
    "        rows[-3] = 2 * num_nodes + 1\n",
    "        rows[-2] = 0.0\n",
    "        rows[-1] = 1.0\n",
    "        cols[-4] = 0.0\n",
    "        cols[-3] = 1.0\n",
    "        cols[-2] = 2 * num_nodes\n",
    "        cols[-1] = 2 * num_nodes + 1\n",
    "\n",
    "        n_K = 2 * num_nodes + 2\n",
    "        return coo_matrix((data, (rows, cols)), shape=(n_K, n_K)).tocsc()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ComplianceComp(om.ExplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('num_elements', types=int)\n",
    "        self.options.declare('force_vector', types=np.ndarray)\n",
    "\n",
    "    def setup(self):\n",
    "        num_nodes = self.options['num_elements'] + 1\n",
    "\n",
    "        self.add_input('displacements', shape=2 * num_nodes)\n",
    "        self.add_output('compliance')\n",
    "\n",
    "    def setup_partials(self):\n",
    "        num_nodes = self.options['num_elements'] + 1\n",
    "        force_vector = self.options['force_vector']\n",
    "        self.declare_partials('compliance', 'displacements',\n",
    "                              val=force_vector.reshape((1, 2 * num_nodes)))\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['compliance'] = np.dot(self.options['force_vector'], inputs['displacements'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VolumeComp(om.ExplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('num_elements', types=int)\n",
    "        self.options.declare('b', default=1.)\n",
    "        self.options.declare('L')\n",
    "\n",
    "    def setup(self):\n",
    "        num_elements = self.options['num_elements']\n",
    "        b = self.options['b']\n",
    "        L = self.options['L']\n",
    "        L0 = L / num_elements\n",
    "\n",
    "        self.add_input('h', shape=num_elements)\n",
    "        self.add_output('volume')\n",
    "\n",
    "        self.declare_partials('volume', 'h', val=b * L0)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        L0 = self.options['L'] / self.options['num_elements']\n",
    "\n",
    "        outputs['volume'] = np.sum(inputs['h'] * self.options['b'] * L0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation: Optimization Script\n",
    "\n",
    "Here is the optimization script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.test_examples.beam_optimization.beam_group import BeamGroup\n",
    "\n",
    "E = 1.\n",
    "L = 1.\n",
    "b = 0.1\n",
    "volume = 0.01\n",
    "\n",
    "num_elements = 50\n",
    "\n",
    "prob = om.Problem(model=BeamGroup(E=E, L=L, b=b, volume=volume, num_elements=num_elements))\n",
    "\n",
    "prob.driver = om.ScipyOptimizeDriver()\n",
    "prob.driver.options['optimizer'] = 'SLSQP'\n",
    "prob.driver.options['tol'] = 1e-9\n",
    "prob.driver.options['disp'] = True\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.run_driver()\n",
    "\n",
    "print(prob['h'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![optimized](images/optimized.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob['h'],\n",
    "                 [0.14915754,  0.14764328,  0.14611321,  0.14456715,  0.14300421,  0.14142417,\n",
    "                  0.13982611,  0.13820976,  0.13657406,  0.13491866,  0.13324268,  0.13154528,\n",
    "                  0.12982575,  0.12808305,  0.12631658,  0.12452477,  0.12270701,  0.12086183,\n",
    "                  0.11898809,  0.11708424,  0.11514904,  0.11318072,  0.11117762,  0.10913764,\n",
    "                  0.10705891,  0.10493903,  0.10277539,  0.10056526,  0.09830546,  0.09599246,\n",
    "                  0.09362243,  0.09119084,  0.08869265,  0.08612198,  0.08347229,  0.08073573,\n",
    "                  0.07790323,  0.07496382,  0.07190453,  0.06870925,  0.0653583,   0.06182632,\n",
    "                  0.05808044,  0.05407658,  0.04975295,  0.0450185,   0.03972912,  0.03363155,\n",
    "                  0.02620192,  0.01610863], 1e-4)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
