优化大规模细胞突变模拟：使用Numba提升Python/NumPy性能

import numpy as np
import numba as nb

@nb.njit('(int64, float64, float64, float64)', parallel=True)
def gen_random_mutations(size, p1, p2, p3):
    """
    生成指定数量的突变类型索引。
    参数:
        size (int64): 要生成的突变数量。
        p1 (float64): 第一种突变类型（-1）的概率。
        p2 (float64): 第二种突变类型（0，野生型）的概率。
        p3 (float64): 第三种突变类型（+1）的概率。
    返回:
        np.ndarray: 包含突变类型（-1, 0, 1）的数组。
    """
    # 确保概率和为1
    assert(np.isclose(p1 + p2 + p3, 1.0))
    # 使用int8以减少内存占用，并直接表示-1, 0, 1
    res = np.empty(size, dtype=np.int8)

    # 定义一个较大的整数范围，以保证足够的精度
    int_max = 1_000_000_000 
    # 计算累积概率的整数阈值
    t1 = np.int32(np.round(p1 * (int_max - 1)))
    t2 = np.int32(np.round((p1 + p2) * (int_max - 1)))

    # 使用Numba的并行循环来加速随机数生成
    for i in nb.prange(size):
        # 生成一个32位整数随机数
        v = np.random.randint(0, int_max)
        # 根据阈值判断突变类型
        # (v > t1) + (v > t2) 会得到0, 1, 2
        # 减去1后，映射为-1, 0, 1
        # v <= t1: 0 (对应p1, 即-1) -> 0 - 1 = -1
        # t1 < v <= t2: 1 (对应p2, 即0) -> 1 - 1 = 0
        # v > t2: 2 (对应p3, 即+1) -> 2 - 1 = 1
        res[i] = (v > t1) + (v > t2) - 1
    return res

import numpy as np
import numba as nb
import pandas as pd

# 假设 gen_random_mutations 函数已定义并编译

@nb.njit(parallel=True)
def _mutation_model_numba_core(cell_arr, total_splits, m_type1_freq, m_type2_freq):
    """
    Numba优化的核心模拟逻辑。
    此函数直接操作cell_arr，避免Python循环和NumPy的临时数组。
    """
    mutation_types = np.array([-1, 0, 1], dtype=np.int8)
    # 计算野生型细胞的频率
    wild_type_freq = 1 - (m_type1_freq + m_type2_freq)

    # 为gen_random_mutations准备概率参数
    # 注意：这里假设m_type1_freq对应-1，m_type2_freq对应+1
    # 那么gen_random_mutations的p1是-1的概率，p2是0的概率，p3是+1的概率
    p_neg1 = m_type1_freq
    p_zero = wild_type_freq
    p_pos1 = m_type2_freq

    exponent = 2 # 当前代细胞数量的倍数，初始为2 (2个细胞)

    for i in range(total_splits - 1):
        # 生成当前代新复制细胞的突变类型
        # gen_random_mutations 返回的是-1, 0, 1
        selection = gen_random_mutations(exponent, p_neg1, p_zero, p_pos1)

        # 直接在循环中更新下一段数组，避免np.add和临时数组
        # cell_arr[exponent + j] = cell_arr[j] + selection[j]
        # 这里的cell_arr[j]是父代细胞的突变累积值
        # selection[j]是本次复制产生的突变值
        for j in nb.prange(exponent):
            cell_arr[exponent + j] = cell_arr[j] + selection[j]

        exponent *= 2 # 下一代的细胞数量翻倍

    return cell_arr

def mutation_model_optimized(total_splits, m_type1_freq, m_type2_freq):
    """
    外部接口函数，调用Numba优化的核心逻辑，并进行结果统计。
    """
    total_cells = 2**total_splits
    # 初始化细胞数组，使用int8以节省内存
    cell_arr = np.zeros(total_cells, dtype=np.int8)

    # 调用Numba优化的核心函数
    final_cell_arr = _mutation_model_numba_core(cell_arr, total_splits, m_type1_freq, m_type2_freq)

    # 统计结果
    dict_data = {
        '+2 mutation': np.count_nonzero(final_cell_arr == 2) / total_cells,
        '+1 mutation': np.count_nonzero(final_cell_arr == 1) / total_cells,
        'Wild type': np.count_nonzero(final_cell_arr == 0) / total_cells,
        '-1 mutation': np.count_nonzero(final_cell_arr == -1) / total_cells,
        '-2 mutation': np.count_nonzero(final_cell_arr == -2) / total_cells,
        '-3 mutation': np.count_nonzero(final_cell_arr == -3) / total_cells,
        '-4 mutation': np.count_nonzero(final_cell_arr == -4) / total_cells,
        '-5 mutation': np.count_nonzero(final_cell_arr == -5) / total_cells
    }
    return dict_data

# 示例运行
if __name__ == "__main__":
    data = []
    # 突变频率调整为gen_random_mutations能接受的顺序
    # 原始问题中，-1的频率是0.0078，+1的频率是0.00079
    # 这里为了匹配gen_random_mutations的p1(-1), p2(0), p3(+1)顺序，需要调整
    # 假设m_type1_freq是-1突变频率，m_type2_freq是+1突变频率
    # 原始问题描述：-1的频率约为+1的10倍，0.0078 vs 0.00079
    # 示例代码给的是0.078和0.0076，这里沿用示例代码的参数
    neg1_freq = 0.078
    pos1_freq = 0.0076

    print("开始进行Numba优化后的模拟...")
    for i in range(10): # 减少迭代次数以便快速演示
        print(f"Working on iteration: {i + 1}")
        # 注意：gen_random_mutations的p1是-1的概率，p2是0的概率，p3是+1的概率
        # 所以_mutation_model_numba_core中的参数传递需要正确映射
        # 这里m_type1_freq作为p_neg1, m_type2_freq作为p_pos1
        mutation_dict = mutation_model_optimized(30, neg1_freq, pos1_freq)
        data.append(mutation_dict)

    df = pd.json_normalize(data)
    print(df.head())
    df.to_csv('mutation_optimized.csv', index=False)
    print("模拟完成，结果已保存到 mutation_optimized.csv")