diff --git a/pyproject.toml b/pyproject.toml index dbd4b661..b52f2c31 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,30 @@ 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. +[tool.pixi.feature.tmol] +pypi-dependencies = {"tmol" = "*", "scikit-build-core" = ">=0.10", "pybind11" = ">=2.12", "packaging" = "*"} +platforms = ["linux-64"] + +# 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" +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 = [ @@ -128,14 +152,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" @@ -146,6 +170,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. diff --git a/src/sampleworks/core/rewards/composite.py b/src/sampleworks/core/rewards/composite.py new file mode 100644 index 00000000..58ef8b9b --- /dev/null +++ b/src/sampleworks/core/rewards/composite.py @@ -0,0 +1,50 @@ +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/core/rewards/protocol.py b/src/sampleworks/core/rewards/protocol.py index 6b5fd72d..079febff 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 new file mode 100644 index 00000000..b0137595 --- /dev/null +++ b/src/sampleworks/core/rewards/tmol_plausibility.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +from collections import Counter +from typing import Any + +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 + + +__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), +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: + logger.warning( + f"tmol: dropped unmapped atoms (residue, atom): {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) + logger.info( + f"tmol plausibility map: {self._m['n_res']} residues, " + f"{len(self._m['sel'])} scored atoms" + ) + 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() diff --git a/src/sampleworks/core/scalers/fk_steering.py b/src/sampleworks/core/scalers/fk_steering.py index 267e4cb4..614f151d 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, @@ -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 isinstance(reward, PreparableRewardProtocol): + 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 46f81787..6a1c45ea 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 @@ -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 isinstance(reward, PreparableRewardProtocol): + 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 45f06862..5448bab5 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,26 @@ 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", "composite"], + help="Reward: experimental density fit, tmol beta2016 plausibility energy, " + "or a weighted combination", + ) + 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 +454,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( @@ -495,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 f477ab73..a7ba2097 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -269,6 +269,90 @@ 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 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, losses: list[Any], @@ -432,14 +516,30 @@ 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), + ) + 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( + "--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__ diff --git a/tests/rewards/test_composite.py b/tests/rewards/test_composite.py new file mode 100644 index 00000000..39f4c762 --- /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