Added default_nonlinear_solver and default_linear_solver options to T…

…rajectory. (#1016)

* Added default_nonlinear_solver and default_linear_solver options to trajectory.\nRemoved automatic recording of inputs in run_problem and under simulation.\nSimulationPhases now automatically set state option `input_initial` to False.

* added tests to ensure warnings are raised/not raised as appropriate.

* jax/jaxlib version bump for CI

* jax version bump for docs CI

* Fixed some instances where the deprecated dm.load_case method was still being used.

* one more test for a changing grid

* Bump Oldesst support OpenMDAO to 3.28 for load_case updated functionality. Fix for parallel phases when MPI is not available or comm.size <= 1.

* A few test cleanup issues

* more test cleanup for the case where MPI is unavailable.

* one more MPI-related skip

* turn off IPOPT printing for some tests where it had slipped in

* Added skips for low-thrust Birkhoff test for now.

* removing redundant if MPI check

.github/workflows/dymos_docs_workflow.yml

@@ -33,7 +33,7 @@ jobs:
             PYOPTSPARSE: 'v2.9.3'
             OPENMDAO: 'latest'
             OPTIONAL: '[docs]'
-            JAX: '0.3.24'
+            JAX: '0.4.14'
             PUBLISH_DOCS: 1
 
           # make sure the latest versions of things don't break the docs

.github/workflows/dymos_tests_workflow.yml

@@ -36,7 +36,7 @@ jobs:
             SNOPT: 7.7
             OPENMDAO: 'latest'
             OPTIONAL: '[all]'
-            JAX: '0.3.24'
+            JAX: '0.4.14'
 
           # baseline versions except no pyoptsparse or SNOPT
           - NAME: no_pyoptsparse
@@ -88,7 +88,7 @@ jobs:
             PETSc: 3.13
             PYOPTSPARSE: 'v2.6.1'
             SNOPT: 7.2
-            OPENMDAO: 3.27.0
+            OPENMDAO: 3.28.0
             OPTIONAL: '[test]'
 
     steps:

dymos/examples/balanced_field/doc/test_doc_balanced_field_length.py

@@ -29,7 +29,7 @@ def test_balanced_field_length_for_docs(self):
         p.driver.declare_coloring()
         p.driver.options['print_results'] = False
         if optimizer == 'IPOPT':
-            p.driver.opt_settings['print_level'] = 5
+            p.driver.opt_settings['print_level'] = 0
             p.driver.opt_settings['derivative_test'] = 'first-order'
 
         # First Phase: Brake release to V1 - both engines operable

dymos/examples/double_integrator/test/test_double_integrator_breakwell.py

@@ -22,7 +22,7 @@ def double_integrator_direct_collocation(transcription='gauss-lobatto', compress
     if optimizer == 'IPOPT':
         p.driver.opt_settings['max_iter'] = 5000
         p.driver.opt_settings['alpha_for_y'] = 'safer-min-dual-infeas'
-        p.driver.opt_settings['print_level'] = 5
+        p.driver.opt_settings['print_level'] = 0
         p.driver.opt_settings['nlp_scaling_method'] = 'gradient-based'
         p.driver.opt_settings['tol'] = 1.0E-3
         p.driver.opt_settings['constr_viol_tol'] = 1.0E-6

dymos/examples/finite_burn_orbit_raise/finite_burn_orbit_raise_problem.py

@@ -8,7 +8,7 @@
 
 
 def make_traj(transcription='gauss-lobatto', transcription_order=3, compressed=False,
-              connected=False):
+              connected=False, default_nonlinear_solver=None, default_linear_solver=None):
     """
     Build a traejctory for the finite burn orbit raise problem.
 
@@ -31,7 +31,8 @@ def make_traj(transcription='gauss-lobatto', transcription_order=3, compressed=F
     t = {'gauss-lobatto': dm.GaussLobatto(num_segments=5, order=transcription_order, compressed=compressed),
          'radau': dm.Radau(num_segments=5, order=transcription_order, compressed=compressed)}
 
-    traj = dm.Trajectory()
+    traj = dm.Trajectory(default_nonlinear_solver=default_nonlinear_solver,
+                         default_linear_solver=default_linear_solver)
 
     traj.add_parameter('c', opt=False, val=1.5, units='DU/TU',
                        targets={'burn1': ['c'], 'burn2': ['c']})
@@ -156,19 +157,14 @@ def make_traj(transcription='gauss-lobatto', transcription_order=3, compressed=F
 
         traj.link_phases(phases=['burn1', 'burn2'], vars=['accel'])
 
-    if connected and MPI:
-        # If running connected and under MPI the phases subsystem requires a Nonlinear Block Jacobi solver.
-        # This is not the most efficient way to actually solve this problem but it demonstrates access
-        # to the traj.phases subsystem before setup.
-        traj.phases.nonlinear_solver = om.NonlinearBlockJac(iprint=0)
-        traj.phases.linear_solver = om.PETScKrylov()
-
     return traj
 
 
 def two_burn_orbit_raise_problem(transcription='gauss-lobatto', optimizer='SLSQP', r_target=3.0,
                                  transcription_order=3, compressed=False, run_driver=True,
-                                 max_iter=300, simulate=True, show_output=True, connected=False, restart=None):
+                                 max_iter=300, simulate=True, show_output=True, connected=False, restart=None,
+                                 solution_record_file='dymos_solution.db', simulation_record_file='dymos_simulation.db',
+                                 default_nonlinear_solver=None, default_linear_solver=None):
     """
     Build and run the finite burn orbit raise problem.
 
@@ -227,10 +223,12 @@ def two_burn_orbit_raise_problem(transcription='gauss-lobatto', optimizer='SLSQP
         p.driver.opt_settings['mu_strategy'] = 'monotone'
         p.driver.opt_settings['derivative_test'] = 'first-order'
         if show_output:
-            p.driver.opt_settings['print_level'] = 5
+            p.driver.opt_settings['print_level'] = 0
 
     traj = make_traj(transcription=transcription, transcription_order=transcription_order,
-                     compressed=compressed, connected=connected)
+                     compressed=compressed, connected=connected,
+                     default_nonlinear_solver=default_nonlinear_solver,
+                     default_linear_solver=default_linear_solver)
     p.model.add_subsystem('traj', subsys=traj)
 
     # Needed to move the direct solver down into the phases for use with MPI.
@@ -293,6 +291,7 @@ def two_burn_orbit_raise_problem(transcription='gauss-lobatto', optimizer='SLSQP
             p.set_val('traj.burn2.controls:u1', val=burn2.interp('u1', [0, 0]))
 
     if run_driver or simulate:
-        dm.run_problem(p, run_driver=run_driver, simulate=simulate, restart=restart, make_plots=True)
+        dm.run_problem(p, run_driver=run_driver, simulate=simulate, restart=restart, make_plots=False,
+                       solution_record_file=solution_record_file, simulation_record_file=simulation_record_file)
 
     return p

dymos/examples/finite_burn_orbit_raise/test/test_multi_phase_restart.py

@@ -1,19 +1,23 @@
 import unittest
+import warnings
 
 import openmdao.api as om
-from openmdao.utils.assert_utils import assert_near_equal
+from openmdao.utils.assert_utils import assert_near_equal, assert_warnings, assert_no_warning
 from openmdao.utils.testing_utils import use_tempdirs, require_pyoptsparse
+from openmdao.utils.mpi import MPI
 import scipy
 
 from dymos.examples.finite_burn_orbit_raise.finite_burn_orbit_raise_problem import two_burn_orbit_raise_problem
 from dymos.utils.testing_utils import assert_cases_equal
 
 
-# This test is separate because connected phases aren't directly parallelizable.
 @require_pyoptsparse(optimizer='IPOPT')
 @use_tempdirs
 class TestExampleTwoBurnOrbitRaiseConnectedRestart(unittest.TestCase):
 
+    N_PROCS = 3
+
+    @unittest.skipUnless(MPI, "MPI is required.")
     def test_ex_two_burn_orbit_raise_connected(self):
         optimizer = 'IPOPT'
 
@@ -39,7 +43,6 @@ def test_ex_two_burn_orbit_raise_connected(self):
         sim_case2 = om.CaseReader('dymos_simulation.db').get_case('final')
 
         # Verify that the second case has the same inputs and outputs
-        assert_cases_equal(case1, p, tol=1.0E-8)
         assert_cases_equal(sim_case1, sim_case2, tol=1.0E-8)
 
     def test_restart_from_solution_radau(self):
@@ -63,21 +66,40 @@ def test_restart_from_solution_radau(self):
         sim_case2 = om.CaseReader('dymos_simulation.db').get_case('final')
 
         # Verify that the second case has the same inputs and outputs
-        assert_cases_equal(case1, p, tol=1.0E-9)
         assert_cases_equal(sim_case1, sim_case2, tol=1.0E-8)
 
 
-# This test is separate because connected phases aren't directly parallelizable.
 @require_pyoptsparse(optimizer='IPOPT')
 @use_tempdirs
 class TestExampleTwoBurnOrbitRaiseConnected(unittest.TestCase):
 
+    N_PROCS = 3
+
+    @unittest.skipUnless(MPI, "MPI is required.")
     def test_ex_two_burn_orbit_raise_connected(self):
         optimizer = 'IPOPT'
 
-        p = two_burn_orbit_raise_problem(transcription='gauss-lobatto', transcription_order=3,
-                                         compressed=False, optimizer=optimizer,
-                                         show_output=False, connected=True)
+        unexpected_warnings = \
+            [(om.OpenMDAOWarning,
+              "'traj' <class Trajectory>: Setting phases.nonlinear_solver to `om.NonlinearBlockJac(iprint=0)`.\n"
+              "Connected phases in parallel require a non-default nonlinear solver.\n"
+              "Use traj.options[\'default_nonlinear_solver\'] to explicitly set the solver."),
+             (om.OpenMDAOWarning,
+              "'traj' <class Trajectory>: Setting phases.linear_solver to `om.PETScKrylov()`.\n"
+              "Connected phases in parallel require a non-default linear solver.\n"
+              "Use traj.options[\'default_linear_solver\'] to explicitly set the solver.")]
+
+        with warnings.catch_warnings(record=True) as w:
+            p = two_burn_orbit_raise_problem(transcription='gauss-lobatto', transcription_order=3,
+                                             compressed=False, optimizer=optimizer,
+                                             show_output=False, connected=True,
+                                             default_nonlinear_solver=om.NonlinearBlockJac(iprint=0),
+                                             default_linear_solver=om.PETScKrylov())
+
+        for category, msg in unexpected_warnings:
+            for warn in w:
+                if (issubclass(warn.category, category) and str(warn.message) == msg):
+                    raise AssertionError(f"Saw unexpected warning {category.__name__}: {msg}")
 
         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,
@@ -90,37 +112,56 @@ def test_ex_two_burn_orbit_raise_connected(self):
         p = two_burn_orbit_raise_problem(transcription='gauss-lobatto', transcription_order=3,
                                          compressed=False, optimizer=optimizer, run_driver=False,
                                          show_output=False, restart='dymos_solution.db',
-                                         connected=True)
-
-        sim_case2 = om.CaseReader('dymos_simulation.db').get_case('final')
+                                         connected=True, solution_record_file='dymos_solution2.db',
+                                         simulation_record_file='dymos_simulation2.db')
+#
 
+        case2 = om.CaseReader('dymos_solution2.db').get_case('final')
+        sim_case2 = om.CaseReader('dymos_simulation2.db').get_case('final')
+#
         # Verify that the second case has the same inputs and outputs
-        assert_cases_equal(case1, p, tol=1.0E-8)
+        assert_cases_equal(case1, case2, tol=1.0E-8)
         assert_cases_equal(sim_case1, sim_case2, tol=1.0E-8)
 
+    @unittest.skipUnless(MPI, "MPI is required.")
     def test_restart_from_solution_radau_to_connected(self):
         optimizer = 'IPOPT'
 
-        p = two_burn_orbit_raise_problem(transcription='radau', transcription_order=3,
-                                         compressed=False, optimizer=optimizer, show_output=False)
-
-        case1 = om.CaseReader('dymos_solution.db').get_case('final')
-        sim_case1 = om.CaseReader('dymos_simulation.db').get_case('final')
-
-        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)
-
-        # Run again without an actual optimzier
-        two_burn_orbit_raise_problem(transcription='radau', transcription_order=3,
-                                     compressed=False, optimizer=optimizer, run_driver=False,
-                                     show_output=False, restart='dymos_solution.db', connected=True)
-
-        sim_case2 = om.CaseReader('dymos_simulation.db').get_case('final')
-
-        # Verify that the second case has the same inputs and outputs
-        assert_cases_equal(case1, p, tol=1.0E-9, require_same_vars=False)
-        assert_cases_equal(sim_case1, sim_case2, tol=1.0E-8)
+        expected_warnings = \
+            [(om.OpenMDAOWarning,
+              "'traj' <class Trajectory>: Setting phases.nonlinear_solver to `om.NonlinearBlockJac(iprint=0)`.\n"
+              "Connected phases in parallel require a non-default nonlinear solver.\n"
+              "Use traj.options[\'default_nonlinear_solver\'] to explicitly set the solver."),
+             (om.OpenMDAOWarning,
+              "'traj' <class Trajectory>: Setting phases.linear_solver to `om.PETScKrylov()`.\n"
+              "Connected phases in parallel require a non-default linear solver.\n"
+              "Use traj.options[\'default_linear_solver\'] to explicitly set the solver.")]
+
+        with assert_warnings(expected_warnings):
+            p = two_burn_orbit_raise_problem(transcription='radau', transcription_order=3,
+                                             compressed=False, optimizer=optimizer,
+                                             show_output=False, connected=True)
+
+            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=4.0E-3)
+
+            case1 = om.CaseReader('dymos_solution.db').get_case('final')
+            sim_case1 = om.CaseReader('dymos_simulation.db').get_case('final')
+
+            # Run again without an actual optimizer
+            p = two_burn_orbit_raise_problem(transcription='radau', transcription_order=3,
+                                             compressed=False, optimizer=optimizer, run_driver=False,
+                                             show_output=False, restart='dymos_solution.db',
+                                             connected=True, solution_record_file='dymos_solution2.db',
+                                             simulation_record_file='dymos_simulation2.db')
+
+            case2 = om.CaseReader('dymos_solution2.db').get_case('final')
+            sim_case2 = om.CaseReader('dymos_simulation2.db').get_case('final')
+
+            # Verify that the second case has the same inputs and outputs
+            assert_cases_equal(case1, case2, tol=1.0E-8)
+            assert_cases_equal(sim_case1, sim_case2, tol=1.0E-8)
 
 
 if __name__ == '__main__':  # pragma: no cover

dymos/examples/hyper_sensitive/test/test_hyper_sensitive.py

@@ -40,7 +40,7 @@ def make_problem(self, transcription=dm.GaussLobatto, optimizer='SLSQP', numseg=
             p.driver.opt_settings['Major feasibility tolerance'] = 1.0E-6
             p.driver.opt_settings['Major optimality tolerance'] = 1.0E-6
         elif optimizer == 'IPOPT':
-            p.driver.opt_settings['print_level'] = 5
+            p.driver.opt_settings['print_level'] = 0
             p.driver.opt_settings['mu_strategy'] = 'adaptive'
             p.driver.opt_settings['bound_mult_init_method'] = 'mu-based'
             p.driver.opt_settings['mu_init'] = 0.01

dymos/examples/low_thrust_spiral/test/test_low_thrust_spiral.py

@@ -10,7 +10,6 @@
 show_plots = False
 
 
-@require_pyoptsparse(optimizer='SNOPT')
 def low_thrust_spiral_direct_collocation(grid_type='lgl'):
 
     optimizer = 'SNOPT'
@@ -66,22 +65,17 @@ def low_thrust_spiral_direct_collocation(grid_type='lgl'):
     return p
 
 
+@require_pyoptsparse(optimizer='SNOPT')
 @use_tempdirs
 class TestLowThrustSpiral(unittest.TestCase):
 
-    # @classmethod
-    # def tearDownClass(cls):
-    #     for filename in ['total_coloring.pkl', 'SLSQP.out', 'SNOPT_print.out']:
-    #         if os.path.exists(filename):
-    #             os.remove(filename)
-
     @staticmethod
-    def _assert_results(p, tol=1.0E-4):
-        # t = p.get_val('traj.phase0.timeseries.time')
-        #
-        # assert_near_equal(t[-1], 228, tolerance=tol)
+    def _assert_results(p, tol=0.05):
+        t = p.get_val('traj.phase0.timeseries.time')
+        assert_near_equal(t[-1], 228, tolerance=tol)
         return
 
+    @unittest.skip('Long running test skipped on CI.')
     def test_low_thrust_spiral_lgl(self):
         p = low_thrust_spiral_direct_collocation(grid_type='lgl')
         dm.run_problem(p)
@@ -97,6 +91,7 @@ def test_low_thrust_spiral_lgl(self):
 
         self._assert_results(p)
 
+    @unittest.skip('Long running test skipped on CI.')
     def test_low_thrust_spiral_cgl(self):
         p = low_thrust_spiral_direct_collocation(grid_type='cgl')
         dm.run_problem(p)

dymos/phase/simulation_phase.py

@@ -91,6 +91,7 @@ def __init__(self, from_phase, times_per_seg=_unspecified, method=_unspecified,
         # Remove invalid options
         for state_name, options in self.state_options.items():
             options['fix_final'] = False  # ExplicitShooting will raise if `fix_final` is True for any states.
+            options['input_initial'] = False  # Only simulate from the initial value, do not connect.
 
         # Remove all but the default timeseries object
         self._timeseries = {ts_name: ts_options for ts_name, ts_options in self._timeseries.items()

dymos/run_problem.py

@@ -75,8 +75,7 @@ def run_problem(problem, refine_method='hp', refine_iteration_limit=0, run_drive
     if solution_record_file not in [rec._filepath for rec in iter(problem._rec_mgr)]:
         recorder = om.SqliteRecorder(solution_record_file)
         problem.add_recorder(recorder)
-        # record_inputs is needed to capture potential input parameters that aren't connected
-        problem.recording_options['record_inputs'] = True
+
         # record_outputs is need to capture the timeseries outputs
         problem.recording_options['record_outputs'] = True