基因组区间窗口平均得分映射的高效实现技术问询

问题描述 #

现有百万行规模的基因组区间DataFrame，已完成区间中心计算、窗口生成、位点得分获取等步骤，最终需要计算所有区间的同一相对中心位置（-250到+250）的得分均值，用于绘制X轴为相对位置、Y轴为平均得分的图表。当前实现代码因嵌套循环和逐行筛选导致速度极慢，需优化。

当前低效代码 #

def get_postional_avgs(promoter_df): 
    promoter_df['center'] = (promoter_df['Start'] + promoter_df['End']) // 2

    # 计算窗口起始/终止位置
    promoter_df['window_start'] = promoter_df['center'] - 250
    promoter_df['window_end'] = promoter_df['center'] + 250
    
    # 准备用于获取得分的区间DataFrame
    conservation_intervals = promoter_df[['Chromosome', 'window_start', 'window_end']]
    conservation_intervals.columns = ['Chromosome', 'Start', 'End']
    
    # 获取每个位点的得分
    base_level_scores = get_base_scores(conservation_intervals)

    average_scores = []
    # 逐行遍历原始区间，筛选对应位点并计算均值
    for index, row in promoter_df.iterrows():
        center = row['center']
        scores_for_cCRE = base_level_scores[
            (base_level_scores['Chromosome'] == row['Chromosome']) &
            (base_level_scores['Start'] >= (center - 250)) &
            (base_level_scores['End'] <= (center + 250))
        ]
        
        avg_scores = {}
        # 遍历每个相对位置，计算单位置均值
        for pos in range(-250, 251):
            current_position = center + pos
            current_scores = scores_for_cCRE[scores_for_cCRE['Start'] == current_position]
            
            if not current_scores.empty:
                avg_scores[pos] = current_scores['Score'].mean()
            else:
                avg_scores[pos] = np.nan
    
        average_scores.append(avg_scores)

    avg_score_df = pd.DataFrame(average_scores)
    return avg_score_df

高效优化方案 #

核心思路：摒弃嵌套循环和逐行筛选，利用pandas向量化运算、关联合并、分组聚合的特性，将时间复杂度从O(N*M)降至O(N)，具体步骤如下：

1. 给原始区间添加唯一ID，用于后续关联位点得分数据
2. 将原始区间的中心信息与位点得分数据关联，定位每个位点所属的原始区间
3. 计算每个位点相对于所属区间中心的位置
4. 按相对位置分组，直接批量计算均值

优化后代码 #

import pandas as pd
import numpy as np

def get_positional_avgs_optimized(promoter_df):
    # 1. 计算中心、窗口，并添加唯一区间ID
    promoter_df = promoter_df.copy()
    promoter_df['center'] = (promoter_df['Start'] + promoter_df['End']) // 2
    promoter_df['window_start'] = promoter_df['center'] - 250
    promoter_df['window_end'] = promoter_df['center'] + 250
    promoter_df['interval_id'] = range(len(promoter_df))

    # 准备带ID的区间请求DataFrame
    conservation_intervals = promoter_df[['Chromosome', 'window_start', 'window_end', 'interval_id']]
    conservation_intervals.columns = ['Chromosome', 'Start', 'End', 'interval_id']

    # 获取位点得分（复用原函数，需传递interval_id）
    base_level_scores = get_base_level_conservation_scores(conservation_intervals)

    # 2. 关联位点得分与原始区间的中心信息
    merged_df = pd.merge(
        base_level_scores,
        promoter_df[['interval_id', 'center']],
        on='interval_id',
        how='left'
    )

    # 3. 计算位点相对于中心的位置（确保坐标索引一致，此处为1-indexed）
    merged_df['relative_pos'] = merged_df['Start'] - merged_df['center']

    # 过滤超出[-250,250]范围的位点（保险操作，理论上不会存在）
    merged_df = merged_df[(merged_df['relative_pos'] >= -250) & (merged_df['relative_pos'] <= 250)]

    # 4. 按相对位置分组计算平均得分
    avg_score_df = merged_df.groupby('relative_pos')['Score'].mean().reset_index()

    # 补充缺失的相对位置，填充NaN
    all_relative_pos = pd.DataFrame({'relative_pos': range(-250, 251)})
    avg_score_df = pd.merge(all_relative_pos, avg_score_df, on='relative_pos', how='left')

    return avg_score_df

额外优化建议 #

1. 得分函数批量优化：原得分函数中逐位点添加数据效率低，可改为批量生成：

# 替换原得分函数内的位点遍历逻辑
if scores_per_base is None or len(scores_per_base) == 0:
    # 批量生成NaN条目
    positions = range(start, end)
    score_data['Chromosome'].extend([chrom]*len(positions))
    score_data['Start'].extend(positions)
    score_data['End'].extend(positions)
    score_data['Score'].extend([np.nan]*len(positions))
    score_data['interval_id'].extend([interval_id]*len(positions))
else:
    # 批量处理连续得分区间
    for start_pos, end_pos, score in scores_per_base:
        length = end_pos - start_pos
        score_data['Chromosome'].extend([chrom]*length)
        score_data['Start'].extend(range(start_pos, end_pos))
        score_data['End'].extend(range(start_pos, end_pos))
        score_data['Score'].extend([score]*length)
        score_data['interval_id'].extend([interval_id]*length)

2. 索引加速：给位点得分DataFrame的Chromosome和Start列建立索引，可大幅提升关联、筛选速度：

base_level_scores = base_level_scores.set_index(['Chromosome', 'Start'])

完整得分函数代码（适配优化逻辑） #

import pyBigWig

def get_base_level_conservation_scores(df):
    """
    获取给定基因组区间内每个位点的保守性得分，输出每行对应一个位点的DataFrame（1-indexed）
    """
    df = df.reset_index(drop=True)
    
    bw = pyBigWig.open("hg38.phyloP100way.bw")

    # 转换为0-indexing（输入为1-indexed）
    df['Start'] = df['Start'] - 1
    
    score_data = {
        'Chromosome': [],
        'Start': [],
        'End': [],
        'Score': [],
        'interval_id': []
    }
    
    for i in range(len(df)):
        chrom = df['Chromosome'][i]
        start = df['Start'][i]
        end = df['End'][i]
        interval_id = df['interval_id'][i]
    
        if start < 0 or end <= 0 or start >= end:
            print(f"无效区间 {chrom}: start={start}, end={end}")
            continue

        try:
            scores_per_base = bw.intervals(chrom, start, end)
        except RuntimeError as e:
            print(f"区间处理错误 {chrom}: start={start}, end={end} - {e}")
            continue

        if scores_per_base is None or len(scores_per_base) == 0:
            positions = range(start, end)
            score_data['Chromosome'].extend([chrom]*len(positions))
            score_data['Start'].extend(positions)
            score_data['End'].extend(positions)
            score_data['Score'].extend([np.nan]*len(positions))
            score_data['interval_id'].extend([interval_id]*len(positions))
        else:
            for start_pos, end_pos, score in scores_per_base:
                length = end_pos - start_pos
                score_data['Chromosome'].extend([chrom]*length)
                score_data['Start'].extend(range(start_pos, end_pos))
                score_data['End'].extend(range(start_pos, end_pos))
                score_data['Score'].extend([score]*length)
                score_data['interval_id'].extend([interval_id]*length)

    score_data_df = pd.DataFrame(score_data)
    score_data_df = score_data_df.drop_duplicates(subset=['Chromosome', 'Start', 'End'])
    score_data_df = score_data_df.dropna(subset=['Score'])
    
    # 转换回1-indexed
    score_data_df['Start'] = score_data_df['Start'] + 1
    score_data_df['End'] = score_data_df['End'] + 1
    
    bw.close()
    return score_data_df

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