from rdkit import Chem from rdkit.Chem import AllChem, MACCSkeys from rdkit.Chem.rdmolops import FastFindRings from rdkit.Chem.rdMolDescriptors import CalcMolFormula import torch import numpy as np import scipy import scipy.sparse as ss import scipy.sparse.linalg import math import json import itertools as it import re from GNN import featurizer as ft import rdkit.RDLogger as rkl logger = rkl.logger() logger.setLevel(rkl.ERROR) import rdkit.rdBase as rkrb rkrb.DisableLog('rdApp.error') # 50w metabolites fpbit relative aboundance > 5% FPBitIdx = [1, 5, 13, 41, 69, 80, 84, 94, 114, 117, 118, 119, 125, 133, 145, 147, 191, 192, 197, 202, 222, 227, 231, 249, 283, 294, 310, 314, 322, 333, 352, 361, 378, 387, 389, 392, 401, 406, 441, 478, 486, 489, 519, 521, 524, 555, 561, 591, 598, 599, 610, 622, 650, 656, 667, 675, 677, 679, 680, 694, 695, 715, 718, 722, 729, 736, 739, 745, 750, 760, 775, 781, 787, 794, 798, 802, 807, 811, 823, 835, 841, 849, 869, 872, 874, 875, 881, 890, 896, 926, 935, 980, 991, 1004, 1009, 1017, 1019, 1027, 1028, 1035, 1037, 1039, 1057, 1060, 1066, 1070, 1077, 1088, 1097, 1114, 1126, 1136, 1142, 1143, 1145, 1152, 1154, 1160, 1162, 1171, 1181, 1195, 1199, 1202, 1218, 1234, 1236, 1243, 1257, 1267, 1274, 1279, 1283, 1292, 1294, 1309, 1313, 1323, 1325, 1349, 1356, 1357, 1366, 1380, 1381, 1385, 1386, 1391, 1399, 1436, 1440, 1441, 1444, 1452, 1454, 1457, 1475, 1476, 1477, 1480, 1487, 1516, 1536, 1544, 1558, 1564, 1573, 1599, 1602, 1604, 1607, 1619, 1648, 1670, 1683, 1693, 1716, 1722, 1737, 1738, 1745, 1747, 1750, 1754, 1755, 1764, 1781, 1803, 1808, 1810, 1816, 1838, 1844, 1847, 1855, 1860, 1866, 1873, 1905, 1911, 1917, 1921, 1923, 1928, 1933, 1950, 1951, 1970, 1977, 1980, 1984, 1991, 2002, 2033, 2034, 2038] class ConfigDict(dict): ''' Makes a dictionary behave like an object,with attribute-style access. ''' def __getattr__(self, name): try: return self[name] except: raise AttributeError(name) def __setattr__(self, name, value): self[name] = value def save(self, fn): json.dump(self, open(fn, 'w'), indent=2) def load_dict(self, dic): for k, v in dic.items(): self[k] = v def load(self, fn): try: d = json.load(open(fn, 'r')) self.load_dict(d) except Exception as e: print(e) def conv_out_dim(length_in, kernel, stride, padding, dilation): length_out = (length_in + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1 return length_out def filter_ms(ms, thr=0.05, max_mz=2000): mz = [] intn = [] maxi = 0 for m, i in ms: if m < max_mz and i > maxi: maxi = i for m, i in ms: if m < max_mz and i/maxi > thr: mz.append(m) intn.append(round(i/maxi*100, 2)) return mz, intn def calc_nls(ms, thr=0.05, max_mz=2000): mz, intn = filter_ms(ms, thr=0.05, max_mz=2000) nlmass = [] nlintn = [] for a, b in it.combinations(mz[::-1], 2): nl = a - b if 0 < nl < 200: nlmass.append(round(nl, 5)) idxa = mz.index(a) idxb = mz.index(b) nlintn.append(round((intn[idxa]+intn[idxb])/2., 5)) nls = sorted(list(zip(nlmass, nlintn))) return nls # --- 2. 辅助函数：匹配诊断离子与中性丢失 --- def check_match(val, targets, tolerance=0.02): """判断 val 是否在 targets 列表中 (带容差)""" if not targets: return 0.0 val_arr = np.array([val]) target_arr = np.array(targets) # 广播计算差值绝对值 diff = np.abs(val_arr.reshape(-1, 1) - target_arr.reshape(1, -1)) # 如果有任意一个差值小于容差，返回 1.0 match = np.any(diff <= tolerance) return 1.0 if match else 0.0 # --- 3. 核心处理函数 (对应图中的整个流程) --- def ms_feature_processor(ms, precursor_mz, metadata_vec=None, max_peaks=100, diagnostic_ions=[102.05, 135.08], # 图片示例：生物碱 neutral_losses=[18.01], # 图片示例：水、羟基 max_mz=2000): """ 输入: ms: list of [mz, intensity] precursor_mz: 前体离子 m/z (用于计算中性丢失) metadata_vec: (25,) 维度的元数据向量 (仪器/加合物/电荷) 输出: node_features: (max_peaks, feature_dim) - 这里的 feature_dim 不包含 m/z嵌入，m/z嵌入通常在模型 forward 中做 mz_values: (max_peaks,) - 用于输入给 SinusoidalMzEmbedding """ # 1. [MS数据采集过滤] - 过滤无效数据 valid_ms = [] for m, i in ms: if m <= max_mz and i > 0: valid_ms.append([m, i]) if not valid_ms: # 如果为空，返回零填充 return torch.zeros(max_peaks, 29), torch.zeros(max_peaks) valid_ms = np.array(valid_ms) # 2. [保留强度 T100 峰] # 按强度降序排序 sort_idx = np.argsort(valid_ms[:, 1])[::-1] top_k_idx = sort_idx[:max_peaks] # 截取 Top K top_ms = valid_ms[top_k_idx] # 为了 Transformer 处理方便，通常按 m/z 重新升序排列 (虽然 Transformer 有位置编码，但有序输入有助于学习) resort_idx = np.argsort(top_ms[:, 0]) top_ms = top_ms[resort_idx] # 解包 mz_vals = top_ms[:, 0] int_vals = top_ms[:, 1] # 归一化强度 (0-1) max_int = int_vals.max() if int_vals.max() > 0 else 1.0 norm_int = int_vals / max_int # --- 特征构建 --- feature_list = [] # 处理元数据 (图片要求：元数据编码25维) if metadata_vec is None: metadata_vec = np.zeros(25) # 默认零向量 else: # 确保是 numpy 且长度正确，这里做简单的截断或填充 metadata_vec = np.array(metadata_vec) if len(metadata_vec) > 25: metadata_vec = metadata_vec[:25] elif len(metadata_vec) < 25: metadata_vec = np.pad(metadata_vec, (0, 25 - len(metadata_vec))) for i in range(len(mz_vals)): m = mz_vals[i] inten = norm_int[i] # 3. [诊断离子与中性丢失匹配] # 诊断标记 (1维) is_diagnostic = check_match(m, diagnostic_ions) # 中性丢失标记 (1维) - 检查 (Precursor - Fragment) 是否在列表中 current_nl = precursor_mz - m is_neutral_loss = check_match(current_nl, neutral_losses) if current_nl > 0 else 0.0 # 前体 m/z 加权 (图片提及 "前体m/z加权") # 这里实现一个简单的注意力加权逻辑：如果 fragment 接近 precursor，权重高 # 或者根据图片意图，可能是指 Precursor m/z 作为一个单独的特征值拼进去 # 这里我们假设它是一个特征维度 precursor_weight = abs(m - precursor_mz) / (precursor_mz + 1e-5) # 4. [特征拼接] # 注意：256维的m/z嵌入通常在 GPU 上通过 nn.Module 计算，这里只准备 inputs # 此时的特征: [强度(1), 诊断(1), 丢失(1), 前体权重(1), 元数据(25)] = 29 dims feat = np.concatenate([ [inten], [is_diagnostic], [is_neutral_loss], [precursor_weight], metadata_vec ]) feature_list.append(feat) # Pad 到 max_peaks (如果不足 100 个峰) num_actual = len(feature_list) pad_len = max_peaks - num_actual features_tensor = torch.FloatTensor(np.array(feature_list)) mz_tensor = torch.FloatTensor(mz_vals) if pad_len > 0: # Padding features with 0 feat_pad = torch.zeros(pad_len, 29) # 29 = 1+1+1+1+25 features_tensor = torch.cat([features_tensor, feat_pad], dim=0) # Padding m/z with 0 (or padding value) mz_pad = torch.zeros(pad_len) mz_tensor = torch.cat([mz_tensor, mz_pad], dim=0) return features_tensor, mz_tensor def ms_binner(ms, nls=[], min_mz=20, max_mz=2000, bin_size=0.05, add_nl=False, binary_intn=False): """ Convert the given spectrum to a binned sparse SciPy vector. Parameters ---------- spectrum_mz : np.ndarray The peak m/z values of the spectrum to be converted to a vector. spectrum_intensity : np.ndarray The peak intensities of the spectrum to be converted to a vector. min_mz : float The minimum m/z to include in the vector. bin_size : float The bin size in m/z used to divide the m/z range. num_bins : int The number of elements of which the vector consists. Returns ------- ss.csr_matrix The binned spectrum vector. """ if add_nl and not nls: nls = calc_nls(ms, max_mz=max_mz) nltensor = None mz, intn = filter_ms(ms) if add_nl: nlmass = [] nlintn = [] if not nls: nls = calc_nls(ms, max_mz=max_mz) for m, i in nls: if m < 200: if binary_intn: i = 1 nlmass.append(m) nlintn.append(i) nlmass = np.array(nlmass) nlintn = np.array(nlintn) if len(nlintn) > 0: nlintn = nlintn/nlintn.max() num_nlbins = math.ceil((200) / bin_size) # print('num_nlbins', num_nlbins) nlbins = (nlmass / bin_size).astype(np.int32) if len(nlmass) > 0: vecnl = ss.csr_matrix( (nlintn, (np.repeat(0, len(nlintn)), nlbins)), shape=(1, num_nlbins), dtype=np.float32) vecnl = (vecnl / scipy.sparse.linalg.norm(vecnl)*100) nltensor = torch.FloatTensor(vecnl.todense()).view(-1) else: nltensor = torch.zeros(num_nlbins) mz = np.array(mz) keepidx = (mz <= max_mz) mz = mz[keepidx] intn = np.array(intn) intn = intn[keepidx] if binary_intn: intn[intn > 0] = 1.0 elif len(intn) > 0: intn = intn/intn.max() num_bins = math.ceil((max_mz - min_mz) / bin_size) # print('num_bins', num_bins) bins = ((mz - min_mz) / bin_size).astype(np.int32) # print(num_bins, intn, bins) if len(mz) > 0: vec = ss.csr_matrix( (intn, (np.repeat(0, len(intn)), bins)), shape=(1, num_bins), dtype=np.float32) if not binary_intn: vec = (vec / scipy.sparse.linalg.norm(vec)*100) mstensor = torch.FloatTensor(vec.todense()).view(-1) else: mstensor = torch.zeros(num_bins) if not nltensor is None: return torch.cat([nltensor, mstensor], dim=0) return mstensor def formula2vec(formula, elements=['C', 'H', 'O', 'N', 'P', 'S', 'P', 'F', 'Cl', 'Br']): formula_p = re.findall(r'([A-Z][a-z]*)(\d*)', formula) vec = np.zeros(len(elements)) for i in range(len(formula_p)): ele = formula_p[i][0] num = formula_p[i][1] if num == '': num = 1 else: num = int(num) if ele in elements: vec[elements.index(ele)] += num return np.array(vec) def mol_fp_encoder0(smiles, tp='rdkit', nbits=2048): mol = Chem.MolFromSmiles(smiles) if mol is None: mol = Chem.MolFromSmiles(smiles, sanitize=False) if not mol is None: mol.UpdatePropertyCache() FastFindRings(mol) if mol is None: return None, None if tp == 'morgan': fp_vec = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=nbits) fp = np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0') fp = fp.tolist() elif tp == 'morgan1': fp_vec = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048) fp = np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0') fp = fp[FPBitIdx].tolist() elif tp == 'macc': # MACCSkeys fp_vec = MACCSkeys.GenMACCSKeys(mol) fp = np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0') fp = fp.tolist() elif tp == 'rdkit': fp_vec = Chem.RDKFingerprint(mol, nBitsPerHash=1) fp = np.frombuffer(fp_vec.ToBitString().encode(), 'u1') - ord('0') fp = fp.tolist() return torch.FloatTensor(fp), mol def mol_fp_encoder(smiles, tp='rdkit', nbits=2048): fpenc, _ = mol_fp_encoder0(smiles, tp, nbits) return fpenc def mol_fp_fm_encoder(smiles, tp='rdkit', nbits=2048): fmenc = None fpenc, mol = mol_fp_encoder0(smiles, tp, nbits) if not mol is None: fm = CalcMolFormula(mol) fmenc = torch.FloatTensor(formula2vec(fm)) return fpenc, fmenc def smi2fmvec(smiles): mol = Chem.MolFromSmiles(smiles) if mol is None: return None fm = CalcMolFormula(mol) fmenc = torch.FloatTensor(formula2vec(fm)) return fmenc def mol_graph_featurizer(smiles): # mol_graph = {V, A, mol_size} '''mol_graph = ft.calc_data_from_smile(smiles, addh=True, with_ring_conj=True, with_atom_feats=True, with_submol_fp=True, radius=2) ''' mol_graph = ft.calc_data_from_smile(smiles, addh=False, with_ring_conj=True, with_atom_feats=True, with_submol_fp=False, radius=2) return mol_graph def pad_V(V, max_n): N, C = V.shape if max_n > N: zeros = torch.zeros(max_n-N, C) V = torch.cat([V, zeros], dim=0) return V def pad_A(A, max_n): N, L, _ = A.shape if max_n > N: zeros = torch.zeros(N, L, max_n-N) A = torch.cat([A, zeros], dim=-1) zeros = torch.zeros(max_n-N, L, max_n) A = torch.cat([A, zeros], dim=0) return A class AvgMeter: def __init__(self, name="Metric"): self.name = name self.reset() def reset(self): self.avg, self.sum, self.count = [0] * 3 def update(self, val, count=1): self.count += count self.sum += val * count self.avg = self.sum / self.count def __repr__(self): text = f"{self.name}: {self.avg:.4f}" return text def get_lr(optimizer): for param_group in optimizer.param_groups: return param_group["lr"] def segment_max(x, size_list): size_list = [int(i) for i in size_list] return torch.stack([torch.max(v, 0).values for v in torch.split(x, size_list)]) def segment_sum(x, size_list): size_list = [int(i) for i in size_list] return torch.stack([torch.sum(v, 0) for v in torch.split(x, size_list)]) def segment_softmax(gate, size_list): segmax = segment_max(gate, size_list) # expand segmax shape to alpha shape segmax_expand = torch.cat([segmax[i].repeat(n, 1) for i, n in enumerate(size_list)], dim=0) subtract = gate - segmax_expand exp = torch.exp(subtract) segsum = segment_sum(exp, size_list) # expand segmax shape to alpha shape segsum_expand = torch.cat([segsum[i].repeat(n, 1) for i, n in enumerate(size_list)], dim=0) attention = exp / (segsum_expand + 1e-16) return attention def pad_ms_list(ms_list, thr=0.05, min_mz=20, max_mz=2000): thr = thr*100 mslst = [] for ms in ms_list: ms = np.array(ms) ms[:, 1] = ms[:, 1]/ms[:, 1].max()*100 if thr > 0: ms = ms[(ms[:, 1] >= thr)] ms = ms[(ms[:, 0] >= min_mz)] ms = ms[(ms[:, 0] <= max_mz)] mslst.append(ms) size_list = [ms.shape[0] for ms in mslst] maxlen = max(size_list) l = [] for ms in mslst: extn = maxlen-len(ms) if extn > 0: l.append(np.concatenate([ms, [[0, 0]]*extn], axis=0)) else: l.append(ms) return torch.FloatTensor(np.stack(l)), torch.IntTensor(size_list)