feature_extract.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import numpy as np
import random
import pandas as pd
from Bio.SeqUtils.ProtParam import ProteinAnalysis
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import RobustScaler # 导入 RobustScaler
import torch
from transformers import T5EncoderModel, T5Tokenizer 



class ProtT5Model:
    """
    从本地加载 ProtT5 模型。如果 finetuned_model_file 不为空，则加载微调后的权重（使用 strict=False）。
    """
    def __init__(self, model_path, finetuned_model_file=None):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        # 尝试加载本地文件，如果失败，transformers库可能会尝试从hub下载（取决于配置）
        try:
            self.tokenizer = T5Tokenizer.from_pretrained(model_path, do_lower_case=False, local_files_only=True)
            self.model = T5EncoderModel.from_pretrained(model_path, local_files_only=True)
        except OSError: # OSError: Can't load tokenizer for '...'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure '...' is the correct path to a directory containing all relevant files for a T5Tokenizer tokenizer.
            print(f"警告: 无法从本地路径 {model_path} 加载ProtT5模型/分词器。尝试从HuggingFace Hub下载（如果transformers配置允许）。")
            self.tokenizer = T5Tokenizer.from_pretrained(model_path.split('/')[-1] if '/' in model_path else model_path, do_lower_case=False) # 尝试使用模型名下载
            self.model = T5EncoderModel.from_pretrained(model_path.split('/')[-1] if '/' in model_path else model_path)


        if finetuned_model_file is not None and os.path.exists(finetuned_model_file):
            try:
                state_dict = torch.load(finetuned_model_file, map_location=self.device)
                missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False)
                print(f"加载微调权重 {finetuned_model_file}：缺失键 {missing_keys}, 意外键 {unexpected_keys}")
            except Exception as e:
                print(f"加载微调权重 {finetuned_model_file} 失败: {e}")
        
        self.model.to(self.device)
        self.model.eval()

    def encode(self, sequence):
        if not sequence or not isinstance(sequence, str): # 增加对空序列或非字符串的检查
            print(f"警告: ProtT5Model.encode 接收到无效序列: {sequence}")
            # 返回一个零向量或根据需要处理错误
            # 假设 ProtT5 输出维度为 1024 (embedding.shape[1])
            # 假设序列处理后平均池化，所以返回 (1024,)
            # 但 encode 返回的是 (seq_len, hidden_dim)，所以这里返回一个模拟的短序列零嵌入
            return np.zeros((1, 1024), dtype=np.float32) # (1, hidden_dim)

        seq_spaced = " ".join(list(sequence)) # 修改变量名以避免覆盖外部seq
        try:
            encoded_input = self.tokenizer(seq_spaced, return_tensors='pt', padding=True, truncation=True, max_length=1022) # ProtT5通常最大长度1024，tokenized后可能更长
        except Exception as e:
            print(f"分词失败序列 '{sequence[:30]}...': {e}")
            return np.zeros((1, 1024), dtype=np.float32)

        encoded_input = {k: v.to(self.device) for k, v in encoded_input.items()}
        with torch.no_grad():
            try:
                embedding = self.model(**encoded_input).last_hidden_state  # (batch_size, seq_len, hidden_dim)
            except Exception as e:
                print(f"ProtT5模型推理失败序列 '{sequence[:30]}...': {e}")
                return np.zeros((1, 1024), dtype=np.float32)

        emb = embedding.squeeze(0).cpu().numpy()  # (seq_len, hidden_dim)
        if emb.shape[0] == 0: # 如果由于某种原因序列长度为0
             return np.zeros((1, 1024), dtype=np.float32)
        return emb


def load_fasta(fasta_file):
    # (您的 load_fasta 实现)
    sequences = []
    try:
        with open(fasta_file, 'r') as f:
            current_seq_lines = []
            for line in f:
                line = line.strip()
                if not line: continue
                if line.startswith(">"):
                    if current_seq_lines: sequences.append("".join(current_seq_lines))
                    current_seq_lines = []
                else: current_seq_lines.append(line)
            if current_seq_lines: sequences.append("".join(current_seq_lines))
    except FileNotFoundError: print(f"文件未找到: {fasta_file}"); return []
    return sequences

def load_fasta_with_labels(fasta_file):
    sequences, labels = [], []
    try:
        with open(fasta_file, 'r') as f:
            current_seq_lines, current_label = [], None
            for line in f:
                line = line.strip()
                if not line: continue
                if line.startswith(">"):
                    if current_seq_lines:
                        sequences.append("".join(current_seq_lines))
                        labels.append(current_label if current_label is not None else 0) # Default label 0
                    current_seq_lines = []
                    current_label = int(line[1]) if len(line) > 1 and line[1] in ['0', '1'] else 0
                else: current_seq_lines.append(line)
            if current_seq_lines:
                sequences.append("".join(current_seq_lines))
                labels.append(current_label if current_label is not None else 0)
    except FileNotFoundError: print(f"文件未找到: {fasta_file}"); return [],[]
    return sequences, labels


def compute_amino_acid_composition(seq):
    if not seq: return {aa: 0.0 for aa in "ACDEFGHIKLMNPQRSTVWY"}
    amino_acids = "ACDEFGHIKLMNPQRSTVWY"
    seq_len = len(seq)
    return {aa: seq.upper().count(aa) / seq_len for aa in amino_acids}


def compute_reducing_aa_ratio(seq):
    if not seq: return 0.0
    reducing = ['C', 'M', 'W']
    return sum(seq.upper().count(aa) for aa in reducing) / len(seq) if len(seq) > 0 else 0.0

def compute_physicochemical_properties(seq):
    if not seq or not all(c.upper() in "ACDEFGHIKLMNPQRSTVWYXUBZ" for c in seq): # ProteinAnalysis might fail on invalid chars
        return 0.0, 0.0, 0.0 # Default values
    try:
        analysis = ProteinAnalysis(str(seq).upper().replace('X','A').replace('U','C').replace('B','N').replace('Z','Q')) # Replace non-standard with common ones for analysis
        return analysis.gravy(), analysis.isoelectric_point(), analysis.molecular_weight()
    except Exception: # Catch any error from ProteinAnalysis
        return 0.0, 7.0, 110.0 * len(seq) # Rough defaults

def compute_electronic_features(seq):
    if not seq: return 0.0, 0.0
    electronegativity = {'A':1.8,'C':2.5,'D':3.0,'E':3.2,'F':2.8,'G':1.6,'H':2.4,'I':4.5,'K':3.0,'L':4.2,'M':4.5,'N':2.0,'P':3.5,'Q':3.5,'R':2.5,'S':1.8,'T':2.5,'V':4.0,'W':5.0,'Y':4.0}
    values = [electronegativity.get(aa.upper(), 2.5) for aa in seq]
    avg_val = sum(values) / len(values) if values else 2.5
    return avg_val + 0.1, avg_val - 0.1


def compute_dimer_frequency(seq):
    if len(seq) < 2: return np.zeros(400) # 20*20
    amino_acids = "ACDEFGHIKLMNPQRSTVWY"
    dimer_counts = {aa1+aa2: 0 for aa1 in amino_acids for aa2 in amino_acids}
    for i in range(len(seq) - 1):
        dimer = seq[i:i+2].upper()
        if dimer in dimer_counts: dimer_counts[dimer] += 1
    total = max(len(seq) - 1, 1)
    for key in dimer_counts: dimer_counts[key] /= total
    return np.array([dimer_counts[d] for d in sorted(dimer_counts.keys())])


def positional_encoding(seq_len_actual, L_fixed=29, d_model=16): 
    pos_enc = np.zeros((L_fixed, d_model))
    for pos in range(L_fixed):
        for i in range(d_model):
            angle = pos / (10000 ** (2 * (i // 2) / d_model))
            pos_enc[pos, i] = np.sin(angle) if i % 2 == 0 else np.cos(angle)
    return pos_enc.flatten()


def perturb_sequence(seq, perturb_rate=0.1, critical=['C', 'M', 'W']):
    # (您的 perturb_sequence 实现)
    if not seq: return ""
    seq_list = list(seq)
    amino_acids = "ACDEFGHIKLMNPQRSTVWY"
    for i, aa in enumerate(seq_list):
        if aa.upper() not in critical and random.random() < perturb_rate:
            seq_list[i] = random.choice([x for x in amino_acids if x != aa.upper()])
    return "".join(seq_list)


def extract_features(seq, prott5_model_instance, L_fixed=29, d_model_pe=16): # Renamed d_model to d_model_pe
    if not seq or not isinstance(seq, str) or len(seq) == 0:
        print(f"警告: extract_features 接收到空或无效序列。返回零特征。")
        return np.zeros(1024 + 20 + 1 + 3 + 2 + 400 + (L_fixed * d_model_pe))


    embedding = prott5_model_instance.encode(seq) # prott5_model is now an instance
    prot_embed = np.mean(embedding, axis=0) if embedding.shape[0] > 0 else np.zeros(embedding.shape[1] if embedding.ndim > 1 else 1024) # Handle empty embedding
    if prot_embed.shape[0] != 1024: # Ensure consistent ProtT5 embedding dim
        # print(f"警告: ProtT5 嵌入维度异常 ({prot_embed.shape[0]}) for seq '{seq[:20]}...'. 使用零向量。")
        prot_embed = np.zeros(1024)


    aa_comp = compute_amino_acid_composition(seq)
    aa_comp_vector = np.array([aa_comp[aa] for aa in "ACDEFGHIKLMNPQRSTVWY"])
    red_ratio = np.array([compute_reducing_aa_ratio(seq)])
    gravy, pI, mol_weight = compute_physicochemical_properties(seq)
    phys_props = np.array([gravy, pI, mol_weight])
    HOMO, LUMO = compute_electronic_features(seq)
    electronic = np.array([HOMO, LUMO])
    dimer_vector = compute_dimer_frequency(seq)
    pos_enc = positional_encoding(len(seq), L_fixed, d_model_pe) # Pass actual length, though current PE uses L_fixed
    
    features = np.concatenate([prot_embed, aa_comp_vector, red_ratio, phys_props, electronic, dimer_vector, pos_enc])
    return features

##############################################
# 主接口函数 prepare_features
##############################################
def prepare_features(neg_fasta, pos_fasta, prott5_model_path, additional_params=None):
    neg_seqs = load_fasta(neg_fasta)
    pos_seqs = load_fasta(pos_fasta)
    
    if not neg_seqs and not pos_seqs:
        raise ValueError("未能从FASTA文件加载任何序列。请检查文件路径和内容。")

    neg_labels = [0] * len(neg_seqs)
    pos_labels = [1] * len(pos_seqs)
    sequences = neg_seqs + pos_seqs
    labels = neg_labels + pos_labels

    combined = list(zip(sequences, labels))
    random.shuffle(combined)
    sequences, labels = zip(*combined)
    sequences = list(sequences)
    labels = list(labels)

    train_seqs, val_seqs, train_labels, val_labels = train_test_split(
        sequences, labels, test_size=0.1, random_state=42, stratify=labels if len(np.unique(labels)) > 1 else None
    )
    print("训练集原始样本数：", len(train_seqs))
    print("验证集原始样本数：", len(val_seqs))

    if additional_params is not None and additional_params.get("augment", False):
        # (数据增强逻辑 - 如果启用)
        augmented_seqs, augmented_labels = [], []
        perturb_rate = additional_params.get("perturb_rate", 0.1)
        for seq, label in zip(train_seqs, train_labels):
            aug_seq = perturb_sequence(seq, perturb_rate=perturb_rate)
            augmented_seqs.append(aug_seq)
            augmented_labels.append(label)
        train_seqs.extend(augmented_seqs)
        train_labels.extend(augmented_labels)
        print("数据增强后训练集样本数：", len(train_seqs))


    finetuned_model_file = additional_params.get("finetuned_model_file") if additional_params else None
    # 创建 ProtT5Model 实例
    prott5_model_instance = ProtT5Model(prott5_model_path, finetuned_model_file=finetuned_model_file)

    def process_data(seqs_list): # Renamed seqs to seqs_list
        feature_list = []
        for s_item in seqs_list: # Renamed s to s_item
            # 将 ProtT5Model 实例传递给 extract_features
            features = extract_features(s_item, prott5_model_instance) 
            feature_list.append(features)
        return np.array(feature_list)

    X_train = process_data(train_seqs)
    X_val = process_data(val_seqs)
    
    if X_train.shape[0] == 0 or X_val.shape[0] == 0:
        raise ValueError("特征提取后训练集或验证集为空。请检查序列数据和特征提取过程。")


    # --- **关键修改：使用 RobustScaler** ---
    scaler = RobustScaler() 
    print("使用 RobustScaler 进行特征归一化。")
    
    X_train_scaled = scaler.fit_transform(X_train)
    X_val_scaled = scaler.transform(X_val)
    
    return X_train_scaled, X_val_scaled, np.array(train_labels), np.array(val_labels), scaler

if __name__ == "__main__":
    # 确保测试时使用的路径是有效的，或者创建虚拟文件
    neg_fasta_test = "dummy_data/test_neg.fasta"
    pos_fasta_test = "dummy_data/test_pos.fasta"
    prott5_path_test = "dummy_prott5_model/" # 需要一个包含config.json, pytorch_model.bin等的目录结构
    
    os.makedirs("dummy_data", exist_ok=True)
    os.makedirs(prott5_path_test, exist_ok=True) # 创建虚拟模型目录

    if not os.path.exists(neg_fasta_test):
        with open(neg_fasta_test, "w") as f: f.write(">neg1\nKALKALKALK\n>neg2\nPEPTPEPT\n")
    if not os.path.exists(pos_fasta_test):
        with open(pos_fasta_test, "w") as f: f.write(">pos1\nAOPPEPTIDE\n>pos2\nTRYTRYTRY\n")
    
    if not os.listdir(prott5_path_test): # 如果目录为空
        print(f"警告: {prott5_path_test} 为空。ProtT5Model可能尝试从HuggingFace Hub下载模型。")
        print(f"请确保您已下载Rostlab/ProstT5-XL-UniRef50或类似模型到该路径，或使用其HuggingFace名称。")
        # 作为演示，我们假设用户会提供一个有效的路径或transformers可以处理它
        # 如果要完全本地运行而不下载，需要填充该目录。

    additional_params_test = {
        "augment": False, 
        "perturb_rate": 0.1, 
        "finetuned_model_file": None # 测试时不使用微调模型
    }
    
    print("开始测试 prepare_features (使用RobustScaler)...")
    try:
        X_train_t, X_val_t, y_train_t, y_val_t, scaler_t = prepare_features(
            neg_fasta_test, 
            pos_fasta_test, 
            "Rostlab/ProstT5-XL-UniRef50", # 使用HuggingFace模型名称，如果本地路径无效
            additional_params_test
        )
        print("prepare_features 测试完成。")
        print("训练集样本数：", X_train_t.shape[0])
        print("验证集样本数：", X_val_t.shape[0])
        if X_train_t.shape[0] > 0:
            print("训练集特征维度:", X_train_t.shape[1])
            print("一个缩放后的训练样本 (前5个特征):", X_train_t[0, :5])
        if scaler_t:
            print("Scaler类型:", type(scaler_t))
    except Exception as e:
        print(f"prepare_features 测试失败: {e}")
        print("这可能是由于无法加载ProtT5模型或FASTA文件处理问题。请检查路径和文件内容。")