Repository: OpenMDAO/dymos
Pull request #1142 (2025-02-03 22:30:16)

Some cleanup of the Birkhoff transcription and tests.
### Summary

Work on the Radau refactor exposed some tests that weren't using the `set_{xxx}_val` interface for setting dymos variables.
Other tests exposed lack of coverage in the Birkhoff transcription. This PR is a cleanup that addresses both of these.

### Related Issues

- Resolves #

### Backwards incompatibilities

- This PR bumps the minimum OpenMDAO version to 3.36.0. As OpenMDAO's jax capability expands, expect new dymos releases to take advantage of it.

### New Dependencies

None


review-notebook-app[bot]: Check out this pull request on&nbsp; <a href="https://app.reviewnb.com/OpenMDAO/dymos/pull/1142"><img align="absmiddle"  alt="ReviewNB" height="28" class="BotMessageButtonImage" src="https://raw.githubusercontent.com/ReviewNB/support/master/images/button_reviewnb.png"/></a> 

 See visual diffs & provide feedback on Jupyter Notebooks. 

---

 <i>Powered by <a href='https://www.reviewnb.com/?utm_source=gh'>ReviewNB</a></i>
coveralls: 
[![Coverage Status](https://coveralls.io/builds/72102682/badge)](https://coveralls.io/builds/72102682)

coverage: 92.63% (-1.2%) from 93.815%
when pulling **820bb8cfdedb6db0f18da7edadf03907cdf703dc on robfalck:birkhoff_updates**
into **cd45db95f36aad8742da632b9aff7723c263555f on OpenMDAO:master**.


Patch:
diff --git a/.github/workflows/dymos_docs_workflow.yml b/.github/workflows/dymos_docs_workflow.yml
index 204b97b4..94dede7a 100644
--- a/.github/workflows/dymos_docs_workflow.yml
+++ b/.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'
diff --git a/benchmark/benchmark_balanced_field.py b/benchmark/benchmark_balanced_field.py
index 6b0f25a2..60d988c6 100644
--- a/benchmark/benchmark_balanced_field.py
+++ b/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)
 
diff --git a/benchmark/benchmark_brachistochrone.py b/benchmark/benchmark_brachistochrone.py
index 87c9c510..1b698ff3 100644
--- a/benchmark/benchmark_brachistochrone.py
+++ b/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()
 
diff --git a/benchmark/benchmark_racecar.py b/benchmark/benchmark_racecar.py
index 413c7c72..411725c9 100644
--- a/benchmark/benchmark_racecar.py
+++ b/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 @@ class BenchmarkRacecar(unittest.TestCase):
 
 
 if __name__ == '__main__':
-    _run_racecar_problem(dm.GaussLobatto, timeseries=False, make_plots=True)
+    _run_racecar_problem(dm.Radau, timeseries=False, make_plots=True)
diff --git a/docs/clean_notebooks.py b/docs/clean_notebooks.py
index 5c034399..14ca3511 100644
--- a/docs/clean_notebooks.py
+++ b/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 @@ if __name__ == '__main__':
                         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)
diff --git a/docs/dymos_book/examples/implicit_duration/brachistochrone_implicit_duration.ipynb b/docs/dymos_book/examples/implicit_duration/brachistochrone_implicit_duration.ipynb
index 0814c855..c269a720 100644
--- a/docs/dymos_book/examples/implicit_duration/brachistochrone_implicit_duration.ipynb
+++ b/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",
diff --git a/dymos/phase/options.py b/dymos/phase/options.py
index 2f3306f4..3a9ef719 100644
--- a/dymos/phase/options.py
+++ b/dymos/phase/options.py
@@ -506,6 +506,9 @@ class TimeOptionsDictionary(om.OptionsDictionary):
         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')
 
diff --git a/dymos/transcriptions/pseudospectral/birkhoff.py b/dymos/transcriptions/pseudospectral/birkhoff.py
index 06a83c9b..8687d2fd 100644
--- a/dymos/transcriptions/pseudospectral/birkhoff.py
+++ b/dymos/transcriptions/pseudospectral/birkhoff.py
@@ -108,11 +108,10 @@ class Birkhoff(TranscriptionBase):
                 src_idxs = self.grid_data.subset_node_indices['all']
                 phase.connect(name, [f'ode_all.{t}' for t in targets], src_indices=src_idxs,
                               flat_src_indices=True)
-                src_idxs = om.slicer[[0, -1], ...]
-                phase.connect(name, [f'boundary_vals.{t}' for t in targets], src_indices=src_idxs)
 
         for name, targets in [('t_initial', options['t_initial_targets']),
-                              ('t_duration', options['t_duration_targets'])]:
+                              ('t_duration', options['t_duration_targets']),
+                              ('t_final', options['t_final_targets'])]:
             for t in targets:
                 shape = ode_inputs[t]['shape']
 
@@ -256,18 +255,24 @@ class Birkhoff(TranscriptionBase):
 
     def configure_defects(self, phase):
         """
-        Configure the continuity_comp and connect the collocation constraints.
+        Connect the collocation constraints.
 
         Parameters
         ----------
         phase : dymos.Phase
             The phase object to which this transcription instance applies.
         """
+        num_nodes = self.grid_data.subset_num_nodes['all']
         for name, options in phase.state_options.items():
             rate_source_type = phase.classify_var(options['rate_source'])
             rate_src_path = self._get_rate_source_path(name, phase)
+            if rate_src_path.startswith('parameter_vals:'):
+                src_idxs = om.slicer[np.zeros(num_nodes, dtype=int), ...]
+            else:
+                src_idxs = None
+
             if rate_source_type not in ('state', 'ode'):
-                phase.connect(rate_src_path, f'f_computed:{name}')
+                phase.connect(rate_src_path, f'f_computed:{name}', src_indices=src_idxs)
 
     def setup_duration_balance(self, phase):
         """
@@ -782,9 +787,11 @@ class Birkhoff(TranscriptionBase):
                 else:
                     src_idxs_raw = np.zeros(self.grid_data.subset_num_nodes['all'], dtype=int)
                     src_idxs = get_src_indices_by_row(src_idxs_raw, options['shape'])
-                    endpoint_src_idxs = om.slicer[[0, -1], ...]
+                    endpoint_src_idxs_raw = np.zeros(2, dtype=int)
+                    endpoint_src_idxs = get_src_indices_by_row(endpoint_src_idxs_raw, options['shape'])
                     if options['shape'] == (1,):
                         src_idxs = src_idxs.ravel()
+                        endpoint_src_idxs = endpoint_src_idxs.ravel()
 
                 connection_info.append((f'ode_all.{tgt}', (src_idxs,)))
                 connection_info.append((f'boundary_vals.{tgt}', (endpoint_src_idxs,)))
diff --git a/dymos/transcriptions/pseudospectral/components/__init__.py b/dymos/transcriptions/pseudospectral/components/__init__.py
index 9b9e01c0..446222f3 100644
--- a/dymos/transcriptions/pseudospectral/components/__init__.py
+++ b/dymos/transcriptions/pseudospectral/components/__init__.py
@@ -3,6 +3,6 @@ from .state_interp_comp import StateInterpComp
 from .state_independents import StateIndependentsComp
 from .control_endpoint_defect_comp import ControlEndpointDefectComp
 from .gauss_lobatto_interleave_comp import GaussLobattoInterleaveComp
-from .birkhoff_collocation_comp import BirkhoffCollocationComp
+from .birkhoff_defect_comp import BirkhoffDefectComp
 from .birkhoff_iter_group import BirkhoffIterGroup
 from .birkhoff_boundary_group import BirkhoffBoundaryGroup
diff --git a/dymos/transcriptions/pseudospectral/components/birkhoff_boundary_group.py b/dymos/transcriptions/pseudospectral/components/birkhoff_boundary_group.py
index 26a79fbb..9d3138fb 100644
--- a/dymos/transcriptions/pseudospectral/components/birkhoff_boundary_group.py
+++ b/dymos/transcriptions/pseudospectral/components/birkhoff_boundary_group.py
@@ -80,7 +80,7 @@ class BirkhoffBoundaryMuxComp(om.ExplicitComponent):
 
 class BirkhoffBoundaryGroup(om.Group):
     """
-    Class definition for the BirkhoffBoundaryEvalGroup.
+    Class definition for the BirkhoffBoundaryGroup.
 
     This group accepts values for initial and final times, states, controls, and parameters
     and evaluates the ODE with those in order to compute the boundary values and
diff --git a/dymos/transcriptions/pseudospectral/components/birkhoff_collocation_comp.py b/dymos/transcriptions/pseudospectral/components/birkhoff_defect_comp.py
similarity index 98%
rename from dymos/transcriptions/pseudospectral/components/birkhoff_collocation_comp.py
rename to dymos/transcriptions/pseudospectral/components/birkhoff_defect_comp.py
index f40ca6aa..32066e67 100644
--- a/dymos/transcriptions/pseudospectral/components/birkhoff_collocation_comp.py
+++ b/dymos/transcriptions/pseudospectral/components/birkhoff_defect_comp.py
@@ -12,11 +12,11 @@ from dymos.utils.cgl import cgl
 from dymos.utils.birkhoff import birkhoff_matrix
 
 
-class BirkhoffCollocationComp(om.ExplicitComponent):
+class BirkhoffDefectComp(om.ExplicitComponent):
     """
-    Class definition for the BirkhoffCollocationComp.
+    Class definition for the BirkhoffDefectComp.
 
-    BirkhoffCollocationComp computes the generalized defects of a segment for implicit collocation.
+    BirkhoffDefectComp computes the generalized defects of a segment for implicit collocation.
     There are four defects to be evaluated; state, state rate, initial state, and final state.
     The state defect is the difference between the state value design variables and the state rate
     design variables.
diff --git a/dymos/transcriptions/pseudospectral/components/birkhoff_iter_group.py b/dymos/transcriptions/pseudospectral/components/birkhoff_iter_group.py
index d26389c7..c4a5ab2b 100644
--- a/dymos/transcriptions/pseudospectral/components/birkhoff_iter_group.py
+++ b/dymos/transcriptions/pseudospectral/components/birkhoff_iter_group.py
@@ -1,8 +1,7 @@
 import numpy as np
 import openmdao.api as om
 
-from .birkhoff_collocation_comp import BirkhoffCollocationComp
-from .birkhoff_state_resid_comp import BirkhoffStateResidComp
+from .birkhoff_defect_comp import BirkhoffDefectComp
 
 from ...grid_data import GridData
 from ....phase.options import TimeOptionsDictionary
@@ -53,13 +52,13 @@ class BirkhoffIterGroup(om.Group):
         self.add_subsystem('ode_all', subsys=ode_class(num_nodes=nn, **ode_init_kwargs))
 
         self.add_subsystem('collocation_comp',
-                           subsys=BirkhoffCollocationComp(grid_data=gd,
-                                                          state_options=state_options,
-                                                          time_units=time_options['units']),
+                           subsys=BirkhoffDefectComp(grid_data=gd,
+                                                     state_options=state_options,
+                                                     time_units=time_options['units']),
                            promotes_inputs=['*'], promotes_outputs=['*'])
 
         if any([opts['solve_segments'] in ('forward', 'backward') for opts in state_options.values()]):
-            self.add_subsystem('states_balance_comp', subsys=BirkhoffStateResidComp(),
+            self.add_subsystem('states_balance_comp', subsys=om.InputResidsComp(),
                                promotes_inputs=['*'], promotes_outputs=['*'])
 
     def _configure_desvars(self, name, options):
@@ -208,20 +207,20 @@ class BirkhoffIterGroup(om.Group):
                                                shape=(nn,) + shape,
                                                units=units)
 
-                states_balance_comp.add_residual_from_input(f'state_defects:{name}',
-                                                            shape=(nn+ns,) + shape,
-                                                            units=units)
+                states_balance_comp.add_input(f'state_defects:{name}',
+                                              shape=(nn+ns,) + shape,
+                                              units=units)
 
                 if ns > 1:
-                    states_balance_comp.add_residual_from_input(f'state_cnty_defects:{name}',
-                                                                shape=(ns - 1,) + shape,
-                                                                units=units)
+                    states_balance_comp.add_input(f'state_cnty_defects:{name}',
+                                                  shape=(ns - 1,) + shape,
+                                                  units=units)
 
             if f'state_rates:{name}' in self._implicit_outputs:
                 states_balance_comp.add_output(f'state_rates:{name}', shape=(nn,) + shape, units=units)
-                states_balance_comp.add_residual_from_input(f'state_rate_defects:{name}',
-                                                            shape=(nn,) + shape,
-                                                            units=units)
+                states_balance_comp.add_input(f'state_rate_defects:{name}',
+                                              shape=(nn,) + shape,
+                                              units=units)
 
             if f'initial_states:{name}' in self._implicit_outputs:
                 states_balance_comp.add_output(f'initial_states:{name}', shape=(1,) + shape, units=units)
@@ -240,102 +239,3 @@ class BirkhoffIterGroup(om.Group):
 
             if var_type == 'ode':
                 self.connect(f'ode_all.{rate_source}', f'f_computed:{name}')
-
-    def _get_rate_source_path(self, state_name, nodes, phase):
-        """
-        Return the rate source location and indices for a given state name.
-
-        Parameters
-        ----------
-        state_name : str
-            Name of the state.
-        nodes : str
-            One of ['col', 'all'].
-        phase : dymos.Phase
-            Phase object containing the rate source.
-
-        Returns
-        -------
-        str
-            Path to the rate source.
-        ndarray
-            Array of source indices.
-        """
-        gd = self.grid_data
-        try:
-            var = phase.state_options[state_name]['rate_source']
-        except RuntimeError:
-            raise ValueError(f"state '{state_name}' in phase '{phase.name}' was not given a "
-                             "rate_source")
-
-        # Note the rate source must be shape-compatible with the state
-        var_type = phase.classify_var(var)
-
-        # Determine the path to the variable
-        if var_type == 't':
-            rate_path = 't'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 't_phase':
-            rate_path = 't_phase'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'state':
-            rate_path = f'states:{var}'
-            # Find the state_input indices which occur at segment endpoints, and repeat them twice
-            state_input_idxs = gd.subset_node_indices['state_input']
-            repeat_idxs = np.ones_like(state_input_idxs)
-            if self.options['compressed']:
-                segment_end_idxs = gd.subset_node_indices['segment_ends'][1:-1]
-                # Repeat nodes that are on segment bounds (but not the first or last nodes in the phase)
-                nodes_to_repeat = list(set(state_input_idxs).intersection(segment_end_idxs))
-                # Now find these nodes in the state input indices
-                idxs_of_ntr_in_state_inputs = np.where(np.isin(state_input_idxs, nodes_to_repeat))[0]
-                # All state input nodes are used once, but nodes_to_repeat are used twice
-                repeat_idxs[idxs_of_ntr_in_state_inputs] = 2
-            # Now we have a way of mapping the state input indices to all nodes
-            map_input_node_idxs_to_all = np.repeat(np.arange(gd.subset_num_nodes['state_input'],
-                                                             dtype=int), repeats=repeat_idxs)
-            # Now select the subset of nodes we want to use.
-            node_idxs = map_input_node_idxs_to_all[gd.subset_node_indices[nodes]]
-        elif var_type == 'indep_control':
-            rate_path = f'control_values:{var}'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'input_control':
-            rate_path = f'control_values:{var}'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'control_rate':
-            control_name = var[:-5]
-            rate_path = f'control_rates:{control_name}_rate'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'control_rate2':
-            control_name = var[:-6]
-            rate_path = f'control_rates:{control_name}_rate2'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'indep_polynomial_control':
-            rate_path = f'control_values:{var}'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'input_polynomial_control':
-            rate_path = f'control_values:{var}'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'polynomial_control_rate':
-            control_name = var[:-5]
-            rate_path = f'control_rates:{control_name}_rate'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'polynomial_control_rate2':
-            control_name = var[:-6]
-            rate_path = f'control_rates:{control_name}_rate2'
-            node_idxs = gd.subset_node_indices[nodes]
-        elif var_type == 'parameter':
-            rate_path = f'parameter_vals:{var}'
-            dynamic = not phase.parameter_options[var]['static_target']
-            if dynamic:
-                node_idxs = np.zeros(gd.subset_num_nodes[nodes], dtype=int)
-            else:
-                node_idxs = np.zeros(1, dtype=int)
-        else:
-            # Failed to find variable, assume it is in the ODE
-            rate_path = f'ode_all.{var}'
-            node_idxs = gd.subset_node_indices[nodes]
-
-        src_idxs = om.slicer[node_idxs, ...]
-
-        return rate_path, src_idxs
diff --git a/dymos/transcriptions/pseudospectral/components/birkhoff_state_resid_comp.py b/dymos/transcriptions/pseudospectral/components/birkhoff_state_resid_comp.py
deleted file mode 100644
index 09cde3b3..00000000
--- a/dymos/transcriptions/pseudospectral/components/birkhoff_state_resid_comp.py
+++ /dev/null
@@ -1,72 +0,0 @@
-import numpy as np
-
-import openmdao.api as om
-
-from openmdao.utils.general_utils import ensure_compatible
-
-from dymos._options import options as dymos_options
-from dymos.utils.misc import om_version
-
-
-class BirkhoffStateResidComp(om.ImplicitComponent):
-    """
-    Class definition for the BirkhoffStateResidComp.
-
-    Generates the residuals for any states that are solved for implicitly.
-
-    Parameters
-    ----------
-    **kwargs : dict
-        Dictionary of optional arguments.
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self._no_check_partials = not dymos_options['include_check_partials']
-        self._io_pairs = []
-
-    def add_residual_from_input(self, name, **kwargs):
-        """
-        Adds a residual whose value is given by resid_input.
-
-        Parameters
-        ----------
-        name : str
-            The name of the input providing the residuals for the given output.
-        **kwargs : dict
-            Additional keyword arguments for add_input and add_residual.
-        """
-        val = kwargs['val'] if 'val' in kwargs else 1.0
-        shape = kwargs['shape'] if 'shape' in kwargs else None
-        val, shape = ensure_compatible(name, value=val, shape=shape)
-        resid_name = 'resid_' + name
-
-        self._io_pairs.append((resid_name, name))
-
-        size = np.prod(shape, dtype=int)
-        ar = np.arange(size, dtype=int)
-
-        self.add_input(name, **kwargs)
-        self.add_residual(resid_name, **kwargs)
-
-        if om_version()[0] > (3, 31, 1):
-            self.declare_partials(of=resid_name, wrt=name, rows=ar, cols=ar, val=1.0)
-        else:
-            self.declare_partials(of='*', wrt='*', method='fd')
-
-    def apply_nonlinear(self, inputs, outputs, residuals):
-        """
-        Compute residuals given inputs and outputs.
-
-        The model is assumed to be in an unscaled state.
-
-        Parameters
-        ----------
-        inputs : Vector
-            Unscaled, dimensional input variables read via inputs[key].
-        outputs : Vector
-            Unscaled, dimensional output variables read via outputs[key].
-        residuals : Vector
-            Unscaled, dimensional residuals written to via residuals[key].
-        """
-        residuals.set_val(inputs.asarray())
diff --git a/joss/joss_demo.py b/joss/joss_demo.py
index b9591930..3cab7f51 100644
--- a/joss/joss_demo.py
+++ b/joss/joss_demo.py
@@ -101,25 +101,14 @@ p.setup()
 
 # Now that the OpenMDAO problem is setup, we can guess the
 # values of time, states, and controls.
-p.set_val('traj.phase0.t_duration', 2.0)
+phase.set_time_val(initial=0.0, duration=2.0)
 
 # States and controls here use a linearly interpolated
 # initial guess along the trajectory.
-p.set_val('traj.phase0.states:x',
-          phase.interp('x', [0, 10]),
-          units='m')
-
-p.set_val('traj.phase0.states:y',
-          phase.interp('y', [10, 5]),
-          units='m')
-
-p.set_val('traj.phase0.states:v',
-          phase.interp('v', [0, 5]),
-          units='m/s')
-
-p.set_val('traj.phase0.controls:theta',
-          phase.interp('theta', [5, 45]),
-          units='deg')
+phase.set_state_val('x', [0, 10], units='m')
+phase.set_state_val('y', [10, 5], units='m')
+phase.set_state_val('v', [0, 5], units='m/s')
+phase.set_control_val('theta', [5, 45], units='deg')
 
 # Run the driver to solve the problem
 p.run_driver()
diff --git a/pyproject.toml b/pyproject.toml
index d75b1c4c..739e5659 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,7 +24,7 @@ classifiers = [
 ]
 dependencies = [
     "numpy",
-    "openmdao>=3.27.0",
+    "openmdao>=3.36.0",
     "scipy",
 ]
 


Test patch:
diff --git a/.github/workflows/dymos_tests_workflow.yml b/.github/workflows/dymos_tests_workflow.yml
index 0d250349..398d922f 100644
--- a/.github/workflows/dymos_tests_workflow.yml
+++ b/.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 }}
diff --git a/dymos/examples/cannonball/doc/test_doc_cannonball_implicit_duration.py b/dymos/examples/cannonball/doc/test_doc_cannonball_implicit_duration.py
index 5c49828c..dd2b958e 100644
--- a/dymos/examples/cannonball/doc/test_doc_cannonball_implicit_duration.py
+++ b/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
diff --git a/dymos/examples/min_time_climb/aero/test/test_aerodynamics.py b/dymos/examples/min_time_climb/aero/test/test_aerodynamics.py
index 0efd3cd9..eff3bb23 100644
--- a/dymos/examples/min_time_climb/aero/test/test_aerodynamics.py
+++ b/dymos/examples/min_time_climb/aero/test/test_aerodynamics.py
@@ -4,6 +4,7 @@ import numpy as np
 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 @@ class TestAeroGroup(unittest.TestCase):
         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 @@ class TestAeroGroup(unittest.TestCase):
         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
diff --git a/dymos/test/test_check_partials.py b/dymos/test/test_check_partials.py
index aea60525..15acb6b0 100644
--- a/dymos/test/test_check_partials.py
+++ b/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):
diff --git a/dymos/test/test_options.py b/dymos/test/test_options.py
index 0c33d995..cce4cdf8 100644
--- a/dymos/test/test_options.py
+++ b/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
diff --git a/dymos/transcriptions/explicit_shooting/test/test_ode_evaluation_group.py b/dymos/transcriptions/explicit_shooting/test/test_ode_evaluation_group.py
index 3afdc944..ab5fbd89 100644
--- a/dymos/transcriptions/explicit_shooting/test/test_ode_evaluation_group.py
+++ b/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 @@ class TestODEEvaluationGroup(unittest.TestCase):
 
         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)
 
 
diff --git a/dymos/transcriptions/explicit_shooting/test/test_ode_integration_comp.py b/dymos/transcriptions/explicit_shooting/test/test_ode_integration_comp.py
index 0002cff1..fdacda83 100644
--- a/dymos/transcriptions/explicit_shooting/test/test_ode_integration_comp.py
+++ b/dymos/transcriptions/explicit_shooting/test/test_ode_integration_comp.py
@@ -4,6 +4,7 @@ import numpy as np
 import openmdao.api as om
 import dymos as dm
 from dymos import options as dymos_options
+from dymos.utils.misc import om_version
 
 from openmdao.utils.assert_utils import assert_check_partials, assert_near_equal
 from dymos.examples.brachistochrone.brachistochrone_ode import BrachistochroneODE
@@ -11,41 +12,13 @@ from dymos.transcriptions.explicit_shooting.ode_integration_comp import ODEInteg
 
 from dymos.phase.options import TimeOptionsDictionary, StateOptionsDictionary, ParameterOptionsDictionary
 from dymos.transcriptions.grid_data import GridData
-
-
-class SimpleODE(om.ExplicitComponent):
-    """
-    A simple ODE from https://math.okstate.edu/people/yqwang/teaching/math4513_fall11/Notes/rungekutta.pdf
-    """
-    def initialize(self):
-        self.options.declare('num_nodes', types=(int,))
-
-    def setup(self):
-        nn = self.options['num_nodes']
-        self.add_input('x', shape=(nn,), units='s**2')
-        self.add_input('t', shape=(nn,), units='s')
-        self.add_input('p', shape=(nn,), units='s**2')
-
-        self.add_output('x_dot', shape=(nn,), units='s')
-
-        ar = np.arange(nn, dtype=int)
-        self.declare_partials(of='x_dot', wrt='x', rows=ar, cols=ar, val=1.0)
-        self.declare_partials(of='x_dot', wrt='t', rows=ar, cols=ar)
-        self.declare_partials(of='x_dot', wrt='p', rows=ar, cols=ar, val=1.0)
-
-    def compute(self, inputs, outputs):
-        x = inputs['x']
-        t = inputs['t']
-        p = inputs['p']
-        outputs['x_dot'] = x - t**2 + p
-
-    def compute_partials(self, inputs, partials):
-        t = inputs['t']
-        partials['x_dot', 't'] = -2*t
+from dymos.utils.testing_utils import SimpleODE
 
 
 class TestODEIntegrationComp(unittest.TestCase):
 
+    @unittest.skipIf(om_version()[0] >= (3, 36, 0) or om_version()[0] < (3, 37, 0),
+                     reason='Test skipped due to an issue in OpenMDAO 3.36.x')
     def test_integrate_scalar_ode(self):
         dymos_options['include_check_partials'] = True
 
@@ -110,6 +83,8 @@ class TestODEIntegrationComp(unittest.TestCase):
 
         dymos_options['include_check_partials'] = False
 
+    @unittest.skipIf(om_version()[0] >= (3, 36, 0) or om_version()[0] < (3, 37, 0),
+                     reason='Test skipped due to an issue in OpenMDAO 3.36.x')
     def test_integrate_with_controls(self):
 
         dymos_options['include_check_partials'] = True
@@ -190,6 +165,8 @@ class TestODEIntegrationComp(unittest.TestCase):
             cpd = p.check_partials(compact_print=False, method='fd')
             assert_check_partials(cpd, atol=1.0E-4, rtol=1.0E-4)
 
+    @unittest.skipIf(om_version()[0] >= (3, 36, 0) or om_version()[0] < (3, 37, 0),
+                     reason='Test skipped due to an issue in OpenMDAO 3.36.x')
     def test_integrate_with_polynomial_controls(self):
 
         dymos_options['include_check_partials'] = True
diff --git a/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_collocation_defect.py b/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_collocation_defect.py
index 2df44cfe..c6b83cd6 100644
--- a/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_collocation_defect.py
+++ b/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_collocation_defect.py
@@ -7,7 +7,7 @@ import openmdao.api as om
 from dymos.utils.testing_utils import assert_check_partials
 
 import dymos as dm
-from dymos.transcriptions.pseudospectral.components.birkhoff_collocation_comp import BirkhoffCollocationComp
+from dymos.transcriptions.pseudospectral.components.birkhoff_defect_comp import BirkhoffDefectComp
 from dymos.transcriptions.grid_data import BirkhoffGrid
 
 # Modify class so we can run it standalone.
@@ -23,7 +23,7 @@ class _PhaseStub():
         return None
 
 
-CollocationComp = CompWrapperConfig(BirkhoffCollocationComp, [_PhaseStub()])
+CollocationComp = CompWrapperConfig(BirkhoffDefectComp, [_PhaseStub()])
 
 
 class TestCollocationComp(unittest.TestCase):
diff --git a/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_iter_group.py b/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_iter_group.py
index 8ffe6070..7ed817f6 100644
--- a/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_iter_group.py
+++ b/dymos/transcriptions/pseudospectral/components/test/test_birkhoff_iter_group.py
@@ -9,56 +9,13 @@ from openmdao.utils.assert_utils import assert_check_partials, assert_near_equal
 from openmdao.utils.testing_utils import use_tempdirs
 
 from dymos.utils.misc import GroupWrapperConfig
+from dymos.utils.testing_utils import PhaseStub, SimpleODE
 from dymos.transcriptions.pseudospectral.components import BirkhoffIterGroup
 from dymos.phase.options import StateOptionsDictionary, TimeOptionsDictionary
 from dymos.transcriptions.grid_data import BirkhoffGrid
 
 
-class _PhaseStub():
-    """
-    A stand-in for the Phase during config_io for testing.
-    It just supports the classify_var method and returns "state", the only value needed for unittests.
-    """
-    def __init__(self):
-        self.nonlinear_solver = None
-        self.linear_solver = None
-
-    def classify_var(self, name):
-        return 'ode'
-
-
-BirkhoffIterGroup = GroupWrapperConfig(BirkhoffIterGroup, [_PhaseStub()])
-
-
-class SimpleODE(om.ExplicitComponent):
-    """
-    A simple ODE from https://math.okstate.edu/people/yqwang/teaching/math4513_fall11/Notes/rungekutta.pdf
-    """
-    def initialize(self):
-        self.options.declare('num_nodes', types=(int,))
-
-    def setup(self):
-        nn = self.options['num_nodes']
-        self.add_input('x', shape=(nn,), units='s**2')
-        self.add_input('t', shape=(nn,), units='s')
-        self.add_input('p', shape=(nn,), units='s**2')
-
-        self.add_output('x_dot', shape=(nn,), units='s')
-
-        ar = np.arange(nn, dtype=int)
-        self.declare_partials(of='x_dot', wrt='x', rows=ar, cols=ar, val=1.0)
-        self.declare_partials(of='x_dot', wrt='t', rows=ar, cols=ar)
-        self.declare_partials(of='x_dot', wrt='p', rows=ar, cols=ar, val=1.0)
-
-    def compute(self, inputs, outputs):
-        x = inputs['x']
-        t = inputs['t']
-        p = inputs['p']
-        outputs['x_dot'] = x - t**2 + p
-
-    def compute_partials(self, inputs, partials):
-        t = inputs['t']
-        partials['x_dot', 't'] = -2*t
+BirkhoffIterGroup = GroupWrapperConfig(BirkhoffIterGroup, [PhaseStub()])
 
 
 @use_tempdirs
diff --git a/dymos/transcriptions/pseudospectral/components/test/test_collocation_defect_solver.py b/dymos/transcriptions/pseudospectral/components/test/test_collocation_defect_solver.py
index f0c41f3c..fd274cc1 100644
--- a/dymos/transcriptions/pseudospectral/components/test/test_collocation_defect_solver.py
+++ b/dymos/transcriptions/pseudospectral/components/test/test_collocation_defect_solver.py
@@ -354,7 +354,9 @@ class TestCollocationBalanceApplyNL(unittest.TestCase):
                     # them (should be some form of -I)
                     continue
                 check_data = data[(of, wrt)]
-                self.assertLess(check_data['abs error'].forward, 1e-8)
+                fwd_abs_error = check_data['abs error'].forward
+                if fwd_abs_error is not None:
+                    self.assertLess(check_data['abs error'].forward, 1e-8)
 
         cpd = p.check_partials(compact_print=True, method='fd', out_stream=None)
         data = cpd['traj0.phases.phase0.indep_states']
@@ -372,7 +374,9 @@ class TestCollocationBalanceApplyNL(unittest.TestCase):
                     # them (should be some form of -I)
                     continue
                 check_data = data[(of, wrt)]
-                self.assertLess(check_data['abs error'].forward, 1e-8)
+                fwd_abs_error = check_data['abs error'].forward
+                if fwd_abs_error is not None:
+                    self.assertLess(check_data['abs error'].forward, 1e-8)
 
         cpd = p.check_partials(compact_print=True, method='fd', out_stream=None)
         data = cpd['traj0.phases.phase0.indep_states']
diff --git a/dymos/utils/testing_utils.py b/dymos/utils/testing_utils.py
index 2d41476a..fda855a9 100644
--- a/dymos/utils/testing_utils.py
+++ b/dymos/utils/testing_utils.py
@@ -1,5 +1,4 @@
 import io
-import os
 
 from packaging.version import Version
 
@@ -361,52 +360,162 @@ def _get_reports_dir(prob):
     return get_reports_dir(prob)
 
 
-# This duplicates OpenMDAO code and is needed for older versions of OpenMDAO (<= 3.19).
-# Once support is dropped for < 3.19 we can get rid of this and use the version from OpenMDAO.
-class set_env_vars(object):
+class PhaseStub():
     """
-    Decorate a function to temporarily set some environment variables.
+    A stand-in for the Phase during config_io for testing.
+
+    It just supports the classify_var method and returns "ode", the only value needed for unittests.
+    """
+    def __init__(self):
+        self.nonlinear_solver = None
+        self.linear_solver = None
+
+    def classify_var(self, name):
+        """
+        A stand-in for classify_var that always sets the variable type to name.
+
+        Parameters
+        ----------
+        name : str
+            The name of the variable to classify.
+
+        Returns
+        -------
+        str
+            The variable classification.
+        """
+        return 'ode'
+
+
+class SimpleODE(om.ExplicitComponent):
+    """
+    A simple ODE for testing purposes.
+
+    Source: https://math.okstate.edu/people/yqwang/teaching/math4513_fall11/Notes/rungekutta.pdf
 
     Parameters
     ----------
-    **envs : dict
-        Keyword args corresponding to environment variables to set.
+    **kwargs : dict of keyword arguments
+        Keyword arguments that will be mapped into the Component options.
+    """
+    def initialize(self):
+        """
+        Declare options for SimpleODE.
+        """
+        self.options.declare('num_nodes', types=(int,))
+
+    def setup(self):
+        """
+        Add inputs and outputs to SimpleODE.
+        """
+        nn = self.options['num_nodes']
+        self.add_input('x', shape=(nn,), units='s**2')
+        self.add_input('t', shape=(nn,), units='s')
+        self.add_input('p', shape=(nn,), units='s**2')
+
+        self.add_output('x_dot', shape=(nn,), units='s')
+
+        ar = np.arange(nn, dtype=int)
+        self.declare_partials(of='x_dot', wrt='x', rows=ar, cols=ar, val=1.0)
+        self.declare_partials(of='x_dot', wrt='t', rows=ar, cols=ar)
+        self.declare_partials(of='x_dot', wrt='p', rows=ar, cols=ar, val=1.0)
+
+    def compute(self, inputs, outputs):
+        """
+        Compute the outputs of SimpleVectorizedODE.
 
-    Attributes
+        Parameters
+        ----------
+        inputs : Vector
+            Vector of inputs.
+        outputs : Vector
+            Vector of outputs.
+        """
+        x = inputs['x']
+        t = inputs['t']
+        p = inputs['p']
+        outputs['x_dot'] = x - t**2 + p
+
+    def compute_partials(self, inputs, partials):
+        """
+        Compute the partials of SimpleVectorizedODE.
+
+        Parameters
+        ----------
+        inputs : Vector
+            Vector of inputs.
+        partials : Dictionary
+            Vector of partials.
+        """
+        t = inputs['t']
+        partials['x_dot', 't'] = -2*t
+
+
+class SimpleVectorizedODE(om.ExplicitComponent):
+    """
+    A simple vector-valued ODE.
+
+    Source: https://math.okstate.edu/people/yqwang/teaching/math4513_fall11/Notes/rungekutta.pdf
+
+    Parameters
     ----------
-    envs : dict
-        Saved mapping of environment var name to value.
+    **kwargs : dict of keyword arguments
+        Keyword arguments that will be mapped into the Component options.
     """
+    def initialize(self):
+        """
+        Declare options for SimpleVectorizedODE.
+        """
+        self.options.declare('num_nodes', types=(int,))
 
-    def __init__(self, **envs):
+    def setup(self):
         """
-        Initialize attributes.
+        Add inputs and outputs to SimpleVectorizedODE.
+        """
+        nn = self.options['num_nodes']
+        self.add_input('z', shape=(nn, 2), units='s**2')
+        self.add_input('t', shape=(nn,), units='s')
+        self.add_input('p', shape=(nn,), units='s**2')
+
+        self.add_output('z_dot', shape=(nn, 2), units='s')
+
+        cs = np.repeat(np.arange(nn, dtype=int), 2)
+        ar2 = np.arange(2 * nn, dtype=int)
+        dzdot_dz_pattern = np.arange(2 * nn, step=2, dtype=int)
+        self.declare_partials(of='z_dot', wrt='z', rows=dzdot_dz_pattern, cols=dzdot_dz_pattern, val=1.0)
+        self.declare_partials(of='z_dot', wrt='t', rows=ar2, cols=cs)
+        dzdot_dp_rows = np.arange(2 * nn, step=2, dtype=int)
+        dzdot_dp_cols = np.arange(nn, dtype=int)
+        self.declare_partials(of='z_dot', wrt='p', rows=dzdot_dp_rows, cols=dzdot_dp_cols, val=1.0)
+
+    def compute(self, inputs, outputs):
+        """
+        Compute the outputs of SimpleVectorizedODE.
+
+        Parameters
+        ----------
+        inputs : Vector
+            Vector of inputs.
+        outputs : Vector
+            Vector of outputs.
         """
-        self.envs = envs
+        z = inputs['z']
+        t = inputs['t']
+        p = inputs['p']
+        outputs['z_dot'][:, 0] = z[:, 0] - t**2 + p
+        outputs['z_dot'][:, 1] = 10 * t
 
-    def __call__(self, fnc):
+    def compute_partials(self, inputs, partials):
         """
-        Apply the decorator.
+        Compute the partials of SimpleVectorizedODE.
 
         Parameters
         ----------
-        fnc : function
-            The function being wrapped.
+        inputs : Vector
+            Vector of inputs.
+        partials : Dictionary
+            Vector of partials.
         """
-        def wrap(*args, **kwargs):
-            saved = {}
-            try:
-                for k, v in self.envs.items():
-                    saved[k] = os.environ.get(k)
-                    os.environ[k] = v  # will raise exception if v is not a string
-
-                return fnc(*args, **kwargs)
-            finally:
-                # put environment back as it was
-                for k, v in saved.items():
-                    if v is None:
-                        del os.environ[k]
-                    else:
-                        os.environ[k] = v
-
-        return wrap
+        t = inputs['t']
+        partials['z_dot', 't'][0::2] = -2*t
+        partials['z_dot', 't'][1::2] = 10
diff --git a/joss/test/test_cannonball_for_joss.py b/joss/test/test_cannonball_for_joss.py
index e4701480..c1e30dc1 100644
--- a/joss/test/test_cannonball_for_joss.py
+++ b/joss/test/test_cannonball_for_joss.py
@@ -136,7 +136,7 @@ class TestCannonballForJOSS(unittest.TestCase):
         traj.add_parameter('mass', units='kg', val=0.01, static_target=True)
         traj.add_parameter('area', units='m**2', static_target=True)
 
-        tx = dm.Radau(num_segments=5, order=3, compressed=True)
+        tx = dm.Radau(num_segments=5, order=3, compressed=False)
         ascent = dm.Phase(transcription=tx, ode_class=CannonballODE)
         traj.add_phase('ascent', ascent)
 
@@ -154,7 +154,7 @@ class TestCannonballForJOSS(unittest.TestCase):
         #     we will set it to zero so that the phase ends at apogee)
         ascent.add_state('gam', units='rad', rate_source='gam_dot',
                          fix_initial=False, fix_final=True)
-        ascent.set_time_options(fix_initial=True, duration_bounds=(1, 100),
+        ascent.set_time_options(fix_initial=True, duration_bounds=(10, 100),
                                 duration_ref=100, units='s')
 
         ascent.add_parameter('mass', units='kg', val=0.01, static_target=True)
@@ -204,23 +204,41 @@ class TestCannonballForJOSS(unittest.TestCase):
 
         # Set Initial guesses for static dvs and ascent
         p.set_val('static_calcs.radius', 0.05, units='m')
-        p.set_val('traj.ascent.t_duration', 10.0)
 
-        p.set_val('traj.ascent.states:r', ascent.interp('r', [0, 100]))
-        p.set_val('traj.ascent.states:h', ascent.interp('h', [0, 100]))
-        p.set_val('traj.ascent.states:v', ascent.interp('v', [200, 150]))
-        p.set_val('traj.ascent.states:gam', ascent.interp('gam', [25, 0]), units='deg')
+        ascent.set_time_val(initial=0, duration=10.0)
+        ascent.set_state_val('r', [0, 100])
+        ascent.set_state_val('h', [0, 100])
+        ascent.set_state_val('v', [200, 150])
+        ascent.set_state_val('gam', [25, 0], units='deg')
 
-        # more intitial guesses for descent
-        p.set_val('traj.descent.t_initial', 10.0)
-        p.set_val('traj.descent.t_duration', 10.0)
+        # These set_val commands applied to previous versions of dymos.
+        # With the updated transcriptions as of dymos 1.14.0, we would also need to set
+        # values for `traj.adscent.initial_states:{r|h|v}` and
+        # `traj.adscent.final_states:{r|h|v}`.  By using the `set_state_val` interface instead,
+        # dymos internally sets all appropriate variables pertaining to the states.
 
-        p.set_val('traj.descent.states:r', descent.interp('r', [100, 200]))
-        p.set_val('traj.descent.states:h', descent.interp('h', [100, 0]))
-        p.set_val('traj.descent.states:v', descent.interp('v', [150, 200]))
-        p.set_val('traj.descent.states:gam', descent.interp('gam', [0, -45]), units='deg')
+        # p.set_val('traj.ascent.t_duration', 10.0)
+        # p.set_val('traj.ascent.states:r', ascent.interp('r', [0, 100]))
+        # p.set_val('traj.ascent.states:h', ascent.interp('h', [0, 100]))
+        # p.set_val('traj.ascent.states:v', ascent.interp('v', [200, 150]))
+        # p.set_val('traj.ascent.states:gam', ascent.interp('gam', [25, 0]), units='deg')
 
-        dm.run_problem(p, simulate=True, make_plots=True)
+        # more intitial guesses for descent
+        descent.set_time_val(initial=10.0, duration=10.0)
+        descent.set_state_val('r', [100, 200])
+        descent.set_state_val('h', [100, 0])
+        descent.set_state_val('v', [150, 200])
+        descent.set_state_val('gam', [0, -45], units='deg')
+
+        # Pre-dymos 1.14.0 code
+        # p.set_val('traj.descent.t_initial', 10.0)
+        # p.set_val('traj.descent.t_duration', 10.0)
+        # p.set_val('traj.descent.states:r', descent.interp('r', [100, 200]))
+        # p.set_val('traj.descent.states:h', descent.interp('h', [100, 0]))
+        # p.set_val('traj.descent.states:v', descent.interp('v', [150, 200]))
+        # p.set_val('traj.descent.states:gam', descent.interp('gam', [0, -45]), units='deg')
+
+        dm.run_problem(p, run_driver=True, simulate=True, make_plots=True)
 
         fig, ax = plt.subplots()
         x0 = p.get_val('traj.ascent.timeseries.r', units='m')
@@ -236,7 +254,7 @@ class TestCannonballForJOSS(unittest.TestCase):
         fig.savefig('cannonball_hr.png', bbox_inches='tight')
         # End code for paper
 
-        assert_near_equal(x1[-1], 3064, tolerance=1.0E-4)
+        assert_near_equal(x1[-1], 3064, tolerance=1.0E-2)
 
 
 if __name__ == '__main__':