msformulator 0.1.0__py3-none-any.whl

msformulator/__init__.py

@@ -0,0 +1,15 @@
+"""msformulator - 基于扩散模型的质谱分子式预测包。
+
+入口函数：main(pr_mz, mz2, type_str, peak_type_str)
+    pr_mz        : 母离子精确质量 (float)
+    mz2          : 二级碎片峰列表，形如 [[mz, intensity], ...]
+    type_str     : 母离子加合类型，如 "[M+H]+"
+    peak_type_str: 碎片峰加合类型，如 "[M+H]+"
+
+返回候选分子式列表（按 ppm 误差升序）。
+"""
+
+from .main import main
+from .model import MSFormulator, GaussianDiffusion
+
+__all__ = ["main", "MSFormulator", "GaussianDiffusion"]

msformulator/chemistry.py

@@ -0,0 +1,627 @@
+"""chemistry.py - 分子式相关的化学 / 质量计算工具函数。
+
+从原始 main.py 抽取，作为 msformulator 包的纯计算模块（不依赖 torch）。
+包含：质量微调 (tiaozheng_f / tiaozheng_f2)、不饱和度校验、ppm 计算、
+分子质量计算、分子式各元素数量的正则提取、离子模式解析等。
+"""
+
+import re
+import numpy as np
+from molmass import Formula
+
+
+def tiaozheng_f2(pr_mass, f_mass,
+                 c_num, o_num, n_num,
+                 cl_num, s_num, p_num,
+                 f_num, h_num):
+
+    # 单同位素质量常量
+    MASS = (
+        12.0,                  # C
+        15.99491461957,        # O
+        14.00307400443,        # N
+        34.968852682,          # Cl
+        31.9720711744,         # S
+        30.97376199842,        # P
+        18.99840316273,        # F
+        1.00782503223          # H
+    )
+
+    max_iter = 20   # 防止极端死循环
+    iteration = 0
+
+    while iteration < max_iter:
+        iteration += 1
+
+        mass_diff = pr_mass - f_mass
+
+        abs_mass = abs(mass_diff)
+
+        changed = False
+        if c_num == 25 and n_num == 1 and o_num == 8 and h_num == 47:
+            print( pr_mass - f_mass)
+            print("--------------------")
+
+        # ===========================
+        # 原逻辑开始
+        # ===========================
+
+        if 59 < abs_mass < 63 and p_num == 1:
+            p_num += -2 if mass_diff < 0 else 2
+            changed = True
+
+        elif 11.9 < abs_mass < 13.0:
+            c_num += -1 if mass_diff < 0 else 1
+            changed = True
+
+        elif abs_mass > 50:
+            break
+
+        elif 29.8 < abs_mass < 31 and p_num == 1:
+            p_num += -1 if mass_diff < 0 else 1
+            changed = True
+
+        elif 10 < abs_mass < 29 and h_num == 100:
+            c_h = int(abs_mass / 1.008) + 1
+            h_num += -c_h if mass_diff < 0 else c_h
+            changed = True
+
+        elif abs_mass > 1 and int(abs_mass) % 12 == 0 and (
+            (mass_diff < 0 and int(abs_mass / 12) < c_num) or mass_diff > 0
+        ):
+            c_c = int(abs_mass / 12)
+            if mass_diff < 0 and c_num > c_c:
+                c_num -= c_c
+            if mass_diff > 0:
+                c_num += c_c
+            changed = True
+
+        elif 15.9 < abs_mass < 16.01 and o_num >= 0:
+
+            o_num += -1 if mass_diff < 0 else 1
+            changed = True
+        elif 1 < abs_mass < 20 and h_num > 1:
+            h_num += -1 if mass_diff < 0 else 1
+
+            changed = True
+
+
+        else:
+            break
+
+        if not changed:
+            break
+
+        # ===========================
+        # 手算新质量（超快）
+        # ===========================
+
+        f_mass = (
+            c_num * MASS[0] +
+            o_num * MASS[1] +
+            n_num * MASS[2] +
+            cl_num * MASS[3] +
+            s_num * MASS[4] +
+            p_num * MASS[5] +
+            f_num * MASS[6] +
+            h_num * MASS[7]
+        )
+
+        # ppm 判断
+        ppm = abs((pr_mass - f_mass) / f_mass * 1e6)
+
+        if ppm < 20:
+            break
+
+
+    return c_num, o_num, n_num, cl_num, s_num, p_num, f_num, h_num
+
+def tiaozheng_f(pr_mass,f_mass,f):
+    mass=pr_mass-f_mass
+
+    c_num = extract_c_numbers(f)[0]
+    o_num = extract_o_numbers(f)
+
+    n_num = extract_n_numbers(f)
+
+    f_num = extract_f_numbers(f)
+
+
+    s_num = extract_s_numbers(f)
+
+
+    p_num = extract_p_numbers(f)
+
+
+    cl_num = extract_cl_numbers(f)
+
+    h_num = extract_h_numbers(f)
+    # print(c_num, o_num, n_num, cl_num, s_num, p_num, f_num, h_num)
+    if np.abs(mass)>50:
+        return f
+    elif np.abs(mass)>29 and np.abs(mass)<31:
+        if mass<0:
+            p_num = p_num - 1
+        else:
+            p_num = p_num + 1
+    elif np.abs(mass)>11.5 and np.abs(mass)<12.5:
+
+        if mass<0:
+
+            c_num = c_num - 1
+
+        else:
+            c_num = c_num + 1
+
+
+    elif np.abs(mass)>59 and np.abs(mass)<63 and p_num==1:
+        if mass<0:
+            p_num = p_num - 2
+        else:
+            p_num = p_num + 2
+    # elif np.abs(mass)>10  and h_num==49:
+    #     c_h=(int(np.abs(mass)/1.007)+1)
+    #     if mass<0:
+    #         h_num = h_num - c_h
+    #     else:
+    #         h_num = h_num + c_h
+    # elif (int(np.abs(mass))%12)==0 :
+    #     c_h=(int(np.abs(mass)/12))
+    #     if mass<0:
+    #         c_num = c_num - c_h
+    #     else:
+    #         c_num = c_num + c_h
+    elif np.abs(mass)>31.5 and np.abs(mass)<33 :
+        c_h=(int(np.abs(mass)/1.007)+1)
+        if mass<0:
+            s_num = s_num - 1
+        else:
+            s_num = s_num + 1
+    elif np.abs(mass)>15.5 and np.abs(mass)<17 :
+        c_h=(int(np.abs(mass)/1.007)+1)
+        if mass<0:
+            o_num = o_num - 1
+        else:
+            o_num = o_num + 1
+    elif np.abs(mass)>1 and np.abs(mass)<11 :
+
+        if mass < 0:
+
+            h_num = h_num - 1
+
+        else:
+            h_num = h_num + 1
+    else:
+        return f
+    new_f = ""
+
+    new_f += "C" + str(c_num)
+    if h_num > 0:
+            new_f += "H" + str(h_num)
+    if cl_num > 0:
+            new_f += "Cl" + str(cl_num)
+    if n_num > 0:
+            new_f += "N" + str(n_num)
+    if o_num > 0:
+            new_f += "O" + str(o_num)
+
+    if s_num > 0:
+            new_f += "S" + str(s_num)
+    if p_num > 0:
+            new_f += "P" + str(p_num)
+    if f_num > 0:
+            new_f += "F" + str(f_num)
+    new_mass=calculate_molecular_mass(new_f)['mass']
+    if calculate_ppm(pr_mass,new_mass,use_absolute=True)<20:
+
+        return new_f
+    else:
+        return tiaozheng_f(pr_mass,new_mass,new_f)
+
+def calculate_unsaturation_and_check(C, O, N, Cl, S, P, F, H, allow_half_integer=False):
+    """
+    计算分子的不饱和度并判断是否合理
+
+    参数:
+    C, O, N, Cl, S, P, F, H: 各原子的数量（整数）
+    allow_half_integer: 是否允许半整数不饱和度（对于离子/自由基），默认False
+
+    返回:
+    (不饱和度, 是否合理): 返回一个元组，包含计算得到的不饱和度和判断结果(1/0)
+    """
+
+    # 计算不饱和度
+    # 公式: Ω = (2C + 2 + N - H - 卤素总数) / 2
+    # 其中卤素包括F, Cl, Br, I。这里我们有F和Cl
+    halogens = F + Cl
+    numerator = 2 * C + 2 + N - H - halogens
+    unsaturation = numerator / 2.0
+
+    # 判断是否合理
+    is_reasonable = 1  # 初始假设合理
+
+    # 条件0: 分子必须有至少一个碳原子（对有机分子而言）
+    if C <= 0:
+        is_reasonable = 0
+        return unsaturation, is_reasonable
+
+    # 条件1: 不饱和度必须是非负数
+    if unsaturation < 0:
+        is_reasonable = 0
+        return unsaturation, is_reasonable
+
+    # 条件2: 检查不饱和度是否为整数
+    # 如果不允许半整数，则必须为整数
+    if not allow_half_integer:
+        if not float(unsaturation).is_integer():
+            is_reasonable = 0
+            return unsaturation, is_reasonable
+    else:
+        # 如果允许半整数，检查2*Ω是否为整数
+        if not float(2 * unsaturation).is_integer():
+            is_reasonable = 0
+            return unsaturation, is_reasonable
+
+    # 条件3: 计算理论最大不饱和度
+    # 当H=0时的不饱和度
+    max_unsaturation = (2 * C + 2 + N - halogens) / 2.0
+    if unsaturation > max_unsaturation:
+        is_reasonable = 0
+        return unsaturation, is_reasonable
+
+    # 条件4: 所有原子数必须非负
+    atom_counts = [C, O, N, Cl, S, P, F, H]
+    for count in atom_counts:
+        if count < 0:
+            is_reasonable = 0
+            return unsaturation, is_reasonable
+
+    # 条件5: 检查氢原子数的奇偶性
+    # 对于中性分子，当氮原子数为奇数时，氢原子数应为奇数
+    # 当氮原子数为偶数时，氢原子数应为偶数
+    if N % 2 == 1:  # 氮原子数为奇数
+        if H % 2 == 0:  # 氢原子数为偶数
+            is_reasonable = 0
+    elif N!=0:  # 氮原子数为偶数
+        if H % 2 == 1:  # 氢原子数为奇数
+            is_reasonable = 0
+
+    return unsaturation, is_reasonable
+
+def check_molecular_formula(C, O, N, Cl, S, P, F, H, allow_half_integer=False):
+    """
+    完整的分子式检查函数
+
+    参数:
+    各原子数量
+    allow_half_integer: 是否允许半整数不饱和度，默认False
+
+    返回:
+    合理返回1，不合理返回0
+    """
+
+    unsaturation, is_reasonable = calculate_unsaturation_and_check(C, O, N, Cl, S, P, F, H, allow_half_integer)
+    # if C == 9 and H == 14 and N == 3 and O == 7 and P == 1:
+    #     print(unsaturation, is_reasonable)
+    return is_reasonable
+
+def calculate_ppm(mass1, mass2, use_absolute=True, use_average=False):
+    """
+    计算两个质量之间的ppm差异
+
+    参数：
+    mass1, mass2: 两个质量值（单位相同）
+    use_absolute: 是否返回绝对值（True）或有符号值（False）
+    use_average: 是否使用平均值作为分母（False时用mass2作为分母）
+
+    返回：
+    ppm差异值
+    """
+    if mass2 == 0:
+        raise ValueError("分母不能为0")
+
+    if use_average:
+        # 使用平均值作为分母
+        denominator = (mass1 + mass2) / 2.0
+    else:
+        # 使用mass2作为分母（通常作为理论值或参考值）
+        denominator = mass2
+
+    ppm = (mass1 - mass2) / denominator * 1e6
+
+    if use_absolute:
+        return abs(ppm)
+    else:
+        return ppm
+
+def calculate_molecular_mass(formula, charge=0, ion_type='M'):
+    """计算分子质量（推荐使用）"""
+    try:
+        # 单同位素质量
+
+        return {
+            'formula': formula,
+            'mass': Formula(formula).isotope.mass,
+        }
+    except Exception as e:
+        raise ValueError(f"计算失败: {e}")
+
+def softmax(x, axis=-1):
+    """
+    Softmax函数实现
+
+    参数:
+    x: 输入数组
+    axis: 沿着哪个轴计算softmax（默认最后一个轴）
+
+    返回:
+    softmax计算结果
+    """
+    # 防止数值溢出，减去最大值
+    x_max = np.max(x, axis=axis, keepdims=True)
+    exp_x = np.exp(x - x_max)
+
+    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
+
+def extract_h_numbers(s):
+    """
+    使用正则表达式在字符串中查找 'H' 及其后的数字。
+
+    参数:
+        s (str): 待检查的字符串。
+
+    返回:
+        int:
+            - 如果未找到 'H'，返回 0
+            - 如果找到 'H' 且其后没有数字字符，返回 1
+            - 如果找到 'H' 且其后有数字字符，则返回这些连续数字字符组成的整数
+    """
+    # 模式解释：H 后面捕获一个或多个数字（\d+）
+    pattern = r'H(\d+)'
+    match = re.search(pattern, s)
+    if match:
+        # 如果找到 'H' 并且后面有数字，返回匹配到的数字
+        return int(match.group(1))
+    else:
+        # 检查字符串中是否有 'H'（不论后面是否有数字）
+        if 'H' in s:  # 使用 in 运算符检查字符是否存在[1,2](@ref)
+            return 1
+        else:
+            return 0
+
+def extract_s_numbers(s):
+    """
+    使用正则表达式在字符串中查找 'H' 及其后的数字。
+
+    参数:
+        s (str): 待检查的字符串。
+
+    返回:
+        int:
+            - 如果未找到 'H'，返回 0
+            - 如果找到 'H' 且其后没有数字字符，返回 1
+            - 如果找到 'H' 且其后有数字字符，则返回这些连续数字字符组成的整数
+    """
+    # 模式解释：H 后面捕获一个或多个数字（\d+）
+    pattern = r'S(\d+)'
+    match = re.search(pattern, s)
+    if match:
+        # 如果找到 'H' 并且后面有数字，返回匹配到的数字
+        return int(match.group(1))
+    else:
+        # 检查字符串中是否有 'H'（不论后面是否有数字）
+        if 'S' in s:  # 使用 in 运算符检查字符是否存在[1,2](@ref)
+            return 1
+        else:
+            return 0
+
+def extract_p_numbers(s):
+    """
+    使用正则表达式在字符串中查找 'H' 及其后的数字。
+
+    参数:
+        s (str): 待检查的字符串。
+
+    返回:
+        int:
+            - 如果未找到 'H'，返回 0
+            - 如果找到 'H' 且其后没有数字字符，返回 1
+            - 如果找到 'H' 且其后有数字字符，则返回这些连续数字字符组成的整数
+    """
+    # 模式解释：H 后面捕获一个或多个数字（\d+）
+    pattern = r'P(\d+)'
+    match = re.search(pattern, s)
+    if match:
+        # 如果找到 'H' 并且后面有数字，返回匹配到的数字
+        return int(match.group(1))
+    else:
+        # 检查字符串中是否有 'H'（不论后面是否有数字）
+        if 'P' in s:  # 使用 in 运算符检查字符是否存在[1,2](@ref)
+            return 1
+        else:
+            return 0
+
+def extract_f_numbers(s):
+    """
+    使用正则表达式在字符串中查找 'H' 及其后的数字。
+
+    参数:
+        s (str): 待检查的字符串。
+
+    返回:
+        int:
+            - 如果未找到 'H'，返回 0
+            - 如果找到 'H' 且其后没有数字字符，返回 1
+            - 如果找到 'H' 且其后有数字字符，则返回这些连续数字字符组成的整数
+    """
+    # 模式解释：H 后面捕获一个或多个数字（\d+）
+    pattern = r'F(\d+)'
+    match = re.search(pattern, s)
+    if match:
+        # 如果找到 'H' 并且后面有数字，返回匹配到的数字
+        return int(match.group(1))
+    else:
+        # 检查字符串中是否有 'H'（不论后面是否有数字）
+        if 'F' in s:  # 使用 in 运算符检查字符是否存在[1,2](@ref)
+            return 1
+        else:
+            return 0
+
+def extract_cl_numbers(s):
+    """
+    使用正则表达式在字符串中查找 'H' 及其后的数字。
+
+    参数:
+        s (str): 待检查的字符串。
+
+    返回:
+        int:
+            - 如果未找到 'H'，返回 0
+            - 如果找到 'H' 且其后没有数字字符，返回 1
+            - 如果找到 'H' 且其后有数字字符，则返回这些连续数字字符组成的整数
+    """
+    # 模式解释：H 后面捕获一个或多个数字（\d+）
+    pattern = r'Cl(\d+)'
+    match = re.search(pattern, s)
+    if match:
+        # 如果找到 'H' 并且后面有数字，返回匹配到的数字
+        return int(match.group(1))
+    else:
+        # 检查字符串中是否有 'H'（不论后面是否有数字）
+        if 'Cl' in s:  # 使用 in 运算符检查字符是否存在[1,2](@ref)
+            return 1
+        else:
+            return 0
+
+def extract_c_numbers(text):
+    pattern = r'C\s*(\d+\.?\d*)'  # 模式
+    matches = re.findall(pattern, text, re.IGNORECASE)  # 全局匹配并忽略大小写
+    return [float(num) if '.' in num else int(num) for num in matches]  # 转为数字类型
+
+def extract_o_numbers(s):
+
+        """
+        使用正则表达式在字符串中查找 'H' 及其后的数字。
+
+        参数:
+            s (str): 待检查的字符串。
+
+        返回:
+            int:
+                - 如果未找到 'H'，返回 0
+                - 如果找到 'H' 且其后没有数字字符，返回 1
+                - 如果找到 'H' 且其后有数字字符，则返回这些连续数字字符组成的整数
+        """
+        # 模式解释：H 后面捕获一个或多个数字（\d+）
+        pattern = r'O(\d+)'
+        match = re.search(pattern, s)
+        if match:
+            # 如果找到 'H' 并且后面有数字，返回匹配到的数字
+            return int(match.group(1))
+        else:
+            # 检查字符串中是否有 'H'（不论后面是否有数字）
+            if 'O' in s:  # 使用 in 运算符检查字符是否存在[1,2](@ref)
+                return 1
+            else:
+                return 0
+
+def extract_n_numbers(s):
+    """
+    使用正则表达式在字符串中查找 'H' 及其后的数字。
+
+    参数:
+        s (str): 待检查的字符串。
+
+    返回:
+        int:
+            - 如果未找到 'H'，返回 0
+            - 如果找到 'H' 且其后没有数字字符，返回 1
+            - 如果找到 'H' 且其后有数字字符，则返回这些连续数字字符组成的整数
+    """
+    # 模式解释：H 后面捕获一个或多个数字（\d+）
+    pattern = r'N(\d+)'
+    match = re.search(pattern, s)
+    if match:
+        # 如果找到 'H' 并且后面有数字，返回匹配到的数字
+        return int(match.group(1))
+    else:
+        # 检查字符串中是否有 'H'（不论后面是否有数字）
+        if 'N' in s:  # 使用 in 运算符检查字符是否存在[1,2](@ref)
+            return 1
+        else:
+            return 0
+
+def parse_ion_pattern(ion_pattern):
+    """
+    解析离子模式字符串，提取各个组成部分
+    格式示例: [2M-2MeOH+6H]3+
+    """
+
+    # 主正则表达式匹配整个模式
+    pattern = r'\[(.*?)\]([-+]?\d*[+-]?)'
+    main_match = re.match(pattern, ion_pattern)
+
+    if not main_match:
+        return None
+
+    inner_content = main_match.group(1)  # 方括号内的内容: 2M-2MeOH+6H
+    charge_str = main_match.group(2)  # 电荷部分: 3+
+
+    # 解析电荷
+    charge_match = re.match(r'([+-]?\d*)([+-])?$', charge_str)
+    if charge_match:
+        charge_num = charge_match.group(1) or '1'
+        charge_sign = charge_match.group(2) or '+'
+    else:
+        charge_num = '1'
+        charge_sign = '+'
+
+    # 解析M前的数字
+    m_pattern = r'^(\d+)?M'
+    m_match = re.search(m_pattern, inner_content)
+    m_coefficient = m_match.group(1) if m_match and m_match.group(1) else '1'
+
+    # 解析所有基团
+    groups = []
+    # 正则表达式匹配各个基团: 可选数字 + 基团名称，前面可能有+或-号
+    group_pattern = r'([-+]?)(\d*)?([A-Za-z0-9]+)'
+
+    # 从第二个字符开始匹配，跳过第一个基团M
+    remaining_content = inner_content[len(m_match.group(0)):] if m_match else inner_content
+
+    for match in re.finditer(group_pattern, remaining_content):
+        if match.group(3):  # 确保有基团名称
+            sign = match.group(1) or '+'  # 如果没有符号，默认为+
+            coefficient = match.group(2) or '1'  # 如果没有数字，默认为1
+            group_name = match.group(3)
+
+            groups.append({
+                'group': group_name,
+                'coefficient': coefficient,
+                'sign': sign
+            })
+
+    return {
+        'M_coefficient': m_coefficient,
+        'groups': groups,
+        'charge': {
+            'number': charge_num,
+            'sign': charge_sign
+        }
+    }
+
+def get_mass(ion_pattern, mz):
+    result = parse_ion_pattern(ion_pattern)
+
+    add_number = 0
+    for group in result['groups']:
+        if group['group'] != 'M':
+            add_number = add_number - float(group['sign'] + str(
+                float(calculate_molecular_mass(group['group'])['mass']) * float(group['coefficient'])))
+    # print(add_number)
+    if result['charge']['number'] == '+' or result['charge']['number'] == '-':
+        result['charge']['number'] = 1
+
+    mz = (float(result['charge']['number']) * mz + add_number) / float(result['M_coefficient'])
+    # print(mz)
+    return mz

msformulator/cli.py

@@ -0,0 +1,43 @@
+"""cli.py - 命令行入口。
+
+安装本包后可通过 ``ms2-formula`` 命令调用，等价于
+``python -m msformulator.cli <pr_mz> <mz2> <type> <peak_type>``。
+"""
+import argparse
+import ast
+import sys
+
+from .main import main as _predict
+from .download import ensure_weights
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(
+        prog="ms2-formula",
+        description="基于扩散模型的质谱（MS2）分子式预测。",
+    )
+    parser.add_argument("pr_mz", type=float, help="母离子精确质量 (float)，如 439.326")
+    parser.add_argument(
+        "mz2",
+        type=str,
+        help='二级碎片峰列表，如 "[]" 或 "[[100.1, 200.0], [150.2, 80.0]]"',
+    )
+    parser.add_argument("type", type=str, help='母离子加合类型，如 "[M+H]+"')
+    parser.add_argument("peak_type", type=str, help='碎片峰加合类型，如 "[M+H]+"')
+    args = parser.parse_args(argv)
+
+    # 提前触发权重下载，给出清晰进度
+    ensure_weights()
+
+    result = _predict(
+        args.pr_mz,
+        ast.literal_eval(args.mz2),
+        args.type,
+        args.peak_type,
+    )
+    print(result)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

msformulator/config.py

@@ -0,0 +1,20 @@
+"""config.py - 包内资源路径配置。
+
+集中管理数据文件与模型权重的路径，避免在各模块中写死相对路径。
+"""
+from pathlib import Path
+
+# 本包目录，例如 .../msformulator
+PACKAGE_DIR = Path(__file__).resolve().parent
+
+# 数据文件目录（embedding_parameters.pth、mass_zhushi.npy），随包分发
+DATA_DIR = PACKAGE_DIR / "data"
+
+# 模型权重目录（两个 .pt 文件）不随包分发，运行时通过 download 模块解析
+# （优先 MS2_WEIGHTS_DIR，否则缓存目录，缺失则自动下载）。
+from .download import ensure_weights, weight_path, WEIGHT_FILES  # noqa: E402
+
+
+def get_weights_dir():
+    """返回权重所在目录（必要时触发下载）。"""
+    return ensure_weights()