Run rigid body against a sandbox so it cannot disturb its caller - #69
Open
kmdalton wants to merge 1 commit into
Open
Run rigid body against a sandbox so it cannot disturb its caller#69kmdalton wants to merge 1 commit into
kmdalton wants to merge 1 commit into
Conversation
`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 HatPdotS#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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Run rigid body against a sandbox so it cannot disturb its caller
RigidBodyRefinementSteprebinds the Refinement to a resolution-truncated dataview at every cutoff.
_rebind_for_dataassignsreflection_data, builds afresh Scaler, and calls
_init_targets+reset_loss_state. Run against thecaller's own Refinement, those assignments are destructive:
_init_targetsreconstructsadp_targetandgeometry_targetfromconstructor defaults, so anything configured on them post-construction is
silently reset.
adp_target['simu'].simu_sigma = 0.25reads back as 2.0.reset_loss_statediscards the LossState, so a weight registered on it isgone. A key present in DEFAULT_GROUP_WEIGHTS is visibly overwritten; a custom
one such as
adp/simusimply returns None afterwards and its target fallsback 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_cutoffdrops every non-xray target from the state beforeoptimizing. 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 itsown 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_xyzswaps itsxyz 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.Modulekeeps submodules in_modules, so the clone copies that dict (and_parameters/_buffers) as well as__dict__; without it a submoduleassignment 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_dictand ensemble paths this change does not touch. Eithercan land without the other; together they cover both the storage and the
rebuild.