{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "active-ipynb",
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "# This cell is mandatory in all Dymos documentation notebooks.\n",
    "missing_packages = []\n",
    "try:\n",
    "    import openmdao.api as om\n",
    "except ImportError:\n",
    "    if 'google.colab' in str(get_ipython()):\n",
    "        !python -m pip install openmdao[notebooks]\n",
    "    else:\n",
    "        missing_packages.append('openmdao')\n",
    "try:\n",
    "    import dymos as dm\n",
    "except ImportError:\n",
    "    if 'google.colab' in str(get_ipython()):\n",
    "        !python -m pip install dymos\n",
    "    else:\n",
    "        missing_packages.append('dymos')\n",
    "try:\n",
    "    import pyoptsparse\n",
    "except ImportError:\n",
    "    if 'google.colab' in str(get_ipython()):\n",
    "        !pip install -q condacolab\n",
    "        import condacolab\n",
    "        condacolab.install_miniconda()\n",
    "        !conda install -c conda-forge pyoptsparse\n",
    "    else:\n",
    "        missing_packages.append('pyoptsparse')\n",
    "if missing_packages:\n",
    "    raise EnvironmentError('This notebook requires the following packages '\n",
    "                           'please install them and restart this notebook\\'s runtime: {\",\".join(missing_packages)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Brachistochrone with externally-sourced initial state values\n",
    "\n",
    "```{admonition} Things you'll learn through this example\n",
    "- How to link phase initial time and duration.\n",
    "```\n",
    "\n",
    "This is another modification of the brachistochrone in which the target initial time and duration of a phase is provided by an external source (an IndepVarComp in this case).\n",
    "\n",
    "Rather than the external value being directly connected to the phase, the values are \"linked\" via `state_ivc`'s `t_initial` and `t_duration`.\n",
    "\n",
    "The following script fully defines the brachistochrone problem with Dymos and solves it. A new `IndepVarComp` is added before the trajectory which provides `t_initial` and `t_duration`. These two outputs are connected directly into the phase. Also, it's important to set `input_duration` and `input_initial` to True inside of `set_time_options`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "import dymos as dm\n",
    "import matplotlib.pyplot as plt\n",
    "from dymos.examples.brachistochrone.brachistochrone_ode import BrachistochroneODE\n",
    "#\n",
    "# Define the OpenMDAO problem\n",
    "#\n",
    "p = om.Problem(model=om.Group())\n",
    "\n",
    "# Instantiate the transcription so we can get the number of nodes from it while\n",
    "# building the problem.\n",
    "tx = dm.GaussLobatto(num_segments=10, order=3)\n",
    "\n",
    "# Add an indep var comp to provide the external control values\n",
    "ivc = p.model.add_subsystem('states_ivc', om.IndepVarComp(), promotes_outputs=['*'])\n",
    "\n",
    "# Add the output to provide the values of theta at the control input nodes of the transcription.\n",
    "# ivc.add_output('x0', shape=(1,), units='m')\n",
    "ivc.add_output('t_initial', val=0.0, units='s')\n",
    "ivc.add_output('t_duration', val=10., units='s')\n",
    "ivc.add_design_var('t_duration', units='s', lower=0.1, upper=10.)\n",
    "\n",
    "# Connect x0 to the state error component so we can constrain the given value of x0\n",
    "# to be equal to the value chosen in the phase.\n",
    "# p.model.connect('x0', 'state_error_comp.x0_target')\n",
    "# p.model.connect('traj.phase0.timeseries.x', 'state_error_comp.x0_actual',\n",
    "#                 src_indices=[0])\n",
    "p.model.connect('t_initial', 'traj.phase0.t_initial')\n",
    "p.model.connect('t_duration', 'traj.phase0.t_duration')\n",
    "#\n",
    "# Define a Trajectory object\n",
    "#\n",
    "traj = dm.Trajectory()\n",
    "p.model.add_subsystem('traj', subsys=traj)\n",
    "\n",
    "#\n",
    "# Define a Dymos Phase object with GaussLobatto Transcription\n",
    "#\n",
    "phase = dm.Phase(ode_class=BrachistochroneODE,\n",
    "                 transcription=tx)\n",
    "traj.add_phase(name='phase0', phase=phase)\n",
    "\n",
    "#\n",
    "# Set the time options\n",
    "# Time has no targets in our ODE.\n",
    "# We fix the initial time so that the it is not a design variable in the optimization.\n",
    "# The duration of the phase is allowed to be optimized, but is bounded on [0.5, 10]\n",
    "# which is set above in state_ivc\n",
    "phase.set_time_options(input_duration=True, input_initial=True, units='s')\n",
    "\n",
    "#\n",
    "# Set the time options\n",
    "# Initial values of positions and velocity are all fixed.\n",
    "# The final value of position are fixed, but the final velocity is a free variable.\n",
    "# The equations of motion are not functions of position, so 'x' and 'y' have no targets.\n",
    "# The rate source points to the output in the ODE which provides the time derivative of the\n",
    "# given state.\n",
    "phase.add_state('x', fix_initial=True, fix_final=True, units='m', rate_source='xdot')\n",
    "phase.add_state('y', fix_initial=True, fix_final=True, units='m', rate_source='ydot')\n",
    "phase.add_state('v', fix_initial=True, fix_final=False, units='m/s',\n",
    "                rate_source='vdot', targets=['v'])\n",
    "\n",
    "# Define theta as a control.\n",
    "# Use opt=False to allow it to be connected to an external source.\n",
    "# Arguments lower and upper are no longer valid for an input control.\n",
    "phase.add_control(name='theta', units='rad', targets=['theta'])\n",
    "# Minimize final time.\n",
    "phase.add_objective('time', loc='final')\n",
    "\n",
    "# Set the driver.\n",
    "p.driver = om.ScipyOptimizeDriver()\n",
    "\n",
    "# Allow OpenMDAO to automatically determine our sparsity pattern.\n",
    "# Doing so can significant speed up the execution of Dymos.\n",
    "p.driver.declare_coloring()\n",
    "\n",
    "# Setup the problem\n",
    "p.setup(check=True)\n",
    "\n",
    "# Now that the OpenMDAO problem is setup, we can set the values of the states.\n",
    "# p.set_val('x0', 0.0, units='m')\n",
    "# Here we're intentially setting the intiial x value to something other than zero, just\n",
    "# to demonstrate that the optimizer brings it back in line with the value of x0 set above.\n",
    "p.set_val('traj.phase0.states:x',\n",
    "          phase.interp('x', [0, 10]),\n",
    "          units='m')\n",
    "p.set_val('traj.phase0.states:y',\n",
    "          phase.interp('y', [10, 5]),\n",
    "          units='m')\n",
    "p.set_val('traj.phase0.states:v',\n",
    "          phase.interp('v', [0, 5]),\n",
    "          units='m/s')\n",
    "p.set_val('traj.phase0.controls:theta',\n",
    "          phase.interp('theta', [90, 90]),\n",
    "          units='deg')\n",
    "\n",
    "# Run the driver to solve the problem\n",
    "dm.run_problem(p, simulate=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting the results\n",
    "\n",
    "In the following cell, we load the solution and simulated results from their respective recorder files and plot the solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the validity of our results by using scipy.integrate.solve_ivp to\n",
    "# integrate the solution.\n",
    "sol = om.CaseReader('dymos_solution.db').get_case('final')\n",
    "sim = om.CaseReader('dymos_simulation.db').get_case('final')\n",
    "\n",
    "# Plot the results\n",
    "fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(12, 4.5))\n",
    "axes[0].plot(sol.get_val('traj.phase0.timeseries.x'),\n",
    "             sol.get_val('traj.phase0.timeseries.y'),\n",
    "             'ro', label='solution')\n",
    "axes[0].plot(sim.get_val('traj.phase0.timeseries.x'),\n",
    "             sim.get_val('traj.phase0.timeseries.y'),\n",
    "             'b-', label='simulation')\n",
    "axes[0].set_xlabel('x (m)')\n",
    "axes[0].set_ylabel('y (m/s)')\n",
    "axes[0].legend()\n",
    "axes[0].grid()\n",
    "axes[1].plot(sol.get_val('traj.phase0.timeseries.time'),\n",
    "             sol.get_val('traj.phase0.timeseries.theta', units='deg'),\n",
    "             'ro', label='solution')\n",
    "axes[1].plot(sim.get_val('traj.phase0.timeseries.time'),\n",
    "             sim.get_val('traj.phase0.timeseries.theta', units='deg'),\n",
    "             'b-', label='simulation')\n",
    "axes[1].set_xlabel('time (s)')\n",
    "axes[1].set_ylabel(r'$\\theta$ (deg)')\n",
    "axes[1].legend()\n",
    "axes[1].grid()\n",
    "plt.show()"
   ]
  },
  {
   "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(p.get_val('traj.phase0.timeseries.time')[-1], 1.8016, tolerance=1.0E-3)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
