Fixed some issues that were causing simulation to fail under MPI.

docs/dymos_book/examples/hull/hull_problem.ipynb

@@ -468,8 +468,7 @@
     "for ax in [x_ax, xL_ax, u_ax]:\n",
     "    ax.grid(True, alpha=0.2)\n",
     "    \n",
-    "plt.figlegend([x_sol_handle, x_sim_handle], ['solution', 'simulation'], ncol=2, loc='lower center');\n",
-    "\n"
+    "plt.figlegend([x_sol_handle, x_sim_handle], ['solution', 'simulation'], ncol=2, loc='lower center');\n"
    ]
   },
   {
@@ -507,7 +506,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.11.0"
   }
  },
  "nbformat": 4,

dymos/examples/finite_burn_orbit_raise/test/test_ex_two_burn_orbit_raise_mpi.py

@@ -1,5 +1,6 @@
 import unittest
 
+import openmdao.api as om
 from openmdao.utils.assert_utils import assert_near_equal
 from openmdao.utils.mpi import MPI
 from openmdao.utils.testing_utils import use_tempdirs, require_pyoptsparse
@@ -16,24 +17,40 @@ class TestExampleTwoBurnOrbitRaiseMPI(unittest.TestCase):
     def test_ex_two_burn_orbit_raise_mpi(self):
         optimizer = 'IPOPT'
 
+        CONNECTED = False
+
         p = two_burn_orbit_raise_problem(transcription='gauss-lobatto', transcription_order=3,
-                                         compressed=False, optimizer=optimizer, simulate=False,
-                                         show_output=False)
+                                         compressed=False, optimizer=optimizer, simulate=True,
+                                         connected=CONNECTED, show_output=False)
+
+        sol_case = om.CaseReader('dymos_solution.db').get_case('final')
+        sim_case = om.CaseReader('dymos_simulation.db').get_case('final')
+
+        # The last phase in this case is run in reverse time if CONNECTED=True,
+        # so grab the correct index to test the resulting delta-V.
+        end_idx = 0 if CONNECTED else -1
 
-        if p.model.traj.phases.burn2 in p.model.traj.phases._subsystems_myproc:
-            assert_near_equal(p.get_val('traj.burn2.states:deltav')[-1], 0.3995,
-                              tolerance=2.0E-3)
+        assert_near_equal(sol_case.get_val('traj.burn2.timeseries.deltav')[end_idx], 0.3995, tolerance=2.0E-3)
+        assert_near_equal(sim_case.get_val('traj.burn2.timeseries.deltav')[end_idx], 0.3995, tolerance=2.0E-3)
 
     def test_ex_two_burn_orbit_raise_connected_mpi(self):
         optimizer = 'IPOPT'
 
+        CONNECTED = True
+
         p = two_burn_orbit_raise_problem(transcription='gauss-lobatto', transcription_order=3,
-                                         compressed=False, optimizer=optimizer, simulate=False,
-                                         connected=True, show_output=False)
+                                         compressed=False, optimizer=optimizer, simulate=True,
+                                         connected=CONNECTED, show_output=False)
+
+        sol_case = om.CaseReader('dymos_solution.db').get_case('final')
+        sim_case = om.CaseReader('dymos_simulation.db').get_case('final')
+
+        # The last phase in this case is run in reverse time if CONNECTED=True,
+        # so grab the correct index to test the resulting delta-V.
+        end_idx = 0 if CONNECTED else -1
 
-        if p.model.traj.phases.burn2 in p.model.traj.phases._subsystems_myproc:
-            assert_near_equal(p.get_val('traj.burn2.states:deltav')[0], 0.3995,
-                              tolerance=2.0E-3)
+        assert_near_equal(sol_case.get_val('traj.burn2.timeseries.deltav')[end_idx], 0.3995, tolerance=2.0E-3)
+        assert_near_equal(sim_case.get_val('traj.burn2.timeseries.deltav')[end_idx], 0.3995, tolerance=2.0E-3)
 
 
 if __name__ == '__main__':  # pragma: no cover

dymos/examples/hull_problem/test/test_hull_problem.py

@@ -61,7 +61,7 @@ def solution(x0, td):
 
     def test_hull_gauss_lobatto(self):
         p = self.make_problem(transcription=GaussLobatto)
-        dm.run_problem(p)
+        dm.run_problem(p, simulate=True)
 
         xf, uf = self.solution(1.5, 10)
 
@@ -75,7 +75,7 @@ def test_hull_gauss_lobatto(self):
 
     def test_hull_radau(self):
         p = self.make_problem(transcription=Radau)
-        dm.run_problem(p)
+        dm.run_problem(p, simulate=True)
 
         xf, uf = self.solution(1.5, 10)
 
@@ -89,7 +89,7 @@ def test_hull_radau(self):
 
     def test_hull_shooting(self):
         p = self.make_problem(transcription=ExplicitShooting)
-        dm.run_problem(p)
+        dm.run_problem(p, simulate=True)
 
         xf, uf = self.solution(1.5, 10)
 

dymos/phase/phase.py

@@ -2363,9 +2363,9 @@ def simulate(self, times_per_seg=10, method=_unspecified, atol=_unspecified, rto
             sim_prob.add_recorder(rec)
 
         sim_prob.setup(check=True)
+        sim_prob.final_setup()
 
-        # sim_phase.set_val_from_phase(from_phase=self)  # TODO: use this for OpenMDAO >= 3.25.1
-        sim_phase.initialize_values_from_phase(prob=sim_prob, from_phase=self)
+        sim_phase.set_vals_from_phase(from_phase=self)
 
         print(f'\nSimulating phase {self.pathname}')
         sim_prob.run_model()

dymos/phase/simulation_phase.py

@@ -93,7 +93,7 @@ def __init__(self, from_phase, times_per_seg=None, method=_unspecified, atol=_un
         self._timeseries = {ts_name: ts_options for ts_name, ts_options in self._timeseries.items()
                             if ts_name == 'timeseries'}
 
-    def set_val_from_phase(self, from_phase):
+    def set_vals_from_phase(self, from_phase):
         """
         Set the necessary values to simulate the phase based on variables in the given phase.
 
@@ -102,27 +102,38 @@ def set_val_from_phase(self, from_phase):
         from_phase : Phase
             The dymos phase from which this simulation phase should pull its values.
         """
+        # The use of `from_src=False` in the get_val calls here is due to the fact that the input/output
+        # vectors are in `from_phase` are already populated and we don't need to track these values
+        # to their ultimate source.
 
-        t_initial = from_phase.get_val('t_initial', units=self.time_options['units'])
+        t_initial = from_phase.get_val('t_initial', units=self.time_options['units'], from_src=False)
         self.set_val('t_initial', t_initial, units=self.time_options['units'])
 
-        t_duration = from_phase.get_val('t_duration', units=self.time_options['units'])
+        t_duration = from_phase.get_val('t_duration', units=self.time_options['units'], from_src=False)
         self.set_val('t_duration', t_duration, units=self.time_options['units'])
 
+        avail_io = {meta['prom_name'] for meta in
+                    from_phase.get_io_metadata(iotypes=('input', 'output'), get_remote=True).values()}
+
         for name, options in self.state_options.items():
-            val = from_phase.get_val(f'states:{name}', units=options['units'])[0, ...]
+            if f'states:{name}' in avail_io:
+                val = from_phase.get_val(f'states:{name}', units=options['units'], from_src=False)[0, ...]
+            elif f'initial_states:{name}' in avail_io:
+                val = from_phase.get_val(f'initial_states:{name}', units=options['units'], from_src=False)
+            else:
+                raise RuntimeError('Unable to find state values in original phase')
             self.set_val(f'initial_states:{name}', val, units=options['units'])
 
         for name, options in self.parameter_options.items():
-            val = from_phase.get_val(f'parameters:{name}', units=options['units'])
+            val = from_phase.get_val(f'parameters:{name}', units=options['units'], from_src=False)
             self.set_val(f'parameters:{name}', val, units=options['units'])
 
         for name, options in self.control_options.items():
-            val = from_phase.get_val(f'controls:{name}', units=options['units'])
+            val = from_phase.get_val(f'controls:{name}', units=options['units'], from_src=False)
             self.set_val(f'controls:{name}', val, units=options['units'])
 
         for name, options in self.polynomial_control_options.items():
-            val = from_phase.get_val(f'polynomial_controls:{name}', units=options['units'])
+            val = from_phase.get_val(f'polynomial_controls:{name}', units=options['units'], from_src=False)
             self.set_val(f'polynomial_controls:{name}', val, units=options['units'])
 
     def initialize_values_from_phase(self, prob, from_phase, phase_path=''):

dymos/trajectory/trajectory.py

@@ -1,3 +1,4 @@
+import warnings
 from collections import OrderedDict
 from collections.abc import Sequence
 import itertools
@@ -352,6 +353,11 @@ def _setup_parameters(self):
                             phs.add_parameter(name, **kwargs)
 
     def _setup_linkages(self):
+
+        if self.options['sim_mode']:
+            # Under simulation, theres no need to enforce any linkages
+            return
+
         has_linkage_constraints = False
 
         err_template = '{traj}: Phase `{phase1}` links variable `{var1}` to phase ' \
@@ -528,8 +534,8 @@ def _configure_phase_options_dicts(self):
         and units options to all procs for all dymos variables.
         """
         for phase in self._phases.values():
-            all_dicts = [phase.state_options, phase.control_options, phase.parameter_options,
-                         phase.polynomial_control_options]
+            all_dicts = [phase.state_options, phase.control_options,
+                         phase.parameter_options, phase.polynomial_control_options]
 
             for opt_dict in all_dicts:
                 for options in opt_dict.values():
@@ -770,6 +776,11 @@ def _is_valid_linkage(self, phase_name_a, phase_name_b, loc_a, loc_b, var_a, var
             return True, ''
 
     def _configure_linkages(self):
+
+        if self.options['sim_mode']:
+            # If this is a simulation trajectory, theres no need to link the phases.
+            return
+
         connected_linkage_inputs = []
 
         def _print_on_rank(rank=0, *args, **kwargs):
@@ -1412,7 +1423,7 @@ def simulate(self, times_per_seg=10, method=_unspecified, atol=_unspecified, rto
 
         sim_traj.parameter_options.update(self.parameter_options)
 
-        sim_prob = om.Problem(model=om.Group(), reports=reports)
+        sim_prob = om.Problem(model=om.Group(), reports=reports, comm=self.comm)
 
         traj_name = self.name if self.name else 'sim_traj'
         sim_prob.model.add_subsystem(traj_name, sim_traj)
@@ -1425,23 +1436,28 @@ def simulate(self, times_per_seg=10, method=_unspecified, atol=_unspecified, rto
             # record_outputs is need to capture the timeseries outputs
             sim_prob.recording_options['record_outputs'] = True
 
-        sim_prob.setup()
+        with warnings.catch_warnings():
+            # Some timeseries options are duplicated (expression options may be provide duplicate shape)
+            # These filters suppress these warnings during simulation when they are not the
+            # fault of the user.
+            warnings.filterwarnings(action='ignore', category=om.UnusedOptionWarning)
+            warnings.filterwarnings(action='ignore', category=om.SetupWarning)
+            sim_prob.setup()
+            sim_prob.final_setup()
 
         # Assign trajectory parameter values
         for name in self.parameter_options:
-            sim_prob_prom_path = f'{traj_name}.parameters:{name}'
-            sim_prob.set_val(sim_prob_prom_path, self.get_val(f'parameters:{name}'))
+            sim_traj.set_val(f'parameters:{name}', self.get_val(f'parameters:{name}'))
 
-        for phase_name, phs in sim_traj._phases.items():
-            # TODO: use the following method once OpenMDAO >= 3.25.1
-            # phs.set_val_from_phase(from_phase=self._phases[phase_name])
-            phs.initialize_values_from_phase(prob=sim_prob,
-                                             from_phase=self._phases[phase_name],
-                                             phase_path=traj_name)
+        for sim_phase_name, sim_phase in sim_traj._phases.items():
+            if sim_phase._is_local:
+                sim_phase.set_vals_from_phase(from_phase=self._phases[sim_phase_name])
 
-        print(f'\nSimulating trajectory {self.pathname}')
+        if sim_traj.comm.rank == 0:
+            print(f'\nSimulating trajectory {self.pathname}')
         sim_prob.run_model(case_prefix=case_prefix, reset_iter_counts=reset_iter_counts)
-        print(f'Done simulating trajectory {self.pathname}')
+        if sim_traj.comm.rank == 0:
+            print(f'Done simulating trajectory {self.pathname}')
         if record_file:
             _case_prefix = '' if case_prefix is None else f'{case_prefix}_'
             sim_prob.record(f'{_case_prefix}final')

dymos/utils/introspection.py

@@ -5,6 +5,7 @@
 
 import openmdao.api as om
 import numpy as np
+from openmdao.utils.units import simplify_unit
 from openmdao.utils.general_utils import ensure_compatible
 from dymos.utils.misc import _unspecified, _none_or_unspecified
 from .._options import options as dymos_options
@@ -334,7 +335,8 @@ def configure_controls_introspection(control_options, ode, time_units='s'):
         if rate_targets:
             if options['units'] is _unspecified:
                 rate_target_units = _get_common_metadata(rate_targets, metadata_key='units')
-                options['units'] = time_units if rate_target_units is None else f'{rate_target_units}*{time_units}'
+                options['units'] = time_units if rate_target_units is None else \
+                    simplify_unit(f'{rate_target_units}*{time_units}')
 
             if options['shape'] in _none_or_unspecified:
                 shape = _get_common_metadata(rate_targets, metadata_key='shape')
@@ -356,7 +358,7 @@ def configure_controls_introspection(control_options, ode, time_units='s'):
             if options['units'] is _unspecified:
                 rate2_target_units = _get_common_metadata(rate_targets, metadata_key='units')
                 options['units'] = f'{time_units**2}' if rate2_target_units is None \
-                    else f'{rate2_target_units}*{time_units}**2'
+                    else simplify_unit(f'{rate2_target_units}*{time_units}**2')
 
             if options['shape'] in _none_or_unspecified:
                 shape = _get_common_metadata(rate2_targets, metadata_key='shape')
@@ -605,7 +607,7 @@ def configure_states_introspection(state_options, time_options, control_options,
                                        f'- Tag the state rate source `{rate_src}` with `dymos.state_units:{{units}}`\n'
                                        f'- Use the `set_state_options(\'{state_name}\', units={{units}})` method on the phase.')
                 else:
-                    options['units'] = f'{rate_src_units}*{time_units}'
+                    options['units'] = simplify_unit(f'{rate_src_units}*{time_units}')
 
 
 def configure_analytic_states_introspection(state_options, ode):

setup.py

@@ -69,7 +69,7 @@
     packages=find_packages(),
     python_requires=">=3.8",
     install_requires=[
-        'openmdao>=3.17.0',
+        'openmdao>=3.26.0',
         'numpy',
         'scipy'
     ],