diff --git a/tests/integration/test_rigid_body_isolation.py b/tests/integration/test_rigid_body_isolation.py new file mode 100644 index 0000000..651ed11 --- /dev/null +++ b/tests/integration/test_rigid_body_isolation.py @@ -0,0 +1,96 @@ +"""refine_rigid_body must not disturb the refinement it is called on. + +Every cutoff rebinds the step's Refinement to a resolution-truncated data view, +which rebuilds the scaler and every target and drops the loss state. Those +rebuilds are needed for the x-ray target, whose data and scaler genuinely +change; they are collateral for the ADP and geometry targets, which are built +from the model alone and are dropped from the loss state before the optimizer +runs. `RigidBodyRefinementStep` therefore runs against a sandbox clone, and +these tests pin that the caller sees none of it. +""" +import pytest +import torch + +from torchref import LBFGSRefinement + + +@pytest.fixture(scope="module") +def refinement(mtz_dir, pdb_dir): + def build(): + return LBFGSRefinement( + data_file=str(mtz_dir / "3E98.mtz"), + pdb=str(pdb_dir / "3E98.pdb"), + device=torch.device("cpu"), + verbose=0, + ) + return build + + +def test_component_restraint_config_survives(refinement): + """A sigma set on a target must still be set afterwards. + + Regression: `_init_targets` rebuilt `TotalADPTarget` per cutoff with no + restraint parameters, so this silently reverted to the ADPSimilarityTarget + default of 2.0 and refinement continued at a restraint weight nobody chose. + """ + ref = refinement() + ref.adp_target["simu"].simu_sigma = 0.25 + ref.get_scales() + + ref.refine_rigid_body(iterations_per_step=10) + + assert ref.adp_target["simu"].simu_sigma == pytest.approx(0.25) + + +def test_custom_loss_state_weight_survives(refinement): + """A weight registered on the LossState must still be registered afterwards. + + `adp/simu` is deliberately a key absent from DEFAULT_GROUP_WEIGHTS: a key + that is present gets visibly overwritten, while a custom one silently + disappears and its target falls back to the group weight. + """ + ref = refinement() + ref.get_scales() + ref.complete_loss_state().set_weight("adp/simu", 0.77) + + ref.refine_rigid_body(iterations_per_step=10) + + assert ref.complete_loss_state().weights.get("adp/simu") == pytest.approx(0.77) + + +def test_targets_and_data_are_not_replaced(refinement): + """Object identity, not just values -- a caller may hold its own references.""" + ref = refinement() + ref.get_scales() + adp, geometry, data = ref.adp_target, ref.geometry_target, ref.reflection_data + + ref.refine_rigid_body(iterations_per_step=10) + + assert ref.adp_target is adp + assert ref.geometry_target is geometry + assert ref.reflection_data is data + + +def test_refined_coordinates_still_reach_the_caller(refinement): + """The sandbox shares the model, so the whole point still has to work.""" + ref = refinement() + ref.get_scales() + before = ref.model.xyz().detach().clone() + + ref.refine_rigid_body(iterations_per_step=30) + + shift = (ref.model.xyz().detach() - before).norm(dim=-1) + assert float(shift.max()) > 0.0, "rigid body moved nothing" + + +def test_refinement_is_usable_afterwards(refinement): + """A normal macrocycle must still run against the caller's own objects.""" + ref = refinement() + ref.get_scales() + + ref.refine_rigid_body(iterations_per_step=10) + ref.refine_scaler() + ref.refine_adp() + + rwork, rfree = ref.get_rfactor() + assert 0.0 < rwork < 1.0 and 0.0 < rfree < 1.0 diff --git a/torchref/refinement/rigid_body_refinement.py b/torchref/refinement/rigid_body_refinement.py index 70bb8d3..3f320d2 100644 --- a/torchref/refinement/rigid_body_refinement.py +++ b/torchref/refinement/rigid_body_refinement.py @@ -8,6 +8,7 @@ below it switches to ``ml`` with the normal Scaler. """ +import copy from typing import List, Optional import torch @@ -88,14 +89,60 @@ def _xray_mode_for_cutoff(cls, d_min: float) -> str: # ----------------------------------------------------------------------- # Run # ----------------------------------------------------------------------- + @staticmethod + def _sandbox(ref): + """A shallow clone of ``ref`` that shares its model but owns its namespace. + + Every cutoff calls :meth:`_rebind_for_data`, which assigns + ``reflection_data``, builds a fresh ``Scaler``, and calls + ``_init_targets`` + ``reset_loss_state``. Run against the real + Refinement those assignments are destructive: ``_init_targets`` + reconstructs ``adp_target`` and ``geometry_target`` from constructor + defaults, so anything configured on them post-construction (for instance + ``adp_target['simu'].simu_sigma``) is silently reset, and + ``reset_loss_state`` discards custom weights registered on the + ``LossState``. Neither survives the step, and neither is wanted by it -- + rigid body drops every non-xray target before optimizing, so those + objects are rebuilt only to be thrown away. + + Directing the step at a clone confines all of it. The real Refinement is + never written to, so there is nothing to restore and no window in which + it is inconsistent. + + The model is deliberately shared, not copied: ``use_rigid_xyz`` swaps its + xyz container in place, so refined coordinates reach the caller by object + identity and need no copy-back. + + ``nn.Module`` keeps submodules in ``_modules``; copying ``__dict__`` + alone would leave that dict shared, and a submodule assignment on the + clone would write straight through to the original. + """ + sandbox = copy.copy(ref) + sandbox.__dict__ = dict(ref.__dict__) + for slot in ("_modules", "_parameters", "_buffers"): + if slot in sandbox.__dict__: + sandbox.__dict__[slot] = dict(ref.__dict__[slot]) + return sandbox + def run(self): """Step through every cutoff coarse to fine and return ``[(d_min, LossState), ...]``. - Restores the original ``reflection_data`` on exit, and bakes the final transform - back - into a plain ``ModelFT`` unless ``commit=False``. + Runs against a sandbox clone of the refinement (see :meth:`_sandbox`), so + the caller's targets, weights and ``reflection_data`` are left untouched. + Refined coordinates still reach the caller: the model is shared. + + Bakes the final transform back into a plain ``ModelFT`` unless + ``commit=False``. """ + real = self.refinement + self.refinement = self._sandbox(real) + try: + return self._run() + finally: + self.refinement = real + + def _run(self): ref = self.refinement original_data = ref.reflection_data