MSFormulator

MSFormulator 是基于扩散模型（diffusion）的质谱二级碎片（MS2）分子式预测方法。 本仓库是对原始 dxzx2/ms2_formula_model 的整理版本：将原本堆在单个 main.py（约 1480 行）中的 化学计算、推理流水线、模型定义拆分成了清晰的 Python 包结构，并把数据/权重与代码分离， 现已可经 pip install 安装与发布（PyPI 包名 msformulator）。

目录结构

ms2_formula_model_organized/        # 工作目录（即本仓库根）
├── run.py                      # 命令行入口：python run.py <pr_mz> <mz2> <type> <peak_type>
├── pyproject.toml             # 打包元数据（包名 msformulator）
├── MANIFEST.in                # 排除大体积权重，避免打进分发包
├── requirements.txt           # 依赖
├── README.md
├── msformulator/              # 核心包（import 名，即方法名）
│   ├── __init__.py            # 暴露 main / MSFormulator / GaussianDiffusion
│   ├── config.py              # 数据/权重的路径配置（不再硬编码相对路径）
│   ├── model.py               # 神经网络：MSFormulator / GaussianDiffusion / UNet1D
│   ├── chemistry.py           # 纯计算：质量微调、不饱和度校验、ppm、分子式解析等
│   ├── inference.py           # 推理流水线：注释查找→峰编码→模型预测→分子式解码
│   ├── cli.py                 # 命令行入口（注册为 `msformulator` 命令）
│   ├── download.py            # 权重解析与按需下载
│   ├── main.py                # 入口函数 main()
│   └── data/                  # 小型数据：embedding_parameters.pth、mass_zhushi.npy
├── weights/                    # 模型权重（两个 .pt，体积较大，不随包分发）
└── scripts/                    # 辅助/遗留脚本
    ├── find_cl.py
    ├── find_pubchem_formula.py
    ├── read_pubchem.py
    ├── add_zhushi.py
    └── legacy/                 # 旧版本 main_quick.py / TEST_BIAOZHUN.py（见下“已知问题”）

入口函数

入口为 main(pr_mz, mz2, type_str, peak_type_str)：

返回按 ppm 误差升序排列的候选分子式列表。

运行方式

1. 安装依赖：pip install -r requirements.txt
2. 命令行（源码方式）：

   python run.py 439.326 "[]" "[M+H]+" "[M+H]+"
   

   或作为模块：

   python -m msformulator.main 439.326 "[]" "[M+H]+" "[M+H]+"
   

3. 在 Python 中调用：

   from msformulator.main import main
   print(main(439.326, [], "[M+H]+", "[M+H]+"))
   

运行推理时需要两个 .pt 权重文件（约 1.5GB）与 msformulator/data/ 下的小型数据文件。 小型数据文件随包分发；权重文件不随包分发，首次运行会自动下载（见下文“安装 / 权重”）。

安装（pip）

本包已配置为标准可安装项目（pyproject.toml，分发名 msformulator）。

1. 本地 / 源码安装

cd ms2_formula_model_organized
pip install .

2. 从 PyPI 安装（发布后）

pip install msformulator

3. 权重获取（开箱即用）

权重较大（约 1.5GB），不打包进 wheel。安装后首次运行会自动从 Hugging Face Hub 下载到 ~/.cache/msformulator/weights，之后复用、无需重复下载。无需任何额外配置即可开箱运行。

如需覆盖默认行为，可设置环境变量：

• MS2_WEIGHTS_DIR：指向已放好两个 .pt 的目录（不触发下载）；
• MS2_WEIGHTS_URL：自定义权重基地址（GitHub Release / HuggingFace / 对象存储均可）；
• MS2_CACHE_DIR：覆盖默认缓存根目录。

# 方式 A：本地已存在权重目录，跳过下载
export MS2_WEIGHTS_DIR="/path/to/your/weights"
msformulator 439.326 "[]" "[M+H]+" "[M+H]+"

# 方式 B：指定自定义下载地址
export MS2_WEIGHTS_URL="https://github.com/<user>/<repo>/releases/download/v1.0.0"
msformulator 439.326 "[]" "[M+H]+" "[M+H]+"

默认下载地址与权重字节数（用于完整性校验）已在 msformulator/download.py 中配置。

发布到 PyPI

cd ms2_formula_model_organized
pip install build twine
python -m build            # 生成 dist/ 下的 sdist 与 wheel
twine check dist/*         # 检查分发包合法性
twine upload dist/*        # 上传到 PyPI（需提前注册 PyPI 账号）

MANIFEST.in 已配置为排除 weights/ 与 scripts/，确保 1.5GB 权重不会被打进 sdist/wheel。 升版时记得同步修改 pyproject.toml 中的 version。

模块说明

• chemistry.py：原 main.py 中的纯计算函数（无 torch 依赖）。包括 tiaozheng_f / tiaozheng_f2（质量微调）、calculate_ppm、calculate_molecular_mass、 check_molecular_formula、extract_*_numbers（分子式各元素数量解析）、parse_ion_pattern / get_mass 等。
• inference.py：推理流水线。get_zhushi（按质量查找注释）、get_peak_mass（碎片峰编码）、 model_test（加载权重并做扩散模型预测）、get_one_formula_tiaozheng / get_formula（解码分子式）。
• model.py：原 model.py 完整保留，仅将 embedding_parameters.pth 的加载路径改为相对于包目录解析。 核心网络类即 MSFormulator（配合 GaussianDiffusion 做扩散采样）。
• download.py：权重解析与下载（支持 MS2_WEIGHTS_DIR / MS2_WEIGHTS_URL / MS2_CACHE_DIR 环境变量）。

已知问题

• scripts/legacy/main_quick.py 与 scripts/legacy/TEST_BIAOZHUN.py 引用了 from msformulator.model import PCTE， 但 model.py 中并不存在 PCTE 类（这是原始代码就存在的问题），因此这两个文件当前无法运行。 它们作为历史版本保留在 legacy/ 中，待补齐 PCTE 定义后可恢复使用。
• 主入口 main() 使用的是 MSFormulator，可正常运行（依赖 weights/ 中的权重文件）。