Fix for Vanderpol example under updated OpenMDAO MPI operation. (#1023)

* changed VanderPolODE to be matrix free

* fixing mpi tests

* skip vanderpol tests under OpenMDAO older than 3.29.0

---------

Co-authored-by: Bret Naylor <[email protected]>

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')

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

dymos/examples/vanderpol/test/test_vanderpol.py

@@ -2,28 +2,33 @@
 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
 
 import dymos as dm
 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):
         # simulate only: with no control, the system oscillates
         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):

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

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)