蒙特卡洛模拟优化批量检测策略：高效寻找最佳批次大小

import numpy as np

def simulate_batch_testing_corrected_loop(N, p, k):
    # 创建总人口，0代表未感染，1代表感染
    population = np.random.choice([0, 1], size=N, p=[1-p, p])
    n_tests = 0
    # 遍历每个k大小的批次
    for j in range(0, N, k):
        n_tests += 1  # 每次批次混合检测
        # 如果批次内有感染者（和大于0），则需要复检
        if population[j:j+k].sum() > 0:
            n_tests += k  # 复检批次内所有个体
    return n_tests

def simulate_batch_testing_optimized(N, p, k):
    # 创建总人口：0代表未感染，1代表感染
    population = np.random.choice([0, 1], size=N, p=[1-p, p])

    # 为了方便Numpy的reshape操作，将人口数组填充至k的整数倍
    padding = (k - (N % k)) % k # 计算需要填充的数量
    N_padded = N + padding

    # 使用np.pad进行填充，填充值设为0（未感染）
    population_padded = np.pad(population, (0, padding), mode='constant', constant_values=0)

    # 将填充后的人口数组重塑为(批次数量, k)的二维数组
    # 这样每一行就是一个批次
    n_batches = N_padded // k
    groups = population_padded.reshape(n_batches, k)

    # 识别需要复检的批次：如果批次内有任何感染者 (即行和大于0)
    # groups.any(axis=1) 返回一个布尔数组，表示每个批次是否含有感染者
    retests_needed = groups.any(axis=1)

    # 计算总检测次数
    # 所有批次都需要一次混合检测：n_batches
    # 需要复检的批次，每个批次额外消耗k次检测：retests_needed.sum() * k
    total_tests = n_batches + retests_needed.sum() * k
    return total_tests

def monte_carlo_simulation(N, p, max_k, num_simulations):
    results = []
    # 批次大小k的范围通常不需要遍历到N，因为k大于N/2时，效率通常很低
    # 且k=1时，相当于不批量，检测次数为N
    for k in range(1, max_k + 1): 
        total_tests_sum = 0
        for _ in range(num_simulations):
            total_tests_sum += simulate_batch_testing_optimized(N, p, k)
        avg_tests = total_tests_sum / num_simulations
        results.append((k, avg_tests))
    return results

import numpy as np
import multiprocessing

# 优化后的批量检测模拟函数
def simulate_batch_testing_optimized(N, p, k):
    # 创建总人口：0代表未感染，1代表感染
    population = np.random.choice([0, 1], size=N, p=[1-p, p])

    # 为了方便Numpy的reshape操作，将人口数组填充至k的整数倍
    padding = (k - (N % k)) % k # 计算需要填充的数量
    N_padded = N + padding

    # 使用np.pad进行填充，填充值设为0（未感染）
    population_padded = np.pad(population, (0, padding), mode='constant', constant_values=0)

    # 将填充后的人口数组重塑为(批次数量, k)的二维数组
    # 这样每一行就是一个批次
    n_batches = N_padded // k
    groups = population_padded.reshape(n_batches, k)

    # 识别需要复检的批次：如果批次内有任何感染者 (即行和大于0)
    retests_needed = groups.any(axis=1)

    # 计算总检测次数
    total_tests = n_batches + retests_needed.sum() * k
    return total_tests

# 蒙特卡洛模拟函数
def monte_carlo_simulation(N, p, max_k, num_simulations):
    results = []
    # 遍历批次大小k
    for k in range(1, max_k + 1): 
        total_tests_sum = 0
        for _ in range(num_simulations):
            total_tests_sum += simulate_batch_testing_optimized(N, p, k)
        avg_tests = total_tests_sum / num_simulations
        results.append((k, avg_tests))
    return results

# 用于并行处理的包装函数
def run_simulation_for_p(args):
    p, N, max_k, num_simulations = args
    print(f"开始模拟 p = {p}...")
    simulation_results = monte_carlo_simulation(N, p, max_k, num_simulations)
    min_avg_tests = min(simulation_results, key=lambda x: x[1])
    print(f"对于 p = {p}, 最佳批次大小 k 是 {min_avg_tests[0]}，平均检测次数为 {min_avg_tests[1]:.2f}。")
    return p, min_avg_tests

if __name__ == "__main__":
    # 参数设置
    N = 10**6  # 总样本数
    num_simulations = 100 # 蒙特卡洛模拟次数，建议减少以提高效率
    max_k = N // 2 # 批次大小k的最大值，建议限制在N//2或更小

    # 感染概率值
    p_values = [10**-1, 10**-2, 10**-3, 10**-4]

    # 准备并行任务
    tasks = [(p, N, max_k, num_simulations) for p in p_values]

    # 使用多进程池进行并行计算
    # 可以根据CPU核心数调整processes参数
    with multiprocessing.Pool(processes=len(p_values)) as pool:
        results = pool.map(run_simulation_for_p, tasks)

    # 打印所有结果
    print("\n--- 最终结果 ---")
    for p, min_avg_tests in results:
        print(f"对于 p = {p}, 最佳批次大小 k 是 {min_avg_tests[0]}，平均检测次数为 {min_avg_tests[1]:.2f}。")