From 9318088524c4b3b1e3dc0ba787a52bf321da0a62 Mon Sep 17 00:00:00 2001 From: Kevin Dalton Date: Sat, 15 Aug 2026 22:48:24 -0400 Subject: [PATCH] Fix ADP restraint config being silently lost on target rebuild `ADPSimilarityTarget.simu_sigma` is public API with a deliberate setter, but there was no way to make a value survive. `Refinement._init_targets` builds `TotalADPTarget(self.model, verbose=self.verbose)` passing no restraint parameters, so every rebuild resets `simu_sigma` / `simu_sigma_aniso` to the constructor defaults (2.0 / 1.0). `refine_rigid_body` rebuilds once per resolution cutoff via `_rebind_for_data` -> `_init_targets`; the ensemble and `create_from_state_dict` paths rebuild too. The result: set a sigma, run rigid body, and refinement silently proceeds at the default. No warning. There was no supported alternative -- no constructor argument, no CLI flag (`--sigma-a-max` is sigma_A, a different quantity) -- so post-construction assignment was the only way in, and it was exactly what got discarded. This is the failure mode already documented on `_xray_target_kwargs`: "a second build site silently reverts whatever it forgets to pass, which once made five CLI flags no-ops." The x-ray targets were given a single source of truth for their construction kwargs; the ADP targets never were. This applies the same pattern. - `CombinedModelTargets` takes `component_config`, `{component: {kwarg: value}}`, set before `_create_targets()` and exposed to subclasses via `_component_kwargs()`. A component name that matches nothing raises rather than no-op'ing, since a silent no-op is the bug being fixed. Config is deep copied so a later mutation of the caller's dict cannot reach the target. - `TotalADPTarget._create_targets` splats the per-component kwargs. - `Refinement` takes `adp_restraints=...`, stores it alongside the other pre-`_init_targets` configuration, and passes it on every rebuild with the same `getattr` fallback `_xray_target_kwargs` uses for the ensemble and state-dict paths. Behaviour is unchanged when no config is passed. LBFGSRefinement(..., adp_restraints={"simu": {"simu_sigma": 0.4}}) Verified against the reported scenario: post-construction assignment reads back as 2.0 after `refine_rigid_body`, constructor config holds at 0.4 through rigid body, a second `get_scales`, and `refine_adp`. `TotalGeometryTarget` has the same latent issue -- its components are built with no configuration path either -- but nothing sets geometry component parameters today, so it is left alone. The base-class mechanism is generic, so wiring it up later is two lines. A CLI flag for `--adp-restraints` would make this reachable from `torchref.refine`; deliberately not bundled here. Tests: tests/unit/test_adp_restraint_config.py, 7 cases covering defaults, propagation, survival across a rebuild, copy-not-alias, and both misspelling paths. Full unit + functional suite passes (1748 passed, 74 skipped), as do the 32 integration tests touching rigid body, ensemble, state-dict and CLI paths. Co-Authored-By: Claude Opus 5 --- tests/unit/test_adp_restraint_config.py | 89 +++++++++++++++++++++++++ torchref/refinement/base_refinement.py | 32 ++++++++- torchref/refinement/targets/combined.py | 57 ++++++++++++++-- 3 files changed, 172 insertions(+), 6 deletions(-) create mode 100644 tests/unit/test_adp_restraint_config.py diff --git a/tests/unit/test_adp_restraint_config.py b/tests/unit/test_adp_restraint_config.py new file mode 100644 index 0000000..b793d99 --- /dev/null +++ b/tests/unit/test_adp_restraint_config.py @@ -0,0 +1,89 @@ +"""ADP restraint configuration must survive target rebuilds. + +`Refinement._init_targets` rebuilds `adp_target` from scratch -- once per +resolution cutoff inside `refine_rigid_body`, and again on the ensemble and +`create_from_state_dict` paths. Configuration passed to the constructor has to +be reapplied on every one of those rebuilds; anything assigned to the target +object afterwards is discarded, silently. + +These tests exercise `TotalADPTarget` directly so they need only a Model, no +reflection data or refinement run. +""" + +import pytest + +from torchref.refinement.targets.combined import TotalADPTarget + + +DEFAULT_SIMU_SIGMA = 2.0 +DEFAULT_SIMU_SIGMA_ANISO = 1.0 + + +def _build(model, **kwargs): + return TotalADPTarget(model, verbose=0, **kwargs) + + +def test_defaults_unchanged_without_config(loaded_model): + """No config means exactly the previous behaviour.""" + adp = _build(loaded_model) + assert adp["simu"].simu_sigma == pytest.approx(DEFAULT_SIMU_SIGMA) + assert adp["simu"].simu_sigma_aniso == pytest.approx(DEFAULT_SIMU_SIGMA_ANISO) + assert adp.component_config == {} + + +def test_config_reaches_the_component(loaded_model): + adp = _build(loaded_model, component_config={"simu": {"simu_sigma": 0.4}}) + assert adp["simu"].simu_sigma == pytest.approx(0.4) + # Untouched kwargs keep their defaults. + assert adp["simu"].simu_sigma_aniso == pytest.approx(DEFAULT_SIMU_SIGMA_ANISO) + + +def test_config_survives_a_rebuild(loaded_model): + """The regression: a rebuilt target must come back configured. + + This is what `refine_rigid_body` does per resolution cutoff, via + `_rebind_for_data` -> `_init_targets`. + """ + config = {"simu": {"simu_sigma": 0.4, "simu_sigma_aniso": 0.2}} + first = _build(loaded_model, component_config=config) + rebuilt = _build(loaded_model, component_config=first.component_config) + + assert rebuilt["simu"].simu_sigma == pytest.approx(0.4) + assert rebuilt["simu"].simu_sigma_aniso == pytest.approx(0.2) + + +def test_post_construction_assignment_does_not_survive_a_rebuild(loaded_model): + """Pin the behaviour that motivates the constructor argument. + + Assigning to the target is still legal and still takes effect immediately -- + it just cannot outlive the object. Documenting that here so the next reader + does not "fix" the setter instead of using `adp_restraints`. + """ + first = _build(loaded_model) + first["simu"].simu_sigma = 0.4 + assert first["simu"].simu_sigma == pytest.approx(0.4) + + rebuilt = _build(loaded_model, component_config=first.component_config) + assert rebuilt["simu"].simu_sigma == pytest.approx(DEFAULT_SIMU_SIGMA) + + +def test_config_is_copied_not_aliased(loaded_model): + """Mutating the caller's dict afterwards must not change the target.""" + config = {"simu": {"simu_sigma": 0.4}} + adp = _build(loaded_model, component_config=config) + config["simu"]["simu_sigma"] = 99.0 + assert adp["simu"].simu_sigma == pytest.approx(0.4) + assert adp.component_config["simu"]["simu_sigma"] == pytest.approx(0.4) + + +def test_unknown_component_raises(loaded_model): + """A name that reaches no component is a silent no-op -- the exact failure + mode this machinery exists to prevent. It must raise instead.""" + with pytest.raises(ValueError, match="no such component"): + _build(loaded_model, component_config={"simuu": {"simu_sigma": 0.4}}) + + +def test_unknown_kwarg_raises(loaded_model): + """A misspelled kwarg must not be swallowed either.""" + with pytest.raises(TypeError): + _build(loaded_model, component_config={"simu": {"simu_sgima": 0.4}}) diff --git a/torchref/refinement/base_refinement.py b/torchref/refinement/base_refinement.py index a5c419d..2e3c1e4 100644 --- a/torchref/refinement/base_refinement.py +++ b/torchref/refinement/base_refinement.py @@ -113,6 +113,7 @@ def __init__( shrink: bool = SHRINK_ENABLED, scale_target: str = DEFAULT_SCALE_TARGET, aniso_selection: Optional[str] = None, + adp_restraints: Optional[Dict[str, Dict[str, Any]]] = None, ): """Initialize Refinement, fully if ``data_file`` and ``pdb`` are given. @@ -165,6 +166,19 @@ def __init__( aniso_selection : str, optional Phenix-style selection of atoms refined anisotropically when ``adp_mode="anisotropic"``. Defaults to all non-water heavy atoms. + adp_restraints : dict, optional + Per-component overrides for the ADP restraints, + ``{component: {kwarg: value}}``. Components are ``'simu'``, + ``'locality'`` and ``'sigd'``; an unknown name raises. For example:: + + LBFGSRefinement(..., adp_restraints={"simu": {"simu_sigma": 0.4}}) + + Set restraint configuration **here** rather than assigning to + ``ref.adp_target['simu'].simu_sigma`` after construction. The + targets are rebuilt by :meth:`_init_targets` -- once per resolution + cutoff inside :meth:`refine_rigid_body`, and again on the ensemble + and ``create_from_state_dict`` paths -- and a rebuild resets any + post-construction assignment to the component defaults, silently. """ super().__init__() # Refinement constructs its own submodules from file paths, so @@ -201,6 +215,15 @@ def __init__( self.xray_mode = xray_mode self.sigma_a_max = sigma_a_max self.shrink = shrink + # Same contract for the ADP restraints: their configuration lives HERE, + # on the Refinement, not on the target. `_init_targets` rebuilds + # `adp_target` from scratch -- once per resolution cutoff during + # `refine_rigid_body`, and again on the ensemble and state-dict paths -- + # so anything set on the target object afterwards is silently discarded + # on the next rebuild. Setting it here is what makes it stick. + self.adp_restraints = { + name: dict(kwargs) for name, kwargs in (adp_restraints or {}).items() + } # A wavelength of 0 means "no anomalous refinement": disable the f'/f'' # correction (model wavelength None) and force a Friedel-merged read so # F(+)/F(-) are not loaded as Bijvoet pairs. @@ -424,7 +447,14 @@ def _init_targets(self, xray_mode: str = None): # Geometry targets now accept model directly instead of refinement self.geometry_target = TotalGeometryTarget(self.model, verbose=self.verbose) - self.adp_target = TotalADPTarget(self.model, verbose=self.verbose) + # `getattr` fallback for the same reason as _xray_target_kwargs: the + # ensemble and create_from_state_dict paths build targets before this + # attribute exists. + self.adp_target = TotalADPTarget( + self.model, + verbose=self.verbose, + component_config=getattr(self, "adp_restraints", None), + ) # Initialize scaler scales (overall scale, anisotropic U, bulk solvent) # so the scaler-regularization targets have valid parameters to read. diff --git a/torchref/refinement/targets/combined.py b/torchref/refinement/targets/combined.py index 545201a..4f013b3 100644 --- a/torchref/refinement/targets/combined.py +++ b/torchref/refinement/targets/combined.py @@ -7,7 +7,7 @@ ``LossState`` via :meth:`add_to_state`. """ -from typing import TYPE_CHECKING, Dict +from typing import TYPE_CHECKING, Any, Dict, Optional import torch from torch import nn @@ -128,9 +128,19 @@ class CombinedModelTargets(ModelTarget): Reference to the Model object. verbose : int, optional Verbosity level. Default is 0. + component_config : dict, optional + Per-component constructor overrides, ``{component_name: {kwarg: value}}`` + -- e.g. ``{'simu': {'simu_sigma': 0.4}}``. Applied when the components are + built, so they survive every rebuild of this object. See + :meth:`_component_kwargs`. """ - def __init__(self, model: "Model" = None, verbose: int = 0): + def __init__( + self, + model: "Model" = None, + verbose: int = 0, + component_config: Optional[Dict[str, Dict[str, Any]]] = None, + ): """ Initialize CombinedModelTargets. @@ -140,10 +150,37 @@ def __init__(self, model: "Model" = None, verbose: int = 0): Reference to Model object. verbose : int, optional Verbosity level. Default is 0. + component_config : dict, optional + Per-component constructor overrides, ``{component_name: {kwarg: value}}``. """ super().__init__(model, verbose) + # Set BEFORE _create_targets(), which reads it through _component_kwargs(). + # A plain dict, not a Module attribute: these are construction inputs, not + # state, and must be re-readable every time the components are rebuilt. + self.component_config = { + name: dict(kwargs) for name, kwargs in (component_config or {}).items() + } self._targets = nn.ModuleDict(self._create_targets()) + # A name that never reached a component is a silent no-op -- exactly the + # failure this config exists to prevent -- so refuse it loudly. + unknown = set(self.component_config) - set(self._targets) + if unknown: + raise ValueError( + f"{type(self).__name__}: no such component(s) " + f"{sorted(unknown)}; known components are {sorted(self._targets)}." + ) + + def _component_kwargs(self, name: str) -> Dict[str, Any]: + """Constructor overrides for one component, empty when unconfigured. + + Subclasses must splat this into every component they build. A component + that forgets it silently ignores its configuration, which is the bug this + machinery exists to prevent -- the same one the note on + ``Refinement._xray_target_kwargs`` describes for the x-ray targets. + """ + return dict(self.component_config.get(name, {})) + def _create_targets(self) -> Dict[str, "Target"]: """Build the ``{name: Target}`` components. Subclasses must override.""" raise NotImplementedError("Subclasses must implement _create_targets() method.") @@ -360,17 +397,27 @@ class TotalADPTarget(CombinedModelTargets): Reference to the Model object. verbose : int, optional Verbosity level. Default is 0. + component_config : dict, optional + Per-component constructor overrides, ``{component_name: {kwarg: value}}``. + For example ``{'simu': {'simu_sigma': 0.4, 'simu_sigma_aniso': 0.2}}`` + tightens the ADP similarity restraint. Reaches the components through + :meth:`CombinedModelTargets._component_kwargs`, so it survives every + rebuild of this target. """ def _create_targets(self) -> Dict[str, Target]: """Build the three ADP component targets.""" print("Initializing TotalADPTarget with component targets...") return { - "simu": ADPSimilarityTarget(self.model, verbose=self.verbose), + "simu": ADPSimilarityTarget( + self.model, verbose=self.verbose, **self._component_kwargs("simu") + ), "locality": ADPLocalityTarget( - self.model, verbose=self.verbose + self.model, verbose=self.verbose, **self._component_kwargs("locality") + ), + "sigd": ADPSigdTarget( + self.model, verbose=self.verbose, **self._component_kwargs("sigd") ), - "sigd": ADPSigdTarget(self.model, verbose=self.verbose), } def print_statistics(self) -> None: