Skip to content

Fix ADP restraint config being silently lost on target rebuild - #68

Closed
kmdalton wants to merge 1 commit into
HatPdotS:devfrom
kmdalton:fix/adp-restraint-config-lost-on-rebuild
Closed

Fix ADP restraint config being silently lost on target rebuild#68
kmdalton wants to merge 1 commit into
HatPdotS:devfrom
kmdalton:fix/adp-restraint-config-lost-on-rebuild

Conversation

@kmdalton

Copy link
Copy Markdown
Contributor

ADPSimilarityTarget.simu_sigma is public API with a deliberate setter, but there was no way to make a value survive. Refinement._init_targets builds TotalADPTarget(self.model, verbose=self.verbose) passing no restraint parameters, so every rebuild resets simu_sigma / simu_sigma_aniso to the constructor defaults (2.0 / 1.0). refine_rigid_body rebuilds once per resolution cutoff via _rebind_for_data -> _init_targets; the ensemble and create_from_state_dict paths rebuild too.

The result: set a sigma, run rigid body, and refinement silently proceeds at the default. No warning. There was no supported alternative -- no constructor argument, no CLI flag (--sigma-a-max is sigma_A, a different quantity) -- so post-construction assignment was the only way in, and it was exactly what got discarded.

This is the failure mode already documented on _xray_target_kwargs: "a second build site silently reverts whatever it forgets to pass, which once made five CLI flags no-ops." The x-ray targets were given a single source of truth for their construction kwargs; the ADP targets never were. This applies the same pattern.

  • CombinedModelTargets takes component_config, {component: {kwarg: value}}, set before _create_targets() and exposed to subclasses via _component_kwargs(). A component name that matches nothing raises rather than no-op'ing, since a silent no-op is the bug being fixed. Config is deep copied so a later mutation of the caller's dict cannot reach the target.
  • TotalADPTarget._create_targets splats the per-component kwargs.
  • Refinement takes adp_restraints=..., stores it alongside the other pre-_init_targets configuration, and passes it on every rebuild with the same getattr fallback _xray_target_kwargs uses for the ensemble and state-dict paths.

Behaviour is unchanged when no config is passed.

LBFGSRefinement(..., adp_restraints={"simu": {"simu_sigma": 0.4}})

Verified against the reported scenario: post-construction assignment reads back as 2.0 after refine_rigid_body, constructor config holds at 0.4 through rigid body, a second get_scales, and refine_adp.

TotalGeometryTarget has the same latent issue -- its components are built with no configuration path either -- but nothing sets geometry component parameters today, so it is left alone. The base-class mechanism is generic, so wiring it up later is two lines. A CLI flag for --adp-restraints would make this reachable from torchref.refine; deliberately not bundled here.

Tests: tests/unit/test_adp_restraint_config.py, 7 cases covering defaults, propagation, survival across a rebuild, copy-not-alias, and both misspelling paths. Full unit + functional suite passes (1748 passed, 74 skipped), as do the 32 integration tests touching rigid body, ensemble, state-dict and CLI paths.

`ADPSimilarityTarget.simu_sigma` is public API with a deliberate setter, but
there was no way to make a value survive. `Refinement._init_targets` builds
`TotalADPTarget(self.model, verbose=self.verbose)` passing no restraint
parameters, so every rebuild resets `simu_sigma` / `simu_sigma_aniso` to the
constructor defaults (2.0 / 1.0). `refine_rigid_body` rebuilds once per
resolution cutoff via `_rebind_for_data` -> `_init_targets`; the ensemble and
`create_from_state_dict` paths rebuild too.

The result: set a sigma, run rigid body, and refinement silently proceeds at
the default. No warning. There was no supported alternative -- no constructor
argument, no CLI flag (`--sigma-a-max` is sigma_A, a different quantity) -- so
post-construction assignment was the only way in, and it was exactly what got
discarded.

This is the failure mode already documented on `_xray_target_kwargs`: "a second
build site silently reverts whatever it forgets to pass, which once made five
CLI flags no-ops." The x-ray targets were given a single source of truth for
their construction kwargs; the ADP targets never were. This applies the same
pattern.

- `CombinedModelTargets` takes `component_config`, `{component: {kwarg: value}}`,
  set before `_create_targets()` and exposed to subclasses via
  `_component_kwargs()`. A component name that matches nothing raises rather
  than no-op'ing, since a silent no-op is the bug being fixed. Config is deep
  copied so a later mutation of the caller's dict cannot reach the target.
- `TotalADPTarget._create_targets` splats the per-component kwargs.
- `Refinement` takes `adp_restraints=...`, stores it alongside the other
  pre-`_init_targets` configuration, and passes it on every rebuild with the
  same `getattr` fallback `_xray_target_kwargs` uses for the ensemble and
  state-dict paths.

Behaviour is unchanged when no config is passed.

    LBFGSRefinement(..., adp_restraints={"simu": {"simu_sigma": 0.4}})

Verified against the reported scenario: post-construction assignment reads back
as 2.0 after `refine_rigid_body`, constructor config holds at 0.4 through rigid
body, a second `get_scales`, and `refine_adp`.

`TotalGeometryTarget` has the same latent issue -- its components are built
with no configuration path either -- but nothing sets geometry component
parameters today, so it is left alone. The base-class mechanism is generic, so
wiring it up later is two lines. A CLI flag for `--adp-restraints` would make
this reachable from `torchref.refine`; deliberately not bundled here.

Tests: tests/unit/test_adp_restraint_config.py, 7 cases covering defaults,
propagation, survival across a rebuild, copy-not-alias, and both misspelling
paths. Full unit + functional suite passes (1748 passed, 74 skipped), as do the
32 integration tests touching rigid body, ensemble, state-dict and CLI paths.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@kmdalton
kmdalton changed the base branch from main to dev August 20, 2026 17:10
@kmdalton

Copy link
Copy Markdown
Contributor Author

hi @HatPdotS , i ran into an issue with the ADP target weights being reset to a default value between macrocycles. Claude proposes this fix. let me know if it sounds right to you.

@codecov-commenter

codecov-commenter commented Aug 20, 2026

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 76.92308% with 3 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
torchref/refinement/targets/combined.py 62.50% 3 Missing ⚠️

📢 Thoughts on this report? Let us know!

@kmdalton

Copy link
Copy Markdown
Contributor Author

one more detail -- i was using a custom refinement loop which runs rigid body during every macrocycle. if i understand correctly, that was the root cause. i can dig in more if needed. just wanted a second set of eyes on it.

@HatPdotS

Copy link
Copy Markdown
Owner

I moved the weights into the refinement loss_state object, and I think these should survive the rigid-body step.
I have been moving away from the combined targets architecture because it quickly becomes very confusing. Can you check if custom loss state weights are lost after the rigid step?

@kmdalton

Copy link
Copy Markdown
Contributor Author

looking into it...

@kmdalton

Copy link
Copy Markdown
Contributor Author

i believe custom loss state is also lost. i guess the issue is that the rigid body implementation sweeps the resolution cutoff and switches the x-ray target depending on the cutoff. i'm a little iffy on the mechanism, but this re-initializes some hyperparameters. i have a more conservative / defensive patch on #69 which runs rigid body on a copy so it doesn't clobber any loss state.

HatPdotS pushed a commit that referenced this pull request Aug 25, 2026
`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.
@kmdalton

Copy link
Copy Markdown
Contributor Author

closing in favor of #69

@kmdalton kmdalton closed this Aug 25, 2026
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.

3 participants