diff --git a/.agents/skills/chem-msms-predict/SKILL.md b/.agents/skills/chem-msms-predict/SKILL.md index 75b02dd0..24c1038b 100644 --- a/.agents/skills/chem-msms-predict/SKILL.md +++ b/.agents/skills/chem-msms-predict/SKILL.md @@ -59,6 +59,7 @@ python .agents/skills/chem-msms-predict/scripts/predict_msms.py \ --inten_ckpt downloads/iceberg_dag_inten_msg_best.ckpt \ --collision_energies 20 40 \ --adduct "[M+H]+" \ + --instrument "Orbitrap" \ --output_dir results/msms_prediction ``` @@ -67,6 +68,7 @@ python .agents/skills/chem-msms-predict/scripts/predict_msms.py \ - `--gen_ckpt` / `--inten_ckpt` — paths to ICEBERG checkpoints - `--collision_energies` — one or more collision energies in eV (e.g. `20 40 60`); model was trained on absolute eV values - `--adduct` — supported adducts: `[M+H]+`, `[M-H]-`, `[M+Na]+`, `[M+NH4]+`, and others from `ms_pred.common.ion2mass` +- `--instrument` — instrument type for intensity prediction (e.g. `"Orbitrap"`, `"QTOF"`) - `--threshold` — confidence cutoff for DAG fragment generator (default `0.1`; lower = more fragments) - `--sparse_k` — maximum number of peaks returned (default `100`) - `--cuda_devices` — GPU device IDs (e.g. `"0"` or `"0,1"`); omit or set to `None` for CPU diff --git a/.agents/skills/chem-msms-predict/scripts/predict_msms.py b/.agents/skills/chem-msms-predict/scripts/predict_msms.py index 32203e24..7aa4a373 100644 --- a/.agents/skills/chem-msms-predict/scripts/predict_msms.py +++ b/.agents/skills/chem-msms-predict/scripts/predict_msms.py @@ -36,6 +36,7 @@ def run_iceberg( inten_ckpt: Path, collision_energies: list, adduct: str, + instrument: str, cuda_devices, batch_size: int, num_workers: int, @@ -51,6 +52,7 @@ def run_iceberg( collision_energies=collision_energies, nce=False, adduct=adduct, + instrument=instrument, exp_name="skill_pred", python_path=sys.executable, gen_ckpt=str(gen_ckpt), @@ -167,6 +169,10 @@ def parse_args() -> argparse.Namespace: help="Collision energies in eV (default: 20 40)", ) p.add_argument("--adduct", default="[M+H]+", help="Adduct type (default: [M+H]+)") + p.add_argument( + "--instrument", default="Orbitrap", choices=["Orbitrap", "QTOF"], + help="Instrument the MSG checkpoints condition on (default: Orbitrap)", + ) p.add_argument( "--output_dir", type=Path, default=Path("results/msms_prediction"), help="Output directory", @@ -215,6 +221,7 @@ def main() -> None: inten_ckpt=args.inten_ckpt, collision_energies=args.collision_energies, adduct=args.adduct, + instrument=args.instrument, cuda_devices=args.cuda_devices, batch_size=args.batch_size, num_workers=args.num_workers, diff --git a/conda-envs/msms-agent/conda_only_env.yaml b/conda-envs/msms-agent/conda_only_env.yaml deleted file mode 100644 index 022bb93e..00000000 --- a/conda-envs/msms-agent/conda_only_env.yaml +++ /dev/null @@ -1,17 +0,0 @@ -name: ms-gen -channels: - - pytorch - - conda-forge - - defaults -dependencies: - - python=3.10 - - cython - - matplotlib - - conda-forge::rdkit - - conda-forge::python-graphviz - - conda-forge::pygraphviz - - pytorch::pytorch=2.4.0 - - pip - - einops - - seaborn - - h5py diff --git a/conda-envs/msms-agent/core_env.yaml b/conda-envs/msms-agent/core_env.yaml index 14f182ed..42ad8b8a 100644 --- a/conda-envs/msms-agent/core_env.yaml +++ b/conda-envs/msms-agent/core_env.yaml @@ -15,8 +15,11 @@ dependencies: - seaborn - h5py - pip: - - dgl -f https://data.dgl.ai/wheels/torch-2.4/repo.html - pytorch-lightning + - "ray[tune]==2.7.2" + - "setuptools<81" + - ipython + - pydantic - scikit-learn - pathos - CairoSVG diff --git a/conda-envs/msms-agent/install.sh b/conda-envs/msms-agent/install.sh index 8bfa256c..3a59a92c 100755 --- a/conda-envs/msms-agent/install.sh +++ b/conda-envs/msms-agent/install.sh @@ -22,18 +22,56 @@ fi source $(conda info --base)/etc/profile.d/conda.sh conda activate $ENV_NAME -echo "Installing dgl from custom index..." -uv pip install dgl --find-links https://data.dgl.ai/wheels/torch-2.4/repo.html +if [[ "$(uname)" == "Darwin" ]]; then + # macOS: DGL 2.2.0 graphbolt dylibs only go up to pytorch 2.3.0. + # torchdata>=0.8 dropped datapipes and pulls torch 2.12 as dep; pin 0.7.1 --no-deps. + # torch_scatter/sparse must match torch 2.3.0 exactly. + echo "macOS: installing torch 2.3.0 (CPU)..." + pip install "torch==2.3.0" --index-url https://download.pytorch.org/whl/cpu -echo "Installing pip dependencies with uv..." -uv pip install -r uv_requirements.txt + echo "macOS: installing DGL from torch-2.3 index..." + uv pip install dgl --find-links https://data.dgl.ai/wheels/torch-2.3/repo.html -# torch-scatter / torch-sparse: no generic arm64 wheels on PyPI. -# Use the PyG find-links index pinned to torch 2.4.0+cpu. -TORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])") -PYG_INDEX="https://data.pyg.org/whl/torch-${TORCH_VERSION}+cpu.html" -echo "Installing torch-scatter and torch-sparse from PyG index (torch ${TORCH_VERSION})..." -uv pip install torch-scatter torch-sparse --find-links "$PYG_INDEX" + echo "macOS: installing pip dependencies..." + uv pip install -r uv_requirements.txt + + # DGL 2.2.0 graphbolt imports torchdata which is incompatible with torch 2.3. + # ICEBERG never uses graphbolt at runtime; stub all torchdata imports in graphbolt. + echo "macOS: patching DGL graphbolt to remove torchdata dependency..." + python - <<'DGLPATCH' +import pathlib, sys, re +site = pathlib.Path(sys.executable).parent.parent / "lib/python3.10/site-packages" +graphbolt = site / "dgl/graphbolt" +for f in graphbolt.glob("*.py"): + txt = f.read_text() + if "torchdata" not in txt: + continue + # Replace all torchdata imports with try/except stubs + patched = re.sub( + r'^((?:from|import) torchdata\S*.*)', + r'try:\n \1\nexcept (ImportError, ModuleNotFoundError):\n pass', + txt, flags=re.MULTILINE + ) + f.write_text(patched) + print(f"Patched {f.name}") +DGLPATCH + + echo "macOS: installing torch-scatter and torch-sparse for torch 2.3.0..." + uv pip install --force-reinstall torch-scatter torch-sparse \ + --find-links https://data.pyg.org/whl/torch-2.3.0+cpu.html +else + echo "Installing dgl from custom index..." + uv pip install dgl --find-links https://data.dgl.ai/wheels/torch-2.4/repo.html + + echo "Installing pip dependencies with uv..." + uv pip install -r uv_requirements.txt + + # torch-scatter / torch-sparse: use the PyG find-links index pinned to torch 2.4.0+cpu. + TORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])") + PYG_INDEX="https://data.pyg.org/whl/torch-${TORCH_VERSION}+cpu.html" + echo "Installing torch-scatter and torch-sparse from PyG index (torch ${TORCH_VERSION})..." + uv pip install torch-scatter torch-sparse --find-links "$PYG_INDEX" +fi # Install ms_pred from GitHub. # setup.py includes a Cython extension for massformer (not used by ICEBERG), @@ -53,8 +91,60 @@ setup( ) SETUP_EOF +# Patch ms_pred source for upstream bugs before installing: +# 1. iceberg_elucidation.py calls predict_smis.py via a cwd-relative path +# 2. predict_smis.py uses pl.utilities.seed.seed_everything (removed in PL 2.0) +# 3. predict_smis.py calls torch.cuda.set_device(gpu_id) unconditionally (fails on CPU) +python - "$TMP_DIR/ms-pred/src/ms_pred/dag_pred" <<'PYPATCH' +import sys, pathlib +d = pathlib.Path(sys.argv[1]) +el = d / "iceberg_elucidation.py" +el.write_text(el.read_text().replace( + '{python_path} src/ms_pred/dag_pred/predict_smis.py', + '{python_path} {Path(__file__).resolve().parent / "predict_smis.py"}')) +ps = d / "predict_smis.py" +t = ps.read_text() +t = t.replace('pl.utilities.seed.seed_everything', 'pl.seed_everything') +t = t.replace('torch.cuda.set_device(gpu_id)', + '(torch.cuda.set_device(gpu_id) if (gpu and avail_gpu_num > 0) else None)') +ps.write_text(t) +print("Patched ms_pred dag_pred (pip-install path / PL2 seed / CPU cuda guard)") +PYPATCH + uv pip install "$TMP_DIR/ms-pred" rm -rf "$TMP_DIR" +# macOS only: create fake torchdata package so DGL graphbolt imports succeed. +# DGL uses graphbolt only for distributed training; ICEBERG never triggers it. +if [[ "$(uname)" == "Darwin" ]]; then + echo "macOS: creating torchdata compatibility shim for DGL..." + python - <<'TDPATCH' +import pathlib, sys +site = pathlib.Path(sys.executable).parent.parent / "lib/python3.10/site-packages" +(site / "torchdata/datapipes/iter").mkdir(parents=True, exist_ok=True) +(site / "torchdata/dataloader2").mkdir(parents=True, exist_ok=True) +(site / "torchdata/__init__.py").write_text("") +(site / "torchdata/datapipes/__init__.py").write_text("from . import iter\n") +(site / "torchdata/datapipes/iter/__init__.py").write_text( + "import torch.utils.data\n" + "class IterDataPipe(torch.utils.data.IterableDataset): pass\n" + "class IterableWrapper(IterDataPipe):\n" + " def __init__(self, iterable): self.iterable = iterable\n" + " def __iter__(self): yield from self.iterable\n" + "class Mapper(IterDataPipe):\n" + " def __init__(self, datapipe=None, fn=None): self.datapipe=datapipe; self.fn=fn\n" + " def __iter__(self):\n" + " for item in self.datapipe: yield self.fn(item)\n" +) +(site / "torchdata/dataloader2/__init__.py").write_text("") +(site / "torchdata/dataloader2/graph.py").write_text( + "def traverse_dps(dp): return {}\n" + "def find_dps(g, t): return []\n" + "def replace_dp(g, old, new): return g\n" +) +print("Created torchdata shim") +TDPATCH +fi + rm -f conda_only_env.yaml uv_requirements.txt echo "Environment $ENV_NAME created successfully!"