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

import numpy as np

def simulate_batch_testing_corrected(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  # 批次检测本身算一次
        # 检查当前批次是否有感染者
        if population[j:j+k].sum() > 0:
            n_tests += k  # 如果批次阳性，需要对批次内所有 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

    # 将人口数据填充，填充的个体视为未感染
    population_padded = np.pad(population, (0, padding), 'constant', constant_values=0)

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

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

    # 识别需要重新检测的批次（即批次中至少有一个感染者）
    # any(axis=1) 会检查每行（每个批次）是否存在非零元素（感染者）
    retests_needed = groups.any(axis=1)

    # 计算总检测次数
    # n_batches 是所有批次的初始检测次数
    # 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):
    population = np.random.choice([0, 1], size=N, p=[1-p, p])
    padding = (k - (N % k)) % k
    population_padded = np.pad(population, (0, padding), 'constant', constant_values=0)
    n_batches = (N + padding) // k
    groups = population_padded.reshape(n_batches, k)
    retests_needed = groups.any(axis=1)
    return n_batches + retests_needed.sum() * k

def monte_carlo_simulation_for_k(args):
    """
    针对单个 k 值进行蒙特卡洛模拟的辅助函数，用于并行化。
    """
    N, p, k, num_simulations = args
    total_tests = 0
    for _ in range(num_simulations):
        total_tests += simulate_batch_testing_optimized(N, p, k)
    avg_tests = total_tests / num_simulations
    return k, avg_tests

def run_monte_carlo_for_p(N, p, num_simulations, max_k_limit):
    """
    针对单个 p 值运行完整的蒙特卡洛模拟，并找出最优 k。
    """
    print(f"开始模拟 p = {p}...")
    results = []

    # k 的范围限制在 N/2 以内，因为 k > N/2 的批次通常效率不高
    # 并且 k=1 和 k=N 的情况是特殊但重要的边界
    k_values_to_test = list(range(1, min(N // 2 + 1, max_k_limit + 1)))

    # 确保 k=N 也在测试范围内，如果它不在 N//2 范围内
    if N not in k_values_to_test and N <= max_k_limit:
        k_values_to_test.append(N)

    # 准备并行任务参数
    tasks = [(N, p, k, num_simulations) for k in k_values_to_test]

    # 使用多进程池并行执行模拟
    # 可以根据CPU核心数调整 processes 参数
    with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
        results = pool.map(monte_carlo_simulation_for_k, tasks)

    min_avg_tests = min(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

# 参数设置
N = 10**6  # 总样本数
num_simulations = 100  # 蒙特卡洛模拟次数（建议降低以加快计算）
max_k_limit = N // 2 # 限制 k 的最大值，通常 k 超过 N/2 效率会降低

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

if __name__ == '__main__':
    start_time = time.time()
    final_optimal_results = []

    # 针对每个 p 值进行模拟
    for p_val in p_values:
        p_result = run_monte_carlo_for_p(N, p_val, num_simulations, max_k_limit)
        final_optimal_results.append(p_result)

    print("\n所有概率值模拟结果:")
    for p, (k_opt, avg_tests_opt) in final_optimal_results:
        print(f"p = {p}: 最优 k = {k_opt}, 平均检测次数 = {avg_tests_opt:.2f}")

    end_time = time.time()
    print(f"\n总模拟时间: {end_time - start_time:.2f} 秒")