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: