OpenMDAO含Fan-In组件的并行导数着色异常问题咨询

工作流概述 #

• 多工况气动弹性稳态求解：气动与结构学科通过NonLinearBlockGS耦合，嵌套在ParallelGroup中，单工况分配单核心求解
• 某学科因需使用FD，仅支持前向模式的compute_jacvec_product；方向导数相比全雅可比矩阵可大幅节省计算资源
• 优化目标：最大化所有工况的总和目标函数，同时满足各工况的输出约束

导数并行问题现象 #

• OpenMDAO的工况并行计算加速效果符合预期，但默认无法识别导数也可并行计算
• 仅定义设计变量和单工况约束时，declare_coloring可正确识别导数并行计算空间
• 添加mux/Fan-In组件的约束与目标函数后，declare_coloring无法找到可并行的导数计算空间
• 手动为设计变量添加parallel_deriv_color参数后，导数可并行计算，但出现数值错误：总和Z值不正确，dYvec_dx1中dYvec0_dx1存在非零值（理论应为零），但分布式组件的导数计算正确

复现代码 #

import numpy as np
import openmdao.api as om

# 模拟仅支持前向FD的结构学科组件
class StructComponent(om.ExplicitComponent):
    def setup(self):
        self.add_input('load', shape=3)
        self.add_output('deformation', shape=3)
        
    def compute(self, inputs, outputs):
        outputs['deformation'] = inputs['load'] * 0.1
        
    def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode):
        # 仅支持前向模式
        if mode == 'fwd':
            if 'load' in d_inputs and 'deformation' in d_outputs:
                d_outputs['deformation'] += d_inputs['load'] * 0.1

# 气动学科组件
class AeroComponent(om.ExplicitComponent):
    def setup(self):
        self.add_input('shape', shape=3)
        self.add_input('design_var', shape=2)
        self.add_output('displacement', shape=3)
        self.add_output('performance', shape=1)
        
    def compute(self, inputs, outputs):
        outputs['displacement'] = inputs['shape'] + inputs['design_var'][0]
        outputs['performance'] = np.sum(inputs['shape']) + inputs['design_var'][1]

# 工况耦合组
class CaseGroup(om.Group):
    def setup(self):
        self.add_subsystem('struct', StructComponent())
        self.add_subsystem('aero', AeroComponent())
        self.connect('aero.displacement', 'struct.load')
        self.connect('struct.deformation', 'aero.shape')
        
        # 非线性求解器
        self.nonlinear_solver = om.NonLinearBlockGS(maxiter=10, rtol=1e-6)

# 自定义mux组件
class CustomMux(om.ExplicitComponent):
    def initialize(self):
        self.options.declare('num_cases', types=int)
    
    def setup(self):
        num_cases = self.options['num_cases']
        for i in range(num_cases):
            self.add_input(f'perf_{i}', shape=1)
        self.add_output('total_perf', shape=num_cases)
    
    def compute(self, inputs, outputs):
        outputs['total_perf'] = np.array([inputs[f'perf_{i}'][0] for i in range(self.options['num_cases'])])

# 目标函数组件
class ObjectiveComponent(om.ExplicitComponent):
    def setup(self):
        self.add_input('total_perf', shape=2)
        self.add_output('obj', shape=1)
    
    def compute(self, inputs, outputs):
        outputs['obj'] = -np.sum(inputs['total_perf'])  # 最大化总和转为最小化负总和

# 顶层问题
if __name__ == "__main__":
    num_cases = 2
    prob = om.Problem()
    model = prob.model
    
    # 添加工况并行组
    cases = model.add_subsystem('cases', om.ParallelGroup())
    for i in range(num_cases):
        cases.add_subsystem(f'case_{i}', CaseGroup())
    
    # 添加mux和目标组件
    model.add_subsystem('mux', CustomMux(num_cases=num_cases))
    model.add_subsystem('objective', ObjectiveComponent())
    
    # 连接端口
    for i in range(num_cases):
        model.connect(f'cases.case_{i}.aero.performance', f'mux.perf_{i}')
    model.connect('mux.total_perf', 'objective.total_perf')
    
    # 声明设计变量
    model.add_design_var('cases.case_0.aero.design_var', lower=-10, upper=10)
    model.add_design_var('cases.case_1.aero.design_var', lower=-10, upper=10)
    
    # 声明目标和约束
    model.add_objective('objective.obj')
    model.add_constraint('mux.total_perf', lower=-5)
    
    # 尝试手动设置并行颜色（出现错误的配置）
    # model.declare_partials('*', '*', parallel_deriv_color='case')
    
    # 尝试自动着色（添加mux后失效）
    # model.declare_coloring(wrt=['cases.case_0.aero.design_var', 'cases.case_1.aero.design_var'], method='fd')
    
    prob.setup()
    prob.run_model()
    prob.check_partials(compact_print=True)

疑问 #

1. 该方案是否无法实现导数并行加速？
2. 是否需要调整顶层与循环组的非线性/线性求解器组合？

解决方案与分析 #

可行性结论 #

该方案完全可以实现导数并行加速，问题出在mux组件的导数逻辑声明、parallel_deriv_color的配置方式上，而非方案本身不可行。

问题根源 #

1. mux组件的导数隐式性：默认mux组件的导数未显式声明，declare_coloring算法无法识别各工况输出的独立性，误以为所有工况输出共享同一导数依赖路径，因此无法拆分并行块。
2. 手动设置parallel_deriv_color的错误：直接给顶层变量设置统一颜色，会导致不同工况的导数计算被混淆，出现交叉依赖的数值错误（如dYvec0_dx1非零）。

具体解决步骤 #

1. 显式声明mux组件的导数逻辑

修改CustomMux组件，添加compute_jacvec_product方法，明确每个输入仅对应输出的特定切片，让着色算法能识别工况间的独立性：

class CustomMux(om.ExplicitComponent):
    def initialize(self):
        self.options.declare('num_cases', types=int)
    
    def setup(self):
        num_cases = self.options['num_cases']
        for i in range(num_cases):
            self.add_input(f'perf_{i}', shape=1)
        self.add_output('total_perf', shape=num_cases)
    
    def compute(self, inputs, outputs):
        outputs['total_perf'] = np.array([inputs[f'perf_{i}'][0] for i in range(self.options['num_cases'])])
    
    def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode):
        num_cases = self.options['num_cases']
        if mode == 'fwd':
            if 'total_perf' in d_outputs:
                for i in range(num_cases):
                    if f'perf_{i}' in d_inputs:
                        d_outputs['total_perf'][i] += d_inputs[f'perf_{i}'][0]
        elif mode == 'rev':
            if 'total_perf' in d_outputs:
                for i in range(num_cases):
                    if f'perf_{i}' in d_inputs:
                        d_inputs[f'perf_{i}'][0] += d_outputs['total_perf'][i]

2. 正确配置parallel_deriv_color

在CaseGroup内部为每个工况的变量设置独立的并行颜色，而非在顶层统一设置：

class CaseGroup(om.Group):
    def initialize(self):
        self.options.declare('case_id', types=int)
    
    def setup(self):
        self.add_subsystem('struct', StructComponent())
        self.add_subsystem('aero', AeroComponent())
        self.connect('aero.displacement', 'struct.load')
        self.connect('struct.deformation', 'aero.shape')
        
        self.nonlinear_solver = om.NonLinearBlockGS(maxiter=10, rtol=1e-6)
        
        # 为当前工况的所有偏导数设置独立颜色
        self.declare_partials('*', '*', parallel_deriv_color=f'case_{self.options["case_id"]}')

在顶层添加工况时传入case_id：

for i in range(num_cases):
    cases.add_subsystem(f'case_{i}', CaseGroup(case_id=i))

3. 调整线性求解器配置

由于使用前向模式方向导数，顶层线性求解器建议使用DirectSolver并关闭雅可比组装，循环组内部使用LinearBlockGS：

# 顶层线性求解器
model.linear_solver = om.DirectSolver(assemble_jac=False)

# 开启自动着色
model.declare_coloring(wrt=['cases.case_0.aero.design_var', 'cases.case_1.aero.design_var'], 
                       method='fd', show_summary=True)

效果验证 #

完成上述调整后，declare_coloring会识别出两个独立的导数计算块，实现并行加速，同时解决dYvec0_dx1的非零错误，总和Z值计算恢复正确。

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