如何修改循环模拟代码以计算notears实验结果的平均值？

现有一段基于Notears的循环模拟实验代码，当前每次循环都会输出单次实验的acc指标结果，现在需要改成获取多次实验结果的平均值，原代码如下：

def run_experiment():
    from notears import utils
    # utils.set_random_seed(1) this line cannot be used to ensure different outcomes in each round of the loop

    n, d, s0, graph_type, sem_type = 1000, 20, 20, 'ER', 'gauss'
    B_true = utils.simulate_dag(d, s0, graph_type)
    W_true = utils.simulate_parameter(B_true)
    np.savetxt('W_true.csv', W_true, delimiter=',')

    X = utils.simulate_linear_sem(W_true, n, sem_type)
    np.savetxt('X.csv', X, delimiter=',')

    W_est = notears_linear(X, lambda1=0.1, loss_type='l2')
    assert utils.is_dag(W_est)
    np.savetxt('W_est.csv', W_est, delimiter=',')
    acc = utils.count_accuracy(B_true, W_est != 0)
    print(acc)

if __name__ == '__main__':
    num_experiments = 3
    start = timer()
    for _ in range(num_experiments):
        run_experiment()
    end = timer()
    print(end - start)

修改后的代码 #

import numpy as np
from timeit import timer
from notears import utils
from notears.linear import notears_linear

def run_experiment():
    # utils.set_random_seed(1) 保持注释，确保每次实验的随机性
    n, d, s0, graph_type, sem_type = 1000, 20, 20, 'ER', 'gauss'
    B_true = utils.simulate_dag(d, s0, graph_type)
    W_true = utils.simulate_parameter(B_true)
    # 可选：给文件加索引后缀避免覆盖，比如 f'W_true_{idx}.csv'，不需要保存可直接注释
    # np.savetxt(f'W_true_{idx}.csv', W_true, delimiter=',')

    X = utils.simulate_linear_sem(W_true, n, sem_type)
    # np.savetxt(f'X_{idx}.csv', X, delimiter=',')

    W_est = notears_linear(X, lambda1=0.1, loss_type='l2')
    assert utils.is_dag(W_est)
    # np.savetxt(f'W_est_{idx}.csv', W_est, delimiter=',')
    
    acc = utils.count_accuracy(B_true, W_est != 0)
    return acc

if __name__ == '__main__':
    num_experiments = 3
    start = timer()
    
    # 存储所有实验的acc结果
    all_accs = []
    for _ in range(num_experiments):
        acc = run_experiment()
        all_accs.append(acc)
        print(f"单次实验结果: {acc}")  # 可选：保留单次结果打印
    
    end = timer()
    
    # 计算各指标的平均值
    avg_acc = {
        key: np.mean([acc[key] for acc in all_accs])
        for key in all_accs[0].keys()
    }
    
    print(f"\n{num_experiments}次实验平均结果:")
    for key, value in avg_acc.items():
        print(f"{key}: {value:.4f}")
    
    print(f"\n总耗时: {end - start:.2f}秒")

关键修改点 #

• 调整函数返回值：将原函数的print(acc)改为return acc，让主函数可以收集每次实验的结果数据。
• 收集实验结果：在主函数中创建列表all_accs，循环执行实验时将每次返回的acc存入列表。
• 计算平均值：通过字典推导式结合np.mean，对acc字典中的每个评价指标分别计算多次实验的平均值。
• 避免文件覆盖（可选）：原代码每次实验会覆盖同名文件，若需要保留单次实验的输出文件，可以给文件名添加循环索引后缀，或者直接注释掉保存文件的代码。

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