Skip to content

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

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

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.

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.

2 participants