合并多个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 和 arr_1 是要保存的 NumPy 数组
arr_0_part1 = np.random.rand(10, 5)
arr_1_part1 = np.random.randint(0, 100, (10, 3))

# 将数据存储在字典中
data_part1 = {'arr_0': arr_0_part1, 'arr_1': arr_1_part1}
np.savez_compressed('path/to/file/filename_part1.npz', **data_part1)

arr_0_part2 = np.random.rand(15, 5)
arr_1_part2 = np.random.randint(0, 100, (15, 3))
data_part2 = {'arr_0': arr_0_part2, 'arr_1': arr_1_part2}
np.savez_compressed('path/to/file/filename_part2.npz', **data_part2)

import numpy as np
import os

def merge_npz_files(file_list, output_filename='merged_data.npz'):
    """
    将多个NPZ文件合并为一个NPZ文件。
    假设所有NPZ文件包含相同的键，且对应数组的维度在拼接轴上保持一致。

    Args:
        file_list (list): 包含所有待合并NPZ文件路径的列表。
        output_filename (str): 合并后输出的NPZ文件名。
    """
    if not file_list:
        print("文件列表为空，无法合并。")
        return

    # 1. 加载所有NPZ文件
    # np.load返回的是NpzFile对象，它表现得像一个字典
    data_all = [np.load(fname) for fname in file_list]

    # 2. 初始化用于存储合并数据的字典
    merged_data = {}

    # 3. 获取所有文件的共同键
    # 假设所有文件都具有相同的键结构，取第一个文件的键即可
    # 如果键可能不一致，需要先找出所有文件的键的交集
    common_keys = data_all[0].keys()

    # 4. 遍历每个共同键，进行数组拼接
    for k in common_keys:
        # 收集所有文件中对应键的数组
        # list(d[k] for d in data_all) 创建一个包含所有NpzFile对象中k键对应数组的列表
        arrays_to_concatenate = list(d[k] for d in data_all)

        # 拼接这些数组
        # 默认情况下，np.concatenate沿着第一个轴（axis=0）进行拼接
        # 这要求除了拼接轴之外的其他轴的维度必须一致
        merged_data[k] = np.concatenate(arrays_to_concatenate, axis=0)
        print(f"键 '{k}' 合并完成，新数组形状：{merged_data[k].shape}")

    # 5. 保存合并后的数据到新的NPZ文件
    np.savez_compressed(output_filename, **merged_data)
    print(f"所有文件已成功合并到 '{output_filename}'")

# 示例使用
if __name__ == "__main__":
    # 创建一些示例NPZ文件
    if not os.path.exists('temp_npz_files'):
        os.makedirs('temp_npz_files')

    for i in range(3):
        arr_0_part = np.random.rand(10 + i, 5) # 模拟不同长度的数据
        arr_1_part = np.random.randint(0, 100, (10 + i, 3))
        data_part = {'arr_0': arr_0_part, 'arr_1': arr_1_part}
        np.savez_compressed(f'temp_npz_files/part_{i}.npz', **data_part)
        print(f"创建文件: temp_npz_files/part_{i}.npz, arr_0_shape: {arr_0_part.shape}")

    # 获取所有待合并的文件名
    filenames = [os.path.join('temp_npz_files', f) for f in os.listdir('temp_npz_files') if f.endswith('.npz')]
    print(f"\n待合并文件: {filenames}")

    # 执行合并操作
    merge_npz_files(filenames, 'merged_output.npz')

    # 验证合并结果
    with np.load('merged_output.npz') as merged_file:
        print("\n验证合并结果:")
        for k, v in merged_file.items():
            print(f"键: {k}, 形状: {v.shape}")

    # 清理示例文件
    import shutil
    shutil.rmtree('temp_npz_files')
    os.remove('merged_output.npz')
    print("\n示例文件已清理。")