diff --git a/.github/workflows/dymos_docs_workflow.yml b/.github/workflows/dymos_docs_workflow.yml index 55e1ce5fc..786ad2ed1 100644 --- a/.github/workflows/dymos_docs_workflow.yml +++ b/.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 diff --git a/.github/workflows/dymos_tests_workflow.yml b/.github/workflows/dymos_tests_workflow.yml index 3dd1c75a3..733db0d57 100644 --- a/.github/workflows/dymos_tests_workflow.yml +++ b/.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: diff --git a/dymos/examples/balanced_field/doc/test_doc_balanced_field_length.py b/dymos/examples/balanced_field/doc/test_doc_balanced_field_length.py index 3b07d2f1c..b86f81781 100644 --- a/dymos/examples/balanced_field/doc/test_doc_balanced_field_length.py +++ b/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 diff --git a/dymos/examples/double_integrator/test/test_double_integrator_breakwell.py b/dymos/examples/double_integrator/test/test_double_integrator_breakwell.py index c46a1b583..0a92957c6 100644 --- a/dymos/examples/double_integrator/test/test_double_integrator_breakwell.py +++ b/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 diff --git a/dymos/examples/finite_burn_orbit_raise/finite_burn_orbit_raise_problem.py b/dymos/examples/finite_burn_orbit_raise/finite_burn_orbit_raise_problem.py index 0541a41f3..c6026d709 100644 --- a/dymos/examples/finite_burn_orbit_raise/finite_burn_orbit_raise_problem.py +++ b/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 diff --git a/dymos/examples/finite_burn_orbit_raise/test/test_multi_phase_restart.py b/dymos/examples/finite_burn_orbit_raise/test/test_multi_phase_restart.py index 7f3b68337..c3d11c8d2 100644 --- a/dymos/examples/finite_burn_orbit_raise/test/test_multi_phase_restart.py +++ b/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' : 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' : 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' : 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' : 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 diff --git a/dymos/examples/hyper_sensitive/test/test_hyper_sensitive.py b/dymos/examples/hyper_sensitive/test/test_hyper_sensitive.py index 8bf49f802..db760dff6 100644 --- a/dymos/examples/hyper_sensitive/test/test_hyper_sensitive.py +++ b/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 diff --git a/dymos/examples/low_thrust_spiral/test/test_low_thrust_spiral.py b/dymos/examples/low_thrust_spiral/test/test_low_thrust_spiral.py index 3b8419dd1..89798d9cf 100644 --- a/dymos/examples/low_thrust_spiral/test/test_low_thrust_spiral.py +++ b/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) diff --git a/dymos/phase/simulation_phase.py b/dymos/phase/simulation_phase.py index ee2c43608..36f33e960 100644 --- a/dymos/phase/simulation_phase.py +++ b/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() diff --git a/dymos/run_problem.py b/dymos/run_problem.py index 625edb901..04e3701c4 100755 --- a/dymos/run_problem.py +++ b/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 diff --git a/dymos/test/test_load_case.py b/dymos/test/test_load_case.py index 0661627ae..f2d84abd3 100644 --- a/dymos/test/test_load_case.py +++ b/dymos/test/test_load_case.py @@ -1,5 +1,7 @@ import os import unittest + +import numpy as np from openmdao.utils.testing_utils import use_tempdirs import openmdao import openmdao.api as om @@ -81,17 +83,16 @@ def test_load_case_unchanged_grid(self): # We unnecessarily call setup again just to make sure we obliterate the previous solution p.setup() + p.set_val('phase0.controls:theta', 0.0) + # Load the values from the previous solution - dm.load_case(p, case) + p.load_case(case) # Run the model to ensure we find the same output values as those that we recorded - p.run_driver() - - inputs = dict([(o[0], o[1]) for o in case.list_inputs(units=True, shape=True, out_stream=None)]) + p.run_model() - assert_near_equal(p['phase0.controls:theta'], - inputs['phase0.control_group.control_interp_comp.controls:theta'] - ['val']) + assert_near_equal(case.get_val('phase0.controls:theta'), + p.get_val('phase0.controls:theta')) def test_load_case_unchanged_grid_polynomial_control(self): import openmdao.api as om @@ -111,16 +112,13 @@ def test_load_case_unchanged_grid_polynomial_control(self): p.setup() # Load the values from the previous solution - dm.load_case(p, case) + p.load_case(case) # Run the model to ensure we find the same output values as those that we recorded - p.run_driver() + p.run_model() - inputs = dict([(o[0], o[1]) for o in case.list_inputs(units=True, shape=True, out_stream=None)]) - - assert_near_equal(p['phase0.polynomial_controls:theta'], - inputs['phase0.polynomial_control_group.interp_comp.polynomial_controls:theta'] - ['val']) + assert_near_equal(p.get_val('phase0.polynomial_controls:theta'), + case.get_val('phase0.polynomial_controls:theta')) def test_load_case_lgl_to_radau(self): import openmdao.api as om @@ -139,7 +137,7 @@ def test_load_case_lgl_to_radau(self): q = setup_problem(dm.Radau(num_segments=20)) # Load the values from the previous solution - dm.load_case(q, case) + q.load_case(case) # Run the model to ensure we find the same output values as those that we recorded q.run_driver() @@ -171,19 +169,19 @@ def test_load_case_radau_to_lgl(self): q = setup_problem(dm.GaussLobatto(num_segments=50)) # Load the values from the previous solution - dm.load_case(q, case) + q.load_case(case) # Run the model to ensure we find the same output values as those that we recorded - q.run_driver() + q.run_model() - outputs = dict([(o[0], o[1]) for o in case.list_outputs(units=True, shape=True, - out_stream=None)]) + time_p = case.get_val('phase0.timeseries.time') + theta_p = case.get_val('phase0.timeseries.theta') - time_val = outputs['phase0.timeseries.timeseries_comp.time']['val'] - theta_val = outputs['phase0.timeseries.timeseries_comp.theta']['val'] + time_q = q.get_val('phase0.timeseries.time') + theta_q = q.get_val('phase0.timeseries.theta') - assert_near_equal(q['phase0.timeseries.timeseries_comp.theta'], - q.model.phase0.interp(xs=time_val, ys=theta_val, nodes='all'), + assert_near_equal(q.model.phase0.interp(xs=time_p, ys=theta_p, nodes='all'), + q.model.phase0.interp(xs=time_q, ys=theta_q, nodes='all'), tolerance=1.0E-2) def test_load_case_warn_fix_final_states(self): @@ -211,7 +209,7 @@ def test_load_case_warn_fix_final_states(self): f" this will overwrite the user-specified value")) with assert_warnings(msgs): - dm.load_case(q, case) + q.load_case(case) def test_load_case_warn_fix_final_control(self): import openmdao.api as om @@ -232,7 +230,7 @@ def test_load_case_warn_fix_final_control(self): f" different final value this will overwrite the user-specified value" with assert_warning(UserWarning, msg): - dm.load_case(q, case) + q.load_case(case) def test_load_case_warn_fix_final_polynomial_control(self): import openmdao.api as om @@ -255,7 +253,7 @@ def test_load_case_warn_fix_final_polynomial_control(self): f" different final value this will overwrite the user-specified value" with assert_warning(UserWarning, msg): - dm.load_case(q, case) + q.load_case(case) if __name__ == '__main__': # pragma: no cover diff --git a/dymos/test/test_run_problem.py b/dymos/test/test_run_problem.py index 5072be9aa..19dc64123 100755 --- a/dymos/test/test_run_problem.py +++ b/dymos/test/test_run_problem.py @@ -831,12 +831,6 @@ def test_simulate_array_param(self): assert_near_equal(sol - sim, np.zeros_like(sol)) - # Test that the parameter is available in the solution and simulation files - sol = sol_results.get_val('traj.phase0.parameters:array') - sim = sim_results.get_val('traj.phase0.parameters:array') - - assert_near_equal(sol - sim, np.zeros_like(sol)) - if __name__ == '__main__': # pragma: no cover unittest.main() diff --git a/dymos/trajectory/trajectory.py b/dymos/trajectory/trajectory.py index a57cf7839..31cce9b7d 100644 --- a/dymos/trajectory/trajectory.py +++ b/dymos/trajectory/trajectory.py @@ -58,6 +58,7 @@ def __init__(self, **kwargs): self._linkages = {} self._phases = {} self._phase_graph = nx.DiGraph() + self._has_connected_phases = False self.parameter_options = {} self.phases = om.ParallelGroup() @@ -68,6 +69,17 @@ def initialize(self): """ self.options.declare('sim_mode', types=bool, default=False, desc='Used internally by Dymos when invoking simulate on a trajectory') + self.options.declare('default_nonlinear_solver', + types=(om.NonlinearBlockJac, om.NewtonSolver, om.BroydenSolver), + default=None, allow_none=True, recordable=False, + desc='A nonlinear solver to be used when Phases are connected but being ' + 'run in parallel. If not specified, Dymos will automatically use ' + 'NonlinearBlockJac in this situation.') + self.options.declare('default_linear_solver', default=None, allow_none=True, recordable=False, + types=(om.PETScKrylov, om.LinearBlockJac, om.NonlinearBlockGS), + desc='A linear solver to be used when Phases are connected but being ' + 'run in parallel. If not specified, Dymos will automatically use ' + 'PETScKrylov in this situation.') def add_phase(self, name, phase, **kwargs): """ @@ -330,10 +342,6 @@ def _setup_parameters(self): 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 ' \ @@ -390,7 +398,7 @@ def setup(self): # This will override the existing phases attribute with the same thing. self.add_subsystem('phases', subsys=self.phases) - if self._linkages: + if self._linkages and not self.options['sim_mode']: self._setup_linkages() def _configure_parameters(self): @@ -759,10 +767,6 @@ def _is_valid_linkage(self, phase_name_a, phase_name_b, loc_a, loc_b, var_a, var def _configure_linkages(self): - if self.options['sim_mode']: - # If this is a simulation trajectory, there's no need to link the phases. - return - connected_linkage_inputs = [] def _print_on_rank(rank=0, *args, **kwargs): @@ -865,6 +869,7 @@ def _get_prefixed_var(var, phase): raise om.OpenMDAOWarning(msg) _print_on_rank(f'{indent * 2}{prefixed_a:<{padding_a}s} [{loc_a}{str_fixed_a}] -> ' f'{prefixed_b:<{padding_b}s} [{loc_b}{str_fixed_b}]') + self._has_connected_phases = True else: if fixed_a and fixed_b: @@ -1002,6 +1007,39 @@ def get_final_tbounds(phase_name, node_data, old_tmin=-INF_BOUND, old_tmax=INF_B err_lines = '\n'.join(errs) if len(errs) > 1 else errs[0] raise RuntimeError(f"{self.msginfo}: {err_lines}") + def _configure_solvers(self): + """ + Automatically configure solvers for the Trajectory. + + If operating under MPI and phases are connected, then + the default nonlinear solver will be a NonlinearBlockJac, + while the default linear solver will be PETScKrylov. + + These solvers can be changed through the + 'default_nonlinear_solver' and 'default_linear_solver' options. + """ + if self._has_connected_phases and MPI and self.comm.size > 1: + + if isinstance(self.phases.nonlinear_solver, om.NonlinearRunOnce): + if self.options['default_nonlinear_solver'] is None: + msg = f'{self.msginfo}: Setting phases.nonlinear_solver to `om.NonlinearBlockJac(iprint=0)`.\n' \ + f'Connected phases in parallel require a non-default nonlinear solver.\n' \ + f'Use {self.pathname}.options[\'default_nonlinear_solver\'] to explicitly set the solver.' + om.issue_warning(msg) + self.phases.nonlinear_solver = om.NonlinearBlockJac(iprint=0) + else: + self.phases.nonlinear_solver = self.options['default_nonlinear_solver'] + + if isinstance(self.phases.linear_solver, om.LinearRunOnce): + if self.options['default_linear_solver'] is None: + msg = f'{self.msginfo}: Setting phases.linear_solver to `om.PETScKrylov()`.\n' \ + f'Connected phases in parallel require a non-default linear solver.\n' \ + f'Use {self.pathname}.options[\'default_linear_solver\'] to explicitly set the solver.' + om.issue_warning(msg) + self.phases.linear_solver = om.PETScKrylov() + else: + self.phases.linear_solver = self.options['default_linear_solver'] + def configure(self): """ Configure the Trajectory Group. @@ -1016,9 +1054,11 @@ def configure(self): if self.parameter_options: self._configure_parameters() - if self._linkages: + if self._linkages and not self.options['sim_mode']: self._configure_linkages() + self._configure_solvers() + self._constraint_report(outstream=sys.stdout) # promote everything else out of phases @@ -1414,9 +1454,7 @@ def simulate(self, times_per_seg=_unspecified, method=_unspecified, atol=_unspec if record_file is not None: rec = om.SqliteRecorder(record_file) sim_prob.add_recorder(rec) - # record_inputs is needed to capture potential input parameters that aren't connected - sim_prob.recording_options['record_inputs'] = True - # record_outputs is need to capture the timeseries outputs + # record_outputs is needed to capture the timeseries outputs sim_prob.recording_options['record_outputs'] = True with warnings.catch_warnings():