feat(rewards): add a prepare hook to the reward protocol and scalers

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__(
@@ -180,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.

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] = []

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] = []

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,

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(
@@ -585,6 +591,83 @@ 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,
+    ):
+        """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)
+        )
+        trajectory_scaler = create_trajectory_scaler_from_type(
+            trajectory_scaler_type,
+            ensemble_size=1,
+            num_steps=3,
+        )
+
+        trajectory_scaler.sample(
+            structure=structure,
+            model=wrapper,
+            sampler=sampler,
+            # 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 == [case.n_model]
+        assert reward.calls > 0
+        assert reward.calls_before_prepare == 0
+
 
 class TestPartialDiffusion:
     """Test partial diffusion (t_start > 0) behavior."""

tests/rewards/test_prepare_hook.py

@@ -0,0 +1,91 @@
+"""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:
+        """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:
+        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():
+    """A reward with prepare() is still an ordinary reward function."""
+    reward = PreparableReward()
+
+    assert isinstance(reward, RewardFunctionProtocol)
+    assert isinstance(reward, PreparableRewardFunctionProtocol)
+
+
+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)
+
+    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():
+    """Callers can prepare unconditionally, so one-phase rewards must be untouched."""
+    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():
+    """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")
+    prepare_reward_if_needed(reward, make_atom_array(5), device="cpu")
+
+    assert reward.prepared_with == [(3, "cpu"), (5, "cpu")]

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