优化Python生物信息学脚本中的__getitem__方法以提升性能

我是一名刚接触编程的生物学家，正在编写Python生物信息学脚本，用于遍历多条序列、查找特定位置的氨基酸并计算其在基因组数据库中的频率。但脚本优化不足，处理普通文件需半分钟，大文件耗时更久。经性能分析发现，Bio.Seq中Seq类的__getitem__函数调用次数过多（达6900万+次）是主要性能瓶颈，尝试过优化但效果甚微，恳请提供可行的优化方案。

原脚本代码 #

def main(self):
    # Process input data and extract necessary information
    mutation_count, mutation_subjects, mutation_positions, alignment_file, seq_ID = self.process_input()

    # Initialize a dictionary to count amino acids at each mutation position
    aa_counts = {pos: {} for pos in mutation_positions}

    # Read the alignment file using Bio.AlignIO
    alignment = AlignIO.read(alignment_file, 'clustal')
    
    # Find the reference sequence in the alignment
    ref_seq = None
    for record in alignment:
        if record.id == seq_ID:
            ref_seq = record
            break

    # Convert sequences in the alignment to strings for easier processing
    ref_seq_str = str(ref_seq.seq)
    alignment_str = {record.id: str(record.seq) for record in alignment}

    # Count amino acids at each position in the alignment
    for seq in alignment_str.values():
        pos_in_ref_seq = 0
        for pos, aa in enumerate(seq):
            if ref_seq_str[pos] != '-':
                pos_in_ref_seq += 1

            if pos_in_ref_seq in aa_counts:
                aa_counts[pos_in_ref_seq][aa] = aa_counts[pos_in_ref_seq].get(aa, 0) + 1
                    
    # If a specific position is provided, calculate and print the amino acid frequencies at that position
    if self.args.position:
        position = self.args.position
        total_count = sum(aa_counts[position].values())
        for aa, count in aa_counts[position].items():
            freq = (count / total_count * 100) if total_count > 0 else 0
            print(f"Amino Acid: {aa} | Frequency: {freq:.2f}% | Count: {count}")
        return
                
    # Analyze mutations and calculate frequencies
    mutation_info = {}
    for mutation, count in sorted(mutation_count.items(), key=lambda x: int(x[0][1:-1])):
        seq_pos = int(mutation[1:-1])
        query_aa = mutation[0]
        subject_aa = mutation[-1]
        
        if query_aa == '-' or subject_aa == '-':
            continue
        
        real_count = aa_counts[seq_pos].get(subject_aa, 0)
        total_count = sum(aa_counts[seq_pos].values())
        query_aa_frequency = (aa_counts[seq_pos].get(query_aa, 0) / total_count * 100) if total_count > 0 else 0
        subject_aa_frequency = (aa_counts[seq_pos].get(subject_aa, 0) / total_count * 100) if total_count > 0 else 0
        
        if subject_aa_frequency <= 10 and real_count > 2:
            mutation_info[mutation] = {
                'query_aa_frequency': query_aa_frequency,
                'subject_aa_frequency': subject_aa_frequency,
                'real_count': real_count,
            }
    
    # Identify strains with specific mutations
    strains_with_mutations = {}
    for mutation in mutation_count:
        query_aa, pos, subject_aa = mutation[0], int(mutation[1:-1]), mutation[-1]
        strains_with_this_mutation = []

        for record in alignment:
            strain_name = record.id[:3] 
            sequence = str(record.seq)
            
            pos_in_ref_seq = 0
            
            for i, aa in enumerate(sequence):
                if ref_seq.seq[i] != '-':
                    pos_in_ref_seq += 1
                
                if pos_in_ref_seq == pos:
                    
                    if aa == subject_aa:
                        strains_with_this_mutation.append(strain_name[:3])
                    break  
        strains_with_mutations[mutation] = strains_with_this_mutation
            
    # Write mutation information to a text file
    with open('gapslist.txt', 'w') as f:
        for mutation, info in mutation_info.items():
            f.write(f"Mutation: {mutation} | {mutation[0]} Frequency: {info['query_aa_frequency']:.2f}% | {mutation[-1]} Frequency: {info['subject_aa_frequency']:.2f}% | Count: {info['real_count']}\n\n")

    # Write mutation and strain information to a CSV file
    with open('rare_mutations.csv', 'w', newline='') as csvfile:
        fieldnames = ['Mutation', 'Strains']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        
        for mutation, info in mutation_info.items():
            writer.writerow({
                'Mutation': mutation,
                'Strains': ';'.join(strains_with_mutations.get(mutation, []))
            })

性能分析结果 #

ncalls  tottime  percall  cumtime  percall filename:lineno(function)

 69006091  25.002  0.000     99.135   0.000   Seq.py:470(__getitem__)

 69098925  6.311   0.000     6.311    0.000   SeqRecord.py:334(<lambda>).

 69008225  10.165  0.000     49.965   0.000   <frozen abc>:117(__instancecheck__)

 69006091  10.330  0.000     22.310   0.000   <frozen abc>:121(__subclasscheck__)

核心瓶颈来自反复调用Seq.__getitem__和重复遍历序列计算位置映射，以下是针对性优化：

1. 预计算比对位置到参考真实位置的映射 #

提前生成数组记录每个比对位置对应的参考序列真实位置，避免每次遍历序列时逐位判断gap：

# 替换原代码中ref_seq_str生成后的部分
ref_seq_str = str(ref_seq.seq)
# 生成比对位置→参考真实位置的映射
align_pos_to_ref_pos = []
current_ref_pos = 0
for aa in ref_seq_str:
    if aa != '-':
        current_ref_pos += 1
    align_pos_to_ref_pos.append(current_ref_pos)
# 生成参考真实位置→比对位置的反向映射（仅保留需要统计的位置）
target_ref_positions = set(mutation_positions)
ref_pos_to_align_positions = {pos: [] for pos in target_ref_positions}
for align_pos, ref_pos in enumerate(align_pos_to_ref_pos):
    if ref_pos in target_ref_positions:
        ref_pos_to_align_positions[ref_pos].append(align_pos)

2. 批量统计氨基酸频率，避免逐序列逐位遍历 #

利用预计算的位置映射，直接针对目标参考位置批量提取氨基酸并统计，结合collections.Counter提升效率：

from collections import Counter

# 替换原代码中的氨基酸统计循环
alignment_str_list = list(alignment_str.values())
for ref_pos, align_positions in ref_pos_to_align_positions.items():
    for align_pos in align_positions:
        # 批量提取所有序列在该比对位置的氨基酸
        aa_list = [seq[align_pos] for seq in alignment_str_list]
        aa_counts[ref_pos] = Counter(aa_list)

3. 提前构建菌株-位置-氨基酸映射，避免重复遍历序列 #

脚本初期就构建所有菌株在目标参考位置的氨基酸映射，后续查找突变菌株直接查询：

# 在生成alignment_str后添加
strain_pos_aa = {}
for strain_id, seq in alignment_str.items():
    strain_short = strain_id[:3]
    strain_pos_aa[strain_short] = {}
    for ref_pos in target_ref_positions:
        # 参考位置对应唯一比对位置
        align_pos = ref_pos_to_align_positions[ref_pos][0]
        strain_pos_aa[strain_short][ref_pos] = seq[align_pos]

# 替换原代码中的菌株查找循环
strains_with_mutations = {}
for mutation in mutation_count:
    query_aa, pos, subject_aa = mutation[0], int(mutation[1:-1]), mutation[-1]
    strains_with_this_mutation = [
        strain for strain, pos_map in strain_pos_aa.items()
        if pos_map.get(pos) == subject_aa
    ]
    strains_with_mutations[mutation] = strains_with_this_mutation

4. 彻底避免Seq对象的__getitem__调用 #

确保所有Seq对象在脚本初期就转换为字符串，后续仅操作字符串，移除所有直接访问record.seq或ref_seq.seq的代码。

from collections import Counter
import csv
from Bio import AlignIO

def main(self):
    # Process input data and extract necessary information
    mutation_count, mutation_subjects, mutation_positions, alignment_file, seq_ID = self.process_input()

    # Initialize a dictionary to count amino acids at each mutation position
    aa_counts = {pos: Counter() for pos in mutation_positions}

    # Read the alignment file using Bio.AlignIO
    alignment = AlignIO.read(alignment_file, 'clustal')
    
    # Find the reference sequence in the alignment
    ref_seq = None
    for record in alignment:
        if record.id == seq_ID:
            ref_seq = record
            break

    # Convert all sequences to strings once and for all
    ref_seq_str = str(ref_seq.seq)
    alignment_str = {record.id: str(record.seq) for record in alignment}
    alignment_str_list = list(alignment_str.values())

    # Precompute alignment position to reference position mapping
    align_pos_to_ref_pos = []
    current_ref_pos = 0
    for aa in ref_seq_str:
        if aa != '-':
            current_ref_pos += 1
        align_pos_to_ref_pos.append(current_ref_pos)

    # Precompute reference position to alignment positions mapping
    target_ref_positions = set(mutation_positions)
    ref_pos_to_align_positions = {pos: [] for pos in target_ref_positions}
    for align_pos, ref_pos in enumerate(align_pos_to_ref_pos):
        if ref_pos in target_ref_positions:
            ref_pos_to_align_positions[ref_pos].append(align_pos)

    # Count amino acids efficiently
    for ref_pos, align_positions in ref_pos_to_align_positions.items():
        for align_pos in align_positions:
            aa_list = [seq[align_pos] for seq in alignment_str_list]
            aa_counts[ref_pos].update(aa_list)

    # If a specific position is provided, calculate and print frequencies
    if self.args.position:
        position = self.args.position
        total_count = sum(aa_counts[position].values())
        for aa, count in aa_counts[position].items():
            freq = (count / total_count * 100) if total_count > 0 else 0
            print(f"Amino Acid: {aa} | Frequency: {freq:.2f}% | Count: {count}")
        return
                
    # Analyze mutations and calculate frequencies
    mutation_info = {}
    for mutation, count in sorted(mutation_count.items(), key=lambda x: int(x[0][1:-1])):
        seq_pos = int(mutation[1:-1])
        query_aa = mutation[0]
        subject_aa = mutation[-1]
        
        if query_aa == '-' or subject_aa == '-':
            continue
        
        total_count = sum(aa_counts[seq_pos].values())
        if total_count == 0:
            continue
        
        query_aa_count = aa_counts[seq_pos].get(query_aa, 0)
        subject_aa_count = aa_counts[seq_pos].get(subject_aa, 0)
        
        query_aa_frequency = (query_aa_count / total_count) * 100
        subject_aa_frequency = (subject_aa_count / total_count) * 100
        
        if subject_aa_frequency <= 10 and subject_aa_count > 2:
            mutation_info[mutation] = {
                'query_aa_frequency': query_aa_frequency,
                'subject_aa_frequency': subject_aa_frequency,
                'real_count': subject_aa_count,
            }
    
    # Prebuild strain to position-amino acid mapping
    strain_pos_aa = {}
    for strain_id, seq in alignment_str.items():
        strain_short = strain_id[:3]
        strain_pos_aa[strain_short] = {}
        for ref_pos in target_ref_positions:
            align_pos = ref_pos_to_align_positions[ref_pos][0]
            strain_pos_aa[strain_short][ref_pos] = seq[align_pos]

    # Identify strains with specific mutations quickly
    strains_with_mutations = {}
    for mutation in mutation_count:
        query_aa, pos, subject_aa = mutation[0], int(mutation[1:-1]), mutation[-1]
        strains_with_this_mutation = [
            strain for strain, pos_map in strain_pos_aa.items()
            if pos_map.get(pos) == subject_aa
        ]
        strains_with_mutations[mutation] = strains_with_this_mutation
            
    # Write mutation information to a text file
    with open('gapslist.txt', 'w') as f:
        for mutation, info in mutation_info.items():
            f.write(f"Mutation: {mutation} | {mutation[0]} Frequency: {info['query_aa_frequency']:.2f}% | {mutation[-1]} Frequency: {info['subject_aa_frequency']:.2f}% | Count: {info['real_count']}\n\n")

    # Write mutation and strain information to a CSV file
    with open('rare_mutations.csv', 'w', newline='') as csvfile:
        fieldnames = ['Mutation', 'Strains']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        
        for mutation, info in mutation_info.items():
            writer.writerow({
                'Mutation': mutation,
                'Strains': ';'.join(strains_with_mutations.get(mutation, []))
            })

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