优化 OpenMDAO Dymos 模拟中的组件数据加载：使用共享数据加载器

import openmdao.api as om
import numpy as np
import time

# 定义一个 DataLoader 类，用于管理数据的加载和缓存
class DataLoader:
    def __init__(self):
        # 使用字典作为内部缓存，键可以是加载数据时的参数组合
        self._arg_cache = {}

    def load(self, **kwargs):
        """
        根据传入的关键字参数加载数据。
        如果数据已在缓存中，则直接返回；否则加载并存入缓存。
        """
        # 将 kwargs 转换为可哈希的元组作为缓存键
        # 注意：kwargs 的顺序可能影响哈希，建议排序
        sorted_kwargs = tuple(sorted(kwargs.items()))

        if sorted_kwargs in self._arg_cache:
            print(f"--- DataLoader: 从缓存中获取数据，参数: {kwargs}")
            return self._arg_cache[sorted_kwargs]

        print(f"+++ DataLoader: 首次加载数据，参数: {kwargs}")
        # 模拟耗时的数据加载过程
        # 实际应用中，这里会调用外部库读取大文件
        time.sleep(0.1) # 模拟IO延迟
        data = f"Loaded data for options: {kwargs}" # 示例数据

        self._arg_cache[sorted_kwargs] = data
        return data

# 在组件类定义之外实例化 DataLoader 对象
# 确保所有 AtmosphereCalculator 实例共享同一个 data_loader
data_loader = DataLoader()

class AtmosphereCalculator(om.ExplicitComponent):
    def initialize(self):
        self.options.declare('time_of_year', default='summer', values=['summer', 'winter'],
                             desc='Time of year for atmospheric model')
        self.options.declare('model_version', default='v1', values=['v1', 'v2'],
                             desc='Version of the atmospheric model data')
        # 可以定义其他影响数据加载的选项

    def setup(self):
        # 根据组件的选项构建加载参数
        load_kwargs = {
            'time_of_year': self.options['time_of_year'],
            'model_version': self.options['model_version']
        }

        # 通过共享的 data_loader 加载数据
        # 即使 setup() 被多次调用，实际的数据加载（如果参数相同）只会发生一次
        self.atmospheric_data = data_loader.load(**load_kwargs)

        # 定义输入和输出
        self.add_input('altitude', val=0.0, units='m', desc='Altitude')
        self.add_output('density', val=1.225, units='kg/m**3', desc='Air density')
        self.add_output('temperature', val=288.15, units='K', desc='Air temperature')

    def compute(self, inputs, outputs):
        altitude = inputs['altitude']
        # 使用已加载的数据进行计算
        # 实际中会根据 altitude 和 self.atmospheric_data 计算密度和温度
        print(f"--- Compute: 使用数据 '{self.atmospheric_data}' 在海拔 {altitude}m 进行计算")
        outputs['density'] = 1.225 * np.exp(-altitude / 10000.0) # 简化模型
        outputs['temperature'] = 288.15 - altitude * 0.0065 # 简化模型

if __name__ == '__main__':
    print("--- 模拟 Dymos 轨迹分段对组件的影响 ---")

    # 模拟第一个分段的问题
    print("\n----- 模拟分段 1 -----")
    prob1 = om.Problem()
    comp1 = AtmosphereCalculator()
    prob1.model.add_subsystem('atm_calc', comp1)
    prob1.setup()
    prob1.set_val('atm_calc.altitude', 1000.0)
    prob1.run_model()
    print(f"分段 1 密度: {prob1.get_val('atm_calc.density'):.4f}")

    # 模拟第二个分段的问题（使用相同选项）
    print("\n----- 模拟分段 2 (相同选项) -----")
    prob2 = om.Problem()
    comp2 = AtmosphereCalculator() # 新的组件实例
    prob2.model.add_subsystem('atm_calc', comp2)
    prob2.setup() # setup() 再次被调用
    prob2.set_val('atm_calc.altitude', 2000.0)
    prob2.run_model()
    print(f"分段 2 密度: {prob2.get_val('atm_calc.density'):.4f}")

    # 模拟第三个分段的问题（使用不同选项）
    print("\n----- 模拟分段 3 (不同选项) -----")
    prob3 = om.Problem()
    comp3 = AtmosphereCalculator(options={'time_of_year': 'winter'}) # 新的组件实例，不同选项
    prob3.model.add_subsystem('atm_calc', comp3)
    prob3.setup() # setup() 再次被调用
    prob3.set_val('atm_calc.altitude', 500.0)
    prob3.run_model()
    print(f"分段 3 密度: {prob3.get_val('atm_calc.density'):.4f}")

    # 模拟第四个分段的问题（再次使用第一个分段的选项）
    print("\n----- 模拟分段 4 (再次使用分段 1 选项) -----")
    prob4 = om.Problem()
    comp4 = AtmosphereCalculator() # 新的组件实例
    prob4.model.add_subsystem('atm_calc', comp4)
    prob4.setup() # setup() 再次被调用
    prob4.set_val('atm_calc.altitude', 3000.0)
    prob4.run_model()
    print(f"分段 4 密度: {prob4.get_val('atm_calc.density'):.4f}")