From 6f88425a945b6e28470aa3c24fcd8ca1e05acde0 Mon Sep 17 00:00:00 2001 From: Utkarsh Singh Date: Mon, 13 Jul 2026 09:44:36 -0700 Subject: [PATCH 1/8] feat(deps): add tmol energy function for guided sampling Add the tmol (Rosetta beta_nov2016_cart) energy function as a pixi feature attached to all model envs, for use as a reward in guided sampling (any strategy). No prebuilt wheel matches this project's torch 2.7.1, so tmol compiles from source; the feature encodes the conda-CUDA cmake workarounds (CUDA root, NVTX3 headers, no build isolation) and a memory-bounded, sm_80;86-targeted build. See TMOL_BUILD.md for the full rationale and the Wynton build/run workflow. pixi.lock is intentionally not included; it regenerates on next `pixi install`. Co-Authored-By: Claude Opus 4.8 --- TMOL_BUILD.md | 133 +++++++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 57 ++++++++++++++++++--- 2 files changed, 183 insertions(+), 7 deletions(-) create mode 100644 TMOL_BUILD.md diff --git a/TMOL_BUILD.md b/TMOL_BUILD.md new file mode 100644 index 0000000..9742824 --- /dev/null +++ b/TMOL_BUILD.md @@ -0,0 +1,133 @@ +# Building & Using tmol on Wynton + +[`tmol`](https://github.com/uw-ipd/tmol) (UW-IPD) is a GPU-accelerated PyTorch +reimplementation of Rosetta's `beta_nov2016_cart` energy function with custom +C++/CUDA kernels. We use it as a reward (energy score) for guided sampling — it's +strategy-agnostic, usable with FK(C) steering, pure guidance, DPS, etc. This note +records how it's wired into the pixi environments and how to build/run it on +Wynton, because the install is non-trivial. + +## TL;DR + +- tmol is a pypi-dependency of the `tmol` pixi **feature**, attached to every + model env (`boltz*`, `protenix*`, `rf3*`). See `[tool.pixi.feature.tmol*]` in + `pyproject.toml`. +- It is **compiled from source** (no usable prebuilt wheel — see below), so a + clean build needs a CUDA toolchain (present in the linux-64 envs via + `cuda-toolkit`) and the cmake/parallelism settings baked into + `[tool.pixi.feature.tmol.activation.env]`. +- Build it on a machine with `nvcc` (e.g. **gpudev1** — a GPU is *not* required + to compile, only to run): + + ```bash + cd /wynton/home/fraserlab//samplew0rks + PIXI_CACHE_DIR=/scratch//pixi-cache pixi install -e boltz-dev + ``` + +- Run on an **sm_80 / sm_86** GPU node (that's what we compile kernels for): + + ```bash + qrsh -q gpu.q -l compute_cap=80 -pe smp 8 + cd /wynton/home/fraserlab//samplew0rks + pixi run -e boltz-dev python -c "import tmol, torch; print(tmol.__version__, torch.cuda.get_device_name())" + ``` + +## Why source-built (no wheel) + +tmol publishes an **sdist only** on PyPI; prebuilt wheels live on GitHub Releases +and are lane-specific (`python × torch-minor × CUDA`). The published GPU lanes +cover torch **2.8–2.12** only. This project is pinned to **torch 2.7.1** +(protenix hard-pins `torch==2.7.1`), so **no prebuilt wheel matches** and tmol +must compile from source. If protenix ever relaxes its torch pin and the project +moves to torch ≥ 2.8, tmol can switch to the fast prebuilt-wheel install and all +the cmake workarounds below can be deleted. + +## The build workarounds (why the config looks the way it does) + +All of these live in `[tool.pixi.feature.tmol.activation.env]` / +`[tool.pixi.pypi-options]` in `pyproject.toml`. They exist because tmol compiles +torch C++/CUDA extensions against a **conda** CUDA toolkit, which is laid out +differently than a system CUDA install. + +1. **Build isolation off** (`[tool.pixi.pypi-options] no-build-isolation`). + tmol needs the env's torch present at build time; an isolated build can't see + it. Its PEP517 build deps (`scikit-build-core`, `pybind11`, `packaging`) are + added to the feature so they're available without isolation. `ninja` comes + from `[tool.pixi.dependencies]`. + +2. **CUDA root** (`CUDA_TOOLKIT_ROOT_DIR` / `CUDAToolkit_ROOT` → + `$CONDA_PREFIX/targets/x86_64-linux`). torch's legacy `FindCUDA` looks for + headers in `/include`, but conda puts them under + `targets/x86_64-linux/include`. Without this it reports + *"cannot find the CUDA libraries."* + +3. **NVTX3 headers** (`USE_SYSTEM_NVTX=ON` + `CMAKE_INCLUDE_PATH` → + `.../site-packages/nvidia/nvtx/include`). CUDA 12 dropped the old + `nvToolsExt` library; torch needs header-only NVTX3, which conda's `cuda-nvtx` + does **not** ship — only torch's bundled `nvidia-nvtx` wheel has the headers. + Without this: *"Could NOT find nvtx3"* → dead `CUDA::nvToolsExt` target. + +4. **Skip the wheel probe** (`TMOL_DISABLE_WHEEL_FETCH=1`). tmol's sdist backend + otherwise probes GitHub Releases for a matching wheel (none exists for torch + 2.7) before falling back to a source build. + +5. **CUDA architectures** (`CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=80;86"`). + - Passed via **`CMAKE_ARGS`, not `SKBUILD_CMAKE_DEFINE`**: the latter splits + entries on `;`, so a `;`-separated arch list crashes its parser + (`not enough values to unpack ... Field cmake.define`). `CMAKE_ARGS` is + shlex-split on spaces, so the `;` survives. + - Targets **sm_80** (A100, 3 nodes) + **sm_86** (RTX 3090 / A40, 29 nodes) — + everything a `-l compute_cap=80` request can land on. + +6. **Bounded parallelism** (`CMAKE_BUILD_PARALLEL_LEVEL=6`). tmol's CUDA kernels + are template-heavy: each `nvcc` front-end (`cicc`) peaks ~3.4 GB, and it + builds **all requested archs concurrently per file**. At ninja's default + (`cores+2`) this OOM-kills the build against **Wynton's 48 GB/user cgroup**. + At 2 archs, `-j6` peaks ~24–34 GB (measured 24 GB). Must live in + `activation.env` — **pixi does not forward shell `export`s into the build + subprocess** (only vars declared in `activation.env` reach it). + +## Gotchas + +- **`PIXI_CACHE_DIR` must be on local disk (`/scratch`).** The default cache on + the NFS home fails uv's atomic renames mid-build + (`Resource busy (os error 16)`). Not committed to `pyproject.toml` because it's + per-machine; set it on the command line or in your shell profile. +- **uv caches the built wheel ignoring these settings.** If you change the arch + list (or any cmake setting) and rebuild, uv may silently reuse the old wheel + and exit 0 without recompiling. Force a real rebuild by purging the cache and + the installed copy: + ```bash + C=/scratch//pixi-cache/uv-cache + find $C/sdists-v9/pypi/tmol -iname 'tmol*.whl' -printf '%h\n' | sort -u | xargs -r rm -rf + find $C/archive-v0 -maxdepth 2 -type d -iname 'tmol*' | xargs -r -n1 dirname | sort -u | xargs -r rm -rf + rm -rf .pixi/envs/boltz-dev/lib/python3.12/site-packages/tmol* + ``` +- **Verify the binary, not the exit code.** Confirm which archs actually got + compiled in: + ```bash + cuobjdump --list-elf .pixi/envs/boltz-dev/lib/python3.12/site-packages/tmol/_C*.so | grep -oE 'sm_[0-9]+' | sort -u + ``` +- **gpudev1 can build but not run tmol.** Its GPUs are GTX 1080 (sm_61), which + isn't in our arch list — `import tmol` works there, but CUDA ops won't. + +## Adding more GPU architectures + +To also run on e.g. L40/RTX 4090 (sm_89) or H100 (sm_90), extend the arch list +in `CMAKE_ARGS` (`-DCMAKE_CUDA_ARCHITECTURES=80;86;89;90`). Each extra arch adds +a concurrent `cicc` (~3.4 GB) per file, so **lower `CMAKE_BUILD_PARALLEL_LEVEL`** +to stay under the 48 GB cap (roughly: `48 GB / (archs × 4 GB)` jobs). Then purge +the uv cache (above) so the rebuild actually happens. + +## Wynton GPU fleet (compute capability → nodes) + +| compute_cap | Arch / cards | # nodes | tmol runs? | +|---|---|---|---| +| 86 | Ampere (RTX 3090 / A40) | 29 | ✅ (built) | +| 80 | Ampere (A100) | 3 | ✅ (built) | +| 75 | Turing (2080 Ti / T4) | 14 | ❌ | +| 61 | Pascal (GTX 1080, gpudev1) | 18 | ❌ | +| ~90 | Hopper (H100) | 1 | ❌ (add sm_90) | + +Request a node with `qrsh -q gpu.q -l compute_cap=80 -pe smp ` (matches +nodes with capability ≥ 8.0; the job gets the node exclusively). diff --git a/pyproject.toml b/pyproject.toml index e370611..39f6f26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,44 @@ platforms = ["osx-arm64"] pypi-dependencies = {"protenix" = "*", "einx" = "*", "triton" = "*"} platforms = ["linux-64"] +# tmol (Rosetta energy fn) as a reward for guided sampling — any strategy, all +# model envs, linux-64. No wheel matches torch 2.7.1, so it builds from source: +# scikit-build-core/pybind11/packaging are PEP517 build deps; build isolation is +# off (see [tool.pixi.pypi-options]) so nvcc sees the env's torch. See TMOL_BUILD.md. +[tool.pixi.feature.tmol] +pypi-dependencies = {"tmol" = "*", "scikit-build-core" = ">=0.10", "pybind11" = ">=2.12", "packaging" = "*"} +platforms = ["linux-64"] + +# tmol compiles torch C++/CUDA extensions from source. Three conda-CUDA quirks +# must be papered over for torch's cmake to find everything: +# 1. torch's legacy FindCUDA wants the CUDA root at $PREFIX/targets/, +# not $PREFIX (headers live under targets/x86_64-linux/include). +# 2. CUDA 12 dropped the nvToolsExt lib; torch needs header-only NVTX3, which +# conda's cuda-nvtx does NOT ship — only torch's bundled nvidia-nvtx wheel +# has the headers. USE_SYSTEM_NVTX=ON + that include dir points cmake at them. +# 3. TMOL_DISABLE_WHEEL_FETCH skips the (futile, no torch-2.7 lane) GitHub probe. +# SKBUILD_CMAKE_DEFINE is scikit-build-core's env channel for -D cmake args. +[tool.pixi.feature.tmol.activation.env] +CUDACXX = "$CONDA_PREFIX/bin/nvcc" +CUDA_PATH = "$CONDA_PREFIX/targets/x86_64-linux" +TMOL_DISABLE_WHEEL_FETCH = "1" +# tmol's CUDA kernels are template-heavy (each nvcc cicc peaks ~3.4 GB) and it +# builds every arch of CMAKE_CUDA_ARCHITECTURES concurrently per file. The default +# 4 archs at ninja's default parallelism OOM-kills the build on memory-capped +# sessions (Wynton caps each user at 48 GB). Fix = few archs + bounded parallelism: +# * CMAKE_CUDA_ARCHITECTURES=80;86 targets sm_80 (A100, 3 nodes) and sm_86 +# (RTX 3090 / A40, 29 nodes) — everything a `-l compute_cap=80` request can +# land on. This is passed via CMAKE_ARGS, NOT SKBUILD_CMAKE_DEFINE: the latter +# splits entries on ';' (a literal '\;' crashes its parser), whereas CMAKE_ARGS +# is shlex-split on spaces so the ';' arch separator survives. Add archs to run +# elsewhere (8.9 L40, 9.0 H100) but each extra arch adds a concurrent cicc per +# file (~3.4 GB), so lower the parallelism below if you do. +# * 6 ninja jobs at 2 archs ~= 34 GB peak (~2 cicc/job). Must live in +# activation.env — pixi does not forward shell exports into the build. +CMAKE_BUILD_PARALLEL_LEVEL = "6" +CMAKE_ARGS = "-DCMAKE_CUDA_ARCHITECTURES=80;86" +SKBUILD_CMAKE_DEFINE = "CUDA_TOOLKIT_ROOT_DIR=$CONDA_PREFIX/targets/x86_64-linux;CUDAToolkit_ROOT=$CONDA_PREFIX/targets/x86_64-linux;USE_SYSTEM_NVTX=ON;CMAKE_INCLUDE_PATH=$CONDA_PREFIX/lib/python3.12/site-packages/nvidia/nvtx/include" + [project] authors = [{email = "karson.chrispens@ucsf.edu", name = "Karson Chrispens"}] classifiers = [ @@ -126,14 +164,14 @@ gxx_linux-64 = ">=9,<13" [tool.pixi.environments] analysis = {features = ["analysis"]} analysis-dev = {features = ["analysis", "dev"]} -boltz = {features = ["boltz"]} -boltz-analysis = {features = ["boltz", "analysis"]} -boltz-dev = {features = ["boltz", "dev"]} +boltz = {features = ["boltz", "tmol"]} +boltz-analysis = {features = ["boltz", "analysis", "tmol"]} +boltz-dev = {features = ["boltz", "dev", "tmol"]} boltz-osx = {features = ["boltz-osx", "dev", "analysis"]} -protenix = {features = ["protenix"]} -protenix-dev = {features = ["protenix", "dev"]} -rf3 = {features = ["rf3"]} -rf3-dev = {features = ["rf3", "dev"]} +protenix = {features = ["protenix", "tmol"]} +protenix-dev = {features = ["protenix", "dev", "tmol"]} +rf3 = {features = ["rf3", "tmol"]} +rf3-dev = {features = ["rf3", "dev", "tmol"]} [tool.pixi.feature.protenix.pypi-options.dependency-overrides] rdkit = ">=2024.3.5,<2025.9" @@ -144,6 +182,11 @@ rc-foundry = {editable = true, extras = ["rf3"], git = "https://github.com/k-chr [tool.pixi.pypi-dependencies] sampleworks = {editable = true, path = "."} +# tmol has no PEP517 build backend on PyPI beyond a shim that expects torch at +# build time; disable build isolation so it compiles against the env's torch. +[tool.pixi.pypi-options] +no-build-isolation = ["tmol"] + # Workspace-level override: reciprocalspaceship (sfcalculator-torch transitive dep) # caps pandas<=2.2.3 in its metadata, but runs fine on 2.3.1; protenix hard-pins # pandas==2.3.1. Standardize all envs on 2.3.1 to avoid silent downgrades. From 8470fdab8cbb2a49d3a64f96c227f47d731fec7c Mon Sep 17 00:00:00 2001 From: Utkarsh Singh Date: Mon, 13 Jul 2026 15:34:25 -0700 Subject: [PATCH 2/8] feat(rewards): add tmol plausibility reward function Add TmolPlausibilityReward, a differentiable physical-plausibility reward using tmol's beta2016 Rosetta energy. Scatters model-order coordinates into tmol's canonical per-residue grid (name-based map built once via prepare()) and scores via pose_stack_from_canonical_form + the beta2016 score function, so gradients flow to input coords (usable for gradient guidance and FK steering). --- .../core/rewards/tmol_plausibility.py | 234 ++++++++++++++++++ 1 file changed, 234 insertions(+) create mode 100644 src/sampleworks/core/rewards/tmol_plausibility.py diff --git a/src/sampleworks/core/rewards/tmol_plausibility.py b/src/sampleworks/core/rewards/tmol_plausibility.py new file mode 100644 index 0000000..0112301 --- /dev/null +++ b/src/sampleworks/core/rewards/tmol_plausibility.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from collections import Counter +from typing import Any + +import numpy as np +import torch +from jaxtyping import Float, Int +from tmol.io.pose_stack_construction import pose_stack_from_canonical_form + +"""Physical-plausibility reward via tmol (differentiable Rosetta ``beta2016`` energy). + +This reward complements the density-fit reward (:mod:`core.rewards.real_space_density`): +where the density term is the *likelihood* (per-atom agreement with the experimental map), +this term is the *prior* (inter-atomic chemistry / physical realism). It scores a conformer's +plausibility by rendering it into a tmol ``PoseStack`` and evaluating the differentiable +``beta2016`` score function, so gradients flow back to the input coordinates for gradient +guidance (and the scalar can equally be used evaluate-only for Feynman-Kac steering). + +tmol needs a fully typed molecular system (residue types + per-residue-type canonical atom +order), whereas the input coordinates arrive in SampleWorks/biotite atom order. The bridge is a +scatter map, built once in :meth:`TmolPlausibilityReward.prepare` and applied per call: + +- ``sel`` : indices, in the incoming coordinate order, of the protein atoms tmol scores +- ``res_of_atom`` : residue index each selected atom belongs to (parallel to ``sel``) +- ``can_of_atom`` : canonical slot each selected atom occupies (parallel to ``sel``) + +Hydrogens and terminal OXT atoms are left as NaN in the canonical grid; tmol's +``pose_stack_from_canonical_form`` builds them differentiably from the heavy atoms. +""" + + + +# Standard amino acids + MSE (selenomethionine). tmol can build missing leaf atoms (H, OXT) +# for these from their heavy atoms, but cannot build e.g. water hydrogens from a lone oxygen, +# so waters / ligands / modified residues are dropped from the plausibility score. +_AA20 = frozenset( + { + "ALA", "ARG", "ASN", "ASP", "CYS", "GLN", "GLU", "GLY", "HIS", "ILE", + "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL", "MSE", + } +) + + +def build_tmol_map(atom_array: Any, co: Any, device: torch.device) -> dict[str, Any]: + """Build the scatter map from a biotite ``AtomArray`` to tmol's canonical grid. + + The map is constant across a trajectory: it records, for each protein atom, where it lives + in the incoming coordinate tensor (``sel``) and where it must be placed in tmol's + per-residue canonical slots (``res_of_atom`` / ``can_of_atom``), plus the per-residue + ``chain_id`` and ``res_types`` tensors tmol needs to construct the pose. + + Parameters + ---------- + atom_array : biotite.structure.AtomArray + Atoms in the order the reward's ``__call__`` will receive coordinates. If the array + carries an ``occupancy`` annotation, zero-occupancy atoms are dropped first (mirroring + the density reward); otherwise all atoms are used. + co : tmol CanonicalOrdering + From ``tmol.io.canonical_ordering.default_canonical_ordering()``. Provides + ``restype_io_equiv_classes`` (residue-type index), ``restypes_atom_index_mapping`` + (atom-name -> canonical slot), and ``max_n_canonical_atoms``. + device : torch.device + Device for the returned tensors. + + Returns + ------- + dict + Keys ``sel``, ``res_of_atom``, ``can_of_atom`` (``long`` [n_protein_atoms]); + ``chain_id``, ``res_types`` (``int32`` [1, n_res]); ``n_res`` (int); + ``max_can`` (int, ``co.max_n_canonical_atoms``). + + Notes + ----- + Atoms whose residue type is unknown to tmol are skipped silently; recognized residues with + an unrecognized atom name are collected and printed once as ``UNMAPPED`` (they are dropped + from the score rather than raising, so a stray atom name does not abort a run). + """ + if "occupancy" in atom_array.get_annotation_categories(): + aa = atom_array[atom_array.occupancy > 0] + else: + aa = atom_array + + resname = [str(x).strip() for x in aa.res_name] + atname = [str(x).strip() for x in aa.atom_name] + chain = [str(x) for x in aa.chain_id] + resid = np.asarray(aa.res_id) + valid = _AA20 & set(co.restype_io_equiv_classes) + + sel: list[int] = [] + res_of_atom: list[int] = [] + can_of_atom: list[int] = [] + res_types_list: list[int] = [] + chain_of_res: list[str] = [] + missing: list[tuple[str, str]] = [] + + prev_key: tuple[str, int] | None = None + r = -1 + for i in range(len(aa)): + if resname[i] not in valid: + continue + key = (chain[i], int(resid[i])) + if key != prev_key: # start of a new protein residue + r += 1 + prev_key = key + res_types_list.append(co.restype_io_equiv_classes.index(resname[i])) + chain_of_res.append(chain[i]) + slot_map = co.restypes_atom_index_mapping[resname[i]] + if atname[i] not in slot_map: + missing.append((resname[i], atname[i])) + continue + sel.append(i) + res_of_atom.append(r) + can_of_atom.append(slot_map[atname[i]]) + n_res = r + 1 + + if missing: + print("UNMAPPED tmol atoms:", Counter(missing).most_common()) + + chains = list(dict.fromkeys(chain_of_res)) + return { + "sel": torch.tensor(sel, dtype=torch.long, device=device), + "res_of_atom": torch.tensor(res_of_atom, dtype=torch.long, device=device), + "can_of_atom": torch.tensor(can_of_atom, dtype=torch.long, device=device), + "chain_id": torch.tensor( + [chains.index(c) for c in chain_of_res], dtype=torch.int32, device=device + ).unsqueeze(0), + "res_types": torch.tensor(res_types_list, dtype=torch.int32, device=device).unsqueeze(0), + "n_res": n_res, + "max_can": co.max_n_canonical_atoms, + } + + +class TmolPlausibilityReward: + """Differentiable physical-plausibility reward (tmol ``beta2016`` energy). + + Satisfies ``RewardFunctionProtocol``: ``__call__`` returns a scalar loss (lower = more + plausible), differentiable w.r.t. the input coordinates. Elements, B-factors, and + occupancies are ignored -- atom identity comes from the scatter map, not the per-call + scattering inputs. + + The scatter map is built lazily via :meth:`prepare`, because it must match the atom order + of the coordinates the sampler actually feeds in (``model_atom_array or atom_array`` from + the processed structure), which is only known after model-specific preprocessing. + + Parameters + ---------- + co : tmol CanonicalOrdering + ``default_canonical_ordering()``. + pbt : tmol PackedBlockTypes + ``default_packed_block_types(device)``. + sfxn : tmol ScoreFunction + ``tmol.beta2016_score_function(device)``. + device : torch.device + Device the pose/score run on. + weight : float, optional + Multiplier on the returned energy (default 1.0). + """ + + def __init__(self, co: Any, pbt: Any, sfxn: Any, device: torch.device, weight: float = 1.0): + self.co = co + self.pbt = pbt + self.sfxn = sfxn + self.device = device + self.weight = weight + self._m: dict[str, Any] | None = None + + def prepare(self, atom_array: Any) -> "TmolPlausibilityReward": + """Build the scatter map from the array whose order matches the sampler coordinates. + + Call once before sampling with ``processed.model_atom_array or processed.atom_array``. + + Parameters + ---------- + atom_array : biotite.structure.AtomArray + The atom array in the order ``__call__`` will receive coordinates. + + Returns + ------- + TmolPlausibilityReward + ``self``, for chaining. + """ + self._m = build_tmol_map(atom_array, self.co, self.device) + return self + + def __call__( + self, + coordinates: Float[torch.Tensor, "batch n_atoms 3"], + elements: Int[torch.Tensor, "batch n_atoms"] | None = None, + b_factors: Float[torch.Tensor, "batch n_atoms"] | None = None, + occupancies: Float[torch.Tensor, "batch n_atoms"] | None = None, + unique_combinations: torch.Tensor | None = None, + inverse_indices: torch.Tensor | None = None, + ) -> Float[torch.Tensor, ""]: + """Score conformer plausibility. Call ``.backward()`` for gradients w.r.t. coordinates. + + Parameters + ---------- + coordinates : Float[torch.Tensor, "batch n_atoms 3"] + Atomic coordinates in the atom order :meth:`prepare` was built against. + elements, b_factors, occupancies, unique_combinations, inverse_indices + Accepted for protocol compatibility; ignored (atom identity comes from the map). + + Returns + ------- + Float[torch.Tensor, ""] + Scalar loss = ``weight * sum_b energy(pose_b)``. Each pose is scored independently, + so the per-particle gradient ``d loss / d coords_b`` depends only on pose ``b``. + """ + m = self._m + if m is None: + raise RuntimeError( + "TmolPlausibilityReward.prepare(atom_array) must be called before __call__." + ) + + batch = coordinates.shape[0] + # NaN-fill the canonical grid; tmol builds missing leaf atoms (H, OXT) from heavy atoms. + coords = coordinates.new_full((batch, m["n_res"], m["max_can"], 3), float("nan")) + # scatter: pick protein atoms (sel) -> place in their canonical slots (differentiable). + coords[:, m["res_of_atom"], m["can_of_atom"]] = coordinates[:, m["sel"]] + + pose = pose_stack_from_canonical_form( + self.co, + self.pbt, + m["chain_id"].expand(batch, -1), + m["res_types"].expand(batch, -1), + coords, + None, + None, + None, + ) + # render_whole_pose_scoring_module returns a per-pose total energy: shape [batch]. + energy = self.sfxn.render_whole_pose_scoring_module(pose)(pose.coords) + return self.weight * energy.sum() From cf81acc3ca4ff09be6953eecd6b76aea004d1108 Mon Sep 17 00:00:00 2001 From: Utkarsh Singh Date: Mon, 13 Jul 2026 23:41:40 -0700 Subject: [PATCH 3/8] feat(guidance): wire tmol plausibility reward into guidance pipeline Make the tmol beta2016 plausibility reward selectable and runnable end-to-end: - Add an optional reward.prepare(atom_array) hook in PureGuidance and FKSteering (structural Protocol, not inheritance): rewards that need the concrete model atom order (tmol) build their scatter map against the same array to_reward_inputs uses; density-style rewards without prepare() are unaffected. - Add get_tmol_reward_and_structure() and a reward_type branch in _run_guidance (no experimental map needed for plausibility). - Add --reward-type {density,plausibility} and --tmol-weight CLI flags plus the matching GuidanceConfig fields; make --density/--resolution optional and required only in density mode. - Log the tmol scatter-map size in prepare() for run-time visibility. Verified end-to-end (CPU): FK steering with plausibility drives beta2016 energy from ~2684 to favorable (negative) with a geometrically sane structure (min Ca-Ca 3.75 A, Rg 15.3 A) when combined with gradient normalization. --- .../core/rewards/tmol_plausibility.py | 5 ++ src/sampleworks/core/scalers/fk_steering.py | 5 ++ src/sampleworks/core/scalers/pure_guidance.py | 6 ++ .../utils/guidance_script_arguments.py | 31 ++++++- .../utils/guidance_script_utils.py | 83 +++++++++++++++++-- 5 files changed, 119 insertions(+), 11 deletions(-) diff --git a/src/sampleworks/core/rewards/tmol_plausibility.py b/src/sampleworks/core/rewards/tmol_plausibility.py index 0112301..2c9007d 100644 --- a/src/sampleworks/core/rewards/tmol_plausibility.py +++ b/src/sampleworks/core/rewards/tmol_plausibility.py @@ -6,6 +6,7 @@ import numpy as np import torch from jaxtyping import Float, Int +from loguru import logger from tmol.io.pose_stack_construction import pose_stack_from_canonical_form """Physical-plausibility reward via tmol (differentiable Rosetta ``beta2016`` energy). @@ -181,6 +182,10 @@ def prepare(self, atom_array: Any) -> "TmolPlausibilityReward": ``self``, for chaining. """ self._m = build_tmol_map(atom_array, self.co, self.device) + logger.info( + f"tmol plausibility map: {self._m['n_res']} residues, " + f"{len(self._m['sel'])} scored atoms" + ) return self def __call__( diff --git a/src/sampleworks/core/scalers/fk_steering.py b/src/sampleworks/core/scalers/fk_steering.py index 267e4cb..26cf52b 100644 --- a/src/sampleworks/core/scalers/fk_steering.py +++ b/src/sampleworks/core/scalers/fk_steering.py @@ -112,6 +112,11 @@ def sample( ensemble_size=self.ensemble_size, ) + # Optional structural hook (see PureGuidance): tmol-style rewards build their + # atom map here against the array to_reward_inputs uses; others skip. + if hasattr(reward, "prepare"): + reward.prepare(processed.model_atom_array or processed.atom_array) + reconciler = processed.reconciler.to(coords.device) reward_inputs = processed.to_reward_inputs(device=coords.device) diff --git a/src/sampleworks/core/scalers/pure_guidance.py b/src/sampleworks/core/scalers/pure_guidance.py index 46f8178..371a762 100644 --- a/src/sampleworks/core/scalers/pure_guidance.py +++ b/src/sampleworks/core/scalers/pure_guidance.py @@ -87,6 +87,12 @@ def sample( ensemble_size=self.ensemble_size, ) + # Optional structural hook (Protocol, not inheritance): rewards that need the + # concrete atom ordering of the coordinates -- e.g. tmol builds its scatter map -- + # prepare against the same array to_reward_inputs uses. Rewards without prepare skip. + if hasattr(reward, "prepare"): + reward.prepare(processed_structure.model_atom_array or processed_structure.atom_array) + reconciler = processed_structure.reconciler.to(coords.device) reward_inputs = processed_structure.to_reward_inputs(device=coords.device) diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index 45f0686..4aa79ab 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -231,6 +231,11 @@ class GuidanceConfig: alignment_reverse_diffusion: bool | None = None recycling_steps: int | None = None num_diffusion_steps: int = 200 + # Reward selection: "density" (experimental map fit) or "plausibility" (tmol beta2016 + # energy prior). tmol_weight scales the plausibility energy; both are ignored in + # density mode. density/resolution are only required in density mode (see add_generic_args). + reward_type: str = "density" + tmol_weight: float = 1.0 # DO NOT remove the **kwargs, it is for compatibility with argparse. def add_argument(self, name: str, default: Any = None, **kwargs): @@ -329,6 +334,8 @@ def from_cli( protein=args.protein, structure=args.structure, density=args.density, + reward_type=getattr(args, "reward_type", "density"), + tmol_weight=getattr(args, "tmol_weight", 1.0), model=model, guidance_type=guidance_type, log_path=getattr(args, "log_path", None) or "", @@ -413,7 +420,25 @@ def as_dict(self) -> dict[str, Any]: def add_generic_args(parser: argparse.ArgumentParser | GuidanceConfig): """Add CLI arguments shared by all models and guidance methods.""" parser.add_argument("--structure", type=str, required=True, help="Input structure") - parser.add_argument("--density", type=str, required=True, help="Input density map") + parser.add_argument( + "--density", + type=str, + default=None, + help="Input density map (required for --reward-type density)", + ) + parser.add_argument( + "--reward-type", + type=str, + default="density", + choices=["density", "plausibility"], + help="Reward: experimental density fit, or tmol beta2016 plausibility energy", + ) + parser.add_argument( + "--tmol-weight", + type=float, + default=1.0, + help="Multiplier on the tmol plausibility energy (--reward-type plausibility)", + ) parser.add_argument("--output-dir", type=str, default="output", help="Output directory") parser.add_argument( "--log-path", type=str, default=None, help="Log file path (default: output-dir/run.log)" @@ -428,8 +453,8 @@ def add_generic_args(parser: argparse.ArgumentParser | GuidanceConfig): parser.add_argument( "--resolution", type=float, - required=True, - help="Map resolution in Angstroms (required for CCP4/MRC/MAP)", + default=None, + help="Map resolution in Angstroms (required for --reward-type density with CCP4/MRC/MAP)", ) parser.add_argument("--device", type=str, default=None, help="Device (cuda/cpu, auto-detect)") parser.add_argument( diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index f477ab7..19012ee 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -269,6 +269,59 @@ def get_reward_function_and_structure( return reward_function, structure +def get_tmol_reward_and_structure( + structure_path: str | Path, + device: torch.device, + weight: float = 1.0, +) -> tuple[Any, dict[str, Any]]: + """Load a structure and build the tmol physical-plausibility reward. + + Mirrors :func:`get_reward_function_and_structure` but needs no experimental map: + the plausibility reward is a prior over chemistry, not a data likelihood. The + atom-order scatter map is built lazily in ``reward.prepare()`` (called by the + scaler once the model-order atom array is known), so it is not needed here. + + Parameters + ---------- + structure_path : str | Path + Path to the structure file (``.cif`` / ``.pdb``) to sample around. + device : torch.device + Device the tmol pose and score function run on. + weight : float, optional + Multiplier on the returned energy (default 1.0). + + Returns + ------- + tuple[TmolPlausibilityReward, dict] + The reward function and the parsed atomworks structure dict. + """ + import tmol + from tmol.io.canonical_ordering import ( + default_canonical_ordering, + default_packed_block_types, + ) + + from sampleworks.core.rewards.tmol_plausibility import TmolPlausibilityReward + + logger.debug(f"Loading structure from {structure_path}") + safe_structure_path = resolve_mixed_hetatm_atom_altlocs(Path(structure_path)) + structure = parse( + safe_structure_path, + hydrogen_policy="remove", + add_missing_atoms=False, + ccd_mirror_path=None, + ) + if str(safe_structure_path) != str(structure_path): + safe_structure_path.unlink() # delete the temporary file if one was created + + logger.info("Creating tmol plausibility reward function") + co = default_canonical_ordering() + pbt = default_packed_block_types(device) + sfxn = tmol.beta2016_score_function(device) + reward_function = TmolPlausibilityReward(co, pbt, sfxn, device, weight=weight) + return reward_function, structure + + def save_everything( args: GuidanceConfig, losses: list[Any], @@ -432,14 +485,28 @@ def _three_state_resolver(value: str | bool | None, default: bool) -> bool: # "guidance_type" is also called "scaler" in many places def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, device): """Run one configured guidance trajectory and save its outputs.""" - reward_function, structure = get_reward_function_and_structure( - args.density, # str/path to a map file. - device, # this needs to come from the global context, not the args object. - args.em, - args.loss_order, - args.resolution, - args.structure, # path/string to a structure file. - ) + # v1 reward selection (becomes a proper --reward flag later): "plausibility" uses the + # tmol beta2016 energy prior (no map needed); anything else falls back to density fit. + if getattr(args, "reward_type", "density") == "plausibility": + reward_function, structure = get_tmol_reward_and_structure( + args.structure, + device, + weight=getattr(args, "tmol_weight", 1.0), + ) + else: + if args.density is None or args.resolution is None: + raise ValueError( + "--reward-type density requires --density and --resolution " + "(they are optional only for --reward-type plausibility)." + ) + reward_function, structure = get_reward_function_and_structure( + args.density, # str/path to a map file. + device, # this needs to come from the global context, not the args object. + args.em, + args.loss_order, + args.resolution, + args.structure, # path/string to a structure file. + ) # Determine model type from wrapper class name wrapper_class_name = model_wrapper.__class__.__name__ From d004e73b11d752ba6bd6b6d9e8b261cb36cebe2a Mon Sep 17 00:00:00 2001 From: etherealsunshine Date: Tue, 14 Jul 2026 11:27:45 -0700 Subject: [PATCH 4/8] feat(rewards): add composite reward combining density and tmol plausibility Add CompositeReward: a weighted sum of RewardFunctionProtocol rewards that is itself a RewardFunctionProtocol, and forwards the optional prepare() hook to any sub-reward that defines it (so tmol's scatter map still gets built inside a composite). Wire it into the pipeline: get_composite_reward_and_structure builds the density and tmol sub-rewards (tmol term scaled by --tmol-weight) and is selectable via --reward-type composite. --- src/sampleworks/core/rewards/composite.py | 52 +++++++++++++++++++ .../utils/guidance_script_arguments.py | 6 +-- .../utils/guidance_script_utils.py | 27 ++++++++++ 3 files changed, 82 insertions(+), 3 deletions(-) create mode 100644 src/sampleworks/core/rewards/composite.py diff --git a/src/sampleworks/core/rewards/composite.py b/src/sampleworks/core/rewards/composite.py new file mode 100644 index 0000000..3f586a1 --- /dev/null +++ b/src/sampleworks/core/rewards/composite.py @@ -0,0 +1,52 @@ +from __future__ import annotations +import torch +from jaxtyping import Float, Int + +class CompositeReward: + """Weighted sum of reward functions using the RewardFunctionProtocol""" + + def __init__(self, + rewards:list, + weights: list[float]|None=None, + ): + + if weights is None: + weights = [1.0] * len(rewards) + if len(weights)!= len(rewards): + raise ValueError( + f"Got {len(rewards)} rewards but {len(weights)} weights." + ) + self.rewards = rewards + self.weights = weights + + def __call__( + self, + coordinates: Float[torch.Tensor, "batch n_atoms 3"], + elements: Int[torch.Tensor, "batch n_atoms"], + b_factors: Float[torch.Tensor, "batch n_atoms"], + occupancies: Float[torch.Tensor, "batch n_atoms"], + unique_combinations: torch.Tensor | None = None, + inverse_indices: torch.Tensor | None = None, + ) -> Float[torch.Tensor,""]: + """Weighted sum of the sub-rewards; differentiable w.r.t coordinates.""" + total = None + for reward, weight in zip(self.rewards,self.weights): + term = weight * reward( + coordinates, + elements, + b_factors, + occupancies, + unique_combinations, + inverse_indices, + ) + total = term if total is None else total+term + return total + + def prepare(self,atom_array): + for rewards in self.rewards: + if hasattr(rewards,"prepare"): + rewards.prepare(atom_array) #prepare from TmolPlausibility + return self + + + diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index 4aa79ab..90ca847 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -430,8 +430,8 @@ def add_generic_args(parser: argparse.ArgumentParser | GuidanceConfig): "--reward-type", type=str, default="density", - choices=["density", "plausibility"], - help="Reward: experimental density fit, or tmol beta2016 plausibility energy", + choices=["density", "plausibility","composite"], + help="Reward: experimental density fit, or tmol beta2016 plausibility energy, or a weighted combination", ) parser.add_argument( "--tmol-weight", @@ -520,7 +520,7 @@ def add_pure_guidance_args(parser: argparse.ArgumentParser | GuidanceConfig): "--step-scaler-type", type=str, default="noisespace", - choices=["dataspace", "noisespace", "none"], + choices=["dataspace", "noisespace", "none",], help="Type of step scaler to use: dataspace (DataSpaceDPSScaler), noisespace " "(NoiseSpaceDPSScaler), or none (NoScalingScaler)", ) diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index 19012ee..a8a0b39 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -321,6 +321,28 @@ def get_tmol_reward_and_structure( reward_function = TmolPlausibilityReward(co, pbt, sfxn, device, weight=weight) return reward_function, structure +def get_composite_reward_and_structure( + args, + device:torch.device, +)-> tuple[Any, dict[str,Any]]: + """Build a CompositeReward = density fit + tmol plausibility, weighted. + + Reuses the two single-reward constructors so we don't reimplement reward + setup; the density term needs a map(--density/--resolution)""" + + from sampleworks.core.rewards.composite import CompositeReward + + if args.density is None or args.resolution is None: + raise ValueError("--reward-type composite needs --density and --resolution for the density term.") + + density_reward, structure = get_reward_function_and_structure( + args.density, device, args.em, args.loss_order, args.resolution, args.structure, + ) + tmol_reward, _ = get_tmol_reward_and_structure(args.structure,device,weight=1.0) + + weights = [1.0, getattr(args, "tmol_weight",1.0)] + composite = CompositeReward([density_reward,tmol_reward],weights) + return composite, structure def save_everything( args: GuidanceConfig, @@ -493,6 +515,11 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic device, weight=getattr(args, "tmol_weight", 1.0), ) + elif getattr(args,"reward_type","density")=="composite": + reward_function, structure = get_composite_reward_and_structure( + args, + device + ) else: if args.density is None or args.resolution is None: raise ValueError( From a71773d12e542358addc2c4b4a729d9737986a07 Mon Sep 17 00:00:00 2001 From: etherealsunshine Date: Tue, 14 Jul 2026 11:40:56 -0700 Subject: [PATCH 5/8] docs: drop TMOL_BUILD.md, keep tmol build notes inline in pyproject The pyproject.toml tmol comments already document the source build and the cmake/CUDA flags next to the config, so the standalone doc was redundant. --- TMOL_BUILD.md | 133 ------------------------------------------------- pyproject.toml | 32 ++++-------- 2 files changed, 9 insertions(+), 156 deletions(-) delete mode 100644 TMOL_BUILD.md diff --git a/TMOL_BUILD.md b/TMOL_BUILD.md deleted file mode 100644 index 9742824..0000000 --- a/TMOL_BUILD.md +++ /dev/null @@ -1,133 +0,0 @@ -# Building & Using tmol on Wynton - -[`tmol`](https://github.com/uw-ipd/tmol) (UW-IPD) is a GPU-accelerated PyTorch -reimplementation of Rosetta's `beta_nov2016_cart` energy function with custom -C++/CUDA kernels. We use it as a reward (energy score) for guided sampling — it's -strategy-agnostic, usable with FK(C) steering, pure guidance, DPS, etc. This note -records how it's wired into the pixi environments and how to build/run it on -Wynton, because the install is non-trivial. - -## TL;DR - -- tmol is a pypi-dependency of the `tmol` pixi **feature**, attached to every - model env (`boltz*`, `protenix*`, `rf3*`). See `[tool.pixi.feature.tmol*]` in - `pyproject.toml`. -- It is **compiled from source** (no usable prebuilt wheel — see below), so a - clean build needs a CUDA toolchain (present in the linux-64 envs via - `cuda-toolkit`) and the cmake/parallelism settings baked into - `[tool.pixi.feature.tmol.activation.env]`. -- Build it on a machine with `nvcc` (e.g. **gpudev1** — a GPU is *not* required - to compile, only to run): - - ```bash - cd /wynton/home/fraserlab//samplew0rks - PIXI_CACHE_DIR=/scratch//pixi-cache pixi install -e boltz-dev - ``` - -- Run on an **sm_80 / sm_86** GPU node (that's what we compile kernels for): - - ```bash - qrsh -q gpu.q -l compute_cap=80 -pe smp 8 - cd /wynton/home/fraserlab//samplew0rks - pixi run -e boltz-dev python -c "import tmol, torch; print(tmol.__version__, torch.cuda.get_device_name())" - ``` - -## Why source-built (no wheel) - -tmol publishes an **sdist only** on PyPI; prebuilt wheels live on GitHub Releases -and are lane-specific (`python × torch-minor × CUDA`). The published GPU lanes -cover torch **2.8–2.12** only. This project is pinned to **torch 2.7.1** -(protenix hard-pins `torch==2.7.1`), so **no prebuilt wheel matches** and tmol -must compile from source. If protenix ever relaxes its torch pin and the project -moves to torch ≥ 2.8, tmol can switch to the fast prebuilt-wheel install and all -the cmake workarounds below can be deleted. - -## The build workarounds (why the config looks the way it does) - -All of these live in `[tool.pixi.feature.tmol.activation.env]` / -`[tool.pixi.pypi-options]` in `pyproject.toml`. They exist because tmol compiles -torch C++/CUDA extensions against a **conda** CUDA toolkit, which is laid out -differently than a system CUDA install. - -1. **Build isolation off** (`[tool.pixi.pypi-options] no-build-isolation`). - tmol needs the env's torch present at build time; an isolated build can't see - it. Its PEP517 build deps (`scikit-build-core`, `pybind11`, `packaging`) are - added to the feature so they're available without isolation. `ninja` comes - from `[tool.pixi.dependencies]`. - -2. **CUDA root** (`CUDA_TOOLKIT_ROOT_DIR` / `CUDAToolkit_ROOT` → - `$CONDA_PREFIX/targets/x86_64-linux`). torch's legacy `FindCUDA` looks for - headers in `/include`, but conda puts them under - `targets/x86_64-linux/include`. Without this it reports - *"cannot find the CUDA libraries."* - -3. **NVTX3 headers** (`USE_SYSTEM_NVTX=ON` + `CMAKE_INCLUDE_PATH` → - `.../site-packages/nvidia/nvtx/include`). CUDA 12 dropped the old - `nvToolsExt` library; torch needs header-only NVTX3, which conda's `cuda-nvtx` - does **not** ship — only torch's bundled `nvidia-nvtx` wheel has the headers. - Without this: *"Could NOT find nvtx3"* → dead `CUDA::nvToolsExt` target. - -4. **Skip the wheel probe** (`TMOL_DISABLE_WHEEL_FETCH=1`). tmol's sdist backend - otherwise probes GitHub Releases for a matching wheel (none exists for torch - 2.7) before falling back to a source build. - -5. **CUDA architectures** (`CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=80;86"`). - - Passed via **`CMAKE_ARGS`, not `SKBUILD_CMAKE_DEFINE`**: the latter splits - entries on `;`, so a `;`-separated arch list crashes its parser - (`not enough values to unpack ... Field cmake.define`). `CMAKE_ARGS` is - shlex-split on spaces, so the `;` survives. - - Targets **sm_80** (A100, 3 nodes) + **sm_86** (RTX 3090 / A40, 29 nodes) — - everything a `-l compute_cap=80` request can land on. - -6. **Bounded parallelism** (`CMAKE_BUILD_PARALLEL_LEVEL=6`). tmol's CUDA kernels - are template-heavy: each `nvcc` front-end (`cicc`) peaks ~3.4 GB, and it - builds **all requested archs concurrently per file**. At ninja's default - (`cores+2`) this OOM-kills the build against **Wynton's 48 GB/user cgroup**. - At 2 archs, `-j6` peaks ~24–34 GB (measured 24 GB). Must live in - `activation.env` — **pixi does not forward shell `export`s into the build - subprocess** (only vars declared in `activation.env` reach it). - -## Gotchas - -- **`PIXI_CACHE_DIR` must be on local disk (`/scratch`).** The default cache on - the NFS home fails uv's atomic renames mid-build - (`Resource busy (os error 16)`). Not committed to `pyproject.toml` because it's - per-machine; set it on the command line or in your shell profile. -- **uv caches the built wheel ignoring these settings.** If you change the arch - list (or any cmake setting) and rebuild, uv may silently reuse the old wheel - and exit 0 without recompiling. Force a real rebuild by purging the cache and - the installed copy: - ```bash - C=/scratch//pixi-cache/uv-cache - find $C/sdists-v9/pypi/tmol -iname 'tmol*.whl' -printf '%h\n' | sort -u | xargs -r rm -rf - find $C/archive-v0 -maxdepth 2 -type d -iname 'tmol*' | xargs -r -n1 dirname | sort -u | xargs -r rm -rf - rm -rf .pixi/envs/boltz-dev/lib/python3.12/site-packages/tmol* - ``` -- **Verify the binary, not the exit code.** Confirm which archs actually got - compiled in: - ```bash - cuobjdump --list-elf .pixi/envs/boltz-dev/lib/python3.12/site-packages/tmol/_C*.so | grep -oE 'sm_[0-9]+' | sort -u - ``` -- **gpudev1 can build but not run tmol.** Its GPUs are GTX 1080 (sm_61), which - isn't in our arch list — `import tmol` works there, but CUDA ops won't. - -## Adding more GPU architectures - -To also run on e.g. L40/RTX 4090 (sm_89) or H100 (sm_90), extend the arch list -in `CMAKE_ARGS` (`-DCMAKE_CUDA_ARCHITECTURES=80;86;89;90`). Each extra arch adds -a concurrent `cicc` (~3.4 GB) per file, so **lower `CMAKE_BUILD_PARALLEL_LEVEL`** -to stay under the 48 GB cap (roughly: `48 GB / (archs × 4 GB)` jobs). Then purge -the uv cache (above) so the rebuild actually happens. - -## Wynton GPU fleet (compute capability → nodes) - -| compute_cap | Arch / cards | # nodes | tmol runs? | -|---|---|---|---| -| 86 | Ampere (RTX 3090 / A40) | 29 | ✅ (built) | -| 80 | Ampere (A100) | 3 | ✅ (built) | -| 75 | Turing (2080 Ti / T4) | 14 | ❌ | -| 61 | Pascal (GTX 1080, gpudev1) | 18 | ❌ | -| ~90 | Hopper (H100) | 1 | ❌ (add sm_90) | - -Request a node with `qrsh -q gpu.q -l compute_cap=80 -pe smp ` (matches -nodes with capability ≥ 8.0; the job gets the node exclusively). diff --git a/pyproject.toml b/pyproject.toml index 39f6f26..0d38b19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,37 +35,23 @@ platforms = ["linux-64"] # tmol (Rosetta energy fn) as a reward for guided sampling — any strategy, all # model envs, linux-64. No wheel matches torch 2.7.1, so it builds from source: # scikit-build-core/pybind11/packaging are PEP517 build deps; build isolation is -# off (see [tool.pixi.pypi-options]) so nvcc sees the env's torch. See TMOL_BUILD.md. +# off (see [tool.pixi.pypi-options]) so nvcc sees the env's torch. [tool.pixi.feature.tmol] pypi-dependencies = {"tmol" = "*", "scikit-build-core" = ">=0.10", "pybind11" = ">=2.12", "packaging" = "*"} platforms = ["linux-64"] -# tmol compiles torch C++/CUDA extensions from source. Three conda-CUDA quirks -# must be papered over for torch's cmake to find everything: -# 1. torch's legacy FindCUDA wants the CUDA root at $PREFIX/targets/, -# not $PREFIX (headers live under targets/x86_64-linux/include). -# 2. CUDA 12 dropped the nvToolsExt lib; torch needs header-only NVTX3, which -# conda's cuda-nvtx does NOT ship — only torch's bundled nvidia-nvtx wheel -# has the headers. USE_SYSTEM_NVTX=ON + that include dir points cmake at them. -# 3. TMOL_DISABLE_WHEEL_FETCH skips the (futile, no torch-2.7 lane) GitHub probe. -# SKBUILD_CMAKE_DEFINE is scikit-build-core's env channel for -D cmake args. +# tmol builds torch C++/CUDA extensions from source against the conda CUDA toolkit. +# The cmake needs: the CUDA root at $CONDA_PREFIX/targets/ (torch's legacy +# FindCUDA looks there, not $CONDA_PREFIX); USE_SYSTEM_NVTX + the nvidia-nvtx include +# dir (CUDA 12 replaced nvToolsExt with header-only NVTX3, shipped only by torch's +# bundled nvidia-nvtx); and TMOL_DISABLE_WHEEL_FETCH to skip the sdist backend's +# prebuilt-wheel probe. CMAKE_CUDA_ARCHITECTURES is passed via CMAKE_ARGS because +# SKBUILD_CMAKE_DEFINE splits values on ';' and cannot carry an arch list. +# CMAKE_BUILD_PARALLEL_LEVEL bounds concurrent nvcc to avoid OOM on limited builders. [tool.pixi.feature.tmol.activation.env] CUDACXX = "$CONDA_PREFIX/bin/nvcc" CUDA_PATH = "$CONDA_PREFIX/targets/x86_64-linux" TMOL_DISABLE_WHEEL_FETCH = "1" -# tmol's CUDA kernels are template-heavy (each nvcc cicc peaks ~3.4 GB) and it -# builds every arch of CMAKE_CUDA_ARCHITECTURES concurrently per file. The default -# 4 archs at ninja's default parallelism OOM-kills the build on memory-capped -# sessions (Wynton caps each user at 48 GB). Fix = few archs + bounded parallelism: -# * CMAKE_CUDA_ARCHITECTURES=80;86 targets sm_80 (A100, 3 nodes) and sm_86 -# (RTX 3090 / A40, 29 nodes) — everything a `-l compute_cap=80` request can -# land on. This is passed via CMAKE_ARGS, NOT SKBUILD_CMAKE_DEFINE: the latter -# splits entries on ';' (a literal '\;' crashes its parser), whereas CMAKE_ARGS -# is shlex-split on spaces so the ';' arch separator survives. Add archs to run -# elsewhere (8.9 L40, 9.0 H100) but each extra arch adds a concurrent cicc per -# file (~3.4 GB), so lower the parallelism below if you do. -# * 6 ninja jobs at 2 archs ~= 34 GB peak (~2 cicc/job). Must live in -# activation.env — pixi does not forward shell exports into the build. CMAKE_BUILD_PARALLEL_LEVEL = "6" CMAKE_ARGS = "-DCMAKE_CUDA_ARCHITECTURES=80;86" SKBUILD_CMAKE_DEFINE = "CUDA_TOOLKIT_ROOT_DIR=$CONDA_PREFIX/targets/x86_64-linux;CUDAToolkit_ROOT=$CONDA_PREFIX/targets/x86_64-linux;USE_SYSTEM_NVTX=ON;CMAKE_INCLUDE_PATH=$CONDA_PREFIX/lib/python3.12/site-packages/nvidia/nvtx/include" From 3f2e8ab7cafd905b8097c3f7a219fc8ec826cf14 Mon Sep 17 00:00:00 2001 From: etherealsunshine Date: Tue, 14 Jul 2026 12:08:57 -0700 Subject: [PATCH 6/8] style: satisfy ruff and ty for the tmol/composite changes Apply ruff import-sorting/formatting and wrap long lines; add a runtime_checkable PreparableRewardProtocol and use isinstance() at the scaler prepare() hooks so ty narrows the optional hook to a callable (replaces the untyped hasattr check). --- src/sampleworks/core/rewards/composite.py | 30 ++++++++--------- src/sampleworks/core/rewards/protocol.py | 14 +++++++- .../core/rewards/tmol_plausibility.py | 27 +++++++++++++--- src/sampleworks/core/scalers/fk_steering.py | 4 +-- src/sampleworks/core/scalers/pure_guidance.py | 4 +-- .../utils/guidance_script_arguments.py | 11 +++++-- .../utils/guidance_script_utils.py | 32 +++++++++++-------- 7 files changed, 81 insertions(+), 41 deletions(-) diff --git a/src/sampleworks/core/rewards/composite.py b/src/sampleworks/core/rewards/composite.py index 3f586a1..58ef8b9 100644 --- a/src/sampleworks/core/rewards/composite.py +++ b/src/sampleworks/core/rewards/composite.py @@ -1,21 +1,22 @@ from __future__ import annotations + import torch from jaxtyping import Float, Int + class CompositeReward: """Weighted sum of reward functions using the RewardFunctionProtocol""" - def __init__(self, - rewards:list, - weights: list[float]|None=None, + def __init__( + self, + rewards: list, + weights: list[float] | None = None, ): if weights is None: weights = [1.0] * len(rewards) - if len(weights)!= len(rewards): - raise ValueError( - f"Got {len(rewards)} rewards but {len(weights)} weights." - ) + if len(weights) != len(rewards): + raise ValueError(f"Got {len(rewards)} rewards but {len(weights)} weights.") self.rewards = rewards self.weights = weights @@ -27,10 +28,10 @@ def __call__( occupancies: Float[torch.Tensor, "batch n_atoms"], unique_combinations: torch.Tensor | None = None, inverse_indices: torch.Tensor | None = None, - ) -> Float[torch.Tensor,""]: + ) -> Float[torch.Tensor, ""]: """Weighted sum of the sub-rewards; differentiable w.r.t coordinates.""" total = None - for reward, weight in zip(self.rewards,self.weights): + for reward, weight in zip(self.rewards, self.weights): term = weight * reward( coordinates, elements, @@ -39,14 +40,11 @@ def __call__( unique_combinations, inverse_indices, ) - total = term if total is None else total+term + total = term if total is None else total + term return total - def prepare(self,atom_array): + def prepare(self, atom_array): for rewards in self.rewards: - if hasattr(rewards,"prepare"): - rewards.prepare(atom_array) #prepare from TmolPlausibility + if hasattr(rewards, "prepare"): + rewards.prepare(atom_array) # prepare from TmolPlausibility return self - - - diff --git a/src/sampleworks/core/rewards/protocol.py b/src/sampleworks/core/rewards/protocol.py index 6b5fd72..079febf 100644 --- a/src/sampleworks/core/rewards/protocol.py +++ b/src/sampleworks/core/rewards/protocol.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Protocol, runtime_checkable, TYPE_CHECKING +from typing import Any, Protocol, runtime_checkable, TYPE_CHECKING import einx import numpy as np @@ -14,6 +14,18 @@ from biotite.structure import AtomArray, AtomArrayStack +@runtime_checkable +class PreparableRewardProtocol(Protocol): + """A reward with an optional one-time ``prepare`` hook. + + Rewards that need the concrete model atom ordering (e.g. tmol builds a scatter + map from it) implement ``prepare``; the sampler calls it once before sampling. + Rewards without it (e.g. density fit) simply don't match this protocol. + """ + + def prepare(self, atom_array: Any) -> Any: ... + + @dataclass class RewardInputs: """Extracted inputs for reward function computation. diff --git a/src/sampleworks/core/rewards/tmol_plausibility.py b/src/sampleworks/core/rewards/tmol_plausibility.py index 2c9007d..210fd65 100644 --- a/src/sampleworks/core/rewards/tmol_plausibility.py +++ b/src/sampleworks/core/rewards/tmol_plausibility.py @@ -9,6 +9,7 @@ from loguru import logger from tmol.io.pose_stack_construction import pose_stack_from_canonical_form + """Physical-plausibility reward via tmol (differentiable Rosetta ``beta2016`` energy). This reward complements the density-fit reward (:mod:`core.rewards.real_space_density`): @@ -31,14 +32,32 @@ """ - # Standard amino acids + MSE (selenomethionine). tmol can build missing leaf atoms (H, OXT) # for these from their heavy atoms, but cannot build e.g. water hydrogens from a lone oxygen, # so waters / ligands / modified residues are dropped from the plausibility score. _AA20 = frozenset( { - "ALA", "ARG", "ASN", "ASP", "CYS", "GLN", "GLU", "GLY", "HIS", "ILE", - "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL", "MSE", + "ALA", + "ARG", + "ASN", + "ASP", + "CYS", + "GLN", + "GLU", + "GLY", + "HIS", + "ILE", + "LEU", + "LYS", + "MET", + "PHE", + "PRO", + "SER", + "THR", + "TRP", + "TYR", + "VAL", + "MSE", } ) @@ -166,7 +185,7 @@ def __init__(self, co: Any, pbt: Any, sfxn: Any, device: torch.device, weight: f self.weight = weight self._m: dict[str, Any] | None = None - def prepare(self, atom_array: Any) -> "TmolPlausibilityReward": + def prepare(self, atom_array: Any) -> TmolPlausibilityReward: """Build the scatter map from the array whose order matches the sampler coordinates. Call once before sampling with ``processed.model_atom_array or processed.atom_array``. diff --git a/src/sampleworks/core/scalers/fk_steering.py b/src/sampleworks/core/scalers/fk_steering.py index 26cf52b..614f151 100644 --- a/src/sampleworks/core/scalers/fk_steering.py +++ b/src/sampleworks/core/scalers/fk_steering.py @@ -12,7 +12,7 @@ from loguru import logger from tqdm import tqdm -from sampleworks.core.rewards.protocol import RewardFunctionProtocol +from sampleworks.core.rewards.protocol import PreparableRewardProtocol, RewardFunctionProtocol from sampleworks.core.samplers.protocol import ( SamplerStepOutput, StepParams, @@ -114,7 +114,7 @@ def sample( # Optional structural hook (see PureGuidance): tmol-style rewards build their # atom map here against the array to_reward_inputs uses; others skip. - if hasattr(reward, "prepare"): + if isinstance(reward, PreparableRewardProtocol): reward.prepare(processed.model_atom_array or processed.atom_array) reconciler = processed.reconciler.to(coords.device) diff --git a/src/sampleworks/core/scalers/pure_guidance.py b/src/sampleworks/core/scalers/pure_guidance.py index 371a762..6a1c45e 100644 --- a/src/sampleworks/core/scalers/pure_guidance.py +++ b/src/sampleworks/core/scalers/pure_guidance.py @@ -6,7 +6,7 @@ from loguru import logger from tqdm import tqdm -from sampleworks.core.rewards.protocol import RewardFunctionProtocol +from sampleworks.core.rewards.protocol import PreparableRewardProtocol, RewardFunctionProtocol from sampleworks.core.samplers.protocol import TrajectorySampler from sampleworks.core.scalers.protocol import GuidanceOutput, StepScalerProtocol from sampleworks.eval.structure_utils import process_structure_to_trajectory_input @@ -90,7 +90,7 @@ def sample( # Optional structural hook (Protocol, not inheritance): rewards that need the # concrete atom ordering of the coordinates -- e.g. tmol builds its scatter map -- # prepare against the same array to_reward_inputs uses. Rewards without prepare skip. - if hasattr(reward, "prepare"): + if isinstance(reward, PreparableRewardProtocol): reward.prepare(processed_structure.model_atom_array or processed_structure.atom_array) reconciler = processed_structure.reconciler.to(coords.device) diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index 90ca847..5448bab 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -430,8 +430,9 @@ def add_generic_args(parser: argparse.ArgumentParser | GuidanceConfig): "--reward-type", type=str, default="density", - choices=["density", "plausibility","composite"], - help="Reward: experimental density fit, or tmol beta2016 plausibility energy, or a weighted combination", + choices=["density", "plausibility", "composite"], + help="Reward: experimental density fit, tmol beta2016 plausibility energy, " + "or a weighted combination", ) parser.add_argument( "--tmol-weight", @@ -520,7 +521,11 @@ def add_pure_guidance_args(parser: argparse.ArgumentParser | GuidanceConfig): "--step-scaler-type", type=str, default="noisespace", - choices=["dataspace", "noisespace", "none",], + choices=[ + "dataspace", + "noisespace", + "none", + ], help="Type of step scaler to use: dataspace (DataSpaceDPSScaler), noisespace " "(NoiseSpaceDPSScaler), or none (NoScalingScaler)", ) diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index a8a0b39..a7ba209 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -321,29 +321,38 @@ def get_tmol_reward_and_structure( reward_function = TmolPlausibilityReward(co, pbt, sfxn, device, weight=weight) return reward_function, structure + def get_composite_reward_and_structure( args, - device:torch.device, -)-> tuple[Any, dict[str,Any]]: + device: torch.device, +) -> tuple[Any, dict[str, Any]]: """Build a CompositeReward = density fit + tmol plausibility, weighted. - + Reuses the two single-reward constructors so we don't reimplement reward setup; the density term needs a map(--density/--resolution)""" from sampleworks.core.rewards.composite import CompositeReward if args.density is None or args.resolution is None: - raise ValueError("--reward-type composite needs --density and --resolution for the density term.") + raise ValueError( + "--reward-type composite needs --density and --resolution for the density term." + ) density_reward, structure = get_reward_function_and_structure( - args.density, device, args.em, args.loss_order, args.resolution, args.structure, + args.density, + device, + args.em, + args.loss_order, + args.resolution, + args.structure, ) - tmol_reward, _ = get_tmol_reward_and_structure(args.structure,device,weight=1.0) + tmol_reward, _ = get_tmol_reward_and_structure(args.structure, device, weight=1.0) - weights = [1.0, getattr(args, "tmol_weight",1.0)] - composite = CompositeReward([density_reward,tmol_reward],weights) + weights = [1.0, getattr(args, "tmol_weight", 1.0)] + composite = CompositeReward([density_reward, tmol_reward], weights) return composite, structure + def save_everything( args: GuidanceConfig, losses: list[Any], @@ -515,11 +524,8 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic device, weight=getattr(args, "tmol_weight", 1.0), ) - elif getattr(args,"reward_type","density")=="composite": - reward_function, structure = get_composite_reward_and_structure( - args, - device - ) + elif getattr(args, "reward_type", "density") == "composite": + reward_function, structure = get_composite_reward_and_structure(args, device) else: if args.density is None or args.resolution is None: raise ValueError( From b16ced339aa56f117965444eb5551b4e76fc11a5 Mon Sep 17 00:00:00 2001 From: etherealsunshine Date: Tue, 14 Jul 2026 12:17:41 -0700 Subject: [PATCH 7/8] test(rewards): add CompositeReward tests and tidy tmol logging Add tests/rewards/test_composite.py covering the weighted sum, default and mismatched weights, gradient flow, and prepare()-forwarding. Replace the stray print() in build_tmol_map with logger.warning to match the codebase's logging. --- .../core/rewards/tmol_plausibility.py | 4 +- tests/rewards/test_composite.py | 88 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 tests/rewards/test_composite.py diff --git a/src/sampleworks/core/rewards/tmol_plausibility.py b/src/sampleworks/core/rewards/tmol_plausibility.py index 210fd65..474775a 100644 --- a/src/sampleworks/core/rewards/tmol_plausibility.py +++ b/src/sampleworks/core/rewards/tmol_plausibility.py @@ -135,7 +135,9 @@ def build_tmol_map(atom_array: Any, co: Any, device: torch.device) -> dict[str, n_res = r + 1 if missing: - print("UNMAPPED tmol atoms:", Counter(missing).most_common()) + logger.warning( + f"tmol: dropped unmapped atoms (residue, atom): {Counter(missing).most_common()}" + ) chains = list(dict.fromkeys(chain_of_res)) return { diff --git a/tests/rewards/test_composite.py b/tests/rewards/test_composite.py new file mode 100644 index 0000000..39f4c76 --- /dev/null +++ b/tests/rewards/test_composite.py @@ -0,0 +1,88 @@ +"""Tests for CompositeReward (weighted sum of reward functions). + +Validates that CompositeReward: +1. Returns the weighted sum of its sub-rewards. +2. Defaults to equal weights of 1.0 and rejects mismatched weight counts. +3. Is differentiable w.r.t. coordinates (gradients flow through every sub-reward). +4. Forwards the optional prepare() hook only to sub-rewards that define it. +""" + +import pytest +import torch +from sampleworks.core.rewards.composite import CompositeReward + + +class _SumSquares: + """Toy reward: sum of squared coordinates.""" + + def __call__(self, coordinates, *args, **kwargs): + return coordinates.pow(2).sum() + + +class _Sum: + """Toy reward: sum of coordinates (no prepare() hook).""" + + def __call__(self, coordinates, *args, **kwargs): + return coordinates.sum() + + +class _Preparable: + """Toy reward that records the atom array passed to its prepare() hook.""" + + def __init__(self): + self.prepared_with = None + + def prepare(self, atom_array): + self.prepared_with = atom_array + return self + + def __call__(self, coordinates, *args, **kwargs): + return coordinates.sum() + + +def _coords(): + return torch.randn(2, 5, 3) + + +def test_weighted_sum_matches_manual(): + coords = _coords() + comp = CompositeReward([_SumSquares(), _Sum()], [2.0, 3.0]) + out = comp(coords, None, None, None) + expected = 2.0 * coords.pow(2).sum() + 3.0 * coords.sum() + assert torch.allclose(out, expected) + + +def test_default_weights_are_one(): + coords = _coords() + comp = CompositeReward([_SumSquares(), _Sum()]) # weights omitted -> all 1.0 + out = comp(coords, None, None, None) + expected = coords.pow(2).sum() + coords.sum() + assert torch.allclose(out, expected) + + +def test_mismatched_weights_raises(): + with pytest.raises(ValueError): + CompositeReward([_SumSquares(), _Sum()], [1.0]) + + +def test_gradient_flows_to_coordinates(): + coords = _coords().requires_grad_(True) + comp = CompositeReward([_SumSquares(), _Sum()], [2.0, 3.0]) + comp(coords, None, None, None).backward() + assert coords.grad is not None + assert torch.isfinite(coords.grad).all() + # d/dx [2*sum(x^2) + 3*sum(x)] = 4x + 3 + assert torch.allclose(coords.grad, 4.0 * coords.detach() + 3.0) + + +def test_prepare_forwards_only_to_preparable_rewards(): + prep = _Preparable() + plain = _Sum() # has no prepare() + comp = CompositeReward([prep, plain]) + atom_array = object() # sentinel + + result = comp.prepare(atom_array) + + assert result is comp # returns self for chaining + assert prep.prepared_with is atom_array # forwarded to the preparable sub-reward + # `plain` has no prepare(); the composite must skip it without error From e65eabc495c3d49197ea591ff37ab514359e5340 Mon Sep 17 00:00:00 2001 From: Utkarsh Singh Date: Tue, 14 Jul 2026 12:30:09 -0700 Subject: [PATCH 8/8] Fix Docstring syntax for increased readability Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/sampleworks/core/rewards/tmol_plausibility.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sampleworks/core/rewards/tmol_plausibility.py b/src/sampleworks/core/rewards/tmol_plausibility.py index 474775a..b013759 100644 --- a/src/sampleworks/core/rewards/tmol_plausibility.py +++ b/src/sampleworks/core/rewards/tmol_plausibility.py @@ -10,7 +10,7 @@ from tmol.io.pose_stack_construction import pose_stack_from_canonical_form -"""Physical-plausibility reward via tmol (differentiable Rosetta ``beta2016`` energy). +__doc__ = """Physical-plausibility reward via tmol (differentiable Rosetta ``beta2016`` energy). This reward complements the density-fit reward (:mod:`core.rewards.real_space_density`): where the density term is the *likelihood* (per-atom agreement with the experimental map),