python - 过滤 numpy 数组中的特定列

Question

我尝试根据自己的需要制作自定义数据结构。我想用一些额外的信息记录模拟变量。我想创建一个预先分配的 2d numpy 数组并在每个时间步填写值。而且我想要一些方便的功能，例如变量可以是一个组的成员，这对于实现能量平衡很有用。

我的问题是我的 get_group 函数太慢了。我是否需要一个以组名作为键并将列索引作为列表的字典，还是有不同的方法？

我也将感谢任何其他改进课程的建议。

class State:

    def __init__(self,time_step_size):
        self.mapping = dict()
        self.group = dict()
        self.description = dict()
        sec_in_year = 31536000
        time = np.linspace(0, sec_in_year, int(np.floor(sec_in_year / time_step_size)))
        self.data = np.empty([len(time),1],dtype=float)
        self.data[:,0] = time
        self.mapping['time'] = 0
        self.group['time'] = None
        self.description['time']="Simulation time in sec"
        self.current_timestep = 0

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

    def get_group(self, group_name):
        vars=[k for k, v in self.group.items() if v == group_name]
        cols=[self.mapping[k] for k in vars]
        return self.data[:,cols]

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

    def add_var(self, var_name, group, description):
        if var_name in self.mapping or var_name in self.group or var_name in self.description:
            raise Exception("Tried to add a variable to the state, that already exists")

        to_append=np.empty([len(self.data),1])
        to_append[:]=np.nan
        self.data=np.append(self.data,to_append,axis=1)

        self.mapping[var_name]=self.data.shape[1]-1
        self.group[var_name]= group
        self.description[var_name] = description

标签： pythonnumpydictionaryoptimizationdata-structures