Fix ADP restraint config being silently lost on target rebuild - #68
Fix ADP restraint config being silently lost on target rebuild#68kmdalton wants to merge 1 commit into
Conversation
`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>
|
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 Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
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. |
|
I moved the weights into the refinement loss_state object, and I think these should survive the rigid-body step. |
|
looking into it... |
|
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. |
`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.
|
closing in favor of #69 |
ADPSimilarityTarget.simu_sigmais public API with a deliberate setter, but there was no way to make a value survive.Refinement._init_targetsbuildsTotalADPTarget(self.model, verbose=self.verbose)passing no restraint parameters, so every rebuild resetssimu_sigma/simu_sigma_anisoto the constructor defaults (2.0 / 1.0).refine_rigid_bodyrebuilds once per resolution cutoff via_rebind_for_data->_init_targets; the ensemble andcreate_from_state_dictpaths 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-maxis 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.CombinedModelTargetstakescomponent_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_targetssplats the per-component kwargs.Refinementtakesadp_restraints=..., stores it alongside the other pre-_init_targetsconfiguration, and passes it on every rebuild with the samegetattrfallback_xray_target_kwargsuses for the ensemble and state-dict paths.Behaviour is unchanged when no config is passed.
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 secondget_scales, andrefine_adp.TotalGeometryTargethas 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-restraintswould make this reachable fromtorchref.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.