[SOLVED] Dymos: How to parallelize compute across collocation nodes when using a black-box function

Dymos: How to parallelize compute across collocation nodes when using a black-box function

I'm trying to use a black-box solver to compute one of the ODE functions in Dymos. For example, the component will compute x_dot = f(x) where f is an existing physics-based solver and x is a vector of state variables of size N.

When using collocation, x now takes on a shape of MxN where M is the number of collocation nodes. The issue is that the physics-based solver cannot take the input in this shape. Currently, I'm having to iterate over M times to compute f(x) over each collocation point in sequence. For example

def compute(self, inputs, outputs):

     M = self.options['num_nodes']

     x = inputs["x"]

     x_dot = np.zeros((M, N))

     for i in range(0, M):
         x_dot[i] = f(x[i])

     outputs["x_dot"] = x_dot

This is obviously computationally inefficient and I'm looking for a way to parallelize this compute method.

Is there any native method that Dymos support? FYI, I tried using multiprocessing (Pool) but got an error TypeError: cannot pickle 'weakref' object.

Thanks.

Solution

OpenMDAO, which dymos is built upon, does parallelization through MPI, through two primary means:

First, you can put multiple components into a ParallelGroup. Each of these components will be executed simultaneously when the code is run via mpirun.

You can then Mux the outputs back togther into a single array.

The second, more advanced, way, is to use distributed inputs and outputs within a single component. When coding it up, it appears to be a single component, but MPI gives copies of that component on different processors different elements in the input/output arrays to compute. There is some explanation of how this works here