Some cleanup of the Birkhoff transcription and tests.

.github/workflows/dymos_docs_workflow.yml

@@ -38,11 +38,12 @@ jobs:
 
           # make sure the latest versions of things don't break the docs
           # sticking with Python 3.12 for now, 3.13 requires NumPy 2.1 which does not work yet with PETSc/pyoptsparse
+          # Pin PETSc back to 3.22.2
           - NAME: latest
             PY: '3.12'
             NUMPY: 1
             SCIPY: 1
-            PETSc: 3
+            PETSc: 3.22.2
             PYOPTSPARSE: 'latest'
             SNOPT: 7.7
             OPENMDAO: 'dev'

.github/workflows/dymos_tests_workflow.yml

@@ -159,7 +159,7 @@ jobs:
             PY: '3.13'
             NUMPY: 2
             SCIPY: 1
-            PETSc: 3
+            PETSc: 3.22.2
             # PYOPTSPARSE: 'latest'
             SNOPT: 7.7
             OPENMDAO: 'dev'
@@ -170,14 +170,14 @@ jobs:
           # oldest supported versions
           - NAME: oldest
             PY: 3.9
-            NUMPY: 1.22
-            SCIPY: 1.7
+            NUMPY: 1.25
+            SCIPY: 1.9
             OPENMPI: '4.0'
             MPI4PY: '3.0'
             PETSc: 3.13
-            PYOPTSPARSE: 'v2.9.0'
+            PYOPTSPARSE: 'v2.10.2'
             SNOPT: 7.2
-            OPENMDAO: 3.28.0
+            OPENMDAO: 3.36.0
             MATPLOTLIB: 3.5
             OPTIONAL: '[test]'
             EXCLUDE: ${{ github.event_name == 'workflow_dispatch'  && ! inputs.oldest }}

benchmark/benchmark_balanced_field.py

@@ -207,37 +207,32 @@ def _run_balanced_field_length_problem(tx=dm.GaussLobatto, timeseries=True, sim=
     #
     p.setup(check=True)
 
-    p.set_val('traj.br_to_v1.t_initial', 0)
-    p.set_val('traj.br_to_v1.t_duration', 35)
-    p.set_val('traj.br_to_v1.states:r', br_to_v1.interp('', [0, 2500.0], nodes='state_input'))
-    p.set_val('traj.br_to_v1.states:v', br_to_v1.interp('', [0, 100.0], nodes='state_input'))
-    p.set_val('traj.br_to_v1.parameters:alpha', 0, units='deg')
-
-    p.set_val('traj.v1_to_vr.t_initial', 35)
-    p.set_val('traj.v1_to_vr.t_duration', 35)
-    p.set_val('traj.v1_to_vr.states:r', v1_to_vr.interp('r', [2500, 300.0]))
-    p.set_val('traj.v1_to_vr.states:v', v1_to_vr.interp('v', [100, 110.0]))
-    p.set_val('traj.v1_to_vr.parameters:alpha', 0.0, units='deg')
-
-    p.set_val('traj.rto.t_initial', 35)
-    p.set_val('traj.rto.t_duration', 35)
-    p.set_val('traj.rto.states:r', rto.interp('r', [2500, 5000.0]))
-    p.set_val('traj.rto.states:v', rto.interp('v', [110, 0]))
-    p.set_val('traj.rto.parameters:alpha', 0.0, units='deg')
-
-    p.set_val('traj.rotate.t_initial', 70)
-    p.set_val('traj.rotate.t_duration', 5)
-    p.set_val('traj.rotate.states:r', rotate.interp('r', [1750, 1800.0]))
-    p.set_val('traj.rotate.states:v', rotate.interp('v', [80, 85.0]))
-    p.set_val('traj.rotate.controls:alpha', 0.0, units='deg')
-
-    p.set_val('traj.climb.t_initial', 75)
-    p.set_val('traj.climb.t_duration', 15)
-    p.set_val('traj.climb.states:r', climb.interp('r', [5000, 5500.0]), units='ft')
-    p.set_val('traj.climb.states:v', climb.interp('v', [160, 170.0]), units='kn')
-    p.set_val('traj.climb.states:h', climb.interp('h', [0, 35.0]), units='ft')
-    p.set_val('traj.climb.states:gam', climb.interp('gam', [0, 5.0]), units='deg')
-    p.set_val('traj.climb.controls:alpha', 5.0, units='deg')
+    br_to_v1.set_time_val(initial=0, duration=35)
+    br_to_v1.set_state_val('r', [0, 2500], units='ft')
+    br_to_v1.set_state_val('v', [0, 100], units='kn')
+    br_to_v1.set_parameter_val('alpha', 0)
+
+    v1_to_vr.set_time_val(initial=0, duration=2)
+    v1_to_vr.set_state_val('r', [2500, 2800], units='ft')
+    v1_to_vr.set_state_val('v', [100, 110], units='kn')
+    v1_to_vr.set_parameter_val('alpha', 0)
+
+    rto.set_time_val(initial=0, duration=2)
+    rto.set_state_val('r', [2500, 5000], units='ft')
+    rto.set_state_val('v', [110, 0], units='kn')
+    rto.set_parameter_val('alpha', 0)
+
+    rotate.set_time_val(initial=37, duration=5)
+    rotate.set_state_val('r', [2800, 2900], units='ft')
+    rotate.set_state_val('v', [110, 115], units='kn')
+    rotate.set_control_val('alpha', [0.0, 10.0], units='deg')
+
+    climb.set_time_val(initial=45, duration=15)
+    climb.set_state_val('r', [5000, 5500], units='ft')
+    climb.set_state_val('v', [115, 150], units='kn')
+    climb.set_state_val('h', [0, 35], units='ft')
+    climb.set_state_val('gam', [0, 5], units='deg')
+    climb.set_control_val('alpha', [10, 5], units='deg')
 
     dm.run_problem(p, run_driver=True, simulate=sim, make_plots=make_plots)
 

benchmark/benchmark_brachistochrone.py

@@ -55,14 +55,12 @@ def brachistochrone_min_time(transcription='gauss-lobatto', num_segments=8, tran
     p.model.linear_solver = om.DirectSolver()
     p.setup(check=True)
 
-    p['phase0.t_initial'] = 0.0
-    p['phase0.t_duration'] = 2.0
-
-    p['phase0.states:x'] = phase.interp('x', ys=[0, 10])
-    p['phase0.states:y'] = phase.interp('y', ys=[10, 5])
-    p['phase0.states:v'] = phase.interp('v', ys=[0, 9.9])
-    p['phase0.controls:theta'] = phase.interp('theta', ys=[5, 100])
-    p['phase0.parameters:g'] = 9.80665
+    phase.set_time_val(initial=0.0, duration=2.0)
+    phase.set_state_val('x', [0, 10])
+    phase.set_state_val('y', [10, 5])
+    phase.set_state_val('v', [0, 9.9])
+    phase.set_control_val('theta', [5, 100])
+    phase.set_parameter_val('g', 9.80665)
 
     p.run_driver()
 

benchmark/benchmark_racecar.py

@@ -140,22 +140,22 @@ def _run_racecar_problem(transcription, timeseries=False, make_plots=False):
     # Now that the OpenMDAO problem is setup, we can set the values of the states.
 
     # States
-    # non-zero velocity in order to protect against 1/0 errors.
-    p.set_val('traj.phase0.states:V', phase.interp('V', [20, 20]), units='m/s')
-    p.set_val('traj.phase0.states:lambda', phase.interp('lambda', [0.0, 0.0]), units='rad')
-    # all other states start at 0
-    p.set_val('traj.phase0.states:omega', phase.interp('omega', [0.0, 0.0]), units='rad/s')
-    p.set_val('traj.phase0.states:alpha', phase.interp('alpha', [0.0, 0.0]), units='rad')
-    p.set_val('traj.phase0.states:ax', phase.interp('ax', [0.0, 0.0]), units='m/s**2')
-    p.set_val('traj.phase0.states:ay', phase.interp('ay', [0.0, 0.0]), units='m/s**2')
-    p.set_val('traj.phase0.states:n', phase.interp('n', [0.0, 0.0]), units='m')
+    # Nonzero velocity to avoid division by zero errors
+    phase.set_state_val('V', 20.0, units='m/s')
+    # All other states start at 0
+    phase.set_state_val('lambda', 0.0, units='rad')
+    phase.set_state_val('omega', 0.0, units='rad/s')
+    phase.set_state_val('alpha', 0.0, units='rad')
+    phase.set_state_val('ax', 0.0, units='m/s**2')
+    phase.set_state_val('ay', 0.0, units='m/s**2')
+    phase.set_state_val('n', 0.0, units='m')
     # initial guess for what the final time should be
-    p.set_val('traj.phase0.states:t', phase.interp('t', [0.0, 100.0]), units='s')
+    phase.set_state_val('t', [0.0, 100.0], units='s')
 
     # Controls
-    p.set_val('traj.phase0.controls:delta', phase.interp('delta', [0.0, 0.0]), units='rad')
-    p.set_val('traj.phase0.controls:thrust', phase.interp('thrust', [0.1, 0.1]), units=None)
     # a small amount of thrust can speed up convergence
+    phase.set_control_val('delta', 0.0, units='rad')
+    phase.set_control_val('thrust', 0.1, units=None)
 
     dm.run_problem(p, run_driver=True, simulate=False, make_plots=make_plots)
     print('Optimization finished')
@@ -183,4 +183,4 @@ def benchmark_radau_timeseries(self):
 
 
 if __name__ == '__main__':
-    _run_racecar_problem(dm.GaussLobatto, timeseries=False, make_plots=True)
+    _run_racecar_problem(dm.Radau, timeseries=False, make_plots=True)

docs/clean_notebooks.py

@@ -31,7 +31,7 @@ def _clear_notebook_outputs(path=os.getcwd(), dry_run=True):
             notebooks.extend(pathlib.Path(dirpath).glob('*.ipynb'))
         if not notebooks:
             print(f'{bcolors.FAIL}No notebooks found.{bcolors.ENDC}')
-            return
+            return 0
     elif path.endswith('.ipynb'):
         notebooks.append(path)
     else:
@@ -74,5 +74,5 @@ def _clear_notebook_outputs(path=os.getcwd(), dry_run=True):
                         help='Print notebooks with outputs but do not clean them.')
     args = parser.parse_args()
     num_cleaned = _clear_notebook_outputs(path=args.path, dry_run=args.dryrun)
-    if num_cleaned >0 and  args.dryrun:
+    if num_cleaned > 0 and  args.dryrun:
         exit(1)

docs/dymos_book/examples/implicit_duration/brachistochrone_implicit_duration.ipynb

@@ -114,7 +114,7 @@
     "        self.declare_partials(of='ydot', wrt='v', rows=ar, cols=ar)\n",
     "        self.declare_partials(of='ydot', wrt='theta', rows=ar, cols=ar)\n",
     "\n",
-    "    def compute(self, v, theta, g):\n",
+    "    def compute_primal(self, v, theta, g):\n",
     "        sin_theta = jnp.sin(theta)\n",
     "        cos_theta = jnp.cos(theta)\n",
     "\n",

dymos/examples/cannonball/doc/test_doc_cannonball_implicit_duration.py

@@ -20,10 +20,9 @@ def initial_guess(t_dur, gam0=np.pi/3):
 
 
 @use_tempdirs
-class TestTwoPhaseCannonballForDocs(unittest.TestCase):
+class TestCannonballImplicitDuration(unittest.TestCase):
 
-    @require_pyoptsparse(optimizer='IPOPT')
-    def test_two_phase_cannonball_for_docs(self):
+    def test_cannonball_implicit_duration(self):
         import openmdao.api as om
         from scipy.interpolate import make_interp_spline
         from openmdao.utils.assert_utils import assert_near_equal

dymos/examples/min_time_climb/aero/test/test_aerodynamics.py

@@ -4,6 +4,7 @@
 from numpy.testing import assert_almost_equal
 
 import openmdao.api as om
+from openmdao.utils.assert_utils import assert_check_partials
 
 from dymos.examples.min_time_climb.aero import AeroGroup
 
@@ -31,7 +32,7 @@ def setUp(self):
         self.prob.model.connect('alpha', 'aero.alpha')
         self.prob.model.connect('sos', 'aero.sos')
 
-        self.prob.setup()
+        self.prob.setup(force_alloc_complex=True)
 
     def test_aero_values(self):
 
@@ -52,15 +53,8 @@ def test_aero_values(self):
         assert_almost_equal(self.prob['aero.f_drag'], drag_expected, decimal=7)
 
     def testAeroDerivs(self):
-        cpd = self.prob.check_partials(compact_print=True)
-
-        for comp in cpd:
-            for (var, wrt) in cpd[comp]:
-                np.testing.assert_almost_equal(actual=cpd[comp][var, wrt]['J_fwd'],
-                                               desired=cpd[comp][var, wrt]['J_fd'],
-                                               decimal=5,
-                                               err_msg='Possible error in partials of component '
-                                                       '{0} for {1} wrt {2}'.format(comp, var, wrt))
+        cpd = self.prob.check_partials(compact_print=True, method='cs', out_stream=None)
+        assert_check_partials(cpd, atol=1.0E-5, rtol=1.0E-5)
 
 
 if __name__ == '__main__':  # pragma: no cover

dymos/phase/options.py

@@ -506,6 +506,9 @@ def __init__(self, read_only=False):
         self.declare(name='t_duration_targets', allow_none=True, default=[],
                      desc='targets in the ODE to which the total duration of the phase is connected')
 
+        self.declare(name='t_final_targets', allow_none=True, default=[],
+                     desc='targets in the ODE to which the final time of the phase is connected')
+
         self.declare(name='t_duration_balance_options', default={},
                      desc='options dictionary for the duration residual')
 

dymos/test/test_check_partials.py

@@ -1,7 +1,7 @@
 import unittest
 import openmdao.api as om
 import dymos as dm
-from dymos.utils.testing_utils import set_env_vars
+from openmdao.utils.testing_utils import set_env_vars
 
 
 class TestCheckPartials(unittest.TestCase):

dymos/test/test_options.py

@@ -1,9 +1,8 @@
 import unittest
 
-from openmdao.utils.testing_utils import use_tempdirs
+from openmdao.utils.testing_utils import use_tempdirs, set_env_vars
 import dymos as dm
 from dymos.examples.brachistochrone.test.ex_brachistochrone import brachistochrone_min_time
-from dymos.utils.testing_utils import set_env_vars
 
 
 @use_tempdirs

dymos/transcriptions/explicit_shooting/test/test_ode_evaluation_group.py

@@ -1,10 +1,11 @@
 import unittest
 
+import numpy as np
 import openmdao.api as om
 from openmdao.utils.assert_utils import assert_check_partials, assert_near_equal
 import dymos as dm
 
-from dymos.transcriptions.explicit_shooting.test.test_ode_integration_comp import SimpleODE
+from dymos.utils.testing_utils import SimpleODE, SimpleVectorizedODE
 from dymos.transcriptions.explicit_shooting.ode_evaluation_group import ODEEvaluationGroup
 
 
@@ -58,7 +59,61 @@ def test_eval(self):
 
         assert_near_equal(p.get_val('ode_eval.state_rate_collector.state_rates:x_rate'), xdot_check)
 
-        cpd = p.check_partials(compact_print=True, method='cs')
+        cpd = p.check_partials(compact_print=True, method='cs', out_stream=None)
+        assert_check_partials(cpd)
+
+    def test_eval_vectorized(self):
+        ode_class = SimpleVectorizedODE
+        time_options = dm.phase.options.TimeOptionsDictionary()
+
+        time_options['targets'] = 't'
+        time_options['units'] = 's'
+
+        state_options = {'z': dm.phase.options.StateOptionsDictionary()}
+
+        state_options['z']['shape'] = (2,)
+        state_options['z']['units'] = 's**2'
+        state_options['z']['rate_source'] = 'z_dot'
+        state_options['z']['targets'] = ['z']
+
+        param_options = {'p': dm.phase.options.ParameterOptionsDictionary()}
+
+        param_options['p']['shape'] = (1,)
+        param_options['p']['units'] = 's**2'
+        param_options['p']['targets'] = ['p']
+
+        control_options = {}
+
+        p = om.Problem()
+
+        igd = dm.GaussLobattoGrid(num_segments=1, nodes_per_seg=3, compressed=False)
+
+        p.model.add_subsystem('ode_eval', ODEEvaluationGroup(ode_class,
+                                                             input_grid_data=igd,
+                                                             time_options=time_options,
+                                                             state_options=state_options,
+                                                             parameter_options=param_options,
+                                                             control_options=control_options,
+                                                             ode_init_kwargs=None))
+        p.setup(check=False, force_alloc_complex=True)
+
+        p.model.ode_eval.set_segment_index(0)
+        p.set_val('ode_eval.states:z', [1.25, 0.0])
+        p.set_val('ode_eval.time', [2.2])
+        p.set_val('ode_eval.parameters:p', [1.0])
+
+        p.run_model()
+
+        z = p.get_val('ode_eval.states:z')
+        p_ = p.get_val('ode_eval.parameters:p')
+        t = p.get_val('ode_eval.time')
+
+        z_rate = p.get_val('ode_eval.state_rate_collector.state_rates:z_rate')
+
+        assert_near_equal(z_rate[0, 0], z[:, 0] - t**2 + p_)
+        assert_near_equal(z_rate[0, 1], 10 * t)
+
+        cpd = p.check_partials(compact_print=True, method='cs', out_stream=None)
         assert_check_partials(cpd)