高效合并多个NumPy NPZ文件教程

import numpy as np
import os

# 假设 file_list 包含所有要合并的 npz 文件路径
# 错误的合并尝试
# data_all = [np.load(fname) for fname in file_list]
# merged_data = {}
# for data in data_all:
#     [merged_data.update({k: v}) for k, v in data.items()]
# np.savez('new_file.npz', **merged_data)

import numpy as np

# 示例数据
arr_0_data = np.random.rand(10, 5)
arr_1_data = np.random.randint(0, 100, size=(10, 3))

# 将数据存储在字典中，并赋予有意义的键名
data_to_save = {
    'feature_array': arr_0_data,
    'label_array': arr_1_data
}

# 使用 np.savez_compressed 保存数据
# np.savez_compressed 会将字典的键作为 npz 文件内部的数组名
np.savez_compressed('path/to/file/my_data_part1.npz', **data_to_save)

# 假设有另一部分数据
arr_0_data_part2 = np.random.rand(15, 5)
arr_1_data_part2 = np.random.randint(0, 100, size=(15, 3))
data_to_save_part2 = {
    'feature_array': arr_0_data_part2,
    'label_array': arr_1_data_part2
}
np.savez_compressed('path/to/file/my_data_part2.npz', **data_to_save_part2)

import numpy as np
import os

def merge_npz_files(filenames, output_filename='merged_data.npz'):
    """
    合并多个 npz 文件中的数据到单个 npz 文件。
    要求所有 npz 文件包含相同的键，且对应数组可在指定轴上拼接。

    Args:
        filenames (list): 包含所有待合并 npz 文件路径的列表。
        output_filename (str): 合并后 npz 文件的输出路径。
    """
    if not filenames:
        print("没有指定文件进行合并。")
        return

    # 1. 加载所有 npz 文件
    # np.load 返回的是一个 NpzFile 对象，行为类似字典
    data_all = [np.load(fname) for fname in filenames]

    merged_data = {}

    # 2. 遍历第一个文件的所有键（假设所有文件具有相同的键）
    # 获取所有键，并确保它们在所有文件中都存在
    first_file_keys = list(data_all[0].keys())

    # 验证所有文件是否包含相同的键
    for i, data_obj in enumerate(data_all):
        current_keys = list(data_obj.keys())
        if sorted(current_keys) != sorted(first_file_keys):
            raise ValueError(
                f"文件 '{filenames[i]}' 的键与第一个文件 '{filenames[0]}' 不一致。"
                f"第一个文件的键: {first_file_keys}, 当前文件的键: {current_keys}"
            )

    # 3. 对每个键，从所有文件中提取对应数组并进行拼接
    for k in first_file_keys:
        # 从每个 NpzFile 对象中获取键 k 对应的数组，并放入一个列表中
        arrays_to_concatenate = [d[k] for d in data_all]

        # 使用 np.concatenate 沿默认轴（通常是0轴）拼接数组
        # 确保所有数组在拼接轴之外的维度是匹配的
        try:
            merged_data[k] = np.concatenate(arrays_to_concatenate, axis=0)
        except ValueError as e:
            print(f"合并键 '{k}' 对应的数组时发生错误: {e}")
            print(f"请检查所有文件中 '{k}' 对应数组的维度是否兼容。")
            # 打印相关数组的形状以便调试
            for i, arr in enumerate(arrays_to_concatenate):
                print(f"文件 '{filenames[i]}' 中 '{k}' 的形状: {arr.shape}")
            raise

    # 4. 保存合并后的数据
    np.savez_compressed(output_filename, **merged_data)
    print(f"成功合并 {len(filenames)} 个 npz 文件到 '{output_filename}'。")

# --- 示例用法 ---
if __name__ == "__main__":
    # 创建一些示例 npz 文件
    # 确保路径存在
    os.makedirs('temp_npz_files', exist_ok=True)

    # 文件 1
    data1 = {
        'features': np.random.rand(10, 5),
        'labels': np.random.randint(0, 2, size=(10,)),
        'metadata': np.array(['A'] * 10)
    }
    np.savez_compressed('temp_npz_files/data_part_1.npz', **data1)

    # 文件 2
    data2 = {
        'features': np.random.rand(15, 5),
        'labels': np.random.randint(0, 2, size=(15,)),
        'metadata': np.array(['B'] * 15)
    }
    np.savez_compressed('temp_npz_files/data_part_2.npz', **data2)

    # 文件 3
    data3 = {
        'features': np.random.rand(5, 5),
        'labels': np.random.randint(0, 2, size=(5,)),
        'metadata': np.array(['C'] * 5)
    }
    np.savez_compressed('temp_npz_files/data_part_3.npz', **data3)

    # 待合并的文件列表
    file_list = [
        'temp_npz_files/data_part_1.npz',
        'temp_npz_files/data_part_2.npz',
        'temp_npz_files/data_part_3.npz'
    ]

    # 执行合并
    merge_npz_files(file_list, 'merged_output.npz')

    # 验证合并结果
    merged_npz = np.load('merged_output.npz')
    print("\n--- 验证合并结果 ---")
    for key in merged_npz.keys():
        print(f"键 '{key}': 形状 {merged_npz[key].shape}")

    # 预期形状: features (10+15+5, 5) = (30, 5)
    # 预期形状: labels (10+15+5,) = (30,)
    # 预期形状: metadata (10+15+5,) = (30,)

    # 清理生成的临时文件
    for f in file_list:
        os.remove(f)
    os.remove('merged_output.npz')
    os.rmdir('temp_npz_files')
    print("\n临时文件已清理。")