使用Pandas高效读取HDF5文件及数据验证实践

import pandas as pd
import numpy as np # 导入numpy以生成示例数据，实际使用时无需

# --- 准备一个示例HDF5文件 (如果已有文件，可跳过此步骤) ---
# 假设我们有一个名为 'file.h5' 的HDF5文件，其中包含 'Mvir_all', 'pid', 'id', 'other_data' 等数据集
# 为了演示，我们先创建一个这样的文件
data_size = 1000
example_data = {
    'Mvir_all': np.random.rand(data_size) * 1e12,
    'pid': np.random.randint(-1, data_size // 2, data_size),
    'id': np.arange(data_size),
    'other_data': np.random.rand(data_size)
}
example_df = pd.DataFrame(example_data)

# 将示例DataFrame保存为HDF5文件
gal_file = "file.h5"
# 注意：Pandas将DataFrame保存到HDF5时，会创建一个默认的键，例如 'df' 或用户指定的键
example_df.to_hdf(gal_file, key='data', mode='w')
print(f"示例HDF5文件 '{gal_file}' 已创建。")
# -----------------------------------------------------------

# 使用 pandas.read_hdf 读取 HDF5 文件
# 指定要读取的列名列表
columns_to_read = ["Mvir_all", "pid", "id"]

try:
    # key 参数是必需的，因为它指定了HDF5文件中DataFrame存储的键
    # 如果不知道键，可以使用 h5py.File(gal_file, 'r').keys() 查看
    df = pd.read_hdf(gal_file, key='data', columns=columns_to_read)

    print("\n成功从HDF5文件读取指定列：")
    print(df.head())
    print(f"\n读取的DataFrame形状: {df.shape}")
    print(f"读取的DataFrame列名: {df.columns.tolist()}")

except KeyError as e:
    print(f"错误：无法找到指定的HDF5键或列。请检查文件结构和列名。错误信息: {e}")
except FileNotFoundError:
    print(f"错误：文件 '{gal_file}' 未找到。请检查文件路径。")
except Exception as e:
    print(f"读取HDF5文件时发生未知错误: {e}")

# 假设df已经通过上述方法正确加载
if 'pid' in df.columns and 'id' in df.columns:
    # 分离光环（pid == -1）和子光环（pid != -1）
    halos = df[df['pid'] == -1]
    subhalos = df[df['pid'] != -1]

    print(f"\n光环数量: {len(halos)}, 子光环数量: {len(subhalos)}")

    # 提取子光环的pid和光环的id
    pid_subhalos = subhalos['pid']
    id_halos = halos['id']

    # 检查所有子光环的pid是否都存在于光环的id中
    # 使用 .isin() 方法进行高效检查
    all_pids_in_halos = pid_subhalos.isin(id_halos).all()

    if all_pids_in_halos:
        print("所有子光环的'pid'值都存在于光环的'id'中。")
    else:
        # 如果存在不匹配的pid，可以找出它们
        mismatched_pids = pid_subhalos[~pid_subhalos.isin(id_halos)]
        print("警告：并非所有子光环的'pid'值都存在于光环的'id'中。")
        print(f"不匹配的'pid'值数量: {len(mismatched_pids.unique())}")
        # print("部分不匹配的'pid'值示例:", mismatched_pids.unique()[:5]) # 打印前5个不匹配的pid

else:
    print("\nDataFrame中缺少'pid'或'id'列，无法执行数据验证。")