From 66f126bf305afc1787e6734dc77ea355884c747b Mon Sep 17 00:00:00 2001 From: Kevin Dalton Date: Fri, 21 Aug 2026 11:24:27 -0400 Subject: [PATCH] Run rigid body against a sandbox so it cannot disturb its caller `RigidBodyRefinementStep` rebinds the Refinement to a resolution-truncated data view at every cutoff. `_rebind_for_data` assigns `reflection_data`, builds a fresh Scaler, and calls `_init_targets` + `reset_loss_state`. Run against the caller's own Refinement, those assignments are destructive: - `_init_targets` reconstructs `adp_target` and `geometry_target` from constructor defaults, so anything configured on them post-construction is silently reset. `adp_target['simu'].simu_sigma = 0.25` reads back as 2.0. - `reset_loss_state` discards the LossState, so a weight registered on it is gone. A key present in DEFAULT_GROUP_WEIGHTS is visibly overwritten; a custom one such as `adp/simu` simply returns None afterwards and its target falls back to the group weight. Neither rebuild is wanted by the step. Only the x-ray target depends on the data and scaler that changed; the ADP and geometry targets are built from the model alone, and `_run_one_cutoff` drops every non-xray target from the state before optimizing. Measured on a 3-cutoff run they are constructed 3 times and evaluated 6 times (registration probe plus loss refresh) purely as overhead, then deleted unused, while the x-ray target takes all 42 gradient evaluations. `run()` now points the step at a shallow clone that shares the model but owns its own attribute namespace, so every one of those assignments lands on the clone. There is nothing to restore afterwards and no window in which the caller's Refinement is inconsistent. The model is deliberately shared rather than copied: `use_rigid_xyz` swaps its xyz container in place, so refined coordinates reach the caller by object identity and no copy-back is needed. That is also what makes the change exactly equivalent rather than approximately so -- on 3E98 the refined coordinates are bit-identical to the previous behaviour, max per-atom difference 0.000e+00. `nn.Module` keeps submodules in `_modules`, so the clone copies that dict (and `_parameters` / `_buffers`) as well as `__dict__`; without it a submodule assignment on the clone would write straight through to the original. Tests: tests/integration/test_rigid_body_isolation.py, five cases -- the sigma survives, a custom LossState weight survives, targets and reflection_data keep their object identity, coordinates still reach the caller, and a normal macrocycle still runs afterwards. Verified that three of them fail when the sandbox is bypassed. Full unit + functional suite passes (1798 passed, 74 skipped). This is independent of #68. That PR gives ADP restraint parameters a constructor-level home so they survive *any* rebuild, including the `create_from_state_dict` and ensemble paths this change does not touch. Either can land without the other; together they cover both the storage and the rebuild. --- .../integration/test_rigid_body_isolation.py | 96 +++++++++++++++++++ torchref/refinement/rigid_body_refinement.py | 53 +++++++++- 2 files changed, 146 insertions(+), 3 deletions(-) create mode 100644 tests/integration/test_rigid_body_isolation.py 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