如何优化numpy数组按组取列功能及仿真数据存储自定义类

优化方案 #

核心问题解决：新增组映射字典完全可行 #

你的思路是对的，原有get_group方法每次都要遍历所有变量的分组字典筛选匹配项，时间复杂度是O(n)，变量数量越多速度越慢。新增一个以组名为键、对应列索引列表为值的字典可以直接把查询复杂度降到O(1)，性能提升非常明显。

具体修改点 #

1. 在__init__方法中初始化分组映射字典
2. 修改变量新增逻辑，同步更新分组映射
3. 改写get_group方法直接查映射字典取值

其他优化建议 #

• 取消逐次np.append操作：该操作每次都会重新申请整个数组的内存，变量多了之后性能极差。你已经明确仿真总时长和步长，建议提前估算最大变量数，一次性预分配足够大的二维数组，用计数器记录已使用列数，新增变量直接填充对应列即可。如果无法提前确定变量数，也可以按块扩容（比如每次新增10列而非1列），大幅降低内存重分配的频次。
• 补充逐时间步赋值方法：现有current_timestep属性没有被利用，可以新增对应方法直接给当前仿真步的变量赋值，更符合常规仿真的执行逻辑。
• 替换通用异常：将新增变量重复时抛出的通用Exception改为ValueError，同时在报错信息中带上重复的变量名，降低排查成本。
• 增加分组存在性校验：调用get_group时先判断组名是否存在，提前抛出明确的报错信息，避免默认KeyError的提示太模糊。

修改后参考代码 #

import numpy as np

class State:
    def __init__(self, time_step_size, pre_allocate_cols=100):
        self.mapping = dict()
        self.group = dict()
        self.group_map = dict() # 新增的分组到列索引的映射
        self.description = dict()
        sec_in_year = 31536000
        self.total_steps = int(np.floor(sec_in_year / time_step_size))
        # 预分配足够大的数组，默认预留100列空间，可根据需求调整
        self.data = np.empty([self.total_steps, pre_allocate_cols], dtype=float)
        self.data[:] = np.nan
        # 初始化时间列
        self.data[:,0] = np.linspace(0, sec_in_year, self.total_steps)
        self.mapping['time'] = 0
        self.group['time'] = None
        self.description['time'] = "Simulation time in sec"
        self.used_cols = 1 # 记录已使用的列数
        self.current_timestep = 0

    def __getitem__(self, key):
        return self.data[:, self.mapping[key]]

    def get_group(self, group_name):
        if group_name not in self.group_map:
            raise ValueError(f"分组{group_name}不存在")
        return self.data[:, self.group_map[group_name]]

    def __setitem__(self, key, value):
        self.data[:, self.mapping[key]] = value

    # 新增的按当前时间步赋值的方法
    def set_at_current_step(self, key, value):
        self.data[self.current_timestep, self.mapping[key]] = value

    def add_var(self, var_name, group, description):
        if var_name in self.mapping:
            raise ValueError(f"变量{var_name}已存在，无法重复添加")
        # 预分配空间不足时按块扩容，每次加50列
        if self.used_cols >= self.data.shape[1]:
            new_cols = np.empty([self.total_steps, 50], dtype=float)
            new_cols[:] = np.nan
            self.data = np.append(self.data, new_cols, axis=1)
        # 赋值
        self.mapping[var_name] = self.used_cols
        self.group[var_name] = group
        self.description[var_name] = description
        # 更新分组映射
        if group not in self.group_map:
            self.group_map[group] = []
        self.group_map[group].append(self.used_cols)
        self.used_cols += 1

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