蒙特卡洛模拟优化批量检测策略的实现与性能提升

import numpy as np

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

    # Mix k samples in a batch - 逻辑缺陷：这里只是从N个样本中随机抽取k个，而非将N个样本分组
    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从1到N，对于大N而言计算量巨大
        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] # 感染概率

# 模拟运行（此代码段因性能问题不建议直接运行）
# for p in p_values:
#     simulation_results = monte_carlo_simulation_original(N, p, num_simulations)
#     min_avg_tests = min(simulation_results, key=lambda x: x[1])
#     print(f"For p = {p}, optimal batch size k is {min_avg_tests[0]} with an average of {min_avg_tests[1]:.2f} tests.")

def simulate_batch_testing_iterative(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
    n_batches = N_padded // k # 计算批次总数

    # 填充人口数组，用0（未感染）填充
    population_padded = np.pad(population, (0, padding), 'constant', constant_values=0)

    # 使用reshape将一维人口数组转换为二维数组，每行代表一个批次
    groups = population_padded.reshape(n_batches, k)

    # 识别需要重新检测的批次（即批次中至少有一个感染者）
    # groups.any(axis=1) 返回一个布尔数组，True表示该批次需要重检
    retests = groups.any(axis=1)

    # 计算并返回所需的总检测次数
    # n_batches 是所有批次检测的次数
    # retests.sum() * k 是需要重检的批次中所有个体进行单独检测的次数
    return n_batches + retests.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
    n_batches = N_padded // k

    population_padded = np.pad(population, (0, padding), 'constant', constant_values=0)
    groups = population_padded.reshape(n_batches, k)

    retests = groups.any(axis=1)

    return n_batches + retests.sum() * k

def run_single_simulation(args):
    """
    用于多进程池的包装函数，运行一次simulate_batch_testing_optimized。
    """
    N, p, k = args
    return simulate_batch_testing_optimized(N, p, k)

def monte_carlo_simulation_optimized(N, p, num_simulations, k_max_limit=None):
    """
    蒙特卡洛模拟，寻找给定p值下的最优批次大小。
    N: 总样本数
    p: 感染概率
    num_simulations: 蒙特卡洛模拟次数
    k_max_limit: k的最大搜索限制，默认为N // 2
    """
    if k_max_limit is None:
        k_max_limit = N // 2 # 默认限制k的最大值为N/2

    # 对于极低的p值，最优k可能接近1/sqrt(p)。可以进一步优化k的搜索范围。
    # 例如，k_upper_bound = min(N // 2, int(2 / np.sqrt(p)))
    k_upper_bound = min(N // 2, max(1, int(2.5 / np.sqrt(p)))) # 动态调整上限，确保覆盖最优值

    results = []

    # 使用多进程池并行运行num_simulations次模拟
    pool_size = multiprocessing.cpu_count() # 获取CPU核心数
    with multiprocessing.Pool(processes=pool_size) as pool:
        for k in range(1, k_upper_bound + 1):
            # 为当前k值生成num_simulations个任务
            tasks = [(N, p, k) for _ in range(num_simulations)]

            # 异步执行任务并收集结果
            # chunksize可以优化任务分发效率
            avg_tests = np.mean(pool.map(run_single_simulation, tasks, chunksize=max(1, num_simulations // pool_size // 10)))
            results.append((k, avg_tests))
            print(f"  p={p:.0e}, k={k}/{k_upper_bound}, Avg Tests: {avg_tests:.2f}")

    return results

if __name__ == "__main__":
    # 参数
    N = 10**6  # 总样本数
    num_simulations = 100  # 减少蒙特卡洛模拟次数
    p_values = [10**-1, 10**-2, 10**-3, 10**-4] # 感染概率

    print(f"开始蒙特卡洛模拟，N={N}, num_simulations={num_simulations}")

    overall_start_time = time.time()

    for p in p_values:
        start_time = time.time()
        print(f"\n--- 正在为 p = {p:.0e} 进行模拟 ---")

        simulation_results = monte_carlo_simulation_optimized(N, p, num_simulations)

        if simulation_results:
            min_avg_tests = min(simulation_results, key=lambda x: x[1])
            print(f"对于 p = {p:.0e}, 最优批次大小 k 是 {min_avg_tests[0]}，平均检测次数为 {min_avg_tests[1]:.2f}。")
        else:
            print(f"对于 p = {p:.0e}, 未找到有效模拟结果。")

        end_time = time.time()
        print(f"p = {p:.0e} 的模拟耗时: {end_time - start_time:.2f} 秒")

    overall_end_time = time.time()
    print(f"\n所有模拟总耗时: {overall_end_time - overall_start_time:.2f} 秒")