ES/02 runner - #175
ES/02 runner#175samarjeet wants to merge 23 commits into
Conversation
…riable Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
…tputs Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
…rough for AFTER_STEP captures Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
…state_id Signed-off-by: Samarjeet Prasad <p.samar.j@gmail.com>
Greptile SummaryThe PR introduces an enhanced-sampling subsystem with conservative and adaptive bias abstractions, a dynamics runner, umbrella and wall biases, periodic CV utilities, tests, documentation, and an example.
Important Files Changed
Reviews (1): Last reviewed commit: "Added a compile regression for per-state..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| self.dynamics.compute(batch) | ||
| self._evaluate_and_apply(batch) |
There was a problem hiding this comment.
When default force priming is used with an AFTER_COMPUTE hook such as MaxForceClampHook or NaNDetectorHook, prime_forces() applies the bias directly without dispatching those hooks, causing the first integrator half-step to consume total forces that were neither clamped nor checked for non-finite values.
Knowledge Base Used: Hook System
|
|
||
| self.dynamics = dynamics | ||
| self.biases: dict[str, BiasPotential] = dict(biases or {}) | ||
| self.steps_per_epoch = int(steps_per_epoch) |
There was a problem hiding this comment.
When steps_per_epoch=0, construction accepts the value and the first identity stamp divides by it, causing prime_forces() or run() to fail immediately with integer division by zero.
| self.steps_per_epoch = int(steps_per_epoch) | |
| self.steps_per_epoch = int(steps_per_epoch) | |
| if self.steps_per_epoch <= 0: | |
| raise ValueError("steps_per_epoch must be a positive integer") |
| > stays functional until the enhanced-sampling runner ships, and constructing | ||
| > it emits a `DeprecationWarning`. | ||
|
|
There was a problem hiding this comment.
The updated skill says BiasedPotentialHook remains functional until the enhanced-sampling runner ships, but this PR ships and documents that runner; refresh the statement so agent guidance accurately describes the deprecated hook's support timeline.
Rule Used: Check `.claude/skills/nvalchemi-dynamics-hooks/SKI... (source)
Knowledge Base Used: Hook System
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| class EnhancedSampling: | ||
| """Run biased dynamics on top of an existing ``BaseDynamics``. | ||
|
|
||
| The runner installs one composite hook and otherwise stays out of the | ||
| way: the model, the integrator, the thermostat, and every other hook | ||
| behave exactly as they would unbiased. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| dynamics: | ||
| Any ``BaseDynamics``. Not subclassed, not wrapped — the runner | ||
| registers a hook on it and calls its ``run``. | ||
| biases: | ||
| Mapping of unique name to :class:`BiasPotential`. May be empty, | ||
| which reduces the runner to identity stamping (useful on its own for | ||
| replica exchange in PR 5). | ||
| steps_per_epoch: | ||
| Steps per consistency epoch, the boundary at which | ||
| :meth:`AdaptivePotentialMixin.commit_epoch` fires. | ||
| compile_biases: | ||
| When ``True``, ``torch.compile`` each conservative bias's | ||
| ``energy()``. Not ``evaluate()`` — that path calls | ||
| ``requires_grad_()``, which ``torch.compile`` cannot trace; see the | ||
| :class:`~nvalchemi.enhanced_sampling.ConservativeBias` docstring. | ||
| prime_after_update: | ||
| When ``True`` (default), re-evaluate biases and rewrite the batch's | ||
| total forces after an ``update()`` bumps a bias's state version, so | ||
| that anything reading ``batch.forces`` between steps sees the current | ||
| bias rather than the previous one. |
There was a problem hiding this comment.
BaseDynamics already owns the stepping loop, so a second run() means every
future workflow has to decide which loop it belongs to, and the two will drift.
The runner's remaining jobs are smaller than they look. It stamps identity
fields onto the batch, fires things on a cadence, and writes a checkpoint. The
toolkit already has mechanisms for the first two — register_bookkeeping_key
(the same mechanism status and system_id use, so the fields survive refill
and graduation) and Hook.frequency. The third is row 8. Note that #176 has to
add state_dict to BaseDynamics anyway for the runner to work, which is
already most of what a strategy would need.
# nvalchemi/dynamics/strategy.py — mirrors TrainingStrategy
class DynamicsStrategy(BaseModel):
"""Declarative recipe that builds and runs a configured BaseDynamics."""
engine: type[BaseDynamics] = NVTLangevin
engine_kwargs: dict[str, Any] = {}
n_steps: int | None = None
extra_hooks: list[Hook] = [] # runtime, not serialized
def build_hooks(self) -> list[Hook]: ...
def build(self, model) -> BaseDynamics: ...
def run(self, batch, model, n_steps=None) -> Batch: ...
def to_spec_dict(self) -> dict[str, Any]: ...
# nvalchemi/enhanced_sampling/strategy.py (~80 lines, was 1414)
class EnhancedSampling(DynamicsStrategy):
biases: dict[str, BiasPotential] = {}
steps_per_epoch: int = 10_000
def build_hooks(self) -> list[Hook]:
return [WalkerIdentityHook(),
*self.biases.values(),
EpochCommitHook(frequency=self.steps_per_epoch)]
BaseDynamics.register_bookkeeping_key(
"walker_id", lambda n, dev: torch.arange(n, device=dev).reshape(n, 1)
)TrainingStrategy is the precedent worth copying: it is a Pydantic model that
configures a loop it does not own, and it serializes to a spec. Under this shape
NEB, Relax and EOSScan become sibling subclasses rather than sibling
runners.
2 participants
ALCHEMI Toolkit Pull Request
EnhancedSampling runner (no exchange): walker identity stamping, force-step ordering, update() exactly-once, evaluate-only force priming, warm_start(); AdaptivePotentialMixin; HarmonicUmbrellaBias; LowerWall / UpperWall / FlatBottomRestraint; periodic_difference; unit tests; gallery example; user guide
Type of Change
Related Issues
Changes Made
Testing
make pytest)make lint)Checklist
Additional Notes
Tip
This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.