Skip to content

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

Open
kmdalton wants to merge 1 commit into
HatPdotS:devfrom
kmdalton:fix/rigid-body-sandbox
Open

Run rigid body against a sandbox so it cannot disturb its caller#69
kmdalton wants to merge 1 commit into
HatPdotS:devfrom
kmdalton:fix/rigid-body-sandbox

Conversation

@kmdalton

Copy link
Copy Markdown
Contributor

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.

`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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant