diff --git a/src/sampleworks/core/rewards/structure_factor.py b/src/sampleworks/core/rewards/structure_factor.py new file mode 100644 index 00000000..c0a52ce2 --- /dev/null +++ b/src/sampleworks/core/rewards/structure_factor.py @@ -0,0 +1,627 @@ +"""Reciprocal-space reward function for structure-factor amplitudes.""" + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import reciprocalspaceship as rs +import torch +from jaxtyping import Bool, Complex, Float, Int +from loguru import logger +from sampleworks.synthetic.synthetic_utils import atomarray_to_gemmi, resolve_mtz_column +from SFC_Torch import SFcalculator +from SFC_Torch.io import PDBParser + + +if TYPE_CHECKING: + from biotite.structure import AtomArray + + +# Loss callable: maps (|Fcalc|, |Fobs|) over the masked reflections to a scalar. +AmplitudeLoss = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + +# Bulk-solvent treatment for the scored amplitude. "off": |Fprotein|. "combined": +# |Ftotal| with one mask from the combined density. "per_conformer": |Ftotal| with +# the mean of per-conformer masks (see the class docstring / _compute_ensemble_ftotal). +_BULK_SOLVENT_MODES = ("off", "combined", "per_conformer") + +# SFcalculator init kwargs owned by this class (derived from other constructor +# arguments or injected in prepare()). They must not be overridden via +# sfcalculator_kwargs, which is an escape hatch for the remaining, unmanaged kwargs. +_RESERVED_SFC_KWARGS = frozenset( + {"dmin", "mode", "anomalous", "mtzdata", "pdbmodel", "expcolumns", "device", "set_experiment"} +) + +# Two thresholds for warnings reflection set too small after dropping outliers and R-free test +# set; see _build_reflection_mask. +# The min fraction is intended to catch an R-free convention mismatch (specify the wrong test +# set value for example). R-free test sets are conventionally 5-10% and outlier should be under +# 1%, so 75% is a safe threshold. +# The min absolute number is intended to catch a reflection set too small to guide the structure +# generation. We should have on the order of 1e4 reflections for 2A resolution macromolecular +# data, so 5e3 is a safe threshold. +_MIN_RETAINED_REFLECTION_FRACTION = 0.75 +_MIN_RETAINED_REFLECTIONS = 5_000 + + +def _resolve_expcolumns(expcolumns: list[str] | None, ds: rs.DataSet) -> tuple[str, str]: + """Resolve the ``[amplitude, sigma]`` columns to read from ``ds``, logging the choice. + + Auto-detection requires the MTZ to hold exactly one amplitude column and one sigma + column. A multi-set MTZ (e.g. ``Fprotein`` + ``Ftotal``) is ambiguous and forces the + caller to pass ``expcolumns`` explicitly. + + Parameters + ---------- + expcolumns + Caller-provided ``[amplitude, sigma]`` to use verbatim, or ``None`` to auto-detect. + ds + The parsed MTZ whose amplitude and sigma columns are selected. + + Returns + ------- + tuple of str + The resolved ``(amplitude, sigma)`` pair, in that order. + + Raises + ------ + ValueError + If a provided ``expcolumns`` is not a length-2 ``[amplitude, sigma]`` pair, or names + a column absent from the MTZ's amplitude/sigma columns; or, when auto-detecting, if + the MTZ has zero or more than one amplitude (or sigma) column. + """ + amplitude_dtype, sigma_dtype = rs.StructureFactorAmplitudeDtype(), rs.StandardDeviationDtype() + if expcolumns is not None: + # A bare string of length 2 could pass the length check and fail in not so obvious ways + # ("FP" -> amplitude "F", sigma "P"), so reject it explicitly. + if isinstance(expcolumns, str) or len(expcolumns) != 2: + raise ValueError(f"expcolumns must be a [amplitude, sigma] pair; got {expcolumns!r}.") + amplitude = resolve_mtz_column(ds, amplitude_dtype, column=expcolumns[0]) + sigma = resolve_mtz_column(ds, sigma_dtype, column=expcolumns[1]) + return amplitude, sigma + + amplitude = resolve_mtz_column(ds, amplitude_dtype) + sigma = resolve_mtz_column(ds, sigma_dtype) + logger.info( + f"No expcolumns provided; auto-detected SFC columns: " + f"amplitude='{amplitude}', sigma='{sigma}'." + ) + return amplitude, sigma + + +class StructureFactorRewardFunction: + def __init__( + self, + mtzfile: str | Path, + *, + expcolumns: list[str] | None = None, + resolution: float | None = None, + scattering_factor_mode: str = "xray", + bulk_solvent: str = "off", + loss: AmplitudeLoss | None = None, + normalize_amplitude: bool = False, + exclude_free_reflections: bool = False, + batch_partition: int = 10, + sfcalculator_kwargs: dict | None = None, + ): + """Reward for fitting structure-factor amplitudes via SFcalculator. + + Scores structures against experimental (or synthetic) structure-factor + amplitudes ``|Fobs|`` using ``SFcalculator`` from ``SFC_Torch``. This is the + reciprocal-space counterpart to :class:`RealSpaceRewardFunction` in + ``real_space_density.py``. + + The reward compares the model amplitudes ``|Fcalc|`` against a target + ``|Fobs|`` loaded from an MTZ. For an ensemble (batch dimension) of + conformers indexed by ``b``, the per-conformer *protein* structure + factors are combined by a *complex* sum in the reciprocal space + (``|sum F| != sum |F|``): ``Fprotein(h) = sum_b F_b(h)``. The atomic + occupancy is accounted for in the calculation of the per-conformer + *protein* structure factor F_b(h). ``|Fcalc|`` is then ``|Fprotein|`` + (``bulk_solvent="off"``) or ``|Ftotal|`` if bulk solvent is folded in. + + ``bulk_solvent`` can take one of the following values: + * ``"off"`` (default): ``|Fprotein|`` — no solvent. + * ``"combined"``: ``|Ftotal|`` with one bulk-solvent mask from the combined + protein density, ``mask()`` — matches the altloc single-structure + ``Ftotal`` that ``generate_synthetic_sf`` writes to the MTZ. + * ``"per_conformer"``: ``|Ftotal|`` with the mean of per-conformer solvent + mask ```` — the ensemble-averaged bulk solvent. Each conformer + contributes a solvent mask at 1/batch_size weight. + + ``"combined"`` and ``"per_conformer"`` differ only for a real ensemble + (batch > 1); the mask operator is nonlinear, so ``mask() != ``. + Both use the default, *unrefined* scales ``kiso=1``, ``kmask=0.35``, small + ``uaniso``; refining them during sampling is left for a later revision. + ``normalize_amplitude`` (``|F|`` vs ``|E|``) is orthogonal and composes with + any ``bulk_solvent`` choice. + + The crystal metadata is taken from the MTZ, which is the only self-consistent + source: ``SFcalculator``'s ``init_mtz`` overwrites the cell and space group of + the structure by the MTZ's whenever they disagree. Naively overriding the cell + and/or the space group of MTZ can be dangerous because the reflections' indexing + could become wrong, so there are no arguments here to override them. + + Unlike the real-space reward function, ``SFcalculator`` needs the full topology + (``PDBParser`` -> ``gemmi.Structure``: atom names/elements, unit cell, space group, + ``resolution`` -> the HKL set) at construction. That information is only available + with the model atom array, i.e. after ``process_structure_to_trajectory_input`` + runs inside ``sample()``. Therefore, construction is two-phase: + * ``__init__`` stores only the up-front config (target MTZ, ``resolution``, + ``scattering_factor_mode``, loss) and resolves the crystal metadata against + the MTZ; it does *not* build ``SFcalculator``. + * :meth:`prepare` builds ``SFcalculator`` from the model atom array, on the + ``device`` given there. The caller (a step scaler) is responsible for + invoking it before the first ``__call__``. + + Parameters + ---------- + mtzfile + Path to the MTZ holding the target amplitudes (and a sigma column). + Loaded with ``set_experiment=True`` so ``sfc.Fo`` / HKL set / bins are + populated and the calculation is aligned to the target reflections. + expcolumns + Column names ``[amplitude, sigma]`` in the MTZ. If ``None`` (default), the + amplitude and sigma columns are auto-detected, which requires the MTZ to hold + exactly one of each; a multi-set MTZ (e.g. ``Fprotein`` + ``Ftotal``) is + ambiguous and raises, forcing an explicit ``[amplitude, sigma]``. If provided, + a ``ValueError`` is raised when it is not a length-2 pair or names a column + absent from the MTZ. + resolution + High-resolution limit (dmin) in Angstrom, or ``None`` (default) to use + the MTZ's own resolution. The HKL set comes from the ``mtzfile``; when + given, ``resolution`` further truncates it. + scattering_factor_mode + SFcalculator scattering mode: ``"xray"`` or ``"cryoem"``. + bulk_solvent + Bulk-solvent treatment, one of ``"off"`` (default; score ``|Fprotein|``), + ``"combined"`` (``|Ftotal|`` with one mask from the combined density), or + ``"per_conformer"`` (``|Ftotal|`` with the plain, unweighted mean of the + per-conformer masks). The per-conformer mean is *not* occupancy-weighted, + meaning each conformer contributes a solvent mask at 1/batch_size weight. + For a single conformer, ``"combined"`` and ``"per_conformer"`` coincide. + For multiple conformers, ``"per_conformer"`` should be a more faithful model. + loss + Callable ``(|Fcalc|, |Fobs|) -> scalar`` over masked reflections. + Defaults to mean-squared error on amplitudes. Pass any callable + (L1, R-factor-style, resolution-weighted, ...) to customize. + normalize_amplitude + If True, score normalized structure factors (E-values) instead of ``|F|``. + exclude_free_reflections + If True, drop the R-free test-set reflections from the loss (use only + the working set). Default False: use all non-outlier reflections. + batch_partition + Ensemble chunk size forwarded to ``SFcalculator.calc_fprotein_batch`` as its + ``PARTITION`` parameter. SFC's own default (20) can still lead to OOM. + Must be a positive integer. + sfcalculator_kwargs + Extra keyword arguments forwarded verbatim to ``SFcalculator(...)`` in + :meth:`prepare` (e.g. ``n_bins``, ``freeflag``, ``testset_value``). Reserved + keys in this class (listed in :data:`_RESERVED_SFC_KWARGS`) cannot be overridden. + """ + self.mtzfile = str(mtzfile) + self.resolution = resolution + self.scattering_factor_mode = scattering_factor_mode + if bulk_solvent not in _BULK_SOLVENT_MODES: + raise ValueError( + f"bulk_solvent must be one of {_BULK_SOLVENT_MODES}, got {bulk_solvent!r}." + ) + self.bulk_solvent = bulk_solvent + self.exclude_free_reflections = exclude_free_reflections + # bool will pass isinstance(..., int) check but not the type check + if type(batch_partition) is not int or batch_partition <= 0: + raise ValueError( + f"batch_partition must be a positive integer, got {batch_partition!r}." + ) + self.batch_partition = batch_partition + self.loss: AmplitudeLoss = loss if loss is not None else torch.nn.MSELoss() + self.normalize_amplitude = normalize_amplitude # |F| vs resolution-bin normalized |E| + + # Resolve crystal metadata / column names against the MTZ. + self._resolve_mtz_metadata(expcolumns) + + # All SFcalculator init kwargs are known except `pdbmodel` (needs the model + # atom array), `mtzdata` (a copy of the parsed dataset), and `device` (resolved + # at prepare() time), all injected in prepare(). + self._sfc_kwargs = dict( + dmin=self.resolution, + mode=self.scattering_factor_mode, + anomalous=False, + set_experiment=True, + expcolumns=self.expcolumns, + ) + if sfcalculator_kwargs: + reserved = _RESERVED_SFC_KWARGS & sfcalculator_kwargs.keys() + if reserved: + raise ValueError( + f"sfcalculator_kwargs may not override reserved keys {sorted(reserved)}; " + "these are managed by this class (via other constructor arguments or " + "injected in prepare())." + ) + self._sfc_kwargs.update(sfcalculator_kwargs) + + # Populated by prepare(); None until then. + self.sfc: SFcalculator | None = None + self._reflection_mask: torch.Tensor | None = None + + def _resolve_mtz_metadata(self, expcolumns: list[str] | None) -> None: + """Set ``self.unit_cell`` / ``space_group`` / ``expcolumns`` from the MTZ. + + The MTZ is the only source: it must be consistent with its own reflections, and + ``SFcalculator`` enforces that by overwriting the cell / space group of the + structure handed to it with the MTZ's (``init_mtz``) before deriving symmetry + operators or the reciprocal cell. The values resolved here are therefore what the + calculation uses; they are stamped onto the gemmi structure in :meth:`prepare` + only to keep it self-consistent. The dataset is parsed once here and retained on + ``self._mtz_dataset`` so :meth:`prepare` can build ``SFcalculator`` from it + directly instead of re-reading the file. ``expcolumns``, unlike the crystal + metadata, is caller-supplied and validated against the MTZ's columns; see + :func:`_resolve_expcolumns`. + + Parameters + ---------- + expcolumns + Caller-supplied ``[amplitude, sigma]`` column names to use verbatim, or + ``None`` to auto-detect from the MTZ; forwarded to :func:`_resolve_expcolumns`. + + Raises + ------ + ValueError + If the MTZ carries no space group gemmi can recognize, or no unit cell (a + missing ``CELL`` header leaves gemmi's placeholder, which would make every + d-spacing meaningless). Also propagated from :func:`_resolve_expcolumns` for a + malformed or unknown ``expcolumns``. + """ + self._mtz_dataset = rs.read_mtz(self.mtzfile) + if self._mtz_dataset.spacegroup is None: + raise ValueError( + f"{self.mtzfile} carries no space group gemmi can recognize (missing or " + "unrecognized SYMINF header)." + ) + # gemmi's own check on if cell_a is exactly 1A (the dummy value gemmi writes) + if not self._mtz_dataset.cell.is_crystal(): + raise ValueError( + f"{self.mtzfile} carries no unit cell (missing CELL header, leaving gemmi's " + f"placeholder {self._mtz_dataset.cell.parameters}). Please re-generate the MTZ." + ) + self.unit_cell = self._mtz_dataset.cell + self.space_group = self._mtz_dataset.spacegroup.hm + + # Stored (and passed to SFcalculator) as a list: SFC_Torch annotates the kwarg + # `List[str]` and swallows any failure into a misleading "columns not in the mtz" error. + self.expcolumns = list(_resolve_expcolumns(expcolumns, self._mtz_dataset)) + + def prepare(self, atom_array: AtomArray, *, device: torch.device | str = "cpu") -> None: + """Build the SFcalculator from the model atom array, on ``device``. + + Constructing the ``SFcalculator`` consumes the MTZ dataset parsed at + ``__init__`` (no second file read) and populates the observed structure + factor amplitudes ``sfc.Fo``, the observed HKL set in the ASU, the resolution + bins, the outlier mask ``sfc.Outlier``, the R-free flags ``sfc.free_flag``, + and the normalized ``|Eo|`` in ``sfc.Eo``. + + Must be called once before the first ``__call__``, with the same atom + array that the sampled coordinates correspond to (model atom space: + ``model_atom_array or atom_array``). The atom ordering of ``atom_array`` + defines the column order of the coordinate tensor passed to ``__call__``. + + This is also where the torch device is set, so pass the device the sampled + coordinates will live on: every tensor built here (the SFcalculator internals and + the reflection mask) is allocated on it, and SFcalculator bakes it in at + construction with no way to migrate afterwards. Re-running ``prepare`` on a new + device is therefore the way to move a built reward — safe to do, since the parsed + MTZ dataset on ``self._mtz_dataset`` is kept pristine for exactly that reason. + + Per-atom B-factors / occupancy are set per ``__call__`` (not here), leaving + the door open to refining them during sampling. + + Parameters + ---------- + atom_array + Biotite AtomArray for the atoms the model operates on. Needs + ``chain_id``, ``res_id``, ``res_name``, ``atom_name``, ``element`` + annotations (its ``b_factor``/``occupancy`` are baked as defaults but + overridden each ``__call__``). A missing ``altloc_id`` is defaulted to + blank inside ``atomarray_to_gemmi``. + device + Torch device to build on. Defaults to CPU rather than auto-selecting a GPU: + the same behavior as how the RewardInputs' device is resolved. + + Raises + ------ + RuntimeError + If ``normalize_amplitude`` is True but SFcalculator's normalization did not + yield finite ``sfc.Eo`` from the supplied MTZ (absent or non-finite). When + ``normalize_amplitude`` is False the same condition only logs a warning, + since it also implies no reflection was flagged as an outlier. + ValueError + If no reflection survives the mask; see :meth:`_build_reflection_mask`. + """ + self.device = torch.device(device) + + gemmi_structure = atomarray_to_gemmi( + atom_array, + unit_cell=self.unit_cell, + space_group=self.space_group, + ) + # SFcalculator mutates its mtzdata in place (dropna / hkl_to_asu on the reference), + # so hand it a copy of the once-parsed dataset; self._mtz_dataset stays pristine and + # prepare() remains safely re-runnable. + sfc_kwargs = { + "mtzdata": self._mtz_dataset.copy(), + "device": self.device, + **self._sfc_kwargs, + } + self.sfc = SFcalculator(pdbmodel=PDBParser(gemmi_structure), **sfc_kwargs) + # inspect_data estimates solvent percentage and grid size from atom positions + # and vdW radii, independent of occupancy / B-factor. + self.sfc.inspect_data() + + # Ftotal uses the default (unrefined) scales kiso=1, kmask=0.35, small uaniso, + # matching generate_synthetic_sf's Ftotal. They don't depend on coordinates, so + # set them once here. _set_scales only needs atom_pos_frac (set at construction) + # and n_bins (set by the experiment init) for dtype/device and bin count. + if self.bulk_solvent != "off": + self.sfc._set_scales(requires_grad=False) + + # |Eo|, epsilon, and Outlier all come from one bare try/except inside SFC's experiment + # init. A missing Eo means that block raised; a non-finite Eo means a resolution bin + # had zero mean intensity. Either way |Eo| is unusable and the sfc.Outlier mask would + # not flag the outlier reflections based on Eo. + Eo = getattr(self.sfc, "Eo", None) + if Eo is None or not torch.isfinite(Eo).all(): + if self.normalize_amplitude: + raise RuntimeError( + "normalize_amplitude=True requires finite sfc.Eo, but SFcalculator's " + "normalization block did not yield them from this MTZ." + ) + logger.warning( + "SFcalculator did not yield finite sfc.Eo from this MTZ; its outlier detection " + "will only flag non-positive amplitudes as outliers instead of also using Eo." + ) + + # True where a reflection contributes to the reward (non-outlier, and not in the + # R-free test set if exclude_free_reflections is True). Shape is [n_hkl]. + mask_np = self._build_reflection_mask( + outlier=self.sfc.Outlier, free_flag=self.sfc.free_flag + ) + self._reflection_mask = torch.from_numpy(mask_np).to(self.device) + + logger.info( + f"Prepared StructureFactorRewardFunction: n_atoms={len(self.sfc.atom_pos_orth)}, " + f"n_reflections={len(self.sfc.Fo)}, n_used={int(mask_np.sum())}, " + f"cell={self.sfc.unit_cell}, space_group={self.sfc.space_group.hm}, " + f"solventpct={self.sfc.solventpct}, gridsize={self.sfc.gridsize}" + ) + + def _build_reflection_mask( + self, + *, + outlier: Bool[np.ndarray, "n_hkl"], # noqa: F821, UP037 + free_flag: Bool[np.ndarray, "n_hkl"], # noqa: F821, UP037 + ) -> Bool[np.ndarray, "n_hkl"]: # noqa: F821, UP037 + """Build a mask of valid training set reflections for reward computation and + sanity check the mask size. + + Outliers are always dropped; the R-free test set is also dropped when user sets + ``exclude_free_reflections``. Raise error when the number of reflections remain + is 0, and log warning when it is too small (below ``_MIN_RETAINED_REFLECTION_FRACTION`` + or ``_MIN_RETAINED_REFLECTIONS``). + + Parameters + ---------- + outlier + ``sfc.Outlier`` ``[n_hkl]``: True where SFcalculator flagged the observed + reflection as an outlier. + free_flag + ``sfc.free_flag`` ``[n_hkl]``: True where the reflection is in the R-free + test set. + + Returns + ------- + numpy.ndarray + Boolean mask over the observed reflections ``[n_hkl]``, True where the + reflection contributes to the loss. + + Raises + ------ + ValueError + If no reflection survives the mask. + """ + mask_np = ~outlier # allocates, so the &= below never touches SFC's own arrays + if self.exclude_free_reflections: + mask_np &= ~free_flag + + n_total, n_used = mask_np.size, int(mask_np.sum()) + composition = ( + f"{n_total} observed reflections, " + f"{int(outlier.sum())} outliers, " + f"{int(free_flag.sum())} flagged free " + f"(exclude_free_reflections={self.exclude_free_reflections})" + ) + convention_hint = ( + "SFcalculator treats mtz[freeflag] == testset_value as the test set; " + "pass a matching `freeflag` / `testset_value` via sfcalculator_kwargs " + "if the MTZ's R-free convention differs from its defaults ('FreeR_flag' / 0)." + ) + if n_used == 0: + raise ValueError(f"No reflections remain: {composition}. {convention_hint}") + if n_used < _MIN_RETAINED_REFLECTIONS: + logger.warning( + f"Only {n_used} reflections remain: {composition}. " + f"Check `resolution` (dmin={self.resolution} A) and the MTZ reflection range." + ) + if n_used < _MIN_RETAINED_REFLECTION_FRACTION * n_total: + logger.warning( + f"Only {n_used}/{n_total} reflections remain: {composition}. {convention_hint}" + ) + return mask_np + + 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 amplitude loss for the (ensemble of) coordinates. + + Call ``.backward()`` on the result to get gradients w.r.t. + ``coordinates``. + + ``elements`` is ignored (topology is fixed in the SFcalculator built by + :meth:`prepare`). ``b_factors`` and ``occupancies`` are reset onto the + SFcalculator each call (mirroring ``RealSpaceRewardFunction``), and must be + broadcast-identical across the batch dim (enforced; SFcalculator has no + per-conformer occupancy/B axis, so non-broadcast input raises ``ValueError``). + + Parameters + ---------- + coordinates + Atomic coordinates ``[batch, n_atoms, 3]`` in model atom space, + matching the atom ordering passed to :meth:`prepare`. + elements + Per-atom element codes ``[batch, n_atoms]``. Ignored: the topology + (including elements) is fixed in the SFcalculator built by + :meth:`prepare`. Present only to satisfy the reward-function signature. + b_factors + Per-atom isotropic B-factors ``[batch, n_atoms]``, written to + ``sfc.atom_b_iso`` (reconciled: real deposited where shared with the + structure, 20.0 for model-only / NaN atoms). + occupancies + Per-atom occupancies ``[batch, n_atoms]`` (uniform ``1/batch_size`` from + the pipeline), written to ``sfc.atom_occ``; the ``1/batch_size`` weighting + makes the complex ensemble sum the multi-conformer total. + unique_combinations + Pre-computed unique (element, b_factor) pairs for vmap compatibility. + Currently unused. + inverse_indices + Pre-computed inverse indices for vmap compatibility. Currently unused. + + Returns + ------- + torch.Tensor + Scalar reward (loss). + + Raises + ------ + RuntimeError + If :meth:`prepare` has not been called. + ValueError + If any input's atom count differs from the SFcalculator topology built by `prepare` + or if ``b_factors`` / ``occupancies`` are not broadcast-identical across batch dim. + """ + if self.sfc is None or self._reflection_mask is None: + raise RuntimeError( + "StructureFactorRewardFunction.prepare() must be called with the model " + "atom array before the reward is evaluated." + ) + + # The topology (atom count included) is fixed by prepare(); a mismatch would otherwise + # surface as a broadcast error inside calc_fprotein_batch. + n_topology_atoms = len(self.sfc.atom_pos_orth) + for name, n_atoms in ( + ("coordinates", coordinates.shape[-2]), + ("b_factors", b_factors.shape[-1]), + ("occupancies", occupancies.shape[-1]), + ): + if n_atoms != n_topology_atoms: + raise ValueError( + f"{name} has {n_atoms} atoms but the SFcalculator topology built by " + f"prepare() has {n_topology_atoms}. Call prepare() with the same atom array " + "the sampled coordinates correspond to (model atom space)." + ) + + # SFcalculator has no per-conformer (batch) occupancy/B axis, so these must be shared + # across the ensemble; row 0 is used. Reject non-broadcast input as a guard. + for name, tensor in (("occupancy", occupancies), ("B-factor", b_factors)): + if not torch.equal(tensor, tensor[:1].expand_as(tensor)): + raise ValueError( + f"StructureFactorRewardFunction requires {name} identical across the " + "batch dim (SFcalculator has no per-conformer occupancy/B axis); got " + "per-conformer values." + ) + self.sfc.atom_b_iso = b_factors[0] + self.sfc.atom_occ = occupancies[0] + + # Multi-conformer combination: complex sum over the ensemble [batch, n_hkl]. + # occ = 1/batch_size (set per call) makes the summed |F| the multi-conformer total. + Fprotein_batch = self.sfc.calc_fprotein_batch( + coordinates, Return=True, PARTITION=self.batch_partition + ) + Fprotein = Fprotein_batch.sum(dim=0) + + fcalc = Fprotein if self.bulk_solvent == "off" else self._compute_ensemble_ftotal(Fprotein) + + mask = self._reflection_mask + if self.normalize_amplitude: + calc = self.sfc.calc_Ec(fcalc).abs() + obs = self.sfc.Eo + else: + calc = torch.abs(fcalc) + obs = self.sfc.Fo + return self.loss(calc[mask], obs[mask]) + + def _compute_ensemble_ftotal( + self, + Fprotein_HKL: Complex[torch.Tensor, "n_hkl"], # noqa: F821, UP037 + ) -> Complex[torch.Tensor, "n_hkl"]: # noqa: F821, UP037 + """Add default-scaled bulk solvent to the ensemble Fprotein to form Ftotal. + + ``Ftotal(h) = kiso * aniso(h) * (Fprotein(h) + kmask * Fmask(h))`` with the + default (unrefined) scales set in :meth:`prepare`, evaluated on the + experimental HKL set. ``Fprotein`` is the ensemble complex sum; the + bulk-solvent ``Fmask`` is combined per :attr:`bulk_solvent`. The two + ``|Ftotal|`` modes are identical for a single conformer. + + ``bulk_solvent="combined"`` (``mask()``) + One mask built from the combined protein density. ``calc_fsolvent`` + builds the mask by FFT over the full ASU set, so the combined + ``Fprotein_asu`` (not just the HKL subset) is fed to it. Matches the + altloc single-structure Ftotal in the synthetic MTZ. + + ``bulk_solvent="per_conformer"`` (````) + Average of the per-conformer bulk solvent masks. ``rsgrid2realmask`` + normalizes the protein density and cuts at a quantile, so each conformer + contributes ``Fmask`` in a scale-invariant way. All conformers are assumed + to contribute the solvent mask equally (1/batch_size weight). + + Two caveats from ``rsgrid2realmask``'s batch path (SFC_Torch 0.3.3): + + - It is NOT permutation-invariant over the ensemble. The quantile ``CUTOFF`` + for bulk solvent mask comes from ``real_grid_norm[0]`` (batch element 0) + and is then used for every conformer. + - Above 5M voxels per conformer the cutoff is a quantile of an unseeded + ``torch.randperm`` subsample, making the mask (and its gradients) + nondeterministic run-to-run. + + Parameters + ---------- + Fprotein_HKL + Ensemble complex-sum Fprotein on the experimental HKL set ``[n_hkl]``. + + Returns + ------- + torch.Tensor + Complex Ftotal on the experimental HKL set ``[n_hkl]``. + """ + assert self.sfc is not None # prepare() built it; __call__ guards before dispatching here + self.sfc.Fprotein_HKL = Fprotein_HKL # drives calc_ftotal on the HKL set + if self.bulk_solvent == "per_conformer": + # calc_fsolvent_batch masks each conformer (from Fprotein_asu_batch, set by + # calc_fprotein_batch); the mean applies the 1/batch_size weight -> . + Fmask_HKL_batch = self.sfc.calc_fsolvent_batch( + Return=True, PARTITION=self.batch_partition + ) + self.sfc.Fmask_HKL = Fmask_HKL_batch.mean(dim=0) + else: + # One mask from the combined density -> mask(). SFC calc_solvent() requires + # having the Fprotein_asu set, but batch mode only sets Fprotein_asu_batch. + 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() diff --git a/tests/rewards/conftest.py b/tests/rewards/conftest.py index 27acb68e..ce9e9747 100644 --- a/tests/rewards/conftest.py +++ b/tests/rewards/conftest.py @@ -8,17 +8,33 @@ `density_map_1vme` deliberately stays in the top-level conftest because `tests/utils/test_density_utils.py` also consumes it; a fixture lives at the lowest common ancestor of everything that uses it. + +The two reward families get their target data differently. The density fixtures read +committed files (a carved .ccp4 and its matching recentered cif). The SF fixtures +generate their (cif, mtz) pair per session via `generate_synthetic_sf.py`. """ from pathlib import Path +from typing import TYPE_CHECKING import pytest import torch from atomworks.io.parser import parse +from biotite.structure import AtomArray +from jaxtyping import Float +from torch import Tensor + + +if TYPE_CHECKING: + from sampleworks.core.forward_models.xray.real_space_density_deps.qfit.volume import ( + XMap, + ) + from sampleworks.core.rewards.real_space_density import RealSpaceRewardFunction + from sampleworks.core.rewards.structure_factor import StructureFactorRewardFunction @pytest.fixture(scope="session") -def structure_1vme_density(resources_dir: Path): +def structure_1vme_density(resources_dir: Path) -> dict: cif_path = resources_dir / "1vme" / "1vme_final_carved_edited_0.5occA_0.5occB.cif" if not cif_path.exists(): pytest.skip(f"Structure not found at {cif_path}") @@ -26,7 +42,10 @@ def structure_1vme_density(resources_dir: Path): @pytest.fixture(scope="session") -def reward_function_1vme(density_map_1vme, structure_1vme_density, device: torch.device): +def reward_function_1vme( + density_map_1vme: "XMap", device: torch.device +) -> "RealSpaceRewardFunction": + """Real-space density reward from the carved 1vme map, in X-ray (non-EM) mode.""" from sampleworks.core.rewards.real_space_density import ( RealSpaceRewardFunction, setup_scattering_params, @@ -38,7 +57,11 @@ def reward_function_1vme(density_map_1vme, structure_1vme_density, device: torch @pytest.fixture(scope="session") -def test_coordinates_1vme(structure_1vme_density, device: torch.device): +def test_coordinates_1vme( + structure_1vme_density: dict, device: torch.device +) -> tuple[Float[Tensor, "n_atoms 3"], AtomArray]: + """Atomic coordinates and the atom array (asym_unit) for the density model, + zero-occupancy atoms dropped.""" atom_array = structure_1vme_density["asym_unit"] # Handle both AtomArray and AtomArrayStack @@ -50,3 +73,124 @@ def test_coordinates_1vme(structure_1vme_density, device: torch.device): atom_array = atom_array[mask] coords = torch.from_numpy(atom_array.coord).to(device=device, dtype=torch.float32) return coords, atom_array + + +@pytest.fixture(scope="session") +def sf_1vme_cif_and_mtz_paths( + resources_dir: Path, tmp_path_factory: pytest.TempPathFactory, device: torch.device +) -> tuple[Path, Path]: + """Generate the synthetic SF data and the exact structure used for the generation, + returns the ``(cif, mtz)`` paths. + + Mirrors the ``--structure`` branch of ``generate_synthetic_sf.main()``, so these + fixtures exercise exactly what the CLI produces. ``--save-structure`` writes back the + post-selection and occupancy modifed model, which is the structure actually used to + generate the synthetic SF data. + + Specifically, we select for the chain A of 1vme with altloc occupancies forced to be + 0.5/0.5 and H + waters stripped, at 1.8 A resolution, with bulk solvent and default + scales. The resulted mtz file carries both Fprotein/SIGFprotein/PHIFprotein and + Ftotal/SIGFtotal/PHIFtotal. + """ + from sampleworks.synthetic.generate_synthetic_sf import ( + _process_single_row, + BatchRowForMTZ, + ) + + source_cif = "1vme_final.cif" + source_dir = resources_dir / "1vme" + if not (source_dir / source_cif).exists(): + pytest.skip(f"Source structure not found at {source_dir / source_cif}") + + resolution = 1.8 + output_dir = tmp_path_factory.mktemp("sf_1vme") + _process_single_row( + row=BatchRowForMTZ( + filename=source_cif, + selection="chain A", + occupancy_values=[0.5, 0.5], # assigned to altlocs in sorted order, A then B + ), + base_dir=source_dir, + output_dir=output_dir, + resolution=resolution, + scattering_factor_mode="xray", + occupancy_mode="custom", + test_fraction=0.1, # needed for test_inverted_testset_value_warns + seed=0, # R-free flag assignment only + device=device, + strip_hydrogens=True, + strip_waters=True, + simulate_solvent_and_scale=True, + save_structure=True, + ) + + # Following the default file names in the generate_synthetic_sf.py script: + # `{stem}_{resolution:.2f}A.mtz`, and the saved structure as `{stem}_sf_input.cif`. + # Ideally we want to modify _process_single_row to return the paths directly in case + # the naming defaults changed or accept naming parameters. + # Because _process_single_row logs and returns on failure instead of raising errors, + # missing output is the only way to check generation success. + stem = Path(source_cif).stem + cif_path = output_dir / f"{stem}_sf_input.cif" + mtz_path = output_dir / f"{stem}_{resolution:.2f}A.mtz" + missing = [p.name for p in (cif_path, mtz_path) if not p.exists()] + if missing: + raise RuntimeError( + f"generate_synthetic_sf wrote no {missing} to {output_dir}; see the " + f"generator's logged error for the cause" + ) + return cif_path, mtz_path + + +@pytest.fixture(scope="session") +def mtz_path_1vme(sf_1vme_cif_and_mtz_paths: tuple[Path, Path]) -> Path: + """Synthetic 1.8 A target: Fprotein and Ftotal sets, each with its SIGF and phase.""" + _, mtz_path = sf_1vme_cif_and_mtz_paths + return mtz_path + + +@pytest.fixture(scope="session") +def structure_1vme_sf(sf_1vme_cif_and_mtz_paths: tuple[Path, Path]) -> AtomArray: + """The exact chain-A model with altlocs in the P2_1 crystal frame that the mtz was + computed from.""" + from sampleworks.utils.atom_array_utils import load_structure_with_altlocs + + cif_path, _ = sf_1vme_cif_and_mtz_paths + return load_structure_with_altlocs(cif_path) + + +@pytest.fixture(scope="session") +def test_coordinates_1vme_sf( + structure_1vme_sf: AtomArray, device: torch.device +) -> tuple[Float[Tensor, "n_atoms 3"], AtomArray]: + """Atomic coordinates and the atom array (asym_unit) which for the SF reward is the + same atom array as the structure_1vme_sf. + + Exists for the ``_REWARD_BUNDLES`` mapping in ``test_reward_function_contract.py``, + which resolves one ``(coords, atom_array)`` fixture per reward so every reward is + exercised through the same shape. SF-specific tests that need only the atom array take + ``structure_1vme_sf`` directly. + """ + # No occupancy>0 filter here (unlike the density fixture) as SFC topology is fixed. + coords = torch.from_numpy(structure_1vme_sf.coord).to(device=device, dtype=torch.float32) + return coords, structure_1vme_sf + + +@pytest.fixture(scope="session") +def reward_function_1vme_sf( + mtz_path_1vme: Path, + structure_1vme_sf: AtomArray, + device: torch.device, +) -> "StructureFactorRewardFunction": + """Structure factor reward (|Fprotein| only) prepared using the 1vme structure's topology.""" + from sampleworks.core.rewards.structure_factor import StructureFactorRewardFunction + + # normalize_amplitude=True scores normalized E-values (|Ec| vs sfc.Eo), which are + # unit-variance per resolution shell, so the MSE can be tested on an absolute scale. + reward_function = StructureFactorRewardFunction( + mtz_path_1vme, + expcolumns=["Fprotein", "SIGFprotein"], + normalize_amplitude=True, + ) + reward_function.prepare(structure_1vme_sf, device=device) + return reward_function diff --git a/tests/rewards/test_reward_function_contract.py b/tests/rewards/test_reward_function_contract.py index f383ed96..19a9f727 100644 --- a/tests/rewards/test_reward_function_contract.py +++ b/tests/rewards/test_reward_function_contract.py @@ -31,9 +31,11 @@ # Every test exercises CUDA-targeted reward code on the `device` fixture (try_gpu), so the -# whole module is gpu-marked. Deliberately NOT `slow`: measured warm per-test time is -# sub-second. The fixed cost is one-time import + session-scoped reward construction, which -# is paid once per pytest invocation and cannot be skipped by `slow`-marking these tests. +# whole module is gpu-marked. Deliberately NOT `slow`: measured warm per-test time is ~1s +# at most (the SFC gradient-descent loop; the rest are sub-0.5s). The fixed cost is a one-time +# import, the session-scoped synthetic SF generation, and session-scoped reward construction +# (real_space + SFC), paid once per pytest invocation and not skippable by `slow`-marking +# these tests. pytestmark = pytest.mark.gpu @@ -61,10 +63,12 @@ def batch(self, n: int = 1, coords: torch.Tensor | None = None) -> dict: # Per-reward bundles: each param resolves its OWN coordinates/atom_array/reward so the # inputs are self-consistent with that reward's target. real_space uses the recentered -# carved 1vme (matches its .ccp4 map frame). New rewards register here alongside a matching -# entry in `_LOSS_THRESHOLDS`. +# carved 1vme (matches its .ccp4 map frame); structure_factor uses the crystal-frame +# chain-A model the synthetic MTZ was generated from (recentering corrupts SF symmetry +# mates). New rewards register here alongside a matching entry in `_LOSS_THRESHOLDS`. _REWARD_BUNDLES = { "real_space": ("test_coordinates_1vme", "reward_function_1vme"), + "structure_factor": ("test_coordinates_1vme_sf", "reward_function_1vme_sf"), } # Absolute loss bar for the TRUE structure, per reward. RealSpace's loss is MSE on @@ -72,8 +76,12 @@ def batch(self, n: int = 1, coords: torch.Tensor | None = None) -> dict: # a 0.5 A perturbation ~0.034, and random coordinates ~O(1). The 0.01 bar sits ~5x above # the true loss (robust to device/precision variance) yet ~3x below the 0.5 A-perturbed # loss, so it comfortably passes the truth while meaningfully failing a wrong structure. +# SFC (normalize_amplitude=True -> normalized E-values) measured on the synthetic MTZ: +# true ~2e-14, a 0.5 A perturbation ~0.34, and random ~0.44, so the 0.1 bar sits well +# above numerical zero and ~3x below the smallest perturbation signal. _LOSS_THRESHOLDS = { "real_space": 0.01, + "structure_factor": 0.1, } diff --git a/tests/rewards/test_structure_factor_reward.py b/tests/rewards/test_structure_factor_reward.py new file mode 100644 index 00000000..e8982c3b --- /dev/null +++ b/tests/rewards/test_structure_factor_reward.py @@ -0,0 +1,661 @@ +"""Tests specific to the structure-factor (reciprocal-space) reward function. + +Reward-agnostic contract tests (interface, relative correlation, coordinate gradients, +batch=1 occupancy gradients, batch/edge handling) live in +``test_reward_function_contract.py``, where they run against every reward. What remains +here is specific to ``StructureFactorRewardFunction`` +(``sampleworks.core.rewards.structure_factor``) or to SFcalculator's forward model. + +Specific features of the SF reward include: +- No per-conformer occupancy/B-factors — SFcalculator batches only coordinates, so + ``__call__`` should reject if not broadcastable across the batch. +- Bulk-solvent modes — ``off`` scores ``|Fprotein|``; ``combined`` or ``per_conformer`` + additionally account for bulk solvent contributions depending on the ensemble. +- MTZ information — unit cell, space group, and amplitude/sigma column; reflection masking + for outliers and R-free test set. +""" + +import logging +from pathlib import Path +from unittest import mock + +import gemmi +import numpy as np +import pytest +import reciprocalspaceship as rs +import torch +from biotite.structure import AtomArray +from reciprocalspaceship.utils import add_rfree +from sampleworks.core.rewards.structure_factor import ( + _MIN_RETAINED_REFLECTION_FRACTION, + _MIN_RETAINED_REFLECTIONS, + StructureFactorRewardFunction, +) +from sampleworks.utils.atom_array_utils import ( + build_pairwise_altloc_arrays, + find_all_altloc_ids, +) + +from tests.rewards.reward_input_helpers import build_reward_input_tensors_without_coords + + +def make_prepared_reward(mtz_path, atom_array, device: torch.device, **kwargs): + """Construct and ``prepare()`` an SF reward with a per-test config. + + Touches the SFcalculator forward model (heavier than a header read), so callers are marked + ``gpu``. ``kwargs`` pass straight to the constructor (e.g. ``bulk_solvent``, + ``normalize_amplitude``, ``exclude_free_reflections``, ``expcolumns``). The single fixed + config lives in the ``reward_function_1vme_sf`` fixture; these tests need varied configs. + """ + rf = StructureFactorRewardFunction(mtz_path, **kwargs) + rf.prepare(atom_array, device=device) + return rf + + +def write_toy_mtz( + path: Path, + *, + cell: gemmi.UnitCell, + spacegroup: str, + column_pairs: tuple[tuple[str, str], ...], +) -> Path: + """Write a minimal MTZ carrying only header metadata, and return ``path``. + + Five toy reflections written for the reward function's construction-time tests. Each entry + of ``column_pairs`` becomes one ``(amplitude, sigma)`` pair, amplitudes scaled per pair so + the sets are distinguishable. + + Parameters + ---------- + path + Destination ``.mtz`` file. + cell + Unit cell written to the CELL header. Pass a default-constructed ``gemmi.UnitCell()`` + to emulate an MTZ whose cell is gemmi's placeholder. + spacegroup + Space group as a Hermann-Mauguin string. Required: gemmi refuses to write an MTZ with + no space group ("Cannot write Mtz which has no space group") and rs raises before that + ("has no space group information"), so there is no ``None`` to pass here. + column_pairs + One ``(amplitude, sigma)`` name pair per structure-factor set to write. + + Returns + ------- + Path + The ``path`` that was written, for convenient use as a fixture return value. + """ + hkl = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 1], [2, 1, 0]], dtype=np.int32) + amplitudes = np.array([100.0, 90.0, 80.0, 70.0, 60.0], dtype=np.float32) + columns = {"H": hkl[:, 0], "K": hkl[:, 1], "L": hkl[:, 2]} + for pair_index, (amplitude_col, sigma_col) in enumerate(column_pairs): + columns[amplitude_col] = amplitudes * (1.0 + 0.1 * pair_index) + columns[sigma_col] = amplitudes / 10.0 + dataset = rs.DataSet(columns, cell=cell, spacegroup=spacegroup).infer_mtz_dtypes() + dataset.set_index(["H", "K", "L"], inplace=True) + dataset.write_mtz(str(path)) + return path + + +@pytest.fixture(scope="module") +def sf_true_inputs(test_coordinates_1vme_sf, device: torch.device) -> dict[str, torch.Tensor]: + """``__call__`` kwargs (batch=1) for the true 1vme structure. + + The reward-agnostic contract tests already exercise the generic interface/gradient/batch + behavior against the SF reward's default config, so the tests here only add the + config-specific deltas and reuse these inputs. + """ + coords, atom_array = test_coordinates_1vme_sf + elements, b_factors, occupancies = build_reward_input_tensors_without_coords(atom_array, device) + return dict( + coordinates=coords.unsqueeze(0), + elements=elements.unsqueeze(0), + b_factors=b_factors.unsqueeze(0), + occupancies=occupancies.unsqueeze(0), + ) + + +@pytest.fixture(scope="module") +def sf_ensemble_inputs( + structure_1vme_sf, device: torch.device +) -> tuple[dict[str, torch.Tensor], AtomArray]: + """A 2-conformer ensemble (altloc A vs B) for the 1vme structure. + + The batch=2 counterpart to `sf_true_inputs`, returned as ``(__call__ kwargs, atom array)`` + tuple. The atom array is used for ``prepare()`` to build the SFcalculator. + + ``build_pairwise_altloc_arrays`` pairs each altloc with the shared blank-altloc atoms and + filters to their common atoms, so both frames share one topology and differ *only* in the + alternate-conformation coordinates. Residues modeled in only one altloc have no counterpart + and are dropped to keep that shared topology. This divergence is what separates behavior + that is nonlinear in the per-conformer density from behavior that is not — for example, the + different bulk solvent modes. + + SFcalculator has no per-conformer occ/B axis, so ``b_factors`` is shared from altloc-A + across the batch, and occupancy shared at the uniform ``1/batch_size = 0.5``. + """ + altloc_ids = sorted(find_all_altloc_ids(structure_1vme_sf)) # ["A", "B"] + altloc_a, altloc_b = altloc_ids + altloc_pairs = build_pairwise_altloc_arrays(structure_1vme_sf, altloc_ids) + array_a, array_b = altloc_pairs[(altloc_a, altloc_b)] + ref_atom_array = array_a[0] # state-A reference; retains element/b_factor/occupancy + + coords = torch.stack( + [ + torch.from_numpy(array_a.coord[0]), # altloc-A conformer [N, 3] + torch.from_numpy(array_b.coord[0]), # altloc-B conformer [N, 3] + ] + ).to(device=device, dtype=torch.float32) # [2, N, 3] + + elements, b_factors, _ = build_reward_input_tensors_without_coords(ref_atom_array, device) + n_atoms = coords.shape[1] + occ = torch.full((n_atoms,), 0.5, device=device) # uniform 1/batch_size + reward_inputs = dict( + coordinates=coords, + elements=elements.unsqueeze(0).expand(2, -1), + b_factors=b_factors.unsqueeze(0).expand(2, -1), + occupancies=occ.unsqueeze(0).expand(2, -1), + ) + return reward_inputs, ref_atom_array + + +class TestStructureFactorConstruction: + """Construction-time behavior, on CPU: no GPU and no ``prepare()``/SF compute. + + ``__init__`` validates its config and reads exactly three things from the MTZ + (``_resolve_mtz_metadata``): the unit cell, the space group, and the amplitude/sigma column + layout. For efficiency, we build a ``toy_multi_set_mtz`` instead of taking the session-scoped + ``mtz_path_1vme`` that on a CPU costs ~18 s, and build it once per class since several tests + here only read it. The MTZ metadata are specified as class constants because they should be + immutable for assertions. + """ + + # a, b, c (A), alpha, beta, gamma (deg) — a gemmi.UnitCell is built from these on demand. + TOY_CELL_PARAMETERS = (11.0, 22.0, 33.0, 90.0, 100.0, 120.0) + TOY_SPACE_GROUP = "P 1 2 1" # Hermann-Mauguin string + PROTEIN_COLUMNS = ("Fprotein", "SIGFprotein") # (amplitude, sigma), protein-only + TOTAL_COLUMNS = ("Ftotal", "SIGFtotal") # (amplitude, sigma), protein + bulk solvent + + @pytest.fixture(scope="class") + @classmethod + def toy_multi_set_mtz(cls, tmp_path_factory: pytest.TempPathFactory) -> Path: + """A minimal multi-set MTZ carrying header metadata on unit cell and space group.""" + return write_toy_mtz( + tmp_path_factory.mktemp("sf_construction") / "toy_multi_set.mtz", + cell=gemmi.UnitCell(*cls.TOY_CELL_PARAMETERS), + spacegroup=cls.TOY_SPACE_GROUP, + column_pairs=(cls.PROTEIN_COLUMNS, cls.TOTAL_COLUMNS), + ) + + # The sibling test for mtz without space-group (_resolve_mtz_metadata raises when MTZ + # has space group None) is not tested, because gemmi / rs prevents writing such an MTZ, + # although such a MTZ could exist from other programs or corrupted files. + def test_mtz_without_unit_cell_raises(self, tmp_path: Path): + """An MTZ left with gemmi's placeholder cell is rejected at construction. + + Unlike the space group, this input *is* constructible: a default ``gemmi.UnitCell`` is + (1, 1, 1, 90, 90, 90), which both writers accept and ``is_crystal()`` reports as False. + """ + mtz_path = write_toy_mtz( + tmp_path / "no_unit_cell.mtz", + cell=gemmi.UnitCell(), + spacegroup=self.TOY_SPACE_GROUP, + column_pairs=(self.PROTEIN_COLUMNS,), + ) + with pytest.raises(ValueError, match="carries no unit cell"): + StructureFactorRewardFunction(mtz_path) + + def test_multi_set_mtz_requires_expcolumns(self, toy_multi_set_mtz): + """A multi-set MTZ with no ``expcolumns`` is ambiguous and fails fast, rather than + silently auto-selecting the first amplitude/sigma pair. + """ + with pytest.raises(ValueError, match="Multiple SFAmplitude columns"): + StructureFactorRewardFunction(toy_multi_set_mtz) + + def test_explicit_expcolumns_override_selection(self, toy_multi_set_mtz): + """Explicit ``expcolumns`` are used verbatim, overriding the auto-selected first pair.""" + reward = StructureFactorRewardFunction( + toy_multi_set_mtz, expcolumns=list(self.TOTAL_COLUMNS) + ) + assert reward.expcolumns == list(self.TOTAL_COLUMNS) + + def test_unknown_expcolumns_raise(self, toy_multi_set_mtz): + """Explicit ``expcolumns`` naming a column absent from the MTZ fail fast at construction.""" + with pytest.raises(ValueError, match="is not among the dataset's SFAmplitude columns"): + StructureFactorRewardFunction( + toy_multi_set_mtz, expcolumns=["Fnonexistent", self.PROTEIN_COLUMNS[1]] + ) + + @pytest.mark.parametrize( + "bad_expcolumns", + [ + ["Fprotein"], # missing the sigma + ["Fprotein", "SIGFprotein", "Ftotal"], # too many + [], + "FP", # a bare string of length 2 would otherwise split into ["F", "P"] + ], + ) + def test_malformed_expcolumns_raise(self, toy_multi_set_mtz, bad_expcolumns): + """``expcolumns`` that is not an ``[amplitude, sigma]`` pair fails fast, including a + bare string, which would otherwise be indexed character-wise into a bogus column pair. + """ + with pytest.raises(ValueError, match=r"expcolumns must be a \[amplitude, sigma\] pair"): + StructureFactorRewardFunction(toy_multi_set_mtz, expcolumns=bad_expcolumns) + + def test_cell_and_space_group_read_from_mtz(self, toy_multi_set_mtz): + """The cell and space group are parsed off the MTZ, the only source for them.""" + reward = StructureFactorRewardFunction( + toy_multi_set_mtz, expcolumns=list(self.TOTAL_COLUMNS) + ) + assert reward.space_group == self.TOY_SPACE_GROUP + assert reward.unit_cell.parameters == pytest.approx(self.TOY_CELL_PARAMETERS, abs=1e-3) + + @pytest.mark.parametrize("bad_partition", [0, -5, 10.0, 2.5, "10", None, True]) + def test_invalid_batch_partition_raises(self, toy_multi_set_mtz, bad_partition): + """A non-integer or non-positive ``batch_partition`` (an OOM knob) fails fast. + + The check precedes column resolution, so no ``expcolumns`` are needed. + """ + with pytest.raises(ValueError, match="batch_partition must be a positive integer"): + StructureFactorRewardFunction(toy_multi_set_mtz, batch_partition=bad_partition) + + def test_unknown_bulk_solvent_mode_raises(self, toy_multi_set_mtz): + """An unrecognized ``bulk_solvent`` mode fails fast at construction.""" + with pytest.raises(ValueError, match="bulk_solvent must be one of"): + StructureFactorRewardFunction(toy_multi_set_mtz, bulk_solvent="not_off") + + def test_reserved_sfcalculator_kwargs_raise(self, toy_multi_set_mtz): + """``sfcalculator_kwargs`` should not override reserved keys for the reward function's + ``__init__ arguments``.""" + with pytest.raises(ValueError, match="may not override reserved keys"): + StructureFactorRewardFunction( + toy_multi_set_mtz, + expcolumns=list(self.TOTAL_COLUMNS), + sfcalculator_kwargs={"dmin": 1.0, "device": "cuda"}, + ) + + def test_call_before_prepare_raises(self, toy_multi_set_mtz): + """Evaluating the reward before ``prepare()`` gives a clear error.""" + reward = StructureFactorRewardFunction( + toy_multi_set_mtz, expcolumns=list(self.TOTAL_COLUMNS) + ) + n_atoms = 4 + with pytest.raises(RuntimeError, match=r"prepare\(\) must be called"): + reward( + coordinates=torch.zeros(1, n_atoms, 3), + elements=torch.zeros(1, n_atoms, dtype=torch.int64), + b_factors=torch.full((1, n_atoms), 20.0), + occupancies=torch.ones(1, n_atoms), + ) + + +@pytest.mark.gpu +class TestStructureFactorOccupancy: + """``__call__``'s input guards, both of which run before any SF compute: the atom count must + match the topology fixed by ``prepare()``, and occupancy/B must be broadcast-identical across + the batch (SFC has no per-conformer occupancy/B axis). + """ + + def test_atom_count_mismatch_raises(self, reward_function_1vme_sf, sf_ensemble_inputs): + """Inputs whose atom count disagrees with the prepared topology fail fast. + + ``reward_function_1vme_sf`` is prepared on the full 1vme topology (both altlocs present), + while ``sf_ensemble_inputs`` is built on the shared-altloc topology (blank plus the atoms + common to A and B), which is strictly smaller. + """ + ensemble_reward_inputs, _ = sf_ensemble_inputs + with pytest.raises(ValueError, match="topology built by"): + reward_function_1vme_sf(**ensemble_reward_inputs) + + @pytest.mark.parametrize("field", ["occupancies", "b_factors"]) + def test_per_conformer_occupancy_or_b_raises( + self, reward_function_1vme_sf, test_coordinates_1vme_sf, device, field + ): + """Per-conformer (non-broadcast) occupancy/B is rejected, not silently dropped. + + SFcalculator batches only coordinates, so __call__ honors a single shared occupancy/B + vector. Production always feeds broadcast-identical rows (the batch=1 and identical-row + ensemble cases are covered by the shared contract tests); this pins that genuinely + per-conformer values raise instead of being silently ignored. The guard runs before any + SF compute, so this also documents that the [batch, n_atoms] signature does NOT mean + per-conformer occupancy/B is supported. + """ + coords, atom_array = test_coordinates_1vme_sf + elements, b_factors, occupancies = build_reward_input_tensors_without_coords( + atom_array, device + ) + + batch = 3 + kwargs = dict( + coordinates=coords.unsqueeze(0).expand(batch, -1, -1), + elements=elements.unsqueeze(0).expand(batch, -1), + b_factors=b_factors.unsqueeze(0).expand(batch, -1), + occupancies=occupancies.unsqueeze(0).expand(batch, -1), + ) + # Make one conformer's values genuinely differ (additive shift works regardless of zeros). + perturbed = kwargs[field].clone() + perturbed[1] += 1.0 + kwargs[field] = perturbed + + with pytest.raises(ValueError, match="identical across the batch"): + reward_function_1vme_sf(**kwargs) + + +@pytest.mark.gpu +class TestStructureFactorBulkSolvent: + """The bulk-solvent modes. ``off`` scores ``|Fprotein|`` (covered by the contract tests); + ``combined``/``per_conformer`` score ``|Ftotal|`` with default-scaled solvent. The two + ``|Ftotal|`` modes differ only for a real ensemble. Comparisons are relative, so they don't + depend on the raw ``|F|`` scale. + """ + + # Four prepared Ftotal rewards: {combined, per_conformer} x {full, ensemble} topology, all + # raw |F| and all scored against the MTZ's Ftotal set. The topology differs on atom counts: + # _full -- a single 1vme structure with both altlocs present + # _ensemble -- topology consists of blank + shared atoms between the two altlocs + @pytest.fixture(scope="class") + @classmethod + def reward_with_solvent_combined_full(cls, mtz_path_1vme, structure_1vme_sf, device): + """``combined`` on the full 1vme topology.""" + return make_prepared_reward( + mtz_path_1vme, + structure_1vme_sf, + device, + expcolumns=["Ftotal", "SIGFtotal"], + bulk_solvent="combined", + normalize_amplitude=False, + ) + + @pytest.fixture(scope="class") + @classmethod + def reward_with_solvent_per_conformer_full(cls, mtz_path_1vme, structure_1vme_sf, device): + """``per_conformer`` on the full 1vme topology.""" + return make_prepared_reward( + mtz_path_1vme, + structure_1vme_sf, + device, + expcolumns=["Ftotal", "SIGFtotal"], + bulk_solvent="per_conformer", + normalize_amplitude=False, + ) + + @pytest.fixture(scope="class") + @classmethod + def reward_with_solvent_combined_ensemble(cls, mtz_path_1vme, sf_ensemble_inputs, device): + """``combined`` on the altloc-ensemble topology.""" + _, ref_atom_array = sf_ensemble_inputs + return make_prepared_reward( + mtz_path_1vme, + ref_atom_array, + device, + expcolumns=["Ftotal", "SIGFtotal"], + bulk_solvent="combined", + normalize_amplitude=False, + ) + + @pytest.fixture(scope="class") + @classmethod + def reward_with_solvent_per_conformer_ensemble(cls, mtz_path_1vme, sf_ensemble_inputs, device): + """``per_conformer`` on the altloc-ensemble topology.""" + _, ref_atom_array = sf_ensemble_inputs + return make_prepared_reward( + mtz_path_1vme, + ref_atom_array, + device, + expcolumns=["Ftotal", "SIGFtotal"], + bulk_solvent="per_conformer", + normalize_amplitude=False, + ) + + def test_ftotal_modes_agree_for_single_conformer( + self, + reward_with_solvent_combined_full, + reward_with_solvent_per_conformer_full, + sf_true_inputs, + ): + """For a single conformer (batch_size=1) ``mask()`` and ```` are the same + mask, so ``combined`` and ``per_conformer`` give the same loss.""" + torch.testing.assert_close( + reward_with_solvent_combined_full(**sf_true_inputs), + reward_with_solvent_per_conformer_full(**sf_true_inputs), + ) + + def test_combined_fits_ftotal_column( + self, + mtz_path_1vme, + structure_1vme_sf, + reward_with_solvent_combined_full, + sf_true_inputs, + device, + ): + """Adding default-scaled bulk solvent (``combined``) fits the synthetic ``Ftotal`` + column far better than protein-only (``off``). It should also reproduce ``Ftotal`` + as the synthetic MTZ should be generated with the same structure and config.""" + reward_with_solvent_off = make_prepared_reward( + mtz_path_1vme, + structure_1vme_sf, + device, + expcolumns=["Ftotal", "SIGFtotal"], + bulk_solvent="off", + normalize_amplitude=False, + ) + loss_combined = reward_with_solvent_combined_full(**sf_true_inputs) + loss_off = reward_with_solvent_off(**sf_true_inputs) + assert loss_combined < loss_off + assert loss_combined.item() < 1e-6, ( + "combined mode should reproduce MTZ's Ftotal with loss less than 1e-6, but the " + f"combined loss here is {loss_combined.item():.6e}" + ) + + def test_per_conformer_averages_masks_over_ensemble( + self, reward_with_solvent_per_conformer_full, sf_true_inputs + ): + """batch_size=2 *identical* conformers (occ 1/2 each) score the same as the single + conformer. + + This guards the equal-population assumption in ``per_conformer``: the protein path bakes + ``atom_occ`` into each conformer and *sums*, while the solvent path averages + scale-invariant per-conformer masks (``Fmask_HKL_batch.mean(dim=0)``) — a hardcoded + uniform ``1/batch_size`` weight. Both are population *averages*, so occ = 1/batch_size + keeps them consistent and the ensemble matches the single conformer; this would break if + ``per_conformer`` summed (rather than averaged) the masks. Non-uniform per-conformer + occupancy is properly rejected by the reward now — see + ``TestStructureFactorOccupancy.test_per_conformer_occupancy_or_b_raises``. + """ + identical_ensemble = dict( + coordinates=sf_true_inputs["coordinates"].expand(2, -1, -1), + elements=sf_true_inputs["elements"].expand(2, -1), + b_factors=sf_true_inputs["b_factors"].expand(2, -1), + occupancies=(sf_true_inputs["occupancies"] / 2).expand(2, -1), + ) + torch.testing.assert_close( + reward_with_solvent_per_conformer_full(**sf_true_inputs), + reward_with_solvent_per_conformer_full(**identical_ensemble), + ) + + def test_ftotal_modes_diverge_for_distinct_conformer_ensemble( + self, + reward_with_solvent_combined_ensemble, + reward_with_solvent_per_conformer_ensemble, + sf_ensemble_inputs, + ): + """For a genuine 2-conformer ensemble (altloc A vs B), ``mask() != ``, so + ``combined`` and ``per_conformer`` give *different* losses. + + This is the behavior that justifies keeping both modes — the agree-cases above + (batch_size=1 and batch_size=2 *identical* conformers) never exercise the nonlinearity. + Here the two frames differ in the alternate-conformation coordinates, so the combined + mask (built from the summed density) and the per-conformer mean of masks genuinely + diverge. + + The mock.patch.object supplements the numerical difference test by checking the dispatch + logic and ensure that the correct SFcalculator method for Fsolvent is called. + """ + ensemble_reward_inputs, _ = sf_ensemble_inputs + loss_combined = reward_with_solvent_combined_ensemble(**ensemble_reward_inputs) + # Specifically checking the dispatch logic in _compute_ensemble_ftotal, where + # ``bulk_solvent="per_conformer"`` should route to a calc_fsolvent_batch call. + with mock.patch.object( + reward_with_solvent_per_conformer_ensemble.sfc, + "calc_fsolvent_batch", + wraps=reward_with_solvent_per_conformer_ensemble.sfc.calc_fsolvent_batch, + ) as sfc_calc_fsolvent_batch: + loss_per_conformer = reward_with_solvent_per_conformer_ensemble( + **ensemble_reward_inputs + ) + assert sfc_calc_fsolvent_batch.call_count == 1 + assert torch.isfinite(loss_combined) and torch.isfinite(loss_per_conformer) + + relative_gap = ((loss_per_conformer - loss_combined) / loss_combined).item() + relative_threshold = 0.2 + # The true relative difference is ~1.24 and the noise from permuting the batch for + # per_conformer is ~0.02. Subjected to change if loss function or normalization changes. + assert abs(relative_gap) > relative_threshold, ( + f"combined and per_conformer losses differ relatively by {relative_gap:+.3f}, " + f"under the {relative_threshold} threshold we set: {loss_combined.item()} vs " + f"{loss_per_conformer.item()}" + ) + + @pytest.mark.parametrize( + "reward_fixture", + ["reward_with_solvent_combined_ensemble", "reward_with_solvent_per_conformer_ensemble"], + ) + def test_bulk_solvent_branch_carries_gradient( + self, request: pytest.FixtureRequest, reward_fixture: str, sf_ensemble_inputs + ): + """Fsolvent under ``combined`` and ``per_conformer`` modes should contribute gradients + to the coordinates. + + ``Ftotal = kiso * aniso * (Fprotein_HKL + kmask * Fmask_HKL)``, and ``Fmask_HKL`` + descends from ``Fprotein_asu_batch`` rather than from ``Fprotein_HKL``. + """ + reward = request.getfixturevalue(reward_fixture) + ensemble_reward_inputs, _ = sf_ensemble_inputs + coordinates = ensemble_reward_inputs["coordinates"].clone().requires_grad_(True) + loss = reward(**{**ensemble_reward_inputs, "coordinates": coordinates}) + + Fprotein_HKL, Fmask_HKL = reward.sfc.Fprotein_HKL, reward.sfc.Fmask_HKL + assert Fprotein_HKL.requires_grad and Fmask_HKL.requires_grad, ( + f"{reward.bulk_solvent} mode should have requires_grad=True on both Fprotein_HKL " + f"and Fmask_HKL, but got {Fprotein_HKL.requires_grad} and {Fmask_HKL.requires_grad}" + ) + (total_gradient,) = torch.autograd.grad(loss, coordinates, retain_graph=True) + d_protein, d_solvent = torch.autograd.grad( + loss, (Fprotein_HKL, Fmask_HKL), retain_graph=True + ) + (via_protein,) = torch.autograd.grad( + Fprotein_HKL, coordinates, grad_outputs=d_protein, retain_graph=True + ) + (via_solvent,) = torch.autograd.grad(Fmask_HKL, coordinates, grad_outputs=d_solvent) + + for branch, gradient in (("Fprotein", via_protein), ("Fmask", via_solvent)): + assert 0 < gradient.norm().item() < 1e6, ( + f"the coordinate gradient through {branch} in {reward.bulk_solvent} mode has " + f"norm {gradient.norm().item():.3e}" + ) + + residual = ( + (via_protein + via_solvent - total_gradient).norm() / total_gradient.norm() + ).item() + assert residual < 1e-3, ( + "coordinate gradients from the Fprotein and Fmask branches do not sum back to " + f"d(loss)/d(coords) (relative residual norm {residual:.2e})" + ) + + +@pytest.mark.gpu +class TestStructureFactorConfig: + """Config knobs beyond the forward model: reflection selection.""" + + def test_exclude_free_set_drops_reflections(self, mtz_path_1vme, structure_1vme_sf, device): + """``exclude_free_reflections=True`` drops the R-free test set from the scored mask. + + The synthetic MTZ's free column is ``R-free-flags`` with the test set flagged 1 + (rs/Phenix convention), so SFcalculator is pointed at it explicitly — its defaults + (``FreeR_flag`` / testset value 0) don't match. Outliers are always excluded regardless. + """ + free_flag_kwargs = {"freeflag": "R-free-flags", "testset_value": 1} + reward_all = make_prepared_reward( + mtz_path_1vme, + structure_1vme_sf, + device, + expcolumns=["Fprotein", "SIGFprotein"], + exclude_free_reflections=False, + sfcalculator_kwargs=free_flag_kwargs, + ) + reward_work = make_prepared_reward( + mtz_path_1vme, + structure_1vme_sf, + device, + expcolumns=["Fprotein", "SIGFprotein"], + exclude_free_reflections=True, + sfcalculator_kwargs=free_flag_kwargs, + ) + assert reward_all.sfc.free_flag.any() # safety: the free set is actually recognized + assert int(reward_work._reflection_mask.sum()) < int(reward_all._reflection_mask.sum()) + + def test_empty_reflection_mask_raises( + self, mtz_path_1vme, structure_1vme_sf, device, tmp_path: Path + ): + """A mask that scores no reflection fails in ``prepare()``. + + The all-free MTZ is built with rs's own ``add_rfree`` (``fraction=1.0``) so the flag + column name and convention match the committed MTZ's (``R-free-flags``, test set = 1). + """ + dataset = add_rfree(rs.read_mtz(str(mtz_path_1vme)), fraction=1.0, seed=0) + mtz_all_free = tmp_path / "1vme_all_free.mtz" + dataset.write_mtz(str(mtz_all_free)) + with pytest.raises(ValueError, match="No reflections remain"): + make_prepared_reward( + mtz_all_free, + structure_1vme_sf, + device, + expcolumns=["Fprotein", "SIGFprotein"], + exclude_free_reflections=True, + sfcalculator_kwargs={"freeflag": "R-free-flags", "testset_value": 1}, + ) + + def test_inverted_testset_value_warns(self, mtz_path_1vme, structure_1vme_sf, device, caplog): + """An inverted ``testset_value`` keeps only the test set, and is warned about. + + Inverting the flag keeps the fraction of valid reflections small although the + absolute size can still be large. + """ + with caplog.at_level(logging.WARNING): + reward = make_prepared_reward( + mtz_path_1vme, + structure_1vme_sf, + device, + expcolumns=["Fprotein", "SIGFprotein"], + exclude_free_reflections=True, + sfcalculator_kwargs={"freeflag": "R-free-flags", "testset_value": 0}, + ) + n_used, n_total = int(reward._reflection_mask.sum()), len(reward.sfc.Fo) + # check that the fraction precondition held and the floor's precondition did not hold + assert n_used < _MIN_RETAINED_REFLECTION_FRACTION * n_total + assert n_used >= _MIN_RETAINED_REFLECTIONS + assert f"Only {n_used}/{n_total} reflections remain" in caplog.text + # check that only the warning message from the fraction threshold is logged + assert "testset_value" in caplog.text + assert "the MTZ reflection range" not in caplog.text + + def test_small_reflection_set_warns(self, mtz_path_1vme, structure_1vme_sf, device, caplog): + """A reflection set too small of absolute size to guide coordinates should warn. + + Truncating to 10 A leaves only a few hundred reflections but the fraction of valid + reflections can still be large. + """ + with caplog.at_level(logging.WARNING): + reward = make_prepared_reward( + mtz_path_1vme, + structure_1vme_sf, + device, + expcolumns=["Fprotein", "SIGFprotein"], + resolution=10.0, + ) + n_used, n_total = int(reward._reflection_mask.sum()), len(reward.sfc.Fo) + # check that the floor precondition held and the fraction's precondition did not hold + assert n_used < _MIN_RETAINED_REFLECTIONS + assert n_used >= _MIN_RETAINED_REFLECTION_FRACTION * n_total + # check that only the warning message from the absolute threshold is logged + assert "the MTZ reflection range" in caplog.text + assert f"{n_used}/{n_total}" not in caplog.text