如何用Numpy遍历CSV文件，按轨迹ID批量计算MSD并导出结果

需求说明 #

现有一个4列的CSV文件，存储轨迹数据：

• 第1列：轨迹ID
• 第2列：时间（范围0-39）
• 第3、4列：X、Y坐标

已实现calc_msd_np函数用于计算单方向的MSD，需要通过嵌套循环按轨迹ID分组，为每个轨迹计算X、Y方向的MSD，并将所有结果写入新的CSV文件。

示例数据 #

ID Time    X           Y
1   0   1601.335544 496.4608855
1   1   1581.838184 497.3440951
1   2   1574.051153 493.2139985
...
2   39  1591.533636 576.2326799

现有MSD计算函数 #

import numpy as np

def calc_msd_np(x):
    msd = []
    # 遍历时间间隔s（从1到轨迹长度-1）
    for s in range(1, len(x)):
        # 计算间隔s的位移差
        dx = x[s:] - x[:-s]
        # 计算该间隔的平均位移平方，加入MSD列表
        msd.append(np.average(dx**2))
    return msd

完整实现代码 #

import numpy as np
import pandas as pd

# 读取CSV数据
df = pd.read_csv("your_trajectory_data.csv")  # 替换为你的实际文件路径

# 获取所有唯一的轨迹ID
unique_ids = df["ID"].unique()

# 准备存储结果的列表
results = []

# 外层循环：遍历每个轨迹ID，完成分组处理
for trajectory_id in unique_ids:
    # 提取当前ID对应的轨迹数据
    trajectory_data = df[df["ID"] == trajectory_id]
    # 按时间排序，确保序列连续
    trajectory_data = trajectory_data.sort_values("Time")
    # 提取X、Y坐标的numpy数组
    x_coords = trajectory_data["X"].values
    y_coords = trajectory_data["Y"].values
    
    # 调用函数计算X、Y方向的MSD
    msd_x = calc_msd_np(x_coords)
    msd_y = calc_msd_np(y_coords)
    
    # 内层循环：拆分每个时间间隔的MSD结果，整理为结构化数据
    for s in range(len(msd_x)):
        results.append({
            "ID": trajectory_id,
            "Time_Interval": s + 1,  # 时间间隔从1开始计数
            "MSD_X": msd_x[s],
            "MSD_Y": msd_y[s]
        })

# 将结果转换为DataFrame并写入CSV
result_df = pd.DataFrame(results)
result_df.to_csv("trajectory_msd_results.csv", index=False)

嵌套循环逻辑说明 #

1. 外层循环：遍历所有唯一轨迹ID，完成数据分组与提取，确保每个轨迹独立处理。
2. 函数内循环：针对当前轨迹的X/Y坐标数组，遍历所有可能的时间间隔s，计算每个间隔下的平均位移平方，生成该方向的MSD序列。
3. 结果整理循环：将每个轨迹的MSD结果按时间间隔拆分，转换为CSV所需的结构化格式。

注意事项 #

• 替换代码中的your_trajectory_data.csv为实际文件路径。
• 若原始数据存在时间乱序，sort_values("Time")步骤可保证MSD计算基于连续的时间序列。
• 最终输出的CSV包含4列：轨迹ID、时间间隔、X方向MSD、Y方向MSD。

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