auto_integrate.py

import openmdao.api as om 
import numpy as np 

# OpenConcept PhaseGroup will be used to hold analysis phases with time integration
# =============== These will go in OpenConcept Core ===============#
def find_duration_variables_in_phase(system, abs_namespace, listofvars):
    durationvar = system._problem_meta['duration_var']

    # check if we are a group or not
    if isinstance(system, om.Group):
        # TODO maybe needs to be allprocs
        for subsys in system._subsystems_myproc:
            if not abs_namespace:
                next_namespace = subsys.name
            else:
                next_namespace = abs_namespace + '.' + subsys.name
            find_duration_variables_in_phase(subsys, next_namespace, listofvars)
    else:
        # if the duration variable shows up we need to add its absolute path to listofvars
        for varname in system._var_rel_names['input']:
            if varname == durationvar:
                listofvars.append(abs_namespace + '.' + varname)

class PhaseGroup(om.Group):
    def __init__(self, **kwargs):
        super(PhaseGroup, self).__init__(**kwargs)
        self._duration_var_name = 'duration'

    def _setup_procs(self, pathname, comm, mode, prob_meta):
        # need to pass down the name of the duration variable via prob_meta
        prob_meta.update({'duration_var': self._duration_var_name})
        super(PhaseGroup, self)._setup_procs(pathname, comm, mode, prob_meta)
    
    def _configure(self):
        super(PhaseGroup, self)._configure()
        # check child subsys for variables to be integrated and add them all
        listofvars = []
        find_duration_variables_in_phase(self, '', listofvars)
        # make connections from duration to integrated vars automatically
        for var_abs_address in listofvars:
            self.connect(self._duration_var_name, var_abs_address)


# --------------------- This is just an example
class RateSource(om.ExplicitComponent):
    def setup(self):
        self.add_input('mass', val=2.0, units='kg')
        self.add_input('force', val=1.0, units='N')
        self.add_input('blah', units='m/s')
        # mark the variable for integration
        self.add_output('accel', val=0.5, units='m/s**2', tags=['integrate', 'intunit:m/s','intname:velocity'])
        self.declare_partials(['*'], ['*'], method='cs')

    def compute(self, inputs, outputs):
        outputs['accel'] = inputs['force'] / inputs['mass']

class Integrator(om.ExplicitComponent):
    # need to update my "real integrator" to support this type of method
    def add_integrand(self, name, units=None):
        self.add_input(name+'_dqdt', units=units)
        self.add_output(name+'_q', units=units)

    def setup(self):
        self.add_input('duration', units='s')
        self.declare_partials(['*'],['*'], method='cs')

    def compute(self, inputs, outputs):
        pass

class DeepComp(om.Group):
    def setup(self):
        self.integrate = True
        self.add_subsystem('rs', RateSource())
        # eventually will add the integrator automatically during setup_procs if self.integrate = True
        # or just use a dedicated subclass any time integration is neeeded
        self.add_subsystem('integrator', Integrator())
        self.integrator.add_integrand('blah')
        self.connect('rs.accel', 'integrator.blah_dqdt')
    
    # def configure(self):
        # add the integrand during configure or maybe _configure()

class MyPhase(PhaseGroup):
    def setup(self):
        self.add_subsystem('ivc', om.IndepVarComp('duration', val=1.0, units='s'), promotes_outputs=['duration'])
        self.add_subsystem('dc', DeepComp())

class TopGroup(om.Group):
    def setup(self):
        self.add_subsystem('phase1', MyPhase()) 
        self.add_subsystem('phase2', MyPhase())
        # todo eventually a method to link integrated variables together automatically from phase to phase. I have a list of them all in prior step

if __name__ == "__main__":
    prob = om.Problem(TopGroup())
    prob.setup()
    prob.run_model()
    prob.model.list_outputs()