利用Numba优化大规模细胞突变模拟：提升NumPy数组操作效率

import numpy as np
import pandas as pd

def mutation_model(total_splits, m_type1_freq, my_type2_freq):
    mutation_types = np.array([-1, 0, 1])
    mutation_freqs = np.array([m_type1_freq, 1-(m_type1_freq + my_type2_freq), my_type2_freq])
    cell_arr = np.zeros((2**total_splits, ), dtype=int)
    exponent = 2

    for i in range(total_splits - 1):
        duplicate_arr = cell_arr[:exponent] # 复制现有细胞状态
        # 随机选择新细胞的突变类型
        random_indices = np.random.choice(len(mutation_types), size=exponent, p=mutation_freqs)
        selection = mutation_types[random_indices] # 获取突变值

        # 更新数组，累加突变
        cell_arr[exponent:(exponent * 2)] = np.add(duplicate_arr, selection)
        exponent *= 2

    # ... (后续统计代码省略)
    return {} # 示例返回空字典

import numba as nb
import numpy as np

@nb.njit('(int64, float64, float64, float64)', parallel=True)
def gen_random_mutations(size, p_neg1, p_zero, p_pos1):
    """
    使用Numba高效生成指定大小的随机突变类型数组。
    参数:
        size (int): 要生成的突变数量。
        p_neg1 (float): -1突变的概率。
        p_zero (float): 0 (野生型) 突变的概率。
        p_pos1 (float): +1突变的概率。
    返回:
        np.ndarray: 包含-1, 0, 1的int8类型数组。
    """
    # 确保概率和为1
    assert(np.isclose(p_neg1 + p_zero + p_pos1, 1.0))

    # 结果数组使用int8，节省内存
    res = np.empty(size, dtype=np.int8)

    # 定义一个大整数范围，用于概率阈值计算
    # 使用32位整数，避免64位整数生成开销，同时保证足够精度
    int_max = 1_000_000_000 

    # 计算概率阈值
    # t1用于区分-1和非-1，t2用于区分0和+1
    t1 = np.int32(np.round(p_neg1 * (int_max - 1)))
    t2 = np.int32(np.round((p_neg1 + p_zero) * (int_max - 1)))

    # 使用nb.prange进行并行循环
    for i in nb.prange(size):
        # 生成32位随机整数
        v = np.random.randint(0, int_max) 

        # 根据随机数v和阈值t1, t2确定突变类型
        # 逻辑：
        # 如果 v <= t1，则 res[i] = -1
        # 如果 t1 < v <= t2，则 res[i] = 0
        # 如果 v > t2，则 res[i] = 1
        # 
        # (v > t1) 结果是 0 或 1
        # (v > t2) 结果是 0 或 1
        # 
        # 情况1: v <= t1  => (0) + (0) - 1 = -1
        # 情况2: t1 < v <= t2 => (1) + (0) - 1 = 0
        # 情况3: v > t2   => (1) + (1) - 1 = 1
        res[i] = (v > t1) + (v > t2) - 1
    return res

# 将原始的 mutation_model 函数进行修改以使用 Numba 优化的随机数生成
def mutation_model_optimized(total_splits, m_type1_freq, my_type2_freq):
    # 假设 m_type1_freq 是 -1 的概率
    # my_type2_freq 是 +1 的概率
    # wild_type_freq 是 0 的概率
    p_neg1 = m_type1_freq
    p_pos1 = my_type2_freq
    p_zero = 1 - (p_neg1 + p_pos1)

    cell_arr = np.zeros((2**total_splits, ), dtype=np.int8) # 使用int8节省内存
    exponent = 2

    for i in range(total_splits - 1):
        # 不再需要复制整个 duplicate_arr
        # 而是直接计算新生成的细胞数量
        current_generation_size = exponent

        # 使用Numba优化的函数生成突变
        # 这里需要注意的是，gen_random_mutations 生成的是新细胞的“突变增量”
        # 而不是它们最终的累积突变值。
        # 原始代码中的 selection = mutation_types[random_indices] 也是突变增量。
        selection_increments = gen_random_mutations(
            current_generation_size, p_neg1, p_zero, p_pos1
        )

        # 更新数组：将父代的突变值与新生成的突变增量相加
        # 这一步也可以在 Numba 中优化，以避免临时数组和复制
        # 例如，可以编写一个 Numba 函数来执行此范围内的元素级加法
        cell_arr[exponent:(exponent * 2)] = cell_arr[:exponent] + selection_increments

        exponent *= 2

    # ... (后续统计代码省略)
    return {} # 示例返回空字典

# 原始参数示例
m_type1_freq = 0.078  # 对应 -1 突变
my_type2_freq = 0.0076 # 对应 +1 突变
# 注意：原始问题中的频率是 0.0078 和 0.00079，这里使用答案中的 0.078 和 0.0076。
# 并且，答案中的 gen_random_mutations 使用的概率是 (0.07843, 0.91373, 0.00784)，
# 这与问题中的 (0.0078, 0.00079) 或 (0.078, 0.0076) 均不匹配。
# 在实际应用中，请确保传递给 gen_random_mutations 的概率与模型需求一致。
# 这里我们使用与 mutation_model_optimized 匹配的参数名。

# 调用优化后的模型 (仅为示例，实际运行时需根据需求调整)
# data_optimized = []
# for i in range(100):
#     print("Working on optimized iteration: ", i + 1)
#     mutation_dict = mutation_model_optimized(30, m_type1_freq, my_type2_freq)
#     data_optimized.append(mutation_dict)