OpenMDAO - CO (Collaborative Optimization) on Sellar test case

Question

An almost similar question was asked but subproblem class was implemented in OpenMDAO to tackle this issue but does not seem to work in my case

I am trying to solve

Answer 1

I've implemented a somewhat reasonable solution to this problem in OpenMDAO 2.0. It was a bit tricky to get to work correctly, with the most notable issue being that I couldn't use ScipyOptimizer on both the top level and sub-optimizations because it appears that its not re-entrant.

The other trick, which was in the above question is that you have to create components that have sub-problems in them. This means you have openmdao running within openmdao. Its not the most efficient setup in the world, and there are numerical challenges because you end up finite-differencing around the optimizations. In theory it would be possible to implement post-optimality sensitivities to get more efficient derivatives here.

NOTE: As expected for CO, the convergence properties are terrible. Solving the Sellar problem this way is horribly inefficient. But it shows the rough approach to setting up MDAO architectures in OpenMDAO 2.0

import numpy as np

from openmdao.api import ExplicitComponent, ImplicitComponent, Group, IndepVarComp, ExecComp


class SellarDis1(ExplicitComponent):


    def setup(self):

        # Global Design Variable
        self.add_input('z', val=np.zeros(2))

        # Local Design Variable
        self.add_input('x', val=0.)

        # Coupling parameter
        self.add_input('y2', val=1.0)

        # Coupling output
        self.add_output('y1', val=1.0)

        # Finite difference all partials.
        self.declare_partials('*', '*')

    def compute(self, inputs, outputs):

        z1 = inputs['z'][0]
        z2 = inputs['z'][1]
        x1 = inputs['x']
        y2 = inputs['y2']

        outputs['y1'] = z1**2 + z2 + x1 - 0.2*y2

    def compute_partials(self, inputs, partials):
        """
        Jacobian for Sellar discipline 1.
        """
        partials['y1', 'y2'] = -0.2
        partials['y1', 'z'] = np.array([[2.0 * inputs['z'][0], 1.0]])
        partials['y1', 'x'] = 1.0


class SellarDis2(ExplicitComponent):


    def setup(self):
        # Global Design Variable
        self.add_input('z', val=np.zeros(2))

        # Coupling parameter
        self.add_input('y1', val=1.0)

        # Coupling output
        self.add_output('y2', val=1.0)

        # Finite difference all partials.
        self.declare_partials('*', '*')

    def compute(self, inputs, outputs):


        z1 = inputs['z'][0]
        z2 = inputs['z'][1]
        y1 = inputs['y1']

        # Note: this may cause some issues. However, y1 is constrained to be
        # above 3.16, so lets just let it converge, and the optimizer will
        # throw it out
        if y1.real < 0.0:
            y1 *= -1

        outputs['y2'] = y1**.5 + z1 + z2

    def compute_partials(self, inputs, J):
        """
        Jacobian for Sellar discipline 2.
        """
        y1 = inputs['y1']
        if y1.real < 0.0:
            y1 *= -1

        J['y2', 'y1'] = .5*y1**-.5
        J['y2', 'z'] = np.array([[1.0, 1.0]])


class SubOpt1(ExplicitComponent):
    ''' minimize differences between target and local variables of the first discipline of the sellar problem '''

    def setup(self):
        self.add_input('z', val=np.array([5.0, 2.0]))
        self.add_input('x_hat', val=1.)
        self.add_input('y1_hat', val=1)
        self.add_input('y2_hat', val=1)

        self.add_output('y1', val=1.0)
        self.add_output('z_hat', val=np.ones(2))
        self.add_output('x', val=1.0)

        # using FD to get derivatives across the sub-optimization
        # note: the sub-optimization itself is using analytic derivatives
        self.declare_partials('y1', ['z', 'x_hat', 'y1_hat', 'y2_hat'], method='fd', step=1e-4, step_calc='abs')
        self.declare_partials('z_hat', ['z', 'x_hat', 'y1_hat', 'y2_hat'], method='fd', step=1e-4, step_calc='abs')
        self.declare_partials('x', ['z', 'x_hat', 'y1_hat', 'y2_hat'], method='fd', step=1e-4, step_calc='abs')


        self.prob = p = Problem()

        # have to define these copies so that OpenMDAO can compute derivs wrt these variables
        params = p.model.add_subsystem('params', IndepVarComp(), promotes=['*'])
        params.add_output('z', val=np.ones(2))
        params.add_output('x_hat', val=1.)
        params.add_output('y1_hat', val=1.)
        params.add_output('y2_hat', val=1.)

        des_vars = p.model.add_subsystem('des_vars', IndepVarComp(), promotes=['*'])
        des_vars.add_output('z_hat', val=np.array([5.0, 2.0]))
        des_vars.add_output('x', val=1.)

        p.model.add_subsystem('d1', SellarDis1())

        # using (global-local)**2 ordering
        p.model.add_subsystem('J', ExecComp('f = sum((z-z_hat)**2) + (x_hat-x)**2 +(y1_hat-y1)**2', z=np.zeros(2), z_hat=np.zeros(2)))
        p.model.add_subsystem('con', ExecComp('c = 3.16 - y1'))

        # data connections in the !!!sub-problem!!!
        p.model.connect('z', 'J.z')
        p.model.connect('x_hat', 'J.x_hat')
        p.model.connect('y2_hat', 'd1.y2')
        p.model.connect('y1_hat', 'J.y1_hat')

        p.model.connect('d1.y1', ['J.y1','con.y1'])
        p.model.connect('z_hat', ['J.z_hat', 'd1.z'])
        p.model.connect('x', ['J.x','d1.x'])

        p.driver = ScipyOptimizer()
        p.driver.options['optimizer'] = 'SLSQP'
        p.driver.options['maxiter'] = 100
        p.driver.options['tol'] = 1e-8

        p.model.add_design_var('x', lower=0, upper=10)
        p.model.add_design_var('z_hat', lower=-10.0, upper=10)
        p.model.add_objective('J.f')
        p.model.add_constraint('con.c', upper=0)

        p.setup()
        p.final_setup()


    def compute(self, inputs, outputs): 

        p = self.prob
        # push any global inputs down, using full absolute path names
        p['y2_hat'] = inputs['y2_hat'] 
        p['z'] = inputs['z'] 
        p['x_hat'] = inputs['x_hat'] 
        p['y1_hat'] = inputs['y1_hat'] 

        #run the optimization 
        print('subopt 1 solve')
        # print('    ', inputs['z'], inputs['x_hat'], inputs['y1_hat'], inputs['y2_hat'], outputs['y1'], outputs['z_hat'])
        p.run_driver()

        # pull the values back up into the output array
        outputs['y1'] = p['d1.y1']
        outputs['z_hat'] = p['z_hat']
        outputs['x'] = p['x']



class SubOpt2(ExplicitComponent):
    ''' minimize differences between target and local variables of the second discipline of the sellar problem '''

    def setup(self):
        self.add_input('z', val=np.array([5.0, 2.0]))
        self.add_input('y1_hat', val=1)
        self.add_input('y2_hat', val=1)

        self.add_output('y2', val=1.0)
        self.add_output('z_hat', val=np.ones(2))


        # using FD to get derivatives across the sub-optimization
        # note: the sub-optimization itself is using analytic derivatives
        self.declare_partials('y2', ['z', 'y1_hat', 'y2_hat'], method='fd', step=1e-4, step_calc='abs')
        self.declare_partials('z_hat', ['z', 'y1_hat', 'y2_hat'], method='fd', step=1e-4, step_calc='abs')

        self.prob = p = Problem()

        # have to define these copies so that OpenMDAO can compute derivs wrt these variables
        params = p.model.add_subsystem('params', IndepVarComp(), promotes=['*'])
        params.add_output('z', val=np.ones(2))
        params.add_output('y1_hat', val=1.)
        params.add_output('y2_hat', val=1.)

        des_vars = p.model.add_subsystem('des_vars', IndepVarComp(), promotes=['*'])
        des_vars.add_output('z_hat', val=np.array([5.0, 2.0]))

        p.model.add_subsystem('d2', SellarDis2())

        # using (global-local)**2 ordering
        p.model.add_subsystem('J', ExecComp('f = sum((z-z_hat)**2) + (y2_hat-y2)**2', z=np.zeros(2), z_hat=np.zeros(2)))
        p.model.add_subsystem('con', ExecComp('c = y2 - 24.0'))

        # data connections in the !!!sub-problem!!!
        p.model.connect('y1_hat', 'd2.y1')
        p.model.connect('z', 'J.z')
        p.model.connect('y2_hat', 'J.y2_hat')

        p.model.connect('d2.y2', ['J.y2','con.y2'])
        p.model.connect('z_hat', ['J.z_hat', 'd2.z'])

        p.driver = ScipyOptimizer()
        p.driver.options['optimizer'] = 'SLSQP'
        p.driver.options['maxiter'] = 100
        p.driver.options['tol'] = 1e-8

        p.model.add_design_var('z_hat', lower=-10.0, upper=10)
        p.model.add_objective('J.f')
        p.model.add_constraint('con.c', upper=0)

        p.setup()
        p.final_setup()


    def compute(self, inputs, outputs): 

        p = self.prob
        # push any global inputs down, using full absolute path names
        p['y1_hat'] = inputs['y1_hat'] 
        p['z'] = inputs['z'] 
        p['y2_hat'] = inputs['y2_hat'] 

        #run the optimization 
        print('subopt 2 solve')
        p.run_driver()

        # pull the values back up into the output array
        outputs['y2'] = p['d2.y2']
        outputs['z_hat'] = p['z_hat']
        # print('    ', inputs['z'], inputs['y1_hat'], inputs['y2_hat'], outputs['y2'], outputs['z_hat'])



class SellarCO(Group):
    ''' optimize top objective function of the sellar problem with the target variables '''

    def setup(self): 


        des_vars = self.add_subsystem('des_vars', IndepVarComp(), promotes=['*'])

        des_vars.add_output('z', val=np.array([5.0, 2.0]))
        des_vars.add_output('x_hat', val=1)

        des_vars.add_output('y1_hat', val=1)
        des_vars.add_output('y2_hat', val=2.5)

        self.add_subsystem('subopt_1', SubOpt1())
        self.add_subsystem('subopt_2', SubOpt2())

        self.add_subsystem('J', ExecComp('c = (sum((z-z1_hat)**2) + sum((z-z2_hat)**2) + (x_hat-x) + (y1_hat-y1)**2 + (y2_hat-y2)**2)**.5', 
                                         z=np.zeros(2), z1_hat=np.zeros(2), z2_hat=np.zeros(2)))
        self.add_subsystem('obj', ExecComp('f = x_hat**2 + z[1] + y1_hat + exp(-y2_hat)', z=np.zeros(2)))

        self.connect('z', ['subopt_1.z', 'subopt_2.z', 'obj.z', 'J.z'])
        self.connect('x_hat', ['obj.x_hat', 'J.x_hat', 'subopt_1.x_hat'])
        self.connect('y1_hat', ['subopt_1.y1_hat', 'subopt_2.y1_hat', 'J.y1_hat', 'obj.y1_hat'])
        self.connect('y2_hat', ['subopt_1.y2_hat', 'subopt_2.y2_hat', 'J.y2_hat', 'obj.y2_hat'])

        self.connect('subopt_1.z_hat', 'J.z1_hat')
        self.connect('subopt_1.y1', 'J.y1')
        self.connect('subopt_1.x', 'J.x')
        self.connect('subopt_2.z_hat', 'J.z2_hat')
        self.connect('subopt_2.y2', 'J.y2')




if __name__ == '__main__':

    from openmdao.api import Problem, ScipyOptimizer, pyOptSparseDriver

    prob = Problem()
    prob.model = SellarCO()

    prob.driver = pyOptSparseDriver() 
    prob.driver.options['optimizer'] = 'SNOPT'
    prob.driver.opt_settings['Major optimality tolerance'] = 1e-1
    prob.driver.opt_settings['Major feasibility tolerance'] = 1e-3

    prob.model.add_design_var('z', lower=np.array([-10.0, 0.0]),upper=np.array([10.0, 10.0]))
    prob.model.add_design_var('x_hat', lower=0.0, upper=10.0)
    prob.model.add_design_var('y1_hat', lower=-10.0, upper=10.0)
    prob.model.add_design_var('y2_hat', lower=-10.0, upper=10.0)

    prob.model.add_objective('obj.f')
    prob.model.add_constraint('J.c', upper=0.005)

    prob.setup()

    prob.run_driver()

    print("\n")
    print( "Minimum target found at (%f, %f, %f)" % (prob['z'][0], prob['z'][1], prob['x_hat']))

    print("Coupling vars target: %f, %f" % (prob['y1_hat'], prob['y2_hat']))
    print("Minimum objective: ", prob['obj.f'])
    # print("constraints: ", prob['con1'] , prob['con2'])