feat(rewards): add tmol plausibility and composite rewards

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/<arch> (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 = "[email protected]", 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.

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

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.