From c9aa02a338426a0b712d4a5d6e2e46d3f0526368 Mon Sep 17 00:00:00 2001
From: xraymemory <me.anzuoni@gmail.com>
Date: Thu, 13 Aug 2026 16:45:26 -0400
Subject: [PATCH 1/7] 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 <me.anzuoni@gmail.com>
Date: Thu, 13 Aug 2026 16:45:26 -0400
Subject: [PATCH 2/7] 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 <me.anzuoni@gmail.com>
Date: Thu, 13 Aug 2026 17:31:16 -0400
Subject: [PATCH 3/7] 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")

From 3a9f5a5b875d7a3fb3185a9a43e608f1a11fc173 Mon Sep 17 00:00:00 2001
From: xraymemory <me.anzuoni@gmail.com>
Date: Thu, 13 Aug 2026 16:45:26 -0400
Subject: [PATCH 4/7] feat(rewards): add reward option schemas and a reward
 registry

Reward arguments lived in add_generic_args as if density were the only
reward we would ever have, and construction was hardwired to
RealSpaceRewardFunction. Each reward now declares its options as a frozen
dataclass and registers a lazily-imported builder, so a new reward is a
schema plus one registry entry.

Builders sit next to their rewards and raise their own missing-input errors,
naming both the flag and the config key. CLI flags derive from option names,
so flag, config key and schema cannot drift apart.

get_reward_function_and_structure splits at its seam: loading the structure
is reward-agnostic and becomes load_guidance_structure, the rest is the
density builder.

Adds Rewards.STRUCTURE_FACTOR, which #324 never got.
---
 src/sampleworks/core/rewards/options.py       | 151 +++++++++++++
 .../core/rewards/real_space_density.py        |  58 +++++
 src/sampleworks/core/rewards/registry.py      | 203 ++++++++++++++++++
 .../core/rewards/structure_factor.py          |  52 +++++
 src/sampleworks/utils/guidance_constants.py   |   1 +
 .../utils/guidance_script_utils.py            |  70 +++---
 tests/rewards/test_reward_registry.py         | 110 ++++++++++
 tests/utils/test_guidance_script_utils.py     |  27 +--
 8 files changed, 621 insertions(+), 51 deletions(-)
 create mode 100644 src/sampleworks/core/rewards/options.py
 create mode 100644 src/sampleworks/core/rewards/registry.py
 create mode 100644 tests/rewards/test_reward_registry.py

diff --git a/src/sampleworks/core/rewards/options.py b/src/sampleworks/core/rewards/options.py
new file mode 100644
index 00000000..89df6a8f
--- /dev/null
+++ b/src/sampleworks/core/rewards/options.py
@@ -0,0 +1,151 @@
+"""Per-reward option schemas.
+
+Each reward type declares its configurable options once, here, as a frozen
+dataclass. Everything else is derived from that declaration: the CLI flags
+(:func:`sampleworks.utils.guidance_script_arguments.add_reward_args`), the
+``reward_options`` schema of a reward configuration file
+(:mod:`sampleworks.core.rewards.config`), container path remapping in
+``GuidanceConfig.as_dict``, and the "unknown option" error messages.
+
+Option names are the configuration-file keys, and the CLI flag for an option is
+always its name with underscores as dashes -- ``loss_order`` is ``--loss-order``
+and nothing else. Keeping that mapping mechanical is what lets one declaration
+serve every consumer.
+
+This module deliberately imports nothing heavy. The registry has to be
+importable in every pixi environment to render ``--help`` or validate a config,
+while the reward implementations themselves pull in environment-specific stacks
+(``SFC_Torch``, the vendored qFit code); those are only imported when a reward is
+actually built.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import typing
+from dataclasses import dataclass, field
+from typing import Any
+
+
+def opt(
+    default: Any,
+    *,
+    help: str,
+    choices: tuple[Any, ...] | None = None,
+    path: bool = False,
+    json_arg: bool = False,
+) -> Any:
+    """Declare one reward option.
+
+    Parameters
+    ----------
+    default
+        Value used when neither the CLI nor a configuration file sets the option.
+        The dataclass owns every default; the argparse layer defaults to ``None``
+        so that "not passed" stays distinguishable from "passed the default".
+    help
+        Help text for the generated CLI flag.
+    choices
+        Allowed values, enforced by argparse and by config-file validation.
+    path
+        The value is a filesystem path, so it is remapped between container and
+        host paths when the run configuration is serialized.
+    json_arg
+        The value is a JSON document on the command line (for dict-valued escape
+        hatches such as ``sfcalculator_kwargs``).
+
+    Returns
+    -------
+    Any
+        A :func:`dataclasses.field` carrying the metadata above.
+    """
+    return field(
+        default=default,
+        metadata={"help": help, "choices": choices, "path": path, "json_arg": json_arg},
+    )
+
+
+@dataclass(frozen=True)
+class RealSpaceDensityOptions:
+    """Options for :class:`~sampleworks.core.rewards.real_space_density.RealSpaceRewardFunction`."""
+
+    density: str | None = opt(None, help="Input density map (CCP4/MRC/MAP or MTZ)", path=True)
+    resolution: float | None = opt(None, help="Map resolution in Angstroms")
+    loss_order: int = opt(2, help="L1 or L2 loss", choices=(1, 2))
+    em: bool = opt(False, help="Use electron (cryo-EM) scattering factors")
+
+
+@dataclass(frozen=True)
+class StructureFactorOptions:
+    """Options for the structure-factor reward.
+
+    See :class:`~sampleworks.core.rewards.structure_factor.StructureFactorRewardFunction`.
+    """
+
+    mtzfile: str | None = opt(None, help="MTZ holding the target amplitudes", path=True)
+    expcolumns: list[str] | None = opt(
+        None,
+        help="MTZ amplitude and sigma column names, e.g. --expcolumns FP SIGFP "
+        "(default: auto-detect, which requires exactly one of each)",
+    )
+    resolution: float | None = opt(
+        None, help="High-resolution limit (dmin) in Angstroms (default: the MTZ's own)"
+    )
+    scattering_factor_mode: str = opt(
+        "xray", help="SFcalculator scattering mode", choices=("xray", "cryoem")
+    )
+    bulk_solvent: str = opt(
+        "off",
+        help="Bulk-solvent treatment: off (score |Fprotein|), combined (one mask from "
+        "the combined density), or per_conformer (mean of per-conformer masks)",
+        choices=("off", "combined", "per_conformer"),
+    )
+    normalize_amplitude: bool = opt(False, help="Score normalized amplitudes (|E|) instead of |F|")
+    exclude_free_reflections: bool = opt(False, help="Score the working set only, dropping R-free")
+    batch_partition: int = opt(10, help="Ensemble chunk size for the SFcalculator batch")
+    sfcalculator_kwargs: dict | None = opt(
+        None,
+        help="Extra SFcalculator keyword arguments as JSON, e.g. '{\"n_bins\": 15}'",
+        json_arg=True,
+    )
+
+
+def option_type(options_cls: type, name: str) -> Any:
+    """Return the declared type of one option, with ``None`` stripped from unions.
+
+    ``str | None`` is reported as ``str``: optionality is expressed by the default,
+    while consumers (argparse, config coercion) need the underlying value type.
+
+    Parameters
+    ----------
+    options_cls
+        A reward's option dataclass.
+    name
+        Option name.
+
+    Returns
+    -------
+    Any
+        The option's value type, e.g. ``float``, ``bool``, ``list[str]``.
+    """
+    hint = typing.get_type_hints(options_cls)[name]
+    args = [arg for arg in typing.get_args(hint) if arg is not type(None)]
+    if not args:
+        return hint
+    return args[0] if len(args) == 1 else hint
+
+
+def path_option_names(options_cls: type) -> tuple[str, ...]:
+    """Return the options of ``options_cls`` that hold filesystem paths.
+
+    Parameters
+    ----------
+    options_cls
+        A reward's option dataclass.
+
+    Returns
+    -------
+    tuple[str, ...]
+        Names of options declared with ``path=True``.
+    """
+    return tuple(f.name for f in dataclasses.fields(options_cls) if f.metadata.get("path"))
diff --git a/src/sampleworks/core/rewards/real_space_density.py b/src/sampleworks/core/rewards/real_space_density.py
index a107b7f4..efa99f35 100644
--- a/src/sampleworks/core/rewards/real_space_density.py
+++ b/src/sampleworks/core/rewards/real_space_density.py
@@ -17,6 +17,8 @@
 from sampleworks.core.forward_models.xray.real_space_density_deps.qfit.volume import (
     XMap,
 )
+from sampleworks.core.rewards.options import RealSpaceDensityOptions
+from sampleworks.core.rewards.registry import RewardBuildContext
 from sampleworks.utils.elements import elements_to_scattering_indices
 from sampleworks.utils.torch_utils import try_gpu
 
@@ -310,3 +312,59 @@ def __call__(
         ).sum(0)  # sum over batch dimension
 
         return self.loss(density, self.transformer.xmap.array)
+
+
+def build_real_space_density_reward(
+    options: RealSpaceDensityOptions, context: RewardBuildContext
+) -> RealSpaceRewardFunction:
+    """Build the real-space density reward from its configured options.
+
+    Parameters
+    ----------
+    options
+        Density reward options: map path, resolution, loss order, and whether to
+        use electron scattering factors.
+    context
+        Run-level inputs; the parsed input structure and the target device.
+
+    Returns
+    -------
+    RealSpaceRewardFunction
+        Reward scoring the model against the given map.
+
+    Raises
+    ------
+    ValueError
+        If the map or the resolution is missing.
+    """
+    if options.density is None:
+        raise ValueError(
+            "The real_space_density reward needs a density map. Pass --density, or set "
+            "reward_options.density in a --reward-config file."
+        )
+    if options.resolution is None:
+        raise ValueError(
+            "The real_space_density reward needs a map resolution. Pass --resolution, or "
+            "set reward_options.resolution in a --reward-config file."
+        )
+
+    device = torch.device(context.device) if isinstance(context.device, str) else context.device
+
+    logger.debug(f"Loading density map from {options.density}")
+    xmap = XMap.fromfile(options.density, resolution=options.resolution)
+
+    logger.debug("Setting up scattering parameters")
+    atom_array = context.structure["asym_unit"]
+    scattering_params = setup_scattering_params(em_mode=options.em, device=device)
+
+    selection_mask = atom_array.occupancy > 0
+    logger.info(f"Selected {selection_mask.sum()} atoms with occupancy > 0")
+
+    return RealSpaceRewardFunction(
+        xmap,
+        scattering_params,
+        selection_mask,
+        em=options.em,
+        loss_order=options.loss_order,
+        device=device,
+    )
diff --git a/src/sampleworks/core/rewards/registry.py b/src/sampleworks/core/rewards/registry.py
new file mode 100644
index 00000000..2b3146e0
--- /dev/null
+++ b/src/sampleworks/core/rewards/registry.py
@@ -0,0 +1,203 @@
+"""Registry of reward types, their option schemas, and their builders.
+
+Adding a reward type is: write the reward, declare its options in
+:mod:`sampleworks.core.rewards.options`, write a ``build_*`` function next to the
+reward, and add one :class:`RewardSpec` here. The CLI, the configuration-file
+schema, and run serialization all follow from that entry -- there is no argparse
+plumbing to edit and no dispatch chain to extend.
+
+Builders are addressed by ``"module:function"`` string and imported only when a
+reward is actually built, so this module stays importable in environments that
+cannot import every reward's dependencies.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import importlib
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any, TYPE_CHECKING
+
+from sampleworks.core.rewards.options import (
+    RealSpaceDensityOptions,
+    StructureFactorOptions,
+)
+from sampleworks.utils.guidance_constants import Rewards
+
+
+if TYPE_CHECKING:
+    import torch
+    from sampleworks.core.rewards.protocol import RewardFunctionProtocol
+
+
+@dataclass(frozen=True)
+class RewardBuildContext:
+    """Run-level inputs a reward builder may need that are not reward options.
+
+    Attributes
+    ----------
+    structure
+        Atomworks-parsed structure dictionary for the input structure, loaded once
+        by the caller and shared by every reward in a run.
+    device
+        Torch device the reward runs on.
+    """
+
+    structure: dict
+    device: torch.device | str = "cpu"
+
+
+@dataclass(frozen=True)
+class RewardSpec:
+    """Everything the framework needs to know about a reward without importing it.
+
+    Attributes
+    ----------
+    name
+        The reward's :class:`~sampleworks.utils.guidance_constants.Rewards` member,
+        which is also its configuration-file key and ``--reward-type`` value.
+    options_cls
+        Option dataclass declaring this reward's configurable options.
+    builder_path
+        ``"module:function"`` address of the builder, resolved lazily.
+    description
+        One-line summary, shown in ``--help``.
+    data_path_option
+        Name of the option holding the experimental data file, if any. Lets
+        callers that resolve data per protein (grid search) inject it without
+        knowing which reward they are configuring.
+    resolution_option
+        Name of the option holding the resolution, if any. Same rationale.
+    """
+
+    name: Rewards
+    options_cls: type
+    builder_path: str
+    description: str
+    data_path_option: str | None = None
+    resolution_option: str | None = None
+
+    def builder(self) -> Callable[[Any, RewardBuildContext], RewardFunctionProtocol]:
+        """Import and return this reward's builder function.
+
+        Returns
+        -------
+        Callable
+            Builder taking ``(options, context)`` and returning a reward function.
+        """
+        module_path, function_name = self.builder_path.split(":")
+        return getattr(importlib.import_module(module_path), function_name)
+
+
+REWARD_SPECS: dict[Rewards, RewardSpec] = {
+    Rewards.REAL_SPACE_DENSITY: RewardSpec(
+        name=Rewards.REAL_SPACE_DENSITY,
+        options_cls=RealSpaceDensityOptions,
+        builder_path="sampleworks.core.rewards.real_space_density:build_real_space_density_reward",
+        description="Real-space density fit (X-ray or cryo-EM map).",
+        data_path_option="density",
+        resolution_option="resolution",
+    ),
+    Rewards.STRUCTURE_FACTOR: RewardSpec(
+        name=Rewards.STRUCTURE_FACTOR,
+        options_cls=StructureFactorOptions,
+        builder_path="sampleworks.core.rewards.structure_factor:build_structure_factor_reward",
+        description="Reciprocal-space structure-factor amplitude fit (MTZ target).",
+        data_path_option="mtzfile",
+        resolution_option="resolution",
+    ),
+}
+
+
+def get_reward_spec(reward: Rewards | str) -> RewardSpec:
+    """Look up the specification for one reward type.
+
+    Parameters
+    ----------
+    reward
+        A :class:`Rewards` member or its string value.
+
+    Returns
+    -------
+    RewardSpec
+        The registered specification.
+
+    Raises
+    ------
+    ValueError
+        If ``reward`` is not a registered reward type.
+    """
+    try:
+        return REWARD_SPECS[Rewards(reward)]
+    except ValueError:
+        raise ValueError(
+            f"Unknown reward type {str(reward)!r}. Available reward types: {reward_type_names()}."
+        ) from None
+
+
+def reward_type_names() -> list[str]:
+    """Return the registered reward type names, for ``choices`` and error messages.
+
+    Returns
+    -------
+    list[str]
+        Reward names in registration order.
+    """
+    return [spec.name.value for spec in REWARD_SPECS.values()]
+
+
+def coerce_options(spec: RewardSpec, raw: Mapping[str, Any]) -> Any:
+    """Validate a raw option mapping and materialize it as the reward's options.
+
+    Parameters
+    ----------
+    spec
+        Specification of the reward the options belong to.
+    raw
+        Option values from a CLI namespace or a configuration file. Options that
+        are absent take the dataclass default.
+
+    Returns
+    -------
+    Any
+        An instance of ``spec.options_cls``.
+
+    Raises
+    ------
+    ValueError
+        If ``raw`` names an option the reward does not have.
+    """
+    valid = {f.name for f in dataclasses.fields(spec.options_cls)}
+    unknown = sorted(set(raw) - valid)
+    if unknown:
+        raise ValueError(
+            f"Unknown option(s) {unknown} for reward '{spec.name.value}'. "
+            f"Valid options: {sorted(valid)}."
+        )
+    return spec.options_cls(**dict(raw))
+
+
+def build_single_reward(
+    reward: Rewards | str,
+    options: Mapping[str, Any],
+    context: RewardBuildContext,
+) -> RewardFunctionProtocol:
+    """Build one reward from its option mapping.
+
+    Parameters
+    ----------
+    reward
+        Reward type to build.
+    options
+        Option values for that reward; missing options take their defaults.
+    context
+        Run-level inputs (parsed structure, device).
+
+    Returns
+    -------
+    RewardFunctionProtocol
+        The constructed reward function.
+    """
+    spec = get_reward_spec(reward)
+    return spec.builder()(coerce_options(spec, options), context)
diff --git a/src/sampleworks/core/rewards/structure_factor.py b/src/sampleworks/core/rewards/structure_factor.py
index c0a52ce2..e22ebe9a 100644
--- a/src/sampleworks/core/rewards/structure_factor.py
+++ b/src/sampleworks/core/rewards/structure_factor.py
@@ -11,6 +11,8 @@
 import torch
 from jaxtyping import Bool, Complex, Float, Int
 from loguru import logger
+from sampleworks.core.rewards.options import StructureFactorOptions
+from sampleworks.core.rewards.registry import RewardBuildContext
 from sampleworks.synthetic.synthetic_utils import atomarray_to_gemmi, resolve_mtz_column
 from SFC_Torch import SFcalculator
 from SFC_Torch.io import PDBParser
@@ -625,3 +627,53 @@ def _compute_ensemble_ftotal(
             self.sfc.Fprotein_asu = self.sfc.Fprotein_asu_batch.sum(dim=0)
             self.sfc.calc_fsolvent()  # sets Fmask_HKL
         return self.sfc.calc_ftotal()  # default scales set in prepare()
+
+
+def build_structure_factor_reward(
+    options: StructureFactorOptions, context: RewardBuildContext
+) -> StructureFactorRewardFunction:
+    """Build the structure-factor reward from its configured options.
+
+    The returned reward is not yet usable: like every two-phase reward it binds to
+    the model topology in :meth:`StructureFactorRewardFunction.prepare`, which the
+    trajectory scaler calls once sampling knows the model atom array. ``context``
+    is therefore unused -- neither the input structure nor the device is settled at
+    build time.
+
+    Parameters
+    ----------
+    options
+        Structure-factor reward options; see
+        :class:`~sampleworks.core.rewards.options.StructureFactorOptions`.
+    context
+        Run-level inputs. Unused here, kept for a uniform builder signature.
+
+    Returns
+    -------
+    StructureFactorRewardFunction
+        Reward scoring the model amplitudes against the MTZ target.
+
+    Raises
+    ------
+    ValueError
+        If no MTZ was given.
+    """
+    del context  # topology and device arrive in prepare()
+
+    if options.mtzfile is None:
+        raise ValueError(
+            "The structure_factor reward needs a target MTZ. Pass --mtzfile, or set "
+            "reward_options.mtzfile in a --reward-config file."
+        )
+
+    return StructureFactorRewardFunction(
+        options.mtzfile,
+        expcolumns=options.expcolumns,
+        resolution=options.resolution,
+        scattering_factor_mode=options.scattering_factor_mode,
+        bulk_solvent=options.bulk_solvent,
+        normalize_amplitude=options.normalize_amplitude,
+        exclude_free_reflections=options.exclude_free_reflections,
+        batch_partition=options.batch_partition,
+        sfcalculator_kwargs=options.sfcalculator_kwargs,
+    )
diff --git a/src/sampleworks/utils/guidance_constants.py b/src/sampleworks/utils/guidance_constants.py
index f3be6965..d4f4fb38 100644
--- a/src/sampleworks/utils/guidance_constants.py
+++ b/src/sampleworks/utils/guidance_constants.py
@@ -53,6 +53,7 @@ class Rewards(StrEnum):
     """Enum for all RewardFunctionProtocol implementations."""
 
     REAL_SPACE_DENSITY = "real_space_density"
+    STRUCTURE_FACTOR = "structure_factor"
 
 
 class Boltz2Method(StrEnum):
diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py
index 9d1a5e61..1de71e67 100644
--- a/src/sampleworks/utils/guidance_script_utils.py
+++ b/src/sampleworks/utils/guidance_script_utils.py
@@ -17,11 +17,12 @@
 from biotite.structure.io import save_structure
 from loguru import logger
 
-from sampleworks.core.forward_models.xray.real_space_density_deps.qfit.volume import XMap
+from sampleworks.core.rewards.options import RealSpaceDensityOptions
 from sampleworks.core.rewards.real_space_density import (
+    build_real_space_density_reward,
     RealSpaceRewardFunction,
-    setup_scattering_params,
 )
+from sampleworks.core.rewards.registry import RewardBuildContext
 from sampleworks.core.samplers.edm import AF3EDMSampler, EDMSamplerConfig
 from sampleworks.core.scalers.fk_steering import FKSteering
 from sampleworks.core.scalers.pure_guidance import PureGuidance
@@ -264,16 +265,23 @@ def get_model_and_device(
     return device, model_wrapper
 
 
-# TODO: further atomize for easier testing.
-def get_reward_function_and_structure(
-    density: str | Path,
-    device: torch.device,
-    em,
-    loss_order,
-    resolution,
-    structure_path: str | Path,
-) -> tuple[RealSpaceRewardFunction, dict[str, Any]]:
-    """Load structure and density inputs and build the real-space reward function."""
+def load_guidance_structure(structure_path: str | Path) -> dict[str, Any]:
+    """Parse the input structure for a guidance run.
+
+    Reward-agnostic: every reward in a run scores the same structure, so it is
+    loaded once here and handed to the reward builders through
+    :class:`~sampleworks.core.rewards.registry.RewardBuildContext`.
+
+    Parameters
+    ----------
+    structure_path : str | Path
+        Path to the structure file (``.cif`` / ``.pdb``) to sample around.
+
+    Returns
+    -------
+    dict[str, Any]
+        Atomworks-parsed structure dictionary.
+    """
     logger.debug(f"Loading structure from {structure_path}")
     safe_structure_path = resolve_mixed_hetatm_atom_altlocs(Path(structure_path))
     structure = parse(
@@ -288,26 +296,32 @@ def get_reward_function_and_structure(
     if str(safe_structure_path) != str(structure_path):
         safe_structure_path.unlink()  # delete the temporary file if it was created
 
-    logger.debug(f"Loading density map from {density}")
-    xmap = XMap.fromfile(density, resolution=resolution)
-
-    logger.debug("Setting up scattering parameters")
+    return structure
 
-    atom_array = structure["asym_unit"]
-    scattering_params = setup_scattering_params(em_mode=em, device=device)
 
-    selection_mask = atom_array.occupancy > 0
-    n_selected = selection_mask.sum()
-    logger.info(f"Selected {n_selected} atoms with occupancy > 0")
+def get_reward_function_and_structure(
+    density: str | Path,
+    device: torch.device,
+    em,
+    loss_order,
+    resolution,
+    structure_path: str | Path,
+) -> tuple[RealSpaceRewardFunction, dict[str, Any]]:
+    """Load structure and density inputs and build the real-space reward function.
 
+    .. deprecated::
+        Build rewards through
+        :func:`sampleworks.core.rewards.registry.build_single_reward` (or
+        :func:`build_reward` for a whole run configuration) instead. Kept for
+        callers that still construct the density reward positionally.
+    """
+    structure = load_guidance_structure(structure_path)
     logger.info("Creating reward function")
-    reward_function = RealSpaceRewardFunction(
-        xmap,
-        scattering_params,
-        selection_mask,
-        em=em,
-        loss_order=loss_order,
-        device=device,
+    reward_function = build_real_space_density_reward(
+        RealSpaceDensityOptions(
+            density=str(density), resolution=resolution, loss_order=loss_order, em=em
+        ),
+        RewardBuildContext(structure=structure, device=device),
     )
     return reward_function, structure
 
diff --git a/tests/rewards/test_reward_registry.py b/tests/rewards/test_reward_registry.py
new file mode 100644
index 00000000..30288a1a
--- /dev/null
+++ b/tests/rewards/test_reward_registry.py
@@ -0,0 +1,110 @@
+"""Tests for the reward registry and its option schemas."""
+
+import dataclasses
+
+import pytest
+from sampleworks.core.rewards.options import (
+    option_type,
+    path_option_names,
+    RealSpaceDensityOptions,
+)
+from sampleworks.core.rewards.registry import (
+    build_single_reward,
+    coerce_options,
+    get_reward_spec,
+    REWARD_SPECS,
+    reward_type_names,
+    RewardBuildContext,
+)
+from sampleworks.utils.guidance_constants import Rewards
+
+
+@pytest.mark.parametrize("reward", list(REWARD_SPECS), ids=lambda r: r.value)
+class TestRegistryEntries:
+    """Contract every registered reward must satisfy."""
+
+    def test_builder_is_importable(self, reward: Rewards):
+        """A registry entry that cannot resolve its builder is a broken entry."""
+        assert callable(REWARD_SPECS[reward].builder())
+
+    def test_spec_is_keyed_by_its_own_name(self, reward: Rewards):
+        assert REWARD_SPECS[reward].name is reward
+
+    def test_declared_data_and_resolution_options_exist(self, reward: Rewards):
+        """The grid-search injection points must name real options."""
+        spec = REWARD_SPECS[reward]
+        option_names = {f.name for f in dataclasses.fields(spec.options_cls)}
+
+        for declared in (spec.data_path_option, spec.resolution_option):
+            if declared is not None:
+                assert declared in option_names
+
+    def test_file_valued_options_are_declared_as_paths(self, reward: Rewards):
+        """A path option that forgets ``path=True`` silently escapes path remapping."""
+        spec = REWARD_SPECS[reward]
+        looks_like_a_path = {
+            f.name
+            for f in dataclasses.fields(spec.options_cls)
+            if f.name.endswith(("file", "path", "density"))
+        }
+
+        assert looks_like_a_path <= set(path_option_names(spec.options_cls))
+
+    def test_options_are_constructible_with_defaults(self, reward: Rewards):
+        assert coerce_options(REWARD_SPECS[reward], {}) == REWARD_SPECS[reward].options_cls()
+
+
+def test_every_reward_enum_member_is_registered():
+    """The enum and the registry must not drift apart."""
+    assert set(REWARD_SPECS) == set(Rewards)
+    assert reward_type_names() == [r.value for r in REWARD_SPECS]
+
+
+def test_get_reward_spec_accepts_the_string_form():
+    assert get_reward_spec("structure_factor") is REWARD_SPECS[Rewards.STRUCTURE_FACTOR]
+
+
+def test_get_reward_spec_names_the_alternatives_for_an_unknown_reward():
+    with pytest.raises(ValueError, match="Unknown reward type 'diffuse_scattering'"):
+        get_reward_spec("diffuse_scattering")
+
+
+def test_coerce_options_rejects_unknown_options_and_lists_the_valid_ones():
+    with pytest.raises(ValueError, match=r"Unknown option\(s\) \['looss_order'\]"):
+        coerce_options(get_reward_spec(Rewards.REAL_SPACE_DENSITY), {"looss_order": 1})
+
+
+def test_coerce_options_materializes_defaults_for_absent_options():
+    options = coerce_options(get_reward_spec(Rewards.REAL_SPACE_DENSITY), {"resolution": 1.8})
+
+    assert options == RealSpaceDensityOptions(resolution=1.8, loss_order=2, em=False)
+
+
+def test_option_type_strips_optionality():
+    assert option_type(RealSpaceDensityOptions, "resolution") is float
+    assert option_type(RealSpaceDensityOptions, "loss_order") is int
+    assert option_type(RealSpaceDensityOptions, "em") is bool
+
+
+class TestBuilderValidation:
+    """Missing required inputs are the reward's own error to raise."""
+
+    def test_density_reward_requires_a_map(self):
+        with pytest.raises(ValueError, match="needs a density map"):
+            build_single_reward(
+                Rewards.REAL_SPACE_DENSITY,
+                {"resolution": 1.8},
+                RewardBuildContext(structure={}),
+            )
+
+    def test_density_reward_requires_a_resolution(self):
+        with pytest.raises(ValueError, match="needs a map resolution"):
+            build_single_reward(
+                Rewards.REAL_SPACE_DENSITY,
+                {"density": "map.ccp4"},
+                RewardBuildContext(structure={}),
+            )
+
+    def test_structure_factor_reward_requires_an_mtz(self):
+        with pytest.raises(ValueError, match="needs a target MTZ"):
+            build_single_reward(Rewards.STRUCTURE_FACTOR, {}, RewardBuildContext(structure={}))
diff --git a/tests/utils/test_guidance_script_utils.py b/tests/utils/test_guidance_script_utils.py
index be5b7191..c74e582d 100644
--- a/tests/utils/test_guidance_script_utils.py
+++ b/tests/utils/test_guidance_script_utils.py
@@ -2,7 +2,6 @@
 
 import json
 from pathlib import Path
-from unittest.mock import MagicMock
 
 import pytest
 import sampleworks.utils.guidance_script_utils as guidance_script_utils
@@ -12,7 +11,7 @@
     _three_state_resolver,
     _write_job_metadata,
     get_model_and_device,
-    get_reward_function_and_structure,
+    load_guidance_structure,
     save_everything,
 )
 
@@ -95,7 +94,7 @@ def test_save_everything_uses_model_atom_array_for_mismatch(tmp_path: Path):
     assert (tmp_path / "refined.cif").exists()
 
 
-def test_get_reward_function_keeps_original_structure_file(
+def test_load_guidance_structure_keeps_original_structure_file(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ):
     """The input structure file must survive when altloc resolution leaves it unchanged.
@@ -119,30 +118,12 @@ def test_get_reward_function_keeps_original_structure_file(
             "asym_unit": build_test_atom_array(n_atoms=3, with_occupancy=True)
         },
     )
-    monkeypatch.setattr(
-        "sampleworks.utils.guidance_script_utils.XMap",
-        MagicMock(fromfile=MagicMock(return_value=MagicMock())),
-    )
-    monkeypatch.setattr(
-        "sampleworks.utils.guidance_script_utils.setup_scattering_params",
-        MagicMock(return_value=MagicMock()),
-    )
-    monkeypatch.setattr(
-        "sampleworks.utils.guidance_script_utils.RealSpaceRewardFunction",
-        MagicMock(return_value=MagicMock()),
-    )
 
     # Pass the path as a string to exercise the str-vs-Path comparison.
-    get_reward_function_and_structure(
-        density="dummy_density.mrc",
-        device=torch.device("cpu"),
-        em=False,
-        loss_order=2,
-        resolution=2.0,
-        structure_path=str(structure_file),
-    )
+    structure = load_guidance_structure(str(structure_file))
 
     assert structure_file.exists(), "original structure file must not be deleted"
+    assert "asym_unit" in structure
 
 
 def test_write_job_metadata_with_job_result_appends_timing_and_status(

From db674718324beda740e24181c3b2393cb6f33b12 Mon Sep 17 00:00:00 2001
From: xraymemory <me.anzuoni@gmail.com>
Date: Thu, 13 Aug 2026 16:45:27 -0400
Subject: [PATCH 5/7] feat(rewards): parse reward configuration files (#358)

RewardConfig reads the {reward: {weight, reward_options}} mapping from #358,
as JSON, YAML or TOML. YAML goes through OmegaConf, already a dependency, so
${oc.env:VAR} works the way it does in the run presets.

Weights are 1/N when none are given and verbatim when all are. Giving only
some is an error: a default quietly disagreeing with a number someone typed
is worse than a complaint.

with_experimental_data lets grid search drop in a per-protein map or MTZ
without knowing which reward it is filling.
---
 src/sampleworks/core/rewards/config.py | 355 +++++++++++++++++++++++++
 tests/rewards/test_reward_config.py    | 212 +++++++++++++++
 2 files changed, 567 insertions(+)
 create mode 100644 src/sampleworks/core/rewards/config.py
 create mode 100644 tests/rewards/test_reward_config.py

diff --git a/src/sampleworks/core/rewards/config.py b/src/sampleworks/core/rewards/config.py
new file mode 100644
index 00000000..c5fbd5f6
--- /dev/null
+++ b/src/sampleworks/core/rewards/config.py
@@ -0,0 +1,355 @@
+"""Reward configuration: which rewards a run uses, with what weights and options.
+
+The configuration is a mapping from reward name to that reward's weight and
+options (issue #358)::
+
+    real_space_density:
+      weight: 0.4
+      reward_options:
+        loss_order: 1
+    structure_factor:
+      weight: 0.6
+      reward_options:
+        mtzfile: /data/1vme.mtz
+        bulk_solvent: combined
+
+Both ways of configuring a run produce this same structure: ``--reward-type``
+with per-option flags produces a single entry, and ``--reward-config FILE``
+produces one entry per reward in the file. Everything downstream -- building the
+rewards, serializing the run, injecting per-protein data in a grid search --
+works on :class:`RewardConfig` and does not care which surface produced it.
+
+Weights are resolved at build time: omitting them all gives every reward ``1/N``,
+so a single-reward run is unweighted and a two-reward run is a plain average
+unless the user says otherwise.
+"""
+
+from __future__ import annotations
+
+import json
+import tomllib
+from collections.abc import Mapping
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+from sampleworks.core.rewards.options import path_option_names
+from sampleworks.core.rewards.registry import coerce_options, get_reward_spec, reward_type_names
+from sampleworks.utils.guidance_constants import Rewards
+
+
+# Key holding a reward's options inside a configuration file entry.
+REWARD_OPTIONS_KEY = "reward_options"
+WEIGHT_KEY = "weight"
+
+_YAML_SUFFIXES = (".yaml", ".yml")
+
+
+@dataclass(frozen=True)
+class RewardEntry:
+    """One reward in a run: which reward, how strongly, and how configured.
+
+    Attributes
+    ----------
+    reward
+        The reward type.
+    weight
+        Multiplier on this reward's value in the combined objective. ``None``
+        means "unspecified", resolved to ``1/N`` by
+        :meth:`RewardConfig.resolved_weights`.
+    options
+        Option values for this reward. Options that are absent take the defaults
+        declared in :mod:`sampleworks.core.rewards.options`.
+    """
+
+    reward: Rewards
+    weight: float | None = None
+    options: dict[str, Any] = field(default_factory=dict)
+
+    def validated(self) -> RewardEntry:
+        """Return this entry with its options checked against the reward's schema.
+
+        Returns
+        -------
+        RewardEntry
+            The same entry; raises rather than returning on invalid options.
+
+        Raises
+        ------
+        ValueError
+            If the options name something the reward does not have, or the weight
+            is negative.
+        """
+        coerce_options(get_reward_spec(self.reward), self.options)
+        if self.weight is not None and self.weight < 0:
+            raise ValueError(
+                f"Weight for reward '{self.reward.value}' must be non-negative, got {self.weight}."
+            )
+        return self
+
+
+@dataclass(frozen=True)
+class RewardConfig:
+    """The full set of rewards a guidance run scores against."""
+
+    entries: tuple[RewardEntry, ...]
+
+    def __post_init__(self):
+        if not self.entries:
+            raise ValueError(
+                f"A reward configuration needs at least one reward. "
+                f"Available reward types: {reward_type_names()}."
+            )
+
+        duplicates = sorted({e.reward.value for e in self.entries if self._count(e.reward) > 1})
+        if duplicates:
+            raise ValueError(f"Reward(s) {duplicates} configured more than once.")
+
+        for entry in self.entries:
+            entry.validated()
+
+    def _count(self, reward: Rewards) -> int:
+        return sum(1 for entry in self.entries if entry.reward is reward)
+
+    @classmethod
+    def single(cls, reward: Rewards | str, **options: Any) -> RewardConfig:
+        """Build a one-reward configuration, as ``--reward-type`` produces.
+
+        Parameters
+        ----------
+        reward
+            The reward type to use.
+        **options
+            Option values for that reward.
+
+        Returns
+        -------
+        RewardConfig
+            Configuration holding exactly this reward.
+        """
+        return cls((RewardEntry(reward=Rewards(reward), options=dict(options)),))
+
+    @classmethod
+    def from_mapping(cls, data: Mapping[str, Any]) -> RewardConfig:
+        """Build a configuration from the ``{reward: {weight, reward_options}}`` mapping.
+
+        A bare ``{reward: {}}`` (or ``{reward: None}``) is accepted and means
+        "this reward, all defaults".
+
+        Parameters
+        ----------
+        data
+            Mapping keyed by reward name.
+
+        Returns
+        -------
+        RewardConfig
+            The parsed configuration.
+
+        Raises
+        ------
+        ValueError
+            If a reward name is unknown, an entry is not a mapping, or an entry
+            holds keys other than ``weight`` and ``reward_options``.
+        """
+        entries = []
+        for name, raw_entry in data.items():
+            reward = get_reward_spec(name).name
+            entry = {} if raw_entry is None else raw_entry
+            if not isinstance(entry, Mapping):
+                raise ValueError(
+                    f"Configuration for reward '{reward.value}' must be a mapping with "
+                    f"'{WEIGHT_KEY}' and/or '{REWARD_OPTIONS_KEY}' keys, got "
+                    f"{type(entry).__name__}."
+                )
+
+            unexpected = sorted(set(entry) - {WEIGHT_KEY, REWARD_OPTIONS_KEY})
+            if unexpected:
+                raise ValueError(
+                    f"Unexpected key(s) {unexpected} in the configuration for reward "
+                    f"'{reward.value}'. Reward options belong under '{REWARD_OPTIONS_KEY}'."
+                )
+
+            weight = entry.get(WEIGHT_KEY)
+            entries.append(
+                RewardEntry(
+                    reward=reward,
+                    weight=None if weight is None else float(weight),
+                    options=dict(entry.get(REWARD_OPTIONS_KEY) or {}),
+                )
+            )
+
+        return cls(tuple(entries))
+
+    @classmethod
+    def from_file(cls, path: str | Path) -> RewardConfig:
+        """Load a reward configuration from a JSON, YAML, or TOML file.
+
+        YAML is read through OmegaConf, so ``${oc.env:VAR}`` interpolation works
+        the same way it does in the run presets.
+
+        Parameters
+        ----------
+        path
+            Path to the configuration file; the format follows its suffix.
+
+        Returns
+        -------
+        RewardConfig
+            The parsed configuration.
+
+        Raises
+        ------
+        FileNotFoundError
+            If the file does not exist.
+        ValueError
+            If the suffix is not a supported format, or the contents are not a
+            mapping of reward names.
+        """
+        path = Path(path)
+        if not path.is_file():
+            raise FileNotFoundError(f"Reward configuration file not found: {path}")
+
+        suffix = path.suffix.lower()
+        if suffix == ".json":
+            data = json.loads(path.read_text())
+        elif suffix in _YAML_SUFFIXES:
+            from omegaconf import OmegaConf
+
+            data = OmegaConf.to_container(OmegaConf.load(path), resolve=True)
+        elif suffix == ".toml":
+            data = tomllib.loads(path.read_text())
+        else:
+            supported = ", ".join((".json", *_YAML_SUFFIXES, ".toml"))
+            raise ValueError(
+                f"Unsupported reward configuration format '{suffix or path.name}'. "
+                f"Supported formats: {supported}."
+            )
+
+        if not isinstance(data, Mapping):
+            raise ValueError(
+                f"Reward configuration in {path} must be a mapping of reward name to "
+                f"{{{WEIGHT_KEY}, {REWARD_OPTIONS_KEY}}}, got {type(data).__name__}."
+            )
+
+        return cls.from_mapping(data)
+
+    def to_mapping(self) -> dict[str, Any]:
+        """Return the configuration as the plain mapping it was parsed from.
+
+        Round-trips through :meth:`from_mapping`. Values are primitives only, so
+        the result is safe to JSON-encode and to pickle across sampleworks
+        versions.
+
+        Returns
+        -------
+        dict[str, Any]
+            Mapping keyed by reward name.
+        """
+        mapping: dict[str, Any] = {}
+        for entry in self.entries:
+            payload: dict[str, Any] = {}
+            if entry.weight is not None:
+                payload[WEIGHT_KEY] = entry.weight
+            if entry.options:
+                payload[REWARD_OPTIONS_KEY] = dict(entry.options)
+            mapping[entry.reward.value] = payload
+        return mapping
+
+    def resolved_weights(self) -> tuple[float, ...]:
+        """Resolve the per-reward weights, filling in the uniform default.
+
+        Returns
+        -------
+        tuple[float, ...]
+            One weight per entry, in order.
+
+        Raises
+        ------
+        ValueError
+            If some but not all entries carry a weight -- the uniform default
+            would silently disagree with the weights that were given.
+        """
+        weighted = [entry for entry in self.entries if entry.weight is not None]
+        if not weighted:
+            return tuple(1.0 / len(self.entries) for _ in self.entries)
+
+        if len(weighted) != len(self.entries):
+            missing = sorted(e.reward.value for e in self.entries if e.weight is None)
+            raise ValueError(
+                f"Reward(s) {missing} have no weight while others do. Give every reward a "
+                f"weight, or none of them (which weights each by 1/{len(self.entries)})."
+            )
+
+        weights = tuple(float(entry.weight) for entry in self.entries)  # ty:ignore[invalid-argument-type]
+        total = sum(weights)
+        if abs(total - 1.0) > 1e-6:
+            logger.warning(
+                f"Reward weights sum to {total:g}, not 1. Using them as given; scale them "
+                "yourself if you meant them to be relative."
+            )
+        return weights
+
+    def with_experimental_data(
+        self, *, path: str | Path | None = None, resolution: float | None = None
+    ) -> RewardConfig:
+        """Fill in per-run experimental data without knowing which reward is configured.
+
+        Grid search resolves a map or MTZ and a resolution per protein, long after
+        the reward type was chosen. Each reward declares which of its options hold
+        those (``data_path_option`` / ``resolution_option``), so they can be
+        injected generically. Options already set are left alone -- an explicit
+        value from the user or a config file wins.
+
+        Parameters
+        ----------
+        path
+            Experimental data file for this run (map or MTZ).
+        resolution
+            Resolution in Angstroms for this run.
+
+        Returns
+        -------
+        RewardConfig
+            A new configuration with the data filled in where it was missing.
+        """
+        entries = []
+        for entry in self.entries:
+            spec = get_reward_spec(entry.reward)
+            options = dict(entry.options)
+            for option_name, value in (
+                (spec.data_path_option, None if path is None else str(path)),
+                (spec.resolution_option, resolution),
+            ):
+                if option_name is not None and value is not None:
+                    options.setdefault(option_name, value)
+            entries.append(replace(entry, options=options))
+
+        return RewardConfig(tuple(entries))
+
+    def remapped_paths(self, remap: Any) -> dict[str, Any]:
+        """Return :meth:`to_mapping` with path-valued options passed through ``remap``.
+
+        Used when writing run metadata, so a run executed in a container records
+        host paths like every other path in the configuration.
+
+        Parameters
+        ----------
+        remap
+            Callable taking a path string and returning the path to record.
+
+        Returns
+        -------
+        dict[str, Any]
+            The configuration mapping, with path options remapped.
+        """
+        mapping = self.to_mapping()
+        for entry in self.entries:
+            options = mapping[entry.reward.value].get(REWARD_OPTIONS_KEY)
+            if not options:
+                continue
+            for option_name in path_option_names(get_reward_spec(entry.reward).options_cls):
+                if options.get(option_name) is not None:
+                    options[option_name] = remap(str(options[option_name]))
+        return mapping
diff --git a/tests/rewards/test_reward_config.py b/tests/rewards/test_reward_config.py
new file mode 100644
index 00000000..b8450d42
--- /dev/null
+++ b/tests/rewards/test_reward_config.py
@@ -0,0 +1,212 @@
+"""Tests for reward configuration parsing (issue #358)."""
+
+import json
+import pickle
+
+import pytest
+from sampleworks.core.rewards.config import RewardConfig, RewardEntry
+from sampleworks.utils.guidance_constants import Rewards
+
+
+ISSUE_358_YAML = """
+real_space_density:
+  weight: 0.4
+  reward_options:
+    loss_order: 1
+structure_factor:
+  weight: 0.6
+  reward_options:
+    mtzfile: /data/1vme.mtz
+    bulk_solvent: combined
+"""
+
+
+class TestParsing:
+    """The configuration shape from issue #358, in each supported format."""
+
+    def test_parses_the_documented_yaml_shape(self, tmp_path):
+        config_file = tmp_path / "rewards.yaml"
+        config_file.write_text(ISSUE_358_YAML)
+
+        config = RewardConfig.from_file(config_file)
+
+        assert config.entries == (
+            RewardEntry(Rewards.REAL_SPACE_DENSITY, 0.4, {"loss_order": 1}),
+            RewardEntry(
+                Rewards.STRUCTURE_FACTOR,
+                0.6,
+                {"mtzfile": "/data/1vme.mtz", "bulk_solvent": "combined"},
+            ),
+        )
+
+    @pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml", ".toml"])
+    def test_formats_agree(self, tmp_path, suffix):
+        """The same configuration parses identically from JSON, YAML, and TOML."""
+        mapping = {
+            "structure_factor": {
+                "weight": 1.0,
+                "reward_options": {"mtzfile": "/data/x.mtz", "batch_partition": 4},
+            }
+        }
+        config_file = tmp_path / f"rewards{suffix}"
+        if suffix == ".toml":
+            config_file.write_text(
+                "[structure_factor]\nweight = 1.0\n"
+                '[structure_factor.reward_options]\nmtzfile = "/data/x.mtz"\nbatch_partition = 4\n'
+            )
+        else:
+            config_file.write_text(json.dumps(mapping))  # valid YAML too
+
+        assert RewardConfig.from_file(config_file) == RewardConfig.from_mapping(mapping)
+
+    def test_yaml_resolves_environment_interpolation(self, tmp_path, monkeypatch):
+        monkeypatch.setenv("SW_TEST_DATA_DIR", "/mnt/data")
+        config_file = tmp_path / "rewards.yaml"
+        config_file.write_text(
+            "structure_factor:\n"
+            "  reward_options:\n"
+            "    mtzfile: ${oc.env:SW_TEST_DATA_DIR}/1vme.mtz\n"
+        )
+
+        config = RewardConfig.from_file(config_file)
+
+        assert config.entries[0].options["mtzfile"] == "/mnt/data/1vme.mtz"
+
+    def test_an_entry_may_be_empty_meaning_all_defaults(self):
+        config = RewardConfig.from_mapping({"real_space_density": None})
+
+        assert config.entries == (RewardEntry(Rewards.REAL_SPACE_DENSITY),)
+
+    def test_cli_and_file_forms_produce_the_same_configuration(self):
+        from_flags = RewardConfig.single(Rewards.REAL_SPACE_DENSITY, density="m.ccp4", loss_order=1)
+        from_file = RewardConfig.from_mapping(
+            {"real_space_density": {"reward_options": {"density": "m.ccp4", "loss_order": 1}}}
+        )
+
+        assert from_flags == from_file
+
+    def test_mapping_round_trips(self):
+        mapping = {
+            "real_space_density": {"weight": 0.4, "reward_options": {"loss_order": 1}},
+            "structure_factor": {"weight": 0.6, "reward_options": {"mtzfile": "/data/x.mtz"}},
+        }
+
+        assert RewardConfig.from_mapping(mapping).to_mapping() == mapping
+
+    def test_mapping_survives_json_and_pickle(self):
+        """Run configurations are JSON-serialized into metadata and pickled into job queues."""
+        mapping = RewardConfig.single(Rewards.STRUCTURE_FACTOR, mtzfile="/data/x.mtz").to_mapping()
+
+        assert json.loads(json.dumps(mapping)) == mapping
+        assert pickle.loads(pickle.dumps(mapping)) == mapping
+
+
+class TestValidation:
+    """Bad configurations fail at parse time, naming what to fix."""
+
+    def test_unknown_reward_name_lists_the_known_ones(self):
+        with pytest.raises(ValueError, match="Unknown reward type 'densty'"):
+            RewardConfig.from_mapping({"densty": {}})
+
+    def test_unknown_option_is_rejected(self):
+        with pytest.raises(ValueError, match=r"Unknown option\(s\) \['mtz_file'\]"):
+            RewardConfig.from_mapping(
+                {"structure_factor": {"reward_options": {"mtz_file": "/data/x.mtz"}}}
+            )
+
+    def test_options_outside_reward_options_are_rejected(self):
+        """A flat entry is the most likely mistake; say where the options go."""
+        with pytest.raises(ValueError, match="Reward options belong under 'reward_options'"):
+            RewardConfig.from_mapping({"real_space_density": {"loss_order": 1}})
+
+    def test_empty_configuration_is_rejected(self):
+        with pytest.raises(ValueError, match="needs at least one reward"):
+            RewardConfig(())
+
+    def test_negative_weight_is_rejected(self):
+        with pytest.raises(ValueError, match="must be non-negative"):
+            RewardConfig((RewardEntry(Rewards.REAL_SPACE_DENSITY, -1.0),))
+
+    def test_unsupported_file_format_is_rejected(self, tmp_path):
+        config_file = tmp_path / "rewards.ini"
+        config_file.write_text("[real_space_density]\n")
+
+        with pytest.raises(ValueError, match="Unsupported reward configuration format"):
+            RewardConfig.from_file(config_file)
+
+    def test_missing_file_is_reported_as_such(self, tmp_path):
+        with pytest.raises(FileNotFoundError, match="Reward configuration file not found"):
+            RewardConfig.from_file(tmp_path / "absent.yaml")
+
+
+class TestWeights:
+    """Weight resolution is where composition semantics live."""
+
+    def test_omitted_weights_are_uniform(self):
+        config = RewardConfig.from_mapping({"real_space_density": {}, "structure_factor": {}})
+
+        assert config.resolved_weights() == (0.5, 0.5)
+
+    def test_a_single_reward_is_unweighted(self):
+        assert RewardConfig.single(Rewards.REAL_SPACE_DENSITY).resolved_weights() == (1.0,)
+
+    def test_given_weights_are_used_as_given(self):
+        config = RewardConfig.from_mapping(
+            {"real_space_density": {"weight": 2.0}, "structure_factor": {"weight": 3.0}}
+        )
+
+        assert config.resolved_weights() == (2.0, 3.0)
+
+    def test_partially_specified_weights_are_rejected(self):
+        config = RewardConfig.from_mapping(
+            {"real_space_density": {"weight": 0.4}, "structure_factor": {}}
+        )
+
+        with pytest.raises(ValueError, match=r"\['structure_factor'\] have no weight"):
+            config.resolved_weights()
+
+    def test_the_same_reward_cannot_be_configured_twice(self):
+        with pytest.raises(ValueError, match="configured more than once"):
+            RewardConfig(
+                (
+                    RewardEntry(Rewards.REAL_SPACE_DENSITY),
+                    RewardEntry(Rewards.REAL_SPACE_DENSITY, options={"loss_order": 1}),
+                )
+            )
+
+
+class TestExperimentalDataInjection:
+    """Grid search resolves the data file per protein, after the reward is chosen."""
+
+    def test_data_lands_in_each_reward_s_own_option(self):
+        config = RewardConfig.from_mapping(
+            {"real_space_density": {}, "structure_factor": {}}
+        ).with_experimental_data(path="/data/1vme.mtz", resolution=1.8)
+
+        by_reward = {entry.reward: entry.options for entry in config.entries}
+        assert by_reward[Rewards.REAL_SPACE_DENSITY] == {
+            "density": "/data/1vme.mtz",
+            "resolution": 1.8,
+        }
+        assert by_reward[Rewards.STRUCTURE_FACTOR] == {
+            "mtzfile": "/data/1vme.mtz",
+            "resolution": 1.8,
+        }
+
+    def test_explicitly_configured_values_win(self):
+        config = RewardConfig.single(
+            Rewards.STRUCTURE_FACTOR, mtzfile="/explicit.mtz"
+        ).with_experimental_data(path="/per-protein.mtz", resolution=2.5)
+
+        assert config.entries[0].options == {"mtzfile": "/explicit.mtz", "resolution": 2.5}
+
+    def test_paths_are_remapped_for_run_metadata(self):
+        config = RewardConfig.single(
+            Rewards.STRUCTURE_FACTOR, mtzfile="/data/x.mtz", resolution=2.0
+        )
+
+        mapping = config.remapped_paths(lambda p: p.replace("/data", "/host"))
+
+        options = mapping["structure_factor"]["reward_options"]
+        assert options["mtzfile"] == "/host/x.mtz"
+        assert options["resolution"] == 2.0

From 46ec0bef5368a12069ffad04e6bd33d366a2f434 Mon Sep 17 00:00:00 2001
From: xraymemory <me.anzuoni@gmail.com>
Date: Thu, 13 Aug 2026 16:45:27 -0400
Subject: [PATCH 6/7] feat(rewards): combine weighted rewards with
 CompositeReward

build_reward turns a configuration into the reward a run scores against. A
single reward at full weight comes back as itself, so current runs keep the
gradients they have today. Anything else becomes a weighted sum.

Weights default to 1/N rather than 1, so adding a term does not quietly
scale the gradient up and change what the step size means. Negative weights
are rejected: against a minimized objective they flip a term instead of
damping it.

prepare() forwards to whichever terms need it.
---
 src/sampleworks/core/rewards/composite.py | 138 +++++++++++++++++++++
 src/sampleworks/core/rewards/config.py    |  54 +++++++-
 tests/rewards/test_composite.py           | 142 ++++++++++++++++++++++
 3 files changed, 332 insertions(+), 2 deletions(-)
 create mode 100644 src/sampleworks/core/rewards/composite.py
 create mode 100644 tests/rewards/test_composite.py

diff --git a/src/sampleworks/core/rewards/composite.py b/src/sampleworks/core/rewards/composite.py
new file mode 100644
index 00000000..54b1ac5f
--- /dev/null
+++ b/src/sampleworks/core/rewards/composite.py
@@ -0,0 +1,138 @@
+"""Weighted combination of several reward functions."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+import torch
+from jaxtyping import Float, Int
+from sampleworks.core.rewards.protocol import prepare_reward_if_needed, RewardFunctionProtocol
+
+
+if TYPE_CHECKING:
+    from biotite.structure import AtomArray
+
+
+class CompositeReward:
+    """Sum of weighted reward functions, itself a reward function.
+
+    Combines terms that score different things about the same coordinates -- a
+    density fit and a physical-plausibility prior, say, or two experimental data
+    sets. Every term follows the package's sign convention (see
+    :class:`~sampleworks.core.rewards.protocol.RewardFunctionProtocol`): values are
+    minimized, so the weighted sum is too.
+
+    Weights are the terms' relative influence on the gradient. They default to
+    ``1/len(rewards)``, which keeps the combined magnitude comparable to a single
+    reward's and so leaves the guidance step size meaning what it meant before.
+
+    Parameters
+    ----------
+    rewards
+        Reward functions to combine. Must not be empty.
+    weights
+        One weight per reward, or None (default) for uniform ``1/N`` weights.
+
+    Raises
+    ------
+    ValueError
+        If ``rewards`` is empty, ``weights`` has a different length, or any
+        weight is negative.
+    """
+
+    def __init__(
+        self,
+        rewards: Sequence[RewardFunctionProtocol],
+        weights: Sequence[float] | None = None,
+    ):
+        if not rewards:
+            raise ValueError(
+                "CompositeReward needs at least one reward function; combining none of "
+                "them has no meaningful value or gradient."
+            )
+
+        if weights is None:
+            weights = [1.0 / len(rewards)] * len(rewards)
+        elif len(weights) != len(rewards):
+            raise ValueError(
+                f"Got {len(weights)} weights for {len(rewards)} rewards; they must correspond "
+                "one to one."
+            )
+
+        negative = [w for w in weights if w < 0]
+        if negative:
+            raise ValueError(
+                f"Reward weights must be non-negative, got {negative}. A negative weight "
+                "inverts that term's sign and steers away from it."
+            )
+
+        self.rewards = list(rewards)
+        self.weights = [float(w) for w in 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, ""]:
+        """Compute the weighted sum of the component rewards.
+
+        Parameters
+        ----------
+        coordinates
+            Atomic coordinates, shape [batch, n_atoms, 3].
+        elements
+            Atomic element indices, shape [batch, n_atoms].
+        b_factors
+            Per-atom B-factors, shape [batch, n_atoms].
+        occupancies
+            Per-atom occupancies, shape [batch, n_atoms].
+        unique_combinations
+            Pre-computed unique (element, b_factor) pairs, forwarded verbatim.
+            Rewards that do not use them ignore them; they exist so a caller can
+            hoist that deduplication out of a vmap, where dynamic shapes are not
+            allowed.
+        inverse_indices
+            Indices reconstructing the per-atom values from
+            ``unique_combinations``, forwarded verbatim.
+
+        Returns
+        -------
+        Float[torch.Tensor, ""]
+            Scalar value to be minimized.
+        """
+        total = torch.zeros((), dtype=coordinates.dtype, device=coordinates.device)
+        for reward, weight in zip(self.rewards, self.weights, strict=True):
+            total = total + weight * reward(
+                coordinates,
+                elements,
+                b_factors,
+                occupancies,
+                unique_combinations,
+                inverse_indices,
+            )
+        return total
+
+    def prepare(self, atom_array: AtomArray, *, device: torch.device | str = "cpu") -> None:
+        """Prepare each component reward that needs the model topology.
+
+        Parameters
+        ----------
+        atom_array
+            Model-order atom array the coordinates will follow.
+        device
+            PyTorch device the prepared state is placed on.
+        """
+        for reward in self.rewards:
+            prepare_reward_if_needed(reward, atom_array, device=device)
+
+    def __repr__(self) -> str:
+        terms = ", ".join(
+            f"{weight:g}*{type(reward).__name__}"
+            for reward, weight in zip(self.rewards, self.weights, strict=True)
+        )
+        return f"CompositeReward({terms})"
diff --git a/src/sampleworks/core/rewards/config.py b/src/sampleworks/core/rewards/config.py
index c5fbd5f6..27998e4d 100644
--- a/src/sampleworks/core/rewards/config.py
+++ b/src/sampleworks/core/rewards/config.py
@@ -31,14 +31,24 @@
 from collections.abc import Mapping
 from dataclasses import dataclass, field, replace
 from pathlib import Path
-from typing import Any
+from typing import Any, TYPE_CHECKING
 
 from loguru import logger
 from sampleworks.core.rewards.options import path_option_names
-from sampleworks.core.rewards.registry import coerce_options, get_reward_spec, reward_type_names
+from sampleworks.core.rewards.registry import (
+    build_single_reward,
+    coerce_options,
+    get_reward_spec,
+    reward_type_names,
+    RewardBuildContext,
+)
 from sampleworks.utils.guidance_constants import Rewards
 
 
+if TYPE_CHECKING:
+    from sampleworks.core.rewards.protocol import RewardFunctionProtocol
+
+
 # Key holding a reward's options inside a configuration file entry.
 REWARD_OPTIONS_KEY = "reward_options"
 WEIGHT_KEY = "weight"
@@ -353,3 +363,43 @@ def remapped_paths(self, remap: Any) -> dict[str, Any]:
                 if options.get(option_name) is not None:
                     options[option_name] = remap(str(options[option_name]))
         return mapping
+
+
+def build_reward(config: RewardConfig, context: RewardBuildContext) -> RewardFunctionProtocol:
+    """Build the reward function a run scores against.
+
+    One configured reward at full weight is built and returned directly, so the
+    single-reward runs that are today's norm keep exactly the values and gradients
+    they had before there was a registry. Anything else becomes a
+    :class:`~sampleworks.core.rewards.composite.CompositeReward`.
+
+    Parameters
+    ----------
+    config
+        The run's reward configuration.
+    context
+        Run-level inputs (the parsed input structure, the device).
+
+    Returns
+    -------
+    RewardFunctionProtocol
+        A single reward or a weighted combination of several.
+    """
+    weights = config.resolved_weights()
+    rewards = [
+        build_single_reward(entry.reward, entry.options, context) for entry in config.entries
+    ]
+
+    if len(rewards) == 1 and weights[0] == 1.0:
+        return rewards[0]
+
+    from sampleworks.core.rewards.composite import CompositeReward
+
+    logger.info(
+        "Combining rewards: "
+        + ", ".join(
+            f"{weight:g}*{entry.reward.value}"
+            for entry, weight in zip(config.entries, weights, strict=True)
+        )
+    )
+    return CompositeReward(rewards, weights)
diff --git a/tests/rewards/test_composite.py b/tests/rewards/test_composite.py
new file mode 100644
index 00000000..3af75117
--- /dev/null
+++ b/tests/rewards/test_composite.py
@@ -0,0 +1,142 @@
+"""Tests for weighted reward combination."""
+
+import pytest
+import torch
+from sampleworks.core.rewards.composite import CompositeReward
+from sampleworks.core.rewards.config import build_reward, RewardConfig
+from sampleworks.core.rewards.protocol import RewardFunctionProtocol
+from sampleworks.core.rewards.registry import RewardBuildContext
+from sampleworks.utils.guidance_constants import Rewards
+
+
+class QuadraticReward:
+    """Loss = 0.5 * scale * ||coords||^2, so the gradient is scale * coords."""
+
+    def __init__(self, scale: float = 1.0):
+        self.scale = scale
+
+    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 0.5 * self.scale * (coordinates**2).sum()
+
+
+class PreparableQuadraticReward(QuadraticReward):
+    """A quadratic reward that also binds to the model topology."""
+
+    def __init__(self, scale: float = 1.0):
+        super().__init__(scale)
+        self.prepared_atom_counts: list[int] = []
+
+    def prepare(self, atom_array, *, device="cpu") -> None:
+        self.prepared_atom_counts.append(atom_array.array_length())
+
+
+def coords(value: float = 2.0) -> torch.Tensor:
+    return torch.full((1, 3, 3), value)
+
+
+def per_atom(n_atoms: int = 3) -> dict:
+    return dict(
+        elements=torch.ones(1, n_atoms, dtype=torch.long),
+        b_factors=torch.full((1, n_atoms), 20.0),
+        occupancies=torch.ones(1, n_atoms),
+    )
+
+
+class TestCompositeValue:
+    def test_is_a_reward_function(self):
+        assert isinstance(CompositeReward([QuadraticReward()]), RewardFunctionProtocol)
+
+    def test_value_is_the_weighted_sum_of_its_terms(self):
+        terms = [QuadraticReward(1.0), QuadraticReward(3.0)]
+        composite = CompositeReward(terms, [0.25, 0.75])
+
+        combined = composite(coords(), **per_atom())
+
+        expected = 0.25 * terms[0](coords()) + 0.75 * terms[1](coords())
+        assert torch.isclose(combined, expected)
+
+    def test_default_weights_average_the_terms(self):
+        composite = CompositeReward([QuadraticReward(1.0), QuadraticReward(3.0)])
+
+        assert torch.isclose(composite(coords(), **per_atom()), QuadraticReward(2.0)(coords()))
+
+    def test_gradient_is_the_weighted_sum_of_gradients(self):
+        composite = CompositeReward([QuadraticReward(1.0), QuadraticReward(3.0)], [0.5, 0.5])
+        x = coords().requires_grad_(True)
+
+        composite(x, **per_atom()).backward()
+
+        assert torch.allclose(x.grad, 2.0 * coords())
+
+    def test_a_single_term_is_returned_unweighted(self):
+        """A one-term composite must not quietly halve the gradient."""
+        composite = CompositeReward([QuadraticReward(2.0)])
+
+        assert torch.isclose(composite(coords(), **per_atom()), QuadraticReward(2.0)(coords()))
+
+
+class TestCompositeValidation:
+    def test_no_rewards_is_rejected(self):
+        with pytest.raises(ValueError, match="needs at least one reward function"):
+            CompositeReward([])
+
+    def test_mismatched_weight_count_is_rejected(self):
+        with pytest.raises(ValueError, match="one to one"):
+            CompositeReward([QuadraticReward()], [0.5, 0.5])
+
+    def test_negative_weight_is_rejected(self):
+        with pytest.raises(ValueError, match="must be non-negative"):
+            CompositeReward([QuadraticReward()], [-1.0])
+
+
+def test_prepare_is_forwarded_only_to_terms_that_need_it():
+    from biotite.structure import AtomArray
+
+    preparable = PreparableQuadraticReward()
+    composite = CompositeReward([QuadraticReward(), preparable])
+
+    composite.prepare(AtomArray(6), device="cpu")
+
+    assert preparable.prepared_atom_counts == [6]
+
+
+class TestBuildReward:
+    """build_reward turns a configuration into the reward a run scores against."""
+
+    def test_a_single_reward_at_full_weight_is_not_wrapped(self, monkeypatch):
+        monkeypatch.setattr(
+            "sampleworks.core.rewards.config.build_single_reward",
+            lambda reward, options, context: QuadraticReward(),
+        )
+        config = RewardConfig.single(Rewards.REAL_SPACE_DENSITY, density="m.ccp4", resolution=1.8)
+
+        reward = build_reward(config, RewardBuildContext(structure={}))
+
+        assert isinstance(reward, QuadraticReward)
+
+    def test_several_rewards_are_combined_with_their_weights(self, monkeypatch):
+        scales = {Rewards.REAL_SPACE_DENSITY: 1.0, Rewards.STRUCTURE_FACTOR: 3.0}
+        monkeypatch.setattr(
+            "sampleworks.core.rewards.config.build_single_reward",
+            lambda reward, options, context: QuadraticReward(scales[reward]),
+        )
+        config = RewardConfig.from_mapping(
+            {
+                "real_space_density": {"weight": 0.25},
+                "structure_factor": {"weight": 0.75},
+            }
+        )
+
+        reward = build_reward(config, RewardBuildContext(structure={}))
+
+        assert isinstance(reward, CompositeReward)
+        assert reward.weights == [0.25, 0.75]
+        assert torch.isclose(reward(coords(), **per_atom()), QuadraticReward(2.5)(coords()))

From 990351c54d4b46e3b1c0c7cac47e4c3cc034b5ed Mon Sep 17 00:00:00 2001
From: xraymemory <me.anzuoni@gmail.com>
Date: Thu, 13 Aug 2026 17:31:43 -0400
Subject: [PATCH 7/7] fix(rewards): keep generic option types whole, reject
 non-mapping reward_options

option_type stripped None from any hint with type args, so a bare list[str]
came back as str and its CLI flag would have lost nargs. Only unions are
unwrapped now.

reward_options holding a list reached dict() and raised TypeError, which the
CLI does not catch, so a typo in a config file printed a traceback. It is a
ValueError naming the reward now.

Both from Copilot on #374 and #375.
---
 src/sampleworks/core/rewards/config.py  | 58 ++++++++++++++++++++++++-
 src/sampleworks/core/rewards/options.py |  8 +++-
 tests/rewards/test_reward_config.py     |  6 +++
 tests/rewards/test_reward_registry.py   | 12 +++++
 4 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/src/sampleworks/core/rewards/config.py b/src/sampleworks/core/rewards/config.py
index 27998e4d..a74cd73e 100644
--- a/src/sampleworks/core/rewards/config.py
+++ b/src/sampleworks/core/rewards/config.py
@@ -26,6 +26,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 import json
 import tomllib
 from collections.abc import Mapping
@@ -181,12 +182,19 @@ def from_mapping(cls, data: Mapping[str, Any]) -> RewardConfig:
                     f"'{reward.value}'. Reward options belong under '{REWARD_OPTIONS_KEY}'."
                 )
 
+            raw_options = entry.get(REWARD_OPTIONS_KEY) or {}
+            if not isinstance(raw_options, Mapping):
+                raise ValueError(
+                    f"'{REWARD_OPTIONS_KEY}' for reward '{reward.value}' must be a mapping of "
+                    f"option name to value, got {type(raw_options).__name__}."
+                )
+
             weight = entry.get(WEIGHT_KEY)
             entries.append(
                 RewardEntry(
                     reward=reward,
                     weight=None if weight is None else float(weight),
-                    options=dict(entry.get(REWARD_OPTIONS_KEY) or {}),
+                    options=dict(raw_options),
                 )
             )
 
@@ -267,6 +275,54 @@ def to_mapping(self) -> dict[str, Any]:
             mapping[entry.reward.value] = payload
         return mapping
 
+    def with_effective_options(self) -> RewardConfig:
+        """Return this configuration with every reward's defaults written out.
+
+        A run's metadata should record what actually ran, not only what was typed:
+        defaults change between versions, and an option that was defaulted is
+        otherwise indistinguishable from one that did not exist. Options left at
+        ``None`` stay absent, so they remain fillable by
+        :meth:`with_experimental_data`.
+
+        Returns
+        -------
+        RewardConfig
+            The same rewards, with defaulted option values materialized.
+        """
+        entries = []
+        for entry in self.entries:
+            options = coerce_options(get_reward_spec(entry.reward), entry.options)
+            effective = {
+                name: value
+                for name, value in dataclasses.asdict(options).items()
+                if value is not None
+            }
+            entries.append(replace(entry, options=effective))
+        return RewardConfig(tuple(entries))
+
+    def missing_required_options(self) -> dict[str, tuple[str, ...]]:
+        """Report configured rewards that are still missing an input they need.
+
+        Lets a caller refuse a configuration before doing expensive work, without
+        knowing anything about individual rewards. The builders check the same
+        thing when they run, which is what protects callers that never come
+        through here.
+
+        Returns
+        -------
+        dict[str, tuple[str, ...]]
+            Reward name to the options it is missing, for rewards missing any.
+        """
+        missing = {}
+        for entry in self.entries:
+            spec = get_reward_spec(entry.reward)
+            absent = tuple(
+                name for name in spec.required_options if entry.options.get(name) is None
+            )
+            if absent:
+                missing[entry.reward.value] = absent
+        return missing
+
     def resolved_weights(self) -> tuple[float, ...]:
         """Resolve the per-reward weights, filling in the uniform default.
 
diff --git a/src/sampleworks/core/rewards/options.py b/src/sampleworks/core/rewards/options.py
index 89df6a8f..5e2bab5e 100644
--- a/src/sampleworks/core/rewards/options.py
+++ b/src/sampleworks/core/rewards/options.py
@@ -22,6 +22,7 @@
 from __future__ import annotations
 
 import dataclasses
+import types
 import typing
 from dataclasses import dataclass, field
 from typing import Any
@@ -115,6 +116,8 @@ def option_type(options_cls: type, name: str) -> Any:
 
     ``str | None`` is reported as ``str``: optionality is expressed by the default,
     while consumers (argparse, config coercion) need the underlying value type.
+    Non-union hints are returned whole, so ``list[str]`` stays ``list[str]`` rather
+    than collapsing to its element type.
 
     Parameters
     ----------
@@ -129,9 +132,10 @@ def option_type(options_cls: type, name: str) -> Any:
         The option's value type, e.g. ``float``, ``bool``, ``list[str]``.
     """
     hint = typing.get_type_hints(options_cls)[name]
-    args = [arg for arg in typing.get_args(hint) if arg is not type(None)]
-    if not args:
+    if typing.get_origin(hint) not in (types.UnionType, typing.Union):
         return hint
+
+    args = [arg for arg in typing.get_args(hint) if arg is not type(None)]
     return args[0] if len(args) == 1 else hint
 
 
diff --git a/tests/rewards/test_reward_config.py b/tests/rewards/test_reward_config.py
index b8450d42..1699735a 100644
--- a/tests/rewards/test_reward_config.py
+++ b/tests/rewards/test_reward_config.py
@@ -114,6 +114,12 @@ def test_unknown_option_is_rejected(self):
                 {"structure_factor": {"reward_options": {"mtz_file": "/data/x.mtz"}}}
             )
 
+    def test_reward_options_must_be_a_mapping(self):
+        with pytest.raises(ValueError, match="must be a mapping of option name to value"):
+            RewardConfig.from_mapping(
+                {"real_space_density": {"reward_options": ["density", "x.ccp4"]}}
+            )
+
     def test_options_outside_reward_options_are_rejected(self):
         """A flat entry is the most likely mistake; say where the options go."""
         with pytest.raises(ValueError, match="Reward options belong under 'reward_options'"):
diff --git a/tests/rewards/test_reward_registry.py b/tests/rewards/test_reward_registry.py
index 30288a1a..e88a539f 100644
--- a/tests/rewards/test_reward_registry.py
+++ b/tests/rewards/test_reward_registry.py
@@ -86,6 +86,18 @@ def test_option_type_strips_optionality():
     assert option_type(RealSpaceDensityOptions, "em") is bool
 
 
+def test_option_type_keeps_a_generic_whole():
+    """Collapsing list[str] to str would silently drop nargs from its CLI flag."""
+
+    @dataclasses.dataclass(frozen=True)
+    class Options:
+        required_columns: list[str] = dataclasses.field(default_factory=list)
+        optional_columns: list[str] | None = None
+
+    assert option_type(Options, "required_columns") == list[str]
+    assert option_type(Options, "optional_columns") == list[str]
+
+
 class TestBuilderValidation:
     """Missing required inputs are the reward's own error to raise."""