Unify subsystem addition and solver control across all of our ODEs.

aviary/mission/base_ode.py

@@ -1,6 +1,8 @@
 import openmdao.api as om
 
+from aviary.subsystems.aerodynamics.aerodynamics_builder import AerodynamicsBuilder
 from aviary.subsystems.atmosphere.atmosphere import Atmosphere
+from aviary.subsystems.propulsion.propulsion_builder import PropulsionBuilder
 from aviary.utils.aviary_values import AviaryValues
 from aviary.variable_info.variable_meta_data import CoreMetaData
 
@@ -52,31 +54,36 @@ def add_atmosphere(self, **kwargs):
             promotes=['*'],
         )
 
-    def add_subsystems(self, solver_group=None):
+    def add_subsystems_and_solver(
+        self, solver_sub=None, couple_propulsion=False, couple_aero=False, aero_solver_sub=False
+    ):
         """
-        Adds all specified subsystems to ODE in their own group.
+        Adds all specified subsystems to this ODE. Subsystems that need a solver due to coupling
+        are instead added to a group called "solver_sub".
 
         Parameters
         ----------
-        solver_group : om.Group
-            If not None, subsystems that require a solver (subsystem.needs_mission_solver() == True)
-            are placed inside solver_group.
-
-            If None, all subsystems are added to BaseODE regardless of if they request a solver.
-            TODO add solver compatibility to all ODEs
-
+        solver_sub: None or om.Group
+            Pre-created group to add the solver.
+        couple_propulsion : bool
+            When True, the ODE couples with any propulsion subsystems via a throttle to commanded
+            thrust balance.
+        couple_aero : bool
+            When True, the ODE couples with any aerodynamics subsystems via a force balance.
+        aero_solver_sub : None or om.Group
+            Some ODEs (like solved 2DOF) place the aerodynamics and propulsion cycles in separate
+            groups. When this is specified, the aerodynamics subsystem is placed in this sub.
         Returns
         -------
-        use_mission_solver : bool
-            Flag that communicates that one or more subsystem requests to be placed inside a solver
-            (independent of the needs of an individual ODE's setup)
+        om.Group
+            Target group for the ODE. This will be self unless a solver is needed, in which case it
+            will be solver_sub.
         """
         nn = self.options['num_nodes']
         aviary_options = self.options['aviary_options']
         all_subsystems = self.options['subsystems']
         all_subsystem_options = self.options['subsystem_options']
         user_options = self.options['user_options']
-        use_mission_solver = False
 
         for subsystem in all_subsystems:
             # check if subsystem_options has entry for a subsystem of this name
@@ -100,9 +107,36 @@ def add_subsystems(self, solver_group=None):
                     subsystem_options=subsystem_options,
                 )
 
-                if needs_solver and solver_group is not None:
-                    target = solver_group
-                    use_mission_solver = True
+                # ODE couples with propulsion.
+                if couple_propulsion and isinstance(subsystem, PropulsionBuilder):
+                    needs_solver = True
+                elif couple_aero and isinstance(subsystem, AerodynamicsBuilder):
+                    needs_solver = True
+
+                if needs_solver:
+                    if solver_sub is None:
+                        solver_sub = self.add_subsystem('solver_sub', om.Group(), promotes=['*'])
+                        solver_sub.options['auto_order'] = True
+
+                        solver_sub.nonlinear_solver = om.NewtonSolver(
+                            solve_subsystems=True,
+                            atol=1.0e-10,
+                            rtol=1.0e-10,
+                            err_on_non_converge=True,
+                            iprint=2,
+                        )
+                        solver_sub.nonlinear_solver.linesearch = om.BoundsEnforceLS()
+
+                        solver_sub.linear_solver = om.DirectSolver(assemble_jac=True)
+
+                    if (
+                        aero_solver_sub
+                        and couple_aero
+                        and isinstance(subsystem, AerodynamicsBuilder)
+                    ):
+                        target = aero_solver_sub
+                    else:
+                        target = solver_sub
 
                 mission_in = subsystem.mission_inputs(
                     aviary_inputs=aviary_options,
@@ -121,4 +155,4 @@ def add_subsystems(self, solver_group=None):
                     promotes_outputs=mission_out,
                 )
 
-        return use_mission_solver
+        return solver_sub if solver_sub else self

aviary/mission/energy_state/ode/energy_state_ODE.py

@@ -58,12 +58,11 @@ def setup(self):
 
         throttle_enforcement = options['throttle_enforcement']
 
-        sub1 = self.add_subsystem('solver_sub', om.Group(), promotes=['*'])
-        sub1.options['auto_order'] = True
-
-        use_mission_solver = self.add_subsystems(solver_group=sub1)
+        ode_sub = self.add_subsystems_and_solver(
+            couple_propulsion=throttle_enforcement != 'control'
+        )
 
-        sub1.add_subsystem(
+        ode_sub.add_subsystem(
             name='mission_EOM',
             subsys=MissionEOM(num_nodes=nn),
             promotes_inputs=[
@@ -89,7 +88,7 @@ def setup(self):
         if num_engine_type > 1:
             # Multi Engine
 
-            sub1.add_subsystem(
+            ode_sub.add_subsystem(
                 name='throttle_balance',
                 subsys=om.BalanceComp(
                     name='aggregate_throttle',
@@ -105,7 +104,7 @@ def setup(self):
                 promotes_outputs=['*'],
             )
 
-            sub1.add_subsystem(
+            ode_sub.add_subsystem(
                 'throttle_allocator',
                 ThrottleAllocator(
                     num_nodes=nn, throttle_allocation=self.options['throttle_allocation']
@@ -136,7 +135,7 @@ def setup(self):
                 self.add_constraint('thrust_residual', ref=thrust_res_ref, equals=0.0)
             else:
                 # Add a balance comp to compute throttle based on the required thrust.
-                sub1.add_subsystem(
+                ode_sub.add_subsystem(
                     name='throttle_balance',
                     subsys=om.BalanceComp(
                         name=Dynamic.Vehicle.Propulsion.THROTTLE,
@@ -163,16 +162,3 @@ def setup(self):
         self.set_input_defaults(Dynamic.Mission.VELOCITY, val=np.ones(nn), units='m/s')
         self.set_input_defaults(Dynamic.Mission.ALTITUDE, val=np.ones(nn), units='m')
         self.set_input_defaults(Dynamic.Mission.ALTITUDE_RATE, val=np.ones(nn), units='m/s')
-
-        if use_mission_solver or throttle_enforcement != 'control':
-            sub1.nonlinear_solver = om.NewtonSolver(
-                solve_subsystems=True,
-                atol=1.0e-10,
-                rtol=1.0e-10,
-            )
-            print_level = 2
-
-            sub1.nonlinear_solver.linesearch = om.BoundsEnforceLS()
-            sub1.linear_solver = om.DirectSolver(assemble_jac=True)
-            sub1.nonlinear_solver.options['err_on_non_converge'] = True
-            sub1.nonlinear_solver.options['iprint'] = print_level

aviary/mission/energy_state/ode/landing_ode.py

@@ -51,7 +51,7 @@ def setup(self):
             promotes_outputs=[('stall_speed', 'v_stall')],
         )
 
-        self.add_subsystems()
+        self.add_subsystems_and_solver()
 
         self.add_subsystem(
             'landing_eom',

aviary/mission/energy_state/ode/takeoff_ode.py

@@ -55,7 +55,7 @@ def setup(self):
             promotes_outputs=[('stall_speed', 'v_stall')],
         )
 
-        self.add_subsystems()
+        self.add_subsystems_and_solver()
 
         kwargs = {
             'num_nodes': nn,

aviary/mission/solved_two_dof/ode/groundroll_ode.py

@@ -23,10 +23,6 @@ def initialize(self):
 
     def setup(self):
         nn = self.options['num_nodes']
-        aviary_options = self.options['aviary_options']
-        subsystems = self.options['subsystems']
-        subsystem_options = self.options['subsystem_options']
-        user_options = self.options['user_options']
 
         self.add_atmosphere()
 
@@ -44,43 +40,7 @@ def setup(self):
             ],
         )
 
-        kwargs = {
-            'method': 'low_speed',
-        }
-        for subsystem in subsystems:
-            # check if subsystem_options has entry for a subsystem of this name
-            if subsystem.name in subsystem_options:
-                kwargs.update(subsystem_options[subsystem.name])
-            system = subsystem.build_mission(
-                num_nodes=nn,
-                aviary_inputs=aviary_options,
-                user_options=user_options,
-                subsystem_options=kwargs,
-            )
-            if system is not None:
-                mission_in = subsystem.mission_inputs(
-                    aviary_inputs=aviary_options,
-                    user_options=user_options,
-                    subsystem_options=kwargs,
-                )
-                mission_out = subsystem.mission_outputs(
-                    aviary_inputs=aviary_options,
-                    user_options=user_options,
-                    subsystem_options=kwargs,
-                )
-                self.add_subsystem(
-                    subsystem.name,
-                    system,
-                    promotes_inputs=mission_in,
-                    promotes_outputs=mission_out,
-                )
-
-            if isinstance(subsystem, AerodynamicsBuilder):
-                self.promotes(
-                    subsystem.name,
-                    inputs=[Dynamic.Vehicle.ANGLE_OF_ATTACK],
-                    src_indices=np.zeros(nn, dtype=int),
-                )
+        self.add_subsystems_and_solver()
 
         self.add_subsystem('groundroll_eom', GroundrollEOM(num_nodes=nn), promotes=['*'])
 

aviary/mission/solved_two_dof/ode/test/test_groundroll_ode.py

@@ -28,10 +28,13 @@ def setUp(self):
             'GASP', [build_engine_deck(aviary_options)]
         )
 
+        subsystem_options = {'aerodynamics': {'method': 'low_speed'}}
+
         self.prob.model = GroundrollODE(
             num_nodes=2,
             aviary_options=get_option_defaults(),
             subsystems=default_mission_subsystems,
+            subsystem_options=subsystem_options,
         )
 
         setup_model_options(self.prob, aviary_options)

aviary/mission/solved_two_dof/ode/unsteady_solved_ode.py

@@ -77,10 +77,6 @@ def setup(self):
         nn = self.options['num_nodes']
         ground_roll = self.options['ground_roll']
         input_speed_type = self.options['input_speed_type']
-        aviary_options = self.options['aviary_options']
-        subsystem_options = self.options['subsystem_options']
-        user_options = self.options['user_options']
-        subsystems = self.options['subsystems']
         throttle_enforcement = self.options['throttle_enforcement']
 
         self.add_subsystem(
@@ -154,60 +150,12 @@ def setup(self):
         throttle_balance_group.linear_solver = om.DirectSolver(assemble_jac=True)
         throttle_balance_group.nonlinear_solver.options['err_on_non_converge'] = True
 
-        kwargs = {
-            'method': 'low_speed',
-        }
-        if self.options['clean']:
-            kwargs['method'] = 'cruise'
-        for subsystem in subsystems:
-            # check if subsystem_options has entry for a subsystem of this name
-            if subsystem.name in subsystem_options:
-                kwargs.update(subsystem_options[subsystem.name])
-            system = subsystem.build_mission(
-                num_nodes=nn,
-                aviary_inputs=aviary_options,
-                user_options=user_options,
-                subsystem_options=kwargs,
-            )
-            if system is not None:
-                mission_in = subsystem.mission_inputs(
-                    aviary_inputs=aviary_options,
-                    user_options=user_options,
-                    subsystem_options=kwargs,
-                )
-                mission_out = subsystem.mission_outputs(
-                    aviary_inputs=aviary_options,
-                    user_options=user_options,
-                    subsystem_options=kwargs,
-                )
-                if isinstance(subsystem, AerodynamicsBuilder):
-                    mission_inputs = mission_in.copy()
-                    if (
-                        subsystem.code_origin is LegacyCode.FLOPS
-                        and 'angle_of_attack' in mission_inputs
-                    ):
-                        mission_inputs.remove('angle_of_attack')
-                        mission_inputs.append(('angle_of_attack', Dynamic.Vehicle.ANGLE_OF_ATTACK))
-                    control_iter_group.add_subsystem(
-                        subsystem.name,
-                        system,
-                        promotes_inputs=mission_inputs,
-                        promotes_outputs=mission_out,
-                    )
-                elif isinstance(subsystem, PropulsionBuilder):
-                    throttle_balance_group.add_subsystem(
-                        subsystem.name,
-                        system,
-                        promotes_inputs=mission_in,
-                        promotes_outputs=mission_out,
-                    )
-                else:
-                    self.add_subsystem(
-                        subsystem.name,
-                        system,
-                        promotes_inputs=mission_in,
-                        promotes_outputs=mission_out,
-                    )
+        self.add_subsystems_and_solver(
+            solver_sub=throttle_balance_group,
+            couple_propulsion=True,
+            couple_aero=True,
+            aero_solver_sub=control_iter_group,
+        )
 
         eom_comp = UnsteadySolvedEOM(num_nodes=nn, ground_roll=ground_roll)
 

aviary/mission/test/test_external_subsystems_in_mission.py

@@ -106,10 +106,9 @@ def test_mission_solver_2DOF(self):
 
         prob.run_model()
 
-        # NOTE currently 2DOF ODEs do not use the solver subsystem
         self.assertTrue(
             hasattr(
-                prob.model.traj.phases.cruise.rhs_all,
+                prob.model.traj.phases.cruise.rhs_all.solver_sub,
                 'solve_me',
             )
         )

aviary/mission/two_dof/ode/accel_ode.py

@@ -26,13 +26,7 @@ def setup(self):
             promotes_outputs=['weight'],
         )
 
-        kwargs = {
-            'method': 'cruise',
-            'output_alpha': True,
-        }
-        self.options['subsystem_options'].setdefault('aerodynamics', {}).update(kwargs)
-
-        self.add_subsystems()
+        self.add_subsystems_and_solver()
 
         self.add_subsystem(
             'accel_eom',