Accelerating Derivative Computations using the Functional Form in OpenMDAO
PhD research, 2023 - 2024
Overview
We proposed new hierarchical linear solution strategies using the functional form representation for coupled multidisciplinary models. This approach accelerates analytic derivative computations (hence gradient-based optimization) for some problems, including trajectory optimization and control co-design problems.
Background: OpenMDAO and MAUD
OpenMDAO is an open-source optimization framework developed by NASA, and it has been widely used to enable state-of-the-art gradient-based optimizations. OpenMDAO uses the modular analysis and unified derivatives (MAUD) architecture, which computes the analytic derivatives by solving a linear system called the unified derivatives equation (UDE):
\[\frac{\partial r}{\partial u} \frac{\mathrm d u}{\mathrm d r} = I = \frac{\partial r}{\partial u}^T \frac{\mathrm d u}{\mathrm d r}^T\]The UDE is a generalization of the implicit analytic methods (direct and adjoint). Solving the UDE linear system is typically the most computationally expensive part of the derivative computations in OpenMDAO. Therefore, an efficient linear solution approach, which should be tailored for the coupling structure, is the key for computational efficiency.
New hierarchical linear solver strategies
In this project, we proposed new linear solver strategies using the functional form representation of the coupled model. In short, this approach solves multiple layers of the UDEs, instead of just one UDE which is the default approach in MAUD and OpenMDAO.
Here is an example of a two-layer case, where we first solve the sub-level UDE for component-level total derivatives, then the top-level UDE for the system-level derivatives to be provided to gradient-based optimizer.
This multi-layer UDE enables new linear solver strategies, which are not possible in the single-layer UDE:
- Bidirectional hierarchical solver, which switches forward (direct) and reverse (adjoint) modes depending on the component.
- System-level direct solver (LU factorization) + component-level iterative solver.
These new approaches achieve significant speedup for certain problems (but not all problems; after all, it is coupling structure-dependent).
Functional form implementation in OpenMDAO
We need a little “hacky” implementation to do the functional form UDE in OpenMDAO. In OpenMDAO, a Problem instance hosts one UDE; therefore, we need nested Problems to enable the multi-layer UDEs. This approach is also called subproblem in the OpenMDAO context.
Results
We demonstrated the derivative accelerations by the proposed approach using a test problem that combines wing aerostructural analysis and climb trajectory.
When solving this problem using the SAND architecture (collocation method), the bidirectional hierarchical solver is the best approach. As a result, the proposed approach achieved 1.7 times speedup in derivative computations compared to the “conventional” approach using the single-layer UDE.
When solving the same problem in the MDF architecture (parallel-in-time shooting), this changes the coupling structure, hence the appropriate linear solver strategy. For the MDF problem, the hierarhical linear solver with the system-level direct solver and component-level parallelized Krylov solver is the most efficient. It achieved up to 10.2 times speedup compare to the conventional approach!
When should consider using the new strategies?
Of course, the proposed functional form UDE approach is not always the best approach; the best approach is always problem structure-dependent. The following decision tree serves as a guideline on when we should consider using the functional form approach.
Relevant publications
Shugo Kaneko, Eytan Adler, Justin Gray, John Jasa, and Joaquim R. R. A. Martins, Accelerating Derivative Computations Using the Functional Form for Coupled Systems, Optimization and Engineering (under review).