feat(cli): select and configure rewards via --reward-type and --reward-config

AGENTS.md

@@ -182,6 +182,43 @@ pixi run -e boltz sampleworks-guidance \
 
 Run `sampleworks-guidance --model <model> --guidance-type <type> --help` to see all available options.
 
+### Choosing rewards
+
+`--reward-type` selects which reward guides the run (default `real_space_density`), and the
+options listed under it in `--help` are that reward's own — `--density`/`--resolution` above
+belong to the density reward, while `--reward-type structure_factor` takes `--mtzfile` and
+friends instead. Passing one reward's option to another is an error rather than a silent
+no-op.
+
+To combine rewards, or to keep a run's reward settings in version control, pass a
+configuration file instead (YAML, JSON, or TOML; `${oc.env:VAR}` interpolation works in YAML):
+
+```yaml
+# rewards.yaml -- sampleworks-guidance ... --reward-config rewards.yaml
+real_space_density:
+  weight: 0.4
+  reward_options:
+    density: /data/1vme.ccp4
+    resolution: 1.8
+structure_factor:
+  weight: 0.6
+  reward_options:
+    mtzfile: /data/1vme.mtz
+    bulk_solvent: combined
+```
+
+Weights default to `1/N`, so combining rewards does not change what the guidance step size
+means. Omit them all or give them all; a partly-weighted configuration is an error.
+
+**Adding a reward type** is three things and no argparse edits: implement the reward in
+`core/rewards/`, declare its options as a frozen dataclass in `core/rewards/options.py`, and
+register it in `core/rewards/registry.py` with a `build_*` function that raises its own
+errors for inputs it cannot do without. CLI flags, configuration-file schema, validation
+messages, and run metadata all follow from that one declaration. A reward that needs the
+model's atom ordering (structure factors, anything with a topology) implements
+`prepare(atom_array, *, device)` from `PreparableRewardFunctionProtocol`; the trajectory
+scalers call it once the model atom array exists.
+
 The `run_guidance()` function in `utils/guidance_script_utils.py` is the central orchestrator. It wires together the model wrapper, sampler (`AF3EDMSampler`), step scaler (`DataSpaceDPSScaler` or `NoiseSpaceDPSScaler`), trajectory scaler (`PureGuidance` or `FKSteering`), and reward function. When adding a new model or guidance strategy, this is the best reference for how components compose in practice.
 
 ## Development Environment
@@ -473,7 +510,7 @@ Proteins exist as thermodynamic ensembles, not static structures. Current genera
 Currently planned:
 - Real-space electron density (X-ray crystallography) *implemented*
 - Cryo-EM density *implemented*
-- Structure factors (reciprocal space)
+- Structure factors (reciprocal space) *implemented*
 - Diffuse scattering
 - Cryo-EM image stacks
 

README.md

@@ -100,11 +100,45 @@ Output files appear in `output/boltz2_pure_guidance/`: `refined.cif` (final ense
 | `--guidance-type` | `pure_guidance` or `fk_steering` |
 | `--protein` | Protein identifier (should match naming used in grid search / evaluation) |
 | `--structure` | Path to input structure file (CIF) |
-| `--density` | Path to density map (CCP4/MRC/MAP) |
-| `--resolution` | Map resolution in Angstroms |
+| `--density` | Path to density map (CCP4/MRC/MAP) — required by the default reward |
+| `--resolution` | Map resolution in Angstroms — required by the default reward |
 
 Model-specific arguments (e.g. `--method` for boltz2, `--msa-path` for rf3) and guidance-type-specific arguments (e.g. `--num-particles` for fk_steering) are included automatically. Run `sampleworks-guidance --model <model> --guidance-type <type> --help` to see all available options.
 
+### Rewards
+
+`--reward-type` chooses what the run is guided by, and each reward brings its own options:
+
+| Reward | Guided by | Its options |
+|---|---|---|
+| `real_space_density` (default) | Fit to a density map | `--density`, `--resolution`, `--loss-order`, `--em` |
+| `structure_factor` | Fit to structure-factor amplitudes from an MTZ | `--mtzfile`, `--expcolumns`, `--resolution`, `--bulk-solvent`, `--scattering-factor-mode`, ... |
+
+```bash
+sampleworks-guidance --model boltz2 --guidance-type pure_guidance \
+    --protein 1VME --structure 1vme.cif \
+    --reward-type structure_factor --mtzfile 1vme.mtz --bulk-solvent combined
+```
+
+To combine rewards, or to keep reward settings in version control, describe them in a file
+(YAML, JSON, or TOML) and pass `--reward-config rewards.yaml`:
+
+```yaml
+real_space_density:
+  weight: 0.4
+  reward_options:
+    density: 1vme.ccp4
+    resolution: 1.8
+structure_factor:
+  weight: 0.6
+  reward_options:
+    mtzfile: 1vme.mtz
+    bulk_solvent: combined
+```
+
+Weights default to `1/N` when omitted, so combining rewards leaves the meaning of the
+guidance step size intact.
+
 
 
 ## Grid Search

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})"