{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cannonball Example with Euler Integration and an External Optimizer\n",
    "\n",
    "This example will show you how to do the following:\n",
    "\n",
    "1. Use an OpenMDAO Problem inside of a loop to perform a broader calculation (in this case, simple integration.)\n",
    "2. Optimize a Problem using an external optimizer with a functional interface.\n",
    "3. Perform a complex step across an OpenMDAO Problem.\n",
    "\n",
    "In the example, we want to find the optimal angle to fire a cannon to maximize the distance it travels down\n",
    "range. We already know that a 45 degree angle will maximize the range in ideal conditions, but the model\n",
    "we will use also includes aerodynamic effects, so we will solve for the firing angle in the presence of a\n",
    "small amount of drag force.\n",
    "\n",
    "## Model\n",
    "\n",
    "A very general set of dynamics for a simplified aircraft are given by these differential equations:\n",
    "\n",
    "A very general set of dynamics for a simplified aircraft are given by these differential equations:\n",
    "\n",
    "$$\n",
    "  \\begin{align}\n",
    "    \\frac{dv}{dt} &= \\frac{T}{m} \\cos \\alpha - \\frac{D}{m} - g \\sin \\gamma \\\\\n",
    "    \\frac{d\\gamma}{dt} &= \\frac{T}{m v} \\sin \\alpha + \\frac{L}{m v} - \\frac{g \\cos \\gamma}{v} \\\\\n",
    "    \\frac{dh}{dt} &= v \\sin \\gamma \\\\\n",
    "    \\frac{dr}{dt} &= v \\cos \\gamma \\\\\n",
    "  \\end{align}\n",
    "$$\n",
    "\n",
    "We will use these for the cannonball dynamics, though angle of attack, lift, and thrust will all be\n",
    "zero.  We implement these equations into an OpenMDAO `ExplicitComponent` whose outputs are the\n",
    "time rates of change of the states which are velocity 'v', flight path angle 'gam', altitude\n",
    "'h' and range 'r'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class FlightPathEOM2D(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Computes the position and velocity equations of motion using a 2D flight path\n",
    "    parameterization of states per equations 4.42 - 4.46 of _[1].\n",
    "\n",
    "    References\n",
    "    ----------\n",
    "    .. [1] Bryson, Arthur Earl. Dynamic optimization. Vol. 1. Prentice Hall, p.172, 1999.\n",
    "    \"\"\"\n",
    "    def initialize(self):\n",
    "        self.options.declare('num_nodes', types=int)\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input(name='m', val=1.0, units='kg',\n",
    "                       desc='aircraft mass')\n",
    "        self.add_input(name='v', val=1.0, units='m/s',\n",
    "                       desc='aircraft velocity magnitude')\n",
    "        self.add_input(name='T', val=0.0, units='N',\n",
    "                       desc='thrust')\n",
    "        self.add_input(name='alpha', val=0.0, units='rad',\n",
    "                       desc='angle of attack')\n",
    "        self.add_input(name='L', val=0.0, units='N',\n",
    "                       desc='lift force')\n",
    "        self.add_input(name='D', val=0.0, units='N',\n",
    "                       desc='drag force')\n",
    "        self.add_input(name='gam', val=0.0, units='rad',\n",
    "                       desc='flight path angle')\n",
    "\n",
    "        self.add_output(name='v_dot', val=0.0, units='m/s**2',\n",
    "                        desc='rate of change of velocity magnitude')\n",
    "        self.add_output(name='gam_dot', val=0.0, units='rad/s',\n",
    "                        desc='rate of change of flight path angle')\n",
    "        self.add_output(name='h_dot', val=0.0, units='m/s',\n",
    "                        desc='rate of change of altitude')\n",
    "        self.add_output(name='r_dot', val=0.0, units='m/s',\n",
    "                        desc='rate of change of range')\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('v_dot', ['T', 'D', 'm', 'gam', 'alpha'])\n",
    "        self.declare_partials('gam_dot', ['T', 'L', 'm', 'gam', 'alpha', 'v'])\n",
    "        self.declare_partials(['h_dot', 'r_dot'], ['gam', 'v'])\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        g = 9.80665\n",
    "        m = inputs['m']\n",
    "        v = inputs['v']\n",
    "        T = inputs['T']\n",
    "        L = inputs['L']\n",
    "        D = inputs['D']\n",
    "        gam = inputs['gam']\n",
    "        alpha = inputs['alpha']\n",
    "\n",
    "        calpha = np.cos(alpha)\n",
    "        salpha = np.sin(alpha)\n",
    "\n",
    "        cgam = np.cos(gam)\n",
    "        sgam = np.sin(gam)\n",
    "\n",
    "        mv = m * v\n",
    "\n",
    "        outputs['v_dot'] = (T * calpha - D) / m - g * sgam\n",
    "        outputs['gam_dot'] = (T * salpha + L) / mv - (g / v) * cgam\n",
    "        outputs['h_dot'] = v * sgam\n",
    "        outputs['r_dot'] = v * cgam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need to compute the aerodynamic forces L and D, which are computed from the coefficients of lift and drag as well as the dynamic pressure.  This was implemented in two components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DynamicPressureComp(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input(name='rho', val=1.0, units='kg/m**3',\n",
    "                       desc='atmospheric density')\n",
    "        self.add_input(name='v', val=1.0, units='m/s',\n",
    "                       desc='air-relative velocity')\n",
    "\n",
    "        self.add_output(name='q', val=1.0, units='N/m**2',\n",
    "                        desc='dynamic pressure')\n",
    "\n",
    "        self.declare_partials(of='q', wrt='rho')\n",
    "        self.declare_partials(of='q', wrt='v')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['q'] = 0.5 * inputs['rho'] * inputs['v'] ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LiftDragForceComp(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Compute the aerodynamic forces on the vehicle in the wind axis frame\n",
    "    (lift, drag, cross) force.\n",
    "    \"\"\"\n",
    "    def initialize(self):\n",
    "        self.options.declare('num_nodes', types=int)\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input(name='CL', val=0.0,\n",
    "                       desc='lift coefficient')\n",
    "        self.add_input(name='CD', val=0.0,\n",
    "                       desc='drag coefficient')\n",
    "        self.add_input(name='q', val=0.0, units='N/m**2',\n",
    "                       desc='dynamic pressure')\n",
    "        self.add_input(name='S', val=0.0, units='m**2',\n",
    "                       desc='aerodynamic reference area')\n",
    "\n",
    "        self.add_output(name='f_lift', shape=(1, ), units='N',\n",
    "                        desc='aerodynamic lift force')\n",
    "        self.add_output(name='f_drag', shape=(1, ), units='N',\n",
    "                        desc='aerodynamic drag force')\n",
    "\n",
    "        self.declare_partials(of='f_lift', wrt=['q', 'S', 'CL'])\n",
    "        self.declare_partials(of='f_drag', wrt=['q', 'S', 'CD'])\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        q = inputs['q']\n",
    "        S = inputs['S']\n",
    "        CL = inputs['CL']\n",
    "        CD = inputs['CD']\n",
    "\n",
    "        qS = q * S\n",
    "\n",
    "        outputs['f_lift'] = qS * CL\n",
    "        outputs['f_drag'] = qS * CD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we put it all together in our top model. Given any cannonball flight path angle, and velocity we can compute the rates of change of all states by running this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CannonballODE(om.Group):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_subsystem(name='dynamic_pressure',\n",
    "                           subsys=DynamicPressureComp(),\n",
    "                           promotes=['*'])\n",
    "\n",
    "        self.add_subsystem(name='aero',\n",
    "                           subsys=LiftDragForceComp(),\n",
    "                           promotes_inputs=['*'])\n",
    "\n",
    "        self.add_subsystem(name='eom',\n",
    "                           subsys=FlightPathEOM2D(),\n",
    "                           promotes=['*'])\n",
    "\n",
    "        self.connect('aero.f_drag', 'D')\n",
    "        self.connect('aero.f_lift', 'L')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Time Integration\n",
    "\n",
    "Given an initial angle and velocity for the cannonball, one way we can compute the range is by integrating the equations of motion that are provided by OpenMDAO. A simple way to do this is to use the [Euler integration](https://en.wikipedia.org/wiki/Euler_method). We can perform this by choosing a time step, running the Problem at the starting location to compute the state rates, and then use Euler's method to compute the new cannonball state.  We can do this sequentially until the height turns negative, which means we have hit the ground sometime between this time step and the previous one. Finally, we use linear interpolation between the final point and the previous one to find the location where the height is zero.\n",
    "\n",
    "Note that Euler integration is a first-order method, so the accuracy will be linearly proportional to the step size. We use a 'dt' of 0.1 seconds here, which isn't particularly accurate, but runs quickly. In practice, you might use a higher order integration method here (e.g., Runge-Kutta.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "def eval_cannonball_range(gam_init, prob, complex_step=False):\n",
    "    \"\"\"\n",
    "    Compute distance given initial speed and angle of cannonball.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    gam_init : float\n",
    "        Initial cannonball firing angle in degrees.\n",
    "    prob : <Problem>\n",
    "        OpenMDAO problem that contains the equations of motion.\n",
    "    complex_step : bool\n",
    "        Set to True to perform complex step.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        Negative of range in m.\n",
    "    \"\"\"\n",
    "    dt = 0.1        # Time step\n",
    "    h_init = 1.0    # Height of cannon.\n",
    "    v_init = 100.0  # Initial cannonball velocity.\n",
    "    h_target = 0.0  #\n",
    "\n",
    "    v = v_init\n",
    "    gam = gam_init\n",
    "    h = h_init\n",
    "    r = 0.0\n",
    "    t = 0.0\n",
    "\n",
    "    if complex_step:\n",
    "        prob.set_complex_step_mode(True)\n",
    "\n",
    "    while h > h_target:\n",
    "\n",
    "        # Set values\n",
    "        prob.set_val('v', v)\n",
    "        prob.set_val('gam', gam, units='deg')\n",
    "\n",
    "        # Run the model\n",
    "        prob.run_model()\n",
    "\n",
    "        # Extract rates\n",
    "        v_dot = prob.get_val('v_dot')\n",
    "        gam_dot = prob.get_val('gam_dot', units='deg/s')\n",
    "        h_dot = prob.get_val('h_dot')\n",
    "        r_dot = prob.get_val('r_dot')\n",
    "\n",
    "        h_last = h\n",
    "        r_last = r\n",
    "\n",
    "        # Euler Integration\n",
    "        v = v + dt * v_dot\n",
    "        gam = gam + dt * gam_dot\n",
    "        h = h + dt * h_dot\n",
    "        r = r + dt * r_dot\n",
    "        t += dt\n",
    "        # print(v, gam, h, r)\n",
    "\n",
    "    # Linear interpolation between last two points to get the landing point accurate.\n",
    "    r_final = r_last + (r - r_last) * h_last / (h_last - h)\n",
    "\n",
    "    if complex_step:\n",
    "        prob.set_complex_step_mode(False)\n",
    "\n",
    "    #print(f\"Distance: {r_final}, Time: {t}, Angle: {gam_init}\")\n",
    "    return -r_final"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we have placed the integration code inside of a function that takes the initial angle as an argument and returns the negative of the computed range. This is the objective we wish to optimize with an external optimizer which requires a function that it can call to evaluate the objective.\n",
    "\n",
    "## Providing Derivatives for an External Optimizer\n",
    "\n",
    "The optimizer also allows you to specify a function to evaluate the gradient. If you do not provide one, it will use finite difference. We can improve the accuracy by performing complex step instead. OpenMDAO allows you to run a model in complex mode. When the mode is enabled on the Problem, you can use 'set_val' to set complex values and 'get_val' to retrieve them. In the code for the objective evaluation above, we turn this feature on by calling `prob.set_complex_step_mode(True)`.  Likewise, it is important to turn it off when not needed, and it should only be used when you are performing a complex step to compute the derivatives.\n",
    "\n",
    "The following function computes the total derivatives that the external optimizer needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient_cannonball_range(gam_init, prob):\n",
    "    \"\"\"\n",
    "    Uses complex step to compute gradient of range wrt initial angle.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    gam_init : float\n",
    "        Initial cannonball firing angle in degrees.\n",
    "    prob : <Problem>\n",
    "        OpenMDAO problem that contains the equations of motion.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        Derivative of range wrt initial angle in m/deg.\n",
    "    \"\"\"\n",
    "    step = 1.0e-14\n",
    "    dr_dgam = eval_cannonball_range(gam_init + step * 1j, prob, complex_step=True)\n",
    "    return dr_dgam.imag / step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the Optimization\n",
    "\n",
    "Now we can put everything together and run the optimization. Our optimizer is scipy.minimize for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "def eval_cannonball_range(gam_init, prob, complex_step=False):\n",
    "    \"\"\"\n",
    "    Compute distance given initial speed and angle of cannonball.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    gam_init : float\n",
    "        Initial cannonball firing angle in degrees.\n",
    "    prob : <Problem>\n",
    "        OpenMDAO problem that contains the equations of motion.\n",
    "    complex_step : bool\n",
    "        Set to True to perform complex step.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        Negative of range in m.\n",
    "    \"\"\"\n",
    "    dt = 0.1        # Time step\n",
    "    h_init = 1.0    # Height of cannon.\n",
    "    v_init = 100.0  # Initial cannonball velocity.\n",
    "    h_target = 0.0  #\n",
    "\n",
    "    v = v_init\n",
    "    gam = gam_init\n",
    "    h = h_init\n",
    "    r = 0.0\n",
    "    t = 0.0\n",
    "\n",
    "    if complex_step:\n",
    "        prob.set_complex_step_mode(True)\n",
    "\n",
    "    while h > h_target:\n",
    "\n",
    "        # Set values\n",
    "        prob.set_val('v', v)\n",
    "        prob.set_val('gam', gam, units='deg')\n",
    "\n",
    "        # Run the model\n",
    "        prob.run_model()\n",
    "\n",
    "        # Extract rates\n",
    "        v_dot = prob.get_val('v_dot')\n",
    "        gam_dot = prob.get_val('gam_dot', units='deg/s')\n",
    "        h_dot = prob.get_val('h_dot')\n",
    "        r_dot = prob.get_val('r_dot')\n",
    "\n",
    "        h_last = h\n",
    "        r_last = r\n",
    "\n",
    "        # Euler Integration\n",
    "        v = v + dt * v_dot\n",
    "        gam = gam + dt * gam_dot\n",
    "        h = h + dt * h_dot\n",
    "        r = r + dt * r_dot\n",
    "        t += dt\n",
    "        # print(v, gam, h, r)\n",
    "\n",
    "    # Linear interpolation between last two points to get the landing point accurate.\n",
    "    r_final = r_last + (r - r_last) * h_last / (h_last - h)\n",
    "\n",
    "    if complex_step:\n",
    "        prob.set_complex_step_mode(False)\n",
    "\n",
    "    #print(f\"Distance: {r_final}, Time: {t}, Angle: {gam_init}\")\n",
    "    return -r_final\n",
    "\n",
    "\n",
    "def gradient_cannonball_range(gam_init, prob):\n",
    "    \"\"\"\n",
    "    Uses complex step to compute gradient of range wrt initial angle.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    gam_init : float\n",
    "        Initial cannonball firing angle in degrees.\n",
    "    prob : <Problem>\n",
    "        OpenMDAO problem that contains the equations of motion.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        Derivative of range wrt initial angle in m/deg.\n",
    "    \"\"\"\n",
    "    step = 1.0e-14\n",
    "    dr_dgam = eval_cannonball_range(gam_init + step * 1j, prob, complex_step=True)\n",
    "    return dr_dgam.imag / step\n",
    "\n",
    "\n",
    "prob = om.Problem(model=CannonballODE())\n",
    "prob.setup(force_alloc_complex=True)\n",
    "\n",
    "# Set constants\n",
    "prob.set_val('CL', 0.0)                          # Lift Coefficient\n",
    "prob.set_val('CD', 0.05)                         # Drag Coefficient\n",
    "prob.set_val('S', 0.25 * np.pi, units='ft**2')   # Wetted Area (1 ft diameter ball)\n",
    "prob.set_val('rho', 1.225)                       # Atmospheric Density\n",
    "prob.set_val('m', 5.5)                           # Cannonball Mass\n",
    "\n",
    "prob.set_val('alpha', 0.0)                       # Angle of Attack (Not Applicable)\n",
    "prob.set_val('T', 0.0)                           # Thrust (Not Applicable)\n",
    "\n",
    "result = minimize(eval_cannonball_range, 27.0,\n",
    "                  method='SLSQP',\n",
    "                  jac=gradient_cannonball_range,\n",
    "                  args=(prob))\n",
    "\n",
    "print(result['x'])"
   ]
  },
  {
   "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(result['x'], 42.3810579, 1e-3)"
   ]
  }
 ],
 "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
}
