{
 "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": [
    "# Revisiting the Beam Problem - Minimizing Stress with KS Constraints and BSplines\n",
    "\n",
    "The following example shows the optimization of a simple beam with rectangular cross section. The beam has been subdivided into\n",
    "elements along the length of the beam, and one end is fixed. The goal is to minimize the volume (and hence the mass of the\n",
    "homogeneous beam) by varying the thickness in each element without exceeding a maximum stress constraint while the beam is\n",
    "subject to multiple load cases, each one being a distributed force load that varies sinusoidally along the span.\n",
    "\n",
    "Constraining the bending stress on each element leads to a more computationally expensive derivative calculation, so we\n",
    "will use the `KSFunction` to reduce the stress vector for each load case to a single scalar value. To do so, we also need\n",
    "to insert an `ExecComp` component that converts the stress into a form where a negative value means it is satisfied, and\n",
    "a positive value means it is violated.\n",
    "\n",
    "The problem presented here is also an example of a multi-point implementation, where we create a separate instance of the\n",
    "parts of the calculation that are impacted by different load cases. This enables our model to take advantage of multiple\n",
    "processors when run in parallel.\n",
    "\n",
    "If we allow the optimizer to vary the thickness of each element, then we have a design variable vector that is as wide as the\n",
    "number of elements in the model. This may perform poorly if we have a large number of elements in the beam. If we assume that\n",
    "the optimal beam thickness is going to have a smooth continuous variation over the length of the beam, then it is a good\n",
    "candidate for using an interpolation component like `SplineComp` to reduce the number of design variables we need.\n",
    "\n",
    "For this example, we have 25 elements, but can reduce that to 5 control points for the optimizer's design variables by\n",
    "including the SplineComp.\n",
    "\n",
    "The code for the top system is this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src9\", get_code(\"openmdao.test_suite.test_examples.beam_optimization.components.local_stiffness_matrix_comp.LocalStiffnessMatrixComp\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `LocalStiffnessMatrixComp` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src9`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src10\", get_code(\"openmdao.test_suite.test_examples.beam_optimization.components.multi_states_comp.MultiStatesComp\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `MultiStatesComp` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src10`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src11\", get_code(\"openmdao.test_suite.test_examples.beam_optimization.components.moment_comp.MomentOfInertiaComp\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `MomentOfInertiaComp` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src11`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src12\", get_code(\"openmdao.test_suite.test_examples.beam_optimization.components.multi_stress_comp.MultiStressComp\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `MultiStressComp` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src12`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src13\", get_code(\"openmdao.test_suite.test_examples.beam_optimization.components.volume_comp.VolumeComp\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `VolumeComp` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src13`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This is a multipoint implementation of the beam optimization problem.\n",
    "\n",
    "This version minimizes volume while satisfying a max bending stress constraint in each element\n",
    "for each loadcase.\n",
    "\"\"\"\n",
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "from openmdao.test_suite.test_examples.beam_optimization.components.local_stiffness_matrix_comp import LocalStiffnessMatrixComp\n",
    "from openmdao.test_suite.test_examples.beam_optimization.components.moment_comp import MomentOfInertiaComp\n",
    "from openmdao.test_suite.test_examples.beam_optimization.components.multi_stress_comp import MultiStressComp\n",
    "from openmdao.test_suite.test_examples.beam_optimization.components.volume_comp import VolumeComp\n",
    "from openmdao.utils.spline_distributions import sine_distribution\n",
    "\n",
    "\n",
    "def divide_cases(ncases, nprocs):\n",
    "    \"\"\"\n",
    "    Divide up load cases among available procs.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    ncases : int\n",
    "        Number of load cases.\n",
    "    nprocs : int\n",
    "        Number of processors.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list of list of int\n",
    "        Integer case numbers for each proc.\n",
    "    \"\"\"\n",
    "    data = []\n",
    "    for j in range(nprocs):\n",
    "        data.append([])\n",
    "\n",
    "    wrap = 0\n",
    "    for j in range(ncases):\n",
    "        idx = j - wrap\n",
    "        if idx >= nprocs:\n",
    "            idx = 0\n",
    "            wrap = j\n",
    "\n",
    "        data[idx].append(j)\n",
    "\n",
    "    return data\n",
    "\n",
    "\n",
    "class MultipointBeamGroup(om.Group):\n",
    "    \"\"\"\n",
    "    System setup for minimization of volume (i.e., mass) subject to KS aggregated bending stress constraints.\n",
    "    \"\"\"\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('E')\n",
    "        self.options.declare('L')\n",
    "        self.options.declare('b')\n",
    "        self.options.declare('volume')\n",
    "        self.options.declare('max_bending')\n",
    "        self.options.declare('num_elements', 5)\n",
    "        self.options.declare('num_cp', 50)\n",
    "        self.options.declare('num_load_cases', 1)\n",
    "        self.options.declare('parallel_derivs', False, types=bool, allow_none=True)\n",
    "        self.options.declare('ks_add_constraint', default=False, types=bool)\n",
    "\n",
    "    def setup(self):\n",
    "        E = self.options['E']\n",
    "        L = self.options['L']\n",
    "        b = self.options['b']\n",
    "        # volume = self.options['volume']\n",
    "        max_bending = self.options['max_bending']\n",
    "        num_elements = self.options['num_elements']\n",
    "        num_nodes = num_elements + 1\n",
    "        num_cp = self.options['num_cp']\n",
    "        num_load_cases = self.options['num_load_cases']\n",
    "        parallel_derivs = self.options['parallel_derivs']\n",
    "\n",
    "        x_interp = sine_distribution(num_elements)\n",
    "        comp = om.SplineComp(method='bsplines', num_cp=num_cp, x_interp_val=x_interp)\n",
    "        comp.add_spline(y_cp_name='h_cp', y_interp_name='h')\n",
    "        self.add_subsystem('interp', comp)\n",
    "\n",
    "        I_comp = MomentOfInertiaComp(num_elements=num_elements, b=b)\n",
    "        self.add_subsystem('I_comp', I_comp)\n",
    "\n",
    "        comp = LocalStiffnessMatrixComp(num_elements=num_elements, E=E, L=L)\n",
    "        self.add_subsystem('local_stiffness_matrix_comp', comp)\n",
    "\n",
    "        # Parallel Subsystem for load cases.\n",
    "        par = self.add_subsystem('parallel', om.ParallelGroup())\n",
    "\n",
    "        # Determine how to split cases up over the available procs.\n",
    "        nprocs = self.comm.size\n",
    "        divide = divide_cases(num_load_cases, nprocs)\n",
    "\n",
    "        for j, this_proc in enumerate(divide):\n",
    "            num_rhs = len(this_proc)\n",
    "\n",
    "            name = 'sub_%d' % j\n",
    "            sub = par.add_subsystem(name, om.Group())\n",
    "\n",
    "            # Load is a sinusoidal distributed force of varying spatial frequency.\n",
    "            force_vector = np.zeros((2 * num_nodes, num_rhs))\n",
    "            for i, k in enumerate(this_proc):\n",
    "\n",
    "                end = 1.5 * np.pi\n",
    "                if num_load_cases > 1:\n",
    "                    end += k * 0.5 * np.pi / (num_load_cases - 1)\n",
    "\n",
    "                x = np.linspace(0, end, num_nodes)\n",
    "                f = - np.sin(x)\n",
    "                force_vector[0:-1:2, i] = f\n",
    "\n",
    "            comp = MultiStatesComp(num_elements=num_elements, force_vector=force_vector,\n",
    "                                   num_rhs=num_rhs)\n",
    "            sub.add_subsystem('states_comp', comp)\n",
    "\n",
    "            comp = MultiStressComp(num_elements=num_elements, E=E, num_rhs=num_rhs)\n",
    "            sub.add_subsystem('stress_comp', comp)\n",
    "\n",
    "            self.connect('local_stiffness_matrix_comp.K_local',\n",
    "                         'parallel.%s.states_comp.K_local' % name)\n",
    "\n",
    "            for k in range(num_rhs):\n",
    "                sub.connect('states_comp.d_%d' % k,\n",
    "                            'stress_comp.displacements_%d' % k,\n",
    "                            src_indices=np.arange(2 *num_nodes))\n",
    "\n",
    "                if parallel_derivs:\n",
    "                    color = 'red_%d' % k\n",
    "                else:\n",
    "                    color = None\n",
    "\n",
    "                comp = om.KSComp(width=num_elements, upper=max_bending,\n",
    "                                 add_constraint=self.options['ks_add_constraint'],\n",
    "                                 parallel_deriv_color=color)\n",
    "\n",
    "                sub.add_subsystem('KS_%d' % k, comp)\n",
    "\n",
    "                sub.connect('stress_comp.stress_%d' % k,\n",
    "                            'KS_%d.g' % k)\n",
    "\n",
    "                if not self.options['ks_add_constraint']:\n",
    "                    sub.add_constraint('KS_%d.KS' % k, upper=0.0,\n",
    "                                       parallel_deriv_color=color)\n",
    "\n",
    "        comp = VolumeComp(num_elements=num_elements, b=b, L=L)\n",
    "        self.add_subsystem('volume_comp', comp)\n",
    "\n",
    "        self.connect('interp.h', 'I_comp.h')\n",
    "        self.connect('interp.h', 'volume_comp.h')\n",
    "        self.connect('I_comp.I', 'local_stiffness_matrix_comp.I')\n",
    "\n",
    "        self.add_design_var('interp.h_cp', lower=1e-2, upper=10.)\n",
    "        self.add_objective('volume_comp.volume')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we run the model, and choose `ScipyOptimizeDriver` SLSQP to be our optimizer. At the conclusion of optimization, we print out the design variable, which is the thickness for each element."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.test_examples.beam_optimization.components.multi_states_comp import MultiStatesComp\n",
    "\n",
    "E = 1.\n",
    "L = 1.\n",
    "b = 0.1\n",
    "volume = 0.01\n",
    "max_bending = 100.0\n",
    "\n",
    "num_cp = 5\n",
    "num_elements = 25\n",
    "num_load_cases = 2\n",
    "\n",
    "model = MultipointBeamGroup(E=E, L=L, b=b, volume=volume, max_bending = max_bending,\n",
    "                            num_elements=num_elements, num_cp=num_cp,\n",
    "                            num_load_cases=num_load_cases)\n",
    "\n",
    "prob = om.Problem(model=model)\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(mode='rev')\n",
    "\n",
    "prob.run_driver()\n",
    "\n",
    "print(prob['interp.h'][0])"
   ]
  },
  {
   "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['interp.h'][0],\n",
    "                 [ 0.45632323,  0.45612552,  0.45543324,  0.45397058,  0.45134629,  0.44714397,\n",
    "                   0.4410258,   0.43283139,  0.42265378,  0.41087801,  0.3981731,   0.3854358,\n",
    "                   0.37369202,  0.36342186,  0.35289066,  0.34008777,  0.32362887,  0.30300358,\n",
    "                   0.27867837,  0.25204063,  0.22519409,  0.20063906,  0.18088818,  0.16807856,\n",
    "                   0.16364104], 1e-4)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "interpreter": {
   "hash": "c3e5a3598b81f1f111a5814b7aa7b9fddabe1e30fc374c51323a6f1030ed4ad3"
  },
  "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
}
