Run rigid body against a sandbox so it cannot disturb its caller

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

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