基于numpy向量化的长区间滑动窗口方法性能优化咨询

我正在为自有基因组学数据集实现覆盖长区间（单区间长度可达50000核苷酸以上）的滑动窗口计算方法。当前方案可正常运行，但执行效率极低：单区间处理耗时可达数秒，当区间长度超过150k bp时，单区间处理耗时可达数分钟。
现有实现代码如下：

import numpy as np
import pandas as pd

# Start、End为参考基因组上的区间起止标记
VectorizedRange = np.arange(Start, End)
# 100为滑动窗口大小
SlidingWindow = np.lib.stride_tricks.sliding_window_view(VectorizedRange, 100)

GroupedDictFrame = pd.DataFrame({"Bins":GenomeRange})
GroupedDictFrame["ReadCov"] = 0
GroupedDictFrame["ReadSeq"] = [list() for _ in range(len(GroupedDictFrame.index.values))]
GroupedDictFrame.set_index(keys=["Bins"], inplace=True, drop=True)

def Appender(Start, End, Width, Seq):
     AvgCov = 0
     SeqList = []
     if End <= Window[-1]:
        AvgCov += 1
        SeqList.append(Seq)
     elif End > Window[-1]:
        AvgCov += (Window[-1] - Start)/Width
        SeqList.append(Seq[0:(Window[-1] - Start)])
            
     GroupedDictFrame.loc[Window[0], "ReadCov"] += AvgCov 
     GroupedDictFrame.loc[Window[0], "ReadSeq"] = SeqList

for Window in SlidingWindow:
        
    SubsetBAM = BAMFrame[(
        (BAMFrame["start_coord"]>=Window[0])&
        (BAMFrame["start_coord"]<=Window[-1])
        )].reset_index(drop=True)
    
    SubsetBAM.apply(
        lambda x: Appender(x.start_coord, 
                           x.end_coord,
                           x.width_lis,
                           x.seq_lis), axis=1
        )

我判断当前实现的向量化逻辑存在优化空间，求可有效提升代码运行速度的改进方案。

当前代码的性能问题是多层冗余低效操作叠加导致的，和“向量化逻辑有优化空间”的判断完全吻合，核心问题点：

• 逐窗口全表扫描过滤：每遍历一个滑动窗口就对全量BAM表做一次布尔索引筛选，总时间复杂度达到O(窗口数 * BAM记录数)。当区间长度到150k时，步长为1的窗口数接近15万个，只要BAM记录数稍大，耗时会指数级上涨。
• 逐行apply加高频索引写入：pandas中axis=1的逐行apply本身比原生numpy向量化操作慢2~3个数量级，再叠加每次赋值都用.loc做标签索引查找，频繁触发pandas内部的索引校验、数据类型转换开销，积少成多延迟极高。
• 冗余内存占用：提前用np.arange生成全区间坐标数组再创建滑动窗口视图，本质只是为了获取每个窗口的起止坐标，完全没必要开辟内存存储完整的坐标序列。
• 逻辑缺陷：Appender函数直接引用外层循环的Window变量，没有作为参数传入，一旦做并行改造就会出现变量污染出脏数据；且当前逻辑只筛选了read起点落在窗口内的记录，完全漏掉了起点在上游、跨到当前窗口的read，最终统计的覆盖度结果是不准确的。

按优先级从高到低改造，完成后150k区间的处理速度可提升100~1000倍：

1. 反转遍历逻辑，避免逐窗口全表过滤 #

不要逐窗口筛选BAM记录，改为遍历数量远少于窗口数的BAM读段，直接计算每条读段覆盖的窗口范围，一次性完成赋值，从根源上降低时间复杂度：

• 步长为1的场景下（和你当前代码逻辑一致），窗口i的坐标范围是[i, i+窗口大小-1]，不需要提前生成滑动窗口数组，直接用窗口起始坐标当窗口唯一ID即可。
• 对任意一条BAM记录，直接计算它覆盖的首尾窗口ID，批量给对应ID区间的窗口加覆盖度权重，不需要逐窗口循环匹配。
• 跨窗口读段的覆盖度按落在窗口内的片段长度占读段总长度的比例加权，和你原逻辑的计算规则保持一致。

2. 用numpy数组替代pandas中间过程操作 #

覆盖度统计直接用一维numpy数组存储，序列集合用原生Python列表初始化，所有计算完成后再一次性组装为DataFrame，避免中间过程的pandas索引开销，参考实现如下：

import numpy as np
import pandas as pd

window_size = 100
step = 1 # 匹配当前代码步长为1的逻辑
total_len = End - Start
n_windows = (total_len - window_size) // step + 1
# 用float32存覆盖度，平衡精度和内存占用
cov_arr = np.zeros(n_windows, dtype=np.float32)
seq_list = [[] for _ in range(n_windows)]

# 提前把BAM表转成numpy数组，避免逐行取pandas对象的开销
read_starts = BAMFrame["start_coord"].to_numpy()
read_ends = BAMFrame["end_coord"].to_numpy()
read_widths = BAMFrame["width_lis"].to_numpy()
read_seqs = BAMFrame["seq_lis"].to_numpy()

# 提前过滤完全不在目标区间的读段，减少无效计算
valid_mask = (read_ends > Start) & (read_starts < End)
read_starts = read_starts[valid_mask] - Start # 转为相对区间起点的偏移量
read_ends = read_ends[valid_mask] - Start
read_widths = read_widths[valid_mask]
read_seqs = read_seqs[valid_mask]

# 批量计算每条读段覆盖的首尾窗口ID，超出窗口范围的直接截断到边界
win_start_per_read = np.clip(read_starts // step, 0, n_windows-1)
win_end_per_read = np.clip((read_ends - 1) // step, 0, n_windows-1)

# 遍历读段（数量远小于窗口数，长区间场景下优势极明显）批量更新结果
for i in range(len(read_starts)):
    rs, re, rw, seq = read_starts[i], read_ends[i], read_widths[i], read_seqs[i]
    ws, we = win_start_per_read[i], win_end_per_read[i]
    if ws == we:
        # 读段完全落在单个窗口内
        cov_arr[ws] += 1
        seq_list[ws].append(seq)
    else:
        # 读段跨多个窗口，分段计算覆盖权重
        first_win_end = (ws + 1)*step + window_size - step
        cov_arr[ws] += (first_win_end - rs)/rw
        seq_list[ws].append(seq[:first_win_end - rs])
        # 中间完全被读段覆盖的窗口直接加1
        if we - ws > 1:
            cov_arr[ws+1:we] += 1
        # 最后一个窗口的覆盖权重
        last_win_start = we * step
        cov_arr[we] += (re - last_win_start)/rw
        seq_list[we].append(seq[-(re - last_win_start):])

# 所有计算完成后一次性组装为DataFrame
res_df = pd.DataFrame({
    "Bins": np.arange(Start, End - window_size +1, step),
    "ReadCov": cov_arr,
    "ReadSeq": seq_list
}).set_index("Bins")

3. 进阶提速选项 #

如果改造完上述逻辑仍达不到性能要求，可按场景选择方案：

• 若输入BAM是按坐标排序好的，直接维护一个滑动读段集合：窗口向前滑动时，弹出已经离开窗口范围的read，加入新进入窗口范围的read，不需要做任何批量过滤，时间复杂度可以降到O(BAM记录数 + 窗口数)，是长区间计算的最优实现逻辑。
• 如果不需要存储每个窗口的序列片段、仅统计覆盖度，可以直接给核心循环加numba JIT装饰器，速度还能再提升10~100倍；也可以直接用samtools、bedtools的现成覆盖度统计命令，不用自己实现逻辑。
• 不建议直接存储每个窗口对应的read序列切片，长区间下这部分内存开销会达到几十上百GB，建议只存对应read的ID，需要取序列时再根据ID到原BAM文件中提取。
• 如果滑动窗口步长不是1，只需要调整窗口ID的计算逻辑即可，整体计算框架不需要改动。

补充：原实现中仅筛选read起点落在窗口内的记录，会遗漏起点在窗口外、但片段跨到当前窗口的read，最终覆盖度统计结果存在系统性偏差，上述优化代码已经修正了这个逻辑问题。

内容的提问来源于stack exchange，提问作者Luca