大型MultiIndex DataFrame时空查询的灵活读取、重组与绘图方法

1 读取逻辑优化 #

原读取逻辑存在硬编码、全量加载的问题，优化后直接按需读取，大幅降低内存占用：

• 先根据传入的ot参数[开始时间, 结束时间, 步长]计算需要保留的时间步序号，跳过所有不在该范围内的时间步对应数据块，不需要加载全量100万时间步数据，内存占用可降低90%以上
• 移除硬编码的skiprows2逻辑，每个时间步对应的数据块起始位置可通过公式直接计算：t_block_start = t_idx * (nblocks + nheaders)，仅跳过非目标时间步的所有行+目标时间步的重复表头行即可
• 读取每个chunk时直接绑定对应的真实t值，不用枚举的k作为时间索引，避免索引和实际时间不匹配

优化后读取代码示例：

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

def arg_parser():
    # 保留原有参数解析逻辑，补充oc/or/ot参数接收即可
    pass

if __name__ == '__main__':
    args = arg_parser()
    in_file = args.input_filename
    out_file = args.output_filename
    c, r, total_t = args.cols, args.rows, args.steps
    oc, or_, ot = args.oc, args.or_, args.ot
    start_t, end_t, step_t = ot
    target_t_list = np.arange(start_t, end_t+1, step_t)
    
    nblocks = c * r
    nheaders = 2
    block_length = nblocks + nheaders
    # 预计算所有需要跳过的行
    skip_rows = []
    for t_idx in range(total_t):
        block_start = t_idx * block_length
        if t_idx not in target_t_list:
            skip_rows.extend(range(block_start, block_start + block_length))
        else:
            skip_rows.append(block_start + 1) # 仅跳过s,A,B,C重复表头行，保留t行读时间值
    
    # 按需分块读取
    chunks = []
    for t in target_t_list:
        chunk = pd.read_csv(
            in_file, 
            skiprows=skip_rows, 
            skipinitialspace=True,
            nrows=nblocks, 
            header=None, 
            names=['s','A','B','C']
        )
        chunk['t'] = t
        chunks.append(chunk)
    
    # 合并+按oc/or过滤空间点（默认s命名规则为s_列号_行号，可根据实际规则调整）
    df = pd.concat(chunks, ignore_index=True)
    df[['prefix','col_idx','row_idx']] = df['s'].str.split('_', expand=True)
    df['col_idx'] = df['col_idx'].astype(int)
    df['row_idx'] = df['row_idx'].astype(int)
    df = df[(df['col_idx'] < oc) & (df['row_idx'] < or_)]
    # 生成便于查询的MultiIndex结构
    df = df.set_index(['t','s']).drop(['prefix','col_idx','row_idx'], axis=1)
    # 输出汇总CSV
    df.to_csv(out_file)

2 灵活绘图功能封装 #

直接封装两类通用绘图函数，适配所有时空查询需求：

2.1 时序类折线图 #

支持单图多空间点对比、多子图单空间点展示两种模式，可自由选择展示的指标：

def plot_timeseries(df, metrics=['A','B','C'], compare_points=None, subplots=False):
    if compare_points is None:
        compare_points = df.index.get_level_values('s').unique()
    if subplots:
        fig, axes = plt.subplots(len(compare_points), 1, figsize=(10, 3*len(compare_points)), sharex=True)
        for i, s in enumerate(compare_points):
            df.loc[(slice(None), s), metrics].plot(ax=axes[i], title=f'空间点{s}时序变化')
            axes[i].set_ylabel('取值')
        axes[-1].set_xlabel('时间步')
    else:
        fig, ax = plt.subplots(figsize=(12,6))
        for s in compare_points:
            for m in metrics:
                df.loc[(slice(None), s), m].plot(ax=ax, label=f'{s}-{m}')
        ax.set_title('多空间点时序对比')
        ax.set_xlabel('时间步')
        ax.set_ylabel('取值')
        ax.legend()
    plt.tight_layout()
    return fig

2.2 空间类热力图 #

支持选择统计类型（均值/峰值）、指标、时间切片：

def plot_heatmap(df, metric='A', stat='mean', t_slice=None):
    # 按时间切片过滤
    if t_slice is not None:
        plot_df = df.loc[t_slice, :]
    else:
        plot_df = df
    # 计算统计值
    if stat == 'mean':
        stat_df = plot_df.groupby('s')[metric].mean()
    elif stat == 'max':
        stat_df = plot_df.groupby('s')[metric].max()
    # 转成空间网格结构
    grid = np.zeros((or_, oc))
    for s in stat_df.index:
        _, col, row = s.split('_')
        grid[int(row), int(col)] = stat_df[s]
    # 绘图
    fig, ax = plt.subplots(figsize=(oc+1, or_+1))
    sns.heatmap(grid, annot=True, cmap='viridis', ax=ax)
    ax.set_title(f'空间{metric}值{stat}热力图')
    ax.set_xlabel('列号')
    ax.set_ylabel('行号')
    return fig

3 调用示例 #

可根据CLI传入的参数动态选择绘图类型，比如：

# 生成指定空间点的多指标时序对比图
fig1 = plot_timeseries(df, compare_points=['s_1_0','s_1_1','s_2_0','s_2_1'])
fig1.savefig('timeseries_compare.png')
# 生成A值全时间范围均值热力图
fig2 = plot_heatmap(df, metric='A', stat='mean')
fig2.savefig('A_mean_heatmap.png')
plt.show()

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