ES/03 checkpoint

[samarjeet]

ALCHEMI Toolkit Pull Request

Exact checkpoint / restore: Zarr checkpoint format; BaseDynamics state_dict / load_state_dict / redistribute_state / apply_thermodynamic_state for NVTLangevin and NVTNoseHoover; transactional manifest; round-trip test

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

Changes Made

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR introduces the enhanced-sampling API, built-in umbrella and wall biases, hook-driven biased dynamics, and transactional Zarr checkpoint/restore support. It also adds integrator state serialization and thermodynamic-state rebinding for the NVT integrators.

• Adds conservative and adaptive bias abstractions, collective variables, bias aggregation, and the EnhancedSampling runner.
• Adds checksummed checkpoint manifests covering runner, bias, dynamics, and walker-batch state.
• Adds documentation, examples, exports, and round-trip/lifecycle tests.

Important Files Changed

Filename	Overview
nvalchemi/enhanced_sampling/_runner.py	Implements enhanced-sampling lifecycle orchestration and checkpoint integration; warm_start supplies empty results to result-dependent adaptive updates.
nvalchemi/enhanced_sampling/_checkpoint.py	Implements manifest-last Zarr checkpoint writes, mandatory checksum coverage, validation, and batch reconstruction.
nvalchemi/enhanced_sampling/_bias.py	Defines bias protocols, detached result aggregation, and autograd-derived conservative forces and stress.
nvalchemi/enhanced_sampling/_adaptive.py	Defines adaptive-bias update, epoch, state-version, and cooperative checkpoint contracts.
nvalchemi/dynamics/base.py	Adds serialization, restoration, redistribution, and thermodynamic-state extension points for per-system integrator state.
nvalchemi/dynamics/integrators/nvt_langevin.py	Adds temperature rebinding and queued velocity rescaling for Langevin dynamics.
nvalchemi/dynamics/integrators/nvt_nose_hoover.py	Adds temperature rebinding with chain-mass, chain-velocity, and atomic-velocity transformations.
nvalchemi/enhanced_sampling/biases/umbrella.py	Adds harmonic umbrella bias with per-state parameter selection and stiffness validation.
nvalchemi/enhanced_sampling/biases/walls.py	Adds upper, lower, and flat-bottom restraint potentials.
nvalchemi/enhanced_sampling/cv/pair_distance.py	Adds differentiable pair-distance evaluation with periodic minimum-image handling.
docs/userguide/enhanced_sampling.md	Documents the new bias, runner, lifecycle, checkpoint, and migration workflows.
.claude/skills/nvalchemi-dynamics-hooks/SKILL.md	Updates hook guidance to deprecate BiasedPotentialHook in favor of enhanced sampling.

_{Reviews (1): Last reviewed commit: "Missing component checksums are now mani..." | Re-trigger Greptile}

[greptile-apps]

Warm-start drops bias results

When an adaptive bias uses energy, forces, stress, or observables from its preceding BiasResult, warm_start() passes an empty result for every historical frame instead of evaluating the bias, causing reconstructed history to differ from live sampling and potentially fail on missing result fields.

[dallasfoster]

Writing state transactionally and refusing to restore a torn store has nothing
to do with enhanced sampling. Training already implements it, MD restart wants
it, and NEB restart will want it. Shipping a second implementation under
enhanced_sampling/ is what guarantees a third.

# nvalchemi/_checkpoint.py
class Stateful(Protocol):
    def state_dict(self) -> Mapping[str, Any]: ...
    def load_state_dict(self, state: Mapping[str, Any]) -> None: ...

def save_checkpoint(path, components: Mapping[str, Stateful], *,
                    batch: Batch | None = None, compatibility=None) -> None: ...

def load_checkpoint(path, components: Mapping[str, Stateful], *, device=None) -> None: ...

Hooks, integrators, biases and strategies all satisfy Stateful already, or
would under row 3. If refactoring training/_checkpoint.py onto this is too
much for the series, the minimum ask is that the new module lands at
nvalchemi/_checkpoint.py rather than inside enhanced_sampling/.