utils.py

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)