From c9aa02a338426a0b712d4a5d6e2e46d3f0526368 Mon Sep 17 00:00:00 2001 From: xraymemory Date: Thu, 13 Aug 2026 16:45:26 -0400 Subject: [PATCH 1/3] docs(rewards): write down that rewards are minimized Guidance backprops the value and FK steering picks with argmin, so every reward here is really a loss. Nothing said so. Now the protocol docstring does, and points at the contract test that catches a term with the wrong sign. --- src/sampleworks/core/rewards/protocol.py | 8 ++++++++ tests/rewards/test_reward_function_contract.py | 4 ++++ 2 files changed, 12 insertions(+) diff --git a/src/sampleworks/core/rewards/protocol.py b/src/sampleworks/core/rewards/protocol.py index 6b5fd72d..4f5b15a7 100644 --- a/src/sampleworks/core/rewards/protocol.py +++ b/src/sampleworks/core/rewards/protocol.py @@ -142,6 +142,14 @@ class RewardFunctionProtocol(Protocol): Any callable that computes a scalar reward from atomic coordinates and properties can implement this protocol. + + Sign convention: the returned scalar is **minimized**. Lower is better, so + every implementation is a loss or a penalty, and a term naturally written as + a score to maximize must negate itself. Guidance backpropagates the value as + a loss and Feynman-Kac steering selects particles with ``argmin``, so a + sign-inverted term steers away from the data rather than towards it. Weighted + combinations (:class:`~sampleworks.core.rewards.composite.CompositeReward`) + are only meaningful when every term agrees on this. """ def __call__( diff --git a/tests/rewards/test_reward_function_contract.py b/tests/rewards/test_reward_function_contract.py index 19a9f727..22b7f300 100644 --- a/tests/rewards/test_reward_function_contract.py +++ b/tests/rewards/test_reward_function_contract.py @@ -6,6 +6,10 @@ so that new reward functions inherit the shared contract by adding a single bundle entry to `_REWARD_BUNDLES`. +`TestRewardCorrelation` is where the package's sign convention is enforced: rewards are +minimized, so moving away from the target must not lower the value (see +`RewardFunctionProtocol`). A term written as a score to maximize fails here. + Only *reward-agnostic* checks live here: - Absolute loss thresholds ARE shared, but the value is per-reward (see `_LOSS_THRESHOLDS`): RealSpace's loss is MSE on normalized density (sigma units), so a wrong/random model From 593e2ab527a4058c65f4a18e8ea57af28fd96632 Mon Sep 17 00:00:00 2001 From: xraymemory Date: Thu, 13 Aug 2026 16:45:26 -0400 Subject: [PATCH 2/3] feat(rewards): add a prepare hook to the reward protocol and scalers The structure-factor reward from #324 is built in two phases, but nothing in src/ ever called the second one, so it could not run from the pipeline at all. Adds PreparableRewardFunctionProtocol and a prepare_reward_if_needed helper, called from both trajectory scalers once the model atom array exists. prepare() mutates the reward and returns None. The tmol reward in #319 and the torchref one in #372 both need this hook. Also replaces an `or` fallback on an AtomArray with a reward_atom_array property. Whether an empty AtomArray is falsy is biotite's call, not ours. --- src/sampleworks/core/rewards/protocol.py | 53 ++++++++++++ src/sampleworks/core/scalers/fk_steering.py | 3 +- src/sampleworks/core/scalers/pure_guidance.py | 5 +- src/sampleworks/eval/structure_utils.py | 20 ++++- .../integration/test_pipeline_integration.py | 51 +++++++++++ tests/rewards/test_prepare_hook.py | 85 +++++++++++++++++++ 6 files changed, 211 insertions(+), 6 deletions(-) create mode 100644 tests/rewards/test_prepare_hook.py diff --git a/src/sampleworks/core/rewards/protocol.py b/src/sampleworks/core/rewards/protocol.py index 4f5b15a7..f40639d2 100644 --- a/src/sampleworks/core/rewards/protocol.py +++ b/src/sampleworks/core/rewards/protocol.py @@ -188,6 +188,59 @@ def __call__( ... +@runtime_checkable +class PreparableRewardFunctionProtocol(RewardFunctionProtocol, Protocol): + """Protocol for reward functions that must see the model topology first. + + Rewards whose forward model needs the atom ordering itself — element symbols, + residue identity, a unit cell — cannot be fully built from the input structure + file, because the model may represent the same protein with a different atom + set (see ``utils/atom_reconciler.py``). Those rewards are constructed in two + phases: ``__init__`` takes the up-front configuration, and :meth:`prepare` + binds the reward to the model atom array once sampling knows it. + """ + + def prepare(self, atom_array: AtomArray, *, device: torch.device | str = "cpu") -> None: + """Bind this reward to the model atom ordering. + + Mutates the reward in place and returns nothing. Implementations must be + re-runnable, so a caller can prepare the same reward again for a different + atom array or device. + + Parameters + ---------- + atom_array + Model-order atom array the subsequent ``__call__`` coordinates follow. + device + PyTorch device the prepared state is placed on. + """ + ... + + +def prepare_reward_if_needed( + reward: RewardFunctionProtocol, + atom_array: AtomArray, + *, + device: torch.device | str = "cpu", +) -> None: + """Prepare ``reward`` against the model topology when it asks to be prepared. + + Rewards that do not implement :class:`PreparableRewardFunctionProtocol` are + left untouched, so callers can apply this unconditionally. + + Parameters + ---------- + reward + Reward function about to be used for guidance. + atom_array + Model-order atom array the reward's coordinates will follow. + device + PyTorch device the reward's prepared state is placed on. + """ + if isinstance(reward, PreparableRewardFunctionProtocol): + reward.prepare(atom_array, device=device) + + @runtime_checkable class PrecomputableRewardFunctionProtocol(RewardFunctionProtocol, Protocol): """Protocol for reward functions with precomputation for vmap compatibility. diff --git a/src/sampleworks/core/scalers/fk_steering.py b/src/sampleworks/core/scalers/fk_steering.py index 267e4cb4..dff947f8 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 prepare_reward_if_needed, RewardFunctionProtocol from sampleworks.core.samplers.protocol import ( SamplerStepOutput, StepParams, @@ -114,6 +114,7 @@ def sample( reconciler = processed.reconciler.to(coords.device) reward_inputs = processed.to_reward_inputs(device=coords.device) + prepare_reward_if_needed(reward, processed.reward_atom_array, device=coords.device) schedule = sampler.compute_schedule(self.num_steps) loss_history: list[torch.Tensor] = [] diff --git a/src/sampleworks/core/scalers/pure_guidance.py b/src/sampleworks/core/scalers/pure_guidance.py index 46f81787..59260b16 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 prepare_reward_if_needed, 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 @@ -89,6 +89,9 @@ def sample( reconciler = processed_structure.reconciler.to(coords.device) reward_inputs = processed_structure.to_reward_inputs(device=coords.device) + prepare_reward_if_needed( + reward, processed_structure.reward_atom_array, device=coords.device + ) trajectory_denoised: list[torch.Tensor] = [] trajectory_next_step: list[torch.Tensor] = [] diff --git a/src/sampleworks/eval/structure_utils.py b/src/sampleworks/eval/structure_utils.py index 4bee8e5c..145cfa37 100644 --- a/src/sampleworks/eval/structure_utils.py +++ b/src/sampleworks/eval/structure_utils.py @@ -51,6 +51,21 @@ class SampleworksProcessedStructure: reconciler: AtomReconciler model_atom_array: AtomArray | None = None + @property + def reward_atom_array(self) -> AtomArray: + """Atom array the reward tensors and reward topology follow. + + ``model_atom_array`` is None when the model conditioning/features don't + expose a separate atom array, i.e. the model operates on the same atom set + as the input structure and the reconciler is an identity mapping. + + Returns + ------- + AtomArray + The model atom array when the model exposes one, else the structure's. + """ + return self.model_atom_array if self.model_atom_array is not None else self.atom_array + def to_reward_inputs(self, device: torch.device | str = "cpu") -> RewardInputs: """Build RewardInputs with model atom count when model atom arrays are available. @@ -70,10 +85,7 @@ def to_reward_inputs(self, device: torch.device | str = "cpu") -> RewardInputs: ------- RewardInputs """ - # model_atom_array is None when the model conditioning/features don't expose a - # separate atom array i.e. the model operates on the same atom set as the - # input structure and the reconciler is an identity mapping. - atom_array_for_rewards = self.model_atom_array or self.atom_array + atom_array_for_rewards = self.reward_atom_array reward_inputs = RewardInputs.from_atom_array( atom_array=atom_array_for_rewards, diff --git a/tests/integration/test_pipeline_integration.py b/tests/integration/test_pipeline_integration.py index 7073d005..c1f01abd 100644 --- a/tests/integration/test_pipeline_integration.py +++ b/tests/integration/test_pipeline_integration.py @@ -585,6 +585,57 @@ def test_trajectory_scaler_handles_multiple_particles( assert result.final_state is not None assert torch.isfinite(torch.as_tensor(result.final_state)).all() + @pytest.mark.parametrize( + "trajectory_scaler_type", get_all_trajectory_scalers(), ids=lambda s: s.value + ) + def test_preparable_reward_is_prepared_with_model_topology_before_first_call( + self, + trajectory_scaler_type: TrajectoryScalers, + device: torch.device, + mock_wrapper: MockFlowModelWrapper, + mock_structure: dict, + mock_step_scaler: MockStepScaler, + ): + """Two-phase rewards see the model atom array before any reward evaluation.""" + + class RecordingPreparableReward(MockGradientRewardFunction): + """Reward that records its preparation, in the order it happened.""" + + def __init__(self): + super().__init__() + self.prepared_atom_counts: list[int] = [] + self.calls_before_prepare = 0 + + def prepare(self, atom_array, *, device="cpu") -> None: + self.prepared_atom_counts.append(atom_array.array_length()) + + def __call__(self, coordinates: Tensor, *args, **kwargs) -> Tensor: + if not self.prepared_atom_counts: + self.calls_before_prepare += 1 + return super().__call__(coordinates, *args, **kwargs) + + reward = RecordingPreparableReward() + sampler = AF3EDMSampler( + EDMSamplerConfig(device=device, augmentation=False, align_to_input=False) + ) + trajectory_scaler = create_trajectory_scaler_from_type( + trajectory_scaler_type, + ensemble_size=1, + num_steps=3, + ) + + trajectory_scaler.sample( + structure=mock_structure, + model=mock_wrapper, + sampler=sampler, + step_scaler=mock_step_scaler, + reward=reward, + num_particles=1, + ) + + assert reward.prepared_atom_counts == [mock_wrapper.num_atoms] + assert reward.calls_before_prepare == 0 + class TestPartialDiffusion: """Test partial diffusion (t_start > 0) behavior.""" diff --git a/tests/rewards/test_prepare_hook.py b/tests/rewards/test_prepare_hook.py new file mode 100644 index 00000000..dadc8720 --- /dev/null +++ b/tests/rewards/test_prepare_hook.py @@ -0,0 +1,85 @@ +"""Tests for the two-phase reward preparation hook.""" + +import torch +from biotite.structure import AtomArray +from sampleworks.core.rewards.protocol import ( + PreparableRewardFunctionProtocol, + prepare_reward_if_needed, + RewardFunctionProtocol, +) + + +class PlainReward: + """Reward that is fully configured at construction time.""" + + def __call__( + self, + coordinates: torch.Tensor, + elements: torch.Tensor | None = None, + b_factors: torch.Tensor | None = None, + occupancies: torch.Tensor | None = None, + unique_combinations: torch.Tensor | None = None, + inverse_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + return (coordinates**2).sum() + + +class PreparableReward(PlainReward): + """Reward that binds to the model topology, recording what it was given.""" + + def __init__(self): + self.prepared_with: list[tuple[int, str]] = [] + self.calls_before_prepare = 0 + + def prepare(self, atom_array: AtomArray, *, device: torch.device | str = "cpu") -> None: + self.prepared_with.append((atom_array.array_length(), str(device))) + + def __call__(self, coordinates: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if not self.prepared_with: + self.calls_before_prepare += 1 + return super().__call__(coordinates) + + +def make_atom_array(n_atoms: int = 4) -> AtomArray: + """Build a minimal AtomArray of carbons at the origin.""" + atom_array = AtomArray(n_atoms) + atom_array.coord = torch.zeros(n_atoms, 3).numpy() + atom_array.element = ["C"] * n_atoms + return atom_array + + +def test_preparable_reward_satisfies_both_protocols(): + reward = PreparableReward() + + assert isinstance(reward, RewardFunctionProtocol) + assert isinstance(reward, PreparableRewardFunctionProtocol) + + +def test_plain_reward_is_not_preparable(): + assert not isinstance(PlainReward(), PreparableRewardFunctionProtocol) + + +def test_prepare_hook_forwards_atom_array_and_device(): + reward = PreparableReward() + atom_array = make_atom_array(7) + + prepare_reward_if_needed(reward, atom_array, device=torch.device("cpu")) + + assert reward.prepared_with == [(7, "cpu")] + + +def test_prepare_hook_is_a_no_op_for_rewards_that_do_not_need_it(): + reward = PlainReward() + + prepare_reward_if_needed(reward, make_atom_array(), device="cpu") + + assert reward(torch.ones(1, 4, 3)) == 12.0 + + +def test_prepare_is_rerunnable_for_a_new_topology(): + reward = PreparableReward() + + prepare_reward_if_needed(reward, make_atom_array(3), device="cpu") + prepare_reward_if_needed(reward, make_atom_array(5), device="cpu") + + assert reward.prepared_with == [(3, "cpu"), (5, "cpu")] From 741255f583b366e75b5bb1cf2df5a766f4e3fd99 Mon Sep 17 00:00:00 2001 From: xraymemory Date: Thu, 13 Aug 2026 17:31:16 -0400 Subject: [PATCH 3/3] test(rewards): make the prepare-hook test prove what it claims The reward was never actually called, so 'no calls before prepare' held trivially. Uses a real DPS step scaler, counts the calls, and drives a mismatch case where the model has four atoms and the structure five, so preparing against the wrong array fails the test. Both from CodeRabbit on #373. --- .../integration/test_pipeline_integration.py | 50 +++++++++++++++---- tests/rewards/test_prepare_hook.py | 6 +++ 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/tests/integration/test_pipeline_integration.py b/tests/integration/test_pipeline_integration.py index c1f01abd..2df06d99 100644 --- a/tests/integration/test_pipeline_integration.py +++ b/tests/integration/test_pipeline_integration.py @@ -46,8 +46,14 @@ STEP_SCALER_REGISTRY, STRUCTURES, ) -from tests.mocks import MockFlowModelWrapper, MockStepScaler +from tests.mocks import ( + MismatchCase, + MismatchCaseWrapper, + MockFlowModelWrapper, + MockStepScaler, +) from tests.mocks.rewards import MockGradientRewardFunction +from tests.utils.atom_array_builders import build_test_atom_array def create_step_context_with_reward( @@ -592,28 +598,51 @@ def test_preparable_reward_is_prepared_with_model_topology_before_first_call( self, trajectory_scaler_type: TrajectoryScalers, device: torch.device, - mock_wrapper: MockFlowModelWrapper, - mock_structure: dict, - mock_step_scaler: MockStepScaler, ): - """Two-phase rewards see the model atom array before any reward evaluation.""" + """Two-phase rewards see the model atom array before any reward evaluation. + + The model and the structure deliberately have different atom counts. A reward + prepared against the input structure would bind the wrong atom ordering, which + is how the structure-factor reward silently scores the wrong thing. + """ class RecordingPreparableReward(MockGradientRewardFunction): """Reward that records its preparation, in the order it happened.""" def __init__(self): + """Start with no preparations and no evaluations recorded.""" super().__init__() self.prepared_atom_counts: list[int] = [] + self.calls = 0 self.calls_before_prepare = 0 def prepare(self, atom_array, *, device="cpu") -> None: + """Record the size of the topology this reward was bound to.""" self.prepared_atom_counts.append(atom_array.array_length()) def __call__(self, coordinates: Tensor, *args, **kwargs) -> Tensor: + """Score as usual, counting evaluations and any that precede prepare().""" + self.calls += 1 if not self.prepared_atom_counts: self.calls_before_prepare += 1 return super().__call__(coordinates, *args, **kwargs) + struct_atom_array = build_test_atom_array( + chain_ids=["A"] * 5, + res_ids=[1, 2, 3, 4, 5], + atom_names=["N", "CA", "C", "O", "CB"], + ) + case = MismatchCase( + id="model_drops_one_atom", + description="Model represents four of the structure's five atoms.", + model_atom_array=struct_atom_array[:4].copy(), + struct_atom_array=struct_atom_array, + expected_n_common=4, + expected_has_mismatch=True, + ) + wrapper = MismatchCaseWrapper(case, device=device) + structure = {"asym_unit": struct_atom_array, "metadata": {"id": "mismatch"}} + reward = RecordingPreparableReward() sampler = AF3EDMSampler( EDMSamplerConfig(device=device, augmentation=False, align_to_input=False) @@ -625,15 +654,18 @@ def __call__(self, coordinates: Tensor, *args, **kwargs) -> Tensor: ) trajectory_scaler.sample( - structure=mock_structure, - model=mock_wrapper, + structure=structure, + model=wrapper, sampler=sampler, - step_scaler=mock_step_scaler, + # A real step scaler, so the reward is actually evaluated and the + # before-first-call assertion below has something to be true about. + step_scaler=DataSpaceDPSScaler(step_size=0.1), reward=reward, num_particles=1, ) - assert reward.prepared_atom_counts == [mock_wrapper.num_atoms] + assert reward.prepared_atom_counts == [case.n_model] + assert reward.calls > 0 assert reward.calls_before_prepare == 0 diff --git a/tests/rewards/test_prepare_hook.py b/tests/rewards/test_prepare_hook.py index dadc8720..87aca09a 100644 --- a/tests/rewards/test_prepare_hook.py +++ b/tests/rewards/test_prepare_hook.py @@ -32,6 +32,7 @@ def __init__(self): self.calls_before_prepare = 0 def prepare(self, atom_array: AtomArray, *, device: torch.device | str = "cpu") -> None: + """Record the topology and device this reward was bound to.""" self.prepared_with.append((atom_array.array_length(), str(device))) def __call__(self, coordinates: torch.Tensor, *args, **kwargs) -> torch.Tensor: @@ -49,6 +50,7 @@ def make_atom_array(n_atoms: int = 4) -> AtomArray: def test_preparable_reward_satisfies_both_protocols(): + """A reward with prepare() is still an ordinary reward function.""" reward = PreparableReward() assert isinstance(reward, RewardFunctionProtocol) @@ -56,10 +58,12 @@ def test_preparable_reward_satisfies_both_protocols(): def test_plain_reward_is_not_preparable(): + """Structural typing must not classify every reward as two-phase.""" assert not isinstance(PlainReward(), PreparableRewardFunctionProtocol) def test_prepare_hook_forwards_atom_array_and_device(): + """The reward is bound to the atom array and device the caller passed.""" reward = PreparableReward() atom_array = make_atom_array(7) @@ -69,6 +73,7 @@ def test_prepare_hook_forwards_atom_array_and_device(): def test_prepare_hook_is_a_no_op_for_rewards_that_do_not_need_it(): + """Callers can prepare unconditionally, so one-phase rewards must be untouched.""" reward = PlainReward() prepare_reward_if_needed(reward, make_atom_array(), device="cpu") @@ -77,6 +82,7 @@ def test_prepare_hook_is_a_no_op_for_rewards_that_do_not_need_it(): def test_prepare_is_rerunnable_for_a_new_topology(): + """Preparing again rebinds, so a reward can move between topologies or devices.""" reward = PreparableReward() prepare_reward_if_needed(reward, make_atom_array(3), device="cpu")