Wrapping an external solver as an component in openMDAO and how to specify Python objects from the solver as output in openMDAO?

Question

I am trying to 'disintegrate' an external Python disciplinary analysis BEM solver into openMDAO components to try to implement semi-analytic adjoints. The problem I am running is that the python object from that solver is to be designated as input and outputs in the openMDAO components. I am not sure how we can specify the complex objects as input, and outputs and how to declare partials for them?

Is there a better way to wrap an external Python solver in openMDAO?

Answer 1

Here is a more concrete example that should provide stronger guidance. Its based on a simple toy "solver", but should get the broad concept across. Here, we have a computational object that has some internal state which is maintained and used across a couple of its methods. For most PDE solvers (e.g. FEA or CFD) you will have one method that converges the implicit states, and another that computes some key quantities --- which I call functionals --- by integrating over the state values (e.g. stress, lift, drag).

So that is what the simple object I wrote mimics. It has two main methods:

• solve_for_state
• compute_functional

It also has one more method R which is needed by the solver inside solve_for_states, which I defined as public method to better integrate with OpenMDAO's APIs.

The model has a very simple structure. One main input a, one implicit variable b and one functional variable c. The N2 diagram shows how its all connected. Notice that nowhere do I pass the actual solver object around as a variable between the components. I only every pass around floating point values (or more generally arrays of floating point values). This is critical, since OpenMDAO's derivative system only understands floats and arrays of floats.

So I have multiple components which all need to use the shared object, but which can't pass it around as a variable. Instead, I define the shared object as an option and pass it to all the components during instantiation.

The last notable feature of the model is that I pass a to both the states and functional components via OpenMDAO. This may seem somewhat counter intuitive since the computational object stores a internally itself and hence shouldn't strictly need it when it comes time to call the compute_functional method. While that's true, it's also true that OpenMDAO needs to know that the functional component has a direct dependency on the value of a. So we have to make sure that connection is present in the model itself.

import numpy as np
from scipy.optimize import root

import openmdao.api as om

class SomeSolver:

    def __init__(self):
        self.a=1
        self.b=2
        self.c=3

    # residual function
    def R(self, b):
        return self.a + b**2 - 10*b

    def solve_for_state(self):
        """Converging a linear or nonlinear function that defines
        the state `b` as an implicit function of the input `a`."""

        # initial guess for `b` is 3.0
        sol = root(self.R, 3.0)

        self.b = sol.x

        return self.b

    def compute_functional(self):
        """Explicit function of input `a` and states `b`"""
        self.c = self.a + 2*self.b

        return self.c


class State(om.ImplicitComponent):

    def initialize(self):
        self.options.declare('solver')

    def setup(self):
        self.add_input('a')
        self.add_output('b')

    def apply_nonlinear(self, inputs, outputs, residuals):
        solver.a = inputs['a']
        residuals['b'] = solver._R(b)

    def solve_nonlinear(self, inputs, outputs):
        solver = self.options['solver']

        solver.a = inputs['a']

        solver.solve_for_state()

        outputs['b'] = solver.b


class Functional(om.ExplicitComponent):
    def initialize(self):
        self.options.declare('solver')

    def setup(self):
        self.add_input('a')
        self.add_input('b')
        self.add_output('c')

    def compute(self, inputs, outputs):
        solver = self.options['solver']

        solver.a = inputs['a']
        solver.b = inputs['b']

        solver.compute_functional()

        outputs['c'] = solver.c

if __name__ == "__main__":

    p = om.Problem()

    some_solver=SomeSolver()

    p.model.add_subsystem('state', State(solver=some_solver), promotes=['*'])
    p.model.add_subsystem('Functional', Functional(solver=some_solver), promotes=['*'])

    p.setup()

    p.set_val('a', 2)

    p.run_model()

    p.model.list_outputs()

Answer 2

You question is pretty vauge, but broadly speaking I can give some guidance.

You really don't want to pass around objects as I/O for anything in OpenMDAO that needs to use derivatives. Instead, you can pass the shared object around during setup so ever instance has access to it. But at the borders of the components you need to pass scalars or arrays of data.

Usually in this case, you have one "solver" and one "functional" class. The solver outputs the state array. The functional component takes the state array as inputs.

You might want to look at some of the wrappers in the Mphys library which do this sort of thing a lot.