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. 权重获取（重要）

权重较大，不打包进 wheel。运行时按以下优先级解析：

• 环境变量 MS2_WEIGHTS_DIR 指向的目录（直接放好两个 .pt 即可，不触发下载）；或
• 本地缓存 ~/.cache/msformulator/weights，首次运行时从 MS2_WEIGHTS_URL 自动下载。

# 方式 A：用自己的托管地址（GitHub Release / HuggingFace / 对象存储）
export MS2_WEIGHTS_URL="https://github.com/<user>/<repo>/releases/download/v1.0.0"
msformulator 439.326 "[]" "[M+H]+" "[M+H]+"

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

发布前请把 msformulator/download.py 中的 DEFAULT_WEIGHTS_URL 改成你的真实地址， 并（建议）在 WEIGHT_SIZES 中填入两个权重文件的字节数用于完整性校验。

发布到 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。

模块说明

• 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/ 中的权重文件）。