diff --git a/src/sampleworks/core/rewards/torchref_rewards.py b/src/sampleworks/core/rewards/torchref_rewards.py new file mode 100644 index 00000000..aa65f425 --- /dev/null +++ b/src/sampleworks/core/rewards/torchref_rewards.py @@ -0,0 +1,972 @@ +"""TorchRef-backed reward functions for reciprocal-space fitting. + +Provides :class:`TorchRefXrayRewardFunction`, which scores coordinates against +experimental structure factors using torchref's scaling and maximum-likelihood stack: +per-resolution-bin scale, anisotropic scale tensor, refined bulk solvent, and a +sigma_A target carrying a model-error term. + +A ``torchref.model.ModelFT`` is built directly rather than through ``Model.load()``. +Cell and space group come from the MTZ, a hand-built ``pdb`` DataFrame carries the +topology, and coordinates and occupancies are caller-owned tensors held in +:class:`_TensorSlot`, so gradients flow from ``target()`` back to the caller's tensor. + +Structure factors are linear over atoms:: + + F(h) = sum_j f_j occ_j exp(2 pi i h . x_j) + +so ``C`` conformers at occupancy ``1/C`` form a single structure-factor calculation +over a ``C * n_atoms`` stack rather than ``C`` separate ones. The stack is +conformer-major, matching the ``[batch, n]`` layout of the reward protocol. + +Scale, bulk solvent, sigma_A and ADPs are nuisance parameters refit periodically by +:meth:`TorchRefXrayRewardFunction.refresh_nuisance_parameters`. The objective changes +discontinuously at each refresh. + +Usage constraints +----------------- +- Returns an unnormalized summed negative log-likelihood, O(1e5) for a real dataset, + not sign-constrained. Guidance step sizes do not transfer from MSE-valued rewards. +- sigma_A is fitted on the free set, so ``R_free`` is a diagnostic here rather than an + independent validation statistic. +- The bulk-solvent mask is built at the current coordinates under ``no_grad``. Use + ``bulk_solvent=False`` when guidance starts from near-noise coordinates. +- ADPs are refined during maintenance, one per asymmetric-unit atom shared across + conformers. ``adp_weight=0`` freezes them at ``b_factor``. +- Geometry restraints are off by default (``geometry_weight=0``). A zero weight means + the target is never constructed. +""" + +from __future__ import annotations + +import functools +from collections import Counter +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import numpy as np +import pandas as pd +import torch +from jaxtyping import Float, Int +from loguru import logger +from sampleworks.utils.elements import elements_to_scattering_indices + + +if TYPE_CHECKING: + from biotite.structure import AtomArray + + +# torchref restricts the scale-fit objective to these two: an alpha-centred mode is +# degenerate with the scale being fitted and drives it to absorb 1/alpha. +_SCALE_TARGETS = ("nll", "ml_noalpha") + +# sigma_A maximum-likelihood rows carry a model-error term; "nll"/"ls" do not but are +# stateless (their maintenance() is a no-op). +_TARGET_MODES = ("ml", "ml_noalpha", "ml_full", "nll_beta", "nll", "ls", "ls_wunit_k1") + +DEFAULT_B_FACTOR = 20.0 +"""Initial isotropic ADP (Ų) assigned to every atom.""" + + +class _TensorSlot: + """Callable holder for a caller-owned tensor, read by ``ModelFT`` as a wrapper. + + Returns the held tensor verbatim rather than detaching it, so autograd reaches the + caller's leaf. Implements the three members ``ModelFT`` requires: the call itself + (``Model.get_iso``), ``refinable_params`` (``ModelFT._check_forward_dtype`` + float-dtype probe) and ``fixed_values`` (``Model.get_aniso`` placeholder sizing). + + Not an ``nn.Module``, so the held tensors stay out of ``model.parameters()`` and + cannot be picked up by an optimizer. + + Parameters + ---------- + t : torch.Tensor + The held tensor. Reassign :attr:`t` to rebind; nothing is copied. + """ + + __slots__ = ("t",) + + def __init__(self, t: torch.Tensor): + self.t = t + + def __call__(self) -> torch.Tensor: + return self.t + + @property + def refinable_params(self) -> torch.Tensor: + """Float-dtype probe for ``ModelFT._check_forward_dtype``.""" + return self.t + + @property + def fixed_values(self) -> torch.Tensor: + """Dtype/device/shape source for ``Model.get_aniso``'s empty placeholders.""" + return self.t + + +@functools.lru_cache(maxsize=1) +def _shared_adp_cls() -> type: + """Build (once) the shared-ADP wrapper. Lazy so importing this module needs no torchref. + + Holds one refinable B per asymmetric-unit atom and expands it across every conformer + in the stack, so the same atom carries one B-factor rather than C independent ones. + + Subclasses ``PositiveMixedTensor`` so the refinable leaf keeps torchref's ADP + parameterisation (log-space, positive by construction) -- the form + ``parameters_of_types(("adp",))`` and the ADP restraint targets expect. + """ + from torchref.model.parameter_wrappers import PositiveMixedTensor + + class _SharedADP(PositiveMixedTensor): + """``(n_asu,)`` refinable log-B, expanded to ``(n_conformers * n_asu,)``.""" + + def __init__(self, *args: Any, n_conformers: int = 1, **kwargs: Any): + super().__init__(*args, **kwargs) + self._n_conformers = int(n_conformers) + + def forward(self) -> torch.Tensor: + """Return the shared B-factors expanded to ``(n_conformers * n_asu,)``.""" + # repeat, not repeat_interleave: the stack is conformer-major. Gradients from + # all C copies sum onto the shared leaf. + # + # Side effect: the ADP restraint targets read model.adp() and so see the + # expanded (C*n_asu,) array, evaluating each restraint C times over identical + # values. The effective ADP weight is C x its nominal value, and adp_weight is + # therefore not comparable across conformer counts. + per_asu = super().forward() + if self._n_conformers == 1: + return per_asu + return per_asu.repeat(self._n_conformers) + + return _SharedADP + + +@functools.lru_cache(maxsize=1) +def _external_model_cls() -> type: + """Build (once) the ``ModelFT`` subclass this module drives. + + Lazy so importing this module needs no torchref. Adds coordinate/occupancy setters + and disables ``CachedForwardMixin``'s forward cache. That cache fingerprints only + ``parameters()`` and ``buffers()``; caller-owned coordinates are neither, so it would + never invalidate and every call after the first would return a stale F_calc. Each + call now builds a fresh graph, which also removes the mixin's ``retain_graph`` + requirement on a second backward. + """ + from torchref.model import ModelFT + + class _ExternalModelFT(ModelFT): + """``ModelFT`` driven by caller-owned tensors, with the forward cache off.""" + + def set_coordinates(self, xyz: torch.Tensor) -> None: + """Bind ``xyz``. The single place coordinates are set.""" + self.xyz.t = xyz + + def set_occupancies(self, occ: torch.Tensor) -> None: + """Bind ``occ``. The single place occupancies are set.""" + self.occupancy.t = occ + + def __call__(self, *args: Any, **kwargs: Any) -> torch.Tensor: + kwargs.pop("recalc", None) # the cache is off; recalc is meaningless + return self.forward(*args, **kwargs) + + return _ExternalModelFT + + +def _conformer_tag(i: int) -> str: + """Single-character conformer label: A..Z then a..z, wrapping past 52.""" + alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + return alphabet[i % len(alphabet)] + + +def _resolve_structure(structure: AtomArray | str | Path) -> AtomArray: + """Return an ``AtomArray``, loading from a PDB/mmCIF path if given one. + + Parameters + ---------- + structure + An ``AtomArray``, or a path to a structure file. + + Returns + ------- + AtomArray + The atom array to take topology from. + """ + if isinstance(structure, (str, Path)): + from sampleworks.utils.atom_array_utils import load_structure_with_altlocs + + logger.info(f"Loading topology from {structure}") + return load_structure_with_altlocs(Path(structure)) + return structure + + +class TorchRefXrayRewardFunction: + """Reciprocal-space reward using torchref's scaling and likelihood stack. + + Scores an ensemble of conformers against experimental amplitudes from an MTZ. + Construction is two-phase: ``__init__`` reads the reflection data, :meth:`prepare` + takes the topology, and the torchref model is built on first call once the ensemble + size is known. + + Parameters + ---------- + mtzfile + Path to the MTZ holding the target amplitudes. The MTZ is the sole authority for + the unit cell and space group; there are no arguments to override them. + structure + Optional topology (``AtomArray`` or path to PDB/mmCIF). When given, + :meth:`prepare` is called immediately. Omit it to defer to an explicit + :meth:`prepare`. + device + Device to place the reflection data and model on. Only used when + ``structure`` is given; otherwise pass it to :meth:`prepare`. + resolution + High-resolution cutoff (dmin, Å). ``None`` (default) keeps the MTZ's own range. + Also sets the density grid spacing. + target_mode + torchref X-ray target. ``"ml"`` (default) is the sigma_A maximum-likelihood row + and carries a model-error term. ``"nll"`` and ``"ls"`` are stateless and cheaper + but do not penalise overfitting. + bulk_solvent + Fit and apply a bulk-solvent contribution. Default True. Costs a real-space mask + build plus an FFT on every refresh, and is non-differentiable with respect to + coordinates. + b_factor + Initial isotropic ADP (Ų) shared by every atom. + nbins + Resolution bins for the scale. torchref may lower this for sparse data. + refresh_interval + Refit the nuisance parameters every this many calls; the first call always + refreshes. + scale_target + Objective for the scale fit; must be one of ``("nll", "ml_noalpha")``. + use_set + Reflection subset the loss is summed over: ``"work"`` (default), ``"free"`` + or ``"val"``. + french_wilson + Derive amplitudes from intensities via French-Wilson. Forwarded to + ``ReflectionData.load_mtz``. + column_names + Explicit MTZ column mapping, e.g. ``{"F": "Fprotein", "SIGF": "SIGFprotein"}``. + Required when the MTZ carries more than one amplitude set. + + Raises + ------ + ValueError + For an unknown ``target_mode``/``scale_target``/``use_set``, a + non-positive ``refresh_interval`` or ``b_factor``, or an MTZ without a + usable cell or space group. + """ + + def __init__( + self, + mtzfile: str | Path, + *, + structure: AtomArray | str | Path | None = None, + device: torch.device | str | None = None, + resolution: float | None = None, + target_mode: str = "ml", + bulk_solvent: bool = True, + b_factor: float = DEFAULT_B_FACTOR, + nbins: int = 20, + refresh_interval: int = 10, + scale_target: str = "nll", + adp_weight: float = 0.02, + geometry_weight: float = 0.0, + use_set: str = "work", + french_wilson: bool = True, + column_names: dict | None = None, + ): + if target_mode not in _TARGET_MODES: + raise ValueError(f"target_mode must be one of {_TARGET_MODES}, got {target_mode!r}.") + if scale_target not in _SCALE_TARGETS: + raise ValueError( + f"scale_target must be one of {_SCALE_TARGETS}, got {scale_target!r}. An " + "alpha-centred mode is degenerate with the scale being fitted." + ) + if use_set not in ("work", "free", "val"): + raise ValueError(f"use_set must be 'work', 'free' or 'val'; got {use_set!r}.") + if refresh_interval < 1: + raise ValueError(f"refresh_interval must be >= 1, got {refresh_interval}.") + if b_factor <= 0: + raise ValueError(f"b_factor must be positive, got {b_factor}.") + if adp_weight < 0 or geometry_weight < 0: + raise ValueError( + f"restraint weights must be >= 0; got adp_weight={adp_weight}, " + f"geometry_weight={geometry_weight}. Use 0 to disable a group entirely." + ) + + self.mtzfile = str(mtzfile) + self.target_mode = target_mode + self.bulk_solvent = bulk_solvent + self.b_factor = float(b_factor) + self.nbins = nbins + self.refresh_interval = refresh_interval + self.scale_target = scale_target + self.use_set = use_set + self.adp_weight = float(adp_weight) + self.geometry_weight = float(geometry_weight) + # ADPs are refinable only when something regularises them; on the x-ray term + # alone they would absorb model error. + self.refine_adp = self.adp_weight > 0 + + self._load_reflection_data(resolution, french_wilson, column_names) + + # Populated by prepare(). + self._prepared = False + self._element_symbols: list[str] | None = None + self._expected_codes: torch.Tensor | None = None + self.n_atoms: int | None = None + self.device: torch.device | None = None + # Keyed on ensemble size; see _stack_for. The call counter is per-stack because + # nuisance parameters belong to a triple, so a newly built one is cold-fitted on + # its first use regardless of any other stack's count. + self._stacks: dict[int, tuple[Any, Any, Any]] = {} + self._calls: dict[int, int] = {} + self._states: dict[int, Any] = {} + + if structure is not None: + self.prepare(structure, device=device or "cpu") + + # ------------------------------------------------------------------ + # Construction + # ------------------------------------------------------------------ + def _load_reflection_data( + self, resolution: float | None, french_wilson: bool, column_names: dict | None + ) -> None: + """Read the MTZ into a ``ReflectionData`` and cache its crystal metadata. + + ``ReflectionData.__post_init__`` runs ``setup_scale()`` and ``setup_anisotropy()``, + so ``get_corrected_data()``, which the sigma_A rows require, works from here on. + """ + from torchref import read_mtz + + data = read_mtz( + self.mtzfile, verbose=0, french_wilson=french_wilson, column_names=column_names + ) + if resolution is not None: + data = data.filter_by_resolution(d_min=resolution) + if data.cell is None: + raise ValueError(f"{self.mtzfile} carries no unit cell; cannot compute d-spacings.") + if data.spacegroup is None: + raise ValueError(f"{self.mtzfile} carries no space group gemmi can recognise.") + + self._data = data + self.unit_cell = data.cell.data.detach().cpu().tolist() + self.space_group = data.spacegroup.hm + # dmin drives the density grid spacing as well as the resolution range, so + # ModelFT's 1.0 A default would build a needlessly fine grid for lower-res data. + self.resolution = float(resolution if resolution is not None else data.d_min) + + n_free = int((data.rfree_flags == 0).sum()) + logger.info( + f"Loaded {self.mtzfile}: n_reflections={len(data.hkl)}, " + f"dmin={self.resolution:.2f}A, cell={[round(c, 2) for c in self.unit_cell]}, " + f"space_group={self.space_group}, n_free={n_free}" + ) + if self.target_mode in ("ml", "ml_noalpha", "ml_full", "nll_beta") and n_free < 100: + logger.warning( + f"Only {n_free} free reflections. sigma_A is fitted on the free set; below " + "~100 the per-shell fit degenerates to a single conservative value, which " + "silently removes the model-error term's resolution dependence." + ) + + def prepare( + self, structure: AtomArray | str | Path, *, device: torch.device | str = "cpu" + ) -> None: + """Take the topology and place the reflection data on ``device``. + + Caches the element symbols, their scattering-table codes, the per-atom + annotations the restraint builder needs, and a coordinate snapshot. + + Idempotent: clears the model cache, so it is also how a built reward is moved to + another device. + + Parameters + ---------- + structure + ``AtomArray`` or path to a PDB/mmCIF. Its atom order defines the column + order of the coordinate tensor passed to :meth:`__call__`. + device + Device for the reflection data and every model built from here. + + Raises + ------ + ValueError + If the structure has no atoms. + """ + atom_array = _resolve_structure(structure) + if len(atom_array) == 0: + raise ValueError("Structure has no atoms.") + + # Resolve to a concrete device so the per-call check compares like with like: + # torch.device("cuda") != torch.device("cuda:0") even though a tensor placed on + # the former reports the latter. + self.device = torch.zeros(0, device=torch.device(device)).device + self._element_symbols = [str(e).strip() for e in atom_array.element] + self.n_atoms = len(self._element_symbols) + + codes = elements_to_scattering_indices(self._element_symbols) + self._expected_codes = torch.tensor(codes, dtype=torch.long, device=self.device) + + # Per-atom annotations the restraint builder needs, cached once. Biotite always + # provides these on a parsed structure; `ins_code` may be absent on arrays built + # by hand, so default it to blank rather than failing. + self._topology = { + "atom_name": [str(v).strip() for v in atom_array.atom_name], + "res_name": [str(v).strip() for v in atom_array.res_name], + "res_id": [int(v) for v in atom_array.res_id], + "chain_id": [str(v).strip() for v in atom_array.chain_id], + "ins_code": [ + str(v).strip() for v in getattr(atom_array, "ins_code", [""] * self.n_atoms) + ], + } + # Snapshot coordinates: the restraint build reads pdb[["x","y","z"]] to count + # each parent's heavy neighbours for the riding-hydrogen topology. + self._input_coords = np.ascontiguousarray(np.asarray(atom_array.coord, dtype=float)) + + self._data = self._data.to(self.device) + + # Invalidate anything built for a previous topology or device. + self._stacks = {} + self._calls = {} + self._prepared = True + + # Index 0 is the zero-scattering '?' row; those atoms contribute nothing to F_calc. + n_unknown = int(self._expected_codes.eq(0).sum()) + if n_unknown: + unknown = sorted({s for s, c in zip(self._element_symbols, codes) if c == 0}) + logger.warning( + f"{n_unknown} atoms have no scattering factors (elements {unknown}) and will " + "contribute zero density." + ) + composition = ", ".join( + f"{el}:{n}" for el, n in sorted(Counter(self._element_symbols).items()) + ) + logger.info( + f"Prepared TorchRefXrayRewardFunction: n_atoms={self.n_atoms}, " + f"device={self.device}, target_mode={self.target_mode}, " + f"bulk_solvent={self.bulk_solvent}, b_factor={self.b_factor}, " + f"composition=({composition})" + ) + + # ------------------------------------------------------------------ + # Model / scaler / target construction + # ------------------------------------------------------------------ + def _build_pdb_dataframe(self, n_conformers: int) -> pd.DataFrame: + """Assemble the atom table restraint building reads, conformer-major. + + Carries only the columns torchref reads on these paths: + ``Model._build_restraints`` needs ``name, index, chainid, resseq, resname, icode, + x, y, z, element, ATOM, altloc``; ``Model.Z`` and ``get_vdw_radii`` need + ``element``. + + Each conformer gets its own chain id and altloc letter. Distinct chain ids keep + residue grouping and peptide links within a conformer, since residues are found by + contiguous run over a ``chainid_resseq`` key and links by ``resseq + 1``. The + altloc letter suppresses van der Waals restraints between conformers, which are + otherwise built for every near-coincident duplicated atom pair. + + ``index`` is an explicit column of positional indices, used by the restraint + builders to index the coordinate tensor directly. ``x, y, z`` hold the + prepare-time coordinate snapshot, read by ``build_hydrogen_topology`` to count + each parent's heavy neighbours. + """ + n = self.n_atoms + base = self._topology # per-ASU-atom annotation arrays, from prepare() + # Chain ids: keep the original for a single conformer so the common case reads + # naturally; suffix per conformer only when there is a stack to disambiguate. + chain_ids, altlocs = [], [] + for c in range(n_conformers): + if n_conformers == 1: + chain_ids.extend(base["chain_id"]) + altlocs.extend([""] * n) + else: + tag = _conformer_tag(c) + chain_ids.extend(f"{cid}{tag}" for cid in base["chain_id"]) + altlocs.extend([tag] * n) + + coords = self._input_coords # (n, 3) float, prepare-time snapshot + xyz = np.tile(coords, (n_conformers, 1)) + df = pd.DataFrame( + { + "ATOM": ["ATOM"] * (n * n_conformers), + "name": list(base["atom_name"]) * n_conformers, + "altloc": altlocs, + "resname": list(base["res_name"]) * n_conformers, + "chainid": chain_ids, + "resseq": list(base["res_id"]) * n_conformers, + "icode": list(base["ins_code"]) * n_conformers, + "x": xyz[:, 0], + "y": xyz[:, 1], + "z": xyz[:, 2], + "element": self._element_symbols * n_conformers, + } + ) + # Positional, NOT the pandas index -- see the docstring. + df["index"] = np.arange(len(df), dtype=int) + return df + + def _build_model(self, n_conformers: int) -> Any: + """Build a ``ModelFT`` over a ``n_conformers * n_atoms`` conformer stack. + + Sets up only what the structure-factor and scaling paths read. ``Model.load()`` is + bypassed because its hydrogen stripping and NaN-row dropping would change the atom + count and de-align the caller's coordinate tensor. + + Parameters + ---------- + n_conformers + Conformers sharing the topology. Occupancy defaults to ``1 / n_conformers`` + per atom, which is what makes the stack equal the ensemble mean. + + Returns + ------- + ModelFT + An initialised model whose coordinate and occupancy slots are + :class:`_TensorSlot` holders. + """ + from torchref.symmetry import Cell + + n_total = n_conformers * self.n_atoms + dev = self.device + + # wavelength=None: ModelFT defaults to 1.0, which applies the dispersive f' + # correction on every forward. Nothing here wants that. + model = _external_model_cls()( + verbose=0, wavelength=None, max_res=self.resolution, device=dev + ) + model.cell = Cell(self.unit_cell, dtype=model.dtype_float, device=dev) + # Setter builds the SfFFT submodule once cell and space group are both set. + model.spacegroup = self.space_group + + model.pdb = self._build_pdb_dataframe(n_conformers) + model.initialized = True # gates Z / _build_parametrization + + model.register_buffer("aniso_flag", torch.zeros(n_total, dtype=torch.bool, device=dev)) + model._rebuild_sf_indices() # _iso_indices / _iso_covers_all / _aniso_is_empty + + # torchref's symbol -> Z map differs from the scattering table used above, so an + # ionic form resolved there can still land on Z=0 here. Checked once per model. + n_unknown_z = int(model.Z.eq(0).sum()) + if n_unknown_z: + missing = sorted({s for s, z in zip(model.pdb["element"], model.Z.tolist()) if z == 0}) + logger.warning( + f"{n_unknown_z} atoms have no atomic number in torchref's scattering table " + f"(elements {missing}) and will contribute zero density to F_calc." + ) + + full = functools.partial(torch.full, (n_total,), dtype=model.dtype_float, device=dev) + if self.refine_adp: + # One refinable B per ASU atom, broadcast across the stack. Constructed with + # n_atoms values -- not n_total -- which is what makes the leaf shared. + model.adp = _shared_adp_cls()( + torch.full((self.n_atoms,), self.b_factor, dtype=model.dtype_float), + name="adp", + device=dev, + n_conformers=n_conformers, + ) + else: + model.adp = _TensorSlot(full(self.b_factor)) + model.occupancy = _TensorSlot(full(1.0 / n_conformers)) + model.xyz = _TensorSlot(torch.zeros(n_total, 3, dtype=model.dtype_float, device=dev)) + + # Called explicitly: _late_symmetry_compatible starts as None and is set only in + # setup_grid(), but compute_structure_factors reads it before build_density_map + # lazily calls setup_grid. Without this the first call takes the early-symmetry + # path and later calls take late symmetry -- same answer, ~5x the cost. + model.setup_grid() + return model + + def _stack_for(self, n_conformers: int) -> tuple[Any, Any, Any]: + """Return the cached ``(model, scaler, target)`` for this ensemble size. + + The triple is cached as a unit so each set stays internally consistent; sharing + one scaler across models would require re-pointing ``Scaler.model``, the + ``SolventModel``'s model and the target's ``_model`` submodule, and rebuilding the + solvent mask. With a constant ensemble size only one triple is ever built. + """ + if n_conformers in self._stacks: + return self._stacks[n_conformers] + + from torchref.refinement.targets import create_xray_target + from torchref.scaling import Scaler + + logger.info(f"Building torchref model stack for {n_conformers} conformer(s)") + model = self._build_model(n_conformers) + scaler = Scaler(model, self._data, nbins=self.nbins, verbose=0, device=self.device) + # The model is attached to the target as well; the solvent mask reads it directly. + target = create_xray_target( + data=self._data, + model=model, + scaler=scaler, + mode=self.target_mode, + use_set=self.use_set, + verbose=0, + device=self.device, + ) + self._stacks[n_conformers] = (model, scaler, target) + self._states[n_conformers] = self._build_loss_state(model, target) + return self._stacks[n_conformers] + + def _build_loss_state(self, model: Any, xray_target: Any) -> Any: + """Assemble the weighted ``LossState`` over x-ray and optional restraint groups. + + Follows ``base_refinement._create_loss_state`` so naming, weighting and + maintenance semantics match torchref's. + + Restraint groups are registered via ``register_targets`` without a name: the leaf + targets self-name ``geometry/bond``, ``adp/simu`` and so on, and passing a name + would double-prefix them, which stops ``"geometry/ramachandran": 0.0`` matching + and re-enables Ramachandran at the group weight. + + ``set_weights`` is required. ``LossState.weights`` defaults to ``{}`` with a + per-name lookup default of 1.0, which would run the ADP prior 50x too strong. + Weights compose multiplicatively down the ``/``-separated path. + + Groups are gated on their weight at registration rather than evaluation, because + ``register_target`` probes each target once and that first call is what builds the + restraint graph, riding-hydrogen topology and van der Waals pair list. + """ + from torchref.refinement.loss_state import LossState + from torchref.refinement.targets import TotalADPTarget, TotalGeometryTarget + + state = LossState(device=self.device) + state.register_target("xray", xray_target) + + if self.adp_weight > 0: + state.register_targets(TotalADPTarget(model, verbose=0)) + if self.geometry_weight > 0: + state.register_targets(TotalGeometryTarget(model, verbose=0)) + + state.set_weights( + { + "xray": 1.0, + "adp": self.adp_weight, + "geometry": self.geometry_weight, + # torchref disables Ramachandran by default; keep that. + "geometry/ramachandran": 0.0, + } + ) + + groups = sorted(state.targets) + logger.info( + f"Loss groups: {groups} " + f"(xray=1.0, adp={self.adp_weight}, geometry={self.geometry_weight})" + ) + if self.adp_weight > 0 or self.geometry_weight > 0: + logger.info(f"Restraint counts: {self._restraint_counts(model)}") + return state + + @staticmethod + def _restraint_counts(model: Any) -> dict[str, int]: + """Per-type restraint counts, for logging. + + A zero count means no restraints were built for that type, which happens silently + when a residue is absent from the monomer library. + + ``model.restraints.restraints`` is a ``_RestraintsAccessor`` rather than a dict -- + it has ``get``/``__getitem__``/``__contains__`` but no ``items()`` -- so this uses + the nested-``get`` idiom from ``RestraintsNew.summary``. + """ + r = model.restraints.restraints + # "vdw" and "chiral" are torchref's _FLAT_TYPES: stored as restraints[t]["indices"] + # with no origin level, unlike the nested restraints[t][origin]["indices"]. + flat = {"vdw", "chiral"} + counts: dict[str, int] = {} + for rtype in ("bond", "angle", "torsion", "plane", "chiral", "vdw"): + total = 0 + try: + group = r.get(rtype, {}) + if rtype in flat: + idx = group.get("indices") + total = 0 if idx is None else int(idx.shape[0]) + else: + for origin in group.keys() if hasattr(group, "keys") else (): + idx = group.get(origin, {}).get("indices") + if idx is not None: + total += int(idx.shape[0]) + except Exception: # noqa: BLE001 - logging must not break the build + total = -1 + counts[rtype] = total + return counts + + # ------------------------------------------------------------------ + # Nuisance parameters + # ------------------------------------------------------------------ + def refresh_nuisance_parameters( + self, n_conformers: int = 1, coordinates: torch.Tensor | None = None + ) -> dict[str, float]: + """Refit the scale, bulk solvent and sigma_A, and report R-factors. + + Called automatically on the first call and every ``refresh_interval`` calls after + it; exposed so a caller can drive the cadence itself. + + Coordinates are swapped for a ``.detach().clone()`` for the duration, giving the + refit its own storage and version counter. A bare ``detach()`` would share the + version counter, so an in-place write during the refit would fail autograd's + version check on the caller's later ``backward()``. Cost is an ``(N, 3)`` copy. + + A warm failure is logged and swallowed, keeping the previous scale and sigma_A. A + cold-start failure raises, since there is no previous state and every subsequent + loss would be scored against an unfitted scale. + + Parameters + ---------- + n_conformers + Which cached stack to refresh. + coordinates + Coordinates to fit against, ``(C, n, 3)`` or ``(C*n, 3)``. Defaults to + whatever is currently bound. + + Returns + ------- + dict of str to float + ``{"r_work": ..., "r_free": ...}``, or empty if a warm refit failed. + + Raises + ------ + RuntimeError + If the very first fit fails; see above. + """ + self._require_prepared() + model, scaler, target = self._stack_for(n_conformers) + + if coordinates is not None: + model.set_coordinates(coordinates.reshape(-1, 3)) + live = model.xyz() + + cold = not hasattr(scaler, "log_scale") + model.set_coordinates(live.detach().clone()) + try: + if cold: + # Cold start. Scaler.initialize is the canonical order -- + # calc_initial_scale -> setup_solvent -> setup_anisotropy_correction -- + # but it always sets up solvent, so do the pair by hand when it is off. + if self.bulk_solvent: + scaler.initialize() + else: + scaler.calc_initial_scale() + scaler.setup_anisotropy_correction() + elif self.bulk_solvent: + # Warm: rebuild the mask at the current coordinates. + scaler.solvent.update_solvent() + # forward() only rebuilds the raw solvent structure factors when this + # cache is None, so without it the new mask never reaches F_calc. + scaler._f_sol_raw = None + + scaler.refine_lbfgs(scale_target=self.scale_target) + target.maintenance() # drops the sigma_A cache; nothing else invalidates it + + if self.refine_adp: + self._refine_adp(n_conformers) + + with torch.no_grad(): + r_work, r_free = target.get_rfactor() + logger.info( + f"Refreshed nuisance parameters (C={n_conformers}, " + f"call {self._calls.get(n_conformers, 0)}): " + f"R_work={r_work:.4f}, R_free={r_free:.4f}" + ) + return {"r_work": float(r_work), "r_free": float(r_free)} + except Exception as exc: # noqa: BLE001 - a failed refit must not kill sampling + if cold: + # No previous state to fall back on, so every loss from here would be + # scored against an unfitted scale. + raise RuntimeError( + "The initial nuisance-parameter fit failed, so the reward has no " + "usable scale. Check that the MTZ and the structure describe the " + f"same crystal and that the free set is non-empty. Cause: {exc}" + ) from exc + logger.warning( + f"Nuisance-parameter refresh failed ({type(exc).__name__}: {exc}); " + "keeping the previous scale and sigma_A." + ) + return {} + finally: + # Reattach even if the refit raised; a bound detached clone would return a + # gradient-free loss for every subsequent call. + model.set_coordinates(live) + + def _refine_adp(self, n_conformers: int) -> None: + """LBFGS over the shared B-factors, coordinates frozen. + + Called from :meth:`refresh_nuisance_parameters`, where the caller's coordinates + are already swapped for a detached clone. + + Refines ``("adp",)`` only, unlike torchref's ``refine_adp`` which also passes + ``("u", "occupancy")``: every atom here is isotropic so ``u`` is a zero-element + leaf, and occupancy is caller-owned. Coordinates need no explicit freeze, since + ``LossState.run`` diffs its loss leaves against the optimizer's parameters and + disables the rest for the duration. + """ + model, _scaler, _target = self._stacks[n_conformers] + state = self._states[n_conformers] + + params = model.parameters_of_types(("adp",)) + if not params: + logger.warning("No refinable ADP leaf found; skipping the ADP refit.") + return + + # The k-NN list behind the ADP locality restraint is cached and is NOT covered by + # maintenance(), so it goes stale as the caller moves coordinates. The simu bond + # pairs are topology-only and stay valid. + locality = getattr(state, "targets", {}).get("adp/locality") + if locality is not None: + try: + locality(recompute_neighbors=True) + except TypeError: + pass # older signature without the kwarg; the stale list still works + + before = float(state.aggregate()) + b_before = model.adp().detach() + optimizer = torch.optim.LBFGS( + params, lr=1.0, max_iter=20, history_size=100, line_search_fn="strong_wolfe" + ) + state.step(optimizer, context="torchref_reward.refine_adp") + after = float(state.aggregate()) + + b_now = model.adp().detach() + logger.info( + f"Refined {params[0].numel()} shared B-factors: loss {before:.4g} -> {after:.4g}, " + f"B {b_before.min():.1f}-{b_before.max():.1f} -> " + f"{b_now.min():.1f}-{b_now.max():.1f} A^2" + ) + + # ------------------------------------------------------------------ + # Evaluation + # ------------------------------------------------------------------ + def _require_prepared(self) -> None: + if not self._prepared: + raise RuntimeError( + "TorchRefXrayRewardFunction.prepare() must be called with the model atom " + "array (or `structure=` passed to __init__) before the reward is evaluated." + ) + + def _validate( + self, + coordinates: torch.Tensor, + elements: torch.Tensor, + b_factors: torch.Tensor, + occupancies: torch.Tensor, + ) -> None: + """Check the incoming batch against the cached topology. + + Atom count is checked before element identity so the common case reports a clear + message rather than a shape error. A mismatch in either is raised rather than + adapted to: element identity and atom ordering are fixed for a given experiment. + """ + for name, n in ( + ("coordinates", coordinates.shape[-2]), + ("elements", elements.shape[-1]), + ("b_factors", b_factors.shape[-1]), + ("occupancies", occupancies.shape[-1]), + ): + if n != self.n_atoms: + raise ValueError( + f"{name} has {n} atoms but prepare() cached a topology of " + f"{self.n_atoms}. Call prepare() with the same atom array the " + "sampled coordinates correspond to (model atom space)." + ) + expected = self._expected_codes.expand_as(elements) + if not torch.equal(elements, expected): + raise ValueError( + "elements do not match the topology cached by prepare(). Element identity " + "and atom ordering are fixed for a given experiment, so this means the " + "coordinate columns no longer correspond to the cached atoms." + ) + + 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 X-ray likelihood for an ensemble of conformers. + + Call ``.backward()`` on the result for gradients with respect to ``coordinates`` + and ``occupancies``. + + The batch dimension is the ensemble: ``C`` conformers form a single + ``C * n_atoms`` stack evaluated in one structure-factor calculation. The reshapes + are already in conformer-major order, so no tiling is needed and both tensors stay + in the autograd graph. + + Nuisance parameters are refreshed before the loss is computed, so the returned + value is scored with a scale at most ``refresh_interval`` calls stale, and the + cold-start fit happens before the reward first returns. + + Parameters + ---------- + coordinates + ``[batch, n_atoms, 3]`` Cartesian coordinates (Å) in the crystal frame. Not + SE(3)-invariant: a rigid translation changes every phase, so these must arrive + aligned to the MTZ's frame. + elements + ``[batch, n_atoms]`` scattering-table codes. Checked against the cached + topology rather than used; see :meth:`_validate`. + b_factors + ``[batch, n_atoms]``. Ignored. ADPs come from the model's own slot, either + held at ``b_factor`` or refined during maintenance when ``adp_weight > 0``. + occupancies + ``[batch, n_atoms]``. Used directly and differentiable. The ``1/batch_size`` + weighting makes the stacked structure factor the multi-conformer total. + unique_combinations, inverse_indices + Accepted for protocol compatibility and unused; torchref's kernels do not + vmap. + + Returns + ------- + torch.Tensor + Scalar summed negative log-likelihood, unnormalised and not sign-constrained. + + Raises + ------ + RuntimeError + If :meth:`prepare` has not been called, or if the cold-start nuisance + fit on the first call fails — see + :meth:`refresh_nuisance_parameters`. + ValueError + On an atom-count or element mismatch against the cached topology, or if + the inputs are on a different device than :meth:`prepare` was given. + """ + self._require_prepared() + if coordinates.ndim != 3: + raise ValueError( + f"coordinates must be [batch, n_atoms, 3]; got shape {tuple(coordinates.shape)}." + ) + if coordinates.device != self.device: + raise ValueError( + f"coordinates are on {coordinates.device} but this reward was prepared on " + f"{self.device}. Re-run prepare(device=...) to move it." + ) + self._validate(coordinates, elements, b_factors, occupancies) + + n_conformers = coordinates.shape[0] + model, _scaler, target = self._stack_for(n_conformers) + + live = coordinates.reshape(-1, 3) + model.set_coordinates(live) + model.set_occupancies(occupancies.reshape(-1)) + + # Per-stack count: 0 means this triple has never been fitted, so a new ensemble + # size cold-starts rather than inheriting another stack's refresh position. + if self._calls.get(n_conformers, 0) % self.refresh_interval == 0: + self.refresh_nuisance_parameters(n_conformers) + self._calls[n_conformers] = self._calls.get(n_conformers, 0) + 1 + + # No fcalc argument: torchref computes F_calc from the model, applies the scaler + # and evaluates the likelihood. + # + # aggregate() is the weighted sum over registered groups; with both restraint + # weights at 0 only "xray" is registered and this reduces to target(). + return self._states[n_conformers].aggregate() + + def __repr__(self) -> str: + if not self._prepared: + return f"TorchRefXrayRewardFunction({self.mtzfile!r}, unprepared)" + return ( + f"TorchRefXrayRewardFunction({self.mtzfile!r}, n_atoms={self.n_atoms}, " + f"mode={self.target_mode!r}, device={self.device}, calls={sum(self._calls.values())})" + ) diff --git a/src/sampleworks/utils/imports.py b/src/sampleworks/utils/imports.py index 1417c46e..ad1f2f71 100644 --- a/src/sampleworks/utils/imports.py +++ b/src/sampleworks/utils/imports.py @@ -11,6 +11,7 @@ PROTENIX_AVAILABLE = False RF3_AVAILABLE = False PROTPARDELLE_AVAILABLE = False +TORCHREF_AVAILABLE = False try: from sampleworks.models.boltz.wrapper import Boltz1Wrapper, Boltz2Wrapper @@ -52,6 +53,17 @@ except (ImportError, ModuleNotFoundError, OSError): pass +try: + # torchref backs the reciprocal-space reward in + # sampleworks.core.rewards.torchref_rewards, which imports it lazily so this + # flag is the only thing that has to know whether it is installed. + from torchref import ModelFT, read_mtz + + TORCHREF_AVAILABLE = True + del ModelFT, read_mtz +except (ImportError, ModuleNotFoundError): + pass + F = TypeVar("F", bound=Callable[..., Any]) diff --git a/tests/rewards/test_torchref_rewards.py b/tests/rewards/test_torchref_rewards.py new file mode 100644 index 00000000..84a1156b --- /dev/null +++ b/tests/rewards/test_torchref_rewards.py @@ -0,0 +1,768 @@ +"""Tests for the torchref-backed reciprocal-space reward. + +Covers two areas: that a manually built model agrees with one built through +``Model.load()``, and that gradients reach the caller's tensors, point the right way +against a finite difference, and survive a nuisance-parameter refresh. + +Finite-difference tolerances are loose because ``SfFFT`` samples atoms onto a grid and +truncates each Gaussian, making its gradient approximate by construction (~4e-2 relative +L2 against the analytic ``SfDS`` oracle, cosine 0.999). Direction is asserted tightly, +magnitude loosely. + +The reward is not registered in ``_REWARD_BUNDLES``: it returns an unnormalized, +sign-unconstrained summed likelihood, so the shared ``TestRewardCorrelation`` +absolute-loss bar does not apply. +""" + +from pathlib import Path + +import pytest +import torch +from biotite.structure import AtomArray +from sampleworks.utils.imports import TORCHREF_AVAILABLE + +from tests.rewards.reward_input_helpers import build_reward_input_tensors_without_coords + + +# Every test drives the torchref structure-factor path on the `device` fixture. +pytestmark = [ + pytest.mark.gpu, + pytest.mark.skipif(not TORCHREF_AVAILABLE, reason="torchref is not installed"), +] + +B_FACTOR = 20.0 + +# A deposited structure with its own experimental MTZ, from torchref's test files. +# A PDB rather than a CIF because `save_structure_to_cif` does not preserve the space +# group, which would make the reference model compute F_calc in P1 and void the +# equivalence check in `test_matches_a_normally_loaded_modelft`. 3GR5 is 1329 atoms at +# 2.05 A with conventional column labels (FP/SIGFP/FreeR_flag). +CASE = "3GR5" + + +@pytest.fixture(scope="module") +def torchref_files() -> Path: + """``TorchRef/tests/files``, or skip if the sibling checkout is not present.""" + import torchref + + root = Path(torchref.__file__).resolve().parent.parent / "tests" / "files" + if not (root / "pdb" / f"{CASE}.pdb").exists(): + pytest.skip(f"torchref test files not found at {root}") + return root + + +@pytest.fixture(scope="module") +def pdb_path(torchref_files: Path) -> Path: + return torchref_files / "pdb" / f"{CASE}.pdb" + + +@pytest.fixture(scope="module") +def mtz_path(torchref_files: Path) -> Path: + return torchref_files / "mtz" / f"{CASE}.mtz" + + +@pytest.fixture(scope="module") +def structure(pdb_path: Path) -> AtomArray: + """Topology in the crystal frame, loaded the way the rest of the repo does.""" + from sampleworks.utils.atom_array_utils import load_structure_with_altlocs + + return load_structure_with_altlocs(pdb_path) + + +def make_reward(mtz_path, structure, device, **overrides): + """Build a reward with the standard test config, overridable per test. + + Solvent off and a large ``refresh_interval`` by default, giving the gradient tests a + stationary objective and skipping the expensive mask build. + """ + from sampleworks.core.rewards.torchref_rewards import TorchRefXrayRewardFunction + + # ADPs frozen by default: most of these tests are about the x-ray path, and + # refining B-factors adds a restraint build per instance (~30 s) and makes two + # independently-refit models incomparable. The ADP tests opt in explicitly. + kwargs = dict( + bulk_solvent=False, + b_factor=B_FACTOR, + refresh_interval=10_000, + adp_weight=0.0, + geometry_weight=0.0, + ) + kwargs.update(overrides) + return TorchRefXrayRewardFunction(mtz_path, structure=structure, device=device, **kwargs) + + +@pytest.fixture(scope="module") +def reward(mtz_path: Path, structure: AtomArray, device: torch.device): + """Reward built from the deposited MTZ and the model it was refined against.""" + return make_reward(mtz_path, structure, device) + + +@pytest.fixture(scope="module") +def inputs(structure: AtomArray, device: torch.device): + """``[1, n_atoms, ...]`` reward inputs for the deposited coordinates.""" + elements, b_factors, _occ = build_reward_input_tensors_without_coords(structure, device) + coords = torch.from_numpy(structure.coord).to(device=device, dtype=torch.float32) + return { + "coordinates": coords.unsqueeze(0), + "elements": elements.unsqueeze(0), + "b_factors": b_factors.unsqueeze(0), + "occupancies": torch.ones(1, len(structure), device=device), + } + + +def _call(reward, inputs, coords=None, occupancies=None): + """Evaluate the reward, optionally overriding coordinates or occupancies.""" + kwargs = dict(inputs) + if coords is not None: + kwargs["coordinates"] = coords + if occupancies is not None: + kwargs["occupancies"] = occupancies + return reward(**kwargs) + + +class TestInitialization: + def test_prepared_state(self, reward, structure): + assert reward.n_atoms == len(structure) + assert reward.space_group # resolved from the MTZ, not the structure + assert len(reward.unit_cell) == 6 + + def test_no_element_resolves_to_zero_scattering(self, reward): + """Index 0 is the '?' row and contributes no density; this case should have none.""" + assert not bool(reward._expected_codes.eq(0).any()) + + def test_model_exposes_no_refinable_parameters(self, reward, inputs): + """With ADPs frozen, nothing on the model belongs to an optimizer. + + Coordinates and occupancy are caller-owned, so neither may appear. Once + ``adp_weight > 0`` exactly one leaf appears — the shared B — which + :class:`TestSharedADP` checks. + """ + _call(reward, inputs) + model, _scaler, _target = reward._stack_for(1) + assert list(model.parameters()) == [] + + def test_matches_a_normally_loaded_modelft(self, mtz_path, pdb_path, structure, device, inputs): + """The manual init agrees with a ``load_pdb``-built ModelFT on Fcalc. + + Checks the construction path bypassing ``Model.load()``: the hand-built ``pdb``, + ``aniso_flag``, ``_rebuild_sf_indices``, the element -> Z -> ITC92 chain and the + explicit ``setup_grid``. + + The reference model's coordinates, ADPs and occupancies are copied into our slots + rather than the reverse, because ``Model.load`` runs occupancies through + ``OccupancyTensor``, which collapses residue sharing groups and renormalizes + altloc groups to 1.0. ``max_res`` is matched since it sets the grid spacing. + """ + from sampleworks.core.rewards.torchref_rewards import _TensorSlot + from torchref.model import ModelFT + from torchref.symmetry import Cell + + # Own instance: this test rebinds the model's ADP slot, which would otherwise + # leak into the module-scoped `reward`. + reward = make_reward(mtz_path, structure, device) + _call(reward, inputs) # forces the stack to be built + model, _scaler, _target = reward._stack_for(1) + + ref = ModelFT( + verbose=0, + strip_H=False, + wavelength=None, + device=model.device, + max_res=reward.resolution, + ) + ref.load_pdb(str(pdb_path)) + + assert len(ref.pdb) == reward.n_atoms, ( + f"reference has {len(ref.pdb)} atoms, reward cached {reward.n_atoms}" + ) + # The symmetry has to match or the comparison is meaningless: a reference built + # from a file that lost its space group computes F_calc in P1 against our full + # symmetry, which looks like a large numerical disagreement rather than the + # category error it is. + assert ref.spacegroup.hm == model.spacegroup.hm, ( + f"reference space group {ref.spacegroup.hm!r} != model {model.spacegroup.hm!r}" + ) + # 3GR5's CRYST1 cell and its MTZ cell differ in the third decimal (90.645 vs + # 90.670 A) -- routine between a deposited header and the processing run it came + # from. The reward takes the MTZ's, so put the reference on the same cell; + # otherwise this test measures that discrepancy instead of the construction path. + ref.cell = Cell(model.cell.data.tolist(), dtype=ref.dtype_float, device=ref.device) + ref.setup_grid() + + with torch.no_grad(): + model.set_coordinates(ref.xyz().detach()) + model.set_occupancies(ref.occupancy().detach()) + model.adp = _TensorSlot(ref.adp().detach()) + + hkl = reward._data.hkl_for_sf() + f_manual = model(hkl) + f_ref = ref(hkl, recalc=True) + + # Amplitudes: the target is amplitude-based, so a global phase convention + # difference would be a red herring. + rel = (f_manual.abs() - f_ref.abs()).norm() / f_ref.abs().norm() + assert rel < 1e-4, f"manual init disagrees with load_pdb: relative L2 {rel:.3e}" + + def test_structure_path_and_atom_array_agree( + self, mtz_path, pdb_path, structure, device, inputs + ): + """Constructing from a pdb path and from an AtomArray must score identically.""" + from_path = make_reward(mtz_path, pdb_path, device) + from_array = make_reward(mtz_path, structure, device) + assert from_path.n_atoms == from_array.n_atoms + with torch.no_grad(): + a, b = _call(from_path, inputs).item(), _call(from_array, inputs).item() + # Not bitwise: each instance cold-fits its own scaler, so the two losses differ + # by the difference between two LBFGS runs (measured ~2e-5 relative). + assert abs(a - b) / abs(a) < 1e-3, f"path vs AtomArray: {a:.6g} vs {b:.6g}" + + def test_call_before_prepare_raises(self, mtz_path, inputs): + from sampleworks.core.rewards.torchref_rewards import TorchRefXrayRewardFunction + + rf = TorchRefXrayRewardFunction(mtz_path, bulk_solvent=False) + with pytest.raises(RuntimeError, match="prepare"): + _call(rf, inputs) + + @pytest.mark.parametrize( + "kwargs, match", + [ + ({"target_mode": "nope"}, "target_mode"), + ({"scale_target": "ml"}, "scale_target"), + ({"use_set": "nope"}, "use_set"), + ({"refresh_interval": 0}, "refresh_interval"), + ({"b_factor": 0.0}, "b_factor"), + ], + ) + def test_invalid_config_raises(self, mtz_path, kwargs, match): + """`scale_target="ml"` in particular: an alpha-centred mode is degenerate + with the scale being fitted, so torchref rejects it and so do we.""" + from sampleworks.core.rewards.torchref_rewards import TorchRefXrayRewardFunction + + with pytest.raises(ValueError, match=match): + TorchRefXrayRewardFunction(mtz_path, **kwargs) + + +class TestTopologyChecks: + def test_pipeline_elements_match_the_cached_codes(self, reward, inputs): + """The conservation assumption, checked end-to-end through the production path.""" + expected = reward._expected_codes.expand_as(inputs["elements"]) + assert torch.equal(inputs["elements"], expected) + + def test_changed_elements_raise(self, reward, inputs): + bad = inputs["elements"].clone() + bad[0, 0] = bad[0, 0] + 1 + with pytest.raises(ValueError, match="do not match the topology"): + reward(**{**inputs, "elements": bad}) + + def test_permuted_elements_raise(self, reward, inputs): + """A permutation preserves the count, so only the per-atom check catches it.""" + perm = torch.randperm(reward.n_atoms, device=inputs["elements"].device) + bad = inputs["elements"][:, perm] + if torch.equal(bad, inputs["elements"]): + pytest.skip("random permutation happened to be identity-equivalent") + with pytest.raises(ValueError, match="do not match the topology"): + reward(**{**inputs, "elements": bad}) + + def test_wrong_atom_count_raises(self, reward, inputs): + truncated = {k: v[:, :-1] for k, v in inputs.items()} + with pytest.raises(ValueError, match="atoms but prepare"): + reward(**truncated) + + +class TestGradients: + def test_gradient_reaches_the_callers_coordinates(self, reward, inputs): + coords = inputs["coordinates"].clone().requires_grad_(True) + _call(reward, inputs, coords=coords).backward() + + assert coords.grad is not None, "gradient did not reach the caller's tensor" + assert coords.grad.shape == coords.shape + assert torch.isfinite(coords.grad).all() + assert coords.grad.abs().max() > 0 + + def test_gradient_reaches_the_callers_occupancies(self, reward, inputs): + """Occupancy passes through rather than being frozen, so it carries gradient.""" + occ = inputs["occupancies"].clone().requires_grad_(True) + _call(reward, inputs, occupancies=occ).backward() + + assert occ.grad is not None + assert torch.isfinite(occ.grad).all() + assert occ.grad.abs().max() > 0 + + def test_gradient_matches_directional_finite_difference(self, reward, inputs): + """Directional derivative g·d vs (L(x+hd) - L(x-hd)) / 2h. + + Directional rather than per-component gradcheck because the FFT route's + gradient is approximate by construction; see the module docstring. + """ + x0 = inputs["coordinates"] + coords = x0.clone().requires_grad_(True) + _call(reward, inputs, coords=coords).backward() + grad = coords.grad.clone() + + torch.manual_seed(7) + d = torch.randn_like(x0) + d /= d.norm() + + # eps is set by float32 cancellation, not by curvature. The loss is a sum of order + # 1e5, so one ULP is ~1.3e-2 and fd is quantized in steps of ULP/(2 eps). Relative + # to |g| ~ 5.3 that floor is 6.1% at eps=2e-2 -- above this tolerance, so the FD + # estimate cannot resolve it -- and 0.61% at 2e-1. Measured relative error: 4.6% at + # 2e-2, 0.14% at 2e-1, 1.5% at 5e-1 where curvature starts to show. + eps = 2e-1 + with torch.no_grad(): + lp = _call(reward, inputs, coords=x0 + eps * d).item() + lm = _call(reward, inputs, coords=x0 - eps * d).item() + fd = (lp - lm) / (2 * eps) + analytic = (grad * d).sum().item() + + scale = max(abs(fd), abs(analytic), 1e-8) + assert abs(fd - analytic) / scale < 0.03, ( + f"directional derivative mismatch: analytic={analytic:.6g} fd={fd:.6g}" + ) + + def test_gradient_points_downhill_from_a_perturbed_structure(self, reward, inputs): + """One small step along -grad must reduce the loss.""" + torch.manual_seed(0) + x0 = inputs["coordinates"] + 0.3 * torch.randn_like(inputs["coordinates"]) + coords = x0.clone().requires_grad_(True) + + loss0 = _call(reward, inputs, coords=coords) + loss0.backward() + step = 0.01 / coords.grad.abs().max() + with torch.no_grad(): + loss1 = _call(reward, inputs, coords=x0 - step * coords.grad) + assert loss1.item() < loss0.item() + + def test_descent_reduces_the_loss(self, reward, inputs): + """A short Adam run must make real progress.""" + torch.manual_seed(0) + coords = ( + inputs["coordinates"] + 0.3 * torch.randn_like(inputs["coordinates"]) + ).requires_grad_(True) + opt = torch.optim.Adam([coords], lr=0.02) + + first = None + for _ in range(20): + opt.zero_grad() + loss = _call(reward, inputs, coords=coords) + loss.backward() + opt.step() + if first is None: + first = loss.item() + assert loss.item() < first + + +class TestFrameSensitivity: + """The reward is not SE(3)-invariant and must notice rigid motion. + + Structure factors are more sensitive to this than real-space density: a rigid + translation within the cell changes every phase — so a reward that did *not* + respond to translation would mean coordinates are not reaching the MTZ's frame, + and nothing downstream could work. + """ + + def test_translation_changes_the_loss(self, reward, inputs): + shift = torch.tensor([1.5, 0.0, 0.0], device=inputs["coordinates"].device) + with torch.no_grad(): + base = _call(reward, inputs).item() + moved = _call(reward, inputs, coords=inputs["coordinates"] + shift).item() + assert abs(moved - base) > 1e-3 * abs(base) + + def test_true_structure_beats_a_perturbation(self, reward, inputs): + torch.manual_seed(0) + noise = 0.5 * torch.randn_like(inputs["coordinates"]) + with torch.no_grad(): + truth = _call(reward, inputs).item() + perturbed = _call(reward, inputs, coords=inputs["coordinates"] + noise).item() + assert truth < perturbed + + +class TestConformerStack: + def test_identical_conformers_reproduce_one_copy(self, reward, inputs, device): + """C copies at occupancy 1/C == one copy at occupancy 1. + + The identity the stacked representation rests on: F is linear over atoms, so + an ensemble is one structure-factor calculation over a bigger stack. + + Compared on ``Fcalc``, not on the loss. Each ensemble size gets its own + ``(model, scaler, target)`` triple, so the two losses are scored under two + independently cold-fitted scales and differ by the difference between two LBFGS + fits (~0.2%) — which says nothing about the linearity being tested here. + """ + c = 3 + n = reward.n_atoms + stacked = { + "coordinates": inputs["coordinates"].expand(c, -1, -1).contiguous(), + "elements": inputs["elements"].expand(c, -1).contiguous(), + "b_factors": inputs["b_factors"].expand(c, -1).contiguous(), + "occupancies": torch.full((c, n), 1.0 / c, device=device), + } + # 1. Fcalc: the exact identity. Nothing here is approximate. + with torch.no_grad(): + _call(reward, inputs) + reward(**stacked) + hkl = reward._data.hkl_for_sf() + f_single = reward._stack_for(1)[0](hkl) + f_multi = reward._stack_for(c)[0](hkl) + rel_f = (f_multi - f_single).abs().norm() / f_single.abs().norm() + assert rel_f < 1e-5, f"conformer stack is not exact in Fcalc: rel L2 {rel_f:.3e}" + + # 2. Loss: close but not identical, because the two triples were cold-fitted + # independently and so carry slightly different scales. + cs = inputs["coordinates"].clone().requires_grad_(True) + cm = stacked["coordinates"].clone().requires_grad_(True) + loss_single = _call(reward, inputs, coords=cs) + loss_multi = reward(**{**stacked, "coordinates": cm}) + rel_l = abs(loss_multi.item() - loss_single.item()) / abs(loss_single.item()) + assert rel_l < 5e-3, f"stacked loss differs by {rel_l:.3%} (scaler fits differ)" + + # 3. Gradients: each of the C identical conformers carries 1/C of the density, + # so its gradient is 1/C of the single-copy one and they sum back to it. + loss_single.backward() + loss_multi.backward() + g_single, g_multi = cs.grad[0], cm.grad.sum(0) + cos = torch.nn.functional.cosine_similarity( + g_single.reshape(1, -1), g_multi.reshape(1, -1) + ).item() + ratio = (g_multi.norm() / g_single.norm()).item() + # Close but not bitwise: the two triples were cold-fitted independently, so the + # per-bin scale, aniso tensor and sigma_A differ slightly and reweight the + # residuals. Measured cosine 0.999992 and magnitude ratio within 0.1%; forcing + # both triples onto one scaler makes it exact (cosine 1.000000). + assert cos > 0.9999, f"stacked gradient points elsewhere: cosine {cos:.6f}" + assert abs(ratio - 1.0) < 0.01, f"stacked gradient magnitude ratio {ratio:.4f}" + + def test_a_new_stack_is_cold_fitted(self, reward, inputs, device): + """A newly built ensemble size is fitted on its first use. + + The refresh counter is per-stack for this reason: with a + single global counter and a large ``refresh_interval``, only call 0 ever + refreshes, so any triple built later would run with an *identity* scaler and + silently score against unscaled Fcalc. That showed up as a 14% gradient + disagreement between the single and stacked models before it was fixed. + """ + c = 4 + n = reward.n_atoms + _call(reward, inputs) # burn calls so the global count is far from 0 + _call(reward, inputs) + assert c not in reward._stacks, "pick an ensemble size no other test built" + + reward( + coordinates=inputs["coordinates"].expand(c, -1, -1).contiguous(), + elements=inputs["elements"].expand(c, -1).contiguous(), + b_factors=inputs["b_factors"].expand(c, -1).contiguous(), + occupancies=torch.full((c, n), 1.0 / c, device=device), + ) + _model, scaler, _target = reward._stack_for(c) + assert hasattr(scaler, "log_scale"), "the new stack was never cold-fitted" + assert reward._calls[c] == 1 + + def test_distinct_conformers_differ_from_the_mean_structure(self, reward, inputs, device): + """A real ensemble is a complex sum, not an average of coordinates.""" + torch.manual_seed(0) + c = 2 + n = reward.n_atoms + coords = torch.stack( + [inputs["coordinates"][0] + 0.4 * torch.randn(n, 3, device=device) for _ in range(c)] + ) + stacked = { + "coordinates": coords, + "elements": inputs["elements"].expand(c, -1).contiguous(), + "b_factors": inputs["b_factors"].expand(c, -1).contiguous(), + "occupancies": torch.full((c, n), 1.0 / c, device=device), + } + with torch.no_grad(): + ensemble = reward(**stacked).item() + mean_structure = _call(reward, inputs, coords=coords.mean(0, keepdim=True)).item() + assert ensemble != pytest.approx(mean_structure, rel=1e-6) + + +class TestCacheAndRefresh: + def test_loss_changes_when_only_coordinates_change(self, reward, inputs): + """End-to-end guard on the disabled forward cache. + + torchref computes Fcalc internally via calls that do not pass ``recalc``, so + with the cache enabled every call after the first would score the *first* + call's coordinates — silently. + """ + with torch.no_grad(): + a = _call(reward, inputs).item() + b = _call(reward, inputs, coords=inputs["coordinates"] + 0.5).item() + c = _call(reward, inputs).item() + assert a != pytest.approx(b, rel=1e-9) + assert a == pytest.approx(c, rel=1e-6), "same coordinates should score the same" + + def test_bare_modelft_would_have_been_stale(self, reward, inputs): + """Documents why the subclass exists, and fails if torchref makes it moot.""" + from torchref.model import ModelFT + + _call(reward, inputs) + model, _scaler, _target = reward._stack_for(1) + hkl = reward._data.hkl_for_sf() + + with torch.no_grad(): + f0 = ModelFT.__call__(model, hkl) # the cached path we bypass + model.set_coordinates(model.xyz() + 1.0) + f_stale = ModelFT.__call__(model, hkl) + f_fresh = model(hkl) # our uncached __call__ + assert torch.allclose(f0, f_stale), "expected the bare mixin to serve a stale hit" + assert not torch.allclose(f0, f_fresh) + + def test_refresh_preserves_the_callers_graph(self, reward, inputs): + """The property the detached clone exists to guarantee.""" + coords = inputs["coordinates"].clone().requires_grad_(True) + version_before = coords._version + + reward.refresh_nuisance_parameters(1, coordinates=coords) + + assert coords._version == version_before, ( + "the refresh bumped the caller's version counter -- a bare detach() instead " + "of detach().clone() would do this, and it makes backward() fail later" + ) + loss = _call(reward, inputs, coords=coords) + assert loss.grad_fn is not None + loss.backward() + assert coords.grad is not None and torch.isfinite(coords.grad).all() + + def test_refresh_reports_rfactors(self, reward, inputs): + _call(reward, inputs) + stats = reward.refresh_nuisance_parameters(1) + assert set(stats) == {"r_work", "r_free"} + assert 0.0 < stats["r_work"] < 1.5 + assert 0.0 < stats["r_free"] < 1.5 + + def test_refresh_is_idempotent(self, reward, inputs): + _call(reward, inputs) + first = reward.refresh_nuisance_parameters(1) + second = reward.refresh_nuisance_parameters(1) + assert second["r_work"] == pytest.approx(first["r_work"], abs=0.02) + + def test_cold_start_failure_raises(self, mtz_path, structure, device, inputs, monkeypatch): + """A *first* refit that fails leaves no usable scale, so it must not be swallowed. + + Every later loss would otherwise be scored against an unfitted scale and be + silently meaningless — worse than an exception. + """ + rf = make_reward(mtz_path, structure, device) + _model, scaler, _target = rf._stack_for(1) + + def boom(*args, **kwargs): + raise RuntimeError("simulated cold-start failure") + + monkeypatch.setattr(scaler, "refine_lbfgs", boom) + with pytest.raises(RuntimeError, match="initial nuisance-parameter fit failed"): + _call(rf, inputs) + + def test_refresh_failure_is_swallowed_and_unwinds(self, reward, inputs, monkeypatch): + """A failed scale fit must not kill a run, and must still reattach.""" + coords = inputs["coordinates"].clone().requires_grad_(True) + _call(reward, inputs, coords=coords) + _model, scaler, _target = reward._stack_for(1) + + def boom(*args, **kwargs): + raise RuntimeError("simulated LBFGS failure") + + monkeypatch.setattr(scaler, "refine_lbfgs", boom) + assert reward.refresh_nuisance_parameters(1, coordinates=coords) == {} + monkeypatch.undo() + + # The finally clause must have rebound the live tensor, or every later call + # silently returns a gradient-free loss. + loss = _call(reward, inputs, coords=coords) + assert loss.grad_fn is not None + loss.backward() + assert coords.grad is not None + + +class TestSharedADP: + """Refinable B-factors shared across conformers, regularised by the ADP restraints. + + These opt into ``adp_weight > 0``, which builds the restraint graph (``adp/simu`` + reads the bond list) and makes the B-factors refinable, so they are slower than the + rest of the module. + """ + + @pytest.fixture(scope="class") + def adp_reward(self, mtz_path, structure, device): + return make_reward(mtz_path, structure, device, adp_weight=0.02) + + def test_one_shared_leaf_expanded_across_the_stack(self, adp_reward, inputs, device): + """``(n_asu,)`` refinable log-B, ``forward()`` expanded to ``(C*n_asu,)``.""" + c, n = 3, adp_reward.n_atoms + adp_reward( + coordinates=inputs["coordinates"].expand(c, -1, -1).contiguous(), + elements=inputs["elements"].expand(c, -1).contiguous(), + b_factors=inputs["b_factors"].expand(c, -1).contiguous(), + occupancies=torch.full((c, n), 1.0 / c, device=device), + ) + model, _s, _t = adp_reward._stack_for(c) + + params = model.parameters_of_types(("adp",)) + assert len(params) == 1, "expected exactly the shared adp leaf" + assert params[0].numel() == n, f"leaf is {params[0].numel()}, want n_asu={n}" + + expanded = model.adp() + assert expanded.shape == (c * n,) + # Conformer-major: the ASU block repeats, so every conformer sees the same B. + per_conf = expanded.reshape(c, n) + for b in range(1, c): + assert torch.allclose(per_conf[0], per_conf[b]) + + # And exactly one parameter on the whole model -- coordinates and occupancy stay + # caller-owned, so they must not have become optimizable. + assert len(list(model.parameters())) == 1 + + def test_gradient_sums_onto_the_shared_leaf(self, adp_reward, inputs, device): + """C copies contribute C gradients to one parameter, which is the sharing point.""" + c, n = 3, adp_reward.n_atoms + stacked = dict( + coordinates=inputs["coordinates"].expand(c, -1, -1).contiguous(), + elements=inputs["elements"].expand(c, -1).contiguous(), + b_factors=inputs["b_factors"].expand(c, -1).contiguous(), + occupancies=torch.full((c, n), 1.0 / c, device=device), + ) + adp_reward(**stacked) + model, _s, _t = adp_reward._stack_for(c) + leaf = model.parameters_of_types(("adp",))[0] + + leaf.grad = None + adp_reward(**stacked).backward() + assert leaf.grad is not None + assert leaf.grad.shape == (n,) + assert torch.isfinite(leaf.grad).all() + assert leaf.grad.abs().max() > 0 + + def test_shared_wrapper_is_a_positive_mixed_tensor(self, adp_reward, inputs): + """The shared wrapper keeps torchref's ADP parameter type. + + ``parameters_of_types(("adp",))`` and the ADP restraint targets both expect the + log-space ``PositiveMixedTensor`` form, so the subclass relationship is part of + the contract rather than an implementation detail. + """ + from torchref.model.parameter_wrappers import PositiveMixedTensor + + c, n = 2, adp_reward.n_atoms + adp_reward( + coordinates=inputs["coordinates"].expand(c, -1, -1).contiguous(), + elements=inputs["elements"].expand(c, -1).contiguous(), + b_factors=inputs["b_factors"].expand(c, -1).contiguous(), + occupancies=torch.full((c, n), 0.5, device=inputs["coordinates"].device), + ) + model, _s, _t = adp_reward._stack_for(c) + + assert isinstance(model.adp, PositiveMixedTensor) + assert (model.adp() > 0).all(), "positivity is what the parameterisation buys" + + def test_restraints_built_and_non_empty(self, adp_reward, inputs): + """The restraint graph is populated. + + A residue missing from the monomer library silently gets no restraints, so zero + counts would mean the regularisation is absent while the weights imply otherwise. + """ + _call(adp_reward, inputs) + model, _s, _t = adp_reward._stack_for(1) + counts = adp_reward._restraint_counts(model) + for rtype in ("bond", "angle", "torsion"): + assert counts[rtype] > 0, f"no {rtype} restraints: {counts}" + + def test_restraints_scale_with_conformers_and_never_cross_them( + self, adp_reward, inputs, device + ): + """C conformers give C independent copies, with no bond spanning two of them. + + With a shared chain/resseq the peptide-link builder produces one cross-conformer + bond per junction instead of C proper ones; distinct chain ids and altloc letters + prevent that. + """ + c, n = 2, adp_reward.n_atoms + _call(adp_reward, inputs) + one = adp_reward._restraint_counts(adp_reward._stack_for(1)[0])["bond"] + + adp_reward( + coordinates=inputs["coordinates"].expand(c, -1, -1).contiguous(), + elements=inputs["elements"].expand(c, -1).contiguous(), + b_factors=inputs["b_factors"].expand(c, -1).contiguous(), + occupancies=torch.full((c, n), 1.0 / c, device=device), + ) + model_c = adp_reward._stack_for(c)[0] + many = adp_reward._restraint_counts(model_c)["bond"] + assert many == pytest.approx(c * one, rel=0.02), ( + f"{c} conformers gave {many} bonds, expected ~{c * one} (one set per conformer)" + ) + + # No bond may join two different conformer blocks. + idx = model_c.restraints.restraints["bond"]["all"]["indices"] + block = (idx // n).to(torch.long) + assert int((block[:, 0] != block[:, 1]).sum()) == 0, "cross-conformer bond found" + + def test_refit_moves_b_and_lowers_the_loss(self, adp_reward, inputs): + """The point of the whole exercise: B stops being a flat 20 A^2.""" + _call(adp_reward, inputs) + model, _s, _t = adp_reward._stack_for(1) + b = model.adp().detach() + assert b.std() > 1.0, f"B did not move off uniform: std={b.std():.3f}" + assert b.min() > 0, "B must stay positive (log parameterisation)" + + def test_refit_leaves_the_callers_coordinates_untouched(self, adp_reward, inputs): + """The refit runs on a detached clone, so the caller's tensor must be pristine.""" + coords = inputs["coordinates"].clone().requires_grad_(True) + before, version = coords.detach().clone(), coords._version + adp_reward.refresh_nuisance_parameters(1, coordinates=coords) + assert coords._version == version + assert torch.equal(coords.detach(), before) + loss = _call(adp_reward, inputs, coords=coords) + assert loss.grad_fn is not None + + +class TestRestraintGating: + """A zero weight must mean *never constructed*, not merely skipped at evaluation.""" + + def test_geometry_off_by_default(self, reward, inputs): + _call(reward, inputs) + state = reward._states[1] + assert not any(k.startswith("geometry/") for k in state.targets), ( + f"geometry registered despite weight 0: {sorted(state.targets)}" + ) + + def test_both_weights_zero_builds_no_restraints(self, mtz_path, structure, device, inputs): + """With nothing needing the bond graph, the restraint build does not happen. + + ``register_target`` probes each target on registration, and for these targets that + first call is what builds the restraints. Gating only at evaluation would pay the + whole build and discard it. + """ + rf = make_reward(mtz_path, structure, device, adp_weight=0.0, geometry_weight=0.0) + _call(rf, inputs) + model, _s, _t = rf._stack_for(1) + assert model._restraints is None, "restraints were built though nothing needs them" + assert sorted(rf._states[1].targets) == ["xray"] + + def test_weights_are_applied_with_the_right_names(self, mtz_path, structure, device, inputs): + """Guards the double-prefix trap that silently re-enables Ramachandran. + + Registering the aggregates *with* a name would key them ``geometry/geometry/bond``, + at which point ``"geometry/ramachandran": 0.0`` stops matching and Ramachandran + runs at the group weight instead of being disabled. + """ + rf = make_reward(mtz_path, structure, device, adp_weight=0.02, geometry_weight=0.2) + _call(rf, inputs) + state = rf._states[1] + assert state.get_effective_weight("xray") == pytest.approx(1.0) + assert state.get_effective_weight("adp/simu") == pytest.approx(0.02) + assert state.get_effective_weight("geometry/bond") == pytest.approx(0.2) + assert state.get_effective_weight("geometry/ramachandran") == pytest.approx(0.0) + # And the keys are single-prefixed. + assert "geometry/geometry/bond" not in state.targets + + +class TestBulkSolvent: + @pytest.mark.slow + def test_solvent_changes_the_loss(self, mtz_path, structure, device, inputs): + """Marked slow: this builds a real-space mask and FFTs it.""" + without = make_reward(mtz_path, structure, device, bulk_solvent=False) + with_solvent = make_reward(mtz_path, structure, device, bulk_solvent=True) + with torch.no_grad(): + a = _call(without, inputs).item() + b = _call(with_solvent, inputs).item() + assert a != pytest.approx(b, rel=1e-6) + assert with_solvent._stack_for(1)[1].solvent is not None