优化OpenMDAO Dymos组件数据加载：共享DataLoader模式

import openmdao.api as om
import numpy as np
import time # 用于模拟耗时的数据加载

class DataLoader:
    """
    一个用于按需加载并缓存数据的类。
    所有OpenMDAO组件实例将共享此DataLoader的单个实例。
    """
    def __init__(self):
        """
        初始化数据加载器，创建内部缓存。
        """
        self._arg_cache = {} # 缓存字典，键是数据加载参数，值是加载的数据

    def load(self, **kwargs):
        """
        根据提供的参数加载数据。如果数据已在缓存中，则直接返回；
        否则，执行实际加载并存入缓存。

        参数:
            **kwargs: 用于唯一标识或配置数据加载的参数。
                      例如，可以是时间、地理位置、模型版本等。
                      kwargs必须是可哈希的（例如，使用元组作为键）。
        """
        # 将kwargs转换为可哈希的元组，作为缓存的键
        # 注意：kwargs的顺序可能影响元组的哈希值，确保一致性
        cache_key = tuple(sorted(kwargs.items()))

        if cache_key in self._arg_cache:
            print(f"DataLoader: 从缓存加载数据，键: {cache_key}")
            return self._arg_cache[cache_key]

        print(f"DataLoader: 首次加载数据，键: {cache_key} (模拟耗时操作...)")
        # 模拟耗时的数据加载操作
        time.sleep(0.1) # 模拟文件读取或复杂计算

        # 实际的数据加载逻辑，根据kwargs决定加载什么数据
        # 这里只是一个示例，实际应根据业务逻辑实现
        data = {
            "property_a": np.random.rand(10) * kwargs.get('factor', 1.0),
            "property_b": np.random.rand(10) + kwargs.get('offset', 0.0)
        }

        self._arg_cache[cache_key] = data
        return data

# 在组件类定义之外实例化DataLoader，使其成为所有组件共享的单例
data_loader = DataLoader()

import openmdao.api as om
import numpy as np

# 假设data_loader已经如上所示被定义并实例化
# data_loader = DataLoader() 

class AtmosphereCalculator(om.ExplicitComponent):
    """
    一个计算大气属性的OpenMDAO组件，使用共享DataLoader加载数据。
    """
    def initialize(self):
        """
        定义组件的选项。
        """
        self.options.declare('time_of_year', default='summer', types=str,
                             desc='指定加载哪个季节的大气数据')
        self.options.declare('altitude_range_max', default=10000.0, types=float,
                             desc='指定大气数据适用的最大高度范围')
        # 其他可能影响数据加载的选项...

    def setup(self):
        """
        在setup中调用共享的DataLoader加载数据。
        """
        # 从组件选项构建用于DataLoader的参数
        load_kwargs = {
            'season': self.options['time_of_year'],
            'max_alt': self.options['altitude_range_max']
            # 可以添加其他影响数据加载的选项
        }

        # 调用共享的DataLoader加载数据
        # 首次调用时数据会被加载并缓存，后续调用直接从缓存获取
        self.atmospheric_data = data_loader.load(**load_kwargs)

        # 定义组件的输入和输出
        self.add_input('altitude', val=0.0, units='m', desc='飞行器高度')
        self.add_output('density', val=1.225, units='kg/m**3', desc='大气密度')
        self.add_output('temperature', val=288.15, units='K', desc='大气温度')

        # 假设大气数据中包含了一些属性计算所需的系数
        self.add_output('property_a_factor', val=1.0)
        self.add_output('property_b_offset', val=0.0)

    def compute(self, inputs, outputs):
        """
        使用加载的数据计算大气属性。
        """
        altitude = inputs['altitude']

        # 实际的计算逻辑会使用 self.atmospheric_data 中的数据
        # 这里仅为示例，简化计算
        outputs['density'] = self.atmospheric_data['property_a'][0] * np.exp(-altitude / 10000.0)
        outputs['temperature'] = self.atmospheric_data['property_b'][0] - (altitude * 0.0065)

        # 示例：将加载数据中的一部分作为输出
        outputs['property_a_factor'] = self.atmospheric_data['property_a'][1]
        outputs['property_b_offset'] = self.atmospheric_data['property_b'][2]


# --- 完整示例：如何在一个OpenMDAO问题中使用此组件 ---
if __name__ == "__main__":
    # 创建一个OpenMDAO问题
    prob = om.Problem()

    # 将AtmosphereCalculator组件添加到问题中
    # 可以创建多个实例，模拟不同分段或不同配置
    prob.model.add_subsystem('atmos_calc_segment1', AtmosphereCalculator(time_of_year='summer', altitude_range_max=10000.0))
    prob.model.add_subsystem('atmos_calc_segment2', AtmosphereCalculator(time_of_year='winter', altitude_range_max=12000.0))
    prob.model.add_subsystem('atmos_calc_segment3', AtmosphereCalculator(time_of_year='summer', altitude_range_max=10000.0)) # 与segment1配置相同

    # 设置驱动器
    prob.driver = om.ScipyOptimizeDriver()
    prob.driver.options['optimizer'] = 'SLSQP'

    # 设置问题
    prob.setup()

    # 运行问题，观察DataLoader的输出
    print("\n--- 第一次运行问题 ---")
    prob.run_model()

    print("\n--- 验证结果 ---")
    print(f"Segment 1 Density: {prob.get_val('atmos_calc_segment1.density')}")
    print(f"Segment 2 Density: {prob.get_val('atmos_calc_segment2.density')}")
    print(f"Segment 3 Density: {prob.get_val('atmos_calc_segment3.density')}")

    # 再次运行问题，验证缓存效果
    print("\n--- 第二次运行问题 (验证缓存) ---")
    prob.run_model()