优化批量测试的蒙特卡洛模拟：寻找最佳批次大小

import numpy as np

def simulate_batch_testing_original(N, p, k):
    infected = np.random.choice([0, 1], size=N, p=[1-p, p])

    # 错误：从N个样本中随机抽取k个，而不是将N个样本分成k个批次
    mixed_sample = np.sum(np.random.choice(infected, size=k))

    if mixed_sample == 0:
        return k
    else:
        return N

def monte_carlo_simulation_original(N, p, num_simulations):
    results = []
    for k in range(1, N + 1): # 错误：k的范围过大
        total_tests = 0
        for _ in range(num_simulations):
            total_tests += simulate_batch_testing_original(N, p, k)
        avg_tests = total_tests / num_simulations
        results.append((k, avg_tests))
    return results

# 示例参数
N = 10**6
num_simulations = 1000
p_values = [10**-1, 10**-2, 10**-3, 10**-4]

def simulate_batch_testing_corrected(N, p, k):
    # 创建总人口的感染状态
    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  # 批次测试计数
        # 如果批次中有人感染
        if population[j:j+k].sum() > 0:
            n_tests += k  # 需要对批次内所有个体进行复检
    return n_tests

def simulate_batch_testing_optimized(N, p, k):
    # 创建总人口的感染状态
    population = np.random.choice([0, 1], size=N, p=[1-p, p])

    # 计算需要填充的长度，使N成为k的整数倍
    padding = (k - (N % k)) % k
    N_padded = N + padding

    # 使用0（未感染）填充人口，使其长度为k的整数倍
    population_padded = np.pad(population, (0, padding), mode='constant', constant_values=0)

    # 计算批次数量
    n_batches = N_padded // k

    # 使用reshape将填充后的人口分组为n_batches个k大小的批次
    groups = population_padded.reshape(n_batches, k)

    # 识别需要复检的批次（即批次中至少有一个人感染）
    # groups.any(axis=1) 返回一个布尔数组，表示每个批次是否含有感染者
    retests_needed = groups.any(axis=1)

    # 计算总检测次数：批次测试次数 + 复检次数
    # n_batches 是所有批次的初始测试次数
    # retests_needed.sum() 是需要复检的批次数量
    # retests_needed.sum() * k 是这些批次所需的个体复检次数
    return n_batches + retests_needed.sum() * k

import numpy as np
import multiprocessing
import time

def simulate_batch_testing_optimized(N, p, k):
    """
    优化后的批量测试模拟函数，使用NumPy向量化操作。
    N: 总样本数
    p: 感染概率
    k: 批次大小
    返回: 总检测次数
    """
    population = np.random.choice([0, 1], size=N, p=[1-p, p])

    padding = (k - (N % k)) % k
    N_padded = N + padding

    population_padded = np.pad(population, (0, padding), mode='constant', constant_values=0)

    n_batches = N_padded // k
    groups = population_padded.reshape(n_batches, k)

    retests_needed = groups.any(axis=1)

    return n_batches + retests_needed.sum() * k

def run_single_simulation(args):
    """
    用于并行化模拟的辅助函数。
    """
    N, p, k = args
    return simulate_batch_testing_optimized(N, p, k)

def monte_carlo_simulation(N, p, num_simulations, max_k, use_multiprocessing=False):
    """
    执行蒙特卡洛模拟以找到给定p值的最优批次大小。
    N: 总样本数
    p: 感染概率
    num_simulations: 蒙特卡洛模拟次数
    max_k: 批次大小k的最大探索值
    use_multiprocessing: 是否使用多进程并行化
    返回: (最优批次大小k, 对应的平均检测次数)
    """
    results = []

    # 限制k的探索范围，避免不必要的计算
    # 考虑到实际效率，k通常不会超过N的某个较小比例，此处为N//2
    k_values_to_test = range(1, min(N + 1, max_k + 1)) 

    for k in k_values_to_test:
        if use_multiprocessing:
            # 准备并行任务的参数列表
            tasks = [(N, p, k) for _ in range(num_simulations)]
            with multiprocessing.Pool() as pool:
                total_tests_list = pool.map(run_single_simulation, tasks)
            avg_tests = np.mean(total_tests_list)
        else:
            total_tests = 0
            for _ in range(num_simulations):
                total_tests += simulate_batch_testing_optimized(N, p, k)
            avg_tests = total_tests / num_simulations
        results.append((k, avg_tests))

    # 找到平均检测次数最少的批次大小
    if results:
        min_avg_tests_result = min(results, key=lambda x: x[1])
        return min_avg_tests_result
    return None, None

# 参数设置
N = 10**6  # 总样本数
num_simulations = 100 # 蒙特卡洛模拟次数，推荐降低以平衡速度和精度
max_k_to_explore = N // 2 # 限制k的最大探索值，例如N/2
p_values = [10**-1, 10**-2, 10**-3, 10**-4]

if __name__ == "__main__":
    print(f"开始蒙特卡洛模拟，N={N}, 模拟次数={num_simulations}, k最大探索值={max_k_to_explore}")

    # 可以选择是否使用多进程并行化
    # 对于每个p值单独运行，如果核心数足够，可以手动启动多个脚本
    # 如果想在单个p值的模拟中并行化num_simulations，则设置use_multiprocessing=True

    for p in p_values:
        start_time = time.time()
        # 针对每个p值，可以在其内部使用多进程来加速num_simulations次模拟
        # 也可以选择不使用，而是在外层通过启动多个独立进程来处理不同的p值
        optimal_k, min_avg_tests = monte_carlo_simulation(N, p, num_simulations, max_k_to_explore, use_multiprocessing=True)
        end_time = time.time()

        if optimal_k is not None:
            print(f"对于 p = {p}, 最优批次大小 k 是 {optimal_k}, 平均检测次数为 {min_avg_tests:.2f}。耗时: {end_time - start_time:.2f} 秒")
        else:
            print(f"对于 p = {p}, 未找到最优批次大小。")