{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "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": [
    "# Hohmann Transfer Example - Optimizing a Spacecraft Manuever\n",
    "This example will demonstrate the use of OpenMDAO for optimizing a simple orbital mechanics problem. We seek the minimum possible delta-V to transfer a spacecraft from Low Earth Orbit (LEO) to geostationary orbit (GEO) using a two-impulse Hohmann Transfer.\n",
    "\n",
    "The Hohmann Transfer is a maneuver which minimizes the delta-V for transferring a spacecraft from one circular orbit to another. Hohmann transfers have a practical application in that they can be used to transfer satellites from LEO parking orbits to geostationary orbit.\n",
    "\n",
    "To do so, the vehicle first imparts a delta-V along the velocity vector while in LEO. This boosts apogee radius to the radius of the geostationary orbit (42164 km). In this model we will model this delta-V as an impulsive maneuver which changes the spacecraft’s velocity instantaneously.\n",
    "\n",
    "We will assume that the first impulse is performed at the ascending node in LEO. Thus perigee of the transfer orbit is coincident with the ascending node of the transfer orbit. Apogee of the transfer orbit is thus coincident with the descending node, where we will perform the second impulse.\n",
    "\n",
    "After the first impulse, the spacecraft coasts to apogee. Once there it performs a second impulsive burn along the velocity vector to raise perigee radius to the radius of GEO, thus circularizing the orbit.\n",
    "\n",
    "Simple, right? The issue is that, unless they launch from the equator, launch vehicles do not put satellites in a low Earth parking orbit with the same inclination as geostationary orbit. For instance, a due-east launch from Kennedy Space Center will result in a parking orbit with an inclination of 28.5 degrees. We therefore need to change the inclination of our satellite during its two impulsive burn maneuvers. The question is, what change in inclination at each burn will result in the minimum possible delta-V?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"hohmann_transfer.png\"/></center>\n",
    "<div align=\"center\"><b>An inclined Hohmann Transfer diagram</b></div>\n",
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The trajectory optimization problem can thus be stated as:\n",
    "\n",
    "$$\n",
    "    \\begin{align*}\n",
    "    Minimize J = {\\Delta} V \\\\\n",
    "    s.t. \\\\\n",
    "    \\Delta i_1 + \\Delta i_2 = 28.5^o \\\\\n",
    "    \\end{align*}\n",
    "$$\n",
    "\n",
    "The total ${\\Delta} V$ is the sum of the two impulsive ${\\Delta Vs}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"hohmann_dv1.png\"/></center>\n",
    "<div align=\"center\"><b>Vector diagram of the first impulsive &Delta;V</b></div>\n",
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The component of the ${\\Delta V}$ in the orbital plane is along the\n",
    "local horizontal plane.  The orbit-normal component is in the\n",
    "direction of the desired inclination change.  Knowing the\n",
    "velocity magnitude before (${v_c}$) and after (${v_p}$) the impulse, and the\n",
    "change in inclination due to the impulse (${\\Delta i}$), the ${\\Delta V}$\n",
    "is then computed from the law of cosines:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ \\Delta V_1 = \\sqrt{v_c^2 + v_p^2 - 2 v_c v_p \\cos{\\Delta i}} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the first impulse, ${v_1}$ is the circular velocity in LEO.  In\n",
    "this case ${v_c}$ refers to the circular velocity in geostationary\n",
    "orbit, and ${v_a}$ is the velocity at apogee of the transfer\n",
    "orbit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the circular velocity in either orbit from the following equation:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ v_c = \\sqrt{\\mu/r} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where ${\\mu}$ is the gravitational parameter of the Earth\n",
    "and ${r}$ is the distance from the center of the Earth."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The velocity after the first impulse is the periapsis velocity\n",
    "of the transfer orbit.  This can be solved for based on what we\n",
    "know about the orbit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The specific angular momentum of the transfer orbit is constant.\n",
    "At periapsis, it is simply the product of the velocity and radius.\n",
    "Therefore, rearranging we have:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$    v_p = \\frac{h}{r_p} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The specific angular momentum can also be computed as:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$    h = \\sqrt{p \\mu} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Where ${p}$ is the semilatus rectum of the orbit and ${\\mu}$ is\n",
    "the gravitational parameter of the central body."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The semilatus rectum is computed as:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$    p = a*(1.0-e^2)  $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Where ${a}$ and ${e}$ are the semi-major axis and eccentricity of the transfer orbit, respectively.\n",
    "Since we know ${r_a}$ and ${r_p}$ of the transfer orbit, it's semimajor axis is simply:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ e = (a-r_p)/a $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The eccentricity is known by the relationship of ${a}$ and ${e}$ to ${r_p}$ (or ${r_a}$):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$    a = (r_a+r_p)/2.0 $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus we can compute periapsis velocity based on the periapsis and apoapsis\n",
    "radii of the transfer orbit, and the gravitational parameter of the central body."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the second impulse, the final velocity is the circular velocity of the\n",
    "final orbit, which can be computed in the same way as the circular velocity\n",
    "of the initial orbit.  The initial velocity at the second impulse is the\n",
    "apoapsis velocity of the transfer orbit, which is:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"hohmann_dv2.png\"/></center>\n",
    "<div align=\"center\"><b>Vector diagram of the second impulsive &Delta;V</b></div>\n",
    "<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$    \\Delta V = \\sqrt{ v_a^2 + v_c^2 - 2 v_a v_c \\cos{\\Delta i} } $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$   v_a = \\frac{h}{r_a} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having already computed the specific angular momentum of the transfer orbit, this is\n",
    "easily computed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we have the necessary calculations to compute the ${\\Delta V}$ of the Hohmann\n",
    "transfer with a plane change."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Components\n",
    "The first component we define computes the circular velocity given the radius from the center of the central body and the gravitational parameter of the central body."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class VCircComp(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Computes the circular orbit velocity given a radius and gravitational\n",
    "    parameter.\n",
    "    \"\"\"\n",
    "    def initialize(self):\n",
    "        pass\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('r',\n",
    "                       val=1.0,\n",
    "                       desc='Radius from central body',\n",
    "                       units='km')\n",
    "\n",
    "        self.add_input('mu',\n",
    "                       val=1.0,\n",
    "                       desc='Gravitational parameter of central body',\n",
    "                       units='km**3/s**2')\n",
    "\n",
    "        self.add_output('vcirc',\n",
    "                        val=1.0,\n",
    "                        desc='Circular orbit velocity at given radius '\n",
    "                             'and gravitational parameter',\n",
    "                        units='km/s')\n",
    "\n",
    "        self.declare_partials(of='vcirc', wrt='r')\n",
    "        self.declare_partials(of='vcirc', wrt='mu')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        r = inputs['r']\n",
    "        mu = inputs['mu']\n",
    "\n",
    "        outputs['vcirc'] = np.sqrt(mu / r)\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        r = inputs['r']\n",
    "        mu = inputs['mu']\n",
    "        vcirc = np.sqrt(mu / r)\n",
    "\n",
    "        partials['vcirc', 'mu'] = 0.5 / (r * vcirc)\n",
    "        partials['vcirc', 'r'] = -0.5 * mu / (vcirc * r ** 2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The transfer orbit component computes the velocity magnitude at periapsis\n",
    "and apoapsis of an orbit, given the radii of periapsis and apoapsis, and\n",
    "the gravitational parameter of the central body."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransferOrbitComp(om.ExplicitComponent):\n",
    "    def initialize(self):\n",
    "        pass\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('mu',\n",
    "                       val=398600.4418,\n",
    "                       desc='Gravitational parameter of central body',\n",
    "                       units='km**3/s**2')\n",
    "        self.add_input('rp', val=7000.0, desc='periapsis radius', units='km')\n",
    "        self.add_input('ra', val=42164.0, desc='apoapsis radius', units='km')\n",
    "\n",
    "        self.add_output('vp', val=0.0, desc='periapsis velocity', units='km/s')\n",
    "        self.add_output('va', val=0.0, desc='apoapsis velocity', units='km/s')\n",
    "\n",
    "        # We're going to be lazy and ask OpenMDAO to approximate our\n",
    "        # partials with finite differencing here.\n",
    "        self.declare_partials(of='*', wrt='*', method='fd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        mu = inputs['mu']\n",
    "        rp = inputs['rp']\n",
    "        ra = inputs['ra']\n",
    "\n",
    "        a = (ra + rp) / 2.0\n",
    "        e = (a - rp) / a\n",
    "        p = a * (1.0 - e ** 2)\n",
    "        h = np.sqrt(mu * p)\n",
    "\n",
    "        outputs['vp'] = h / rp\n",
    "        outputs['va'] = h / ra\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The delta-V component is used to compute the delta-V performed in changing\n",
    "the velocity vector, giving the magnitudes of the initial and final velocities\n",
    "and the angle between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeltaVComp(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Compute the delta-V performed given the magnitude of two velocities\n",
    "    and the angle between them.\n",
    "    \"\"\"\n",
    "    def initialize(self):\n",
    "        pass\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('v1', val=1.0, desc='Initial velocity', units='km/s')\n",
    "        self.add_input('v2', val=1.0, desc='Final velocity', units='km/s')\n",
    "        self.add_input('dinc', val=1.0, desc='Plane change', units='rad')\n",
    "\n",
    "        # Note:  We're going to use trigonometric functions on dinc.  The\n",
    "        # automatic unit conversion in OpenMDAO comes in handy here.\n",
    "\n",
    "        self.add_output('delta_v', val=0.0, desc='Delta-V', units='km/s')\n",
    "\n",
    "        self.declare_partials(of='delta_v', wrt='v1')\n",
    "        self.declare_partials(of='delta_v', wrt='v2')\n",
    "        self.declare_partials(of='delta_v', wrt='dinc')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        v1 = inputs['v1']\n",
    "        v2 = inputs['v2']\n",
    "        dinc = inputs['dinc']\n",
    "\n",
    "        outputs['delta_v'] = np.sqrt(v1 ** 2 + v2 ** 2 - 2.0 * v1 * v2 * np.cos(dinc))\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        v1 = inputs['v1']\n",
    "        v2 = inputs['v2']\n",
    "        dinc = inputs['dinc']\n",
    "\n",
    "        delta_v = np.sqrt(v1 ** 2 + v2 ** 2 - 2.0 * v1 * v2 * np.cos(dinc))\n",
    "\n",
    "        partials['delta_v', 'v1'] = 0.5 / delta_v * (2 * v1 - 2 * v2 * np.cos(dinc))\n",
    "        partials['delta_v', 'v2'] = 0.5 / delta_v * (2 * v2 - 2 * v1 * np.cos(dinc))\n",
    "        partials['delta_v', 'dinc'] = 0.5 / delta_v * (2 * v1 * v2 * np.sin(dinc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Putting it all together\n",
    "Now we assemble the model for our problem.\n",
    "\n",
    "Two instances of VCircComp are used to compute the velocity of the\n",
    "spacecraft in the initial and final circular orbits.\n",
    "\n",
    "The TransferOrbitComp is used to compute the periapsis and apoapsis velocity\n",
    "of the spacecraft in the transfer orbit.\n",
    "\n",
    "Now we can use the DeltaVComp to provide the magnitude of the delta-V\n",
    "at each of the two impulses.\n",
    "\n",
    "We use two ExecComps to provide some simple calculations.  One\n",
    "sums the delta-Vs of the two impulses to provide the total delta-V of the\n",
    "transfer.  We will use this as the objective for the optimization.\n",
    "\n",
    "The other ExecComp sums up the inclination change at each impulse.  We\n",
    "will provide this to the driver as a constraint to ensure that our total\n",
    "inclination change meets our requirements.\n",
    "\n",
    "Lastly, we provide unambiguous values and units for the gravitational parameter, the radii\n",
    "of the two circular orbits, and the delta-V to be performed at each of the\n",
    "two impulses.\n",
    "\n",
    "We will use the initial and final radii of the orbits, and the inclination\n",
    "change at each of the two impulses as our design variables.\n",
    "\n",
    "To run the model, we provide values for the design variables and invoke `run_model`.\n",
    "\n",
    "To find the optimal solution for the model, we invoke `run_driver`, where we have\n",
    "defined the driver of the problem to be `ScipyOptimizeDriver`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = om.Problem()\n",
    "\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('leo', subsys=VCircComp(), promotes_inputs=[('r', 'r1'), 'mu'])\n",
    "model.add_subsystem('geo', subsys=VCircComp(), promotes_inputs=[('r', 'r2'), 'mu'])\n",
    "\n",
    "model.add_subsystem('transfer', subsys=TransferOrbitComp(),\n",
    "                    promotes_inputs=[('rp', 'r1'), ('ra', 'r2'), 'mu'])\n",
    "\n",
    "model.add_subsystem('dv1', subsys=DeltaVComp(), promotes_inputs=[('dinc', 'dinc1')])\n",
    "\n",
    "model.connect('leo.vcirc', 'dv1.v1')\n",
    "model.connect('transfer.vp', 'dv1.v2')\n",
    "\n",
    "model.add_subsystem('dv2', subsys=DeltaVComp(), promotes_inputs=[('dinc', 'dinc2')])\n",
    "\n",
    "model.connect('transfer.va', 'dv2.v1')\n",
    "model.connect('geo.vcirc', 'dv2.v2')\n",
    "\n",
    "model.add_subsystem('dv_total',\n",
    "                    subsys=om.ExecComp('delta_v=dv1+dv2',\n",
    "                                       delta_v={'units': 'km/s'},\n",
    "                                       dv1={'units': 'km/s'},\n",
    "                                       dv2={'units': 'km/s'}),\n",
    "                    promotes=['delta_v'])\n",
    "\n",
    "model.connect('dv1.delta_v', 'dv_total.dv1')\n",
    "model.connect('dv2.delta_v', 'dv_total.dv2')\n",
    "\n",
    "model.add_subsystem('dinc_total',\n",
    "                    subsys=om.ExecComp('dinc=dinc1+dinc2',\n",
    "                                       dinc={'units': 'deg'},\n",
    "                                       dinc1={'units': 'deg'},\n",
    "                                       dinc2={'units': 'deg'}),\n",
    "                    promotes=['dinc', 'dinc1', 'dinc2'])\n",
    "\n",
    "prob.driver = om.ScipyOptimizeDriver()\n",
    "\n",
    "model.add_design_var('dinc1', lower=0, upper=28.5)\n",
    "model.add_design_var('dinc2', lower=0, upper=28.5)\n",
    "model.add_constraint('dinc', lower=28.5, upper=28.5, scaler=1.0)\n",
    "model.add_objective('delta_v', scaler=1.0)\n",
    "\n",
    "# set defaults for our promoted variables to remove ambiguities in value and/or units\n",
    "model.set_input_defaults('r1', val=6778.0)\n",
    "model.set_input_defaults('r2', val=42164.0)\n",
    "model.set_input_defaults('mu', val=398600.4418)\n",
    "model.set_input_defaults('dinc1', val=0., units='deg')\n",
    "model.set_input_defaults('dinc2', val=28.5, units='deg')\n",
    "\n",
    "# Setup the problem\n",
    "prob.setup()\n",
    "\n",
    "# Execute the model with the given inputs\n",
    "prob.run_model()\n",
    "\n",
    "print('Delta-V (km/s):', prob['delta_v'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Inclination change split (deg):', prob['dinc1'][0], prob['dinc2'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob.run_driver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Optimized Delta-V (km/s):', prob['delta_v'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Inclination change split (deg):', prob['dinc1'][0], prob['dinc2'][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "We built a model representing a Hohmann transfer with a plane change.  This model\n",
    "utilized components with both analytic partial derivatives and approximated partials\n",
    "using finite differencing.  We utilized ExecComps for some simple calculations\n",
    "to reduce the amount of code we needed to write.  Finally, we used this model\n",
    "to demonstrate that performing the necessary plane change entirely at apoapsis is\n",
    "somewhat less optimal, from a delta-V standpoint, than performing some of the plane\n",
    "change at the first impulse."
   ]
  }
 ],
 "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
}
