OpenMDAO 3.x：ParallelGroup近似总导数的梯度计算效率优化问询

问题背景 #

当前模型使用单个向量输入x_vec扇出到ParallelGroup下的多个子组件，在两种梯度计算场景下存在冗余开销：

• FD近似总导数：扰动x_vec的单个索引时，所有子组件都会执行计算
• 矩阵自由前向模式：种子化x_vec的单个索引时，所有子组件都会触发compute_jacvec_product

当将每个子组件的输入单独设为设计变量时，矩阵自由模式能正常工作（仅调用对应组件的compute_jacvec_product），但FD近似总导数仍存在冗余。

问题1：向量输入扇出模型适配矩阵自由前向导数 #

原因 #

直接使用src_indices从向量中提取单个元素连接到子组件时，OpenMDAO的矩阵自由模式默认会将整个向量的种子广播到所有子组件，导致所有组件都执行导数计算。

解决方案：拆分向量输入为独立变量子集 #

使用SubsetComp将原向量x_vec拆分为多个独立的单元素变量，再分别连接到对应子组件。这样矩阵自由模式能精准识别每个子组件的依赖，仅为受影响的组件传递种子。

核心修改代码段：

# 添加SubsetComp拆分向量为单个变量
subset = p.model.add_subsystem('subset', om.SubsetComp())
for i in range(4):
    subset.add_subset(f'x_{i}', 'x_vec', indices=[i])

# 替换原连接方式，使用拆分后的单个变量
for i in range(4):
    par.add_subsystem('case%02d' % i, MultComp(mult=float(i+1)))
    p.model.connect(f'subset.x_{i}', f'par.case%02d.x' % i)
    p.model.connect(f'par.case%02d.y' % i, f'mux_comp.y_vec_%i' % i)

设置完成后，矩阵自由前向模式下，种子只会传递给对应索引的子组件，仅触发该组件的compute_jacvec_product。

问题2：FD近似总导数的相关性缩减 #

原因 #

当前在ParallelGroup上直接调用approx_totals时，OpenMDAO默认会重新运行整个组内的所有组件，即使只有一个组件的输入被扰动。这是因为原连接方式未明确告知OpenMDAO每个子组件仅依赖x_vec的单个元素。

解决方案：启用相关性缩减并明确依赖关系 #

1. 启用相关性缩减：在approx_totals中设置correlation_reduction=True，让OpenMDAO自动识别受扰动变量影响的组件，仅重新计算这些组件。
2. 明确组件依赖：配合SubsetComp的拆分方式，让OpenMDAO清晰识别每个子组件的输入依赖，确保相关性缩减生效。

核心修改代码段：

# 在par组上设置approx_totals时启用相关性缩减
p.model.par.approx_totals(method='fd', step=1e-4, form='central', correlation_reduction=True)

额外优化：并行FD计算 #

如果环境支持并行计算，可开启use_parallel_fd=True，让多个FD扰动的计算并行执行，进一步提升效率：

p.model.par.approx_totals(method='fd', step=1e-4, form='central', 
                          correlation_reduction=True, use_parallel_fd=True)

完整修改后的代码示例 #

import numpy as np
import openmdao.api as om


class MultComp(om.ExplicitComponent):
    def __init__(self, delay=1.0, size=3, mult=2.0):
        super().__init__()
        self.delay = delay
        self.size = size
        self.mult = mult

    def setup(self):
        self.add_input('x', val=self.mult)
        self.add_output('y', val=0.0)
        self.declare_partials(of='*', wrt='*')

    def compute(self, inputs, outputs):
        outputs['y'] = inputs['x'] * self.mult
        print('hello from comp', self.mult, inputs['x'], outputs['y'])

    def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode):
        d_outputs['y'] += d_inputs['x'] * self.mult
        print('hello from jacvec comp', self.mult, inputs['x'])


p = om.Problem(group_by_pre_opt_post=True)

indeps = p.model.add_subsystem('indeps', om.IndepVarComp())
indeps.add_output('x_vec', np.linspace(1., 4., 4))

# 拆分向量为单个变量子集
subset = p.model.add_subsystem('subset', om.SubsetComp())
for i in range(4):
    subset.add_subset(f'x_{i}', 'x_vec', indices=[i])

par = p.model.add_subsystem('par', om.ParallelGroup())

mux = p.model.add_subsystem('mux_comp', om.MuxComp(vec_size=4))
mux.add_var("y_vec", shape=(1,))

for i in range(4):
    par.add_subsystem('case%02d' % i, MultComp(mult=float(i+1)))
    p.model.connect(f'subset.x_{i}', f'par.case%02d.x' % i)
    p.model.connect(f'par.case%02d.y' % i, f'mux_comp.y_vec_%i' % i)

p.model.add_subsystem('sum_comp', om.ExecComp('ysum=sum(y_vec)', y_vec=np.zeros(4)))
p.model.connect('mux_comp.y_vec', 'sum_comp.y_vec')
p.model.add_design_var('indeps.x_vec', upper=20.)
p.model.add_objective('sum_comp.ysum', scaler=-1.)

p.driver = om.ScipyOptimizeDriver()
p.driver.options["optimizer"] = "SLSQP"

# 启用相关性缩减的FD近似总导数
p.model.par.approx_totals(method='fd', step=1e-4, form='central', correlation_reduction=True)

p.setup(mode='fwd')
p.run_model()

totals = p.compute_totals()

内容的提问来源于stack exchange，提问作者frza