Fix ADP restraint config being silently lost on target rebuild

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

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.

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: