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 ad55821bc..8657a0637 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 @@ -4,8 +4,7 @@ import openmdao.api as om 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 openmdao.utils.mpi import MPI, multi_proc_exception_check 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 @@ -25,9 +24,10 @@ def test_ex_two_burn_orbit_raise_connected(self): 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) + with multi_proc_exception_check(p.comm): + 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') @@ -54,9 +54,10 @@ def test_restart_from_solution_radau(self): 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) + with multi_proc_exception_check(p.comm): + 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, @@ -101,9 +102,10 @@ def test_ex_two_burn_orbit_raise_connected(self): 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, - tolerance=4.0E-3) + with multi_proc_exception_check(p.comm): + 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') @@ -117,7 +119,7 @@ def test_ex_two_burn_orbit_raise_connected(self): 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) @@ -130,9 +132,10 @@ def test_restart_from_solution_radau_to_connected(self): 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) + with multi_proc_exception_check(p.comm): + 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') diff --git a/dymos/examples/vanderpol/doc/test_doc_vanderpol.py b/dymos/examples/vanderpol/doc/test_doc_vanderpol.py index 18d8d364e..20abb1907 100644 --- a/dymos/examples/vanderpol/doc/test_doc_vanderpol.py +++ b/dymos/examples/vanderpol/doc/test_doc_vanderpol.py @@ -1,10 +1,14 @@ import os import unittest +import openmdao from openmdao.utils.testing_utils import use_tempdirs from openmdao.utils.mpi import MPI +om_version = tuple([int(s) for s in openmdao.__version__.split('-')[0].split('.')]) + + @use_tempdirs class TestVanderpolForDocs(unittest.TestCase): def tearDown(self): @@ -21,6 +25,7 @@ def test_vanderpol_for_docs_simulation(self): dm.run_problem(p, run_driver=False, simulate=True, make_plots=True) + @unittest.skipIf(om_version < (3, 29, 0), 'Test requires OpenMDAO 3.29.0 or later') def test_vanderpol_for_docs_optimize(self): import dymos as dm from dymos.examples.vanderpol.vanderpol_dymos import vanderpol @@ -31,6 +36,7 @@ def test_vanderpol_for_docs_optimize(self): dm.run_problem(p, simulate=True, make_plots=True) + @unittest.skipIf(om_version < (3, 29, 0), 'Test requires OpenMDAO 3.29.0 or later') def test_vanderpol_for_docs_optimize_refine(self): import dymos as dm from dymos.examples.vanderpol.vanderpol_dymos import vanderpol @@ -56,6 +62,7 @@ def test_vanderpol_for_docs_optimize_refine(self): class TestVanderpolDelayMPI(unittest.TestCase): N_PROCS = 2 + @unittest.skipIf(om_version < (3, 29, 0), 'Test requires OpenMDAO 3.29.0 or later') def test_vanderpol_delay_mpi(self): import openmdao.api as om import dymos as dm diff --git a/dymos/examples/vanderpol/test/test_vanderpol.py b/dymos/examples/vanderpol/test/test_vanderpol.py index 79cd92d0a..351704674 100644 --- a/dymos/examples/vanderpol/test/test_vanderpol.py +++ b/dymos/examples/vanderpol/test/test_vanderpol.py @@ -2,6 +2,7 @@ from numpy.testing import assert_almost_equal import numpy as np +import openmdao from openmdao.utils.testing_utils import use_tempdirs, require_pyoptsparse from openmdao.utils.mpi import MPI @@ -9,6 +10,9 @@ from dymos.examples.vanderpol.vanderpol_dymos import vanderpol +om_version = tuple([int(s) for s in openmdao.__version__.split('-')[0].split('.')]) + + @use_tempdirs class TestVanderpolExample(unittest.TestCase): def test_vanderpol_simulate(self): @@ -16,14 +20,15 @@ def test_vanderpol_simulate(self): p = vanderpol(transcription='gauss-lobatto', num_segments=75) p.run_model() + @unittest.skipIf(om_version < (3, 29, 0), 'Test requires OpenMDAO 3.29.0 or later') @require_pyoptsparse(optimizer='IPOPT') def test_vanderpol_simulate_true(self): - # simulate true p = vanderpol(transcription='radau-ps', num_segments=30, transcription_order=3, compressed=True, optimizer='IPOPT', delay=0.005, distrib=True, use_pyoptsparse=True) dm.run_problem(p, run_driver=True, simulate=True) + @unittest.skipIf(om_version < (3, 29, 0), 'Test requires OpenMDAO 3.29.0 or later') def test_vanderpol_optimal(self): p = vanderpol(transcription='gauss-lobatto', num_segments=75) dm.run_problem(p) # find optimal control solution to stop oscillation @@ -34,6 +39,7 @@ def test_vanderpol_optimal(self): assert_almost_equal(p.get_val('traj.phase0.states:x1')[-1, ...], np.zeros(1)) assert_almost_equal(p.get_val('traj.phase0.controls:u')[-1, ...], np.zeros(1), decimal=3) + @unittest.skipIf(om_version < (3, 29, 0), 'Test requires OpenMDAO 3.29.0 or later') def test_vanderpol_optimal_grid_refinement(self): # enabling grid refinement gives a faster and better solution with fewer segments p = vanderpol(transcription='gauss-lobatto', num_segments=15) @@ -53,6 +59,7 @@ class TestVanderpolExampleMPI(unittest.TestCase): N_PROCS = 4 + @unittest.skipIf(om_version < (3, 29, 0), 'Test requires OpenMDAO 3.29.0 or later') @require_pyoptsparse(optimizer='IPOPT') @unittest.skipUnless(MPI, 'this test requires MPI') def test_vanderpol_optimal_mpi(self): diff --git a/dymos/examples/vanderpol/vanderpol_dymos.py b/dymos/examples/vanderpol/vanderpol_dymos.py index 34a5e2572..26612dbfd 100644 --- a/dymos/examples/vanderpol/vanderpol_dymos.py +++ b/dymos/examples/vanderpol/vanderpol_dymos.py @@ -20,7 +20,7 @@ def vanderpol(transcription='gauss-lobatto', num_segments=40, transcription_orde if optimizer == 'SNOPT': p.driver.opt_settings['iSumm'] = 6 # show detailed SNOPT output elif optimizer == 'IPOPT': - p.driver.opt_settings['print_level'] = 0 + p.driver.opt_settings['print_level'] = 1 p.driver.declare_coloring() # define a Trajectory object and add to model @@ -76,7 +76,7 @@ def vanderpol(transcription='gauss-lobatto', num_segments=40, transcription_orde phase.add_objective('J', loc='final') # setup the problem - p.setup(check=True) + p.setup(check=True, force_alloc_complex=True) p['traj.phase0.t_initial'] = 0.0 p['traj.phase0.t_duration'] = t_final diff --git a/dymos/examples/vanderpol/vanderpol_ode.py b/dymos/examples/vanderpol/vanderpol_ode.py index 791912ac2..4ba52225f 100644 --- a/dymos/examples/vanderpol/vanderpol_ode.py +++ b/dymos/examples/vanderpol/vanderpol_ode.py @@ -2,6 +2,7 @@ import openmdao.api as om import time from openmdao.utils.array_utils import evenly_distrib_idxs +from openmdao.utils.mpi import MPI class VanderpolODE(om.ExplicitComponent): @@ -72,16 +73,66 @@ def compute(self, inputs, outputs): outputs['x1dot'] = x0 outputs['Jdot'] = x0**2 + x1**2 + u**2 - def compute_partials(self, inputs, jacobian): - time.sleep(self.options['delay'] * self.io_size) - - x0 = inputs['x0'][self.start_idx:self.end_idx] - x1 = inputs['x1'][self.start_idx:self.end_idx] - u = inputs['u'][self.start_idx:self.end_idx] - - jacobian['x0dot', 'x0'] = 1.0 - x1 * x1 - jacobian['x0dot', 'x1'] = -2.0 * x1 * x0 - 1.0 - - jacobian['Jdot', 'x0'] = 2.0 * x0 - jacobian['Jdot', 'x1'] = 2.0 * x1 - jacobian['Jdot', 'u'] = 2.0 * u + def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode): + + # it's necessary to make this component matrix free because the inputs are non-distributed + # and the outputs are distributed, and the framework doesn't know how to populate the full + # nondistributed inputs on each rank in reverse mode. + + # FIXME: this delay will be applied on every call to compute_jacvec_product, which may be + # more often than was originally intended before this component was converted to + # matrix free (it was originally done in compute_partials). + time.sleep(self.options['delay']) + + myslice = slice(self.start_idx, self.end_idx) + + locx0 = inputs['x0'][myslice] + locx1 = inputs['x1'][myslice] + locu = inputs['u'][myslice] + + locdx0 = d_inputs['x0'][myslice] + locdx1 = d_inputs['x1'][myslice] + locdu = d_inputs['u'][myslice] + + if mode == 'fwd': + if 'x0dot' in d_outputs: + if 'x0' in d_inputs: + d_outputs['x0dot'] += (1.0 - locx1**2) * locdx0 + if 'x1' in d_inputs: + d_outputs['x0dot'] += (-2.0 * locx1 * locx0 - 1.) * locdx1 + if 'u' in d_inputs: + d_outputs['x0dot'] += locdu + if 'x1dot' in d_outputs: + if 'x0' in d_inputs: + d_outputs['x1dot'] += locdx0 + if 'Jdot' in d_outputs: + if 'x0' in d_inputs: + d_outputs['Jdot'] += 2.0 * locx0 * locdx0 + if 'x1' in d_inputs: + d_outputs['Jdot'] += 2.0 * locx1 * locdx1 + if 'u' in d_inputs: + d_outputs['Jdot'] += 2.0 * locu * locdu + + elif mode == 'rev': + if 'x0dot' in d_outputs: + if 'x0' in d_inputs: + d_inputs['x0'][myslice] += (1.0 - locx1**2) * d_outputs['x0dot'] + if 'x1' in d_inputs: + d_inputs['x1'][myslice] += (-2.0 * locx1 * locx0 - 1.) * d_outputs['x0dot'] + if 'u' in d_inputs: + d_inputs['u'][myslice] += d_outputs['x0dot'] + if 'x1dot' in d_outputs: + if 'x0' in d_inputs: + d_inputs['x0'][myslice] += d_outputs['x1dot'] + if 'Jdot' in d_outputs: + if 'x0' in d_inputs: + d_inputs['x0'][myslice] += 2.0 * locx0 * d_outputs['Jdot'] + if 'x1' in d_inputs: + d_inputs['x1'][myslice] += 2.0 * locx1 * d_outputs['Jdot'] + if 'u' in d_inputs: + d_inputs['u'][myslice] += 2.0 * locu * d_outputs['Jdot'] + + if self.comm.size > 1 and self.options['distrib']: + d_inputs['x0'] = self.comm.allreduce(d_inputs['x0'], op=MPI.SUM) + d_inputs['x1'] = self.comm.allreduce(d_inputs['x1'], op=MPI.SUM) + d_inputs['u'] = self.comm.allreduce(d_inputs['u'], op=MPI.SUM)