Python中实现字典内外层级的类似np.rollaxis交换操作

嘿，我完全get到你的需求了——就是要把字典里的epoch和dataset这两层的位置给交换过来，有点像numpy里np.rollaxis调换数组维度的感觉对吧？这种结构转换在处理实验数据的时候特别常见，我来给你分享两种实用的实现方式。

首先咱们先明确一下原结构和目标结构的具体样子，用代码示例会更清楚：

原结构示例 #

import numpy as np

data_group_test = {
    "Example1": {
        "modelA": {
            20: {"dataset_a": np.array([1,2]), "dataset_b": np.array([3,4])},
            40: {"dataset_a": np.array([5,6]), "dataset_b": np.array([7,8])},
            60: {"dataset_a": np.array([9,10]), "dataset_b": np.array([11,12])}
        }
    }
}

目标结构示例 #

data_group_test_new = {
    "Example1": {
        "modelA": {
            "dataset_a": {20: np.array([1,2]), 40: np.array([5,6]), 60: np.array([9,10])},
            "dataset_b": {20: np.array([3,4]), 40: np.array([7,8]), 60: np.array([11,12])}
        }
    }
}

方法一：简洁的嵌套字典推导式 #

如果你喜欢简洁的代码，这种一行式的推导式非常适合，利用你提到的“所有dataset的epoch完全一致”这个前提，可以快速完成转换：

data_group_test_new = {
    example: {
        model: {
            dataset: {epoch: epoch_data[epoch][dataset] for epoch in epoch_data}
            for dataset in next(iter(epoch_data.values())).keys()
        }
        for model, epoch_data in models.items()
    }
    for example, models in data_group_test.items()
}

逻辑解释： #

1. 从最外层的example开始遍历，逐层深入到model层级
2. 对于每个model对应的epoch字典，我们取第一个epoch的dataset键（因为所有epoch的dataset都一样，所以可以安全这么做）
3. 为每个dataset构建新字典，键是原来的epoch，值是原结构中对应epoch下的dataset数组

方法二：直观的嵌套循环（适合新手理解） #

如果你觉得推导式有点绕，用嵌套循环的方式会更直观，每一步的逻辑都清晰可见：

data_group_test_new = {}

# 遍历外层的Example
for example, models in data_group_test.items():
    data_group_test_new[example] = {}
    # 遍历每个model
    for model, epoch_data in models.items():
        data_group_test_new[example][model] = {}
        # 获取所有dataset名称（取第一个epoch的键即可，因为所有epoch的dataset都一致）
        all_datasets = next(iter(epoch_data.values())).keys()
        # 遍历每个dataset
        for dataset in all_datasets:
            data_group_test_new[example][model][dataset] = {}
            # 遍历每个epoch，把对应的数据填充进去
            for epoch, ds_content in epoch_data.items():
                data_group_test_new[example][model][dataset][epoch] = ds_content[dataset]

验证转换结果 #

你可以用下面的代码验证转换是否正确：

print(data_group_test_new["Example1"]["modelA"]["dataset_a"][20])
# 输出: [1 2]，和原结构对应位置的数据一致

备注：内容来源于stack exchange，提问作者Diego Alejandro Gómez Pardo