import torch import torch.nn as nn from typing import List, Optional import numpy as np from sklearn.preprocessing import StandardScaler # class PhysicochemicalEncoder(nn.Module): # """Amino Acid Physicochemical Property Encoder (AAindex版本)""" # def __init__(self, device, use_aaindex=True, selected_features=None): # """ # Args: # device: torch device # use_aaindex: 是否使用AAindex特征（True）还是简单的5特征（False） # selected_features: 选择使用哪些AAindex特征（None=使用全部） # """ # super().__init__() # self.device = device # self.use_aaindex = use_aaindex # if use_aaindex: # # 从AAindex加载特征 # self.aa_properties, self.feature_names = self._load_aaindex_features(selected_features) # self.n_features = len(list(self.aa_properties['A'].values())) # print(f"✓ Loaded {self.n_features} AAindex features") # else: # # 使用简单的5特征 # self.aa_properties = self._get_basic_properties() # self.n_features = 5 # print(f"✓ Using {self.n_features} basic features") # # 标准化（重要！不同特征范围差异大） # self.scaler = self._fit_scaler() # def _load_aaindex_features(self, selected_features=None): # """从AAindex加载特征""" # try: # # 尝试导入生成的文件 # from aa_properties_aaindex import AA_PROPERTIES_AAINDEX, FEATURE_DESCRIPTIONS # if selected_features is not None: # # 只选择指定的特征 # filtered_props = {} # for aa, props in AA_PROPERTIES_AAINDEX.items(): # filtered_props[aa] = {k: v for k, v in props.items() # if k in selected_features} # return filtered_props, selected_features # else: # # 使用所有特征 # feature_names = list(AA_PROPERTIES_AAINDEX['A'].keys()) # return AA_PROPERTIES_AAINDEX, feature_names # except ImportError: # print("⚠ Warning: aa_properties_aaindex.py not found!") # print(" Falling back to basic 5 features") # print(" Run 'python aaindex_downloader.py' to download AAindex features") # return self._get_basic_properties(), ['hydro', 'charge', 'volume', 'flex', 'aroma'] # def _get_basic_properties(self): # """基础的5特征（作为fallback）""" # return { # 'A': [1.8, 0.0, 88.6, 0.36, 0.0], # 'C': [2.5, 0.0, 108.5, 0.35, 0.0], # 'D': [-3.5, -1.0, 111.1, 0.51, 0.0], # 'E': [-3.5, -1.0, 138.4, 0.50, 0.0], # 'F': [2.8, 0.0, 189.9, 0.31, 1.0], # 'G': [-0.4, 0.0, 60.1, 0.54, 0.0], # 'H': [-3.2, 0.5, 153.2, 0.32, 0.5], # 'I': [4.5, 0.0, 166.7, 0.46, 0.0], # 'K': [-3.9, 1.0, 168.6, 0.47, 0.0], # 'L': [3.8, 0.0, 166.7, 0.37, 0.0], # 'M': [1.9, 0.0, 162.9, 0.30, 0.0], # 'N': [-3.5, 0.0, 114.1, 0.46, 0.0], # 'P': [-1.6, 0.0, 112.7, 0.51, 0.0], # 'Q': [-3.5, 0.0, 143.8, 0.49, 0.0], # 'R': [-4.5, 1.0, 173.4, 0.53, 0.0], # 'S': [-0.8, 0.0, 89.0, 0.51, 0.0], # 'T': [-0.7, 0.0, 116.1, 0.44, 0.0], # 'V': [4.2, 0.0, 140.0, 0.39, 0.0], # 'W': [-0.9, 0.0, 227.8, 0.31, 1.0], # 'Y': [-1.3, 0.0, 193.6, 0.42, 1.0], # 'X': [0.0, 0.0, 120.0, 0.40, 0.0], # } # def _fit_scaler(self): # """拟合标准化器""" # # 收集所有氨基酸的特征 # all_features = [] # for aa in 'ARNDCQEGHILKMFPSTWYV': # 20种标准氨基酸 # if isinstance(self.aa_properties[aa], dict): # # AAindex格式 # features = list(self.aa_properties[aa].values()) # else: # # 列表格式 # features = self.aa_properties[aa] # all_features.append(features) # all_features = np.array(all_features) # # Z-score标准化 # scaler = StandardScaler() # scaler.fit(all_features) # return scaler # def _get_aa_features(self, aa: str) -> List[float]: # """获取单个氨基酸的特征""" # aa = aa.upper() # if aa not in self.aa_properties: # aa = 'X' # Unknown # if isinstance(self.aa_properties[aa], dict): # # AAindex格式：字典 # features = list(self.aa_properties[aa].values()) # else: # # 基础格式：列表 # features = self.aa_properties[aa] # return features # def forward(self, sequences: List[str]) -> torch.Tensor: # """ # Args: # sequences: List of amino acid sequences # Returns: # [B, max_len, n_features] 标准化后的特征 # """ # batch_size = len(sequences) # max_len = max(len(seq) for seq in sequences) # # 收集特征 # properties = [] # for seq in sequences: # seq_props = [] # for aa in seq: # props = self._get_aa_features(aa) # seq_props.append(props) # # Padding # while len(seq_props) < max_len: # seq_props.append([0.0] * self.n_features) # properties.append(seq_props) # properties = np.array(properties) # [B, L, n_features] # # 标准化（除了padding位置） # batch_size, seq_len, n_feat = properties.shape # properties_flat = properties.reshape(-1, n_feat) # # 标准化 # properties_normalized = self.scaler.transform(properties_flat) # properties_normalized = properties_normalized.reshape(batch_size, seq_len, n_feat) # # 转为tensor # properties_tensor = torch.tensor( # properties_normalized, # dtype=torch.float32, # device=self.device # ) # return properties_tensor # [B, L, n_features] import torch import torch.nn as nn import numpy as np from sklearn.preprocessing import StandardScaler from typing import List class PhysicochemicalEncoder(nn.Module): """Amino Acid Physicochemical Property Encoder (AAindex版本, 向量化优化版)""" def __init__(self, device, use_aaindex=True, selected_features=None): super().__init__() self.device = device self.use_aaindex = use_aaindex # 加载特征 if use_aaindex: self.aa_properties, self.feature_names = self._load_aaindex_features(selected_features) self.n_features = len(list(self.aa_properties['A'].values())) print(f"✓ Loaded {self.n_features} AAindex features") else: self.aa_properties = self._get_basic_properties() self.n_features = 5 print(f"✓ Using {self.n_features} basic features") # 拟合标准化器 self.scaler = self._fit_scaler() # ======================== 🔥 预处理部分 ======================== # # 1. 构建 lookup table aa_list = list(self.aa_properties.keys()) aa_list.sort() # 保证稳定顺序 self.aa_to_idx = {aa: i for i, aa in enumerate(aa_list)} self.pad_idx = len(self.aa_to_idx) # padding index aa_feature_table = [] for aa in aa_list: feats = self._get_aa_features(aa) aa_feature_table.append(feats) aa_feature_table.append([0.0] * self.n_features) # padding vector self.aa_feature_table = torch.tensor( np.array(aa_feature_table), dtype=torch.float32 ).to(self.device) # [n_aa+1, n_feat] # 2. 标准化参数预存成 GPU tensor self.mean_tensor = torch.tensor(self.scaler.mean_, dtype=torch.float32, device=self.device) self.scale_tensor = torch.tensor(self.scaler.scale_, dtype=torch.float32, device=self.device) # 下面这些函数和你原来的完全一致，不动 def _load_aaindex_features(self, selected_features=None): try: from aa_properties_aaindex import AA_PROPERTIES_AAINDEX, FEATURE_DESCRIPTIONS if selected_features is not None: filtered_props = {} for aa, props in AA_PROPERTIES_AAINDEX.items(): filtered_props[aa] = {k: v for k, v in props.items() if k in selected_features} return filtered_props, selected_features else: feature_names = list(AA_PROPERTIES_AAINDEX['A'].keys()) return AA_PROPERTIES_AAINDEX, feature_names except ImportError: print("⚠ Warning: aa_properties_aaindex.py not found!") return self._get_basic_properties(), ['hydro', 'charge', 'volume', 'flex', 'aroma'] def _get_basic_properties(self): # 这里同你原来的 return { 'A': [1.8, 0.0, 88.6, 0.36, 0.0], 'C': [2.5, 0.0, 108.5, 0.35, 0.0], 'D': [-3.5, -1.0, 111.1, 0.51, 0.0], 'E': [-3.5, -1.0, 138.4, 0.50, 0.0], 'F': [2.8, 0.0, 189.9, 0.31, 1.0], 'G': [-0.4, 0.0, 60.1, 0.54, 0.0], 'H': [-3.2, 0.5, 153.2, 0.32, 0.5], 'I': [4.5, 0.0, 166.7, 0.46, 0.0], 'K': [-3.9, 1.0, 168.6, 0.47, 0.0], 'L': [3.8, 0.0, 166.7, 0.37, 0.0], 'M': [1.9, 0.0, 162.9, 0.30, 0.0], 'N': [-3.5, 0.0, 114.1, 0.46, 0.0], 'P': [-1.6, 0.0, 112.7, 0.51, 0.0], 'Q': [-3.5, 0.0, 143.8, 0.49, 0.0], 'R': [-4.5, 1.0, 173.4, 0.53, 0.0], 'S': [-0.8, 0.0, 89.0, 0.51, 0.0], 'T': [-0.7, 0.0, 116.1, 0.44, 0.0], 'V': [4.2, 0.0, 140.0, 0.39, 0.0], 'W': [-0.9, 0.0, 227.8, 0.31, 1.0], 'Y': [-1.3, 0.0, 193.6, 0.42, 1.0], 'X': [0.0, 0.0, 120.0, 0.40, 0.0], } def _fit_scaler(self): all_features = [] for aa in 'ARNDCQEGHILKMFPSTWYV': if isinstance(self.aa_properties[aa], dict): features = list(self.aa_properties[aa].values()) else: features = self.aa_properties[aa] all_features.append(features) all_features = np.array(all_features) scaler = StandardScaler() scaler.fit(all_features) return scaler def _get_aa_features(self, aa: str): aa = aa.upper() if aa not in self.aa_properties: aa = 'X' if isinstance(self.aa_properties[aa], dict): return list(self.aa_properties[aa].values()) else: return self.aa_properties[aa] def forward(self, sequences: List[str]) -> torch.Tensor: batch_size = len(sequences) max_len = max(len(seq) for seq in sequences) # 1) encode sequences to indices with padding idx_batch = np.full((batch_size, max_len), self.pad_idx, dtype=np.int64) for i, seq in enumerate(sequences): idx_seq = [self.aa_to_idx.get(aa.upper(), self.pad_idx) for aa in seq] idx_batch[i, :len(idx_seq)] = idx_seq idx_tensor = torch.tensor(idx_batch, dtype=torch.long, device=self.device) # [B, L] # 2) lookup properties props = self.aa_feature_table[idx_tensor] # [B, L, n_feat] props = (props - self.mean_tensor) / self.scale_tensor return props