diff --git a/.claude/skills/nvalchemi-dynamics-hooks/SKILL.md b/.claude/skills/nvalchemi-dynamics-hooks/SKILL.md index a9a6c738..cda96b11 100644 --- a/.claude/skills/nvalchemi-dynamics-hooks/SKILL.md +++ b/.claude/skills/nvalchemi-dynamics-hooks/SKILL.md @@ -164,6 +164,14 @@ MaxForceClampHook( **BiasedPotentialHook** — add an external bias potential for enhanced sampling. +> **Deprecated.** Use `nvalchemi.enhanced_sampling` (`ConservativeBias`) for new +> biases: it derives forces and stress from one energy definition by autograd. +> A `bias_fn` bias contributes no stress, so it is invisible to an NPT/NPH +> barostat — safe under NVE/NVT, silently wrong under a barostat. Run new +> biases through `EnhancedSampling`, which covers everything this hook does. +> The hook stays functional and no removal date is set; constructing it emits +> a `DeprecationWarning`. + ```python def my_bias(batch: Batch) -> tuple[torch.Tensor, torch.Tensor]: """Return (bias_energy [B, 1], bias_forces [V, 3]).""" diff --git a/CHANGELOG.md b/CHANGELOG.md index bb251a50..59ec210f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,212 @@ ### Added +- `EnhancedSampling` runner for biased dynamics, plus the first built-in + biases. The runner installs one internal hook on an existing `BaseDynamics` + and owns what a bias cannot: walker identity stamping (`walker_id`, + `thermodynamic_state_id`, `sampling_step`, `exchange_segment`, + `sampling_epoch`), the force-step ordering, exactly-once `update()` + delivery, and force priming. Every bias is evaluated against the same + unmodified model output and the contributions summed once, so no bias can + observe another's forces and the total is independent of registration + order. Built-ins: `HarmonicUmbrellaBias` (per-window centers and stiffness + selected by `thermodynamic_state_id`, validated symmetric + positive-semidefinite), `UpperWall`, `LowerWall`, and + `FlatBottomRestraint`. `AdaptivePotentialMixin` supplies the + `update`/`commit_epoch`/state-version battery for biases whose state + evolves during sampling; it must precede `nn.Module` in the base list, and + raises `TypeError` otherwise rather than letting `nn.Module.state_dict` + shadow it and drop bias history from checkpoints. `periodic_difference` + wraps CV differences onto a circle. `warm_start()` gives approximate + continuation from prior frames; for exact resumption see the checkpoint + entry below. + +- Synchronous replica exchange. `ReplicaExchange` and `ThermodynamicState` + advance a ladder of states as one batch and periodically swap which walker + holds which rung, with an even/odd pair schedule so every pair in a segment + is disjoint and decidable simultaneously. Exchange permutes + `thermodynamic_state_id`; coordinates never move between rows. The + acceptance rule is inferred from the ladder rather than declared — varying + temperatures select the Metropolis temperature rule, equal ones the + umbrella rule — because a declared rule that disagreed with the ladder + would break detailed balance silently. An accepted swap is indivisible: + label, integrator target temperature, velocity rescaling, and forces move + together, and an integrator that cannot rebind (`NVE`) is rejected at + construction rather than sampling the state it just left. Acceptance draws + derive from `random_seed + exchange_id`, so a restored run reproduces the + same decisions and the checkpoint stores two integers rather than an RNG + blob; exchange state lives under `sampling/exchange/`. The manifest records + the ladder (mode, acceptance rule, interval, temperatures) and `restore()` + refuses a mismatch, including exchange-versus-none in either direction — + the ladder decides what a swap means, so restoring into a different one + would keep the assignment and counters while silently changing the + acceptance exponent. Per-pair acceptance + rates are reported for ladder tuning. Asynchronous exchange and force-only + (ABF-style) biases are rejected explicitly, as is the unimplemented + combined temperature-plus-window rule: a bias declaring + `state_dependent_for_exchange` is refused at construction, and the runner + additionally probes every bias empirically at prime time by evaluating it + under a permuted assignment, which catches a user bias that declares + nothing. A single-window `HarmonicUmbrellaBias` now ignores + `thermodynamic_state_id` rather than indexing it, so one shared restraint + can run alongside a multi-rung temperature ladder. + +- Metadynamics, in two flavours. `WellTemperedMetaDynamicsBias` deposits + Gaussian hills along any differentiable CV, one per walker per deposition, + with the well-tempered height damping + `h_t = h_0 exp(-V(s_t) / (k_B T (gamma - 1)))` that makes the sum converge; + `bias_factor=None` gives standard metadynamics. `free_energy()` returns + `-(gamma / (gamma - 1)) V(s)`. Three storage policies, chosen rather than + defaulted: `preallocated` keeps tensor shapes fixed for the whole run and + **raises** when capacity is exhausted, because silently dropping hills would + change the physics of a converging run with nothing to show for it; `grow` + allocates another chunk and recompiles; `fifo` bounds memory by discarding + the oldest hill, which is a scientific choice and not a cache policy — the + well-tempered convergence argument no longer applies, so `free_energy()` + refuses under it rather than returning a plausible number. Three history + modes: `shared` (the multiple-walker scheme — `B` walkers fill a basin + roughly `B` times faster in one batched force evaluation), `walker` (`B` + independent runs in one batch), and `state` (per-rung history for a + replica-exchange ladder, which declares `state_dependent_for_exchange`). + The latter two require `walker_id` / `thermodynamic_state_id` on the batch + and raise if it is missing or the wrong length, rather than defaulting to a + single owner — that fallback filed every hill under one key and produced + energies numerically identical to `history="shared"`, silently delivering + the opposite of what was configured. The runner stamps both fields on every + step, so only a directly evaluated bias has to supply them. + `periods` wraps the hill difference onto a circle so a hill near a branch + cut repels from both sides. `sigma` and `periods` are checked against the + CV on first evaluation rather than broadcast against it: a mismatched + length would widen the hill table and silently change the Gaussian + exponent. `sigma` may be a scalar shared across components or one entry + per component; `periods` must be per-component, since a `0` entry is what + marks a component non-periodic and one value cannot carry that + distinction. The hill table takes its width from the CV, so a scalar + `sigma` works with a multi-component CV. + + `RMSDMetaDynamicsBias` is the xTB/CREST-style variant, whose history is a + set of retained structures rather than CV values, and which therefore needs + no collective variable at all. Optimal translation/rotation alignment is + solved by the quaternion characteristic-polynomial route rather than an SVD + Kabsch: the proper-rotation constraint is built in instead of needing a + non-differentiable `det` correction, and only the largest eigenvalue is + taken, which stays well-conditioned for symmetric-top and linear molecules + where singular-vector gradients blow up. The squared RMSD is used + throughout — `sqrt` has infinite derivative at zero, and a reference is + visited at RMSD zero every time one is deposited. Consequences: the energy + is invariant to rigid motion and the bias forces sum to exactly zero. + Non-periodic systems only — a periodic batch is rejected, because an atom + crossing a cell face is physically unmoved but Cartesian-displaced by a + lattice vector, which would inject a large spurious force. Periodicity is + read from `batch.pbc` rather than from the presence of a cell, matching + `pair_distance` and the rest of the toolkit: a bounding-box cell with `pbc` + all-False is the non-periodic case this bias is for and is accepted, a slab + is rejected, and a cell carrying no `pbc` flags is refused as undeclared. + Atom correspondence is fixed, `atom_indices` selects a per-graph + subset, warm-start references seed the history, and there is deliberately no + `free_energy()` — it is a structure generator, not an estimator. + + Both deposit at `AFTER_STEP`, so a hill marks the configuration the walker + reached; both bump the state version so the runner re-primes forces and the + new hill is felt on the next step rather than one late; and neither deposits + during `prime_forces()`. + +- `AdaptiveBiasingForce` — measures the mean force along a pair distance and + applies its negative, so a well-sampled bin leaves no residual force along + the CV and the walker diffuses across it. What it accumulates already *is* + the free-energy gradient, so `free_energy()` integrates it directly with no + hills to deconvolve and no histogram to reweight. + + The estimator is + `dA/dr = <-(F_j - F_i).u / 2 - 2 kB T / r>`, and the second term — the + metric correction — is not a refinement. A naive Cartesian projection gives + the mean force in the *constrained* ensemble; the unconstrained free energy + differs by the Jacobian of the coordinate change, which for a distance in + three dimensions is `-2 kB T / r`. Omitting it produces a smoothly wrong + answer rather than noise: two non-interacting particles come out with a flat + PMF when the true one is the purely entropic `-2 kB T ln r`. This is why the + class takes an atom **pair** rather than a general `cv` callable, unlike + every other bias in the subpackage — the correction belongs to this + coordinate, and accepting an arbitrary CV would apply a distance-shaped + correction to something that is not a distance. + + At or below `min_samples` a bin applies nothing, and between `min_samples` + and `full_samples` the applied fraction ramps linearly to one, so no bin + switches on with a jump. `full_samples == min_samples` asks for a step + instead — nothing until the threshold, full force at it — which is the + classic hard-threshold form and what `min_samples=0` gives by default, so it + is handled as its own case rather than falling through the linear formula + and arriving a sample late; `max_force` optionally caps what a bin visited once + at an awkward geometry can do. `observation_stage` is `AFTER_COMPUTE`, where + `batch.forces` still holds the unbiased physical force — an estimator shown + its own output converges to whatever it had already decided. An update + landing in a bin still below its threshold does not bump the state version, + since the applied force has not changed. + + `evaluate()` returns forces with `energy=None`: there is no potential to + report, which is what non-conservative means here. `supplies_exchange_energy` + is `False`, so `ReplicaExchange` refuses the combination at construction + rather than dropping the bias from the acceptance exponent. `mean_force()` + and `free_energy()` report unvisited bins as `nan` rather than zero, and + `free_energy()` raises on an interior gap because integrating across a hole + would leave every value beyond it wrong by an unknown constant. Per-step + diagnostics are `cv`, `bin`, `applied_gradient`, `samples`, `ramp`, and + `in_range`; a walker outside `cv_range` reports `bin = -1` and zero for + every per-bin quantity, so the diagnostics agree with the zero force it + receives rather than reporting the nearest edge bin's statistics. + +- Adaptive biases now record a configuration fingerprint in their + `state_dict()` and reject a mismatched restore. The manifest records each + bias's class, but a class name says nothing about the settings its saved + state depends on: an ABF histogram restored under a different `cv_range` is + shape-compatible and silently relabels every bin, `bias_factor` changes the + ratio `free_energy()` applies to hills deposited under a different gamma, + and `k_push`/`alpha` change what stored RMSD references do. Worse for + settings held as buffers — `sigma`, `atom_indices` — + `nn.Module.load_state_dict` *overwrote* the caller's value with the + checkpoint's rather than leaving it unvalidated. `AdaptivePotentialMixin` + gains `config_fingerprint()` (empty by default, so other biases are + unaffected) and checks it before delegating, so the guard covers a bias + restored directly as well as through the runner. Capacity is deliberately + excluded, since `storage="grow"` legitimately reaches a size the + constructor never had. + +- `pair_displacement` — the vector form of `pair_distance`, exposed for + methods that work with the CV gradient rather than its value. + `pair_distance` is now its norm, so the two cannot drift apart in their + validation, device handling, or minimum-image convention. + +- Exact checkpoint and restore for enhanced sampling. + `EnhancedSampling.checkpoint()` writes a transactional Zarr store that + extends the existing `AtomicData` layout with a `sampling/` group holding + integrator, bias, and runner state; `restore()` reads it back and returns a + force-primed batch that reproduces the identical trajectory. The manifest is + written last and is the commit marker, so an interrupted write has none and + is refused rather than half-restored. Integrity cover is total: each + `sampling/` component is SHA-256 checksummed, and a separate + `batch_checksum` covers `meta/`, `core/`, and `custom/` — the positions, + velocities, pointer arrays, and walker identity that `AtomicDataZarrWriter` + writes outside the component path. All are verified on read, and cover is + mandatory: a manifest with a gap — a declared component lacking a checksum, + a checksum naming no component, or no batch checksum — is rejected as + invalid rather than read unverified. State is Zarr + arrays and JSON attributes with **no pickle payloads** — an unsupported + value type raises rather than falling back, so loading a checkpoint cannot + execute code. Checkpoints are permitted + only at a consistency-epoch boundary, the one point with no pending + `update()` or in-flight epoch commit; the error names the next valid step. + `checkpoint()` also drains the completed epoch's `commit_epoch()` before + collecting state, since that normally fires lazily on the next step — so a + shared-history bias is saved merged rather than mid-merge. The drain is + tracked per epoch index and cannot double-count. + `BaseDynamics` gains `state_dict()`, `load_state_dict()`, + `redistribute_state()`, and `apply_thermodynamic_state()`, the last + implemented for `NVTLangevin` (velocity rescaling) and `NVTNoseHoover` + (chain masses and velocities transformed with kT, leaving the chain kinetic + energy invariant). Model weights are never restored from a checkpoint; the + manifest records model, dynamics, and bias classes and `restore()` refuses a + mismatch. + - Domain decomposition for distributed inference and dynamics: a spatial halo strategy and a graph-parallel strategy, both driven by a declarative `MLIPSpec` a model wrapper publishes as `distribution_spec`. Ewald, PME, @@ -93,6 +299,14 @@ ### Fixed +- **Zero-dimensional tensors in enhanced-sampling checkpoints** — Zarr reads a + 0-d array back as shape `(1,)`, so a component holding a scalar buffer (a + step counter, a deposition count — the kind of state a compile-safe bias + keeps as a tensor rather than a Python int) no longer matched the digest + taken when it was written, and `restore()` failed the component's own + checksum. The true shape is now recorded alongside the dtype and reapplied + on decode; checkpoints written before this are read exactly as before. + - **Ewald charge gradients and cell derivatives** — the reciprocal term was only ever differentiated with respect to positions and charges, so a non-hybrid Ewald returned a wrong `dE/dq`, and strain-autograd through the detached @@ -136,6 +350,20 @@ ### Deprecated +- `BiasedPotentialHook`, superseded by the `nvalchemi.enhanced_sampling` + subpackage. Its `bias_fn(batch) -> (energy, forces)` contract has no slot + for a cell response, so a bias applied through the hook contributes no + stress and is invisible to the NPT/NPH barostat — the cell evolves as if + the bias were absent, with no error raised. It also requires bias forces + to be written by hand (nothing checks they are `-dE/dr`), and composes + several biases by sequential in-place mutation of `batch.forces` rather + than summing them against the unmodified model output. Constructing the + hook now emits a `DeprecationWarning`. It remains functional so existing + code keeps working, and no removal date is set; `EnhancedSampling` (also in + this release) covers everything it does. No adapter is provided: + bridging a `BiasPotential` onto `bias_fn` would have to discard + `BiasResult.stress`, reintroducing the exact failure the new API removes. + - `cells_inv` argument on `_cell_kinetic_energy`. Cell kinetic energy is computed directly from the strain rate `ε̇` and no longer needs the cell inverse. The argument is retained for backwards diff --git a/docs/userguide/enhanced_sampling.md b/docs/userguide/enhanced_sampling.md new file mode 100644 index 00000000..081a3042 --- /dev/null +++ b/docs/userguide/enhanced_sampling.md @@ -0,0 +1,930 @@ +# Enhanced Sampling + +Molecular dynamics follows the natural motion of atoms, which means it spends +almost all of its time in free-energy minima. Barrier crossings — diffusion +events, reactions, nucleation, conformational change — are rare on MD +timescales. Enhanced sampling adds a bias potential that pushes the system +into regions it would not visit on its own, so that a fixed budget of model +evaluations buys more of the physics you actually care about. + +`nvalchemi.enhanced_sampling` provides the bias abstractions, a set of +built-in biases, and the `EnhancedSampling` runner that wires them into an +existing dynamics object. + +```{contents} +:local: +:depth: 2 +``` + +## Quick start + +```python +import torch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.enhanced_sampling import ( + EnhancedSampling, HarmonicUmbrellaBias, pair_distance, +) + +device = "cuda" if torch.cuda.is_available() else "cpu" + +pair = torch.tensor([0, 5], device=device) +umbrella = HarmonicUmbrellaBias( + cv=lambda batch: pair_distance(batch, pair), + centers=torch.tensor([[2.0], [2.5], [3.0]]), # three windows + stiffness=10.0, # eV/A^2 + name="umbrella", +) + +dynamics = NVTLangevin(model=model, dt=0.5, temperature=300.0, friction=0.05) +sampling = EnhancedSampling(dynamics, {"umbrella": umbrella}) + +# One window per graph. batch already carries forces/energy buffers — see +# "Batch requirements" below. +batch["thermodynamic_state_id"] = torch.tensor([0, 1, 2], device=device) +batch = sampling.run(batch, n_steps=10_000) +``` + +Every window is a row of one batch, so all three are advanced by a single +batched force evaluation per step rather than three separate simulations. + +## Collective variables + +A CV is **any differentiable callable** `cv(batch) -> Tensor[B, D]`. There is +no base class, no registration, and nothing to subclass: + +```python +pair = torch.tensor([0, 5], device=device) + +def bond(batch): + return pair_distance(batch, pair) +``` + +`atom_indices` is moved to the batch's device for you, so a CV closure built +before the batch reaches the GPU still works — but hoisting the tensor out of +the closure and placing it explicitly avoids reallocating it on every call. +The same applies to a bias: `ConservativeBias` moves its buffers to the +batch's device on first evaluation, so `HarmonicUmbrellaBias(...)` built on +CPU evaluates correctly against a CUDA batch without an explicit `.to()`. + +`pair_distance` is the built-in geometric CV. It handles non-periodic systems +and the minimum-image convention for **Minkowski-reduced** cells. + +:::{warning} +`pair_distance` is not a general triclinic MIC. The 27-image search it uses is +correct only for reduced cells; it raises `ValueError` in eager mode when the +cell violates the reduction condition. Under `torch.compile` that check is +skipped and supplying a reduced cell is the caller's responsibility — pre-reduce +with `atoms.get_cell().niggli_reduce()` or equivalent. +::: + +For a CV that lives on a circle (a dihedral, say), use `periodic_difference` +so a restraint at `+3.0 rad` does not pull a configuration at `-3.0 rad` the +long way round: + +```python +from nvalchemi.enhanced_sampling import periodic_difference + +delta = periodic_difference(values, centers, periods=torch.tensor([2 * math.pi])) +``` + +## Built-in biases + +| Bias | Energy | Use for | +|------|--------|---------| +| `HarmonicUmbrellaBias` | `0.5 * delta^T K delta` | Umbrella sampling, restrained MD | +| `UpperWall` | `(k/p) * max(s - s0, 0)^p` | Stop a CV rising past a bound | +| `LowerWall` | `(k/p) * max(s0 - s, 0)^p` | Stop a CV falling below a bound | +| `FlatBottomRestraint` | both of the above | Confine a CV to an interval | +| `WellTemperedMetaDynamicsBias` | `sum_i h_i exp(-(s - c_i)^2 / 2 sigma^2)` | Free energy along a CV you chose | +| `RMSDMetaDynamicsBias` | `sum_r k_push exp(-alpha RMSD(x, x_r)^2)` | Structure search with no CV at all | +| `AdaptiveBiasingForce` | *(no energy — force only)* | Free-energy profile along a pair distance | + +The first four are static: their energy depends only on the current +configuration. The last three are **history-dependent** — they accumulate +state as sampling proceeds, and so mix in `AdaptivePotentialMixin`. +`AdaptiveBiasingForce` is further the only one that is **non-conservative**: +it applies a measured force that is not the gradient of anything it holds. + +Walls contribute **exactly zero** energy and force inside the allowed region, +and their default quadratic form has zero force at the boundary, so switching +one on does not deliver an impulse. + +### Per-window parameters + +`HarmonicUmbrellaBias` accepts `centers` of shape `[D]` (shared) or `[S, D]` +(one row per thermodynamic state). Each graph selects its row via +`batch.thermodynamic_state_id`; without that field every graph uses state `0`. +`stiffness` accepts a scalar, `[D]`, `[D, D]`, or `[S, D, D]`, and is validated +to be symmetric positive-semidefinite — a negative eigenvalue would turn the +restraint into an unbounded repulsion. + +## Metadynamics + +Umbrella sampling needs you to name the windows in advance. Metadynamics does +not: it deposits a Gaussian hill wherever the system currently is, so the +accumulated bias pushes it towards wherever it has not yet been. + +### Well-tempered metadynamics + +`WellTemperedMetaDynamicsBias` deposits one hill per walker at its current CV +value, every `update_frequency` steps. In the well-tempered scheme each hill +is damped by the bias already standing at that point, + +```text +h_t = h_0 * exp(-V(s_t) / (k_B T (gamma - 1))) +``` + +so the sum converges rather than filling forever, and the converged bias is a +free-energy estimate: + +```python +from nvalchemi.enhanced_sampling import WellTemperedMetaDynamicsBias + +metad = WellTemperedMetaDynamicsBias( + cv=bond_distance, + height=0.005, # h_0, eV + sigma=0.25, # hill width, CV units + temperature=300.0, + bias_factor=8.0, # gamma; None gives standard metadynamics + update_frequency=500, + storage="preallocated", + max_hills=2000, + name="metad", +) +... +profile = metad.free_energy(grid) # -(gamma / (gamma - 1)) * V(s) +``` + +`bias_factor=None` is the `gamma -> infinity` limit: every hill keeps height +`h_0` and `F(s) = -V(s)`. That is standard metadynamics, which does not +converge — the bias keeps growing and oscillates about the true profile. + +Pass `periods` for a CV that lives on a circle. A hill at `+3.10 rad` must +repel a configuration at `-3.10 rad`, which is `0.083 rad` away the short way +and `6.20 rad` the long way; without a period the bias sees the second number +and does nothing. + +### Storage policies + +The hill table has to be bounded somehow, and the three ways of bounding it +are not interchangeable. + +| `storage` | On reaching `max_hills` | Compile | `free_energy()` | +|-----------|-------------------------|---------|-----------------| +| `preallocated` | **raises** | Shapes fixed for the whole run | Valid | +| `grow` | Allocates another chunk | Recompiles on each resize | Valid | +| `fifo` | Overwrites the oldest hill | Shapes fixed | **Raises** | + +`preallocated` raises rather than evicting because silently dropping hills +would change the physics of a converging run with nothing to show for it. The +error names the ways out. + +`grow` carries a limit that is easy to miss. Under torch's default of static +parameter shapes, each resize changes the hill-tensor shape and a compiled +`energy()` retraces — and Dynamo caps retraces per code object at +`torch._dynamo.config.recompile_limit`, 8 by default. The next growth +therefore **hard-fails mid-run**, once the trajectory is already underway. + +Three ways out, in rough order of preference: + +1. Use `preallocated`, which holds a single trace for the whole run. +2. Size `max_hills` so the number of growths stays under the limit. +3. Set `force_parameter_static_shapes = False` on `torch._dynamo.config`, + which makes Dynamo trace the hill-table dimension symbolically so growth + stops triggering a retrace at all. + +Both settings are process-global, and other code in this toolkit changes them: +`DistributedModel` raises the limit to 64 *and* disables static parameter +shapes, so a domain-decomposed run will not see this failure. Read the live +config rather than assuming the defaults. + +`fifo` is not merely a cache policy: once hills are +discarded the accumulated bias is no longer the integral of everything +deposited, the well-tempered convergence argument no longer applies, and +`free_energy()` refuses rather than returning a number that looks fine. + +### Multi-walker history + +`history` decides which hills a given walker feels. + +- `"shared"` (default) — every walker feels every hill. This is the + multiple-walker scheme: `B` walkers fill a basin roughly `B` times faster, + at the cost of one batched force evaluation per step rather than `B` + separate runs. +- `"walker"` — each walker feels only its own hills, so one batch runs `B` + genuinely independent metadynamics simulations. +- `"state"` — hills belong to the `thermodynamic_state_id` that deposited + them, which is what a replica-exchange ladder needs. + +`"state"` sets `state_dependent_for_exchange = True`, so combining it with a +temperature ladder is rejected rather than run under an acceptance rule that +does not cover it; see [Acceptance](#acceptance). + +`"state"` and `"walker"` read `thermodynamic_state_id` and `walker_id` off the +batch, and **raise** if the field is absent or the wrong length. The runner +stamps both on every step, so this only affects a bias you evaluate directly. +The alternative — defaulting a missing field to zero — would file every hill +under one key and silently collapse the per-owner histories into a single +shared one, which is the opposite of what was asked for and produces energies +numerically identical to `history="shared"`. + +### xTB-style RMSD metadynamics + +`RMSDMetaDynamicsBias` drops the collective variable entirely. Its history is +a set of retained *structures*, and it pushes away from all of them at once: + +```text +V(x, t) = sum_r f_r(t) * k_push * exp(-alpha * RMSD(x, x_r)^2) +``` + +RMSD is measured after optimal translation and rotation, so the bias is +invariant to rigid motion and its forces sum to exactly zero. This is the +scheme xTB/CREST uses for conformer and isomer searching, and it is the right +tool when you cannot say in advance which coordinate matters. + +```python +from nvalchemi.enhanced_sampling import RMSDMetaDynamicsBias + +explorer = RMSDMetaDynamicsBias( + k_push=0.08, # eV + alpha=10.0, # A^-2 + update_frequency=500, + max_references=64, # FIFO by default + atom_indices=torch.tensor([0, 4, 7]), # heavy atoms only + name="explorer", +) +``` + +Choosing `alpha` is the main decision, and it must match the RMSD scale the +system actually explores: the kernel only has usable gradient where +`alpha * RMSD^2` is of order one. Set it far too small and +`exp(-alpha * RMSD^2)` sits at ~1 for every structure, leaving the bias nearly +constant and nearly forceless. A rigid cluster moving 0.3 A wants `alpha` +around 10; a floppy molecule sampling 1 A wants `alpha` around 1. + +Three constraints are worth knowing before reaching for it: + +- **Non-periodic systems only.** Cartesian RMSD against a stored reference is + not defined under periodic boundary conditions — an atom crossing a cell face + is physically unmoved but Cartesian-displaced by a lattice vector, which would + inject a large spurious force. Bias a periodic-aware CV with + `WellTemperedMetaDynamicsBias` instead. + + Periodicity is read from `batch.pbc`, not from the presence of a cell. A cell + is a box; only the flags say whether atoms wrap through its faces. So a + molecular batch carrying a **bounding box** with `pbc` all-False is accepted + — the common case for a boxed or solvated molecule — while a slab + (`pbc=[True, True, False]`) is rejected, since wrapping along any one axis is + enough to break the metric. A non-zero cell with no `pbc` flags at all is + refused as undeclared rather than assumed harmless. +- **Fixed atom correspondence.** Atom `i` is always compared with atom `i` of + the reference; there is no permutation search, so two structures identical + up to relabelling of equivalent atoms count as distinct. +- **No free energy.** This is a structure generator, not an estimator, and + there is deliberately no `free_energy()` method. What it produces is a set + of structures worth optimising or re-scoring at a higher level of theory. + +`atom_indices` are **per-graph local** indices. Restricting to heavy atoms is +the usual choice: methyl hydrogens spinning freely generate RMSD that says +nothing about the conformer. + +### Deposition timing + +Both biases deposit at `AFTER_STEP`, so a hill marks the configuration the +walker actually reached rather than the one it started from. A deposition +bumps the bias state version, and the runner re-primes forces in response, so +a new hill is felt on the very next step rather than one step late. + +Neither deposits during `prime_forces()`: priming evaluates forces, it does +not advance the trajectory, and depositing there would double-count the +starting configuration. + +## Adaptive biasing force + +Umbrella sampling restrains, metadynamics fills. ABF does neither: it +*measures* the mean force along the CV in each bin and applies its negative, +so once a bin is well sampled the residual force along the CV averages to +zero and the walker diffuses across the coordinate. + +The accumulated quantity already **is** the free-energy gradient, so there is +nothing to deconvolve or reweight at the end — `free_energy()` integrates it. + +```python +from nvalchemi.enhanced_sampling import AdaptiveBiasingForce + +abf = AdaptiveBiasingForce( + atom_indices=torch.tensor([0, 1]), + temperature=300.0, # must match the thermostat + cv_range=(2.0, 6.0), # angstrom + n_bins=40, + min_samples=200, + full_samples=400, + name="abf", +) +... +profile = abf.free_energy() # PMF at bin_centers, in eV +``` + +### The metric correction + +For a pair distance the estimator is + +```text +dA/dr = < -(F_j - F_i).u / 2 - 2 kB T / r > +``` + +That second term is not a refinement. Projecting Cartesian forces onto the CV +gradient and averaging gives the mean force in the *constrained* ensemble; the +free energy of the unconstrained one differs by the Jacobian of the coordinate +change. For a distance in three dimensions the number of configurations at +separation `r` grows as the sphere surface `4 pi r^2`, contributing +`-2 kB T / r`. + +Omitting it does not add noise — it produces a smoothly wrong answer. Two +non-interacting particles have zero Cartesian force, so a naive projection +reports a **flat** PMF when the true one is `-2 kB T ln r`, a purely entropic +profile that drives the pair apart. + +This is why `AdaptiveBiasingForce` takes an atom **pair** rather than a +general `cv` callable, unlike every other bias here. The correction is +specific to this coordinate, and accepting an arbitrary CV would mean applying +a distance-shaped correction to something that is not a distance. Other CVs +need their own correction and are P1. + +### Sample threshold and ramp + +A mean force estimated from a handful of samples is noise, and applying it +would drive the walker on the strength of that noise. At or below +`min_samples` a bin applies nothing; between `min_samples` and `full_samples` +the applied fraction ramps linearly to one, so no bin switches on with a jump. + +Setting `full_samples == min_samples` asks for a step instead — nothing until +the threshold, full force at it. That is the classic hard-threshold form, and +it is what `min_samples=0` gives by default, so it is supported rather than +rejected. A bin with no samples reports a ramp of zero whatever the thresholds +are, since it has no estimate to apply. + +`max_force` optionally caps `|dA/dr|`, bounding what a bin visited once at an +awkward geometry can do while its average settles. + +### Observation ordering + +`observation_stage` is `AFTER_COMPUTE`, where `batch.forces` still holds the +**unbiased** physical force. This is load-bearing: an estimator shown its own +output converges to whatever it had already decided, and the resulting profile +looks perfectly smooth. The runner captures the frame before applying any bias +contribution, so this holds even with several biases registered. + +For the same reason, an `update()` that only lands in a bin still below its +threshold does **not** bump the state version — the applied force has not +changed, so re-priming would be pure cost. + +### No energy, and therefore no replica exchange + +`evaluate()` returns `forces` with `energy=None`. There is genuinely no +potential to report, which is what makes ABF non-conservative. The Metropolis +acceptance rule needs each bias's energy evaluated under both states being +swapped, so `supplies_exchange_energy` is `False` and `ReplicaExchange` +refuses the combination at construction rather than dropping the bias from the +exponent and breaking detailed balance silently. + +### Reading the profile + +`mean_force()` returns the per-bin estimate and `free_energy()` its integral. +The per-step diagnostic is named `bias//applied_gradient` rather than +`mean_force`, because it is the ramped and capped value actually used — a +threshold-suppressed zero there is not a measured zero mean force. + +The per-step observables are `cv`, `bin`, `applied_gradient`, `samples`, +`ramp`, and `in_range`. A walker outside `cv_range` reports `bin = -1` and +zero for every per-bin quantity, matching the zero force it receives. Only +`cv` is still reported, since the coordinate is genuinely measured wherever +the walker is. +Bins never visited come back as `nan` rather than zero — a bin with no samples +has no estimate, and zero is a perfectly plausible value that would hide that. +`free_energy()` **raises** on an interior gap: integration carries the profile +across a hole, so every value beyond it would be wrong by an unknown constant. + +## Writing your own bias + +### The boundary: `BiasPotential` + +`BiasPotential` is a `@runtime_checkable` Protocol that **inherits nothing**. +A bias needs a `name` and an `evaluate`: + +```python +class MyBias: + name = "my_bias" + + def evaluate(self, current): + return BiasResult(energy=..., forces=...) +``` + +That class satisfies the protocol with no base class at all. + +#### Why `BiasResult` is narrower than `ModelOutputs` + +`ModelOutputs` is an open mapping; `BiasResult` carries a fixed set of physics +fields. The split is between what the runner **applies** and what it merely +**reports**. + +An applied output needs a destination (`batch.energy`, `batch.forces`, +`batch.stress`), a per-graph or per-atom reshape rule, a rule for combining it +across biases, and a conversion — stress and virial are the same physics in two +conventions, and moving between them needs the cell volume. An unrecognised key +has none of that, and `_check_destinations` already raises for any produced +output with no buffer to receive it. So an open payload would be open only up +to the first key the runner could not apply. + +`observables` is the open half: an arbitrary `Mapping[str, Tensor]`, no shape +checks, surfaced as `bias//`. A per-atom energy decomposition +belongs there — it is a diagnostic, not a contribution to `batch.energy`. + +The cost is real: a method producing a genuinely new *applied* output cannot +express it without editing the framework. That is accepted, because the +alternative is a contribution the runner silently drops. + +### The batteries: mixins + +Capability is opt-in per bias, supplied as composable mixins: + +```python +class MyRestraint(ConservativeBias): ... # energy -> forces + stress +class MyMetaD(AdaptivePotentialMixin, ConservativeBias): ... # ...and evolving state +class MyABF(AdaptivePotentialMixin): ... # adaptive, no energy +``` + +:::{important} +`AdaptivePotentialMixin` must come **first** in the base list. `ConservativeBias` +inherits `nn.Module`, whose `state_dict` would otherwise shadow the mixin's and +silently drop the bias history from every checkpoint. Getting the order wrong +raises `TypeError` at class-creation time. +::: + +### `ConservativeBias` + +Override `energy()` and get forces and stress by autograd: + +```python +class MyRestraint(ConservativeBias): + def __init__(self, k): + super().__init__(name="my_restraint") # required + self.k = k + + def energy(self, current): + return 0.5 * self.k * my_cv(current) ** 2 # [B, 1] +``` + +Notes: + +- **Stress, not virial.** `ConservativeBias` emits tensile-positive Cauchy + stress, matching every model wrapper in the toolkit, so bias output sums + directly with model output. `BiasResult.virial` exists for hand-written + biases that produce a virial directly, but the runner will reject it — + convert with `sigma = -W/V` first. +- **Partial dependence is fine.** An energy that depends only on the cell (a + volume restraint) yields zero forces and real stress; one that returns a + constant on some branch yields zeros for both. Neither is an error. +- **`torch.compile` boundary is `energy()`, not `evaluate()`.** + `evaluate()` calls `requires_grad_()`, which `torch.compile` cannot trace. + `EnhancedSampling(compile_biases=True)` compiles each bias's `energy()`. + +:::{important} +Because `compile_biases=True` hands `energy()` to `torch.compile`, **keep +data-dependent Python branches out of it**. A `bool(tensor.any())` there — +a bounds check, a "did anything violate this" guard — breaks `fullgraph=True` +outright with a "Could not guard on data-dependent expression" error. + +Put such validation in an override of `evaluate()` instead, which is eager by +construction, then call `super().evaluate(current)`. `HarmonicUmbrellaBias` +validates `thermodynamic_state_id` this way. That placement is strictly better +than an eager-only `torch.compiler.is_compiling()` guard: the check still runs +when `energy()` is compiled, rather than being skipped exactly when a mistake +is hardest to diagnose. +::: + +### Adaptive biases + +`AdaptivePotentialMixin` separates read-only evaluation from state mutation: + +```python +class MyMetaD(AdaptivePotentialMixin, ConservativeBias): + update_frequency = 100 + observation_stage = DynamicsStage.AFTER_STEP + + def energy(self, current): ... # read-only, compile-friendly + + def update(self, frames, result): # called once per due step + self.deposit_hill(frames) + self.bump_state_version() # tells the runner forces are stale +``` + +`observation_stage` decides which frame `update` receives: + +- `AFTER_STEP` — post-step coordinates. What metadynamics wants. +- `AFTER_COMPUTE` — captured while `batch.forces` still holds the **unbiased** + physical forces. What ABF requires; an estimator fed its own output diverges. + +## The runner + +`EnhancedSampling` installs one internal hook on the dynamics and otherwise +leaves it alone — the model, integrator, thermostat, and every other hook +behave exactly as they would unbiased. + +### What it guarantees + +1. **Every bias sees the same unmodified physical output.** Contributions are + summed once and applied together, so no bias can observe another's forces + and the total does not depend on registration order. +2. **`update()` is delivered exactly once per due step**, after integration. +3. **Observables are namespaced** `bias//`, so two biases of the + same type cannot collide. +4. **Forces are primed** before the first step. A velocity-Verlet-style + integrator reads `batch.forces` in its first half-step, before any model + call; without priming, step 0 would be the one step that ignores the bias. + +The bias hook is inserted at the **front** of the hook list, so a safety hook +such as `MaxForceClampHook` clamps the *total* force rather than the model +force alone. + +### Diagnostics + +```python +sampling.last_outputs["physical/forces"] # model only, before any bias +sampling.last_outputs["bias/umbrella/energy"] # one bias's contribution +sampling.last_outputs["bias_total/forces"] # sum across all biases +sampling.last_outputs["total/forces"] # physical + bias +``` + +`total/*` is read back from the batch after the bias is applied, so +`total == physical + bias_total`. Note that this is the state as the *runner* +leaves it, not necessarily what the integrator consumed: the runner's hook runs +first at `AFTER_COMPUTE` (so a force clamp acts on the total rather than the +model force alone), which means a later hook can still modify `batch.forces`. +Read the batch directly if you need the exact value the integrator used. + +For WHAM or MBAR you want `physical/energy` and the per-bias energies +separately, not `total/energy` — the reweighting needs the unbiased potential. +Free-energy reconstruction is deliberately not built in; use `pymbar` or an +equivalent. + +### Batch requirements + +Dynamics writes model outputs back **in place**, so the buffers must exist: + +```python +AtomicData( + positions=..., atomic_numbers=..., atomic_masses=..., + forces=torch.zeros(n_atoms, 3), + energy=torch.zeros(1, 1), + stress=torch.zeros(1, 3, 3), # required whenever a bias produces stress +) +``` + +The runner raises a named `ValueError` naming the field, the biases that +produced it, and how to allocate the buffer — rather than skipping the field +and letting the contribution vanish. Because `run()` primes before the first +step, this surfaces at setup, not part-way through a trajectory. + +:::{warning} +`stress` is the one that matters most. `ConservativeBias` produces stress +whenever the batch has a cell and at least one periodic dimension, so a +periodic run needs the buffer even under NVT. A stress contribution dropped on +the floor is invisible to an NPT/NPH barostat — the cell evolves as if the +bias were absent, with nothing to indicate it. If a run genuinely has no use +for a cell response, pass `compute_stress=False` to the bias, which drops +`"stress"` from its `active_outputs` and skips the strain leaf entirely. That +is a deliberate choice; a missing buffer is not. +::: + +### Walker identity + +The runner stamps five graph-level fields each step. Batch *position* is not +an identity — selection and refill can move a walker to a different row — so +anything that must follow a physical configuration is carried as data: + +| Field | Meaning | +|-------|---------| +| `walker_id` | Immutable identity, assigned once | +| `thermodynamic_state_id` | Window / temperature / energy-function state | +| `sampling_step` | Dynamics force-evaluation step | +| `exchange_segment` | Exchange segment, `step // attempt_interval` (see [Replica exchange](#replica-exchange)) | +| `sampling_epoch` | Consistency epoch, `step // steps_per_epoch` | + +A `thermodynamic_state_id` you set yourself is preserved, never overwritten. +Without replica exchange, `exchange_segment` falls back to the epoch length, +since there are no exchange segments to count. + +## Checkpoint and restore + +```python +sampling.checkpoint("run.zarr") # only at an epoch boundary + +# ... later, in a fresh process ... +sampling2 = EnhancedSampling(dynamics, {"umbrella": umbrella}) +batch = sampling2.restore("run.zarr") # returns a force-primed batch +batch = sampling2.run(batch, n_steps=10_000, prime=False) +``` + +Resuming reproduces the **identical trajectory**. `NVTLangevin` derives its +noise from `random_seed + step_count` rather than a stateful generator, so +restoring those two integers restores the noise sequence exactly — there is no +RNG object to serialise. + +### Only at an epoch boundary + +`checkpoint()` raises unless `step_count % steps_per_epoch == 0`, and the error +names the next valid step. This is not bookkeeping fussiness: an epoch boundary +is the only point with no pending `update()` and no in-flight epoch commit, so +anywhere else risks capturing a bias mid-mutation. + +Being *at* a boundary is not the same as being quiescent, though. Both the +epoch commit and the replica exchange fire **lazily** — the first step of the +next epoch or segment is what notices the boundary was crossed — so +immediately after `run(..., n_steps=N)` neither has happened. + +`checkpoint()` therefore drains both itself before collecting any state, in +the same order the runtime uses: **exchange first, then commit**, because the +commit publishes shared history and doing it before the swap would publish +under labels that are about to change. A shared-history bias is recorded with +its deposits merged rather than still pending, and the labels are post-swap. + +Both drains are idempotent — tracked per epoch index and per segment index — +so the lazy path on the next step sees them as already done and cannot +double-count. + +:::{note} +`checkpoint()` is therefore **not a passive snapshot**: it can advance the +exchange assignment as part of reaching a quiescent point. Read +`batch.thermodynamic_state_id` *after* checkpointing if you want the value +that was saved. +::: + +### Transactional by construction + +The store is written walker batch → components → **manifest last**. The +manifest is the commit marker: + +- No manifest ⇒ the write was interrupted ⇒ `read_checkpoint` refuses it. +- **Everything** is checksummed (SHA-256) and verified on read, so damage + *after* the manifest landed is caught too: each `sampling/` component + individually, plus a `batch_checksum` covering `meta/`, `core/`, and + `custom/`. Cover is mandatory, not best-effort — a manifest that declares a + component without a checksum, carries a checksum for no component, or omits + the batch checksum is rejected as invalid — otherwise deleting one key from + the manifest would be enough to leave that component free to modify. +- The batch checksum is the one most easily forgotten: the walker batch is + written by `AtomicDataZarrWriter`, outside the component path, so covering + only the sampling state would attest to the bias and integrator while + restoring corrupted positions or a scrambled walker identity in silence. +- **No pickle payloads.** State is Zarr arrays and JSON attributes, so a + checkpoint is readable by anything that reads Zarr and loading one cannot + execute code. An unsupported value type raises rather than falling back. + +```text +run.zarr/ + meta/, core/, custom/ walker batch (custom/ carries walker identity) + sampling/ + manifest written last — the commit; holds every checksum + dynamics/ step counter, RNG seed, per-system integrator state + biases// each bias's state_dict() + runner/ walker-id allocation, epoch counters +``` + +### Bias configuration is validated, not just bias class + +The manifest records each bias's *class*, but a class name says nothing about +the settings its saved state depends on. An ABF histogram is state; the +`cv_range` that decides what its bins mean is configuration. Restoring the +first without the second leaves the counts shape-compatible and silently +relabels every bin — bin 5 stops meaning `r = 1.55` and starts meaning +`r = 3.1`, carrying its accumulated mean force with it. + +Every adaptive bias therefore records a `config_fingerprint()` inside its own +`state_dict()`, checked on load: + +| Bias | Checked | +|------|---------| +| `AdaptiveBiasingForce` | `atom_indices`, `cv_range`, `n_bins`, `temperature`, `min_samples`, `full_samples`, `max_force` | +| `WellTemperedMetaDynamicsBias` | `height`, `sigma`, `temperature`, `bias_factor`, `storage`, `history`, `ramp_depositions`, `periods` | +| `RMSDMetaDynamicsBias` | `k_push`, `alpha`, `storage`, `history`, `ramp_depositions`, `atom_indices` | + +Because the check lives in `load_state_dict` rather than in the manifest, it +covers a bias restored directly as well as one restored through the runner. + +Capacity (`max_hills`, `max_references`) is deliberately **not** checked: +`storage="grow"` legitimately reaches a size the constructor never had, and +`load_state_dict` already resizes to match. + +The check runs *before* delegating to `nn.Module.load_state_dict`, which +matters for configuration held as a buffer — `sigma`, `atom_indices`. Loading +overwrites buffers, so a check afterwards would come too late to stop the +caller's value being replaced by the checkpoint's, which is the opposite of +what asking for it meant. + +### Model weights are not restored + +Reconstruct the model — including loading its weights through its own API — +before calling `restore()`. The manifest records the model class, dynamics +class, and bias set, and `restore()` refuses a mismatch; but that proves the +*architecture* agrees, not the weights. + +### `warm_start` vs `restore` + +| | `warm_start(frames)` | `restore(path)` | +|---|---|---| +| Bias history | replayed, approximately | exact | +| Velocities, RNG, integrator state | not restored | exact | +| Use when | continuing from a trajectory snapshot | resuming a run exactly | + +They are mutually exclusive: `warm_start()` after `restore()` raises, because +replaying history the restored state already contains would corrupt it. + +## Replica exchange + +A ladder of thermodynamic states, all advanced as one batch, with periodic +swaps of which walker sits on which rung: + +```python +from nvalchemi.enhanced_sampling import ReplicaExchange, ThermodynamicState + +states = [ + ThermodynamicState(state_id=i, temperature=300.0 * 1.15 ** i) + for i in range(4) +] +exchange = ReplicaExchange( + states=states, + initial_state_ids=torch.arange(4), # must be a permutation + attempt_interval=100, # steps per exchange segment + random_seed=2024, +) +sampling = EnhancedSampling(dynamics, biases={}, replica_exchange=exchange) +batch = sampling.run(batch, n_steps=100_000) +``` + +### One walker per rung + +Exchange presumes a bijection: every walker holds exactly one state and every +state exactly one walker, because pairing looks up "which walker holds state +*k*". The runner validates that on the first step, whether the assignment came +from `initial_state_ids` or was already on the batch: + +```text +ReplicaExchange: the ladder has 4 state(s) but the batch has 2 walker(s). +ReplicaExchange: batch.thermodynamic_state_id must be a permutation of 0..3, + got [0, 0, 1, 2]. +``` + +Both are configuration errors that would otherwise surface much later and much +less clearly — a size mismatch as `Length mismatch: 4 vs 2` from inside the +batch storage, a duplicate as a bare `KeyError` from the pair lookup. + +### Labels move, coordinates do not + +An accepted swap permutes `thermodynamic_state_id`. The walker keeps its row, +its velocities, and its integrator arrays; the temperature assigned to it +changes. Nothing is copied between rows, which is what makes the move viable +inside a batched GPU step. + +The swap is **indivisible**: the label, the integrator's target temperature, +the velocity rescaling, and the forces all move together. A walker labelled +one rung while its thermostat targets another samples the wrong ensemble with +no symptom, so the runner refuses at construction any integrator that cannot +rebind: + +```text +TypeError: replica exchange needs NVE to implement +apply_thermodynamic_state(), so an accepted swap can rebind temperature, +velocities, and thermostat state together. +``` + +`NVTLangevin` rescales velocities by `sqrt(T_new / T_old)`. `NVTNoseHoover` +additionally transforms its chain: `Q` scales with `kT` and `eta_dot` with +`1/sqrt(kT)`, which leaves the chain kinetic energy invariant — injecting +thermostat energy on a swap is exactly what breaks detailed balance. + +### Acceptance + +The rule is **inferred from the ladder**, never declared, because a declared +rule that disagreed with the ladder would be silent and wrong acceptance +breaks detailed balance without any symptom a run would show. + +| Ladder | Rule | Formula | +|--------|------|---------| +| Temperatures differ | temperature | `log a = min(0, (β_i − β_j)(U_i − U_j))` | +| Temperatures equal | umbrella | `log a = min(0, u_i(x_i) + u_j(x_j) − u_i(x_j) − u_j(x_i))` | + +Umbrella acceptance needs the bias evaluated under swapped labels, which +costs one extra bias evaluation per attempt. + +A ladder that varies temperature *and* bias window at once needs a combined +rule that is not implemented. The temperature rule alone omits the cross-state +bias terms, so running it anyway breaks detailed balance with no symptom — +it is therefore **rejected**, twice over: + +- A bias that sets `state_dependent_for_exchange` is refused at construction. + `HarmonicUmbrellaBias` sets it whenever it has more than one window. +- At prime time the runner **probes** every bias empirically: it evaluates + each one under the current assignment and under a rotated one, at identical + coordinates. A bias whose energy is independent of the assignment returns + the same number twice; one that reads `thermodynamic_state_id` does not. + That catches a user-written bias which declares nothing. + +Vary one or the other. A **single-window** `HarmonicUmbrellaBias` is fine +alongside a temperature ladder — it ignores `thermodynamic_state_id` rather +than treating it as a window index, so the ids are free to address the +ladder. + +### Pair schedule + +Segments alternate even and odd offsets: `(0,1),(2,3)` then `(1,2),(3,4)`. +No state appears twice in one segment, which is what lets every pair be +decided simultaneously; two segments cover every neighbouring pair. + +A segment's pairs are attempted when it **completes** — entering segment *s* +decides segment *s−1*, the same way entering epoch *e* commits epoch *e−1*. +So the first swap lands at `attempt_interval`, using segment 0's pairs. + +### Tuning the ladder + +```python +exchange.acceptance_rate # overall +exchange.pair_acceptance_rates() # per neighbouring pair +``` + +Per-pair rates are what a ladder is tuned on. A pair far below the others is +a gap the walkers cannot cross and needs another rung; uniformly high rates +mean the rungs are closer than they need to be. + +### Restoring an exchange run + +The manifest records the ladder — mode, acceptance rule, `attempt_interval`, +and temperatures — and `restore()` refuses a mismatch. That includes both +directions of exchange-versus-none: + +```text +EnhancedSampling.restore: the checkpoint was written by a different configuration: + exchange temperatures: checkpoint has [300.0, 350.0, 400.0], + this runner has [100.0, 200.0, 900.0] +``` + +This is not pedantry. The ladder decides what a swap *means*: restoring into +different temperatures would keep the walker assignment and the acceptance +counters while silently changing the exponent every future swap is decided +on. `initial_state_ids` is deliberately *not* checked — it seeds the +assignment only when the batch does not already carry one, and a restored +batch always does. + +### Reproducibility + +Acceptance draws come from `random_seed + exchange_id` rather than a +long-lived generator — the same counter-based scheme `NVTLangevin` uses for +its noise. A checkpoint therefore stores two integers instead of an RNG blob, +and a restored run reproduces the same accept/reject decisions. Exchange +state lives under `sampling/exchange/`. + +### Not supported + +Asynchronous exchange (pair-local rendezvous, non-blocking workers) is not +implemented; `mode="asynchronous"` raises. A force-only bias such as adaptive +biasing force cannot participate — the acceptance rule needs a cross-state +bias energy — and is rejected rather than silently excluded. + +## Relationship to `BiasedPotentialHook` + +{class}`~nvalchemi.hooks.BiasedPotentialHook` covers similar ground and is +**deprecated**. Its `bias_fn(batch) -> (energy, forces)` contract has no slot +for a cell response, so a bias applied through it contributes no stress and is +invisible to an NPT/NPH barostat — the cell evolves as if the bias were absent, +with no error raised. It also cannot check that the returned forces are +`-dE/dr`, and composes several biases by sequential in-place mutation. + +With `EnhancedSampling` now available, the migration path is complete — +anything the hook does, this subpackage does. Existing hook-based code is +correct under NVE and NVT, where nothing reads the stress, so it can be +migrated when convenient rather than urgently. The hook remains functional and +no removal date is set. + +No adapter is provided: bridging a `BiasPotential` onto `bias_fn` would have +to discard `BiasResult.stress`, reintroducing the exact failure the new API +removes. + +## Not yet implemented + +- Adaptive biasing force over any CV other than a pair distance; each + coordinate needs its own metric correction. +- Asynchronous replica exchange (pair-local rendezvous, non-blocking + workers). `mode="asynchronous"` raises; synchronous exchange is available. +- The combined temperature-plus-window acceptance rule; see + [Acceptance](#acceptance) for what is rejected and why. +- General triclinic MIC for unreduced cells. +- Domain decomposition. `ConservativeBias.distribution_spec()` returns `None`, + which makes `DistributedModel` raise rather than shard a bias whose + cross-rank semantics are undefined. A bias that genuinely is local can + override it; a CV like `pair_distance` across the cell is not. + +## See also + +- {doc}`Conventions ` — virial, stress, and pressure signs. +- {doc}`Hooks ` — the hook protocol the runner builds on. +- {doc}`Dynamics ` — integrators and the step sequence. diff --git a/docs/userguide/index.md b/docs/userguide/index.md index 57aebc04..81933d17 100644 --- a/docs/userguide/index.md +++ b/docs/userguide/index.md @@ -44,6 +44,7 @@ For install options, refer to the install guide below. - {doc}`Hooks: Observe & Modify ` - {doc}`Reporting: Summaries and Dashboards ` - [Dynamics: Optimization and MD](dynamics) +- {doc}`Enhanced Sampling: Biases and Free Energy ` ## Distributed Simulations @@ -84,6 +85,7 @@ serialization hooks reporting dynamics +enhanced_sampling ``` ```{toctree} diff --git a/examples/advanced/01_biased_potential.py b/examples/advanced/01_biased_potential.py index a521eb1c..481eb613 100644 --- a/examples/advanced/01_biased_potential.py +++ b/examples/advanced/01_biased_potential.py @@ -30,6 +30,17 @@ sweep the target position along a reaction coordinate and post-process the windowed histograms with WHAM or MBAR. +.. note:: + + :class:`~nvalchemi.hooks.BiasedPotentialHook` is **deprecated** in favour + of :mod:`nvalchemi.enhanced_sampling`, which derives bias forces *and* + stress from a single energy definition by autograd. A ``bias_fn`` bias + contributes no stress, so it is invisible to an NPT/NPH barostat — fine + for the NVT run below, silently wrong under a barostat. The equivalent + with the new API is in ``11_umbrella_sampling.py``; see + :mod:`nvalchemi.enhanced_sampling` for the full comparison. This hook + stays functional and no removal date is set. + Key concepts demonstrated ------------------------- * Implementing a ``bias_fn(batch) -> (energy, forces)`` closure. diff --git a/examples/advanced/11_umbrella_sampling.py b/examples/advanced/11_umbrella_sampling.py new file mode 100644 index 00000000..74ddfd96 --- /dev/null +++ b/examples/advanced/11_umbrella_sampling.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Batched Umbrella Sampling with EnhancedSampling +=============================================== + +Umbrella sampling computes a free-energy profile along a reaction coordinate +by running several simulations, each restrained to a different value of that +coordinate. The restraints keep the system in regions it would otherwise +never visit; the resulting biased histograms are recombined afterwards with +WHAM or MBAR. + +The usual cost is that every window is a separate simulation. Here all +windows are **rows of one batch**: each graph carries a +``thermodynamic_state_id`` selecting its own window center, so a single +:class:`~nvalchemi.enhanced_sampling.HarmonicUmbrellaBias` serves all of them +in one batched GPU force evaluation. + +This example restrains the distance between two argon atoms in a +Lennard-Jones cluster across five windows, and adds an +:class:`~nvalchemi.enhanced_sampling.UpperWall` to stop any window wandering +off to dissociation. + +Key concepts demonstrated +------------------------- +* A collective variable as a plain callable over + :func:`~nvalchemi.enhanced_sampling.pair_distance`. +* Per-window centers selected by ``thermodynamic_state_id``. +* Composing two biases; they are summed against the *same* unmodified model + output, so neither observes the other's forces. +* Reading per-bias diagnostics from ``sampling.last_outputs``. + +Applications +------------ +* Free-energy profiles along a bond, a distance, or a coordination number. +* Potential of mean force for ion transport or ligand unbinding. +* Restrained sampling to generate training data in a targeted region. +""" + +from __future__ import annotations + +import logging +import os + +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.enhanced_sampling import ( + EnhancedSampling, + HarmonicUmbrellaBias, + UpperWall, + pair_distance, +) +from nvalchemi.models.lj import LennardJonesModelWrapper + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger(__name__) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +# Keep the doc build fast; a real run needs orders of magnitude more steps. +N_STEPS = 20 if os.environ.get("NVALCHEMI_SPHINX_BUILD") else 200 + +# %% +# Build one window per restraint center +# ------------------------------------- +# Every window is an independent copy of the same cluster. The output +# buffers (``forces``, ``energy``) must exist up front: dynamics writes model +# outputs back in place rather than creating the fields. + +WINDOW_CENTERS = [3.0, 3.5, 4.0, 4.5, 5.0] # angstrom +N_ATOMS = 8 + +torch.manual_seed(0) + + +def make_cluster() -> AtomicData: + """Return one argon cluster with the buffers dynamics writes into. + + Atoms sit on a 2x2x2 cube at the LJ minimum separation. Random + positions would place some pairs well inside sigma = 3.4 A, where the + repulsive wall is steep enough that the first step diverges — the + restraint is then blamed for what is really an overlapping start. + """ + spacing = 4.0 # ~2^(1/6) * sigma, near the LJ minimum + grid = torch.tensor( + [[x, y, z] for x in (0.0, 1.0) for y in (0.0, 1.0) for z in (0.0, 1.0)] + ) + positions = grid[:N_ATOMS] * spacing + # Atoms 0 and 1 are a cube edge apart, so the CV starts at `spacing`. + data = AtomicData( + positions=positions, + atomic_numbers=torch.full((N_ATOMS,), 18, dtype=torch.long), + atomic_masses=torch.full((N_ATOMS,), 39.948), + forces=torch.zeros(N_ATOMS, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(N_ATOMS, 3)) + return data + + +batch = Batch.from_data_list([make_cluster() for _ in WINDOW_CENTERS]).to(DEVICE) + +# Each graph gets its own window. This is the field HarmonicUmbrellaBias +# indexes to pick a center, and it is what makes one batch five windows. +batch["thermodynamic_state_id"] = torch.arange(len(WINDOW_CENTERS), device=DEVICE) + +# %% +# Define the collective variable +# ------------------------------ +# A CV is any differentiable ``cv(batch) -> Tensor[B, D]``. No base class, no +# registration — a closure over :func:`pair_distance` is a complete CV. + +ATOM_PAIR = torch.tensor([0, 1], device=DEVICE) + + +def bond_distance(atoms: Batch) -> torch.Tensor: + """Distance between atoms 0 and 1, shape ``[B, 1]``.""" + return pair_distance(atoms, ATOM_PAIR) + + +# %% +# Build the biases +# ---------------- +# ``centers`` has shape ``[S, D]`` — one row per window, one column per CV +# dimension. ``stiffness`` broadcasts across windows here, but could equally +# be per-window with shape ``[S, D, D]``. + +umbrella = HarmonicUmbrellaBias( + cv=bond_distance, + centers=torch.tensor(WINDOW_CENTERS).unsqueeze(-1), # [5, 1] + stiffness=5.0, # eV/A^2 + name="umbrella", +) + +# A one-sided wall stops a window that escapes its basin from running away to +# dissociation. It contributes exactly zero energy and force while the +# distance stays below the threshold. +wall = UpperWall( + cv=bond_distance, + threshold=8.0, + stiffness=20.0, + name="dissociation_wall", +) + +# %% +# Run biased dynamics +# ------------------- +# The runner registers one internal hook on the dynamics and otherwise leaves +# it alone: the model, the integrator, and the thermostat are unchanged. + +model = LennardJonesModelWrapper(sigma=3.4, epsilon=0.0104, cutoff=8.5).to(DEVICE) +dynamics = NVTLangevin(model=model, dt=0.5, temperature=120.0, friction=0.05) + +# A cutoff model needs its neighbour list rebuilt at BEFORE_COMPUTE. The +# runner fires that stage during priming too, so the first force evaluation +# is as valid as every later one. +for hook in model.make_neighbor_hooks(): + dynamics.register_hook(hook) + +sampling = EnhancedSampling( + dynamics=dynamics, + biases={"umbrella": umbrella, "dissociation_wall": wall}, +) + +logger.info("Initial CV per window: %s", bond_distance(batch).flatten().tolist()) + +batch = sampling.run(batch, n_steps=N_STEPS) + +final_cv = bond_distance(batch).flatten().tolist() +logger.info("Target centers: %s", WINDOW_CENTERS) +logger.info("Final CV per window: %s", [round(v, 3) for v in final_cv]) + +# %% +# Read the diagnostics +# -------------------- +# ``last_outputs`` separates the physical model contribution from each bias +# and from the total. These are the tensors a WHAM/MBAR post-processing step +# needs: the CV value, the physical energy, and the bias energy per window. +# +# Free-energy reconstruction is deliberately **not** built in — the raw +# per-window data is returned for ``pymbar`` or an equivalent tool. + +for key in sorted(sampling.last_outputs): + value = sampling.last_outputs[key] + logger.info(" %-28s shape=%s", key, tuple(value.shape)) + +umbrella_energy = sampling.last_outputs["bias/umbrella/energy"].flatten() +logger.info( + "Umbrella energy per window (eV): %s", [round(float(v), 4) for v in umbrella_energy] +) + +# The wall should be dormant: every window is well inside 8 A. +wall_energy = sampling.last_outputs["bias/dissociation_wall/energy"].flatten() +logger.info( + "Wall energy per window (eV): %s", [round(float(v), 4) for v in wall_energy] +) diff --git a/examples/advanced/12_replica_exchange.py b/examples/advanced/12_replica_exchange.py new file mode 100644 index 00000000..c59f0f1f --- /dev/null +++ b/examples/advanced/12_replica_exchange.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Temperature Replica Exchange (REMD) +=================================== + +A single low-temperature trajectory gets stuck: the barrier it needs to cross +is many :math:`k_BT` high, so it may never be crossed in the time available. +Replica exchange runs a ladder of temperatures at once and periodically swaps +which walker sits at which temperature. A configuration that wanders up the +ladder crosses barriers easily at high temperature, and when it comes back +down it lands in a basin the cold replica might never have reached on its own. + +The whole ladder is **one batch**: each graph is a walker, and +``thermodynamic_state_id`` records which rung it currently occupies. One +batched force evaluation advances every replica per step. + +Exchange swaps **labels, not coordinates**. A walker keeps its execution +slot, its velocities, and its integrator arrays; the temperature assigned to +it changes. That makes the move local — nothing is copied between rows — +which is what allows it inside a batched GPU step. + +Key concepts demonstrated +------------------------- +* Building a geometric temperature ladder with + :class:`~nvalchemi.enhanced_sampling.ThermodynamicState`. +* The even/odd pair schedule and the Metropolis acceptance rule. +* Reading per-pair acceptance rates, which is how a ladder is tuned. +* Confirming the integrator's target temperature follows the assignment. + +Applications +------------ +* Conformer and reaction-path exploration where barriers exceed a few kT. +* Crystal polymorph search and nucleation. +* Generating diverse training configurations for a potential. +""" + +from __future__ import annotations + +import logging +import os + +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling import ( + EnhancedSampling, + ReplicaExchange, + ThermodynamicState, +) +from nvalchemi.models.lj import LennardJonesModelWrapper + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger(__name__) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +N_STEPS = 80 if os.environ.get("NVALCHEMI_SPHINX_BUILD") else 400 + +# %% +# Build the temperature ladder +# ---------------------------- +# A geometric spacing keeps the acceptance rate roughly uniform across rungs: +# the acceptance depends on the *ratio* of neighbouring temperatures, not +# their difference, so evenly spaced temperatures would exchange readily at +# the cold end and almost never at the hot end. + +N_REPLICAS = 4 +BASE_TEMPERATURE = 80.0 # K — argon is liquid around here +LADDER_FACTOR = 1.25 + +states = [ + ThermodynamicState(state_id=i, temperature=BASE_TEMPERATURE * LADDER_FACTOR**i) + for i in range(N_REPLICAS) +] +logger.info("Temperature ladder (K): %s", [round(s.temperature, 1) for s in states]) + +# %% +# One walker per rung +# ------------------- +# Output buffers must exist up front — dynamics writes model outputs back in +# place rather than creating the fields. + +N_ATOMS = 8 + + +def make_cluster(seed: int) -> AtomicData: + """Return an argon cluster on a cube at the LJ minimum separation.""" + spacing = 4.0 + grid = torch.tensor( + [[x, y, z] for x in (0.0, 1.0) for y in (0.0, 1.0) for z in (0.0, 1.0)] + ) + torch.manual_seed(seed) + positions = grid[:N_ATOMS] * spacing + 0.05 * torch.randn(N_ATOMS, 3) + data = AtomicData( + positions=positions, + atomic_numbers=torch.full((N_ATOMS,), 18, dtype=torch.long), + atomic_masses=torch.full((N_ATOMS,), 39.948), + forces=torch.zeros(N_ATOMS, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(N_ATOMS, 3)) + return data + + +batch = Batch.from_data_list([make_cluster(seed) for seed in range(N_REPLICAS)]).to( + DEVICE +) + +# %% +# Configure the exchange +# ---------------------- +# ``initial_state_ids`` must be a permutation: replica exchange presumes one +# walker per rung, and a duplicate would let two walkers claim the same +# temperature. Acceptance randomness is derived from ``random_seed`` plus an +# attempt counter, so a run is reproducible and a checkpoint needs only two +# integers rather than an RNG blob. + +exchange = ReplicaExchange( + states=states, + initial_state_ids=torch.arange(N_REPLICAS), + attempt_interval=10, # dynamics steps per exchange segment + random_seed=2024, +) +logger.info("Acceptance rule inferred from the ladder: %s", exchange.acceptance) +logger.info("Segment 0 pairs: %s", exchange.pair_schedule(0)) +logger.info("Segment 1 pairs: %s", exchange.pair_schedule(1)) + +# %% +# Run +# --- +# The runner validates up front that the integrator can rebind a +# thermodynamic state. An integrator that could only accept the new label +# would keep sampling the old temperature — wrong, and with no symptom the +# run would show — so ``NVE`` and friends are rejected rather than silently +# mis-sampled. + +model = LennardJonesModelWrapper(sigma=3.4, epsilon=0.0104, cutoff=8.5).to(DEVICE) +dynamics = NVTLangevin(model=model, dt=0.5, temperature=BASE_TEMPERATURE, friction=0.05) +for hook in model.make_neighbor_hooks(): + dynamics.register_hook(hook) + +sampling = EnhancedSampling( + dynamics=dynamics, + biases={}, # pure temperature REMD; biases would compose here + replica_exchange=exchange, + steps_per_epoch=100, +) + +batch = sampling.run(batch, n_steps=N_STEPS) + +# %% +# Read the acceptance statistics +# ------------------------------ +# Per-pair rates are what a ladder is tuned on. A pair far below the others +# is a gap the walkers cannot cross, and the ladder needs another rung there; +# uniformly high rates mean the rungs are closer than they need to be and the +# ladder is wasting replicas. + +logger.info("") +logger.info( + "Overall acceptance: %d/%d = %.2f", + exchange.accepted, + exchange.attempts, + exchange.acceptance_rate, +) +for index, rate in enumerate(exchange.pair_acceptance_rates()): + logger.info( + " states %d <-> %d (%.0f K <-> %.0f K): %.2f", + index, + index + 1, + states[index].temperature, + states[index + 1].temperature, + rate, + ) + +# Expect near-unit acceptance here: an 8-atom cluster has energy fluctuations +# far smaller than the spread between rungs, so the Metropolis exponent stays +# close to zero and almost everything is accepted. That is the "rungs closer +# than they need to be" regime — a production system of thousands of atoms +# has much larger fluctuations, and the ladder would be widened until the +# rates land near the usual 20-30% target. + +# %% +# Confirm the swap was indivisible +# -------------------------------- +# The label, the integrator's target temperature, the velocity scaling, and +# the forces all have to move together. A walker whose label says one rung +# while its thermostat still targets another would sample the wrong ensemble +# silently, so it is worth asserting rather than assuming. + +assignment = batch.thermodynamic_state_id.reshape(-1).tolist() +targets = (dynamics._state.temperature.reshape(-1) / KB_EV).tolist() + +logger.info("") +logger.info("Final assignment (walker -> state): %s", assignment) +for walker, state_id in enumerate(assignment): + expected = states[state_id].temperature + logger.info( + " walker %d holds state %d: target %.1f K (ladder says %.1f K)", + walker, + state_id, + targets[walker], + expected, + ) diff --git a/examples/advanced/13_metadynamics.py b/examples/advanced/13_metadynamics.py new file mode 100644 index 00000000..fd6c4287 --- /dev/null +++ b/examples/advanced/13_metadynamics.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Multiple-Walker Well-Tempered Metadynamics +========================================== + +Metadynamics fills the free-energy basin the system is sitting in with +Gaussian hills, until the barrier out of it is no longer a barrier. Where +umbrella sampling needs you to know in advance which values of the reaction +coordinate to restrain to, metadynamics discovers them: it goes wherever it +has not already been. + +The *well-tempered* variant shrinks each new hill in proportion to the bias +already accumulated at that point, + +.. math:: h_t = h_0 \\exp\\!\\left(-\\frac{V(s_t)}{k_B T (\\gamma - 1)}\\right) + +so the sum converges instead of filling forever, and the converged bias can +be turned back into a free-energy profile. + +This example runs the **multiple-walker** scheme: several walkers explore the +same collective variable at once, all depositing into one shared hill +history, so a basin is filled roughly ``B`` times faster. Because the +walkers are rows of a single batch, that costs one batched GPU force +evaluation per step, not ``B`` separate simulations. + +Key concepts demonstrated +------------------------- +* Depositing hills on a schedule with ``update_frequency``, driven by + :class:`~nvalchemi.enhanced_sampling.EnhancedSampling`. +* Shared multi-walker history, and how it differs from private history. +* Choosing a storage policy, and why ``preallocated`` raises rather than + silently discarding hills. +* Recovering a free-energy profile with + :meth:`~nvalchemi.enhanced_sampling.WellTemperedMetaDynamicsBias.free_energy`. + +Applications +------------ +* Free-energy profiles when the relevant windows are not known in advance. +* Escaping metastable states that plain MD would stay trapped in. +* Barrier crossing for conformational change, dissociation, or diffusion. +""" + +from __future__ import annotations + +import logging +import os + +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.enhanced_sampling import ( + EnhancedSampling, + UpperWall, + WellTemperedMetaDynamicsBias, + pair_distance, +) +from nvalchemi.models.lj import LennardJonesModelWrapper + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger(__name__) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +# Keep the doc build fast; a converged run needs orders of magnitude more. +N_STEPS = 40 if os.environ.get("NVALCHEMI_SPHINX_BUILD") else 600 + +# %% +# Build the walkers +# ----------------- +# Four independent copies of the same argon cluster. They are not windows: +# no walker is restrained anywhere, and they differ only in the random +# velocities the thermostat gives them. + +N_WALKERS = 4 +N_ATOMS = 8 +TEMPERATURE = 120.0 # K + +torch.manual_seed(0) + + +def make_cluster() -> AtomicData: + """Return one argon cluster with the buffers dynamics writes into. + + Atoms sit on a 2x2x2 cube near the LJ minimum separation; random + positions would place some pairs inside sigma, where the repulsive wall + is steep enough that the first step diverges. + """ + spacing = 4.0 + grid = torch.tensor( + [[x, y, z] for x in (0.0, 1.0) for y in (0.0, 1.0) for z in (0.0, 1.0)] + ) + data = AtomicData( + positions=grid[:N_ATOMS] * spacing, + atomic_numbers=torch.full((N_ATOMS,), 18, dtype=torch.long), + atomic_masses=torch.full((N_ATOMS,), 39.948), + forces=torch.zeros(N_ATOMS, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(N_ATOMS, 3)) + return data + + +batch = Batch.from_data_list([make_cluster() for _ in range(N_WALKERS)]).to(DEVICE) + +# %% +# Define the collective variable +# ------------------------------ +# The same plain-callable contract as umbrella sampling: any differentiable +# ``cv(batch) -> Tensor[B, D]``. + +ATOM_PAIR = torch.tensor([0, 1], device=DEVICE) + + +def bond_distance(atoms: Batch) -> torch.Tensor: + """Distance between atoms 0 and 1, shape ``[B, 1]``.""" + return pair_distance(atoms, ATOM_PAIR) + + +# %% +# Configure the bias +# ------------------ +# Three choices carry most of the physics: +# +# ``sigma`` +# Hill width, in CV units. Roughly the resolution of the resulting free +# energy: too wide smears out real features, too narrow needs far more +# hills to fill anything. +# +# ``bias_factor`` (gamma) +# Sets how high the bias is allowed to climb — the effective temperature +# of the CV is ``gamma * T``. Larger explores further and converges +# slower. Passing ``None`` gives standard, non-converging metadynamics. +# +# ``update_frequency`` +# Steps between depositions. Depositing faster than the system +# decorrelates biases the estimate; the usual choice is a few hundred +# steps. +# +# ``storage="preallocated"`` keeps the hill tensors a fixed shape for the +# whole run, which is what keeps ``energy()`` compilable without retracing. +# It **raises** when ``max_hills`` is exhausted rather than dropping hills: +# silently discarding them would change the physics of a converging run with +# no error to show for it. + +metad = WellTemperedMetaDynamicsBias( + cv=bond_distance, + height=0.005, # eV + sigma=0.25, # angstrom + temperature=TEMPERATURE, + bias_factor=8.0, + update_frequency=25, + storage="preallocated", + max_hills=512, + history="shared", # multiple-walker: every walker feels every hill + name="metad", +) + +# Metadynamics pushes outward by construction, so an unbounded CV eventually +# dissociates the pair. A wall bounds the explored region without touching +# the interior, where it contributes exactly zero. +wall = UpperWall(cv=bond_distance, threshold=9.0, stiffness=20.0, name="wall") + +# %% +# Run +# --- +# The runner calls ``update()`` exactly once per due step, at ``AFTER_STEP`` +# so each hill marks the configuration the walker actually reached. A +# deposition bumps the bias state version, and the runner re-primes forces in +# response, so a new hill is felt on the very next step rather than one step +# late. + +model = LennardJonesModelWrapper(sigma=3.4, epsilon=0.0104, cutoff=8.5).to(DEVICE) +dynamics = NVTLangevin(model=model, dt=0.5, temperature=TEMPERATURE, friction=0.05) + +for hook in model.make_neighbor_hooks(): + dynamics.register_hook(hook) + +sampling = EnhancedSampling(dynamics=dynamics, biases={"metad": metad, "wall": wall}) + +logger.info( + "Initial CV per walker: %s", + [round(v, 3) for v in bond_distance(batch).flatten().tolist()], +) + +batch = sampling.run(batch, n_steps=N_STEPS) + +logger.info( + "Final CV per walker: %s", + [round(v, 3) for v in bond_distance(batch).flatten().tolist()], +) +logger.info( + "Depositions: %d, hills stored: %d of %d", + int(metad.deposits), + int(metad.hill_count), + metad.capacity, +) + +# %% +# Hill heights decay +# ------------------ +# This is the well-tempered signature. As the bias accumulates, each new hill +# is shorter than the last; a run whose recent hills are still near ``h_0`` +# has not begun to converge. + +heights = metad.hill_heights[: int(metad.hill_count)] +logger.info("First hill height: %.6f eV", float(heights[0])) +logger.info("Last hill height: %.6f eV", float(heights[-1])) + +# %% +# Recover the free-energy profile +# ------------------------------- +# ``F(s) = -(gamma / (gamma - 1)) V(s)``, up to an additive constant. The +# profile is only meaningful where hills were actually deposited; the flat +# regions beyond simply mean the walkers never went there. + +grid = torch.linspace(3.0, 8.0, 26, device=DEVICE).unsqueeze(-1) +free_energy = metad.free_energy(grid) +free_energy = free_energy - free_energy.min() + +logger.info("Free-energy profile (eV, shifted to zero minimum):") +for value, energy in zip(grid.flatten().tolist(), free_energy.tolist(), strict=True): + logger.info(" d = %.2f A F = %.4f eV", value, energy) + +# %% +# Shared versus private history +# ----------------------------- +# ``history="shared"`` is what makes this *multiple-walker* metadynamics +# rather than four independent runs: one history, filled four times as fast. +# +# The alternatives are ``"walker"``, which gives each walker its own private +# history and so runs ``B`` genuinely independent simulations in one batch, +# and ``"state"``, which keys the history by ``thermodynamic_state_id`` for a +# replica-exchange ladder. +# +# Under ``"shared"`` every hill is unowned, which is what the ``-1`` here +# records. + +logger.info("Hill owners (-1 means shared): %s", metad.hill_owner[:4].tolist()) + +# %% +# Diagnostics +# ----------- +# As with every bias, the runner separates the physical contribution from +# each bias and from the total. + +for key in sorted(sampling.last_outputs): + logger.info(" %-28s shape=%s", key, tuple(sampling.last_outputs[key].shape)) diff --git a/examples/advanced/14_rmsd_metadynamics.py b/examples/advanced/14_rmsd_metadynamics.py new file mode 100644 index 00000000..88980c09 --- /dev/null +++ b/examples/advanced/14_rmsd_metadynamics.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +xTB-Style RMSD Metadynamics for Structure Exploration +===================================================== + +Every method so far has needed a collective variable: a scalar you already +know matters. That is the assumption RMSD metadynamics drops. Its history +is a set of retained *structures*, and the bias pushes the system away from +all of them at once: + +.. math:: + + V(x, t) = \\sum_r f_r(t)\\, k_\\mathrm{push} + \\exp\\!\\left(-\\alpha\\, \\mathrm{RMSD}(x, x_r)^2\\right) + +Nothing here names a reaction coordinate. The "coordinate" is the whole +geometry, compared after optimal translation and rotation, which is what +makes this the method of choice for conformer and isomer searching — the +scheme xTB/CREST uses. + +The trade is explicit: there is no free energy at the end. This is a +structure *generator*, and :class:`RMSDMetaDynamicsBias` deliberately has no +``free_energy`` method to suggest otherwise. + +This example explores isomers of an 8-atom Lennard-Jones cluster, a classic +benchmark with several distinct minima separated by real barriers. + +Key concepts demonstrated +------------------------- +* Biasing without choosing a collective variable. +* Optimal-alignment RMSD, and why the bias exerts no net force. +* FIFO retention, and why it is the natural policy here rather than a + compromise. +* Selecting a subset of atoms to compare. + +Applications +------------ +* Conformer generation for a flexible molecule. +* Isomer and polymorph searching. +* Escaping a minimum when you cannot say in advance what direction "out" is. +""" + +from __future__ import annotations + +import logging +import os + +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.enhanced_sampling import EnhancedSampling, RMSDMetaDynamicsBias +from nvalchemi.models.lj import LennardJonesModelWrapper + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger(__name__) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +N_STEPS = 40 if os.environ.get("NVALCHEMI_SPHINX_BUILD") else 800 + +N_WALKERS = 2 +N_ATOMS = 8 +TEMPERATURE = 40.0 # K — cold enough that plain MD stays near one basin + +torch.manual_seed(0) + +# %% +# Build the cluster +# ----------------- +# A non-periodic molecular batch. This bias **rejects** a periodic one: +# Cartesian RMSD against a stored reference is not defined under periodic +# boundary conditions, because an atom crossing a cell face is physically +# unmoved but Cartesian-displaced by a lattice vector. Rather than return a +# plausible wrong number, ``evaluate()`` raises. +# +# Periodicity is read from ``batch.pbc``, not from the presence of a cell. A +# bounding box with ``pbc=False`` — a boxed or solvated molecule — is fine; +# this cluster simply carries no cell at all. + + +def make_cluster() -> AtomicData: + """Return one argon cluster with the buffers dynamics writes into.""" + spacing = 3.8 + grid = torch.tensor( + [[x, y, z] for x in (0.0, 1.0) for y in (0.0, 1.0) for z in (0.0, 1.0)] + ) + positions = grid[:N_ATOMS] * spacing + positions = positions + 0.05 * torch.randn_like(positions) + data = AtomicData( + positions=positions, + atomic_numbers=torch.full((N_ATOMS,), 18, dtype=torch.long), + atomic_masses=torch.full((N_ATOMS,), 39.948), + forces=torch.zeros(N_ATOMS, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(N_ATOMS, 3)) + return data + + +batch = Batch.from_data_list([make_cluster() for _ in range(N_WALKERS)]).to(DEVICE) + +# %% +# Configure the bias +# ------------------ +# ``k_push`` +# How hard to push away from a retained structure, in eV. It sets the +# barrier height the bias can overcome. +# +# ``alpha`` +# Kernel width in ``A^-2``. It sets how *different* a structure has to +# be before it stops feeling the reference, and it is the parameter most +# worth thinking about, because it must match the RMSD scale your system +# actually explores. The kernel only has usable gradient where +# ``alpha * RMSD^2`` is of order one: pick ``alpha`` far too small and +# ``exp(-alpha * RMSD^2)`` sits at ~1 for every structure, so the bias is +# nearly constant and exerts almost no force at all. A rigid cluster +# moving 0.3 A wants ``alpha`` around 10; a floppy molecule sampling +# 1 A wants ``alpha`` around 1. +# +# ``storage="fifo"`` +# The default here, and unlike well-tempered metadynamics it is not a +# compromise. With no free energy to reconstruct, discarding the oldest +# reference is a deliberate choice: it keeps the system pushing outward +# instead of accumulating an ever-stiffer cage that eventually freezes +# it. This is the xTB-compatible policy. +# +# ``ramp_depositions`` +# A new reference lands at RMSD zero — exactly where the system is +# standing. Without a ramp it would switch on at full amplitude at the +# worst possible moment, so the default ramps it in over one deposition. + +rmsd_bias = RMSDMetaDynamicsBias( + k_push=0.08, # eV + alpha=10.0, # A^-2, matched to the ~0.3 A scale this cluster explores + update_frequency=40, + storage="fifo", + max_references=24, + ramp_depositions=1, + name="rmsd", +) + +# %% +# Comparing a subset of atoms +# --------------------------- +# Passing ``atom_indices`` restricts the comparison to chosen atoms, given as +# *per-graph local* indices. For a molecule the usual choice is heavy atoms +# only: methyl hydrogens spinning freely generate RMSD that says nothing +# about the conformer. Correspondence is fixed — atom ``i`` is always +# compared with atom ``i`` of the reference, with no permutation search — so +# two structures identical up to relabelling count as distinct. +# +# Every atom is compared here, which is what ``atom_indices=None`` means. + +# %% +# Run +# --- + +model = LennardJonesModelWrapper(sigma=3.4, epsilon=0.0104, cutoff=8.5).to(DEVICE) +dynamics = NVTLangevin(model=model, dt=0.5, temperature=TEMPERATURE, friction=0.02) + +for hook in model.make_neighbor_hooks(): + dynamics.register_hook(hook) + +sampling = EnhancedSampling(dynamics=dynamics, biases={"rmsd": rmsd_bias}) + +initial = sampling.prime_forces(batch) +logger.info( + "Initial potential energy: %s", + [round(v, 4) for v in sampling.last_outputs["physical/energy"].flatten().tolist()], +) + +batch = sampling.run(batch, n_steps=N_STEPS, prime=False) + +logger.info( + "Depositions: %d, references retained: %d of %d (written: %d)", + int(rmsd_bias.deposits), + int(rmsd_bias.reference_count), + rmsd_bias.capacity, + int(rmsd_bias.references_written), +) +logger.info( + "Final potential energy: %s", + [round(v, 4) for v in sampling.last_outputs["physical/energy"].flatten().tolist()], +) + +# %% +# The retained structures are distinct +# ------------------------------------ +# The point of the bias is that what it collects are genuinely different +# geometries, not the same one re-recorded. Pairwise RMSD between the +# retained references shows that directly. + +from nvalchemi.enhanced_sampling.biases.rmsd_metad import _squared_rmsd + + +def spread(structures: torch.Tensor) -> tuple[float, float]: + """Return the mean and max pairwise RMSD over a set of structures.""" + pairwise = _squared_rmsd(structures, structures).clamp(min=0.0).sqrt() + mask = ~torch.eye(pairwise.shape[0], dtype=torch.bool, device=pairwise.device) + off_diagonal = pairwise[mask] + return float(off_diagonal.mean()), float(off_diagonal.max()) + + +references = rmsd_bias.reference_coords[: int(rmsd_bias.reference_count)] +biased_mean, biased_max = spread(references) +logger.info("Pairwise RMSD between retained structures (A):") +logger.info(" mean %.3f max %.3f", biased_mean, biased_max) + +# %% +# The comparison that matters +# --------------------------- +# Structural spread on its own proves nothing: a hot enough thermostat +# produces scatter without visiting anything new. Repeating the identical +# run with the bias removed separates the two. + +control_batch = Batch.from_data_list([make_cluster() for _ in range(N_WALKERS)]).to( + DEVICE +) +control_model = LennardJonesModelWrapper(sigma=3.4, epsilon=0.0104, cutoff=8.5).to( + DEVICE +) +control_dynamics = NVTLangevin( + model=control_model, dt=0.5, temperature=TEMPERATURE, friction=0.02 +) +for hook in control_model.make_neighbor_hooks(): + control_dynamics.register_hook(hook) + +# A "bias" that only records where the unbiased trajectory went. Same +# deposition schedule, no k_push acting on the dynamics. +recorder = RMSDMetaDynamicsBias( + k_push=1e-12, + alpha=10.0, + update_frequency=40, + max_references=24, + name="rmsd", +) +control = EnhancedSampling(dynamics=control_dynamics, biases={"rmsd": recorder}) +control.run(control_batch, n_steps=N_STEPS) + +control_refs = recorder.reference_coords[: int(recorder.reference_count)] +control_mean, control_max = spread(control_refs) +logger.info("Unbiased control, same schedule and temperature:") +logger.info(" mean %.3f max %.3f", control_mean, control_max) +logger.info("Bias widened the explored set by %.1fx", biased_mean / control_mean) + +# %% +# The bias cannot move the cluster bodily +# --------------------------------------- +# RMSD is measured after optimal translation and rotation, so the energy is +# invariant to rigid motion — and therefore the bias forces sum to zero. A +# bias that failed this would slowly translate the system, and the drift +# would look like physics. + +bias_forces = sampling.last_outputs["bias/rmsd/forces"] +net = bias_forces.reshape(N_WALKERS, N_ATOMS, 3).sum(dim=1) +logger.info("Net bias force per walker: %.2e eV/A", float(net.abs().max())) + +# %% +# No free energy here +# ------------------- +# Well-tempered metadynamics converges to a bias that *is* a free-energy +# estimate. This method does not: the references are discarded as the FIFO +# ring wraps, and the kernel is not a probability model of anything. What it +# produces is a set of structures worth optimising or re-scoring with a +# higher level of theory. +# +# ``WellTemperedMetaDynamicsBias.free_energy`` refuses under ``"fifo"`` for +# exactly this reason, and ``RMSDMetaDynamicsBias`` has no such method at all. + +logger.info("Retained %d structures for downstream optimisation.", references.shape[0]) diff --git a/examples/advanced/15_adaptive_biasing_force.py b/examples/advanced/15_adaptive_biasing_force.py new file mode 100644 index 00000000..3fbbae00 --- /dev/null +++ b/examples/advanced/15_adaptive_biasing_force.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Adaptive Biasing Force Along a Pair Distance +============================================ + +Umbrella sampling restrains, and metadynamics fills. ABF does neither: it +*measures* the mean force along the collective variable in each bin, and +applies its negative. Once a bin is well sampled the residual force along +the CV averages to zero, so the walker diffuses across the coordinate +instead of being held or pushed. + +The payoff is that the accumulated quantity already **is** the free-energy +gradient: + +.. math:: + + \\frac{\\partial A}{\\partial r} = \\left\\langle + -\\frac{(\\mathbf{F}_j - \\mathbf{F}_i)\\cdot\\hat{\\mathbf{u}}}{2} + - \\frac{2 k_B T}{r} \\right\\rangle_r + +Integrating it gives the PMF directly — no hills to deconvolve, no +histograms to reweight. + +That second term is the **metric correction**, and it is the part worth +understanding before using this method. This example demonstrates it +numerically against a case with a known answer. + +Key concepts demonstrated +------------------------- +* The mean-force estimator, checked against an analytic PMF. +* Why a naive Cartesian force projection is wrong, shown by measuring it. +* The minimum-sample threshold and force ramp. +* Force-only biases: no energy, and therefore no replica exchange. + +Applications +------------ +* Potentials of mean force for bond breaking, ion pairing, or unbinding. +* Flattening a known coordinate so that orthogonal degrees of freedom relax. +* Free-energy profiles where hill deposition would be too slow to converge. +""" + +from __future__ import annotations + +import logging +import math +import os + +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling import ( + AdaptiveBiasingForce, + BiasResult, + EnhancedSampling, +) +from nvalchemi.models.lj import LennardJonesModelWrapper + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger(__name__) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +N_STEPS = 40 if os.environ.get("NVALCHEMI_SPHINX_BUILD") else 1500 + +TEMPERATURE = 120.0 # K +KT = KB_EV * TEMPERATURE + +# %% +# The metric correction, measured +# ------------------------------- +# Take two particles that do not interact at all. Their Cartesian force is +# zero, so a naive projection reports a mean force of zero and therefore a +# *flat* free-energy profile. +# +# The true profile is not flat. The number of ways to place two particles a +# distance ``r`` apart grows as the surface of a sphere, ``4 pi r^2``, so +# ``A(r) = -2 k_B T ln r`` — purely entropic, and it drives the pair apart. +# Missing it is not noise; it is a smooth, plausible, wrong answer. +# +# ABF includes the corresponding ``-2 k_B T / r`` term. Here is the +# difference, on a system whose answer we know exactly. + +FREE_PAIR = torch.tensor([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]]) +free_batch = Batch.from_data_list( + [ + AtomicData( + positions=FREE_PAIR, + atomic_numbers=torch.ones(2, dtype=torch.long), + forces=torch.zeros(2, 3), # non-interacting + ) + ] +).to(DEVICE) + +demo = AdaptiveBiasingForce( + atom_indices=torch.tensor([0, 1]), + temperature=TEMPERATURE, + cv_range=(2.0, 5.0), + n_bins=30, + min_samples=0, + full_samples=0, + name="demo", +).to(DEVICE) +demo.update(free_batch, BiasResult()) + +measured = float(demo.mean_force()[int(demo.bin_index(torch.tensor([3.0]))[0])]) +logger.info("Two non-interacting particles at r = 3.0 A:") +logger.info(" naive projection dA/dr = %+.6f eV/A (flat PMF — wrong)", 0.0) +logger.info(" ABF with correction dA/dr = %+.6f eV/A", measured) +logger.info(" analytic -2 kT / r = %+.6f eV/A", -2 * KT / 3.0) + +# %% +# Build the system +# ---------------- +# An argon dimer inside a small cluster. The CV is the distance between +# atoms 0 and 1; the rest of the cluster is the "environment" whose +# rearrangement the PMF integrates over. + +N_ATOMS = 6 +N_WALKERS = 4 + +torch.manual_seed(0) + + +def make_cluster() -> AtomicData: + """Return one argon cluster with the buffers dynamics writes into.""" + spacing = 3.9 + grid = torch.tensor( + [[x, y, z] for x in (0.0, 1.0) for y in (0.0, 1.0) for z in (0.0, 1.0)] + ) + positions = grid[:N_ATOMS] * spacing + 0.05 * torch.randn(N_ATOMS, 3) + data = AtomicData( + positions=positions, + atomic_numbers=torch.full((N_ATOMS,), 18, dtype=torch.long), + atomic_masses=torch.full((N_ATOMS,), 39.948), + forces=torch.zeros(N_ATOMS, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(N_ATOMS, 3)) + return data + + +batch = Batch.from_data_list([make_cluster() for _ in range(N_WALKERS)]).to(DEVICE) + +# %% +# Configure the bias +# ------------------ +# ``min_samples`` / ``full_samples`` +# A mean force from three samples is noise, and applying it would drive +# the walker on the strength of that noise. Below ``min_samples`` a bin +# applies nothing; between the two the applied fraction ramps linearly to +# one, so no bin ever switches on with a jump. +# +# ``cv_range`` +# Outside it, no force is applied and no sample is recorded. Choose it +# to cover the region you want the profile over — ABF will not tell you +# about coordinates it was never allowed to bin. +# +# ``max_force`` +# A bin visited once at an awkward geometry can hold a large estimate. +# The cap bounds what that can do to the trajectory while the average +# settles. + +abf = AdaptiveBiasingForce( + atom_indices=torch.tensor([0, 1]), + temperature=TEMPERATURE, + cv_range=(3.0, 6.5), # angstrom + n_bins=28, + min_samples=20, + full_samples=60, + max_force=0.5, # eV/A + update_frequency=1, # every uncorrelated sample helps + name="abf", +) + +# %% +# Run +# --- +# The runner hands ABF its observation at ``AFTER_COMPUTE``, where +# ``batch.forces`` still holds the **unbiased** physical force. This is not +# a detail: an estimator shown its own output converges to whatever it had +# already decided, and the resulting profile looks perfectly smooth. + +model = LennardJonesModelWrapper(sigma=3.4, epsilon=0.0104, cutoff=8.5).to(DEVICE) +dynamics = NVTLangevin(model=model, dt=0.5, temperature=TEMPERATURE, friction=0.05) + +for hook in model.make_neighbor_hooks(): + dynamics.register_hook(hook) + +sampling = EnhancedSampling(dynamics=dynamics, biases={"abf": abf}) +batch = sampling.run(batch, n_steps=N_STEPS) + +logger.info("") +logger.info( + "Samples: %d across %d of %d bins", + int(abf.bin_counts.sum()), + int((abf.bin_counts > 0).sum()), + abf.n_bins, +) +logger.info( + "Bins past the threshold (applying force): %d", + int((abf.bin_counts > abf.min_samples).sum()), +) + +# %% +# ABF adds force but no energy +# ---------------------------- +# ``last_outputs`` carries no ``bias/abf/energy`` key, and the total energy +# equals the physical energy exactly. That is not an omission — the applied +# force is genuinely not the gradient of any function the bias holds, which +# is what "non-conservative" means here. + +outputs = sampling.last_outputs +logger.info("") +logger.info("Diagnostics: %s", sorted(k for k in outputs if k.startswith("bias/abf/"))) +logger.info("Has a bias energy: %s", "bias/abf/energy" in outputs) +logger.info( + "total/energy == physical/energy: %s", + bool(torch.allclose(outputs["total/energy"], outputs["physical/energy"])), +) +logger.info( + "Applied bias force magnitude: %.3e eV/A", + float(outputs["bias/abf/forces"].abs().max()), +) + +# %% +# The free-energy profile +# ----------------------- +# Integrating the mean force gives the PMF directly. Bins that were never +# visited come back as ``nan`` rather than zero — a bin with no samples has +# no estimate, and zero is a perfectly plausible free energy that would hide +# that. +# +# A short demo run will not have sampled contiguously, and ``free_energy()`` +# **raises** on an interior gap rather than integrating across it: the +# profile beyond a hole would be wrong by an unknown constant. + +try: + profile = abf.free_energy() + centers = abf.bin_centers + logger.info("") + logger.info("Free-energy profile (eV, shifted to zero minimum):") + for r, value in zip(centers.tolist(), profile.tolist(), strict=True): + if not math.isnan(value): + logger.info(" r = %.2f A A = %+.4f eV", r, value) + + # A free check that the estimator is not fooling itself: the PMF minimum + # should sit near the Lennard-Jones minimum, 2^(1/6) * sigma. + sampled = ~torch.isnan(profile) + minimum = float(centers[sampled][profile[sampled].argmin()]) + logger.info( + "PMF minimum at %.2f A; Lennard-Jones minimum 2^(1/6) sigma = %.2f A", + minimum, + 2 ** (1 / 6) * 3.4, + ) +except RuntimeError as error: + logger.info("") + logger.info("free_energy() declined, as it should on a short run:") + logger.info(" %s", error) + +# %% +# Why ABF cannot join a replica-exchange ladder +# --------------------------------------------- +# The Metropolis acceptance rule needs each bias's energy evaluated under +# *both* states being swapped. A force-only bias has no such energy, so the +# rule cannot be formed. Rather than silently drop the bias from the +# exponent — which would break detailed balance with nothing to show for it +# — the combination is refused at construction. + +logger.info("") +logger.info("supplies_exchange_energy: %s", abf.supplies_exchange_energy) +logger.info("A ReplicaExchange ladder will refuse this bias at construction.") diff --git a/examples/advanced/README.rst b/examples/advanced/README.rst index 1a7decd6..34e59edb 100644 --- a/examples/advanced/README.rst +++ b/examples/advanced/README.rst @@ -33,3 +33,26 @@ molecules via task selection on ``UMAWrapper.from_checkpoint``. **10 — MACE Training**: Training a ScaleShiftMACE model with the ALCHEMI training stack; Zarr dataloading, scheduled Huber losses, EMA, checkpointing, validation, and distributed launch patterns. + +**11 — Umbrella Sampling**: Batched umbrella sampling with +``EnhancedSampling``; per-window centers selected by +``thermodynamic_state_id``, composing a restraint with a wall, and reading +per-bias diagnostics for WHAM/MBAR. + +**12 — Replica Exchange**: Temperature REMD over a geometric ladder; the +even/odd pair schedule, Metropolis acceptance, per-pair acceptance rates for +ladder tuning, and confirming the integrator target follows the assignment. + +**13 — Metadynamics**: Multiple-walker well-tempered metadynamics along a +pair-distance CV; shared hill history across walkers, storage-policy choice, +the well-tempered height decay, and free-energy reconstruction. + +**14 — RMSD Metadynamics**: xTB-style structure exploration with no +collective variable; optimal-alignment RMSD over retained references, FIFO +retention, and an unbiased control run to separate exploration from thermal +scatter. + +**15 — Adaptive Biasing Force**: Measuring and cancelling the mean force +along a pair distance; the metric correction demonstrated against an analytic +answer, sample thresholds and force ramps, and why a force-only bias cannot +join a replica-exchange ladder. diff --git a/nvalchemi/dynamics/base.py b/nvalchemi/dynamics/base.py index 410dd0d0..7b6afc22 100644 --- a/nvalchemi/dynamics/base.py +++ b/nvalchemi/dynamics/base.py @@ -51,7 +51,7 @@ import sys import warnings from collections import OrderedDict -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from enum import Enum from typing import ( TYPE_CHECKING, @@ -1657,6 +1657,179 @@ def _make_new_state(self, n: int, template_batch: Batch) -> "Batch | None": """ return None + def state_dict(self) -> dict[str, Any]: + """Return the integrator state needed to resume this run exactly. + + Covers the step counter, the RNG seed, and every per-system tensor in + ``self._state`` — thermostat chain variables, per-system timesteps, + barostat auxiliaries, whatever the subclass put there. + + Nothing here is a Python object graph: values are tensors and + scalars, so the result is directly representable in Zarr without a + pickle payload. + + Returns + ------- + dict[str, Any] + ``step_count``, ``random_seed`` (when the integrator has one), + and a ``state`` submapping of per-system tensors. ``state`` is + empty for an integrator that keeps none, and for one whose lazy + initialisation has not run yet. + + Notes + ----- + Reproducibility of stochastic integrators + ``NVTLangevin`` derives its noise from ``random_seed + + step_count`` rather than advancing a stateful generator, so + restoring those two integers reproduces the identical noise + sequence. There is no generator object to serialise. + """ + state: dict[str, torch.Tensor] = {} + internal = getattr(self, "_state", None) + if internal is not None: + # Batch yields (name, value) pairs from __iter__; it has no .items(). + for key, value in internal: + if isinstance(value, torch.Tensor): + state[key] = value.detach().clone() + + out: dict[str, Any] = { + "step_count": int(self.step_count), + "state": state, + } + seed = getattr(self, "_random_seed", None) + if seed is not None: + out["random_seed"] = int(seed) + return out + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore integrator state produced by :meth:`state_dict`. + + Parameters + ---------- + state: + The mapping previously returned by :meth:`state_dict`. + + Raises + ------ + RuntimeError + If the checkpoint carries per-system state but this integrator + has not initialised its own — restoring into an uninitialised + integrator would leave the two silently out of step. Prime the + dynamics once (or run a step) before restoring. + KeyError + If a restored key is absent from the live state, which means the + checkpoint came from a differently-configured integrator. + """ + self.step_count = int(state.get("step_count", 0)) + if "random_seed" in state and hasattr(self, "_random_seed"): + self._random_seed = int(state["random_seed"]) + + saved: Mapping[str, torch.Tensor] = state.get("state", {}) or {} + if not saved: + return + + internal = getattr(self, "_state", None) + if internal is None: + raise RuntimeError( + f"{type(self).__name__}.load_state_dict: the checkpoint holds " + f"per-system integrator state ({sorted(saved)}) but this " + "instance has not initialised its own yet. Prime the dynamics " + "against the restored batch before loading, so the shapes are " + "known." + ) + for key, value in saved.items(): + target = getattr(internal, key, None) + if target is None: + raise KeyError( + f"{type(self).__name__}.load_state_dict: checkpoint key " + f"{key!r} has no counterpart in the live integrator state " + f"({sorted(k for k, _ in internal)}). The " + "checkpoint was written by a differently-configured " + "integrator." + ) + with torch.no_grad(): + target.copy_(value.reshape(target.shape).to(target.device)) + + def redistribute_state(self, walker_ids: torch.Tensor) -> None: + """Reorder per-system state to match a new walker layout. + + Needed only when a checkpoint is restored into a batch whose rows are + ordered differently from the one it was written from. Batch position + is not an identity, so the caller supplies the permutation. + + Parameters + ---------- + walker_ids : torch.Tensor + Row permutation, shape ``[B]``: entry *i* is the index in the + current state that should become row *i*. + + Raises + ------ + RuntimeError + If the integrator has no per-system state to reorder. + """ + internal = getattr(self, "_state", None) + if internal is None: + raise RuntimeError( + f"{type(self).__name__}.redistribute_state: no per-system " + "state to reorder." + ) + # Index on the state's own device: ``self.device`` reports the + # process compute device, which is CUDA whenever a GPU is visible even + # for state that was never moved off the CPU. + index = walker_ids.reshape(-1).to(dtype=torch.long) + for key, value in list(internal): + if isinstance(value, torch.Tensor) and value.shape[0] == index.numel(): + internal[key] = value[index.to(value.device)].contiguous() + + def apply_thermodynamic_state( + self, state_ids: torch.Tensor, temperatures: torch.Tensor + ) -> None: + """Rebind each walker to a new thermodynamic state. + + Used by replica exchange after an accepted swap: the walker stays on + its execution slot while its assigned temperature changes. The + change must be indivisible — target temperature, velocity scaling, + and any thermostat private state have to move together, or detailed + balance is broken. + + Parameters + ---------- + state_ids : torch.Tensor + New state id per graph, shape ``[B]``. + temperatures : torch.Tensor + Temperature in Kelvin per state, shape ``[S]``. + + Raises + ------ + NotImplementedError + Always, on the base class. An integrator that cannot rebind + must fail rather than silently accept a label-only swap that + leaves its velocities and thermostat at the old temperature. + """ + raise NotImplementedError( + f"{type(self).__name__} does not support thermodynamic-state " + "rebinding. Temperature replica exchange requires an integrator " + "that can rescale velocities and transform its thermostat state; " + "NVTLangevin and NVTNoseHoover implement this." + ) + + def _rescale_velocities(self, scale_per_graph: torch.Tensor, batch: Batch) -> None: + """Scale each graph's velocities by a per-graph factor, in place. + + Parameters + ---------- + scale_per_graph : torch.Tensor + Multiplicative factor per graph, shape ``[B]``. + batch : Batch + The live batch; ``velocities`` is modified in place. + """ + velocities = getattr(batch, "velocities", None) + if velocities is None: + return + with torch.no_grad(): + velocities.mul_(scale_per_graph[batch.batch_idx].unsqueeze(-1)) + def _ensure_state_initialized(self, batch: Batch) -> None: """Lazily initialize per-system integrator state on the first call. diff --git a/nvalchemi/dynamics/integrators/nvt_langevin.py b/nvalchemi/dynamics/integrators/nvt_langevin.py index 4e111795..fd7318ad 100644 --- a/nvalchemi/dynamics/integrators/nvt_langevin.py +++ b/nvalchemi/dynamics/integrators/nvt_langevin.py @@ -132,6 +132,73 @@ def _init_state(self, batch: Batch) -> None: # (e.g. after an inflight-batching refill). self._batch_int32: torch.Tensor = batch.batch_idx.int() + def apply_thermodynamic_state( + self, state_ids: torch.Tensor, temperatures: torch.Tensor + ) -> None: + """Rebind each walker's target temperature and rescale its velocities. + + Langevin has no thermostat memory to transform — its noise amplitude + is read from the per-system ``temperature`` every step — so the + rebinding is the target update plus the velocity rescaling that keeps + the kinetic energy consistent with the new target. + + Velocities scale by ``sqrt(T_new / T_old)``, applied to the batch + recorded by the last :meth:`_init_state` or step. + + Parameters + ---------- + state_ids : torch.Tensor + New state id per graph, shape ``[B]``. + temperatures : torch.Tensor + Temperature in Kelvin per state, shape ``[S]``. + + Raises + ------ + RuntimeError + If called before the integrator state exists. + IndexError + If a state id has no corresponding temperature. + """ + state = getattr(self, "_state", None) + if state is None: + raise RuntimeError( + "NVTLangevin.apply_thermodynamic_state: integrator state is " + "not initialised; run or prime the dynamics first." + ) + index = state_ids.reshape(-1).to(device=state.temperature.device) + table = temperatures.reshape(-1).to( + device=state.temperature.device, dtype=state.temperature.dtype + ) + if bool(((index < 0) | (index >= table.numel())).any()): + raise IndexError( + f"NVTLangevin.apply_thermodynamic_state: state id out of range " + f"for {table.numel()} temperature(s): {index.tolist()}." + ) + + old_kT = state.temperature.reshape(-1).clone() + new_kT = table[index.to(torch.long)] * KB_EV + with torch.no_grad(): + state.temperature.copy_(new_kT.reshape(state.temperature.shape)) + self._pending_velocity_scale = torch.sqrt(new_kT / old_kT) + + def rescale_velocities_for_state(self, batch: Batch) -> None: + """Apply the velocity scaling queued by the last rebinding. + + Kept separate because the scaling needs the live batch, which + :meth:`apply_thermodynamic_state` does not receive — the two together + are the indivisible change replica exchange requires. + + Parameters + ---------- + batch : Batch + The live batch; ``velocities`` is modified in place. + """ + scale = getattr(self, "_pending_velocity_scale", None) + if scale is None: + return + self._rescale_velocities(scale, batch) + self._pending_velocity_scale = None + def _make_new_state(self, n: int, template_batch: Batch) -> Batch: dev = template_batch.device dtype = template_batch.positions.dtype diff --git a/nvalchemi/dynamics/integrators/nvt_nose_hoover.py b/nvalchemi/dynamics/integrators/nvt_nose_hoover.py index e6100d0c..901ea47b 100644 --- a/nvalchemi/dynamics/integrators/nvt_nose_hoover.py +++ b/nvalchemi/dynamics/integrators/nvt_nose_hoover.py @@ -173,6 +173,80 @@ def _init_state(self, batch: Batch) -> None: dev, ) + def apply_thermodynamic_state( + self, state_ids: torch.Tensor, temperatures: torch.Tensor + ) -> None: + """Rebind target temperature, chain masses, and velocity scaling. + + Unlike Langevin, a Nosé-Hoover chain carries memory: the chain masses + ``Q`` are proportional to ``kT tau^2`` and the chain velocities + ``eta_dot`` are conjugate to them. Rebinding the target without + transforming both leaves the thermostat driving toward the old + temperature and breaks detailed balance, so all three move together: + + * ``temperature`` takes the new ``kT``. + * ``Q`` scales by ``kT_new / kT_old``, preserving ``Q ∝ kT tau^2``. + * ``eta_dot`` scales by ``sqrt(kT_old / kT_new)``, which keeps the + chain kinetic energy ``Q eta_dot^2 / 2`` invariant under the mass + change rather than injecting or removing thermostat energy. + * atomic velocities scale by ``sqrt(kT_new / kT_old)``, queued for + :meth:`rescale_velocities_for_state`. + + ``eta`` itself is a position-like variable and is left untouched. + + Parameters + ---------- + state_ids : torch.Tensor + New state id per graph, shape ``[B]``. + temperatures : torch.Tensor + Temperature in Kelvin per state, shape ``[S]``. + + Raises + ------ + RuntimeError + If called before the integrator state exists. + IndexError + If a state id has no corresponding temperature. + """ + state = getattr(self, "_state", None) + if state is None: + raise RuntimeError( + "NVTNoseHoover.apply_thermodynamic_state: integrator state is " + "not initialised; run or prime the dynamics first." + ) + index = state_ids.reshape(-1).to(device=state.temperature.device) + table = temperatures.reshape(-1).to( + device=state.temperature.device, dtype=state.temperature.dtype + ) + if bool(((index < 0) | (index >= table.numel())).any()): + raise IndexError( + f"NVTNoseHoover.apply_thermodynamic_state: state id out of " + f"range for {table.numel()} temperature(s): {index.tolist()}." + ) + + old_kT = state.temperature.reshape(-1).clone() + new_kT = table[index.to(torch.long)] * KB_EV + ratio = new_kT / old_kT + with torch.no_grad(): + state.temperature.copy_(new_kT.reshape(state.temperature.shape)) + state.nhc_Q.mul_(ratio.reshape(-1, 1)) + state.nhc_eta_dot.mul_(torch.rsqrt(ratio).reshape(-1, 1)) + self._pending_velocity_scale = torch.sqrt(ratio) + + def rescale_velocities_for_state(self, batch: Batch) -> None: + """Apply the velocity scaling queued by the last rebinding. + + Parameters + ---------- + batch : Batch + The live batch; ``velocities`` is modified in place. + """ + scale = getattr(self, "_pending_velocity_scale", None) + if scale is None: + return + self._rescale_velocities(scale, batch) + self._pending_velocity_scale = None + def _make_new_state(self, n: int, template_batch: Batch) -> Batch: dev = template_batch.device dtype = template_batch.positions.dtype diff --git a/nvalchemi/enhanced_sampling/__init__.py b/nvalchemi/enhanced_sampling/__init__.py new file mode 100644 index 00000000..e4ec1e0d --- /dev/null +++ b/nvalchemi/enhanced_sampling/__init__.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Enhanced-sampling subpackage for nvalchemi-toolkit. + +Public surface +-------------- +* :class:`BiasResult` — frozen dataclass; fully-detached bias outputs. +* :class:`BiasPotential` — ``@runtime_checkable`` Protocol; structural + interface every bias must satisfy. +* :class:`ConservativeBias` — autograd helper; subclass and override + :meth:`~ConservativeBias.energy` to get forces and tensile-positive + Cauchy stress for free. +* :class:`AdaptivePotentialMixin` — battery for biases whose state evolves + during sampling; supplies ``update`` / ``commit_epoch`` / state versioning. +* :func:`aggregate_bias_results` — sums a list of ``BiasResult`` objects. +* :func:`pair_distance` — differentiable pair-distance CV; supports + nonperiodic and Minkowski-reduced triclinic MIC. General triclinic MIC + (unreduced cells via LLL) is not yet implemented. +* :func:`periodic_difference` — CV differences wrapped onto a circle. +* :class:`EnhancedSampling` — the runner: walker identity, force-step + ordering, exactly-once ``update()``, and force priming. +* :class:`ThermodynamicState`, :class:`ReplicaExchange` — synchronous + replica exchange; swaps state labels, not coordinates. +* Built-in static biases: :class:`HarmonicUmbrellaBias`, + :class:`UpperWall`, :class:`LowerWall`, :class:`FlatBottomRestraint`. +* :class:`WellTemperedMetaDynamicsBias` — Gaussian hills along a chosen CV, + with well-tempered height damping and a free-energy estimator. +* :class:`RMSDMetaDynamicsBias` — xTB-style repulsion from retained + reference geometries, for exploring when the interesting coordinates are + not known in advance. Non-periodic systems only. + +* :class:`AdaptiveBiasingForce` — measures and cancels the mean force along + a pair distance, including the metric correction that a naive Cartesian + projection omits. Force-only, so it is excluded from replica exchange. + +Not yet implemented +------------------- +* Asynchronous replica exchange +* General triclinic MIC for unreduced cells +* Replica exchange over a temperature ladder combined with a bias whose + energy depends on the thermodynamic state (per-state umbrella windows, + per-state metadynamics history); such biases declare + ``state_dependent_for_exchange`` and the runner rejects the combination + rather than applying an acceptance rule that does not cover it. + +Relationship to ``BiasedPotentialHook`` +--------------------------------------- +:class:`~nvalchemi.hooks.BiasedPotentialHook` covers the same ground with a +narrower contract and is **deprecated** in favour of this subpackage. With +:class:`EnhancedSampling` now available, the migration path is complete: +anything the hook can do, this subpackage does, and it carries a cell +response the hook has no slot for. The hook remains functional so existing +code keeps working; no removal date is set. + +============================ ============================== ========================================== +Concern ``BiasedPotentialHook`` ``enhanced_sampling`` +============================ ============================== ========================================== +Contract ``bias_fn(batch) -> (E, F)`` ``BiasPotential.evaluate -> BiasResult`` +Forces written by hand autograd, from one energy definition +Cell response none symmetric-strain ``stress`` +Composing several biases in-place, sequential summed against unmodified model output +Diagnostics none namespaced ``observables`` +Evolving bias state closure-held, ad hoc ``update()`` exactly once per due step +============================ ============================== ========================================== + +Which to use + :class:`ConservativeBias` (or :class:`BiasPotential` directly), run + through :class:`EnhancedSampling`, for everything new. The cell + response is the substantive difference: a ``bias_fn`` bias contributes + no stress, so under NPT/NPH the barostat reads a ``batch.stress`` the + bias never touched and the cell evolves as if the bias were absent — + with no error raised. Existing hook-based code is correct under NVE and + NVT, where nothing reads the stress, and can be migrated when convenient + rather than urgently. + +No adapter is provided + Bridging a :class:`BiasPotential` onto ``bias_fn`` would have to drop + :attr:`BiasResult.stress` on the floor, since the hook has nowhere to + put it — reintroducing the exact failure the new API exists to remove. + A silent adapter would be worse than none. +""" + +from nvalchemi.enhanced_sampling._adaptive import AdaptivePotentialMixin +from nvalchemi.enhanced_sampling._bias import ( + BiasPotential, + BiasResult, + ConservativeBias, + aggregate_bias_results, +) +from nvalchemi.enhanced_sampling._exchange import ( + ReplicaExchange, + ThermodynamicState, +) +from nvalchemi.enhanced_sampling._runner import EnhancedSampling +from nvalchemi.enhanced_sampling.biases import ( + AdaptiveBiasingForce, + FlatBottomRestraint, + HarmonicUmbrellaBias, + LowerWall, + RMSDMetaDynamicsBias, + UpperWall, + WellTemperedMetaDynamicsBias, +) +from nvalchemi.enhanced_sampling.cv import ( + pair_displacement, + pair_distance, + periodic_difference, +) + +__all__ = [ + # Core abstractions + "BiasResult", + "BiasPotential", + "ConservativeBias", + "AdaptivePotentialMixin", + "aggregate_bias_results", + # Runner + "EnhancedSampling", + # Replica exchange + "ThermodynamicState", + "ReplicaExchange", + # Collective variables + "pair_distance", + "pair_displacement", + "periodic_difference", + # Built-in biases + "HarmonicUmbrellaBias", + "UpperWall", + "LowerWall", + "FlatBottomRestraint", + "WellTemperedMetaDynamicsBias", + "RMSDMetaDynamicsBias", + "AdaptiveBiasingForce", +] diff --git a/nvalchemi/enhanced_sampling/_adaptive.py b/nvalchemi/enhanced_sampling/_adaptive.py new file mode 100644 index 00000000..7c91073e --- /dev/null +++ b/nvalchemi/enhanced_sampling/_adaptive.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Battery for biases whose internal state evolves during sampling. + +``AdaptivePotentialMixin`` is one of the composable mixins described in the +``BiasPotential`` docstring: a bias mixes in only what applies to it. It +carries no energy, no forces, and no model half — a bias that is adaptive +*and* conservative mixes in both this and +:class:`~nvalchemi.enhanced_sampling.ConservativeBias`; a bias that is +adaptive and non-conservative (ABF) mixes in only this one. + +The mixin exists to make one guarantee enforceable by the runner: **evaluation +never mutates state**. ``evaluate()`` is read-only and compile-friendly; +every history-dependent change happens in :meth:`update`, which the runner +calls exactly once per due step, after the integration step has finished. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Any + +from torch import nn + +from nvalchemi.dynamics.base import DynamicsStage + +if TYPE_CHECKING: + from collections.abc import Mapping + + from nvalchemi.data import Batch + from nvalchemi.enhanced_sampling._bias import BiasResult + +__all__ = ["AdaptivePotentialMixin"] + + +def _values_agree(left: Any, right: Any) -> bool: + """Return whether two fingerprint entries describe the same setting. + + Floats are compared with a relative tolerance so that storing a value + through float32 and reading it back does not read as a configuration + change, while a real difference (0.2 against 0.9) still does. + + Parameters + ---------- + left, right: + Fingerprint entries: scalars, strings, ``None``, or flat sequences. + + Returns + ------- + bool + ``True`` when the two agree. + """ + if isinstance(left, (list, tuple)) and isinstance(right, (list, tuple)): + return len(left) == len(right) and all( + _values_agree(a, b) for a, b in zip(left, right, strict=True) + ) + if isinstance(left, bool) or isinstance(right, bool): + return left is right + if isinstance(left, (int, float)) and isinstance(right, (int, float)): + return math.isclose(float(left), float(right), rel_tol=1e-6, abs_tol=1e-12) + return bool(left == right) + + +class AdaptivePotentialMixin: + """Mixin for a bias whose state changes as sampling proceeds. + + Mix in alongside whatever else applies, **this mixin first**:: + + class WellTemperedMetaDynamicsBias(AdaptivePotentialMixin, ConservativeBias): + ... # conservative and adaptive + + class AdaptiveBiasingForce(AdaptivePotentialMixin): + ... # adaptive, but no energy to differentiate + + Order matters and is enforced. :class:`ConservativeBias` inherits + ``nn.Module``, which already defines ``state_dict`` / ``load_state_dict`` + for buffers; putting this mixin second would let those shadow the ones + here and silently drop the bias history from every checkpoint. With the + mixin first, :meth:`state_dict` calls up the MRO and merges both. + ``__init_subclass__`` raises ``TypeError`` on the wrong order rather than + letting a checkpoint quietly lose data. + + The runner detects the capability with ``hasattr(bias, "update")``; there + is no registration and no requirement to inherit this class. A bias may + implement ``update`` structurally instead. + + Attributes + ---------- + update_frequency: + Dynamics steps between :meth:`update` calls. ``1`` means every step. + observation_stage: + Which stage the frame handed to :meth:`update` is captured at. + + * ``AFTER_STEP`` (default) — post-step coordinates. What + metadynamics wants: a hill is deposited at the configuration the + system actually reached. + * ``AFTER_COMPUTE`` — captured while ``batch.forces`` still holds the + **unbiased** physical forces, before any bias contribution is + added. What ABF requires; observing biased forces would feed the + estimator its own output. + + Notes + ----- + State version + :meth:`bump_state_version` records that the bias changed. The runner + reads :attr:`state_version` to decide whether forces in the batch are + stale and need re-priming. It is also checkpointed per bias, ready + for validating that an accepted replica-exchange state assignment is + coherent — which the exchange does not consume yet. A bias that + mutates state inside :meth:`update` must call it. + """ + + update_frequency: int = 1 + observation_stage: DynamicsStage = DynamicsStage.AFTER_STEP + + # Incremented by bump_state_version(); never reset. + _state_version: int = 0 + + def __init_subclass__(cls, **kwargs: Any) -> None: + """Reject an MRO that would let ``nn.Module`` shadow this mixin. + + Parameters + ---------- + **kwargs: + Forwarded cooperatively up the MRO. + + Raises + ------ + TypeError + If ``nn.Module`` precedes this mixin, which would make + :meth:`state_dict` unreachable and drop bias history from + checkpoints without any error. + """ + super().__init_subclass__(**kwargs) + mro = cls.__mro__ + if nn.Module in mro and mro.index(AdaptivePotentialMixin) > mro.index( + nn.Module + ): + raise TypeError( + f"{cls.__name__}: AdaptivePotentialMixin must come before " + f"nn.Module (and therefore before ConservativeBias) in the " + f"base list, otherwise nn.Module.state_dict shadows the " + f"mixin's and the bias history is silently dropped from " + f"checkpoints. Write " + f"'class {cls.__name__}(AdaptivePotentialMixin, ConservativeBias)'." + ) + + @property + def state_version(self) -> int: + """Monotonic counter of state-changing updates applied to this bias.""" + return self._state_version + + def bump_state_version(self) -> None: + """Record that the bias state changed. + + Call from :meth:`update` whenever the change affects the energy the + next :meth:`evaluate` will return. An update that only accumulates + statistics without changing the applied bias (an ABF bin below its + minimum-sample threshold, say) should **not** bump — bumping forces + the runner to re-prime forces for no reason. + """ + self._state_version += 1 + + def update(self, frames: Batch, result: BiasResult) -> None: + """Consume a captured frame after the integration step finishes. + + Called by the runner exactly once per due step. Free to mutate + state, allocate, grow storage, and communicate — none of this is on + the compiled path. + + Parameters + ---------- + frames: + The captured observation, stamped by the runner with + ``walker_id``, ``thermodynamic_state_id``, and ``sampling_step``. + Captured at :attr:`observation_stage`. + result: + The detached :class:`BiasResult` this bias returned during the + preceding force evaluation. + + Raises + ------ + NotImplementedError + If the subclass does not override. + """ + raise NotImplementedError( + f"{type(self).__name__} mixes in AdaptivePotentialMixin but does " + "not implement update(frames, result)." + ) + + def commit_epoch(self) -> None: + """Synchronise pending state at a consistency-epoch boundary. + + Default is a no-op: a bias whose history is local to one walker needs + no synchronisation. Multi-walker shared-history biases override this + to publish and merge. Called only at epoch boundaries, never on the + hot path. + """ + return None + + def config_fingerprint(self) -> Mapping[str, Any]: + """Return the constructor configuration the saved state depends on. + + Default is empty, which disables the check. A bias whose persisted + state is only meaningful under the settings that produced it should + override this and list them. + + The distinction is between *state* and *configuration*. An ABF + histogram is state; the ``cv_range`` that decides what its bins mean + is configuration. Restoring the first without the second silently + reinterprets every bin — bin 5 stops meaning ``r = 1.55`` and starts + meaning ``r = 3.1``, with the same numbers in it. + + Include configuration held as a **buffer** too. ``nn.Module`` + restores buffers by overwriting, so without the check a mismatched + setting is not merely unvalidated: the caller's value is silently + replaced by the checkpoint's, which is the opposite of what asking + for it meant. + + Returns + ------- + Mapping[str, Any] + Zarr-representable scalars, strings, ``None``, or flat sequences + of those. Tensors must be converted to lists. + """ + return {} + + def _check_config_fingerprint(self, saved: Mapping[str, Any]) -> None: + """Reject a saved fingerprint that disagrees with this bias. + + Parameters + ---------- + saved: + The fingerprint recorded when the state was written. + + Raises + ------ + ValueError + If any recorded setting differs from the live one. + """ + live = dict(self.config_fingerprint()) + differences = [] + for key in sorted(set(saved) | set(live)): + was, now = saved.get(key, ""), live.get(key, "") + if not _values_agree(was, now): + differences.append(f" {key}: checkpoint {was!r} vs bias {now!r}") + if differences: + raise ValueError( + f"{type(self).__name__} " + f"{getattr(self, 'name', '')!r}: the saved state was " + f"produced under a different configuration:\n" + + "\n".join(differences) + + "\nThe stored state is only meaningful under the settings " + "that produced it, so restoring it here would reinterpret it " + "rather than continue it. Rebuild the bias with the " + "checkpoint's configuration, or start a fresh run." + ) + + def state_dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + """Return the bias state for checkpointing, merged up the MRO. + + Cooperative: when mixed into an ``nn.Module`` bias this calls + ``nn.Module.state_dict`` first, so buffers (umbrella centers, wall + thresholds) and bias history land in one mapping. A bias with real + history (hills, reference frames, ABF bins) overrides this, calls + ``super().state_dict()``, and adds its own entries. + + Parameters + ---------- + *args, **kwargs: + Forwarded to the next ``state_dict`` in the MRO, if any. + + Returns + ------- + dict[str, Any] + Zarr-representable state, always including ``state_version`` and, + when :meth:`config_fingerprint` is overridden, the configuration + that state is only meaningful under. + """ + parent = getattr(super(), "state_dict", None) + state: dict[str, Any] = dict(parent(*args, **kwargs)) if parent else {} + state["state_version"] = self._state_version + fingerprint = dict(self.config_fingerprint()) + if fingerprint: + state["config_fingerprint"] = fingerprint + return state + + def load_state_dict( + self, state: Mapping[str, Any], *args: Any, **kwargs: Any + ) -> Any: + """Restore bias state produced by :meth:`state_dict`. + + Strips the mixin's own keys before delegating, so that + ``nn.Module.load_state_dict`` does not reject them as unexpected + under its default ``strict=True``. A recorded + :meth:`config_fingerprint` is checked *before* delegating, since the + delegate overwrites buffers and a check afterwards would come too + late to stop the caller's configuration being replaced. + + Parameters + ---------- + state: + The mapping previously returned by :meth:`state_dict`. + *args, **kwargs: + Forwarded to the next ``load_state_dict`` in the MRO, if any. + + Returns + ------- + Any + Whatever the next ``load_state_dict`` returns, or ``None``. + + Raises + ------ + ValueError + If the state was written under a different + :meth:`config_fingerprint`. + """ + remaining = dict(state) + self._state_version = int(remaining.pop("state_version", 0)) + saved = remaining.pop("config_fingerprint", None) + if saved is not None: + # Before delegating: nn.Module.load_state_dict overwrites buffers, + # so a mismatch caught afterwards would already have replaced the + # caller's configuration with the checkpoint's. + self._check_config_fingerprint(saved) + parent = getattr(super(), "load_state_dict", None) + if parent is not None: + return parent(remaining, *args, **kwargs) + return None diff --git a/nvalchemi/enhanced_sampling/_bias.py b/nvalchemi/enhanced_sampling/_bias.py new file mode 100644 index 00000000..1f07b20e --- /dev/null +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -0,0 +1,912 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Core bias abstractions: ``BiasPotential`` protocol, ``BiasResult``, and +``ConservativeBias`` autograd helper. + +This module is the foundation of the enhanced-sampling subpackage. Every +downstream built-in bias depends on these three objects. + +Design guarantees +----------------- +* ``BiasResult`` is a frozen dataclass; all tensor fields are detached + (``requires_grad=False``, ``grad_fn is None``). Validation is enforced + in eager mode; the check is skipped inside ``torch.compile`` to avoid + graph breaks on attribute inspection. +* ``BiasPotential`` is a ``@runtime_checkable`` Protocol. Bias authors + may satisfy it structurally without inheriting from any base class. +* ``ConservativeBias`` encapsulates the autograd subgraph that derives + atomic forces and the tensile-positive Cauchy stress from a scalar + energy function. The subgraph is isolated from the live ``Batch`` so + that no ``requires_grad`` leaf ever escapes into model state, batch + storage, or ``BiasResult``. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +import torch +from torch import Tensor, nn + +from nvalchemi.models._utils import ( + autograd_forces, + autograd_forces_and_stresses, + prepare_strain, + sum_outputs, +) +from nvalchemi.models.base import BaseModelMixin, ModelConfig + +if TYPE_CHECKING: + from pathlib import Path + + from nvalchemi._typing import ModelOutputs + from nvalchemi.data import AtomicData, Batch + from nvalchemi.distributed.config import StrategyKind + from nvalchemi.distributed.spec import MLIPSpec + +__all__ = [ + "BiasResult", + "BiasPotential", + "ConservativeBias", + "aggregate_bias_results", +] + +# --------------------------------------------------------------------------- +# BiasResult +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class BiasResult: + """Immutable, fully-detached output of a single bias evaluation. + + All tensor fields must be detached (``requires_grad=False`` and + ``grad_fn is None``). Energy, forces, stress, and virial are + independently optional. Provide **either** ``stress`` or ``virial``, + not both; the runner converts stress to virial or vice-versa as needed. + + Parameters + ---------- + energy: + Per-graph bias energy, shape ``[B, 1]``, unit eV. + forces: + Per-atom bias forces, shape ``[N_atoms, 3]``, unit eV/Å. + stress: + Tensile-positive Cauchy stress ``σ = −W/V``, shape ``[B, 3, 3]``, + unit eV/ų. This is the toolkit-wide convention and what + :class:`ConservativeBias` produces. Mutually exclusive with + ``virial``. + virial: + Virial ``W = −dE/dε`` with ``ε`` the symmetric infinitesimal + strain tensor, shape ``[B, 3, 3]``, unit eV. Provided for biases + that compute a virial directly. Mutually exclusive with + ``stress``. + state_version: + Integer version IDs used by ``ReplicaExchange`` to validate that + accepted state assignments are coherent, shape ``[B]``. + observables: + Named diagnostic tensors exposed as ``bias//`` in the + runner's output dict. All tensors must be detached. + + Notes + ----- + Why the physics fields are a closed set + :data:`~nvalchemi._typing.ModelOutputs` is an open mapping, and it is + reasonable to ask why the type that extends it is narrower. The + asymmetry is deliberate, and it follows from what the two are for. + + A model's outputs are its *product*: the caller decides what a + Hessian or a dipole is for, and a key the caller does not recognise + costs nothing. A bias's outputs are a *contribution* the runner must + add into a specific batch buffer. For each one the runner has to + know the destination (``batch.energy``, ``batch.forces``, + ``batch.stress``), whether it is per-graph or per-atom so it can + reshape, how it combines across biases, and how it converts — + stress and virial are the same physics in two conventions, and + moving between them needs the cell volume. An unrecognised key has + none of that, so ``EnhancedSampling._check_destinations`` raises for + any produced output with no buffer to receive it. An open payload + would therefore be open only up to the first key the runner could + not apply. + + The extension point that *is* open is :attr:`observables`: an + arbitrary ``Mapping[str, Tensor]``, no shape checks, surfaced as + ``bias//``. The split is between quantities the runner + applies to the dynamics — closed, because each needs a defined + destination and combination rule — and quantities it merely reports, + which are unconstrained. A per-atom energy decomposition belongs to + the second group: it is not a contribution to ``batch.energy``. + + What this does cost is worth stating plainly. A third party whose + novel method produces a genuinely new *applied* output — one the + integrator should act on — cannot express it without editing this + module. That is a real limit, accepted because the alternative is a + contribution the runner silently drops. Two escape hatches exist + short of a framework change: report it through ``observables``, or, + for cross-rank semantics, override + :meth:`ConservativeBias.distribution_spec`. + + Three restrictions follow from the same reasoning rather than from + the field set itself: ``energy`` is ``[B, 1]`` because it is added to + a per-graph buffer; ``stress`` and ``virial`` are mutually exclusive + because carrying both invites two answers that can disagree; and + ``state_version`` is integral because + :class:`~nvalchemi.enhanced_sampling.ReplicaExchange` compares + versions for identity, which a float does not support. + """ + + energy: Tensor | None = None + forces: Tensor | None = None + stress: Tensor | None = None + virial: Tensor | None = None + state_version: Tensor | None = None + observables: Mapping[str, Tensor] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not torch.compiler.is_compiling(): + _validate_bias_result(self) + + +def _validate_bias_result(result: BiasResult) -> None: + """Eager-only validation of a ``BiasResult`` (skipped under compile). + + Checks (in order): + + 1. Mutual exclusion of ``stress`` and ``virial``. + 2. All tensor fields are detached (``requires_grad=False``, ``grad_fn is None``). + 3. Shapes match the documented conventions: + + * ``energy`` — ndim=2, shape ``[B, 1]`` + * ``forces`` — ndim=2, shape ``[N, 3]`` + * ``stress`` — ndim=3, shape ``[B, 3, 3]`` + * ``virial`` — ndim=3, shape ``[B, 3, 3]`` + * ``state_version`` — ndim=1, integer dtype + + 4. Batch-size consistency: all present system-level fields + (``energy``, ``stress``, ``virial``, ``state_version``) must agree + on the leading dimension ``B``. + 5. All floating-point tensors (including ``observables``) are finite + (no NaN or Inf). + """ + # 1. stress / virial mutual exclusion + if result.stress is not None and result.virial is not None: + raise ValueError("BiasResult: provide either 'stress' or 'virial', not both.") + + # 2. Detachment check for every tensor field + tensor_fields: dict[str, Tensor | None] = { + "energy": result.energy, + "forces": result.forces, + "stress": result.stress, + "virial": result.virial, + "state_version": result.state_version, + } + for name, t in tensor_fields.items(): + if t is None: + continue + if t.requires_grad: + raise ValueError( + f"BiasResult.{name} must be detached " + f"(requires_grad=False), got requires_grad=True." + ) + if t.grad_fn is not None: + raise ValueError( + f"BiasResult.{name} must be detached " + f"(grad_fn is None), got grad_fn={t.grad_fn}." + ) + for key, t in result.observables.items(): + if t.requires_grad: + raise ValueError( + f"BiasResult.observables[{key!r}] must be detached " + f"(requires_grad=False)." + ) + if t.grad_fn is not None: + raise ValueError( + f"BiasResult.observables[{key!r}] must be detached (grad_fn is None)." + ) + + # 3. Shape checks + if result.energy is not None: + e = result.energy + if e.ndim != 2 or e.shape[1] != 1: + raise ValueError( + f"BiasResult.energy must have shape [B, 1], got {tuple(e.shape)}." + ) + + if result.forces is not None: + f = result.forces + if f.ndim != 2 or f.shape[1] != 3: + raise ValueError( + f"BiasResult.forces must have shape [N, 3], got {tuple(f.shape)}." + ) + + for name in ("stress", "virial"): + t = getattr(result, name) + if t is not None and (t.ndim != 3 or t.shape[1] != 3 or t.shape[2] != 3): + raise ValueError( + f"BiasResult.{name} must have shape [B, 3, 3], got {tuple(t.shape)}." + ) + + if result.state_version is not None: + sv = result.state_version + if sv.ndim != 1: + raise ValueError( + f"BiasResult.state_version must have shape [B], got {tuple(sv.shape)}." + ) + if sv.dtype not in ( + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + ): + raise ValueError( + f"BiasResult.state_version must be an integer dtype, got {sv.dtype}." + ) + + # 4. Batch-size consistency across system-level fields + b_sizes: dict[str, int] = {} + for name in ("energy", "stress", "virial", "state_version"): + t = getattr(result, name) + if t is not None: + b_sizes[name] = t.shape[0] + if len(set(b_sizes.values())) > 1: + raise ValueError( + f"BiasResult: leading batch dimension B is inconsistent across fields: " + f"{b_sizes}." + ) + + # 5. Finiteness — NaN and Inf are never valid output values + for name, t in tensor_fields.items(): + if t is None or not t.is_floating_point(): + continue + if not t.isfinite().all(): + raise ValueError(f"BiasResult.{name} contains NaN or Inf values.") + for key, t in result.observables.items(): + if t.is_floating_point() and not t.isfinite().all(): + raise ValueError( + f"BiasResult.observables[{key!r}] contains NaN or Inf values." + ) + + +# --------------------------------------------------------------------------- +# BiasPotential Protocol +# --------------------------------------------------------------------------- + + +@runtime_checkable +class BiasPotential(Protocol): + """Structural protocol for all enhanced-sampling bias potentials. + + This is the **boundary**, and it deliberately inherits nothing. A third + party implementing a novel method satisfies it structurally, with no + base class, no registration, and no dependency on anything in this + package beyond :class:`BiasResult`. + + Batteries are supplied as **composable mixins that satisfy this protocol, + not as parallel hierarchies**. A bias mixes in only what applies to it: + + * :class:`ConservativeBias` — for a bias defined by a differentiable + energy. It mixes in :class:`~nvalchemi.models.base.BaseModelMixin`, + because a conservative bias genuinely is an additive potential in this + toolkit's terms (``DFTD3ModelWrapper`` is the closest existing + analogue: a pure additive energy correction with no network and no + embeddings). That inheritance is what supplies ``model_config``, + ``active_outputs``, ``distribution_spec``, and ``+`` composition. + * A non-conservative bias such as ABF applies forces directly with no + energy to differentiate, so it must **not** be forced through + :class:`ConservativeBias`. It implements this protocol directly, or + mixes in the adaptive/checkpointing batteries only. + + The rule: inheritance is opt-in per capability. Nothing here makes "a + bias is a model" true by fiat for cases where it is not. + + Attributes + ---------- + name: + Unique string identifier used as a dict key in + ``EnhancedSampling(biases={...})`` and as a Zarr group name in + checkpoints. + + Methods + ------- + evaluate(current) + Read-only evaluation. Must not mutate bias state, write to + storage, or communicate. Called every force evaluation by + default. + + Adaptive biases additionally implement ``update()``, + ``commit_epoch()``, ``state_dict()``, and ``load_state_dict()``. + These are optional extensions that the runner detects via + ``hasattr``; they are not part of this base protocol. + """ + + name: str + + def evaluate(self, current: Batch) -> BiasResult: + """Evaluate the bias on the current batch. + + Must be **read-only**: it must not mutate bias internal state, + deposit hills, write any storage, or communicate across workers. + It is safe to call ``evaluate`` multiple times on the same batch + without side effects. + + Parameters + ---------- + current: + The live ``Batch`` from the dynamics step. Treat as + read-only; do not modify any field. + + Returns + ------- + BiasResult + Fully detached outputs. All tensor fields must satisfy + ``requires_grad=False`` and ``grad_fn is None``. + """ + ... + + +# --------------------------------------------------------------------------- +# ConservativeBias — autograd helper +# --------------------------------------------------------------------------- + + +class ConservativeBias(nn.Module, BaseModelMixin): + """Autograd helper that derives atomic forces and stress from energy. + + Composed as ``nn.Module, BaseModelMixin`` — the house multiple-inheritance + idiom (``LennardJonesModelWrapper(nn.Module, BaseModelMixin)``, + ``TrainingStrategy(BaseModel, HookRegistryMixin)``). + ``BaseModelMixin.__init_subclass__`` is cooperative (it calls + ``super().__init_subclass__(**kwargs)``), so it composes correctly with + further mixins added later for adaptive or checkpointable biases. + + A conservative bias is an additive potential with no learned parameters, + which is exactly the shape of ``DFTD3ModelWrapper`` and + ``LennardJonesModelWrapper``. The cost of the abstraction is the two + ``BaseModelMixin`` abstract methods (:attr:`embedding_shapes` and + :meth:`compute_embeddings`), stubbed here the same way those two wrappers + stub them. What it buys: + + * ``model_config.active_outputs`` declares which outputs this bias + produces, replacing an ad-hoc private flag. + * :meth:`distribution_spec` gives domain decomposition a defined answer + instead of an undefined one — see that method for why the default is + deliberately ``None``. + * ``+`` composition with a model via ``PipelineModelWrapper``, and + ``state_dict``/``load_state_dict`` from ``nn.Module`` for checkpointing. + + .. note:: + + Subclasses **must** call ``super().__init__(name=...)``. This is the + ``nn.Module`` requirement (attribute assignment before + ``Module.__init__`` raises), and ``BaseModelMixin.__init_subclass__`` + additionally verifies ``self.model_config`` is set after construction. + + Subclass ``ConservativeBias`` and override :meth:`energy` to return a + differentiable per-graph bias energy ``[B, 1]``. The base class + provides :meth:`evaluate`, which: + + 1. Enters a local ``torch.enable_grad()`` region (safe inside + ``torch.no_grad()`` outer contexts). + 2. Creates a detached positions leaf ``pos_leaf`` (for forces) and, + for periodic batches, a per-graph strain leaf via + :func:`~nvalchemi.models._utils.prepare_strain` (for stress). + 3. Substitutes the strained positions and cell into the batch, calls + :meth:`energy`, and restores the original tensors unconditionally + in a ``finally`` block. + 4. Derives forces and stress in one ``autograd.grad`` call through + :func:`~nvalchemi.models._utils.autograd_forces_and_stresses`. + 5. Constructs a ``BiasResult`` from fully detached output tensors. + + The framework must never place a tensor with ``requires_grad=True`` or + a non-null ``grad_fn`` into the live ``Batch``, ``BiasResult``, + retained history, bias state, observables, or a checkpoint. + + Notes + ----- + Stress rather than virial + ``evaluate`` populates :attr:`BiasResult.stress`, never + :attr:`BiasResult.virial`. Tensile-positive Cauchy stress is the + toolkit-wide public convention: every model wrapper emits + ``"stress"``, ``sum_outputs`` treats it as additive, and the + NPT/NPH integrators read ``batch.stress``. Emitting stress lets a + bias contribution be summed directly with model outputs with no + volume conversion at the boundary. ``BiasResult.virial`` remains + available for hand-written biases that produce a virial directly. + + Strain convention + The strain leaf comes from + :func:`~nvalchemi.models._utils.prepare_strain`, which applies only + the symmetric part of the leaf as strain. The resulting gradient + is therefore symmetric by construction, matching the project + definition ``W_ab = −dE/dε_ab`` with ``ε`` the symmetric + infinitesimal strain tensor (see + :doc:`/userguide/about/conventions`). Differentiating with respect + to the full (unsymmetrised) deformation gradient instead yields an + asymmetric tensor for any bias that is not a central pair + interaction. + + Stress computation + Stress is computed only when the batch carries a cell, has at least + one periodic dimension, **and** ``"stress"`` is in + ``model_config.active_outputs``. A bias that never needs a cell + response passes ``compute_stress=False`` to :meth:`__init__`, which + drops ``"stress"`` from ``active_outputs``; callers may also flip it + at runtime the same way they would on any model wrapper. + + torch.compile compatibility + :meth:`evaluate` runs in eager mode. It uses + ``pos_leaf = positions.detach().requires_grad_(True)``, which is + not supported by ``torch.compile`` (``Unsupported + Tensor.requires_grad_() call``). The documented fallback is: + compile :meth:`energy` independently (the user's hot path); + keep :meth:`evaluate` as the eager orchestration wrapper. + ``EnhancedSampling(compile_biases=True)`` applies + ``torch.compile`` to each bias's ``energy()`` override only. + """ + + def __init__(self, name: str, *, compute_stress: bool = True) -> None: + """Initialise the bias and declare its output capabilities. + + Parameters + ---------- + name: + Unique identifier, used as a dict key in + ``EnhancedSampling(biases={...})`` and as a checkpoint group + name. Satisfies the :class:`BiasPotential` protocol's ``name``. + compute_stress: + When ``False``, ``"stress"`` is dropped from + ``model_config.active_outputs`` and the strain leaf is skipped + entirely. Use for force-only biases. + """ + super().__init__() + self.name = name + outputs = {"energy", "forces", "stress"} + self.model_config = ModelConfig( + outputs=frozenset(outputs), + autograd_outputs=frozenset({"forces", "stress"}), + autograd_inputs=frozenset({"positions", "cell"}), + supports_pbc=True, + needs_pbc=False, + active_outputs=outputs if compute_stress else {"energy", "forces"}, + ) + + # ------------------------------------------------------------------ + # BaseModelMixin required surface + # ------------------------------------------------------------------ + + @property + def embedding_shapes(self) -> dict[str, tuple[int, ...]]: + """No embeddings: a bias potential is a closed-form energy term.""" + return {} + + def compute_embeddings( + self, data: AtomicData | Batch, **kwargs: Any + ) -> AtomicData | Batch: + """Not implemented — a bias potential produces no embeddings. + + Follows the same stub as ``LennardJonesModelWrapper`` and + ``DFTD3ModelWrapper``, which are pure-physics potentials with no + learned representation. + + Parameters + ---------- + data: + The input system. + **kwargs: + Unused; accepted for interface compatibility. + + Returns + ------- + AtomicData | Batch + Never returns. + + Raises + ------ + NotImplementedError + Always. + """ + raise NotImplementedError(f"{type(self).__name__} does not produce embeddings.") + + def export_model(self, path: Path, as_state_dict: bool = False) -> None: + """Not implemented — a bias has no underlying model to export. + + Parameters + ---------- + path: + Unused. + as_state_dict: + Unused. + + Raises + ------ + NotImplementedError + Always. Use ``state_dict()`` to checkpoint bias state. + """ + raise NotImplementedError( + f"{type(self).__name__} has no exportable model; use state_dict()." + ) + + def distribution_spec( + self, strategy: StrategyKind | None = None + ) -> MLIPSpec | None: + """Return ``None``: a bias does not claim domain-decomposition support. + + ``None`` is not an oversight, and it is not the same as "unsupported + forever" — it makes ``DistributedModel`` raise ``DistributionError`` + rather than shard a bias whose cross-rank semantics are undefined + (an explicit ``DistributedModel(bias, cfg, spec=...)`` remains the + escape hatch for a caller who knows better). Before this class mixed + in ``BaseModelMixin`` there was no way to express any of this, so a + biased simulation under domain decomposition had no defined + behaviour. Now it fails loudly. + + The default cannot be a halo preset, because a bias is not + necessarily local the way a cutoff potential is. A CV can couple + atoms in different domains by construction: ``pair_distance`` over + two atoms on opposite sides of the cell has no cutoff, and an RMSD + bias reads every atom. ``SPEC_LJ_HALO`` is correct for LJ precisely + because a halo exchange covers its interaction range; nothing + guarantees that for an arbitrary CV. + + A bias that *is* local should override this and declare its outputs, + e.g.:: + + def distribution_spec(self, strategy=None): + return MLIPSpec( + distribution=DistributionSpec(policy=HaloStoragePolicy()), + outputs={ + "energy": OutputSpec(OutputKind.PER_GRAPH, Reduce.ALL_REDUCE), + "forces": OutputSpec(OutputKind.PER_NODE, Reduce.OWNED_ONLY), + "stress": OutputSpec(OutputKind.PER_GRAPH, Reduce.ALL_REDUCE), + }, + ) + + Parameters + ---------- + strategy: + Accepted for the framework contract; ignored by the default. + + Returns + ------- + MLIPSpec | None + Always ``None`` unless a subclass overrides. + """ + return None + + # ------------------------------------------------------------------ + # Bias surface + # ------------------------------------------------------------------ + + def _align_device(self, reference: Tensor) -> None: + """Move this bias's buffers to *reference*'s device if they differ. + + A bias holds its parameters as buffers (window centers, stiffness, + wall thresholds), and a user who builds the bias before moving the + batch to GPU would otherwise hit a bare "expected all tensors to be + on the same device" from inside the energy expression, naming + neither the bias nor the fix. Moving once here is cheaper than + making every bias author remember ``.to(device)``. + + Eager-only, like the rest of :meth:`evaluate`. + + Parameters + ---------- + reference: + Any tensor from the live batch; its device is the target. + """ + buffer = next(self.buffers(), None) + if buffer is not None and buffer.device != reference.device: + self.to(reference.device) + + def energy(self, current: Batch) -> Tensor: + """Return bias energy ``[B, 1]`` (eV). + + Must be differentiable w.r.t. ``current.positions`` and/or + ``current.cell``. Depending on only one of the two is allowed: a + position-independent term (a volume restraint, say) yields zero + forces rather than an error. + + Parameters + ---------- + current: + A *read-only view* of the live batch whose ``positions`` and + ``cell`` have been replaced by their strained counterparts from + :func:`~nvalchemi.models._utils.prepare_strain`. Do not assign + to any batch field inside this method. + """ + raise NotImplementedError( + f"{type(self).__name__} must implement energy(self, current: Batch) -> Tensor" + ) + + def evaluate(self, current: Batch) -> BiasResult: + """Compute energy, forces, and Cauchy stress via autograd. + + This method runs in eager mode. See the class docstring for the + compile boundary note and the chosen fallback. + + Stress derivation + ----------------- + Under a homogeneous strain ``ε`` (ASE row-vector convention), both + atomic positions and the cell deform together:: + + r_n → r_n @ (I + ε) + cell_b → cell_b @ (I + ε) + + :func:`~nvalchemi.models._utils.prepare_strain` applies exactly this + deformation through a leaf whose symmetric part is used, so a single + ``autograd.grad`` call yields: + + * ``dE/d(pos_leaf[n])`` at ``ε=0`` → forces (negated). + * ``dE/dε`` at ``ε=0`` → ``σ = (dE/dε) / V``, tensile-positive + Cauchy stress, equivalently ``σ = −W/V``. + + Straining positions and cell together is what makes this correct for + position-dependent biases built on MIC displacements: the position + term and the cell gradient term are combined automatically. A + strain leaf applied to the cell alone misses the position + contribution and gives the wrong answer for pair restraints that + span an image boundary. + + Partial dependence + ------------------ + :meth:`energy` need not depend on both positions and strain. A term + that depends only on the cell — a volume restraint, a barostat-style + term — returns zero forces and a non-zero stress; a term that returns + a constant on some branch returns zeros for both. These are ordinary + outcomes, not errors: the gradient of an energy with respect to + something it does not use is zero. + + Returns + ------- + BiasResult + With ``energy``, ``forces``, and — for periodic batches — + ``stress``. ``virial`` is always ``None``; see the class + docstring for why stress is the chosen field. + """ + original_positions = current.positions + self._align_device(original_positions) + original_cell = getattr(current, "cell", None) + pbc = getattr(current, "pbc", None) + cell_is_4d = original_cell is not None and original_cell.dim() == 4 + + # Stress needs a cell to strain and a non-zero volume to divide by. A + # batch may carry a placeholder cell with pbc all-False, whose zero + # volume would turn the stress into Inf, so gate on periodicity too. + # bool(pbc.any()) is a data-dependent branch, which is safe here + # because evaluate() is eager-only. + wants_stress = "stress" in (self.model_config.active_outputs or set()) + strain_cell: Tensor | None = None + if ( + wants_stress + and original_cell is not None + and (pbc is None or bool(pbc.any())) + ): + # Detach the stored cell so only the strain leaf carries the + # gradient; a [B, 1, 3, 3] cell is squeezed to [B, 3, 3]. + strain_cell = original_cell.detach() + if cell_is_4d: + strain_cell = strain_cell.squeeze(1) + + with torch.enable_grad(): + # requires_grad_() is not supported by torch.compile; evaluate() + # is intentionally kept eager (see class docstring). + pos_leaf = original_positions.detach().requires_grad_(True) # [N, 3] + + displacement: Tensor | None = None + pos_for_energy: Tensor = pos_leaf + cell_for_energy: Tensor | None = None + + if strain_cell is not None: + pos_for_energy, cell_for_energy, displacement = prepare_strain( + pos_leaf, strain_cell, current.batch_idx + ) + + try: + current["positions"] = pos_for_energy + if cell_for_energy is not None: + current["cell"] = ( + cell_for_energy.unsqueeze(1) if cell_is_4d else cell_for_energy + ) + + bias_energy: Tensor = self.energy(current) # [B, 1] + + stress: Tensor | None = None + if not bias_energy.requires_grad: + # An energy with no graph at all — a bias returning a + # constant on this branch. Every gradient is zero, but + # autograd rejects such an output outright ("does not + # require grad and does not have a grad_fn"), so fill the + # zeros directly rather than calling it. + forces = torch.zeros_like(pos_leaf) + if strain_cell is not None: + stress = torch.zeros_like(strain_cell) + elif strain_cell is not None and displacement is not None: + # allow_unused: a bias need not depend on both positions + # and strain. A pure volume restraint has no position + # dependence, and its zero force is an answer, not an error. + forces, stress = autograd_forces_and_stresses( + bias_energy, + pos_leaf, + displacement, + strain_cell, + current.num_graphs, + allow_unused=True, + ) + else: + forces = autograd_forces(bias_energy, pos_leaf, allow_unused=True) + + finally: + current["positions"] = original_positions + if original_cell is not None: + current["cell"] = original_cell + + return BiasResult( + energy=bias_energy.detach(), + forces=forces.detach(), + stress=None if stress is None else stress.detach(), + ) + + def forward(self, data: AtomicData | Batch, **kwargs: Any) -> ModelOutputs: + """Evaluate the bias and return its contribution as ``ModelOutputs``. + + The ``BaseModelMixin`` view of :meth:`evaluate`. Keys are restricted + to ``model_config.active_outputs``, so this composes with + ``sum_outputs`` and with a model via ``+``. + + :meth:`evaluate` remains the :class:`BiasPotential` entry point and + is what the enhanced-sampling runner calls: ``BiasResult`` carries + ``observables`` and ``state_version``, which have no place in + ``ModelOutputs``, and those are dropped here. + + Parameters + ---------- + data: + The current batch. + **kwargs: + Unused; accepted for interface compatibility. + + Returns + ------- + ModelOutputs + ``energy``, ``forces``, and — when active and the batch is + periodic — ``stress``. + """ + result = self.evaluate(data) # type: ignore[arg-type] + active = self.model_config.active_outputs or set() + outputs: ModelOutputs = {} + for key in ("energy", "forces", "stress"): + value = getattr(result, key) + if value is not None and key in active: + outputs[key] = value + return outputs + + +# --------------------------------------------------------------------------- +# Bias aggregation +# --------------------------------------------------------------------------- + + +def aggregate_bias_results(results: list[BiasResult]) -> BiasResult: + """Sum a list of ``BiasResult`` objects into a single combined result. + + All biases are evaluated against the **same unmodified physical + outputs**; their contributions are summed once here and applied + together. A bias cannot accidentally observe the force contribution + of another bias. + + Rules + ----- + * ``None`` fields are skipped (treated as zero contribution). + * All results must agree on which cell-response field they use: every + result that carries a cell response must use **either** ``stress`` + **or** ``virial`` — never a mix of both across the list. Mixing + raises ``ValueError`` at aggregation time (not inside ``BiasResult``) + with a message identifying which indices contributed each field. + Converting between the two requires the cell volume and is the + caller's responsibility before aggregation. + * ``observables`` dicts are merged, not summed; duplicate keys raise + ``ValueError`` so that namespacing (``bias//``) must be + applied before calling this function. + + Relationship to :func:`~nvalchemi.models._utils.sum_outputs` + ----------------------------------------------------------- + The element-wise tensor sum is delegated to ``sum_outputs``, which + already implements exactly this contract for model composition. The + two rules above stay here because they are stricter than what + ``sum_outputs`` can offer its own callers: + + * ``sum_outputs`` resolves a non-additive key collision by + last-write-wins, which the model pipeline depends on (two composed + models may both emit ``charges``). Silently dropping one bias's + observable is not acceptable, so the collision is an error here. + * ``observables`` is a nested mapping, not part of ``ModelOutputs``. + Flattening it into the same dict would let an observable named + ``energy`` be summed into the bias energy, so it is merged + separately. + + The stress/virial mixing check stays here for a different reason: it + would guard nothing in ``sum_outputs``. Model wrappers normalise a + virial to a tensile-positive stress at the adapter boundary (dividing by + the cell volume), so the ``ModelOutputs`` that reach ``sum_outputs`` + carry ``stress`` and never ``virial`` — which is why its default + ``additive_keys`` omits virial entirely. ``BiasResult`` deliberately + allows either field, so only this layer can be handed both. + + Parameters + ---------- + results: + List of ``BiasResult`` objects from individual biases. May be + empty, in which case an empty ``BiasResult()`` is returned. + + Returns + ------- + BiasResult + Aggregated result with summed contributions. + """ + if not results: + return BiasResult() + + # Detect stress/virial mixing up-front so the error is raised here with + # a clear message, not inside BiasResult.__post_init__ with a generic + # mutual-exclusion message that doesn't identify which results mixed them. + has_stress = any(r.stress is not None for r in results) + has_virial = any(r.virial is not None for r in results) + if has_stress and has_virial: + stress_indices = [i for i, r in enumerate(results) if r.stress is not None] + virial_indices = [i for i, r in enumerate(results) if r.virial is not None] + raise ValueError( + f"aggregate_bias_results: results[{stress_indices}] provide 'stress' " + f"and results[{virial_indices}] provide 'virial' — cannot mix both in " + "the same aggregation. Make all biases return the same field. " + "Converting between stress and virial requires the cell volume and is " + "the caller's responsibility before aggregation." + ) + + summed = sum_outputs( + *( + { + "energy": r.energy, + "forces": r.forces, + "stress": r.stress, + "virial": r.virial, + } + for r in results + ), + additive_keys={"energy", "forces", "stress", "virial"}, + ) + + observables_total: dict[str, Tensor] = {} + observable_source: dict[str, int] = {} + for i, r in enumerate(results): + for key, val in r.observables.items(): + if key in observables_total: + raise ValueError( + f"aggregate_bias_results: duplicate observable key {key!r} from " + f"results[{observable_source[key]}] and results[{i}]. Apply " + "'bias//' namespacing before aggregation." + ) + observables_total[key] = val + observable_source[key] = i + + return BiasResult( + energy=summed.get("energy"), + forces=summed.get("forces"), + stress=summed.get("stress"), + virial=summed.get("virial"), + observables=observables_total, + ) diff --git a/nvalchemi/enhanced_sampling/_checkpoint.py b/nvalchemi/enhanced_sampling/_checkpoint.py new file mode 100644 index 00000000..890d8bbb --- /dev/null +++ b/nvalchemi/enhanced_sampling/_checkpoint.py @@ -0,0 +1,641 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Transactional Zarr checkpoints for enhanced sampling. + +Layout, extending the existing ``AtomicData`` Zarr record in place:: + + checkpoint.zarr/ + meta/, core/, custom/ walker batch, via AtomicDataZarrWriter + sampling/ + manifest committed metadata — WRITTEN LAST + dynamics/ integrator, thermostat, and RNG counters + biases// each bias's state_dict() + runner/ walker-id allocation and epoch counters + exchange/ ladder config, counters, acceptance-RNG position + +State is stored as Zarr arrays (tensors) and group attributes (scalars, +strings, nested mappings). There are **no pickle payloads**: a checkpoint is +readable by anything that can read Zarr, and loading one cannot execute code. + +Transactionality +---------------- +Components are written first, each checksummed, and the manifest is written +last. A checkpoint interrupted at any point therefore has no manifest, and +:func:`read_checkpoint` rejects a store without one rather than restoring a +torn half-state. Checksums are verified on read, so a store that was +truncated *after* the manifest landed is also caught. + +The cover is total. Every ``sampling/`` component carries its own digest, +and ``batch_checksum`` covers ``meta/``, ``core/``, and ``custom/`` — the +positions, velocities, forces, pointer arrays, and walker identity that +``AtomicDataZarrWriter`` writes outside the component path. Checksumming +only the sampling state would attest to the bias and integrator while +leaving the coordinates unguarded, which is the half of the checkpoint a +reader is most likely to trust blindly. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +import zarr +from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import ValidationError as PydanticValidationError + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.data.datapipes.backends.zarr import ( + AtomicDataZarrReader, + AtomicDataZarrWriter, +) + +if TYPE_CHECKING: + from pathlib import Path + +__all__ = [ + "CHECKPOINT_FORMAT_VERSION", + "CheckpointManifest", + "read_checkpoint", + "write_checkpoint", +] + +CHECKPOINT_FORMAT_VERSION = 1 + +_SAMPLING = "sampling" +_MANIFEST = "sampling/manifest" + +# Graph-level fields the runner owns. ``AtomicDataZarrWriter.write`` only +# persists fields it recognises, so these are added explicitly through +# ``add_custom`` — without that, walker identity would be silently dropped and +# a "restored" run would come back with fresh ids and default state +# assignments. +_IDENTITY_FIELDS = ("walker_id", "thermodynamic_state_id") + +_SCALAR_TYPES = (bool, int, float, str) + +# Zarr groups that together hold the walker batch. These are written by +# AtomicDataZarrWriter rather than through _encode_state, so they need their +# own integrity cover — without it the manifest would attest only to the +# sampling/* state and a corrupted core/positions would restore silently. +_BATCH_GROUPS = ("meta", "core", "custom") + + +class CheckpointManifest(BaseModel): + """Committed checkpoint metadata, written after every component. + + Its presence is the commit marker: a store without a manifest is an + incomplete write and is refused. + + Attributes + ---------- + format_version: + Layout version, for forward migration. + sampling_step: + Dynamics step count the checkpoint was taken at. + sampling_epoch: + Consistency epoch the checkpoint was taken at. + steps_per_epoch: + Epoch length in force evaluations. + num_graphs: + Walker count, validated against the restored batch. + components: + Names of the ``sampling/`` groups written. + checksums: + SHA-256 per ``sampling/`` component, verified on read. Must have an + entry for every name in :attr:`components`, and no others. + batch_checksum: + SHA-256 over every array in ``meta/``, ``core/``, and ``custom/`` — + the walker batch itself. Kept separate from :attr:`checksums` + because those name ``sampling/`` groups the reader walks, while this + covers arrays written by ``AtomicDataZarrWriter``. + model_class: + Fully-qualified model wrapper class, validated on restore. + dynamics_class: + Fully-qualified dynamics class, validated on restore. + bias_classes: + Bias name to fully-qualified class, validated on restore. + exchange_config: + Replica-exchange configuration fingerprint, or ``None`` when the run + had no exchange. Validated on restore: the ladder decides what a + swap *means*, so restoring into a different one — or into no exchange + at all — has to be refused rather than silently accepted. + """ + + model_config = ConfigDict(extra="forbid") + + format_version: int = CHECKPOINT_FORMAT_VERSION + sampling_step: int + sampling_epoch: int + steps_per_epoch: int + num_graphs: int + components: list[str] = Field(default_factory=list) + checksums: dict[str, str] = Field(default_factory=dict) + batch_checksum: str = "" + model_class: str = "" + dynamics_class: str = "" + bias_classes: dict[str, str] = Field(default_factory=dict) + exchange_config: dict[str, Any] | None = None + + @model_validator(mode="after") + def _every_component_is_covered(self) -> CheckpointManifest: + """Reject a manifest whose integrity cover has gaps. + + In a manifest-gated format the manifest is the authority on what the + store contains, so a declared component without a checksum is not + "unverified" — it is invalid. Treating a missing entry as permission + to skip verification would make the cover opt-out: deleting one key + from the manifest attributes is enough to leave that component free to + modify. The same goes for the walker batch. + + Returns + ------- + CheckpointManifest + The validated manifest. + + Raises + ------ + ValueError + If a component has no checksum, a checksum names no component, or + the batch checksum is absent. + """ + declared = set(self.components) + covered = set(self.checksums) + + uncovered = sorted(declared - covered) + if uncovered: + raise ValueError( + f"Checkpoint manifest declares component(s) {uncovered} with " + "no checksum. Every declared component must be covered; a " + "missing entry is a tampered or truncated manifest, not a " + "component that may be read unverified." + ) + orphaned = sorted(covered - declared) + if orphaned: + raise ValueError( + f"Checkpoint manifest has checksum(s) for {orphaned}, which it " + "does not declare as components. The manifest is inconsistent " + "with itself." + ) + if not self.batch_checksum: + raise ValueError( + "Checkpoint manifest has no batch_checksum. The walker batch " + "— positions, velocities, and walker identity — would then be " + "restored unverified." + ) + return self + + +def _qualified_name(obj: Any) -> str: + """Return ``module.ClassName`` for *obj*'s type. + + Parameters + ---------- + obj: + Any object. + + Returns + ------- + str + Fully-qualified class name. + """ + cls = type(obj) + return f"{cls.__module__}.{cls.__qualname__}" + + +def _torch_dtype(name: str) -> torch.dtype: + """Resolve ``"torch.float32"`` back to the dtype object. + + Parameters + ---------- + name: + The ``str(dtype)`` form. + + Returns + ------- + torch.dtype + The resolved dtype. + + Raises + ------ + ValueError + If the name does not resolve to a dtype. + """ + candidate = getattr(torch, name.rsplit(".", 1)[-1], None) + if not isinstance(candidate, torch.dtype): + raise ValueError(f"Checkpoint: unknown tensor dtype {name!r}.") + return candidate + + +def _encode_state(group: zarr.Group, state: Mapping[str, Any]) -> None: + """Write a nested state mapping into *group*. + + Tensors become arrays; scalars, strings, ``None``, and flat sequences + become attributes; nested mappings become subgroups. A per-key kind tag + is stored so the decoder never has to guess. + + Parameters + ---------- + group: + Destination Zarr group. + state: + Mapping of tensors, scalars, and nested mappings. + + Raises + ------ + TypeError + If a value is none of the supported kinds. Refusing here is + deliberate: the alternative is a pickle payload, which would make a + checkpoint executable and unreadable outside Python. + """ + kinds: dict[str, str] = {} + scalars: dict[str, Any] = {} + dtypes: dict[str, str] = {} + shapes: dict[str, list[int]] = {} + + for key, value in state.items(): + if isinstance(value, torch.Tensor): + kinds[key] = "tensor" + dtypes[key] = str(value.dtype) + array = value.detach().cpu().contiguous().numpy() + # Zarr materialises a 0-d array as shape (1,), so the true shape + # is recorded separately and reapplied on decode. Without it a + # scalar buffer (a step counter, a deposition count) comes back + # rank-1 and fails its own component checksum on restore. + shapes[key] = list(array.shape) + group.create_array(key, shape=array.shape, dtype=array.dtype) + if array.size: + group[key][...] = array + elif isinstance(value, Mapping): + kinds[key] = "group" + _encode_state(group.require_group(key), value) + elif value is None or isinstance(value, _SCALAR_TYPES): + kinds[key] = "scalar" + scalars[key] = value + elif isinstance(value, (list, tuple)) and all( + v is None or isinstance(v, _SCALAR_TYPES) for v in value + ): + kinds[key] = "scalar" + scalars[key] = list(value) + else: + raise TypeError( + f"Checkpoint: cannot store {key!r} of type " + f"{type(value).__name__}. State must be tensors, scalars, " + "strings, flat sequences of those, or nested mappings — a " + "checkpoint carries no pickle payloads." + ) + + group.attrs["kinds"] = kinds + group.attrs["scalars"] = scalars + group.attrs["dtypes"] = dtypes + group.attrs["shapes"] = shapes + + +def _decode_state(group: zarr.Group, device: torch.device | str) -> dict[str, Any]: + """Read back a mapping written by :func:`_encode_state`. + + Parameters + ---------- + group: + Source Zarr group. + device: + Device to place restored tensors on. + + Returns + ------- + dict[str, Any] + The restored mapping. + """ + kinds = dict(group.attrs.get("kinds", {})) + scalars = dict(group.attrs.get("scalars", {})) + dtypes = dict(group.attrs.get("dtypes", {})) + shapes = dict(group.attrs.get("shapes", {})) + + state: dict[str, Any] = {} + for key, kind in kinds.items(): + if kind == "tensor": + array = np.asarray(group[key][...]) + tensor = torch.from_numpy(np.ascontiguousarray(array)) + # Missing for checkpoints written before shapes were recorded; + # falling back to the stored shape reads those exactly as before. + if key in shapes: + tensor = tensor.reshape(tuple(shapes[key])) + state[key] = tensor.to(device=device, dtype=_torch_dtype(dtypes[key])) + elif kind == "group": + state[key] = _decode_state(group[key], device) + else: + state[key] = scalars.get(key) + return state + + +def _component_checksum(state: Mapping[str, Any]) -> str: + """Return a SHA-256 over a component's contents. + + Order-independent by construction (keys are walked sorted), so the digest + depends on the state and not on dict insertion order. + + Parameters + ---------- + state: + The component state. + + Returns + ------- + str + Hex digest. + """ + digest = hashlib.sha256() + + def _walk(mapping: Mapping[str, Any], prefix: str) -> None: + for key in sorted(mapping): + value = mapping[key] + digest.update(f"{prefix}{key}".encode()) + if isinstance(value, torch.Tensor): + digest.update(str(value.dtype).encode()) + digest.update(str(tuple(value.shape)).encode()) + digest.update(value.detach().cpu().contiguous().numpy().tobytes()) + elif isinstance(value, Mapping): + _walk(value, f"{prefix}{key}/") + else: + digest.update(json.dumps(value, sort_keys=True, default=str).encode()) + + _walk(state, "") + return digest.hexdigest() + + +def _batch_checksum(root: zarr.Group) -> str: + """Return a SHA-256 over every array holding the walker batch. + + Covers ``meta/``, ``core/``, and ``custom/`` — positions, velocities, + forces, the CSR pointer arrays, and the runner's identity fields. These + are written by ``AtomicDataZarrWriter``, not by :func:`_encode_state`, so + they are outside the per-component checksum path and need this. + + Reads the arrays back from the store rather than hashing the in-memory + batch, so the write-side and read-side digests are computed over exactly + the same bytes. The cost is one extra full read of the batch on write; + integrity that only sometimes holds is not worth the saving. + + Parameters + ---------- + root: + The opened checkpoint root group. + + Returns + ------- + str + Hex digest, empty-string-safe if the groups are absent. + """ + digest = hashlib.sha256() + for group_name in _BATCH_GROUPS: + if group_name not in root: + continue + group = root[group_name] + for key in sorted(group.array_keys()): + array = group[key] + digest.update(f"{group_name}/{key}".encode()) + digest.update(str(array.dtype).encode()) + digest.update(str(tuple(array.shape)).encode()) + digest.update(np.ascontiguousarray(array[...]).tobytes()) + return digest.hexdigest() + + +def write_checkpoint( + path: str | Path, + batch: Batch, + components: Mapping[str, Mapping[str, Any]], + *, + sampling_step: int, + sampling_epoch: int, + steps_per_epoch: int, + model_class: str = "", + dynamics_class: str = "", + bias_classes: Mapping[str, str] | None = None, + exchange_config: Mapping[str, Any] | None = None, +) -> CheckpointManifest: + """Write a transactional checkpoint. + + Order matters and is the whole guarantee: walker batch, then each + component, then the manifest. An interruption anywhere before the last + step leaves a store with no manifest, which :func:`read_checkpoint` + refuses. + + Parameters + ---------- + path: + Destination store. + batch: + The live walker batch. + components: + Mapping of component name (``"dynamics"``, ``"biases/umbrella"``, + ``"runner"``) to its state mapping. + sampling_step: + Step count at the checkpoint. + sampling_epoch: + Epoch at the checkpoint. + steps_per_epoch: + Epoch length. + model_class: + Fully-qualified model class, recorded for restore-time validation. + dynamics_class: + Fully-qualified dynamics class, likewise. + bias_classes: + Bias name to fully-qualified class, likewise. + exchange_config: + Replica-exchange configuration fingerprint, or ``None``. + + Returns + ------- + CheckpointManifest + The manifest that was committed. + """ + writer = AtomicDataZarrWriter(str(path)) + writer.write(batch) + + # AtomicDataZarrWriter.write persists only the fields it recognises, so + # the runner's identity fields have to be added through the custom-array + # API or they are dropped without complaint. + for field in _IDENTITY_FIELDS: + value = getattr(batch, field, None) + if value is not None: + writer.add_custom(field, value.reshape(-1), level="system") + + root = zarr.open_group(str(path), mode="a") + # Computed now, while the store holds only the batch: the digest must not + # depend on the sampling/* groups written next. + batch_checksum = _batch_checksum(root) + sampling = root.require_group(_SAMPLING) + + checksums: dict[str, str] = {} + for name, state in components.items(): + group = sampling + for part in name.split("/"): + group = group.require_group(part) + _encode_state(group, state) + checksums[name] = _component_checksum(state) + + manifest = CheckpointManifest( + sampling_step=sampling_step, + sampling_epoch=sampling_epoch, + steps_per_epoch=steps_per_epoch, + num_graphs=batch.num_graphs, + components=sorted(components), + checksums=checksums, + batch_checksum=batch_checksum, + model_class=model_class, + dynamics_class=dynamics_class, + bias_classes=dict(bias_classes or {}), + exchange_config=dict(exchange_config) if exchange_config else None, + ) + + # Written last: this is the commit. + manifest_group = root.require_group(_MANIFEST) + manifest_group.attrs["manifest"] = manifest.model_dump() + return manifest + + +def read_checkpoint( + path: str | Path, device: torch.device | str = "cpu" +) -> tuple[Batch, dict[str, dict[str, Any]], CheckpointManifest]: + """Read a checkpoint, refusing anything not fully committed. + + Parameters + ---------- + path: + Source store. + device: + Device to place restored tensors on. + + Returns + ------- + tuple[Batch, dict[str, dict[str, Any]], CheckpointManifest] + The walker batch, the component states, and the manifest. + + Raises + ------ + ValueError + If the store has no committed manifest; if the manifest is internally + inconsistent (a declared component with no checksum, a checksum for no + component, or no batch checksum); if a declared component is missing; + or if any checksum, component or batch, does not match. + """ + root = zarr.open_group(str(path), mode="r") + if _MANIFEST not in root: + raise ValueError( + f"Checkpoint at {path} has no committed manifest, so it was never " + "finished — the manifest is written last, after every component. " + "Treat this store as an interrupted write and discard it." + ) + try: + manifest = CheckpointManifest(**dict(root[_MANIFEST].attrs["manifest"])) + except PydanticValidationError as exc: + # Surface the same way as every other checkpoint failure — one + # ValueError naming the store — rather than a nested pydantic report. + reasons = "; ".join(str(err["msg"]) for err in exc.errors()) + raise ValueError( + f"Checkpoint at {path} has an invalid manifest: {reasons}" + ) from exc + + if manifest.format_version != CHECKPOINT_FORMAT_VERSION: + raise ValueError( + f"Checkpoint at {path} has format_version " + f"{manifest.format_version}, but this build reads version " + f"{CHECKPOINT_FORMAT_VERSION}." + ) + + sampling = root[_SAMPLING] + states: dict[str, dict[str, Any]] = {} + for name in manifest.components: + group: Any = sampling + for part in name.split("/"): + if part not in group: + raise ValueError( + f"Checkpoint at {path} declares component {name!r} in its " + "manifest but the group is missing; the store is corrupt." + ) + group = group[part] + state = _decode_state(group, device) + actual = _component_checksum(state) + # Unconditional: the manifest validator guarantees the entry exists, + # so no path reads a component without checking it. + expected = manifest.checksums[name] + if actual != expected: + raise ValueError( + f"Checkpoint at {path}: component {name!r} failed its checksum " + f"(expected {expected[:12]}…, got {actual[:12]}…). The store " + "was modified or truncated after the manifest was written." + ) + states[name] = state + + actual = _batch_checksum(root) + if actual != manifest.batch_checksum: + raise ValueError( + f"Checkpoint at {path}: the walker batch failed its checksum " + f"(expected {manifest.batch_checksum[:12]}…, got {actual[:12]}…). " + "One of meta/, core/, or custom/ was modified or truncated after " + "the manifest was written — positions, velocities, or walker " + "identity can no longer be trusted." + ) + + return _read_batch(path, manifest, device), states, manifest + + +def _read_batch( + path: str | Path, manifest: CheckpointManifest, device: torch.device | str +) -> Batch: + """Reconstruct the walker batch, identity fields included. + + Parameters + ---------- + path: + Source store. + manifest: + The committed manifest, read for the expected walker count. + device: + Device for the restored batch. + + Returns + ------- + Batch + The restored batch. + + Raises + ------ + ValueError + If the store holds a different number of walkers than the manifest + recorded. + """ + reader = AtomicDataZarrReader(str(path)) + if len(reader) != manifest.num_graphs: + raise ValueError( + f"Checkpoint at {path} holds {len(reader)} walker(s) but its " + f"manifest records {manifest.num_graphs}; the store is corrupt." + ) + + data_list = [AtomicData(**reader.read(i)[0]) for i in range(len(reader))] + batch = Batch.from_data_list(data_list).to(device) + + root = zarr.open_group(str(path), mode="r") + custom = root["custom"] if "custom" in root else None + for field in _IDENTITY_FIELDS: + if custom is not None and field in custom: + values = np.asarray(custom[field][...]) + batch[field] = torch.from_numpy(np.ascontiguousarray(values)).to( + device=device, dtype=torch.long + ) + return batch diff --git a/nvalchemi/enhanced_sampling/_exchange.py b/nvalchemi/enhanced_sampling/_exchange.py new file mode 100644 index 00000000..aa60d673 --- /dev/null +++ b/nvalchemi/enhanced_sampling/_exchange.py @@ -0,0 +1,790 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Synchronous replica exchange over a batch of walkers. + +Exchange swaps **thermodynamic-state labels, not atomic coordinates**. A +walker keeps its execution slot, its history, and its integrator arrays; what +changes is the temperature or bias window assigned to it. That keeps the +move local — no coordinate traffic, no reallocation — which is what makes it +viable inside one batched GPU step. + +Every walker holds exactly one state and every state exactly one walker, so +an exchange is a permutation of :attr:`Batch.thermodynamic_state_id`. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Any, Literal + +import torch +from pydantic import BaseModel, ConfigDict, Field + +from nvalchemi.dynamics.hooks._utils import KB_EV + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + +__all__ = ["ReplicaExchange", "ThermodynamicState"] + + +class ThermodynamicState(BaseModel): + """One set of conditions a walker can be assigned to. + + Attributes + ---------- + state_id: + Index into the ladder. Must be dense and start at zero across the + set of states, because pairing walks neighbouring indices. + temperature: + Temperature in Kelvin. Equal across all states means the ladder + varies by bias window instead, which selects umbrella acceptance. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + state_id: int = Field(ge=0) + temperature: float = Field(gt=0.0) + + +class ReplicaExchange: + r"""Synchronous replica exchange with an even/odd pair schedule. + + Parameters + ---------- + states: + The ladder. ``state_id`` values must be exactly ``0..S-1``. + initial_state_ids: + Assignment of states to walkers, shape ``[B]``. Must be a + permutation of ``0..S-1``: replica exchange presumes a bijection + between walkers and states, and a duplicate would let two walkers + claim the same rung. + mode: + ``"synchronous"`` only. Asynchronous exchange is not implemented. + attempt_interval: + Dynamics steps per exchange segment. + random_seed: + Base seed for acceptance draws. Randomness is derived per attempt as + ``random_seed + exchange_id`` rather than from a long-lived generator, + so a checkpoint needs two integers instead of an opaque RNG blob — + the same counter-based scheme ``NVTLangevin`` uses for its noise. + + Raises + ------ + ValueError + If the mode is unsupported, ``attempt_interval`` is below 1, the + ladder is not dense, or the assignment is not a permutation. + + Notes + ----- + Acceptance + Which rule applies is inferred from the ladder and validated, rather + than being a free parameter that can silently disagree with it. + + *Temperature exchange* — temperatures differ: + + .. math:: + + \log a = \min\bigl(0,\ (\beta_i - \beta_j)(U_i - U_j)\bigr) + + A cold replica holding anomalously high energy therefore moves up the + ladder with probability one, which is the point of the method. + + *Umbrella exchange* — temperatures are equal and the states differ by + bias window: + + .. math:: + + \log a = \min\bigl(0,\ u_i(x_i) + u_j(x_j) + - u_i(x_j) - u_j(x_i)\bigr) + + with :math:`u_k` the reduced bias potential of state *k*. The two + cross terms need the bias re-evaluated under swapped labels, which + costs one extra bias evaluation per attempt. + + Not supported + A ladder that varies temperature *and* bias window at once needs a + combined acceptance rule that is not implemented. The temperature + rule alone omits the cross-state bias terms, so running it anyway + would break detailed balance with no symptom — it is therefore + **rejected**, twice over: a bias that declares + ``state_dependent_for_exchange`` is refused at construction, and the + runner additionally probes empirically at prime time by evaluating + every bias under a permuted assignment. Vary one or the other. + """ + + def __init__( + self, + states: Sequence[ThermodynamicState], + initial_state_ids: torch.Tensor, + *, + mode: Literal["synchronous"] = "synchronous", + attempt_interval: int = 100, + random_seed: int = 1234, + ) -> None: + if mode != "synchronous": + raise ValueError( + f"ReplicaExchange: mode={mode!r} is not supported. Only " + "'synchronous' is implemented; asynchronous exchange " + "(pair-local rendezvous, non-blocking workers) is future work." + ) + if len(states) < 2: + raise ValueError( + f"ReplicaExchange: need at least 2 states to exchange, got " + f"{len(states)}." + ) + ladder = sorted(states, key=lambda s: s.state_id) + if [s.state_id for s in ladder] != list(range(len(ladder))): + raise ValueError( + f"ReplicaExchange: state_id values must be exactly 0..{len(ladder) - 1}, " + f"got {sorted(s.state_id for s in states)}. Pairing walks " + "neighbouring indices, so a sparse ladder has no defined " + "neighbours." + ) + if int(attempt_interval) < 1: + raise ValueError( + f"ReplicaExchange: attempt_interval must be at least 1, got " + f"{attempt_interval}. A non-positive interval has no meaning " + "as a segment length — it would be clamped to every-step " + "exchange while the checkpoint recorded the value you passed, " + "so the run and its metadata would disagree." + ) + self.states = tuple(ladder) + self.mode = mode + self.attempt_interval = int(attempt_interval) + self.random_seed = int(random_seed) + + self._acceptance = self._infer_acceptance() + self.initial_state_ids = self.validate_assignment( + initial_state_ids, source="initial_state_ids" + ) + self.exchange_id = 0 + self.attempts = 0 + self.accepted = 0 + # Per neighbouring-state-pair tallies, for the acceptance-rate + # diagnostics a REMD run is tuned on. + self.pair_attempts = [0] * (len(self.states) - 1) + self.pair_accepted = [0] * (len(self.states) - 1) + + # ------------------------------------------------------------------ + # Configuration + # ------------------------------------------------------------------ + + def validate_assignment( + self, + state_ids: torch.Tensor, + num_graphs: int | None = None, + *, + source: str = "assignment", + ) -> torch.Tensor: + """Return *state_ids* as a validated ``[B]`` long tensor. + + The bijection between walkers and states is the assumption every + other piece rests on: pairing looks up "which walker holds state k", + and a duplicate or a wrong-length assignment makes that lookup + meaningless. Without this the failures surface far from their cause + — a length mismatch as ``ValueError: Length mismatch: 4 vs 2`` from + inside the batch storage, a duplicate as ``KeyError: 3`` from the + pair lookup — neither of which names the ladder or the batch. + + Used for both the constructor argument and whatever assignment the + batch actually carries, so the rule lives in one place. + + Parameters + ---------- + state_ids : torch.Tensor + Candidate assignment, any shape reshapeable to ``[B]``. + num_graphs : int | None + Walker count to check the ladder against, when known. + source : str + Name of the thing being validated, used in the error. + + Returns + ------- + torch.Tensor + The assignment as a 1-D long tensor. + + Raises + ------ + ValueError + If the walker count disagrees with the ladder, or the assignment + is not a permutation of ``0..S-1``. + """ + n_states = len(self.states) + if num_graphs is not None and num_graphs != n_states: + raise ValueError( + f"ReplicaExchange: the ladder has {n_states} state(s) but the " + f"batch has {num_graphs} walker(s). Replica exchange presumes " + "one walker per state — pairing looks up which walker holds " + "each rung, which has no answer when the counts differ. Build " + "the batch with one graph per ThermodynamicState." + ) + + ids = state_ids.reshape(-1).to(torch.long) + if ids.numel() != n_states: + raise ValueError( + f"ReplicaExchange: {source} has {ids.numel()} entr(ies) but the " + f"ladder has {n_states} state(s); they must agree." + ) + if sorted(ids.tolist()) != list(range(n_states)): + raise ValueError( + f"ReplicaExchange: {source} must be a permutation of " + f"0..{n_states - 1}, got {ids.tolist()}. Replica exchange " + "presumes one walker per state; a duplicate would let two " + "walkers claim the same rung of the ladder, and leave another " + "rung held by none." + ) + return ids + + def _infer_acceptance(self) -> Literal["temperature", "umbrella"]: + """Return which acceptance rule this ladder implies. + + A :class:`ThermodynamicState` carries only a temperature, so the + rule is read from the ladder: varying temperatures mean temperature + exchange, equal ones mean the states can only differ by which bias + window they select. + + Inferring rather than accepting a parameter is deliberate — a + mismatch between a declared rule and the ladder it runs on would be + silent, and wrong acceptance breaks detailed balance without any + symptom a run would show. + + Returns + ------- + Literal["temperature", "umbrella"] + The applicable rule. + """ + temperatures = [s.temperature for s in self.states] + varies_temperature = max(temperatures) - min(temperatures) > 1e-12 + return "temperature" if varies_temperature else "umbrella" + + @property + def acceptance(self) -> str: + """Return the inferred acceptance rule name.""" + return self._acceptance + + @property + def temperatures(self) -> torch.Tensor: + """Return the ladder temperatures in Kelvin, shape ``[S]``.""" + return torch.tensor([s.temperature for s in self.states]) + + def validate_for(self, biases: Mapping[str, Any]) -> None: + """Reject bias/ladder combinations whose acceptance is undefined. + + Parameters + ---------- + biases: + The runner's bias mapping. + + Raises + ------ + ValueError + If umbrella exchange is configured with no bias to exchange over, + if any bias cannot supply the energy the acceptance rule needs, or + if a state-dependent bias is combined with a temperature ladder. + """ + if self._acceptance == "umbrella" and not biases: + raise ValueError( + "ReplicaExchange: every state has the same temperature, so the " + "ladder can only differ by bias window — but no biases were " + "registered. Either vary the temperatures for temperature " + "exchange, or register the bias whose windows the states select." + ) + for name, bias in biases.items(): + energy_less = getattr(bias, "supplies_exchange_energy", None) + if energy_less is False: + raise ValueError( + f"ReplicaExchange: bias {name!r} declares that it supplies " + "no exchange energy (a force-only bias such as adaptive " + "biasing force). The acceptance rule needs a cross-state " + "bias energy, so such a bias cannot participate; run it " + "without replica exchange." + ) + if self._acceptance == "temperature" and ( + getattr(bias, "state_dependent_for_exchange", False) is True + ): + raise ValueError( + f"ReplicaExchange: the ladder varies temperature, which " + f"selects temperature acceptance, but bias {name!r} has " + "per-state parameters. The combined temperature-plus-window " + "acceptance rule is not implemented, and the temperature " + "rule alone omits the cross-state bias terms — so detailed " + "balance would be wrong with nothing to show for it. Use a " + "single-window bias with a temperature ladder, or equal " + "temperatures with a multi-window bias." + ) + + # ------------------------------------------------------------------ + # Scheduling + # ------------------------------------------------------------------ + + def pair_schedule(self, segment: int) -> list[tuple[int, int]]: + """Return the neighbouring state pairs attempted in *segment*. + + Alternating even/odd offsets means every rung of the ladder is + exchangeable with both neighbours over two segments, while no state + appears in two pairs of the same segment — which is what lets all + pairs be decided simultaneously. + + Parameters + ---------- + segment: + Exchange segment index. + + Returns + ------- + list[tuple[int, int]] + Neighbouring ``(state_id, state_id + 1)`` pairs. + """ + offset = segment % 2 + return [(index, index + 1) for index in range(offset, len(self.states) - 1, 2)] + + def _uniforms(self, count: int, device: torch.device) -> torch.Tensor: + """Draw acceptance uniforms for one attempt, reproducibly. + + Parameters + ---------- + count: + Number of draws. + device: + Device for the result. + + Returns + ------- + torch.Tensor + Shape ``[count]`` in ``[0, 1)``. + + Notes + ----- + Drawn on the CPU from a generator seeded with + ``random_seed + exchange_id``, then moved. Seeding per attempt makes + the sequence a pure function of two checkpointed integers, and drawing + on the CPU keeps it independent of the device the run happens to use — + so a restored run reproduces the same accept/reject decisions. + """ + generator = torch.Generator() + generator.manual_seed(self.random_seed + self.exchange_id) + return torch.rand(count, generator=generator).to(device) + + # ------------------------------------------------------------------ + # Acceptance + # ------------------------------------------------------------------ + + def _log_acceptance_temperature( + self, + pairs: list[tuple[int, int]], + walker_of_state: dict[int, int], + energies: torch.Tensor, + ) -> torch.Tensor: + """Return ``log a`` per pair for temperature exchange. + + Parameters + ---------- + pairs: + Neighbouring state pairs. + walker_of_state: + State id to the walker row currently holding it. + energies: + Per-walker potential energy ``U``, shape ``[B]``. + + Returns + ------- + torch.Tensor + ``log a`` per pair, shape ``[len(pairs)]``, capped at zero. + """ + temperatures = self.temperatures.to(energies.device, energies.dtype) + beta = 1.0 / (KB_EV * temperatures) # [S] + + values = [] + for state_i, state_j in pairs: + walker_i = walker_of_state[state_i] + walker_j = walker_of_state[state_j] + delta = (beta[state_i] - beta[state_j]) * ( + energies[walker_i] - energies[walker_j] + ) + values.append(delta) + return torch.clamp(torch.stack(values), max=0.0) + + def _log_acceptance_umbrella( + self, + pairs: list[tuple[int, int]], + walker_of_state: dict[int, int], + bias_current: torch.Tensor, + bias_swapped: torch.Tensor, + ) -> torch.Tensor: + """Return ``log a`` per pair for umbrella exchange. + + Parameters + ---------- + pairs: + Neighbouring state pairs. + walker_of_state: + State id to the walker row currently holding it. + bias_current: + Reduced bias potential per walker under its current state, + shape ``[B]``. + bias_swapped: + Reduced bias potential per walker under the proposed state, + shape ``[B]``. + + Returns + ------- + torch.Tensor + ``log a`` per pair, shape ``[len(pairs)]``, capped at zero. + """ + values = [] + for state_i, state_j in pairs: + walker_i = walker_of_state[state_i] + walker_j = walker_of_state[state_j] + before = bias_current[walker_i] + bias_current[walker_j] + after = bias_swapped[walker_i] + bias_swapped[walker_j] + values.append(before - after) + return torch.clamp(torch.stack(values), max=0.0) + + def decide( + self, + segment: int, + state_ids: torch.Tensor, + energies: torch.Tensor, + bias_current: torch.Tensor | None = None, + bias_swapped: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[tuple[int, int]], torch.Tensor]: + """Decide this segment's swaps and return the new assignment. + + Pure with respect to the batch: it takes energies and returns a + permutation, so the caller owns every side effect (rebinding the + integrator, rescaling velocities, re-priming forces). That keeps the + acceptance rule testable against hand-computed numbers. + + Parameters + ---------- + segment: + Exchange segment index, which selects the even/odd pairing. + state_ids: + Current assignment, shape ``[B]``. + energies: + Per-walker potential energy ``U`` in eV, shape ``[B]``. Used by + temperature acceptance. + bias_current: + Reduced bias potential per walker under its current state, shape + ``[B]``. Required for umbrella acceptance. + bias_swapped: + Reduced bias potential per walker under its proposed state, shape + ``[B]``. Required for umbrella acceptance. + + Returns + ------- + tuple[torch.Tensor, list[tuple[int, int]], torch.Tensor] + The new state assignment ``[B]``, the pairs attempted, and the + boolean accept mask over those pairs. + + Raises + ------ + ValueError + If *state_ids* is not a permutation of the ladder, or if umbrella + acceptance is in force but the bias energies were not supplied. + """ + # Validate before touching any counter. Pairing looks up "which walker + # holds state k", which a duplicate or short assignment answers with a + # bare KeyError — and by then attempts/pair_attempts have already been + # incremented, leaving the tallies corrupted by a call that failed. + ids = self.validate_assignment(state_ids, source="state_ids") + pairs = self.pair_schedule(segment) + walker_of_state = {int(state): row for row, state in enumerate(ids.tolist())} + + if not pairs: + empty = torch.zeros(0, dtype=torch.bool, device=ids.device) + return ids.clone(), pairs, empty + + if self._acceptance == "temperature": + log_alpha = self._log_acceptance_temperature( + pairs, walker_of_state, energies.reshape(-1) + ) + else: + if bias_current is None or bias_swapped is None: + raise ValueError( + "ReplicaExchange: umbrella acceptance needs the bias " + "energy under both the current and the proposed state " + "assignment, but one was not supplied." + ) + log_alpha = self._log_acceptance_umbrella( + pairs, + walker_of_state, + bias_current.reshape(-1), + bias_swapped.reshape(-1), + ) + + uniforms = self._uniforms(len(pairs), log_alpha.device) + accepted = log_acceptance_is_accepted(log_alpha, uniforms) + + new_ids = ids.clone() + for index, ((state_i, state_j), take) in enumerate( + zip(pairs, accepted.tolist(), strict=True) + ): + self.attempts += 1 + self.pair_attempts[state_i] += 1 + if take: + walker_i = walker_of_state[state_i] + walker_j = walker_of_state[state_j] + new_ids[walker_i] = state_j + new_ids[walker_j] = state_i + self.accepted += 1 + self.pair_accepted[state_i] += 1 + del index + self.exchange_id += 1 + return new_ids, pairs, accepted + + def proposed_assignment( + self, segment: int, state_ids: torch.Tensor + ) -> torch.Tensor: + """Return the assignment that would result if every pair swapped. + + Umbrella acceptance needs the bias evaluated under the proposed + labels *before* the decision is made, so the caller needs the + proposal separately from the outcome. + + Parameters + ---------- + segment: + Exchange segment index. + state_ids: + Current assignment, shape ``[B]``. + + Returns + ------- + torch.Tensor + The all-swaps-accepted assignment, shape ``[B]``. + + Raises + ------ + ValueError + If *state_ids* is not a permutation of the ladder. + """ + ids = self.validate_assignment(state_ids, source="state_ids") + walker_of_state = {int(state): row for row, state in enumerate(ids.tolist())} + proposed = ids.clone() + for state_i, state_j in self.pair_schedule(segment): + walker_i = walker_of_state[state_i] + walker_j = walker_of_state[state_j] + proposed[walker_i] = state_j + proposed[walker_j] = state_i + return proposed + + # ------------------------------------------------------------------ + # State + # ------------------------------------------------------------------ + + def config_fingerprint(self) -> dict[str, Any]: + """Return the configuration a restored run must match. + + Exchange semantics live in the ladder, not in the counters: the + temperatures set the acceptance exponent, the rule follows from them, + and the interval sets the segment cadence. Restoring a run into a + different ladder would keep the counters and the walker assignment + while silently changing what a swap *means* — the walkers would carry + on labelled 0..S-1 against temperatures they were never sampled at. + + ``initial_state_ids`` is deliberately excluded. It seeds the + assignment only when the batch does not already carry one, and a + restored batch always does, so it has no effect after step zero and + would be a false mismatch. + + Returns + ------- + dict[str, Any] + JSON-representable configuration. + """ + return { + "mode": self.mode, + "acceptance": self._acceptance, + "attempt_interval": int(self.attempt_interval), + "temperatures": [float(state.temperature) for state in self.states], + } + + @staticmethod + def describe_config_mismatch( + saved: Mapping[str, Any] | None, actual: Mapping[str, Any] | None + ) -> list[str]: + """Return human-readable differences between two fingerprints. + + Parameters + ---------- + saved: + The checkpoint's fingerprint, or ``None`` when it had no exchange. + actual: + The live runner's fingerprint, or ``None`` when it has none. + + Returns + ------- + list[str] + One line per difference; empty when they agree. + """ + if saved is None and actual is None: + return [] + if saved is None: + return [ + " exchange: the checkpoint was written without replica " + "exchange, but this runner has one configured" + ] + if actual is None: + return [ + " exchange: the checkpoint was written with replica exchange " + f"({saved.get('acceptance')}, " + f"{len(saved.get('temperatures', []))} states), but this " + "runner has replica_exchange=None" + ] + + problems: list[str] = [ + f" exchange {key}: checkpoint has {saved.get(key)!r}, " + f"this runner has {actual.get(key)!r}" + for key in ("mode", "acceptance", "attempt_interval") + if saved.get(key) != actual.get(key) + ] + saved_temps = [float(t) for t in saved.get("temperatures", [])] + actual_temps = [float(t) for t in actual.get("temperatures", [])] + if len(saved_temps) != len(actual_temps): + problems.append( + f" exchange ladder: checkpoint has {len(saved_temps)} state(s), " + f"this runner has {len(actual_temps)}" + ) + elif any( + not math.isclose(a, b, rel_tol=1e-9, abs_tol=1e-9) + for a, b in zip(saved_temps, actual_temps, strict=True) + ): + problems.append( + f" exchange temperatures: checkpoint has {saved_temps}, " + f"this runner has {actual_temps}" + ) + return problems + + def state_dict(self) -> dict[str, Any]: + """Return exchange state for checkpointing. + + Returns + ------- + dict[str, Any] + Counters and the acceptance-RNG position. The position is two + integers rather than a generator blob; see :meth:`_uniforms`. + """ + return { + "exchange_id": int(self.exchange_id), + "attempts": int(self.attempts), + "accepted": int(self.accepted), + "random_seed": int(self.random_seed), + "pair_attempts": list(self.pair_attempts), + "pair_accepted": list(self.pair_accepted), + # Carried so the component is self-describing: loading it into a + # different ladder is refused rather than silently accepted. + "config": self.config_fingerprint(), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore exchange state produced by :meth:`state_dict`. + + Parameters + ---------- + state: + The mapping previously returned by :meth:`state_dict`. + + Raises + ------ + ValueError + If the saved configuration disagrees with this instance's. + """ + saved_config = state.get("config") + if saved_config is not None: + problems = self.describe_config_mismatch( + saved_config, self.config_fingerprint() + ) + if problems: + detail = "\n".join(problems) + raise ValueError( + "ReplicaExchange.load_state_dict: the saved exchange was " + f"configured differently:\n{detail}\n" + "Counters and the acceptance-RNG position only mean " + "anything against the ladder they were produced on." + ) + + self.exchange_id = int(state.get("exchange_id", 0)) + self.attempts = int(state.get("attempts", 0)) + self.accepted = int(state.get("accepted", 0)) + # random_seed is RNG *position* and is restored; attempt_interval is + # configuration and is validated above, never silently overwritten. + self.random_seed = int(state.get("random_seed", self.random_seed)) + pair_attempts = state.get("pair_attempts") + if pair_attempts is not None: + self.pair_attempts = [int(v) for v in pair_attempts] + pair_accepted = state.get("pair_accepted") + if pair_accepted is not None: + self.pair_accepted = [int(v) for v in pair_accepted] + + @property + def acceptance_rate(self) -> float: + """Return the fraction of attempted pair swaps that were accepted.""" + return self.accepted / self.attempts if self.attempts else 0.0 + + def pair_acceptance_rates(self) -> list[float]: + """Return the per-neighbouring-pair acceptance rates. + + A REMD ladder is tuned on these: a pair far below the others is a gap + the walkers cannot cross, and the ladder needs another rung there. + + Returns + ------- + list[float] + One rate per neighbouring pair, ``0.0`` where never attempted. + """ + return [ + (accepted / attempts if attempts else 0.0) + for accepted, attempts in zip( + self.pair_accepted, self.pair_attempts, strict=True + ) + ] + + def __repr__(self) -> str: + """Return a concise description of the exchange.""" + return ( + f"{type(self).__name__}(states={len(self.states)}, " + f"acceptance={self._acceptance!r}, " + f"attempt_interval={self.attempt_interval}, " + f"accepted={self.accepted}/{self.attempts})" + ) + + +def log_acceptance_is_accepted( + log_alpha: torch.Tensor, uniforms: torch.Tensor +) -> torch.Tensor: + """Return the accept mask for the given log-acceptance values. + + ``log_alpha`` is capped at zero, so ``log_alpha == 0`` means accept with + probability one. Comparing ``log(u) < log_alpha`` rather than + ``u < exp(log_alpha)`` keeps a very negative ``log_alpha`` from + underflowing to exactly zero and turning a rare-but-possible swap into an + impossible one. + + Parameters + ---------- + log_alpha : torch.Tensor + Log acceptance probability per pair, ``<= 0``. + uniforms : torch.Tensor + Draws in ``[0, 1)``, same shape. + + Returns + ------- + torch.Tensor + Boolean accept mask. + """ + safe = torch.clamp(uniforms, min=torch.finfo(uniforms.dtype).tiny) + return torch.log(safe) < log_alpha diff --git a/nvalchemi/enhanced_sampling/_runner.py b/nvalchemi/enhanced_sampling/_runner.py new file mode 100644 index 00000000..6c9f505b --- /dev/null +++ b/nvalchemi/enhanced_sampling/_runner.py @@ -0,0 +1,1414 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The ``EnhancedSampling`` runner: orchestration around an existing dynamics. + +The runner owns what a bias cannot: walker identity, the ordering of the +force step, exactly-once ``update()`` delivery, and force priming after a +bias changes. Integration itself is delegated entirely to the wrapped +``BaseDynamics`` — the runner never touches an integrator. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, Any + +import torch + +from nvalchemi.dynamics.base import DynamicsStage +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling._bias import BiasResult, aggregate_bias_results +from nvalchemi.enhanced_sampling._checkpoint import ( + CheckpointManifest, + _qualified_name, + read_checkpoint, + write_checkpoint, +) +from nvalchemi.enhanced_sampling._exchange import ReplicaExchange + +if TYPE_CHECKING: + from collections.abc import Mapping + from enum import Enum + from pathlib import Path + + from nvalchemi.data import Batch + from nvalchemi.dynamics.base import BaseDynamics + from nvalchemi.enhanced_sampling._bias import BiasPotential + from nvalchemi.hooks import HookContext + +__all__ = ["EnhancedSampling"] + + +class _BiasCompositeHook: + """The single internal hook the runner installs on the dynamics. + + An implementation detail, not a public API. It defines + ``_runs_on_stage`` so the registry lets one hook object serve three + stages, which is what keeps the ordering guarantees in one place instead + of spread across three separately-registered hooks whose relative order + would then depend on registration sequence. + """ + + def __init__(self, runner: EnhancedSampling) -> None: + self._runner = runner + self.stage: Enum | None = None + self.frequency = 1 + + def _runs_on_stage(self, stage: Enum) -> bool: + """Return whether this hook fires at *stage*. + + Parameters + ---------- + stage: + The stage being dispatched. + + Returns + ------- + bool + ``True`` for the three stages the runner needs. + """ + return stage in ( + DynamicsStage.BEFORE_STEP, + DynamicsStage.AFTER_COMPUTE, + DynamicsStage.AFTER_STEP, + ) + + def __call__(self, ctx: HookContext, stage: Enum) -> None: + """Dispatch to the runner phase for *stage*. + + Parameters + ---------- + ctx: + The dynamics hook context. + stage: + The stage being dispatched. + """ + if stage is DynamicsStage.BEFORE_STEP: + self._runner._stamp_identity(ctx.batch) + elif stage is DynamicsStage.AFTER_COMPUTE: + self._runner._evaluate_and_apply(ctx.batch) + elif stage is DynamicsStage.AFTER_STEP: + self._runner._observe_and_update(ctx.batch) + + +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 — what pure temperature + replica exchange needs, since the ladder alone drives the sampling. + 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. + + Raises + ------ + TypeError + If any value in *biases* does not satisfy :class:`BiasPotential`. + ValueError + If ``steps_per_epoch`` is below 1, or a bias's ``name`` disagrees + with its key in *biases*. + + Examples + -------- + >>> sampling = EnhancedSampling( # doctest: +SKIP + ... dynamics=md, + ... biases={"umbrella": umbrella, "wall": lower_wall}, + ... ) + >>> batch = sampling.run(batch, n_steps=1000) # doctest: +SKIP + + Notes + ----- + Hook ordering + The composite hook is inserted at the **front** of the dynamics hook + list, so that at ``AFTER_COMPUTE`` the bias contribution is applied + before any other hook runs. A safety hook such as + ``MaxForceClampHook`` therefore clamps the *total* force, which is + the physically meaningful quantity, rather than the model force + alone. Bias observations that need unbiased physical forces are + captured inside the runner before the contribution is applied, so + they are unaffected by this ordering. + """ + + def __init__( + self, + dynamics: BaseDynamics, + biases: Mapping[str, BiasPotential] | None = None, + *, + steps_per_epoch: int = 10_000, + compile_biases: bool = False, + prime_after_update: bool = True, + replica_exchange: ReplicaExchange | None = None, + ) -> None: + from nvalchemi.enhanced_sampling._bias import BiasPotential as _Protocol + + if int(steps_per_epoch) < 1: + raise ValueError( + f"EnhancedSampling: steps_per_epoch must be at least 1, got " + f"{steps_per_epoch}. It is a divisor — the epoch index is " + "step // steps_per_epoch and the checkpoint boundary is " + "step % steps_per_epoch — so zero raises deep in a run and a " + "negative value makes both meaningless." + ) + self.dynamics = dynamics + self.biases: dict[str, BiasPotential] = dict(biases or {}) + self.steps_per_epoch = int(steps_per_epoch) + self.prime_after_update = bool(prime_after_update) + self.replica_exchange = replica_exchange + + for key, bias in self.biases.items(): + if not isinstance(bias, _Protocol): + raise TypeError( + f"EnhancedSampling: biases[{key!r}] is a " + f"{type(bias).__name__}, which does not satisfy the " + "BiasPotential protocol (needs a 'name' attribute and an " + "'evaluate(batch) -> BiasResult' method)." + ) + if getattr(bias, "name", None) != key: + raise ValueError( + f"EnhancedSampling: biases[{key!r}] has name=" + f"{getattr(bias, 'name', None)!r}. The key and the bias " + "name must agree — both are used as identifiers, in the " + "output dict and in checkpoint group names respectively." + ) + + if replica_exchange is not None: + replica_exchange.validate_for(self.biases) + self._validate_exchange_capability() + + if compile_biases: + self._compile_bias_energies() + + # Diagnostics from the most recent force evaluation: + # physical/ model only, before any bias + # bias// one bias's contribution + # bias_total/ the sum across biases + # total/ physical + bias, read back from the batch + self.last_outputs: dict[str, torch.Tensor] = {} + + # Per-bias observation captured at its observation_stage, consumed by + # the next update() call. + self._pending: dict[str, tuple[Batch, BiasResult]] = {} + + # Per-bias results from the most recent force evaluation. Held because + # an AFTER_STEP capture happens after that evaluation has returned, + # and update() is documented to receive the result its bias produced + # during it. + self._last_results: dict[str, BiasResult] = {} + self._last_update_step: dict[str, int] = {} + self._last_seen_version: dict[str, int] = {} + self._sync_seen_versions() + self._physical: dict[str, torch.Tensor] = {} + self._next_walker_id = 0 + self._last_epoch = -1 + self._committed_epoch = -1 + self._last_segment = -1 + self._attempted_segment = -1 + # One-shot: the walker/state bijection is checked on the first stamp. + self._validated_assignment = False + # Set while prime_forces runs. Priming evaluates forces at fixed + # coordinates; it must not also advance the sampling state. + self._priming = False + # One-shot: the empirical state-dependence probe runs at prime time. + self._probed_state_dependence = False + self._restored = False + # Set on every stamp so checkpoint() matches the proposal's + # `sampling.checkpoint(path)` signature without a batch argument. + self._current_batch: Batch | None = None + + self._hook = _BiasCompositeHook(self) + dynamics.register_hook(self._hook, stage=DynamicsStage.AFTER_COMPUTE) + # Move to the front: see the "Hook ordering" note in the class docstring. + dynamics.hooks.remove(self._hook) + dynamics.hooks.insert(0, self._hook) + + # ------------------------------------------------------------------ + # Setup helpers + # ------------------------------------------------------------------ + + def _compile_bias_energies(self) -> None: + """Compile each conservative bias's ``energy()`` in place. + + Assigning the compiled callable as an instance attribute shadows the + bound method, so ``evaluate()`` picks it up without an indirection + the eager path would also pay. + """ + for bias in self.biases.values(): + if hasattr(bias, "energy"): + bias.energy = torch.compile(bias.energy) # type: ignore[method-assign] + + def _validate_exchange_capability(self) -> None: + """Reject a dynamics that cannot rebind a thermodynamic state. + + An accepted swap must change the target temperature, rescale + velocities, and transform any thermostat memory as one indivisible + move. An integrator that only accepts the new label would keep + sampling the old temperature, which breaks detailed balance with no + symptom the run would show — so this fails at construction rather + than producing a plausible-looking wrong trajectory. + + Raises + ------ + TypeError + If the dynamics does not implement the rebinding adapters. + """ + for method in ("apply_thermodynamic_state", "rescale_velocities_for_state"): + if not callable(getattr(self.dynamics, method, None)): + raise TypeError( + f"EnhancedSampling: replica exchange needs " + f"{type(self.dynamics).__name__} to implement {method}(), " + "so an accepted swap can rebind temperature, velocities, " + "and thermostat state together. NVTLangevin and " + "NVTNoseHoover implement this; other integrators can run " + "biased dynamics without exchange." + ) + # BaseDynamics defines apply_thermodynamic_state only to raise, so + # presence is not enough — probe it. + try: + self.dynamics.apply_thermodynamic_state( + torch.zeros(0, dtype=torch.long), torch.zeros(0) + ) + except NotImplementedError as exc: + raise TypeError( + f"EnhancedSampling: {type(self.dynamics).__name__} does not " + "support thermodynamic-state rebinding, so it cannot take part " + "in replica exchange." + ) from exc + except Exception: # noqa: S110 - any other failure means it is implemented + pass + + def _reduced_bias_energy( + self, batch: Batch, state_ids: torch.Tensor + ) -> torch.Tensor: + """Return the reduced bias potential per walker under *state_ids*. + + Umbrella acceptance needs the bias evaluated under both the current + and the proposed labels, so the assignment is swapped in, the biases + re-evaluated, and the original restored in a ``finally``. + + Parameters + ---------- + batch: + The live batch. + state_ids: + Assignment to evaluate under, shape ``[B]``. + + Returns + ------- + torch.Tensor + ``beta * E_bias`` per walker, shape ``[B]``. + """ + exchange = self.replica_exchange + original = batch.thermodynamic_state_id + try: + batch["thermodynamic_state_id"] = state_ids + total = torch.zeros( + batch.num_graphs, + dtype=batch.positions.dtype, + device=batch.positions.device, + ) + for bias in self.biases.values(): + energy = bias.evaluate(batch).energy + if energy is not None: + total = total + energy.reshape(-1) + finally: + batch["thermodynamic_state_id"] = original + + temperatures = exchange.temperatures.to(total.device, total.dtype) + beta = 1.0 / (KB_EV * temperatures[state_ids.reshape(-1).to(torch.long)]) + return beta * total + + def _probe_state_dependence(self, batch: Batch) -> None: + """Reject a state-dependent bias under a temperature ladder. + + The declaration checked at construction only covers biases that know + to declare — ``HarmonicUmbrellaBias`` does, an arbitrary user bias + does not. This probes instead of asking: evaluate every bias under + the current assignment and under a rotated one, at identical + coordinates. A bias whose energy is independent of the assignment + returns the same number twice; one that reads + ``thermodynamic_state_id`` does not. + + Temperature acceptance uses only ``U``, so a bias whose energy varies + with the assignment contributes cross-state terms that the rule never + computes — detailed balance would be wrong with no symptom. Runs + once, at prime time, before any exchange has been attempted. + + Parameters + ---------- + batch: + The live batch, already carrying a state assignment. + + Raises + ------ + ValueError + If any bias's energy changes when the assignment is permuted. + """ + exchange = self.replica_exchange + if ( + exchange is None + or exchange.acceptance != "temperature" + or not self.biases + or self._probed_state_dependence + ): + return + self._probed_state_dependence = True + + current = batch.thermodynamic_state_id.reshape(-1).to(torch.long) + if current.numel() < 2: + return + # Rotate by one: every walker sees a different state, and the result + # is still a permutation, so a per-state lookup stays in range. + rotated = torch.roll(current, shifts=1) + + original = batch.thermodynamic_state_id + try: + offenders: list[str] = [] + for name, bias in self.biases.items(): + batch["thermodynamic_state_id"] = current + before = bias.evaluate(batch).energy + batch["thermodynamic_state_id"] = rotated + after = bias.evaluate(batch).energy + if before is None or after is None: + continue + if not torch.allclose(before, after, rtol=1e-9, atol=1e-12): + offenders.append(name) + finally: + batch["thermodynamic_state_id"] = original + + if offenders: + raise ValueError( + f"EnhancedSampling: bias(es) {sorted(offenders)} produce a " + "different energy when the thermodynamic-state assignment is " + "permuted, so they depend on the state — but the ladder varies " + "temperature, which selects temperature acceptance. That rule " + "uses only the total energy and omits the cross-state bias " + "terms, so detailed balance would be wrong with nothing to " + "show for it. The combined acceptance rule is not implemented: " + "use a state-independent bias with a temperature ladder, or " + "equal temperatures with a multi-window bias." + ) + + def _attempt_segment(self, batch: Batch, segment: int) -> None: + """Attempt *segment*'s pairs, at most once. + + Idempotent by segment index, for the same reason + :meth:`_commit_epoch` is: an accepted swap re-primes forces, which + runs the identity stamp again, and a nested pass must not decide the + same segment twice. + + Parameters + ---------- + batch: + The live batch. + segment: + The completed segment whose pairs are due. Negative, or already + attempted, is a no-op. + """ + if segment < 0 or segment <= self._attempted_segment: + return + self._attempted_segment = segment + self._attempt_exchange(batch, segment) + + def _attempt_exchange(self, batch: Batch, segment: int) -> None: + """Attempt one round of swaps and apply the accepted ones. + + Applying is the indivisible half: the batch labels, the integrator's + target temperature, the velocity rescaling, and the forces all move + together. Leaving any of them behind would sample a state the + assignment says the walker is no longer in. + + Parameters + ---------- + batch: + The live batch. + segment: + Exchange segment index, which selects the even/odd pairing. + """ + exchange = self.replica_exchange + if exchange is None: + return + current = batch.thermodynamic_state_id.reshape(-1).to(torch.long) + + bias_current = bias_swapped = None + if exchange.acceptance == "umbrella": + proposed = exchange.proposed_assignment(segment, current) + bias_current = self._reduced_bias_energy(batch, current) + bias_swapped = self._reduced_bias_energy(batch, proposed) + + energies = getattr(batch, "energy", None) + if energies is None: + energies = torch.zeros(batch.num_graphs, device=batch.positions.device) + + new_ids, _pairs, accepted = exchange.decide( + segment, + current, + energies.reshape(-1), + bias_current=bias_current, + bias_swapped=bias_swapped, + ) + if not bool(accepted.any()): + return + + batch["thermodynamic_state_id"] = new_ids + self.dynamics.apply_thermodynamic_state(new_ids, exchange.temperatures) + self.dynamics.rescale_velocities_for_state(batch) + # Forces in the batch were produced under the previous labels. The + # integrator reads them in its next half-step before any model call, + # so without re-priming the first step after a swap would integrate + # the state the walker just left. + self.prime_forces(batch) + + def _sync_seen_versions(self) -> None: + """Re-baseline the cached per-bias state versions from the live biases. + + :attr:`_last_seen_version` answers "has this bias changed since the + runner last looked", which is what decides whether forces need + re-priming. It must be re-baselined after anything that changes a + bias's version without the runner observing it — construction, and + :meth:`restore`, which loads a saved version straight onto the bias. + + Called from both, rather than inlined, so the two cannot drift: a + restore that skipped this would leave the cache at ``0`` against a + restored version of ``N``, and the first post-restore ``update()`` + would re-prime on a change that never happened. + """ + self._last_seen_version = { + name: int(getattr(bias, "state_version", 0)) + for name, bias in self.biases.items() + } + + def _adaptive_biases(self) -> dict[str, BiasPotential]: + """Return the biases that implement ``update``. + + Detection is structural (``hasattr``), so a bias satisfies the + adaptive contract without inheriting ``AdaptivePotentialMixin``. + + Returns + ------- + dict[str, BiasPotential] + Name to bias, for adaptive biases only. + """ + return { + name: bias + for name, bias in self.biases.items() + if callable(getattr(bias, "update", None)) + } + + # ------------------------------------------------------------------ + # Phase 1 — identity stamping (BEFORE_STEP) + # ------------------------------------------------------------------ + + def _stamp_identity(self, batch: Batch) -> None: + """Stamp walker identity and counters onto the live batch. + + ``walker_id`` and ``thermodynamic_state_id`` are assigned once and + then preserved; the counters are refreshed every step. + + Parameters + ---------- + batch: + The live batch. + """ + n_graphs = batch.num_graphs + device = batch.positions.device + step = self.dynamics.step_count + + if getattr(batch, "walker_id", None) is None: + batch["walker_id"] = torch.arange( + self._next_walker_id, + self._next_walker_id + n_graphs, + dtype=torch.long, + device=device, + ) + self._next_walker_id += n_graphs + + existing = getattr(batch, "thermodynamic_state_id", None) + if existing is None: + if self.replica_exchange is not None: + # Check before attaching: a ladder-sized tensor on a + # differently-sized batch would otherwise fail as an opaque + # "Length mismatch" from inside the batch storage. + self.replica_exchange.validate_assignment( + self.replica_exchange.initial_state_ids, + n_graphs, + source="initial_state_ids", + ) + batch["thermodynamic_state_id"] = ( + self.replica_exchange.initial_state_ids.to(device) + ) + else: + batch["thermodynamic_state_id"] = torch.zeros( + n_graphs, dtype=torch.long, device=device + ) + elif self.replica_exchange is not None and not self._validated_assignment: + # A batch may arrive carrying its own assignment, which never went + # through the constructor's check. Validate it once — a duplicate + # would surface later as a KeyError from the pair lookup. + self.replica_exchange.validate_assignment( + existing, n_graphs, source="batch.thermodynamic_state_id" + ) + self._validated_assignment = True + + self._current_batch = batch + full = torch.full((n_graphs,), step, dtype=torch.long, device=device) + batch["sampling_step"] = full + batch["sampling_epoch"] = full // self.steps_per_epoch + interval = ( + self.replica_exchange.attempt_interval + if self.replica_exchange is not None + else self.steps_per_epoch + ) + segment = step // interval + batch["exchange_segment"] = torch.full( + (n_graphs,), segment, dtype=torch.long, device=device + ) + + # Ordering at a boundary is exchange, then bias commit — the epoch + # commit publishes shared history, and doing it before the swap would + # publish under labels that are about to change. + # + # Entering segment s means segment s-1 has completed, and it is *that* + # one whose pairs are due — exactly as entering epoch e commits epoch + # e-1 below. Attempting the segment being entered instead would skip + # segment 0's pairs entirely; on a two-state ladder, where segment 0 + # holds the only pair that exists, the first swap would be delayed to + # 2 * attempt_interval. + if self.replica_exchange is not None and not self._priming: + self._attempt_segment(batch, segment - 1) + self._last_segment = segment + + epoch = step // self.steps_per_epoch + if epoch != self._last_epoch: + # Entering epoch e means epoch e-1 has completed. + self._commit_epoch(epoch - 1) + self._last_epoch = epoch + + def _commit_epoch(self, epoch: int) -> None: + """Run every adaptive bias's ``commit_epoch``, at most once per epoch. + + Idempotent by design. The commit is reached from two directions — + lazily, when a step observes that the epoch advanced, and eagerly, + when :meth:`checkpoint` drains a completed epoch — and a + shared-history bias that merged its pending deposits twice would + double-count them. + + Parameters + ---------- + epoch: + Index of the epoch that has completed. Negative, or already + committed, is a no-op. + """ + if epoch < 0 or epoch <= self._committed_epoch: + return + for bias in self._adaptive_biases().values(): + commit = getattr(bias, "commit_epoch", None) + if callable(commit): + commit() + self._committed_epoch = epoch + + # ------------------------------------------------------------------ + # Phase 2 — evaluate and apply (AFTER_COMPUTE) + # ------------------------------------------------------------------ + + def _evaluate_and_apply(self, batch: Batch) -> None: + """Evaluate every bias against the unmodified batch, then apply the sum. + + Implements steps 1-7 of the documented force-step ordering. The + ordering is the whole point: every bias sees the same physical + outputs, so no bias can observe another's contribution and the result + does not depend on registration order. + + Parameters + ---------- + batch: + The live batch, immediately after the model forward pass. + """ + if not self.biases: + return + + # 1. Capture the physical outputs before anything is added. + self._physical = { + key: value.clone() + for key in ("energy", "forces", "stress") + if (value := getattr(batch, key, None)) is not None + } + + # 2-3. Evaluate each bias against the same unmodified batch. + # BiasResult validates itself on construction. + results = {name: bias.evaluate(batch) for name, bias in self.biases.items()} + self._last_results = results + + # 4. Capture AFTER_COMPUTE observations while batch.forces still + # holds unbiased physical forces. ABF depends on this: an + # estimator fed its own output diverges. + self._capture(batch, results, DynamicsStage.AFTER_COMPUTE) + + # 5-6. Namespace observables, then sum once. + total = aggregate_bias_results( + [self._namespace(name, r) for name, r in results.items()] + ) + self._record_diagnostics(results, total) + + # 7. Apply the total contribution, then record the combined result. + self._apply(batch, total, results) + self._record_totals(batch) + + def _namespace(self, name: str, result: BiasResult) -> BiasResult: + """Return *result* with its observables prefixed ``bias//``. + + Namespacing has to happen before aggregation, because + ``aggregate_bias_results`` rejects duplicate observable keys rather + than silently dropping one — two biases of the same type would + otherwise collide on identical names. + + Parameters + ---------- + name: + The bias name. + result: + The bias's result. + + Returns + ------- + BiasResult + A copy with namespaced observables, or *result* unchanged when it + has none. + """ + if not result.observables: + return result + return dataclasses.replace( + result, + observables={ + f"bias/{name}/{key}": value for key, value in result.observables.items() + }, + ) + + def _record_diagnostics( + self, results: dict[str, BiasResult], total: BiasResult + ) -> None: + """Populate :attr:`last_outputs` with the physical and bias views. + + Writes ``physical/*``, ``bias//*``, and ``bias_total/*``. + ``total/*`` is *not* written here: at this point the bias has not + been applied yet, so there is no combined value to record. + :meth:`_record_totals` adds it afterwards. + + Parameters + ---------- + results: + Per-bias results, un-namespaced. + total: + The aggregated bias result — the sum across biases, **not** + physical plus bias. + """ + outputs: dict[str, torch.Tensor] = {} + for key, value in self._physical.items(): + outputs[f"physical/{key}"] = value + for name, result in results.items(): + for key in ("energy", "forces", "stress", "virial"): + value = getattr(result, key) + if value is not None: + outputs[f"bias/{name}/{key}"] = value + for key, value in result.observables.items(): + outputs[f"bias/{name}/{key}"] = value + for key in ("energy", "forces", "stress", "virial"): + value = getattr(total, key) + if value is not None: + outputs[f"bias_total/{key}"] = value + self.last_outputs = outputs + + def _record_totals(self, batch: Batch) -> None: + """Record ``total/*`` — physical plus bias — from the live batch. + + Must run *after* :meth:`_apply`. Reading the batch rather than + adding ``physical/*`` and ``bias_total/*`` back together keeps the + record faithful to what was actually written, including any reshape + :meth:`_apply` performed. + + Parameters + ---------- + batch: + The live batch, immediately after the bias has been applied. + + Notes + ----- + This is the state as the runner leaves it, not necessarily the state + the integrator sees. The runner's hook is deliberately first at + ``AFTER_COMPUTE`` (so a force clamp acts on the total rather than on + the model force alone), which means any later hook at that stage can + still modify ``batch.forces`` afterwards. Read the batch directly if + you need the value the integrator consumed. + """ + for key in ("energy", "forces", "stress"): + value = getattr(batch, key, None) + if value is not None: + self.last_outputs[f"total/{key}"] = value.detach().clone() + + def _apply( + self, + batch: Batch, + total: BiasResult, + results: dict[str, BiasResult] | None = None, + ) -> None: + """Add the aggregated bias contribution to the batch, in place. + + Every non-``None`` output must have a destination buffer. Skipping a + field whose buffer is absent would discard that contribution in + silence — for ``stress`` that is precisely the barostat-invisibility + failure this API exists to remove, arrived at from a different + direction: the bias is computed correctly, applied nowhere, and the + cell evolves as though it did not exist. + + Parameters + ---------- + batch: + The live batch. + total: + The aggregated bias result. + results: + Per-bias results, used only to name the contributors in an error. + + Raises + ------ + ValueError + If the aggregate carries a virial, or if any non-``None`` output + has no destination buffer on the batch. + """ + if total.virial is not None: + raise ValueError( + "EnhancedSampling: a bias returned 'virial', but the runner " + "applies 'stress' to the batch. Convert W -> sigma = -W/V in " + "the bias before returning it; the cell volume is the bias's " + "to supply." + ) + self._check_destinations(batch, total, results or {}) + with torch.no_grad(): + if total.energy is not None: + batch.energy.add_(total.energy.reshape(batch.energy.shape)) + if total.forces is not None: + batch.forces.add_(total.forces) + if total.stress is not None: + batch.stress.add_(total.stress.reshape(batch.stress.shape)) + + @staticmethod + def _check_destinations( + batch: Batch, total: BiasResult, results: dict[str, BiasResult] + ) -> None: + """Raise if any produced output has nowhere to go on the batch. + + Parameters + ---------- + batch: + The live batch. + total: + The aggregated bias result. + results: + Per-bias results, used to name which biases produced each field. + + Raises + ------ + ValueError + Listing every missing destination, the biases responsible, and + both ways to resolve it. + """ + _ALLOCATION_HINT = { + "energy": "energy=torch.zeros(1, 1)", + "forces": "forces=torch.zeros(n_atoms, 3)", + "stress": "stress=torch.zeros(1, 3, 3)", + } + missing = [ + key + for key in ("energy", "forces", "stress") + if getattr(total, key) is not None and getattr(batch, key, None) is None + ] + if not missing: + return + + lines = [] + for key in missing: + contributors = sorted( + name + for name, result in results.items() + if getattr(result, key, None) is not None + ) + who = f" (from {contributors})" if contributors else "" + lines.append(f" '{key}'{who}: add {_ALLOCATION_HINT[key]} to AtomicData") + detail = "\n".join(lines) + extra = "" + if "stress" in missing: + extra = ( + "\nA bias that produces stress with nowhere to put it is " + "invisible to an NPT/NPH barostat — the cell would evolve as " + "if the bias were absent. If this run genuinely has no use for " + "a cell response (NVE/NVT), pass compute_stress=False to those " + "biases instead of leaving the output to be discarded." + ) + raise ValueError( + f"EnhancedSampling: bias output has no destination buffer on the " + f"batch, so it would be silently discarded:\n{detail}{extra}" + ) + + # ------------------------------------------------------------------ + # Phase 3 — observe and update (AFTER_STEP) + # ------------------------------------------------------------------ + + def _capture( + self, + batch: Batch, + results: dict[str, BiasResult], + stage: DynamicsStage, + ) -> None: + """Snapshot the batch for adaptive biases observing at *stage*. + + The stored pair is exactly what :meth:`AdaptivePotentialMixin.update` + is documented to receive: the frame at the bias's + ``observation_stage``, and the :class:`BiasResult` that bias returned + during the preceding force evaluation. An ``AFTER_STEP`` capture + happens after that evaluation has returned, so the results are read + from :attr:`_last_results` rather than recomputed — a metadynamics + bias sizing its next hill from the bias energy it just applied needs + the real value, not an empty placeholder. + + Parameters + ---------- + batch: + The live batch. + results: + Per-bias results from the preceding force evaluation. + stage: + The stage being captured. + """ + step = self.dynamics.step_count + for name, bias in self._adaptive_biases().items(): + if ( + getattr(bias, "observation_stage", DynamicsStage.AFTER_STEP) + is not stage + ): + continue + if step % max(1, getattr(bias, "update_frequency", 1)) != 0: + continue + self._pending[name] = (batch.clone(), results.get(name, BiasResult())) + + def _observe_and_update(self, batch: Batch) -> None: + """Capture post-step frames, deliver ``update()``, then re-prime. + + Implements steps 10-12 of the force-step ordering. + + Parameters + ---------- + batch: + The live batch, after the integrator has finished. + """ + adaptive = self._adaptive_biases() + if not adaptive: + return + + step = self.dynamics.step_count + self._capture(batch, self._last_results, DynamicsStage.AFTER_STEP) + + changed = False + for name, bias in adaptive.items(): + if step % max(1, getattr(bias, "update_frequency", 1)) != 0: + continue + # Exactly once per step, even if this hook is dispatched twice. + if self._last_update_step.get(name) == step: + continue + frames, result = self._pending.pop( + name, (batch, self._last_results.get(name, BiasResult())) + ) + bias.update(frames, result) # type: ignore[attr-defined] + self._last_update_step[name] = step + + version = getattr(bias, "state_version", 0) + if version != self._last_seen_version.get(name, 0): + self._last_seen_version[name] = version + changed = True + + if changed and self.prime_after_update: + self._reprime(batch) + + def _reprime(self, batch: Batch) -> None: + """Rewrite total forces from cached physical outputs and current biases. + + A bias that just deposited a hill leaves ``batch.forces`` describing + the bias as it was *before* the deposition. Anything reading the + batch between steps — a reporter, a convergence check — would see + stale values. This restores the cached physical outputs and re-adds + a freshly evaluated bias contribution. + + The physical part is reused rather than recomputed: this is + *evaluate-only* priming, so it costs one bias evaluation and no model + forward pass. The physical forces are therefore the ones from the + start of the step, not from the post-step coordinates. That is exact + only if the model forward were repeated, which is precisely the cost + this avoids; the next step recomputes them anyway. + + Parameters + ---------- + batch: + The live batch. + """ + if not self._physical: + return + with torch.no_grad(): + for key, value in self._physical.items(): + target = getattr(batch, key, None) + if target is not None: + target.copy_(value.reshape(target.shape)) + results = {name: bias.evaluate(batch) for name, bias in self.biases.items()} + self._last_results = results + total = aggregate_bias_results( + [self._namespace(name, r) for name, r in results.items()] + ) + self._record_diagnostics(results, total) + self._apply(batch, total, results) + self._record_totals(batch) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def prime_forces(self, batch: Batch) -> Batch: + """Run one evaluate-only force evaluation without advancing dynamics. + + Populates ``batch.energy`` / ``forces`` / ``stress`` with the total + (physical plus bias) values at the current coordinates. Needed + before the first step of a warm-started run, where a consumer may + read forces before any step has happened. + + Parameters + ---------- + batch: + The batch to prime. + + Returns + ------- + Batch + The same batch, primed in place. + + Does not advance sampling state: no exchange is attempted and no + epoch is committed from here, so priming after a restore leaves the + run exactly where the checkpoint left it. + + Raises + ------ + ValueError + If the batch has no ``forces`` field to write into. + + Notes + ----- + ``compute()`` writes its outputs back with ``copy_`` into fields that + must already exist — a model output whose batch field is absent is + silently discarded rather than created. That is the toolkit's + contract (``AtomicData(..., forces=torch.zeros(n, 3))``), so the + check below turns a silent no-op into a named error. + + This reproduces the prefix of ``BaseDynamics.step`` up to the model + call, including the ``BEFORE_COMPUTE`` hooks. Those are not + optional: a cutoff model reaches ``adapt_input`` expecting a neighbor + list that ``NeighborListHook`` builds at exactly that stage, so + calling ``compute()`` bare would fail on the first primed evaluation. + """ + if getattr(batch, "forces", None) is None: + raise ValueError( + "EnhancedSampling.prime_forces: batch has no 'forces' field. " + "Model outputs are written back in place, so the buffer must " + "exist first — construct AtomicData with " + "forces=torch.zeros(n_atoms, 3) and energy=torch.zeros(1, 1), " + "plus stress=torch.zeros(1, 3, 3) whenever a bias produces " + "stress (any periodic batch, unless the bias was built with " + "compute_stress=False)." + ) + self.dynamics._ensure_state_initialized(batch) + was_priming = self._priming + self._priming = True + try: + self._stamp_identity(batch) + self.dynamics._call_hooks(DynamicsStage.BEFORE_COMPUTE, batch) + self.dynamics.compute(batch) + self._evaluate_and_apply(batch) + self._probe_state_dependence(batch) + finally: + self._priming = was_priming + return batch + + def warm_start(self, frames: Batch) -> None: + """Replay prior frames into every adaptive bias, in order. + + Approximate by construction: it reconstructs bias history but not + velocities, RNG, or integrator state. Use + :meth:`restore` when exact reproducibility matters. + + Parameters + ---------- + frames: + Prior frames in chronological order, one graph per frame. + + Raises + ------ + RuntimeError + If called after :meth:`restore`; the two are mutually exclusive + and applying a warm start over a restored state would silently + corrupt it. + """ + if getattr(self, "_restored", False): + raise RuntimeError( + "EnhancedSampling: warm_start() and restore() are mutually " + "exclusive. This runner has already been restored from a " + "checkpoint; warm-starting over it would replay history the " + "restored state already contains." + ) + adaptive = self._adaptive_biases() + if not adaptive: + return + for index in range(frames.num_graphs): + frame = frames.index_select( + torch.tensor([index], device=frames.positions.device) + ) + for bias in adaptive.values(): + bias.update(frame, BiasResult()) # type: ignore[attr-defined] + + def run( + self, batch: Batch, n_steps: int | None = None, *, prime: bool = True + ) -> Batch: + """Run biased dynamics. + + Primes forces first by default. A velocity-Verlet-style integrator + reads ``batch.forces`` in its *first* half-step, before any model + call — so without priming, step 0 integrates against whatever the + buffer happened to hold (zeros, for a freshly built batch), making it + the one step in the run that ignores the bias. + + Parameters + ---------- + batch: + The initial batch. + n_steps: + Number of steps; falls back to the dynamics' own ``n_steps``. + prime: + Set ``False`` to skip priming when the caller has already + evaluated forces at these coordinates. + + Returns + ------- + Batch + The batch after all steps. + """ + if prime: + self.prime_forces(batch) + return self.dynamics.run(batch, n_steps=n_steps) + + def _components(self) -> dict[str, dict[str, Any]]: + """Collect every component's state for a checkpoint. + + Returns + ------- + dict[str, dict[str, Any]] + Component name to state mapping. + """ + components: dict[str, dict[str, Any]] = { + "dynamics": dict(self.dynamics.state_dict()), + "runner": { + "steps_per_epoch": self.steps_per_epoch, + "next_walker_id": self._next_walker_id, + "last_epoch": self._last_epoch, + "committed_epoch": self._committed_epoch, + "last_segment": self._last_segment, + "attempted_segment": self._attempted_segment, + "last_update_step": { + name: int(step) for name, step in self._last_update_step.items() + }, + }, + } + for name, bias in self.biases.items(): + getter = getattr(bias, "state_dict", None) + if callable(getter): + components[f"biases/{name}"] = dict(getter()) + if self.replica_exchange is not None: + components["exchange"] = dict(self.replica_exchange.state_dict()) + return components + + def checkpoint(self, path: str | Path, batch: Batch | None = None) -> None: + """Write a transactional checkpoint at a consistency-epoch boundary. + + The boundary is not a convention — it is the only point where there + are no pending ``update()`` calls and no in-flight epoch commit, so a + checkpoint taken anywhere else could capture a bias mid-mutation. + + Parameters + ---------- + path: + Destination Zarr store. + batch: + The batch to save. Defaults to the one last seen by the runner. + + Raises + ------ + RuntimeError + If no batch is available, meaning nothing has been run or primed. + ValueError + If the current step is not an epoch boundary; the message names + the next valid step. + """ + target = batch if batch is not None else self._current_batch + if target is None: + raise RuntimeError( + "EnhancedSampling.checkpoint: no batch to save. Run or prime " + "the sampler first, or pass batch= explicitly." + ) + + step = self.dynamics.step_count + if step % self.steps_per_epoch != 0: + next_step = ((step // self.steps_per_epoch) + 1) * self.steps_per_epoch + raise ValueError( + f"EnhancedSampling.checkpoint: step {step} is not a consistency " + f"epoch boundary (steps_per_epoch={self.steps_per_epoch}). The " + f"next valid checkpoint step is {next_step}. Only at a boundary " + "are there no pending update() calls or in-flight epoch commits " + "to capture mid-mutation." + ) + + # Boundary-aligned is not the same as quiescent. Both the exchange + # and the epoch commit normally fire lazily, when the *next* step + # observes that the boundary was crossed — so at step N neither has + # run, and a checkpoint taken here would record pre-exchange labels + # and a shared-history bias with its deposits still pending rather + # than merged. Drain both, in the same order the runtime stamp uses: + # exchange first, because the commit publishes shared history and + # doing it before the swap would publish under labels that are about + # to change. + if self.replica_exchange is not None: + interval = self.replica_exchange.attempt_interval + self._attempt_segment(target, step // interval - 1) + self._commit_epoch(step // self.steps_per_epoch - 1) + + write_checkpoint( + path, + target, + self._components(), + sampling_step=step, + sampling_epoch=step // self.steps_per_epoch, + steps_per_epoch=self.steps_per_epoch, + model_class=_qualified_name(self.dynamics.model), + dynamics_class=_qualified_name(self.dynamics), + bias_classes={ + name: _qualified_name(bias) for name, bias in self.biases.items() + }, + exchange_config=( + self.replica_exchange.config_fingerprint() + if self.replica_exchange is not None + else None + ), + ) + + def restore( + self, path: str | Path, device: torch.device | str | None = None + ) -> Batch: + """Restore a checkpoint exactly, and prime forces before returning. + + The caller must have reconstructed the same model, dynamics, and + biases first; this validates that they match what was saved. **Model + weights are not restored** — load them through the model's own API + before calling here. The compatibility metadata proves the + architecture agrees, not that the weights do. + + Parameters + ---------- + path: + Source Zarr store. + device: + Device for the restored batch. Defaults to the dynamics' device. + + Returns + ------- + Batch + The restored batch, force-primed and ready to run. + + Raises + ------ + ValueError + If the checkpoint is uncommitted, fails a checksum, or was + written by a different model, dynamics, or bias set. + """ + target_device = device if device is not None else self._model_device() + batch, states, manifest = read_checkpoint(path, target_device) + self._validate_compatibility(manifest) + + self.steps_per_epoch = int(manifest.steps_per_epoch) + runner_state = states.get("runner", {}) + self._next_walker_id = int(runner_state.get("next_walker_id", 0)) + self._last_epoch = int(runner_state.get("last_epoch", -1)) + self._committed_epoch = int(runner_state.get("committed_epoch", -1)) + self._last_segment = int(runner_state.get("last_segment", -1)) + self._attempted_segment = int(runner_state.get("attempted_segment", -1)) + + exchange_state = states.get("exchange") + if exchange_state is not None and self.replica_exchange is not None: + self.replica_exchange.load_state_dict(exchange_state) + self._last_update_step = { + name: int(step) + for name, step in (runner_state.get("last_update_step") or {}).items() + } + + for name, bias in self.biases.items(): + state = states.get(f"biases/{name}") + loader = getattr(bias, "load_state_dict", None) + if state is not None and callable(loader): + loader(state) + # Loading wrote each bias's saved state_version straight onto it, which + # the runner never observed as a change; re-baseline so the first + # post-restore update() does not read it as one. + self._sync_seen_versions() + + # The integrator's per-system state must exist before it can be + # restored into, and its shapes come from the batch — so initialise + # against the restored batch first, then overwrite. + self.dynamics._ensure_state_initialized(batch) + self.dynamics.load_state_dict(states.get("dynamics", {})) + + self._restored = True + self.prime_forces(batch) + return batch + + def _model_device(self) -> torch.device: + """Return the device the model's own tensors live on. + + ``BaseDynamics.device`` reports the process's compute device, which + is CUDA whenever a GPU is visible — even for a model that was never + moved off the CPU. Restoring a batch there would put the batch and + the model on different devices. The model's own parameters are the + authority. + + Returns + ------- + torch.device + The model's device, falling back to the dynamics' device when the + model holds no tensors (a pure-physics wrapper such as LJ). + """ + model = self.dynamics.model + for tensor in list(model.parameters()) + list(model.buffers()): + return tensor.device + return self.dynamics.device + + def _validate_compatibility(self, manifest: CheckpointManifest) -> None: + """Reject a checkpoint written by a different configuration. + + Parameters + ---------- + manifest: + The committed manifest. + + Raises + ------ + ValueError + If the model class, dynamics class, or bias set disagrees. + """ + problems: list[str] = [] + actual_model = _qualified_name(self.dynamics.model) + if manifest.model_class and manifest.model_class != actual_model: + problems.append( + f" model: checkpoint has {manifest.model_class}, " + f"this runner has {actual_model}" + ) + actual_dynamics = _qualified_name(self.dynamics) + if manifest.dynamics_class and manifest.dynamics_class != actual_dynamics: + problems.append( + f" dynamics: checkpoint has {manifest.dynamics_class}, " + f"this runner has {actual_dynamics}" + ) + actual_biases = { + name: _qualified_name(bias) for name, bias in self.biases.items() + } + if manifest.bias_classes != actual_biases: + problems.append( + f" biases: checkpoint has {manifest.bias_classes}, " + f"this runner has {actual_biases}" + ) + # The ladder decides what a swap means, so a mismatch — including + # exchange-versus-none in either direction — changes the semantics of + # every future swap while the counters and assignment carry on looking + # valid. + problems.extend( + ReplicaExchange.describe_config_mismatch( + manifest.exchange_config, + self.replica_exchange.config_fingerprint() + if self.replica_exchange is not None + else None, + ) + ) + if problems: + detail = "\n".join(problems) + raise ValueError( + "EnhancedSampling.restore: the checkpoint was written by a " + f"different configuration:\n{detail}\n" + "Reconstruct the same model, dynamics, and biases before " + "restoring. Note that model *weights* are never restored from " + "a checkpoint — load them through the model's own API." + ) + + def __repr__(self) -> str: + """Return a concise description of the runner.""" + names = ", ".join(self.biases) or "none" + exchange = ( + f", exchange={self.replica_exchange!r}" + if self.replica_exchange is not None + else "" + ) + return ( + f"{type(self).__name__}(dynamics={type(self.dynamics).__name__}, " + f"biases=[{names}], steps_per_epoch={self.steps_per_epoch}" + f"{exchange})" + ) + + def state_dict(self) -> Mapping[str, Any]: + """Return runner counters plus each adaptive bias's state. + + Returns + ------- + Mapping[str, Any] + Nested mapping; bias state lives under ``biases/``. + """ + state: dict[str, Any] = { + "steps_per_epoch": self.steps_per_epoch, + "next_walker_id": self._next_walker_id, + "last_epoch": self._last_epoch, + "biases": {}, + } + for name, bias in self.biases.items(): + getter = getattr(bias, "state_dict", None) + if callable(getter): + state["biases"][name] = getter() + return state diff --git a/nvalchemi/enhanced_sampling/biases/__init__.py b/nvalchemi/enhanced_sampling/biases/__init__.py new file mode 100644 index 00000000..73679759 --- /dev/null +++ b/nvalchemi/enhanced_sampling/biases/__init__.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Built-in bias implementations. + +Static biases: :class:`HarmonicUmbrellaBias`, :class:`UpperWall`, +:class:`LowerWall`, :class:`FlatBottomRestraint`. + +History-dependent biases: :class:`WellTemperedMetaDynamicsBias` (Gaussian +hills along a chosen collective variable) and :class:`RMSDMetaDynamicsBias` +(xTB-style repulsion from retained reference geometries, for when the +interesting coordinates are not known in advance). + +Force-only biases: :class:`AdaptiveBiasingForce`, which applies a measured +mean force with no energy behind it. +""" + +from nvalchemi.enhanced_sampling.biases.abf import AdaptiveBiasingForce +from nvalchemi.enhanced_sampling.biases.metadynamics import ( + WellTemperedMetaDynamicsBias, +) +from nvalchemi.enhanced_sampling.biases.rmsd_metad import RMSDMetaDynamicsBias +from nvalchemi.enhanced_sampling.biases.umbrella import HarmonicUmbrellaBias +from nvalchemi.enhanced_sampling.biases.walls import ( + FlatBottomRestraint, + LowerWall, + UpperWall, +) + +__all__ = [ + "AdaptiveBiasingForce", + "FlatBottomRestraint", + "HarmonicUmbrellaBias", + "LowerWall", + "RMSDMetaDynamicsBias", + "UpperWall", + "WellTemperedMetaDynamicsBias", +] diff --git a/nvalchemi/enhanced_sampling/biases/abf.py b/nvalchemi/enhanced_sampling/biases/abf.py new file mode 100644 index 00000000..66cc3f96 --- /dev/null +++ b/nvalchemi/enhanced_sampling/biases/abf.py @@ -0,0 +1,645 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Adaptive biasing force over a pair-distance collective variable.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from torch import Tensor, nn + +from nvalchemi.dynamics.base import DynamicsStage +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling._adaptive import AdaptivePotentialMixin +from nvalchemi.enhanced_sampling._bias import BiasResult +from nvalchemi.enhanced_sampling.cv.pair_distance import pair_displacement + +if TYPE_CHECKING: + from nvalchemi.data import Batch + +__all__ = ["AdaptiveBiasingForce"] + + +class AdaptiveBiasingForce(AdaptivePotentialMixin, nn.Module): + r"""Estimate and cancel the mean force along a pair distance. + + Metadynamics fills a basin with hills; ABF instead measures the mean + force :math:`\partial A/\partial \xi` in each bin of the collective + variable and applies its negative, so the residual force along the CV + averages to zero and the walker diffuses freely. The estimate *is* the + free-energy gradient, so no reweighting is needed at the end — + :meth:`free_energy` integrates it directly. + + For the pair distance :math:`r = |\mathbf{r}_j - \mathbf{r}_i|` the + estimator is + + .. math:: + + \frac{\partial A}{\partial r} = \left\langle + -\frac{(\mathbf{F}_j - \mathbf{F}_i)\cdot\hat{\mathbf{u}}}{2} + \;-\; \frac{2 k_B T}{r} \right\rangle_r + + with :math:`\hat{\mathbf{u}} = (\mathbf{r}_j - \mathbf{r}_i)/r`. The + second term is the **metric correction**, and it is not optional: see + Notes. + + Parameters + ---------- + atom_indices: + The pair ``(i, j)``, shape ``[2]`` for the same pair in every graph + or ``[B, 2]`` for one per graph. Indices are local to each graph. + temperature: + Simulation temperature in Kelvin. Enters the metric correction, so + it must match the thermostat; a mismatch biases the estimate. + cv_range: + ``(r_min, r_max)`` in angstrom, the binned interval. Outside it no + force is applied and no sample is recorded. + name: + Unique bias identifier. + n_bins: + Number of uniform bins across ``cv_range``. + min_samples: + Samples a bin needs before *any* bias force is applied from it. An + estimate from a handful of samples is noise, and applying it would + drive the walker on the strength of that noise. + full_samples: + Samples at which the applied fraction reaches 1. Defaults to + ``2 * min_samples``. Between the two the force ramps linearly; + equal to ``min_samples`` it is a step, applying nothing until the + threshold and full force at it. + max_force: + Optional cap on ``|dA/dr|`` in eV/A. A bin that has been visited + once at a bad geometry can hold a large estimate; the cap bounds + what that can do to the trajectory. + update_frequency: + Steps between :meth:`update` calls. ``1`` (every step) is the usual + choice — ABF wants every uncorrelated sample it can get. + + Raises + ------ + ValueError + For a malformed ``atom_indices``, a non-positive temperature, an + empty or inverted ``cv_range``, fewer than one bin, a negative + ``min_samples``, ``full_samples < min_samples``, or a non-positive + ``max_force``. + + Notes + ----- + The metric correction is not optional + Projecting Cartesian forces onto :math:`\nabla\xi` and averaging + gives the mean force in the *constrained* ensemble, not the gradient + of the free energy of the unconstrained one. The two differ by the + Jacobian of the coordinate change — for a distance in three + dimensions the shell volume grows as :math:`r^2`, contributing + :math:`-2k_B T/r`. + + Omitting it does not produce noise, it produces a smoothly wrong + answer: two non-interacting particles would be reported as having a + flat PMF when the true one is :math:`-2k_B T\ln r`. That is why + this class takes an atom **pair** rather than a general ``cv`` + callable — the correction is specific to this coordinate, and + accepting an arbitrary CV would mean applying a distance-shaped + correction to something that is not a distance. + + Force-only, and therefore excluded from replica exchange + :meth:`evaluate` returns forces with ``energy=None``. There is no + potential to report: the applied force is not the gradient of any + function the bias holds, which is exactly what makes ABF + non-conservative. ``supplies_exchange_energy`` is ``False``, and + :class:`~nvalchemi.enhanced_sampling.ReplicaExchange` refuses such a + bias rather than evaluating an acceptance rule that needs a + cross-state bias energy. + + Observation ordering + ``observation_stage`` is ``AFTER_COMPUTE``, where ``batch.forces`` + still holds the **unbiased** physical force. Observing after the + bias is applied would feed the estimator its own output, and it + would converge to whatever it had already decided. + + Examples + -------- + >>> import torch + >>> from nvalchemi.enhanced_sampling import AdaptiveBiasingForce + >>> abf = AdaptiveBiasingForce( + ... atom_indices=torch.tensor([0, 1]), + ... temperature=300.0, + ... cv_range=(2.0, 6.0), + ... n_bins=40, + ... ) + >>> int(abf.bin_counts.sum()) + 0 + """ + + #: The acceptance rule needs a cross-state bias energy; this has none. + supplies_exchange_energy: bool = False + + def __init__( + self, + atom_indices: Tensor, + temperature: float, + cv_range: tuple[float, float], + *, + name: str = "abf", + n_bins: int = 100, + min_samples: int = 200, + full_samples: int | None = None, + max_force: float | None = None, + update_frequency: int = 1, + ) -> None: + super().__init__() + + indices = torch.as_tensor(atom_indices, dtype=torch.long) + if indices.dim() == 1: + if indices.numel() != 2: + raise ValueError( + f"AdaptiveBiasingForce: atom_indices must name exactly two " + f"atoms, got {indices.numel()}." + ) + elif indices.dim() != 2 or indices.shape[-1] != 2: + raise ValueError( + f"AdaptiveBiasingForce: atom_indices must have shape [2] or " + f"[B, 2], got {tuple(indices.shape)}." + ) + if bool((indices < 0).any()): + raise ValueError( + f"AdaptiveBiasingForce: atom_indices must be non-negative " + f"per-graph indices, got {indices.tolist()}." + ) + if bool((indices[..., 0] == indices[..., 1]).any()): + raise ValueError( + f"AdaptiveBiasingForce: atom_indices names the same atom twice " + f"({indices.tolist()}); the distance would be identically zero." + ) + + if temperature <= 0.0: + raise ValueError( + f"AdaptiveBiasingForce: temperature must be positive, got " + f"{temperature}. It scales the metric correction, so it must " + "match the thermostat." + ) + + low, high = (float(v) for v in cv_range) + if not high > low: + raise ValueError( + f"AdaptiveBiasingForce: cv_range must be increasing, got " + f"({low}, {high})." + ) + if low < 0.0: + raise ValueError( + f"AdaptiveBiasingForce: cv_range lower bound must be " + f"non-negative for a distance, got {low}." + ) + if int(n_bins) < 1: + raise ValueError( + f"AdaptiveBiasingForce: n_bins must be at least 1, got {n_bins}." + ) + if int(min_samples) < 0: + raise ValueError( + f"AdaptiveBiasingForce: min_samples must be non-negative, got " + f"{min_samples}." + ) + ramp_end = 2 * int(min_samples) if full_samples is None else int(full_samples) + if ramp_end < int(min_samples): + raise ValueError( + f"AdaptiveBiasingForce: full_samples ({ramp_end}) must be at " + f"least min_samples ({int(min_samples)}) — the force ramps up " + "between the two." + ) + if max_force is not None and max_force <= 0.0: + raise ValueError( + f"AdaptiveBiasingForce: max_force must be positive, got {max_force}." + ) + if int(update_frequency) < 1: + raise ValueError( + f"AdaptiveBiasingForce: update_frequency must be at least 1, " + f"got {update_frequency}." + ) + + self.name = name + self.temperature = float(temperature) + self.cv_range = (low, high) + self.n_bins = int(n_bins) + self.min_samples = int(min_samples) + self.full_samples = ramp_end + self.max_force = None if max_force is None else float(max_force) + self.update_frequency = int(update_frequency) + self.observation_stage = DynamicsStage.AFTER_COMPUTE + + dtype = torch.get_default_dtype() + self.register_buffer("atom_indices", indices) + self.register_buffer("bin_counts", torch.zeros(self.n_bins, dtype=torch.long)) + self.register_buffer("force_sum", torch.zeros(self.n_bins, dtype=dtype)) + + # ------------------------------------------------------------------ + # Geometry and binning + # ------------------------------------------------------------------ + + @property + def bin_width(self) -> float: + """Return the width of one CV bin, in angstrom.""" + low, high = self.cv_range + return (high - low) / self.n_bins + + @property + def bin_centers(self) -> Tensor: + """Return the CV value at the middle of each bin, shape ``[n_bins]``.""" + low, _ = self.cv_range + offsets = torch.arange( + self.n_bins, dtype=self.force_sum.dtype, device=self.force_sum.device + ) + return low + (offsets + 0.5) * self.bin_width + + def bin_index(self, values: Tensor) -> Tensor: + """Return the bin each CV value falls in, shape ``[...]``. + + Values outside ``cv_range`` are reported as ``-1``; they contribute + no sample and receive no force, so they belong to no bin. + + Parameters + ---------- + values: + CV values in angstrom. + + Returns + ------- + Tensor + Integer bin indices, or ``-1`` for out-of-range values. + """ + index, in_range = self._bin_of(torch.as_tensor(values)) + return torch.where(in_range, index, torch.full_like(index, -1)) + + def _align_device(self, reference: Tensor) -> None: + """Move this bias's buffers to *reference*'s device if they differ. + + Parameters + ---------- + reference: + Any tensor from the live batch; its device is the target. + """ + if self.force_sum.device != reference.device: + self.to(reference.device) + + def _geometry(self, current: Batch) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Return ``(distance, unit_vector, global_i, global_j)``. + + Parameters + ---------- + current: + The live batch. + + Returns + ------- + tuple[Tensor, Tensor, Tensor, Tensor] + Distances ``[B]``, unit displacements ``[B, 3]``, and the two + global row indices ``[B]`` the bias force is written to. + """ + # pair_displacement validates atom_indices and applies the minimum + # image convention, so the resolution repeated below is known good. + delta = pair_displacement(current, self.atom_indices) # [B, 3] + distance = torch.linalg.vector_norm(delta, dim=-1) # [B] + unit = delta / distance.clamp(min=torch.finfo(delta.dtype).tiny).unsqueeze(-1) + + offsets = current.batch_ptr[:-1] # [B] + indices = self.atom_indices.to(offsets.device) + if indices.dim() == 1: + indices = indices.unsqueeze(0).expand(offsets.numel(), 2) + return distance, unit, offsets + indices[:, 0], offsets + indices[:, 1] + + def _bin_of(self, distance: Tensor) -> tuple[Tensor, Tensor]: + """Return ``(bin_index, in_range)`` for CV values *distance*. + + Parameters + ---------- + distance: + CV values, shape ``[B]``. + + Returns + ------- + tuple[Tensor, Tensor] + Bin indices clamped into range, and a boolean mask marking which + samples actually fell inside ``cv_range``. + + The clamp makes the index always safe to gather with, but it + also means an out-of-range value silently names the nearest edge + bin. Every caller must therefore mask what it gathers with the + second return value, not just the quantities that reach the + force. + """ + low, high = self.cv_range + raw = ((distance - low) / self.bin_width).floor().to(torch.long) + in_range = (distance >= low) & (distance < high) + return raw.clamp(0, self.n_bins - 1), in_range + + # ------------------------------------------------------------------ + # Estimate + # ------------------------------------------------------------------ + + def mean_force(self) -> Tensor: + """Return the per-bin mean-force estimate, shape ``[n_bins]``. + + Returns + ------- + Tensor + ``dA/dr`` in eV/A per bin. Unvisited bins are ``nan`` rather + than zero: a bin with no samples has no estimate, and zero is a + perfectly plausible mean force that would hide that. + """ + counts = self.bin_counts.to(self.force_sum.dtype) + estimate = self.force_sum / counts + return torch.where(self.bin_counts > 0, estimate, torch.nan) + + def _applied_gradient(self) -> Tensor: + """Return the ramped, capped ``dA/dr`` actually applied, ``[n_bins]``. + + Returns + ------- + Tensor + Zero wherever no force should be applied yet, so this is safe to + index with any bin. + """ + counts = self.bin_counts.to(self.force_sum.dtype) + estimate = torch.where( + self.bin_counts > 0, self.force_sum / counts.clamp(min=1.0), 0.0 + ) + if self.max_force is not None: + estimate = estimate.clamp(-self.max_force, self.max_force) + # Reuses ramp_fraction rather than repeating it: two copies of the + # same schedule are two places for the endpoint to disagree. + return estimate * self.ramp_fraction() + + def ramp_fraction(self) -> Tensor: + """Return the applied fraction per bin, shape ``[n_bins]``. + + Zero at or below ``min_samples``, one at or above ``full_samples``, + linear between. Ramping rather than switching on at the threshold + avoids a force discontinuity of exactly the size the threshold exists + to prevent. + + ``full_samples == min_samples`` is a step: nothing until the + threshold, full force at it. That is a legitimate choice — it is the + classic hard-threshold form, and it is what ``min_samples=0`` gives + by default — so it is handled rather than rejected. It needs its own + branch because the linear form would divide by a zero span, and + clamping that span to 1 would delay full force by one sample. + + A bin with no samples has no estimate, so it reports zero whatever + the thresholds are. + + Returns + ------- + Tensor + Fractions in ``[0, 1]``. + """ + counts = self.bin_counts.to(self.force_sum.dtype) + span = self.full_samples - self.min_samples + if span <= 0: + fraction = (self.bin_counts >= self.full_samples).to(counts.dtype) + else: + fraction = ((counts - self.min_samples) / span).clamp(0.0, 1.0) + return torch.where(self.bin_counts > 0, fraction, torch.zeros_like(fraction)) + + # ------------------------------------------------------------------ + # BiasPotential + # ------------------------------------------------------------------ + + def evaluate(self, current: Batch) -> BiasResult: + """Return the bias force, with no energy. + + Read-only: the estimate is applied but never updated here. + + Parameters + ---------- + current: + The live batch. + + Returns + ------- + BiasResult + ``forces`` only, plus per-walker diagnostics: the CV value, its + ``bin`` (``-1`` outside ``cv_range``, matching + :meth:`bin_index`), the ``applied_gradient`` actually used + (ramped and capped, so not the same as :meth:`mean_force`), the + bin's sample count, its ramp fraction, and ``in_range``. Every + per-bin quantity reads zero for a walker outside the range, so + the diagnostics agree with the force rather than reporting the + nearest edge bin's statistics. ``energy`` is ``None`` because + there is none to report. + """ + self._align_device(current.positions) + + with torch.no_grad(): + distance, unit, global_i, global_j = self._geometry(current) + bins, in_range = self._bin_of(distance) + + gradient = self._applied_gradient()[bins] # [B] + gradient = torch.where(in_range, gradient, torch.zeros_like(gradient)) + + # V_bias = -A, so F_bias = +grad A = (dA/dr) grad r, and + # grad_i r = -u while grad_j r = +u. + contribution = gradient.unsqueeze(-1) * unit # [B, 3] + forces = torch.zeros_like(current.positions) + forces.index_add_(0, global_j, contribution) + forces.index_add_(0, global_i, -contribution) + + # Every per-bin diagnostic is masked the same way the gradient is. + # _bin_of clamps the index so it is always safe to gather with, + # which means an out-of-range walker would otherwise read the + # nearest edge bin's statistics and report them as its own — a + # converged-looking sample count and ramp next to in_range == 0. + zero_long = torch.zeros_like(self.bin_counts[bins]) + ramp = self.ramp_fraction()[bins] + + # Named applied_gradient, not mean_force: this is the ramped and + # capped value actually used, which is not what mean_force() + # returns. Reusing that name would invite reading a + # threshold-suppressed zero as a measured zero mean force. + observables = { + "cv": distance.unsqueeze(-1), + "bin": torch.where(in_range, bins, torch.full_like(bins, -1)).unsqueeze( + -1 + ), + "applied_gradient": gradient.unsqueeze(-1), + "samples": torch.where( + in_range, self.bin_counts[bins], zero_long + ).unsqueeze(-1), + "ramp": torch.where(in_range, ramp, torch.zeros_like(ramp)).unsqueeze( + -1 + ), + "in_range": in_range.to(forces.dtype).unsqueeze(-1), + } + + return BiasResult(forces=forces, observables=observables) + + def update(self, frames: Batch, result: BiasResult) -> None: + """Accumulate one mean-force sample per walker. + + Parameters + ---------- + frames: + The ``AFTER_COMPUTE`` capture, whose ``forces`` are the unbiased + physical forces. + result: + This bias's own preceding result; unused, since the estimator + must not see its own output. + + Raises + ------ + ValueError + If the captured frame carries no forces to project. + """ + physical = getattr(frames, "forces", None) + if physical is None: + raise ValueError( + f"AdaptiveBiasingForce {self.name!r}: the observed frame has no " + "forces, so there is no mean force to sample. This bias " + "observes at AFTER_COMPUTE, where the batch must already carry " + "the physical forces from the model." + ) + + with torch.no_grad(): + self._align_device(physical) + distance, unit, global_i, global_j = self._geometry(frames) + bins, in_range = self._bin_of(distance) + + # Projected force along the CV, then the Jacobian term. See the + # class Notes for why the second is not optional. + projected = ( + -((physical[global_j] - physical[global_i]) * unit).sum(dim=-1) / 2.0 + ) + metric = 2.0 * KB_EV * self.temperature / distance + sample = (projected - metric).to(self.force_sum.dtype) + + keep = in_range & torch.isfinite(sample) + if not bool(keep.any()): + return + + bins = bins[keep] + self.force_sum.index_add_(0, bins, sample[keep]) + self.bin_counts.index_add_(0, bins, torch.ones_like(bins)) + + # Only bump when the applied force actually changed. A bin still + # below its threshold contributes nothing, so re-priming forces + # over it would be pure cost — this is the case + # AdaptivePotentialMixin.bump_state_version documents. Asking + # ramp_fraction rather than re-deriving the threshold keeps the + # two from disagreeing at the endpoint. + changed = bool((self.ramp_fraction()[bins] > 0.0).any()) + + if changed: + self.bump_state_version() + + # ------------------------------------------------------------------ + # Analysis + # ------------------------------------------------------------------ + + def config_fingerprint(self) -> dict[str, Any]: + """Return the settings the accumulated histogram is only valid under. + + ``cv_range`` and ``n_bins`` decide what each bin *means*; restoring a + histogram under a different pair silently relabels every bin. + ``temperature`` scales the metric correction already folded into + ``force_sum``, so samples taken at 300 K cannot be extended at 900 K. + ``atom_indices`` names the coordinate itself. ``min_samples`` / + ``full_samples`` / ``max_force`` change what the same counts apply. + + Returns + ------- + dict[str, Any] + The checked configuration. + """ + return { + "atom_indices": self.atom_indices.reshape(-1).tolist(), + "cv_range": list(self.cv_range), + "n_bins": self.n_bins, + "temperature": self.temperature, + "min_samples": self.min_samples, + "full_samples": self.full_samples, + "max_force": self.max_force, + } + + def free_energy(self) -> Tensor: + """Return the PMF at each bin center, shape ``[n_bins]``. + + Integrates the mean-force estimate by the trapezoid rule, shifted so + the minimum over sampled bins is zero. This is the payoff of the + method: unlike metadynamics there is nothing to deconvolve, because + the accumulated quantity already *is* the free-energy gradient. + + Returns + ------- + Tensor + Free energy in eV. Bins never visited are ``nan``. + + Raises + ------ + RuntimeError + If no bin has been sampled, or if an unvisited bin sits between + two visited ones. Integration carries the profile across the + gap, so a hole in the middle would silently contaminate every + value beyond it. + """ + visited = self.bin_counts > 0 + if not bool(visited.any()): + raise RuntimeError( + f"AdaptiveBiasingForce {self.name!r}: no bin has been sampled, " + "so there is no free energy to report. Check that cv_range " + "covers the CV values the run actually visits." + ) + + indices = visited.nonzero(as_tuple=False).squeeze(-1) + first, last = int(indices[0]), int(indices[-1]) + if not bool(visited[first : last + 1].all()): + missing = (~visited[first : last + 1]).nonzero(as_tuple=False).squeeze( + -1 + ) + first + raise RuntimeError( + f"AdaptiveBiasingForce {self.name!r}: bin(s) " + f"{missing.tolist()} were never visited but lie between bins " + f"{first} and {last} that were. Integrating the mean force " + "carries the profile across the gap, so every value beyond it " + "would be wrong by an unknown constant. Run longer, widen the " + "bins, or narrow cv_range." + ) + + gradient = self.force_sum[first : last + 1] / self.bin_counts[ + first : last + 1 + ].to(self.force_sum.dtype) + + profile = torch.full_like(self.force_sum, torch.nan) + if gradient.numel() == 1: + profile[first] = 0.0 + return profile + + midpoints = 0.5 * (gradient[:-1] + gradient[1:]) * self.bin_width + integrated = torch.cat( + [ + torch.zeros(1, dtype=gradient.dtype, device=gradient.device), + torch.cumsum(midpoints, dim=0), + ] + ) + profile[first : last + 1] = integrated - integrated.min() + return profile + + def __repr__(self) -> str: + """Return a concise description of the bias.""" + low, high = self.cv_range + sampled = int((self.bin_counts > 0).sum()) + return ( + f"{type(self).__name__}(name={self.name!r}, " + f"cv_range=({low:g}, {high:g}), n_bins={self.n_bins}, " + f"sampled_bins={sampled}/{self.n_bins}, " + f"samples={int(self.bin_counts.sum())})" + ) diff --git a/nvalchemi/enhanced_sampling/biases/metadynamics.py b/nvalchemi/enhanced_sampling/biases/metadynamics.py new file mode 100644 index 00000000..45b72464 --- /dev/null +++ b/nvalchemi/enhanced_sampling/biases/metadynamics.py @@ -0,0 +1,788 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Well-tempered metadynamics over one or more collective variables.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +import torch +from torch import Tensor + +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling._adaptive import AdaptivePotentialMixin +from nvalchemi.enhanced_sampling._bias import ConservativeBias +from nvalchemi.enhanced_sampling.cv._periodic import periodic_difference + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + from nvalchemi.data import Batch + from nvalchemi.enhanced_sampling._bias import BiasResult + +__all__ = ["WellTemperedMetaDynamicsBias"] + +_HISTORY_MODES = ("shared", "state", "walker") +_STORAGE_POLICIES = ("preallocated", "grow", "fifo") + + +class WellTemperedMetaDynamicsBias(AdaptivePotentialMixin, ConservativeBias): + r"""Gaussian hills deposited along a CV, with well-tempered damping. + + .. math:: + + V(s, t) = \sum_i h_i \exp\!\left( + -\sum_d \frac{(s_d - c_{i,d})^2}{2\sigma_d^2}\right) + + Each deposition adds one hill per walker at that walker's current CV + value. In the well-tempered scheme the height decays where the bias has + already accumulated, + + .. math:: h_t = h_0 \exp\!\left(-\frac{V(s_t)}{k_B T (\gamma - 1)}\right) + + which makes the sum converge instead of filling forever. At convergence + the free energy is recovered as + :math:`F(s) = -\frac{\gamma}{\gamma - 1} V(s)` (Barducci, Bussi & + Parrinello 2008), available from :meth:`free_energy`. + + Passing ``bias_factor=None`` gives standard metadynamics — the + :math:`\gamma \to \infty` limit, where every hill has height ``h_0`` and + :math:`F(s) = -V(s)`. + + Parameters + ---------- + cv: + Differentiable ``cv(batch) -> Tensor[B, D]``. + height: + Initial hill height ``h_0`` in eV. + sigma: + Hill width, either a scalar shared by every CV component or shape + ``[D]``. Any other length is rejected on the first evaluation + rather than broadcast into the hill table. + temperature: + Simulation temperature in Kelvin, used for the well-tempered damping. + bias_factor: + ``gamma > 1``, or ``None`` for standard metadynamics. + name: + Unique bias identifier. + update_frequency: + Dynamics steps between depositions (the deposition pace). + storage: + ``"preallocated"``, ``"grow"``, or ``"fifo"`` — see Notes. + max_hills: + Capacity. Required for ``"preallocated"`` and ``"fifo"``; the + initial chunk for ``"grow"``. + history: + ``"shared"`` (every walker sees every hill), ``"state"`` (hills + belong to the thermodynamic state that deposited them), or + ``"walker"`` (each walker sees only its own). + + ``"state"`` and ``"walker"`` require the batch to carry + ``thermodynamic_state_id`` / ``walker_id`` respectively, and raise if + it does not. :class:`~nvalchemi.enhanced_sampling.EnhancedSampling` + stamps both on every step, so a runner-driven bias never sees this; + a bias evaluated directly must supply the field itself. + periods: + Period per CV component, ``0`` for non-periodic, shape ``[D]``. + Applied to the ``s - c`` difference so a hill near a branch cut + still repels from both sides. Unlike ``sigma`` a scalar is not + accepted for a multi-component CV: the entries carry per-component + meaning, so broadcasting one across all of them would quietly make + every component periodic. + ramp_depositions: + Number of *deposition events* over which a freshly added hill ramps + from zero to full height. ``0`` activates immediately. Counted in + depositions rather than dynamics steps because ``energy()`` sees the + hill table, not the step counter. + compute_stress: + Passed through to :class:`ConservativeBias`. + + Raises + ------ + ValueError + For an invalid storage policy, history mode, non-positive height or + width, ``bias_factor <= 1``, or a missing capacity. A ``sigma`` or + ``periods`` length that disagrees with the CV is raised on first + evaluation, since the CV dimension is not knowable at construction + from a plain callable. + + Notes + ----- + Storage policies + ``preallocated`` keeps tensor shapes fixed for the whole run and + **raises** when capacity is exhausted, rather than silently changing + the physics. It is the compile-stable choice. + + ``grow`` allocates in chunks of ``max_hills``; each growth changes + the hill-tensor shape, which forces a recompile if ``energy()`` is + compiled. Under torch's default of static parameter shapes, Dynamo + caps retraces per code object at + ``torch._dynamo.config.recompile_limit`` (8), so a compiled run + **hard-fails once past it** — mid-trajectory, after the run is + already underway. + + Both settings are process-global and other code changes them: + ``DistributedModel`` raises the limit to 64 *and* sets + ``force_parameter_static_shapes = False``. The second matters more + than the first — with dynamic parameter shapes the hill-table + dimension is traced symbolically, growth stops triggering a retrace, + and the limit is never reached. So the three ways out are: size + ``max_hills`` so growths stay under the limit, enable dynamic + parameter shapes, or use ``preallocated``, which holds one trace for + the whole run regardless. + + ``fifo`` bounds memory by discarding the oldest hill. This is + **scientifically meaningful, not merely a cache eviction**: the + accumulated bias is no longer the integral of everything deposited, + so the well-tempered convergence argument no longer applies and + :meth:`free_energy` is not a valid estimator. Chosen deliberately + for exploration, and the natural policy for the RMSD variant. + + Multi-walker history + ``shared`` is the multiple-walker scheme: every walker feels every + hill, so ``B`` walkers fill a basin roughly ``B`` times faster. + ``state`` keeps a separate history per thermodynamic state, which is + what a replica-exchange ladder needs. ``walker`` runs ``B`` + independent metadynamics simulations in one batch. + + Examples + -------- + >>> import torch + >>> from nvalchemi.enhanced_sampling import ( + ... WellTemperedMetaDynamicsBias, pair_distance, + ... ) + >>> idx = torch.tensor([0, 5]) + >>> bias = WellTemperedMetaDynamicsBias( + ... cv=lambda b: pair_distance(b, idx), + ... height=0.01, sigma=0.1, temperature=300.0, bias_factor=10.0, + ... max_hills=1000, + ... ) + >>> bias.hill_count.item() + 0 + """ + + def __init__( + self, + cv: Callable[[Batch], Tensor], + height: float, + sigma: Tensor | float, + temperature: float, + *, + bias_factor: float | None = None, + name: str = "metadynamics", + update_frequency: int = 500, + storage: Literal["preallocated", "grow", "fifo"] = "preallocated", + max_hills: int | None = None, + history: Literal["shared", "state", "walker"] = "shared", + periods: Tensor | None = None, + ramp_depositions: int = 0, + compute_stress: bool = True, + ) -> None: + super().__init__(name=name, compute_stress=compute_stress) + + if storage not in _STORAGE_POLICIES: + raise ValueError( + f"WellTemperedMetaDynamicsBias: storage must be one of " + f"{list(_STORAGE_POLICIES)}, got {storage!r}." + ) + if history not in _HISTORY_MODES: + raise ValueError( + f"WellTemperedMetaDynamicsBias: history must be one of " + f"{list(_HISTORY_MODES)}, got {history!r}." + ) + if height <= 0.0: + raise ValueError( + f"WellTemperedMetaDynamicsBias: height must be positive, got " + f"{height}. A non-positive hill would attract the walker to " + "where it has already been." + ) + if bias_factor is not None and bias_factor <= 1.0: + raise ValueError( + f"WellTemperedMetaDynamicsBias: bias_factor must be greater " + f"than 1, got {bias_factor}. gamma = 1 divides by zero in the " + "well-tempered height; pass None for standard metadynamics." + ) + if int(update_frequency) < 1: + raise ValueError( + f"WellTemperedMetaDynamicsBias: update_frequency must be at " + f"least 1, got {update_frequency}." + ) + if int(ramp_depositions) < 0: + raise ValueError( + f"WellTemperedMetaDynamicsBias: ramp_depositions must be " + f"non-negative, got {ramp_depositions}." + ) + if storage in ("preallocated", "fifo") and max_hills is None: + raise ValueError( + f"WellTemperedMetaDynamicsBias: storage={storage!r} needs an " + "explicit max_hills — it is the whole point of the policy, " + "either the ceiling that raises or the ring that overwrites." + ) + capacity = int(max_hills) if max_hills is not None else 256 + if capacity < 1: + raise ValueError( + f"WellTemperedMetaDynamicsBias: max_hills must be at least 1, " + f"got {max_hills}." + ) + + sigma_t = torch.as_tensor(sigma, dtype=torch.get_default_dtype()).reshape(-1) + if bool((sigma_t <= 0).any()): + raise ValueError( + f"WellTemperedMetaDynamicsBias: sigma must be positive, got " + f"{sigma_t.tolist()}. A zero width is a delta function with no " + "gradient anywhere but one point." + ) + + self.cv = cv + self.storage = storage + self.history = history + self.update_frequency = int(update_frequency) + self.ramp_depositions = int(ramp_depositions) + self.height = float(height) + self.temperature = float(temperature) + self.bias_factor = None if bias_factor is None else float(bias_factor) + self._capacity = capacity + + self.register_buffer("sigma", sigma_t) + if periods is None: + self.periods: Tensor | None = None + else: + self.register_buffer( + "periods", + torch.as_tensor(periods, dtype=torch.get_default_dtype()).reshape(-1), + ) + self._allocate(capacity, dim=sigma_t.numel()) + + # This bias reads thermodynamic_state_id only in "state" history, and + # even then the hills follow the state rather than the energy varying + # by state at fixed history — but the assignment does change which + # hills a walker feels, so a temperature ladder would need the + # combined acceptance rule that is not implemented. + self.state_dependent_for_exchange = history == "state" + + # ------------------------------------------------------------------ + # Storage + # ------------------------------------------------------------------ + + def _allocate(self, capacity: int, dim: int) -> None: + """Create or replace the hill buffers at *capacity*. + + Parameters + ---------- + capacity: + Number of hill slots. + dim: + CV dimension. + """ + dtype = self.sigma.dtype + device = self.sigma.device + self.register_buffer( + "hill_centers", torch.zeros(capacity, dim, dtype=dtype, device=device) + ) + self.register_buffer( + "hill_heights", torch.zeros(capacity, dtype=dtype, device=device) + ) + self.register_buffer( + "hill_owner", torch.full((capacity,), -1, dtype=torch.long, device=device) + ) + self.register_buffer( + "hill_step", torch.full((capacity,), -1, dtype=torch.long, device=device) + ) + # hill_count saturates at capacity (how many slots are live); + # hills_written keeps counting, which is what makes the FIFO ring + # position unambiguous once it has wrapped. + self.register_buffer( + "hill_count", torch.zeros((), dtype=torch.long, device=device) + ) + self.register_buffer( + "hills_written", torch.zeros((), dtype=torch.long, device=device) + ) + self.register_buffer( + "deposits", torch.zeros((), dtype=torch.long, device=device) + ) + + def _grow(self) -> None: + """Double-buffer the hill tensors by one more chunk.""" + extra = self._capacity + device = self.sigma.device + + def _extend(buffer: Tensor, fill: float | int) -> Tensor: + pad = torch.full( + (extra, *buffer.shape[1:]), fill, dtype=buffer.dtype, device=device + ) + return torch.cat([buffer, pad], dim=0) + + self.hill_centers = _extend(self.hill_centers, 0.0) + self.hill_heights = _extend(self.hill_heights, 0.0) + self.hill_owner = _extend(self.hill_owner, -1) + self.hill_step = _extend(self.hill_step, -1) + + @property + def capacity(self) -> int: + """Return the current number of hill slots.""" + return int(self.hill_centers.shape[0]) + + # ------------------------------------------------------------------ + # Energy + # ------------------------------------------------------------------ + + def _owner_key(self, current: Batch, n_graphs: int) -> Tensor: + """Return the history key per graph, shape ``[B]``. + + Parameters + ---------- + current: + The live batch. + n_graphs: + Number of graphs. + + Returns + ------- + Tensor + Per-graph key matched against ``hill_owner``. All ``-1`` under + ``"shared"``, which the mask then ignores. + """ + device = self.hill_owner.device + if self.history == "state": + field = "thermodynamic_state_id" + elif self.history == "walker": + field = "walker_id" + else: + return torch.full((n_graphs,), -1, dtype=torch.long, device=device) + + ids = getattr(current, field, None) + if ids is None: + raise ValueError( + f"WellTemperedMetaDynamicsBias {self.name!r}: history=" + f"{self.history!r} needs batch.{field}, which this batch " + f"does not carry. Falling back to a single owner would put " + f"every hill under one key and silently collapse the " + f"per-{field} histories into one shared history — the " + f"opposite of what history={self.history!r} asks for. " + "EnhancedSampling stamps this field on every step; a bias " + "driven directly must set it, or use history='shared' if " + "one history really is intended." + ) + if ids.numel() != n_graphs: + raise ValueError( + f"WellTemperedMetaDynamicsBias {self.name!r}: batch.{field} has " + f"{ids.numel()} entries but the batch has {n_graphs} " + f"graph(s). A shorter tensor broadcasts across graphs, which " + f"would file every hill under one walker's key without " + "raising." + ) + return ids.reshape(-1).to(device=device, dtype=torch.long) + + def _hill_scale(self, step: Tensor) -> Tensor: + """Return each hill's ramp factor, shape ``[capacity]``. + + A hill that switched on at full height would deliver a force + discontinuity at the deposition event; ramping over + ``ramp_depositions`` spreads that over a window instead. + + Parameters + ---------- + step: + Current deposition counter. + + Returns + ------- + Tensor + Factors in ``[0, 1]``. + """ + if self.ramp_depositions <= 0: + return torch.ones_like(self.hill_heights) + age = (step - self.hill_step).to(self.hill_heights.dtype) + return torch.clamp(age / float(self.ramp_depositions), min=0.0, max=1.0) + + def _validate_cv(self, values: Tensor) -> int: + """Check the CV output against ``sigma`` and ``periods``. + + Every mismatch here is a *broadcast*, not an error: a CV of width one + against a two-component ``sigma`` broadcasts into a two-column hill + table, and the exponent then sums two terms instead of one. The + Gaussian silently narrows and nothing in the run reports it. + + Only tensor ranks and element counts are inspected, so this is a + compile-time guard rather than a data-dependent branch and + :meth:`gaussian_sum` still traces with ``fullgraph=True``. + + Parameters + ---------- + values: + CV values, expected shape ``[B, D]``. + + Returns + ------- + int + The CV dimension ``D``. + + Raises + ------ + ValueError + If *values* is not rank 2, or if ``sigma`` or ``periods`` has an + element count that would broadcast against ``D`` rather than + match it. + """ + if values.ndim != 2: + raise ValueError( + f"WellTemperedMetaDynamicsBias {self.name!r}: cv must return " + f"shape [B, D], got {tuple(values.shape)}. A CV returning [B] " + "for a scalar collective variable needs an explicit trailing " + "dimension — return values.unsqueeze(-1)." + ) + + dim = int(values.shape[1]) + if self.sigma.numel() not in (1, dim): + raise ValueError( + f"WellTemperedMetaDynamicsBias {self.name!r}: cv returns " + f"{dim} component(s) but sigma has {self.sigma.numel()}. " + "sigma must be a scalar (one width shared by every component) " + f"or have exactly {dim} entries; any other count broadcasts " + "into the hill table and silently changes the Gaussian." + ) + if self.periods is not None and self.periods.numel() != dim: + raise ValueError( + f"WellTemperedMetaDynamicsBias {self.name!r}: cv returns " + f"{dim} component(s) but periods has " + f"{self.periods.numel()}. periods is per-component — a 0 entry " + "marks that component non-periodic — so it must have exactly " + f"{dim} entries rather than broadcast." + ) + return dim + + def gaussian_sum(self, values: Tensor, owner_key: Tensor) -> Tensor: + """Return the accumulated bias at *values*, shape ``[B]``. + + Contains no data-dependent Python branch, so it compiles with + ``fullgraph=True``: the live-hill mask comes from comparing an + ``arange`` against the count tensor rather than slicing by an int. + + Parameters + ---------- + values: + CV values, shape ``[B, D]``. + owner_key: + History key per graph, shape ``[B]``. + + Returns + ------- + Tensor + Bias energy per graph, shape ``[B]``. + + Raises + ------ + ValueError + If *values* does not agree with ``sigma`` and ``periods``; see + :meth:`_validate_cv`. + """ + dim = self._validate_cv(values) + if self.hill_centers.shape[1] != dim: + # The hill table is still at its provisional width, which only + # happens before the first deposition resolves the CV dimension. + # An empty history contributes nothing whatever the width is. + return values.new_zeros(values.shape[0]) + + slots = torch.arange(self.hill_centers.shape[0], device=values.device) + live = slots < self.hill_count + + delta = periodic_difference( + values.unsqueeze(1), # [B, 1, D] + self.hill_centers.unsqueeze(0), # [1, H, D] + self.periods, + ) # [B, H, D] + exponent = 0.5 * ((delta / self.sigma) ** 2).sum(dim=-1) # [B, H] + gaussians = torch.exp(-exponent) + + scale = self._hill_scale(self.deposits) # [H] + weights = self.hill_heights * scale # [H] + + mask = live.unsqueeze(0) # [1, H] + if self.history != "shared": + mask = mask & (self.hill_owner.unsqueeze(0) == owner_key.unsqueeze(1)) + return (gaussians * weights.unsqueeze(0) * mask).sum(dim=-1) + + def energy(self, current: Batch) -> Tensor: + """Return the accumulated metadynamics bias, shape ``[B, 1]``. + + Parameters + ---------- + current: + Batch with strained positions supplied by + :meth:`ConservativeBias.evaluate`. + + Returns + ------- + Tensor + Shape ``[B, 1]`` in eV. + """ + values = self.cv(current) + owner = self._owner_key(current, values.shape[0]) + return self.gaussian_sum(values, owner).unsqueeze(-1) + + # ------------------------------------------------------------------ + # Deposition + # ------------------------------------------------------------------ + + def _well_tempered_height(self, bias_at_cv: Tensor) -> Tensor: + """Return the damped hill height per walker. + + Parameters + ---------- + bias_at_cv: + Accumulated bias at the deposition point, shape ``[B]``. + + Returns + ------- + Tensor + Heights, shape ``[B]``. + """ + if self.bias_factor is None: + return torch.full_like(bias_at_cv, self.height) + damping = KB_EV * self.temperature * (self.bias_factor - 1.0) + return self.height * torch.exp(-bias_at_cv / damping) + + def _next_slots(self, count: int) -> Tensor: + """Return the slot indices *count* new hills will occupy. + + Parameters + ---------- + count: + Number of hills about to be deposited. + + Returns + ------- + Tensor + Slot indices, shape ``[count]``. + + Raises + ------ + RuntimeError + If ``preallocated`` capacity is exhausted. Raising rather than + evicting is the point of the policy: silently dropping hills + would change the physics of a converging run without saying so. + """ + written = int(self.hills_written) + capacity = self.capacity + + if self.storage == "fifo": + # Ring buffer: hill j always lands in slot j % capacity, so the + # oldest is the one overwritten however many times it has wrapped. + return ( + torch.arange(count, device=self.hill_heights.device) + written + ) % capacity + + if written + count > capacity: + if self.storage == "preallocated": + raise RuntimeError( + f"WellTemperedMetaDynamicsBias {self.name!r}: hill storage " + f"is full ({capacity} hills) and storage='preallocated'. " + f"Raise max_hills, lengthen update_frequency, or switch to " + "storage='grow' (recompiles when it resizes) or " + "storage='fifo' (bounded memory, but discards the oldest " + "hills and so is no longer a converging well-tempered run)." + ) + while written + count > self.capacity: + self._grow() + + return torch.arange(count, device=self.hill_heights.device) + written + + def update(self, frames: Batch, result: BiasResult) -> None: + """Deposit one hill per walker at its current CV value. + + Parameters + ---------- + frames: + Post-step frame captured by the runner. + result: + The bias's own result from the preceding force evaluation. + + Raises + ------ + ValueError + If the CV output disagrees with ``sigma`` or ``periods``. + RuntimeError + If the CV changes dimension part-way through a run, which would + silently invalidate every hill already deposited. + """ + with torch.no_grad(): + values = self.cv(frames).detach() # [B, D] + owner = self._owner_key(frames, values.shape[0]) + bias_at_cv = self.gaussian_sum(values, owner) + heights = self._well_tempered_height(bias_at_cv) + + count = values.shape[0] + dim = int(values.shape[1]) + if self.hill_centers.shape[1] != dim: + # The constructor sizes the hill table from sigma, which is + # only right when sigma is per-component; a scalar sigma says + # nothing about the CV width. The first deposition is where + # the true dimension becomes known. + if int(self.hill_count) != 0: + raise RuntimeError( + f"WellTemperedMetaDynamicsBias {self.name!r}: cv " + f"returned {dim} component(s) but the " + f"{int(self.hill_count)} hill(s) already deposited " + f"have {self.hill_centers.shape[1]}. A collective " + "variable must keep its dimension for the whole run — " + "the existing history is not comparable otherwise." + ) + self._allocate(self.capacity, dim) + + slots = self._next_slots(count) + + self.hill_centers[slots] = values.to(self.hill_centers.dtype) + self.hill_heights[slots] = heights.to(self.hill_heights.dtype) + self.hill_owner[slots] = owner + self.hill_step[slots] = self.deposits + + self.deposits += 1 + self.hills_written += count + self.hill_count = torch.clamp(self.hills_written, max=self.capacity) + self.bump_state_version() + + # ------------------------------------------------------------------ + # Analysis + # ------------------------------------------------------------------ + + def free_energy(self, values: Tensor, owner_key: Tensor | None = None) -> Tensor: + """Return the free-energy estimate at *values*. + + ``F(s) = -(gamma / (gamma - 1)) V(s)`` for well-tempered, or + ``F(s) = -V(s)`` for standard metadynamics. + + Parameters + ---------- + values: + CV values, shape ``[B, D]``. + owner_key: + History key per graph. Defaults to the shared history. + + Returns + ------- + Tensor + Free energy in eV, shape ``[B]``, up to an additive constant. + + Raises + ------ + ValueError + If *values* does not match the deposited hills' dimension, or + disagrees with ``sigma`` or ``periods``. + RuntimeError + Under ``storage="fifo"``, where discarded hills mean the + accumulated bias is no longer the integral of everything + deposited and the well-tempered relation does not hold. + """ + if self.storage == "fifo": + raise RuntimeError( + f"WellTemperedMetaDynamicsBias {self.name!r}: free_energy() is " + "not valid under storage='fifo'. Discarding the oldest hills " + "means the accumulated bias is no longer the integral of " + "everything deposited, so the well-tempered relation between " + "bias and free energy does not hold. Use 'preallocated' or " + "'grow' for a run you intend to reconstruct a profile from." + ) + # gaussian_sum answers a width mismatch with zeros, which is right for + # an empty history on the force path but reads as a flat profile here. + dim = self._validate_cv(values) + if int(self.hill_count) and self.hill_centers.shape[1] != dim: + raise ValueError( + f"WellTemperedMetaDynamicsBias {self.name!r}: free_energy() " + f"was given {dim}-component values but the deposited hills " + f"have {self.hill_centers.shape[1]}. Evaluating the profile " + "on a grid of the wrong width would return zeros, which is " + "indistinguishable from a flat free energy." + ) + if owner_key is None: + owner_key = torch.full( + (values.shape[0],), + -1 if self.history == "shared" else 0, + dtype=torch.long, + device=values.device, + ) + bias = self.gaussian_sum(values, owner_key) + if self.bias_factor is None: + return -bias + return -bias * (self.bias_factor / (self.bias_factor - 1.0)) + + def config_fingerprint(self) -> dict[str, Any]: + """Return the settings the deposited hills are only valid under. + + ``sigma`` and ``periods`` are held as buffers, so without this check + ``nn.Module.load_state_dict`` would replace the caller's values with + the checkpoint's rather than reject the mismatch. ``bias_factor`` + and ``temperature`` set the well-tempered damping the stored heights + were computed under, and the ratio ``free_energy`` applies to them. + ``storage`` and ``history`` decide retention and ownership. + + ``max_hills`` is deliberately absent: ``grow`` legitimately reaches a + capacity the constructor never had, and :meth:`load_state_dict` + already resizes to match. + + Returns + ------- + dict[str, Any] + The checked configuration. + """ + return { + "height": self.height, + "sigma": self.sigma.reshape(-1).tolist(), + "temperature": self.temperature, + "bias_factor": self.bias_factor, + "storage": self.storage, + "history": self.history, + "ramp_depositions": self.ramp_depositions, + "periods": None + if self.periods is None + else self.periods.reshape(-1).tolist(), + } + + def load_state_dict( + self, state: Mapping[str, Any], *args: Any, **kwargs: Any + ) -> Any: + """Restore state, resizing the hill buffers to match the checkpoint. + + ``storage="grow"`` means the saved capacity is whatever the run had + reached, which is almost never the constructor's initial chunk — + ``nn.Module.load_state_dict`` would reject the size mismatch. + + Parameters + ---------- + state: + Mapping from :meth:`state_dict`. + *args, **kwargs: + Forwarded up the MRO. + + Returns + ------- + Any + Whatever the next ``load_state_dict`` returns. + """ + centers = state.get("hill_centers") + if centers is not None and tuple(centers.shape) != tuple( + self.hill_centers.shape + ): + self._allocate(int(centers.shape[0]), int(centers.shape[1])) + return super().load_state_dict(state, *args, **kwargs) + + def __repr__(self) -> str: + """Return a concise description of the bias.""" + gamma = "None" if self.bias_factor is None else f"{self.bias_factor:g}" + return ( + f"{type(self).__name__}(name={self.name!r}, " + f"hills={int(self.hill_count)}/{self.capacity}, " + f"storage={self.storage!r}, history={self.history!r}, " + f"bias_factor={gamma})" + ) diff --git a/nvalchemi/enhanced_sampling/biases/rmsd_metad.py b/nvalchemi/enhanced_sampling/biases/rmsd_metad.py new file mode 100644 index 00000000..2e48c8b6 --- /dev/null +++ b/nvalchemi/enhanced_sampling/biases/rmsd_metad.py @@ -0,0 +1,838 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""xTB-style Cartesian RMSD metadynamics for conformer exploration.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +import torch +from torch import Tensor + +from nvalchemi.enhanced_sampling._adaptive import AdaptivePotentialMixin +from nvalchemi.enhanced_sampling._bias import ConservativeBias + +if TYPE_CHECKING: + from collections.abc import Mapping + + from nvalchemi.data import Batch + from nvalchemi.enhanced_sampling._bias import BiasResult + +__all__ = ["RMSDMetaDynamicsBias"] + +_HISTORY_MODES = ("shared", "state", "walker") +_STORAGE_POLICIES = ("preallocated", "grow", "fifo") + + +def _squared_rmsd(coords: Tensor, references: Tensor) -> Tensor: + r"""Return the optimally aligned squared RMSD, shape ``[B, R]``. + + Alignment is solved by the quaternion characteristic-polynomial route + (Coutsias, Seok & Dill 2004) rather than an SVD Kabsch: the maximum of + :math:`\mathrm{tr}(\mathbf{R}\mathbf{H})` over proper rotations is the + largest eigenvalue of a symmetric ``4 x 4`` key matrix built from the + covariance ``H``. Two properties matter here: + + * The proper-rotation constraint is built in. An SVD needs an explicit + ``det`` correction, and the sign flip is a non-differentiable branch on + the compiled path. + * Only an eigen*value* is needed, never an eigenvector or a rotation + matrix. Eigenvector gradients blow up when singular values are nearly + degenerate, which is exactly what happens for symmetric-top and linear + molecules; the largest eigenvalue stays simple. + + The result is the *squared* RMSD throughout. ``sqrt`` has infinite + derivative at zero, and a reference structure is visited at RMSD zero + every time one is deposited, so taking the square root would put a force + singularity at the most frequently sampled point of the run. + + Parameters + ---------- + coords: + Current coordinates, shape ``[B, M, 3]``. + references: + Reference coordinates, shape ``[R, M, 3]``. Assumed already + centered on their centroid. + + Returns + ------- + Tensor + Squared RMSD in ``A^2``, shape ``[B, R]``, clamped at zero. + """ + n_atoms = coords.shape[1] + centered = coords - coords.mean(dim=1, keepdim=True) # [B, M, 3] + + # Inner products, both invariant to rotation. + g_x = (centered**2).sum(dim=(1, 2)) # [B] + g_y = (references**2).sum(dim=(1, 2)) # [R] + + # Covariance for every (walker, reference) pair. + cov = torch.einsum("bmi,rmj->brij", centered, references) # [B, R, 3, 3] + xx, xy, xz = cov[..., 0, 0], cov[..., 0, 1], cov[..., 0, 2] + yx, yy, yz = cov[..., 1, 0], cov[..., 1, 1], cov[..., 1, 2] + zx, zy, zz = cov[..., 2, 0], cov[..., 2, 1], cov[..., 2, 2] + + key = torch.stack( + [ + torch.stack([xx + yy + zz, yz - zy, zx - xz, xy - yx], dim=-1), + torch.stack([yz - zy, xx - yy - zz, xy + yx, zx + xz], dim=-1), + torch.stack([zx - xz, xy + yx, -xx + yy - zz, yz + zy], dim=-1), + torch.stack([xy - yx, zx + xz, yz + zy, -xx - yy + zz], dim=-1), + ], + dim=-2, + ) # [B, R, 4, 4] + + lambda_max = torch.linalg.eigvalsh(key)[..., -1] # [B, R] + msd = (g_x.unsqueeze(1) + g_y.unsqueeze(0) - 2.0 * lambda_max) / n_atoms + # Exactly-aligned pairs land a rounding step below zero. + return torch.clamp(msd, min=0.0) + + +class RMSDMetaDynamicsBias(AdaptivePotentialMixin, ConservativeBias): + r"""Repulsive Gaussian bias over retained reference structures. + + .. math:: + + V(x, t) = \sum_r f_r(t)\, k_\mathrm{push} + \exp\!\left(-\alpha\, \mathrm{RMSD}_A(x, x_r)^2\right) + + where :math:`\mathrm{RMSD}_A` is the Cartesian RMSD after optimal + translation and rotation. Each deposition appends the current geometry + to the reference set, so the bias pushes the system away from everywhere + it has already been. This is the conformer-exploration scheme used by + xTB/CREST metadynamics. + + Unlike :class:`~nvalchemi.enhanced_sampling.WellTemperedMetaDynamicsBias`, + there is no collective variable to choose: the "CV" is the whole + (selected) geometry, which is what makes it a general-purpose explorer + for molecules whose interesting degrees of freedom are not known in + advance. The price is that it has no free-energy interpretation — it is + a structure-generation method, not an estimator, and there is + deliberately no ``free_energy`` here. + + Parameters + ---------- + k_push: + Per-reference bias amplitude in eV. Must be positive; a negative + amplitude would attract the system to structures it has already + visited. + alpha: + Gaussian width parameter in ``A^-2``. Larger is narrower. + name: + Unique bias identifier. + atom_indices: + Atom indices *within each graph* to align and compare, shape ``[M]``. + ``None`` uses every atom. The usual choice is heavy atoms only: + hydrogens rotating on a methyl group produce RMSD the exploration + does not care about. + update_frequency: + Dynamics steps between depositions. + storage: + ``"fifo"`` (default), ``"preallocated"``, or ``"grow"``. FIFO is the + xTB-compatible default and is not a compromise here the way it is for + well-tempered metadynamics: with no free energy to reconstruct, + discarding the oldest references is a deliberate choice to keep + pushing outward rather than accumulating an ever-stiffer cage. + max_references: + Capacity. Required for ``"preallocated"`` and ``"fifo"``; the + initial chunk for ``"grow"``. + history: + ``"shared"``, ``"state"``, or ``"walker"``, as for well-tempered + metadynamics. The latter two require the batch to carry + ``thermodynamic_state_id`` / ``walker_id`` and raise if it does not; + :class:`~nvalchemi.enhanced_sampling.EnhancedSampling` stamps both on + every step, so only direct evaluation has to supply them. + ramp_depositions: + Deposition events over which a new reference ramps from zero to full + amplitude. A freshly deposited reference is at RMSD zero, so without + a ramp it switches on at its full value exactly where the system is + standing, which is the worst possible place for a force + discontinuity. Defaults to ``1``. + references: + Optional warm start, shape ``[R, M, 3]``, in chronological order. + The xTB convention of seeding with a displaced copy of the initial + structure is expressed by passing it here. + compute_stress: + Passed through to :class:`ConservativeBias`. Defaults to ``False``: + this bias rejects periodic systems, so there is no cell to + differentiate with respect to. + + Raises + ------ + ValueError + For a non-positive ``k_push`` or ``alpha``, an invalid storage + policy or history mode, a missing capacity, a malformed + ``atom_indices``, or warm-start references of the wrong shape. + + Notes + ----- + Periodic systems are rejected + Cartesian RMSD against a stored reference is not well defined under + periodic boundary conditions: an atom that diffuses across a cell + face is physically unmoved but Cartesian-displaced by a lattice + vector, so the RMSD jumps and the bias delivers a large spurious + force. Making this correct needs a minimum-image-aware, + correspondence-resolving metric. Rather than return a plausible + wrong number, :meth:`evaluate` raises. + + Periodicity is read from ``batch.pbc``, not from the presence of a + cell. A molecular batch carrying a **bounding box** with ``pbc`` + all-False is accepted, which is the common case for a solvated or + boxed molecule; a slab (``pbc=[True, True, False]``) is rejected, + since wrapping along any axis is enough to break the metric. A + non-zero cell with no ``pbc`` flags at all is refused as undeclared + rather than assumed non-periodic. + + Fixed atom correspondence + Atom ``i`` is always compared against atom ``i`` of the reference. + No permutation search is performed, so two structures identical up to + relabelling of equivalent atoms register as distinct. + + Examples + -------- + >>> import torch + >>> from nvalchemi.enhanced_sampling import RMSDMetaDynamicsBias + >>> bias = RMSDMetaDynamicsBias( + ... k_push=0.02, alpha=0.5, max_references=50, + ... atom_indices=torch.tensor([0, 1, 2]), + ... ) + >>> bias.reference_count.item() + 0 + """ + + def __init__( + self, + k_push: float, + alpha: float, + *, + name: str = "rmsd_metadynamics", + atom_indices: Tensor | None = None, + update_frequency: int = 500, + storage: Literal["preallocated", "grow", "fifo"] = "fifo", + max_references: int | None = None, + history: Literal["shared", "state", "walker"] = "shared", + ramp_depositions: int = 1, + references: Tensor | None = None, + compute_stress: bool = False, + ) -> None: + super().__init__(name=name, compute_stress=compute_stress) + + if storage not in _STORAGE_POLICIES: + raise ValueError( + f"RMSDMetaDynamicsBias: storage must be one of " + f"{list(_STORAGE_POLICIES)}, got {storage!r}." + ) + if history not in _HISTORY_MODES: + raise ValueError( + f"RMSDMetaDynamicsBias: history must be one of " + f"{list(_HISTORY_MODES)}, got {history!r}." + ) + if k_push <= 0.0: + raise ValueError( + f"RMSDMetaDynamicsBias: k_push must be positive, got {k_push}. " + "A non-positive amplitude would attract the system back to " + "structures it has already visited." + ) + if alpha <= 0.0: + raise ValueError( + f"RMSDMetaDynamicsBias: alpha must be positive, got {alpha}." + ) + if int(update_frequency) < 1: + raise ValueError( + f"RMSDMetaDynamicsBias: update_frequency must be at least 1, " + f"got {update_frequency}." + ) + if int(ramp_depositions) < 0: + raise ValueError( + f"RMSDMetaDynamicsBias: ramp_depositions must be non-negative, " + f"got {ramp_depositions}." + ) + if storage in ("preallocated", "fifo") and max_references is None: + raise ValueError( + f"RMSDMetaDynamicsBias: storage={storage!r} needs an explicit " + "max_references — it is the whole point of the policy, either " + "the ceiling that raises or the ring that overwrites." + ) + capacity = int(max_references) if max_references is not None else 64 + if capacity < 1: + raise ValueError( + f"RMSDMetaDynamicsBias: max_references must be at least 1, got " + f"{max_references}." + ) + + if atom_indices is None: + selection: Tensor | None = None + else: + selection = torch.as_tensor(atom_indices, dtype=torch.long).reshape(-1) + if selection.numel() == 0: + raise ValueError( + "RMSDMetaDynamicsBias: atom_indices is empty; pass None to " + "use every atom." + ) + if bool((selection < 0).any()): + raise ValueError( + f"RMSDMetaDynamicsBias: atom_indices must be non-negative " + f"per-graph indices, got {selection.tolist()}." + ) + if selection.unique().numel() != selection.numel(): + raise ValueError( + f"RMSDMetaDynamicsBias: atom_indices contains duplicates " + f"({selection.tolist()}). A repeated atom is silently " + "weighted twice in the RMSD." + ) + + self.k_push = float(k_push) + self.alpha = float(alpha) + self.storage = storage + self.history = history + self.update_frequency = int(update_frequency) + self.ramp_depositions = int(ramp_depositions) + self._capacity = capacity + + if selection is None: + self.atom_indices: Tensor | None = None + else: + self.register_buffer("atom_indices", selection) + + n_sites = 0 if references is None else int(references.shape[1]) + if selection is not None: + n_sites = selection.numel() + self._allocate(capacity, max(n_sites, 1)) + + if references is not None: + self._seed(references) + + # The reference set is per-state under "state" history, so an accepted + # swap changes which references a walker feels. + self.state_dependent_for_exchange = history == "state" + + # ------------------------------------------------------------------ + # Storage + # ------------------------------------------------------------------ + + def _allocate(self, capacity: int, n_sites: int) -> None: + """Create or replace the reference buffers. + + Parameters + ---------- + capacity: + Number of reference slots. + n_sites: + Number of compared atoms per reference. + """ + dtype = torch.get_default_dtype() + device = getattr(self, "reference_coords", torch.empty(0)).device + self.register_buffer( + "reference_coords", + torch.zeros(capacity, n_sites, 3, dtype=dtype, device=device), + ) + self.register_buffer( + "reference_owner", + torch.full((capacity,), -1, dtype=torch.long, device=device), + ) + self.register_buffer( + "reference_step", + torch.full((capacity,), -1, dtype=torch.long, device=device), + ) + # reference_count saturates at capacity; references_written keeps + # counting, which fixes the FIFO ring position after it has wrapped. + self.register_buffer( + "reference_count", torch.zeros((), dtype=torch.long, device=device) + ) + self.register_buffer( + "references_written", torch.zeros((), dtype=torch.long, device=device) + ) + self.register_buffer( + "deposits", torch.zeros((), dtype=torch.long, device=device) + ) + + def _seed(self, references: Tensor) -> None: + """Install warm-start references in chronological order. + + Parameters + ---------- + references: + Shape ``[R, M, 3]``, oldest first. + + Raises + ------ + ValueError + If the shape is not ``[R, M, 3]``, conflicts with + ``atom_indices``, or exceeds the capacity. + """ + refs = torch.as_tensor(references, dtype=self.reference_coords.dtype) + if refs.ndim != 3 or refs.shape[-1] != 3: + raise ValueError( + f"RMSDMetaDynamicsBias: references must have shape [R, M, 3], " + f"got {tuple(refs.shape)}." + ) + if self.atom_indices is not None and refs.shape[1] != self.atom_indices.numel(): + raise ValueError( + f"RMSDMetaDynamicsBias: references have {refs.shape[1]} atoms " + f"but atom_indices selects {self.atom_indices.numel()}." + ) + if refs.shape[0] > self._capacity: + raise ValueError( + f"RMSDMetaDynamicsBias: {refs.shape[0]} warm-start references " + f"exceed max_references={self._capacity}." + ) + if refs.shape[1] != self.reference_coords.shape[1]: + self._allocate(self._capacity, int(refs.shape[1])) + + count = int(refs.shape[0]) + # Stored centered: the metric is translation invariant, so removing the + # centroid once at deposition keeps it out of every later comparison. + centered = refs - refs.mean(dim=1, keepdim=True) + self.reference_coords[:count] = centered.to(self.reference_coords.device) + self.reference_owner[:count] = -1 if self.history == "shared" else 0 + # Warm-start references are already fully active; a ramp is only for + # references deposited into a running trajectory. + self.reference_step[:count] = -self.ramp_depositions - 1 + self.reference_count += count + self.references_written += count + + def _grow(self) -> None: + """Extend the reference buffers by one more chunk.""" + extra = self._capacity + device = self.reference_coords.device + + def _extend(buffer: Tensor, fill: float | int) -> Tensor: + pad = torch.full( + (extra, *buffer.shape[1:]), fill, dtype=buffer.dtype, device=device + ) + return torch.cat([buffer, pad], dim=0) + + self.reference_coords = _extend(self.reference_coords, 0.0) + self.reference_owner = _extend(self.reference_owner, -1) + self.reference_step = _extend(self.reference_step, -1) + + @property + def capacity(self) -> int: + """Return the current number of reference slots.""" + return int(self.reference_coords.shape[0]) + + def _next_slots(self, count: int) -> Tensor: + """Return the slot indices *count* new references will occupy. + + Parameters + ---------- + count: + Number of references about to be deposited. + + Returns + ------- + Tensor + Slot indices, shape ``[count]``. + + Raises + ------ + RuntimeError + If ``preallocated`` capacity is exhausted. + """ + written = int(self.references_written) + device = self.reference_owner.device + + if self.storage == "fifo": + return (torch.arange(count, device=device) + written) % self.capacity + + if written + count > self.capacity: + if self.storage == "preallocated": + raise RuntimeError( + f"RMSDMetaDynamicsBias {self.name!r}: reference storage is " + f"full ({self.capacity} references) and " + "storage='preallocated'. Raise max_references, lengthen " + "update_frequency, or switch to storage='fifo', which is " + "the xTB-compatible policy and discards the oldest " + "reference instead." + ) + while written + count > self.capacity: + self._grow() + + return torch.arange(count, device=device) + written + + # ------------------------------------------------------------------ + # Energy + # ------------------------------------------------------------------ + + def _gather_sites(self, current: Batch) -> Tensor: + """Return the compared coordinates per graph, shape ``[B, M, 3]``. + + Parameters + ---------- + current: + The live batch. + + Returns + ------- + Tensor + Selected positions, differentiable with respect to + ``current.positions``. + """ + positions = current.positions + batch_ptr = current.batch_ptr.to(positions.device) + offsets = batch_ptr[:-1] # [B] + + if self.atom_indices is None: + # Uniform atom counts are checked in evaluate(), so the -1 here + # resolves without a data-dependent int() that would break the graph. + return positions.reshape(offsets.numel(), -1, 3) + + selection = self.atom_indices.to(positions.device) + flat = offsets.unsqueeze(1) + selection.unsqueeze(0) # [B, M] + return positions[flat.reshape(-1)].reshape(offsets.numel(), -1, 3) + + def _owner_key(self, current: Batch, n_graphs: int) -> Tensor: + """Return the history key per graph, shape ``[B]``. + + Parameters + ---------- + current: + The live batch. + n_graphs: + Number of graphs. + + Returns + ------- + Tensor + Per-graph key matched against ``reference_owner``. + """ + device = self.reference_owner.device + if self.history == "state": + field = "thermodynamic_state_id" + elif self.history == "walker": + field = "walker_id" + else: + return torch.full((n_graphs,), -1, dtype=torch.long, device=device) + + ids = getattr(current, field, None) + if ids is None: + raise ValueError( + f"RMSDMetaDynamicsBias {self.name!r}: history=" + f"{self.history!r} needs batch.{field}, which this batch " + f"does not carry. Falling back to a single owner would put " + f"every reference under one key and silently collapse the " + f"per-{field} histories into one shared history — the " + f"opposite of what history={self.history!r} asks for. " + "EnhancedSampling stamps this field on every step; a bias " + "driven directly must set it, or use history='shared' if " + "one history really is intended." + ) + if ids.numel() != n_graphs: + raise ValueError( + f"RMSDMetaDynamicsBias {self.name!r}: batch.{field} has " + f"{ids.numel()} entries but the batch has {n_graphs} " + f"graph(s). A shorter tensor broadcasts across graphs, which " + f"would file every reference under one walker's key without " + "raising." + ) + return ids.reshape(-1).to(device=device, dtype=torch.long) + + def _reference_scale(self) -> Tensor: + """Return each reference's ramp factor, shape ``[capacity]``. + + Returns + ------- + Tensor + Factors in ``[0, 1]``. + """ + if self.ramp_depositions <= 0: + return torch.ones( + self.capacity, + dtype=self.reference_coords.dtype, + device=self.reference_coords.device, + ) + age = (self.deposits - self.reference_step).to(self.reference_coords.dtype) + return torch.clamp(age / float(self.ramp_depositions), min=0.0, max=1.0) + + def repulsion(self, coords: Tensor, owner_key: Tensor) -> Tensor: + """Return the accumulated repulsion, shape ``[B]``. + + Contains no data-dependent Python branch, so it compiles with + ``fullgraph=True``. + + Parameters + ---------- + coords: + Compared coordinates, shape ``[B, M, 3]``. + owner_key: + History key per graph, shape ``[B]``. + + Returns + ------- + Tensor + Bias energy per graph in eV, shape ``[B]``. + """ + if self.reference_coords.shape[1] != coords.shape[1]: + # A bias built without atom_indices or warm-start references does + # not learn its site count until the first deposition; until then + # the reference set is empty and the bias is identically zero. + # This is a shape guard, not a data-dependent branch. + return coords.new_zeros(coords.shape[0]) + + slots = torch.arange(self.capacity, device=coords.device) + live = slots < self.reference_count + + msd = _squared_rmsd(coords, self.reference_coords) # [B, R] + kernel = torch.exp(-self.alpha * msd) + + weights = self.k_push * self._reference_scale() # [R] + mask = live.unsqueeze(0) + if self.history != "shared": + mask = mask & (self.reference_owner.unsqueeze(0) == owner_key.unsqueeze(1)) + return (kernel * weights.unsqueeze(0) * mask).sum(dim=-1) + + def evaluate(self, current: Batch) -> BiasResult: + """Reject periodic batches, then derive energy and forces. + + Parameters + ---------- + current: + The live batch. + + Returns + ------- + BiasResult + As :meth:`ConservativeBias.evaluate`. + + Raises + ------ + ValueError + If the batch is periodic according to ``batch.pbc``, if it + carries a non-zero cell with no ``pbc`` flags, or if graphs in + the batch have differing atom counts while ``atom_indices`` is + ``None``. + """ + self._reject_periodic(current) + self._validate_sites(current) + return super().evaluate(current) + + def _reject_periodic(self, current: Batch) -> None: + """Raise if the batch is under periodic boundary conditions. + + Periodicity is read from ``batch.pbc`` when it is present, matching + :func:`~nvalchemi.enhanced_sampling.pair_distance` and the rest of + the toolkit: a cell is a box, and only the flags say whether atoms + wrap through its faces. A molecular batch carrying a bounding box + with ``pbc`` all-False is therefore accepted — it is exactly the + non-periodic case this bias is for. + + Without ``pbc`` there is nothing to read, and a non-zero cell is + refused rather than assumed harmless. A batch that declares a cell + but no boundary condition has not said which case it is, and the + failure this guard exists to prevent is silent. + + Runs in ``evaluate`` rather than ``energy`` so the flag reduction + stays off the compiled path. + + Parameters + ---------- + current: + The live batch. + + Raises + ------ + ValueError + If any graph is periodic along any axis, or if a non-zero cell + is present with no ``pbc`` flags to interpret it. + """ + cell = getattr(current, "cell", None) + if cell is None or not bool((cell != 0).any()): + return + + pbc = getattr(current, "pbc", None) + if pbc is None: + raise ValueError( + f"RMSDMetaDynamicsBias {self.name!r}: the batch carries a " + "non-zero cell but no pbc flags, so whether atoms wrap " + "through its faces is undeclared. Set pbc=False on a " + "molecular system with a bounding box; Cartesian RMSD is " + "only defined for the non-periodic case." + ) + if not bool(pbc.any()): + return + + raise ValueError( + f"RMSDMetaDynamicsBias {self.name!r}: the batch is periodic " + f"(pbc={pbc.reshape(-1, pbc.shape[-1])[0].tolist()}), and " + "Cartesian RMSD against a stored reference is not defined under " + "periodic boundary conditions — an atom crossing a cell face is " + "physically unmoved but Cartesian-displaced by a lattice vector, " + "which would inject a large spurious force. Use this bias on " + "non-periodic (molecular) systems — a bounding-box cell with " + "pbc=False is fine — or bias a periodic-aware CV with " + "WellTemperedMetaDynamicsBias instead." + ) + + def _validate_sites(self, current: Batch) -> None: + """Check that the compared-atom selection is well defined. + + Parameters + ---------- + current: + The live batch. + + Raises + ------ + ValueError + If atom counts differ across graphs without an explicit + selection, or a selected index is out of range for some graph. + """ + batch_ptr = current.batch_ptr + counts = batch_ptr[1:] - batch_ptr[:-1] + + if self.atom_indices is None: + if bool((counts != counts[0]).any()): + raise ValueError( + f"RMSDMetaDynamicsBias {self.name!r}: graphs in the batch " + f"have differing atom counts {counts.tolist()}, so there " + "is no fixed atom correspondence to compare against a " + "shared reference set. Pass atom_indices to select a " + "common set of atoms per graph." + ) + return + + selection = self.atom_indices.to(counts.device) + smallest = int(counts.min()) + if int(selection.max()) >= smallest: + raise ValueError( + f"RMSDMetaDynamicsBias {self.name!r}: atom_indices requests " + f"local index {int(selection.max())} but the smallest graph in " + f"the batch has only {smallest} atoms. atom_indices are " + "per-graph local indices, not global ones." + ) + + def energy(self, current: Batch) -> Tensor: + """Return the accumulated RMSD bias, shape ``[B, 1]``. + + Parameters + ---------- + current: + Batch with strained positions supplied by + :meth:`ConservativeBias.evaluate`. + + Returns + ------- + Tensor + Shape ``[B, 1]`` in eV. + """ + coords = self._gather_sites(current) + owner = self._owner_key(current, coords.shape[0]) + return self.repulsion(coords, owner).unsqueeze(-1) + + # ------------------------------------------------------------------ + # Deposition + # ------------------------------------------------------------------ + + def update(self, frames: Batch, result: BiasResult) -> None: + """Append the current geometry to the reference set, one per walker. + + Parameters + ---------- + frames: + Post-step frame captured by the runner. + result: + The bias's own result from the preceding force evaluation. + """ + with torch.no_grad(): + coords = self._gather_sites(frames).detach() # [B, M, 3] + centered = coords - coords.mean(dim=1, keepdim=True) + owner = self._owner_key(frames, coords.shape[0]) + + count = coords.shape[0] + if self.reference_coords.shape[1] != centered.shape[1]: + # First deposition of a bias built without atom_indices or + # warm-start references: the site count is only known now. + self._allocate(self.capacity, int(centered.shape[1])) + slots = self._next_slots(count) + + self.reference_coords[slots] = centered.to(self.reference_coords.dtype) + self.reference_owner[slots] = owner + self.reference_step[slots] = self.deposits + + self.deposits += 1 + self.references_written += count + self.reference_count = torch.clamp( + self.references_written, max=self.capacity + ) + self.bump_state_version() + + def config_fingerprint(self) -> dict[str, Any]: + """Return the settings the retained references are only valid under. + + ``k_push`` and ``alpha`` set how hard and how far the stored + structures repel; restoring a reference set under different values + keeps the geometries but changes what they do. ``atom_indices`` is + held as a buffer and names which atoms the stored coordinates are, + so without this check the caller's selection would be silently + overwritten by the checkpoint's. + + ``max_references`` is absent for the same reason ``max_hills`` is + absent from the well-tempered fingerprint. + + Returns + ------- + dict[str, Any] + The checked configuration. + """ + return { + "k_push": self.k_push, + "alpha": self.alpha, + "storage": self.storage, + "history": self.history, + "ramp_depositions": self.ramp_depositions, + "atom_indices": ( + None + if self.atom_indices is None + else self.atom_indices.reshape(-1).tolist() + ), + } + + def load_state_dict( + self, state: Mapping[str, Any], *args: Any, **kwargs: Any + ) -> Any: + """Restore state, resizing the reference buffers to the checkpoint. + + Both ``storage="grow"`` and a bias whose site count was fixed at its + first deposition can carry buffer shapes that the constructor did not + produce, which ``nn.Module.load_state_dict`` would reject. + + Parameters + ---------- + state: + Mapping from :meth:`state_dict`. + *args, **kwargs: + Forwarded up the MRO. + + Returns + ------- + Any + Whatever the next ``load_state_dict`` returns. + """ + coords = state.get("reference_coords") + if coords is not None and tuple(coords.shape) != tuple( + self.reference_coords.shape + ): + self._allocate(int(coords.shape[0]), int(coords.shape[1])) + return super().load_state_dict(state, *args, **kwargs) + + def __repr__(self) -> str: + """Return a concise description of the bias.""" + sites = "all" if self.atom_indices is None else int(self.atom_indices.numel()) + return ( + f"{type(self).__name__}(name={self.name!r}, " + f"references={int(self.reference_count)}/{self.capacity}, " + f"storage={self.storage!r}, history={self.history!r}, " + f"k_push={self.k_push:g}, alpha={self.alpha:g}, atoms={sites})" + ) diff --git a/nvalchemi/enhanced_sampling/biases/umbrella.py b/nvalchemi/enhanced_sampling/biases/umbrella.py new file mode 100644 index 00000000..2f942fcc --- /dev/null +++ b/nvalchemi/enhanced_sampling/biases/umbrella.py @@ -0,0 +1,352 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Harmonic umbrella bias for umbrella sampling.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch import Tensor + +from nvalchemi.enhanced_sampling._bias import BiasResult, ConservativeBias +from nvalchemi.enhanced_sampling.cv._periodic import periodic_difference + +if TYPE_CHECKING: + from collections.abc import Callable + + from nvalchemi.data import Batch + +__all__ = ["HarmonicUmbrellaBias"] + + +class HarmonicUmbrellaBias(ConservativeBias): + r"""Multi-dimensional harmonic restraint on one or more collective variables. + + .. math:: + + E_b = \tfrac{1}{2}\,\Delta^\top K \,\Delta, + \qquad \Delta = s(x) - s_0 + + with ``s(x)`` the CV value, ``s_0`` the window center, and ``K`` the + stiffness matrix. Forces and stress come from + :class:`~nvalchemi.enhanced_sampling.ConservativeBias`, so they are + guaranteed consistent with this energy. + + Per-state parameters + -------------------- + ``centers`` and ``stiffness`` may carry a leading state dimension ``S``. + Each graph then selects its row by ``batch.thermodynamic_state_id``, + which is what makes a batch of umbrella windows a single batched run + rather than ``S`` separate simulations. Without that field every graph + uses state ``0``. + + A **single-window** bias ignores the field entirely rather than treating + it as an index. That is what lets one shared restraint run alongside a + multi-rung temperature ladder, where the ids address the ladder and have + nothing to do with windows. + + Parameters + ---------- + cv: + Any differentiable ``cv(batch) -> Tensor[B, D]``. A plain callable; + no base class, no registration. + centers: + Window centers. Shape ``[D]`` (shared) or ``[S, D]`` (per state). + stiffness: + Force constants, in energy per CV-unit squared. Accepted as: + + * scalar — isotropic, ``k·I`` + * ``[D]`` — diagonal + * ``[D, D]`` — full matrix, shared across states + * ``[S, D, D]`` — full matrix per state + name: + Unique bias identifier. + periods: + Period per CV component, shape ``[D]``; ``0`` marks a non-periodic + component. See + :func:`~nvalchemi.enhanced_sampling.cv._periodic.periodic_difference`. + compute_stress: + Passed through to :class:`ConservativeBias`. + + Raises + ------ + ValueError + If shapes are inconsistent, or if a stiffness matrix is not + symmetric positive-semidefinite. + + Examples + -------- + >>> import torch + >>> from nvalchemi.enhanced_sampling import HarmonicUmbrellaBias, pair_distance + >>> idx = torch.tensor([0, 5]) + >>> bias = HarmonicUmbrellaBias( + ... cv=lambda b: pair_distance(b, idx), + ... centers=torch.tensor([[2.0], [2.5], [3.0]]), # 3 windows + ... stiffness=10.0, # eV/A^2 + ... ) + >>> bias.centers.shape + torch.Size([3, 1]) + """ + + def __init__( + self, + cv: Callable[[Batch], Tensor], + centers: Tensor | float, + stiffness: Tensor | float, + *, + name: str = "umbrella", + periods: Tensor | None = None, + compute_stress: bool = True, + ) -> None: + super().__init__(name=name, compute_stress=compute_stress) + self.cv = cv + + centers_t = torch.as_tensor(centers, dtype=torch.get_default_dtype()) + if centers_t.ndim == 0: + centers_t = centers_t.reshape(1, 1) + elif centers_t.ndim == 1: + centers_t = centers_t.unsqueeze(0) # [D] -> [1, D] + elif centers_t.ndim != 2: + raise ValueError( + f"HarmonicUmbrellaBias: centers must be [D] or [S, D], got " + f"shape {tuple(centers_t.shape)}." + ) + n_states, dim = centers_t.shape + + stiffness_t = self._expand_stiffness(stiffness, n_states, dim) + self._validate_stiffness(stiffness_t) + + # A multi-window bias reads thermodynamic_state_id, so its energy + # depends on the assignment. Replica exchange needs to know: combining + # it with a temperature ladder would need cross-state bias terms that + # the temperature acceptance rule does not compute. + self.state_dependent_for_exchange = n_states > 1 + + # Buffers, not plain attributes: nn.Module then moves them with .to() + # and round-trips them through state_dict. + self.register_buffer("centers", centers_t) + self.register_buffer("stiffness", stiffness_t) + if periods is None: + self.periods: Tensor | None = None + else: + periods_t = torch.as_tensor( + periods, dtype=torch.get_default_dtype() + ).reshape(-1) + if periods_t.numel() != dim: + raise ValueError( + f"HarmonicUmbrellaBias: periods must have {dim} entries to " + f"match the CV dimension, got {periods_t.numel()}." + ) + self.register_buffer("periods", periods_t) + + @staticmethod + def _expand_stiffness(stiffness: Tensor | float, n_states: int, dim: int) -> Tensor: + """Broadcast any accepted stiffness form to ``[S, D, D]``. + + Parameters + ---------- + stiffness: + Scalar, ``[D]``, ``[D, D]``, or ``[S, D, D]``. + n_states: + Number of thermodynamic states ``S``. + dim: + CV dimension ``D``. + + Returns + ------- + Tensor + Shape ``[S, D, D]``. + + Raises + ------ + ValueError + If the shape is none of the accepted forms. + """ + k = torch.as_tensor(stiffness, dtype=torch.get_default_dtype()) + eye = torch.eye(dim, dtype=k.dtype) + + if k.ndim == 0: + full = k * eye + elif k.ndim == 1: + if k.numel() != dim: + raise ValueError( + f"HarmonicUmbrellaBias: diagonal stiffness must have {dim} " + f"entries to match the CV dimension, got {k.numel()}." + ) + full = torch.diag(k) + elif k.ndim == 2: + if k.shape != (dim, dim): + raise ValueError( + f"HarmonicUmbrellaBias: matrix stiffness must be " + f"[{dim}, {dim}], got {tuple(k.shape)}." + ) + full = k + elif k.ndim == 3: + if k.shape != (n_states, dim, dim): + raise ValueError( + f"HarmonicUmbrellaBias: per-state stiffness must be " + f"[{n_states}, {dim}, {dim}] to match centers, got " + f"{tuple(k.shape)}." + ) + return k.clone() + else: + raise ValueError( + f"HarmonicUmbrellaBias: stiffness must be scalar, [D], [D, D], " + f"or [S, D, D]; got shape {tuple(k.shape)}." + ) + return full.unsqueeze(0).expand(n_states, dim, dim).clone() + + @staticmethod + def _validate_stiffness(stiffness: Tensor) -> None: + """Reject a stiffness that is not symmetric positive-semidefinite. + + An asymmetric ``K`` makes the quadratic form ambiguous, and a + negative eigenvalue turns the restraint into a repulsion that drives + the CV away without bound — a runaway that is far cheaper to catch + here than to diagnose from a diverging trajectory. + + Parameters + ---------- + stiffness: + Shape ``[S, D, D]``. + + Raises + ------ + ValueError + If any state's matrix is asymmetric or has a negative eigenvalue. + """ + if not torch.allclose(stiffness, stiffness.mT, atol=1e-8): + raise ValueError( + "HarmonicUmbrellaBias: stiffness must be symmetric; got a " + "matrix that differs from its transpose." + ) + eigenvalues = torch.linalg.eigvalsh(stiffness.double()) + if bool((eigenvalues < -1e-8).any()): + worst = float(eigenvalues.min()) + raise ValueError( + f"HarmonicUmbrellaBias: stiffness must be positive-semidefinite; " + f"smallest eigenvalue is {worst:.6g}. A negative eigenvalue makes " + "the restraint repulsive along that direction." + ) + + def _validate_state_ids(self, current: Batch) -> None: + """Raise if any ``thermodynamic_state_id`` is out of range. + + Called from :meth:`evaluate`, never from :meth:`energy`. That + placement is deliberate: ``energy()`` is the path + ``EnhancedSampling(compile_biases=True)`` hands to ``torch.compile``, + and ``bool(tensor.any())`` there is a data-dependent Python branch + that breaks ``fullgraph=True`` outright. ``evaluate()`` is eager by + construction, so hoisting the check keeps it running in **every** + mode, rather than skipping it under compile the way the eager-only + guards in ``pair_distance`` must. + + Parameters + ---------- + current: + The batch; read for ``thermodynamic_state_id`` if present. + + Raises + ------ + IndexError + If a state id is negative or beyond the configured windows. + """ + state_ids = getattr(current, "thermodynamic_state_id", None) + if state_ids is None or self.centers.shape[0] == 1: + # A single-window bias selects nothing, so the field is not an + # index into it. This is the legitimate combination of one shared + # restraint with a multi-rung temperature ladder, where the ids + # address the ladder rather than the windows. + return + index = state_ids.reshape(-1).to(torch.long) + n_states = self.centers.shape[0] + out_of_range = (index < 0) | (index >= n_states) + if bool(out_of_range.any()): + bad = out_of_range.nonzero(as_tuple=False).squeeze(-1).tolist() + raise IndexError( + f"HarmonicUmbrellaBias: thermodynamic_state_id out of range for " + f"{n_states} configured window(s). Graph(s) {bad} have " + f"{index[out_of_range].tolist()}; valid ids are 0..{n_states - 1}." + ) + + def evaluate(self, current: Batch) -> BiasResult: + """Validate the state ids, then derive energy, forces, and stress. + + Parameters + ---------- + current: + The live batch. + + Returns + ------- + BiasResult + As :meth:`ConservativeBias.evaluate`. + + Raises + ------ + IndexError + If a ``thermodynamic_state_id`` is out of range. + """ + self._validate_state_ids(current) + return super().evaluate(current) + + def _select_per_graph( + self, current: Batch, values: Tensor + ) -> tuple[Tensor, Tensor]: + """Return ``(centers, stiffness)`` broadcast to this batch. + + Contains no data-dependent Python branch, so it compiles with + ``fullgraph=True``. Bounds checking lives in + :meth:`_validate_state_ids`, which :meth:`evaluate` runs first. + + Parameters + ---------- + current: + The batch; read for ``thermodynamic_state_id`` if present. + values: + CV values ``[B, D]``, used for shape and device. + + Returns + ------- + tuple[Tensor, Tensor] + ``centers`` ``[B, D]`` and ``stiffness`` ``[B, D, D]``. + """ + state_ids = getattr(current, "thermodynamic_state_id", None) + if state_ids is None or self.centers.shape[0] == 1: + index = torch.zeros(values.shape[0], dtype=torch.long, device=values.device) + else: + index = state_ids.reshape(-1).to(torch.long) + return self.centers[index], self.stiffness[index] + + def energy(self, current: Batch) -> Tensor: + """Return the harmonic restraint energy ``[B, 1]`` in eV. + + Parameters + ---------- + current: + Batch with strained positions supplied by + :meth:`ConservativeBias.evaluate`. + + Returns + ------- + Tensor + Shape ``[B, 1]``. + """ + values = self.cv(current) # [B, D] + centers, stiffness = self._select_per_graph(current, values) + delta = periodic_difference(values, centers, self.periods) # [B, D] + quadratic = torch.einsum("bi,bij,bj->b", delta, stiffness, delta) + return 0.5 * quadratic.unsqueeze(-1) # [B, 1] diff --git a/nvalchemi/enhanced_sampling/biases/walls.py b/nvalchemi/enhanced_sampling/biases/walls.py new file mode 100644 index 00000000..8107ff0e --- /dev/null +++ b/nvalchemi/enhanced_sampling/biases/walls.py @@ -0,0 +1,372 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""One-sided walls and flat-bottom restraints on a collective variable. + +All three classes here share one energy shape — a power-law penalty on how +far the CV has strayed past a threshold, and exactly zero inside the allowed +region. They differ only in which side is penalised. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch import Tensor + +from nvalchemi.enhanced_sampling._bias import ConservativeBias + +if TYPE_CHECKING: + from collections.abc import Callable + + from nvalchemi.data import Batch + +__all__ = ["FlatBottomRestraint", "LowerWall", "UpperWall"] + + +class _WallBase(ConservativeBias): + r"""Shared machinery for one- and two-sided CV penalties. + + The energy is built from ``clamp(excess, min=0) ** exponent``, which is + what keeps a wall usable as a bias at all: + + * **The graph stays connected inside the wall.** A wall implemented as + ``if inside: return zeros(B, 1)`` returns a tensor with no ``grad_fn``, + and autograd rejects such an output outright. ``clamp`` returns a + genuine zero *that is still attached*, so forces come back as zeros + rather than as an error. + * **The force is continuous at the boundary** for ``exponent >= 2``. + With ``exponent = 1`` the force jumps from ``0`` to ``k`` at the wall, + which a finite time step turns into an impulse; hence the default of 2 + and the warning below. + + Parameters + ---------- + cv: + Differentiable ``cv(batch) -> Tensor[B, D]``. + threshold: + Wall position, broadcastable to the CV shape ``[D]``. + stiffness: + Penalty prefactor ``k``, broadcastable to ``[D]``. + name: + Unique bias identifier. + exponent: + Power of the penalty. Must be ``>= 1``. + compute_stress: + Passed through to :class:`ConservativeBias`. + + Raises + ------ + ValueError + If ``exponent < 1`` or ``stiffness`` is negative. + """ + + def __init__( + self, + cv: Callable[[Batch], Tensor], + threshold: Tensor | float, + stiffness: Tensor | float, + *, + name: str, + exponent: float = 2.0, + compute_stress: bool = True, + ) -> None: + super().__init__(name=name, compute_stress=compute_stress) + if exponent < 1: + raise ValueError( + f"{type(self).__name__}: exponent must be >= 1, got {exponent}. " + "A sub-linear wall has unbounded force at the boundary." + ) + stiffness_t = torch.as_tensor( + stiffness, dtype=torch.get_default_dtype() + ).reshape(-1) + if bool((stiffness_t < 0).any()): + raise ValueError( + f"{type(self).__name__}: stiffness must be non-negative, got " + f"{stiffness_t.tolist()}. A negative wall pushes the system out " + "of the allowed region instead of back into it." + ) + self.cv = cv + self.exponent = float(exponent) + self.register_buffer( + "threshold", + torch.as_tensor(threshold, dtype=torch.get_default_dtype()).reshape(-1), + ) + self.register_buffer("stiffness", stiffness_t) + + def _excess(self, values: Tensor) -> Tensor: + """Return the signed distance past the wall, before clamping. + + Parameters + ---------- + values: + CV values ``[B, D]``. + + Returns + ------- + Tensor + Positive where the wall is violated, shape ``[B, D]``. + + Raises + ------ + NotImplementedError + If the subclass does not override. + """ + raise NotImplementedError + + def energy(self, current: Batch) -> Tensor: + """Return the wall energy ``[B, 1]`` in eV. + + Parameters + ---------- + current: + Batch with strained positions supplied by + :meth:`ConservativeBias.evaluate`. + + Returns + ------- + Tensor + Shape ``[B, 1]``; exactly zero for configurations inside the wall. + """ + values = self.cv(current) # [B, D] + excess = torch.clamp(self._excess(values), min=0.0) + penalty = self.stiffness * excess**self.exponent + return penalty.sum(dim=-1, keepdim=True) / self.exponent # [B, 1] + + +class UpperWall(_WallBase): + r"""Penalise the CV for rising above a threshold. + + .. math:: E_b = \frac{k}{p}\,\max(s - s_0,\,0)^p + + Parameters + ---------- + cv: + Differentiable ``cv(batch) -> Tensor[B, D]``. + threshold: + Upper bound ``s_0``. + stiffness: + Penalty prefactor ``k``. Default ``10.0``. + name: + Unique bias identifier. Default ``"upper_wall"``. + exponent: + Power ``p``. Default ``2.0``. + compute_stress: + Passed through to :class:`ConservativeBias`. + + Examples + -------- + >>> import torch + >>> from nvalchemi.enhanced_sampling import UpperWall, pair_distance + >>> idx = torch.tensor([0, 1]) + >>> wall = UpperWall(cv=lambda b: pair_distance(b, idx), threshold=5.0) + >>> wall.name + 'upper_wall' + """ + + def __init__( + self, + cv: Callable[[Batch], Tensor], + threshold: Tensor | float, + stiffness: Tensor | float = 10.0, + *, + name: str = "upper_wall", + exponent: float = 2.0, + compute_stress: bool = True, + ) -> None: + super().__init__( + cv, + threshold, + stiffness, + name=name, + exponent=exponent, + compute_stress=compute_stress, + ) + + def _excess(self, values: Tensor) -> Tensor: + """Return ``values - threshold``. + + Parameters + ---------- + values: + CV values ``[B, D]``. + + Returns + ------- + Tensor + Shape ``[B, D]``. + """ + return values - self.threshold + + +class LowerWall(_WallBase): + r"""Penalise the CV for falling below a threshold. + + .. math:: E_b = \frac{k}{p}\,\max(s_0 - s,\,0)^p + + Parameters + ---------- + cv: + Differentiable ``cv(batch) -> Tensor[B, D]``. + threshold: + Lower bound ``s_0``. + stiffness: + Penalty prefactor ``k``. Default ``10.0``. + name: + Unique bias identifier. Default ``"lower_wall"``. + exponent: + Power ``p``. Default ``2.0``. + compute_stress: + Passed through to :class:`ConservativeBias`. + """ + + def __init__( + self, + cv: Callable[[Batch], Tensor], + threshold: Tensor | float, + stiffness: Tensor | float = 10.0, + *, + name: str = "lower_wall", + exponent: float = 2.0, + compute_stress: bool = True, + ) -> None: + super().__init__( + cv, + threshold, + stiffness, + name=name, + exponent=exponent, + compute_stress=compute_stress, + ) + + def _excess(self, values: Tensor) -> Tensor: + """Return ``threshold - values``. + + Parameters + ---------- + values: + CV values ``[B, D]``. + + Returns + ------- + Tensor + Shape ``[B, D]``. + """ + return self.threshold - values + + +class FlatBottomRestraint(_WallBase): + r"""Confine the CV to an interval, with no force inside it. + + .. math:: + + E_b = \frac{k}{p}\left[ + \max(s - s_\mathrm{hi},\,0)^p + \max(s_\mathrm{lo} - s,\,0)^p + \right] + + Equivalent to registering a :class:`LowerWall` and an :class:`UpperWall` + with the same stiffness, but as one bias — which matters because the + proposal's rule is that intentionally coupled terms belong in a single + bias object rather than being summed by the runner. + + Parameters + ---------- + cv: + Differentiable ``cv(batch) -> Tensor[B, D]``. + lower: + Lower bound ``s_lo``. + upper: + Upper bound ``s_hi``. + stiffness: + Penalty prefactor ``k``. Default ``10.0``. + name: + Unique bias identifier. Default ``"flat_bottom"``. + exponent: + Power ``p``. Default ``2.0``. + compute_stress: + Passed through to :class:`ConservativeBias`. + + Raises + ------ + ValueError + If any ``lower`` bound is not strictly below its ``upper`` bound. + """ + + def __init__( + self, + cv: Callable[[Batch], Tensor], + lower: Tensor | float, + upper: Tensor | float, + stiffness: Tensor | float = 10.0, + *, + name: str = "flat_bottom", + exponent: float = 2.0, + compute_stress: bool = True, + ) -> None: + lower_t = torch.as_tensor(lower, dtype=torch.get_default_dtype()).reshape(-1) + upper_t = torch.as_tensor(upper, dtype=torch.get_default_dtype()).reshape(-1) + if bool((lower_t >= upper_t).any()): + raise ValueError( + f"FlatBottomRestraint: every lower bound must be strictly below " + f"its upper bound, got lower={lower_t.tolist()} and " + f"upper={upper_t.tolist()}." + ) + super().__init__( + cv, + upper_t, + stiffness, + name=name, + exponent=exponent, + compute_stress=compute_stress, + ) + self.register_buffer("lower", lower_t) + + def energy(self, current: Batch) -> Tensor: + """Return the two-sided confinement energy ``[B, 1]`` in eV. + + Parameters + ---------- + current: + Batch with strained positions supplied by + :meth:`ConservativeBias.evaluate`. + + Returns + ------- + Tensor + Shape ``[B, 1]``; exactly zero inside ``[lower, upper]``. + """ + values = self.cv(current) # [B, D] + above = torch.clamp(values - self.threshold, min=0.0) + below = torch.clamp(self.lower - values, min=0.0) + penalty = self.stiffness * (above**self.exponent + below**self.exponent) + return penalty.sum(dim=-1, keepdim=True) / self.exponent + + def _excess(self, values: Tensor) -> Tensor: + """Return the distance above the upper bound. + + Unused — :meth:`energy` is overridden to handle both sides — but + defined so the class is not abstract in spirit. + + Parameters + ---------- + values: + CV values ``[B, D]``. + + Returns + ------- + Tensor + Shape ``[B, D]``. + """ + return values - self.threshold diff --git a/nvalchemi/enhanced_sampling/cv/__init__.py b/nvalchemi/enhanced_sampling/cv/__init__.py new file mode 100644 index 00000000..5a16fd06 --- /dev/null +++ b/nvalchemi/enhanced_sampling/cv/__init__.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Collective-variable functions for enhanced sampling. + +CVs are plain callables — no class hierarchy, no registration. Any +differentiable function ``cv(batch: Batch) -> Tensor[B, D]`` satisfies +the CV interface. + +Available CVs: :func:`pair_distance`, with :func:`pair_displacement` for +methods that work with the CV *gradient* rather than its value. +:func:`periodic_difference` is a helper for comparing CV values that live +on a circle. +""" + +from nvalchemi.enhanced_sampling.cv._periodic import periodic_difference +from nvalchemi.enhanced_sampling.cv.pair_distance import ( + pair_displacement, + pair_distance, +) + +__all__ = ["pair_displacement", "pair_distance", "periodic_difference"] diff --git a/nvalchemi/enhanced_sampling/cv/_periodic.py b/nvalchemi/enhanced_sampling/cv/_periodic.py new file mode 100644 index 00000000..49bb20f8 --- /dev/null +++ b/nvalchemi/enhanced_sampling/cv/_periodic.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Periodic-aware differences between collective-variable values.""" + +from __future__ import annotations + +import torch +from torch import Tensor + +__all__ = ["periodic_difference"] + + +def periodic_difference( + values: Tensor, centers: Tensor, periods: Tensor | None = None +) -> Tensor: + """Return ``values - centers``, wrapped into ``(-period/2, period/2]``. + + A dihedral restrained at ``+3.0 rad`` must not see a configuration at + ``-3.0 rad`` as ``6.0 rad`` away; the true separation is ``0.28 rad`` the + other way round. Taking a raw difference makes the restraint pull the + long way round the circle, which is both wrong and violently + discontinuous at the branch cut. + + Parameters + ---------- + values: + Current CV values, shape ``[B, D]``. + centers: + Reference values, shape ``[B, D]`` (or broadcastable to it). + periods: + Period per CV component, shape ``[D]``. A component whose period is + ``0`` (or non-finite) is treated as non-periodic and its difference + is returned unwrapped. ``None`` means every component is + non-periodic. + + Returns + ------- + Tensor + Wrapped difference, shape ``[B, D]``. + + Notes + ----- + Differentiability + ``round`` has zero gradient almost everywhere, so the wrap + contributes nothing to ``d(delta)/d(values)`` — the derivative is the + same as for an unwrapped difference, which is what a harmonic + restraint needs. The wrap is discontinuous exactly at + ``delta = period/2``, the antipode of the center; that is inherent to + a periodic CV, not an artefact here. + """ + delta = values - centers + if periods is None: + return delta + + periods = periods.to(device=delta.device, dtype=delta.dtype) + # A zero or non-finite period marks a non-periodic component; guard the + # division so those components produce no wrap rather than NaN/Inf. + active = torch.isfinite(periods) & (periods != 0) + safe = torch.where(active, periods, torch.ones_like(periods)) + wrapped = delta - safe * torch.round(delta / safe) + return torch.where(active, wrapped, delta) diff --git a/nvalchemi/enhanced_sampling/cv/pair_distance.py b/nvalchemi/enhanced_sampling/cv/pair_distance.py new file mode 100644 index 00000000..a8c8b759 --- /dev/null +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -0,0 +1,477 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Differentiable pair-distance collective variable. + +:func:`pair_distance` supports: + +* Non-periodic systems (``batch.cell`` is ``None`` or ``batch.pbc`` is all + ``False``). +* Periodic and mixed-periodic systems via the minimum-image convention (MIC) + for **Minkowski-reduced** triclinic cells (see requirement below). + +Scope: Minkowski-reduced MIC, not general triclinic MIC +-------------------------------------------------------- +This is a **reduced-cell MIC implementation**. It is *not* a general +triclinic MIC implementation. The 27-image exhaustive search (offsets in +``{−1, 0, +1}³``) is correct only when the cell satisfies the Minkowski +reduction conditions. For unreduced cells the minimum-image offset can +exceed ±1 in one or more fractional components, and the search silently +returns a longer-than-minimum image. + +True general triclinic MIC (arbitrary unreduced cells, implemented via LLL +lattice reduction or an extended image search with a data-dependent radius) +is **not yet implemented** — its interaction with the strain-based virial +computation in :class:`ConservativeBias` adds non-trivial complexity. + +Minkowski reduction condition +----------------------------- +For every pair of periodic lattice vectors ``(aᵢ, aⱼ)`` with ``i ≠ j``: + +.. math:: + + |\\mathbf{a}_i \\cdot \\mathbf{a}_j| \\le + \\tfrac{1}{2}\\,\\min(|\\mathbf{a}_i|^2,\\,|\\mathbf{a}_j|^2) + +When this fails, the search returns the wrong image. Counter-example: +cell ``[[1,0,0],[10,0.1,0],[0,0,10]]``, fractional displacement +``[0,0.49,0]`` — the 27-image search returns ≈ 3.9 Å, but the true image +(offset ``[−5,0,0]``) is ≈ 0.11 Å. + +:func:`pair_distance` checks this condition at call time **in eager mode +only** and raises ``ValueError`` for non-reduced cells. + +.. warning:: + + Under ``torch.compile`` the check is skipped (guarded by + ``torch.compiler.is_compiling()``). In compiled mode the caller is + **solely responsible** for supplying Minkowski-reduced cells. Passing + an unreduced cell in compiled mode produces wrong distances with no + error. Pre-reduce cells with a Niggli or LLL algorithm (e.g. + ``ASE: atoms.get_cell().niggli_reduce()``) before simulation. + +Triclinic MIC algorithm +----------------------- +For a reduced cell with lattice matrix ``A`` (rows = lattice vectors, +ASE convention):: + + df = (r_j − r_i) @ A⁻¹ # fractional displacement + df_rounded = df − round(df) × pbc_mask # map to (−0.5, 0.5] + candidates = df_rounded + n, n ∈ {−1,0,+1}³ × pbc_mask + dr_mic = argmin_n |candidates @ A| # shortest image + +torch.compile compatibility +--------------------------- +Shape-based branches (periodicity, cell presence) resolve at trace time. +Gradient flow through ``pair_distance`` for use inside +:class:`ConservativeBias` is fully supported. The Minkowski check and +bounds check are guarded by ``torch.compiler.is_compiling()`` and do not +appear in the compiled graph. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch import Tensor + +if TYPE_CHECKING: + from nvalchemi.data import Batch + +__all__ = ["pair_displacement", "pair_distance"] + + +def pair_displacement(batch: Batch, atom_indices: Tensor) -> Tensor: + """Differentiable pair distance(s) as a collective variable. + + Parameters + ---------- + batch: + Current ``Batch`` containing atomic positions and (optionally) + cell and PBC flags. When a periodic cell is present it must be + **Minkowski-reduced** — see module docstring for details. + atom_indices: + * Shape ``[2]`` — selects the same atom pair ``(i, j)`` in every + graph of the batch. + * Shape ``[B, 2]`` — selects a different pair per graph. + + Indices are **local to each graph** (0-based within the graph, not + global row indices in the batched position tensor). + + May live on any device; it is moved to the batch's device, so a CV + closure built before the batch reaches the GPU still works. + + Returns + ------- + Tensor + Shape ``[B, 3]`` — displacement ``r_j - r_i`` in the same length + unit as ``batch.positions`` (Å). Fully differentiable w.r.t. + ``batch.positions`` and ``batch.cell``. + + Raises + ------ + ValueError + If ``atom_indices`` is not shape ``[2]`` or ``[B, 2]``, or is not + an integer dtype (eager mode only). + IndexError + If any local atom index is negative or >= the graph's atom count + (eager mode only). + ValueError + If any periodic cell is not Minkowski-reduced (eager mode only). + This check is **skipped under** ``torch.compile``; see module + docstring for the compiled-mode caller responsibility. + + Notes + ----- + Non-periodic graphs with an explicit cell + In **eager mode**, MIC is skipped entirely when ``batch.pbc`` is + all-False, so a degenerate cell (e.g. zeros) is safe. In + **compiled mode**, ``bool(pbc.any())`` cannot be evaluated without + a graph break, so MIC is entered whenever ``cell`` and ``pbc`` are + both present. Compiled callers must therefore supply a + non-degenerate cell (or omit it entirely, ``cell=None``) for + non-periodic graphs. + """ + positions = batch.positions # [N_total, 3] + batch_ptr = batch.batch_ptr # [B+1] + B = batch.num_graphs + + # --- Eager-only shape / dtype check on atom_indices ------------------ + # Must come BEFORE the dim()==1 broadcast so wrong shapes are caught, + # not silently coerced. E.g. [1] would expand to [[0,0]] (self-pair) + # and [B,3] would silently drop the third column. + if not torch.compiler.is_compiling(): + _validate_atom_indices(atom_indices, B) + + # --- Follow the batch's device --------------------------------------- + # atom_indices names atoms; it is configuration, not data, so it is + # normally built once — before the batch is moved to GPU — and closed over + # by a CV lambda. Combining it with batch_ptr would then raise a bare + # "expected all tensors to be on the same device" naming neither the CV + # nor the fix. A no-op when it is already on the right device, and the + # same treatment ConservativeBias gives its buffers. + atom_indices = atom_indices.to(device=positions.device) + + # --- Resolve atom_indices to global row indices ----------------------- + if atom_indices.dim() == 1: + atom_indices = atom_indices.unsqueeze(0).expand(B, 2) # [B, 2] + + # --- Eager-only bounds validation ------------------------------------ + if not torch.compiler.is_compiling(): + # Bounds check: catch silent cross-graph wrapping before any indexing. + atoms_per_graph = batch_ptr[1:] - batch_ptr[:-1] # [B] + for col, label in ((0, "atom_indices[…, 0]"), (1, "atom_indices[…, 1]")): + idx = atom_indices[:, col] + neg = idx < 0 + if neg.any(): + bad = neg.nonzero(as_tuple=False).squeeze(-1).tolist() + raise IndexError( + f"pair_distance: {label} has negative values for graph(s) " + f"{bad}: {idx[neg].tolist()}" + ) + oob = idx >= atoms_per_graph + if oob.any(): + bad = oob.nonzero(as_tuple=False).squeeze(-1).tolist() + raise IndexError( + f"pair_distance: {label} is out of range for graph(s) " + f"{bad} — index {idx[oob].tolist()} >= " + f"graph size {atoms_per_graph[oob].tolist()}" + ) + + offsets = batch_ptr[:-1] # [B] + global_i = offsets + atom_indices[:, 0] # [B] + global_j = offsets + atom_indices[:, 1] # [B] + + pos_i = positions[global_i] # [B, 3] + pos_j = positions[global_j] # [B, 3] + dr = pos_j - pos_i # [B, 3], raw Cartesian displacement + + # --- Apply MIC for periodic systems ---------------------------------- + has_cell = getattr(batch, "cell", None) is not None and batch.cell is not None + has_pbc = getattr(batch, "pbc", None) is not None and batch.pbc is not None + + # In eager mode, also require at least one True pbc flag before calling + # _apply_mic. Without this guard, a batch with cell= and + # pbc=all-False would reach torch.linalg.inv and raise LinAlgError. + # + # In compiled mode, bool(pbc.any()) would force a data-dependent Python + # branch that breaks fullgraph=True. We skip the guard there and rely + # on pbc_mask (all-zeros for all-False pbc) to make the MIC computation + # a mathematical identity for non-periodic graphs. Compiled callers + # must therefore supply a non-degenerate cell (or cell=None) for + # non-periodic graphs; a degenerate cell still causes LinAlgError. + any_periodic = ( + has_cell + and has_pbc + and (torch.compiler.is_compiling() or bool(batch.pbc.any())) + ) + + if any_periodic: + if not torch.compiler.is_compiling(): + _check_minkowski_reduced(batch.cell, batch.pbc) + dr = _apply_mic(dr, batch.cell, batch.pbc) + + return dr # [B, 3] + + +def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: + """Return the displacement between two atoms per graph, shape ``[B, 3]``. + + The vector ``r_j - r_i``, minimum-image corrected when the batch is + periodic. :func:`pair_distance` is its norm; the vector form is what a + method needs when it works with the CV *gradient* rather than the value + — adaptive biasing force projects atomic forces onto it. + + The norm of :func:`pair_displacement`, which carries the validation, + device handling, and minimum-image convention documented there. + + Parameters + ---------- + batch: + Current ``Batch``; see :func:`pair_displacement`. + atom_indices: + Atom pair, shape ``[2]`` or ``[B, 2]``; see + :func:`pair_displacement`. + + Returns + ------- + Tensor + Pair distances in angstrom, shape ``[B, 1]``, differentiable with + respect to ``batch.positions`` and ``batch.cell``. + + Raises + ------ + IndexError + If an atom index is out of range for its graph. + ValueError + If *atom_indices* has the wrong shape or dtype, or if a periodic + cell is not Minkowski-reduced (eager mode only). + + Examples + -------- + >>> import torch + >>> from nvalchemi.enhanced_sampling import pair_distance + >>> from nvalchemi.data import AtomicData, Batch + >>> data = AtomicData( + ... positions=torch.tensor([[0.0, 0.0, 0.0], [3.0, 4.0, 0.0]]), + ... atomic_numbers=torch.ones(2, dtype=torch.long), + ... ) + >>> pair_distance(Batch.from_data_list([data]), torch.tensor([0, 1])) + tensor([[5.]]) + """ + dr = pair_displacement(batch, atom_indices) + return torch.linalg.vector_norm(dr, dim=-1, keepdim=True) # [B, 1] + + +# --------------------------------------------------------------------------- +# atom_indices validation +# --------------------------------------------------------------------------- + +_INTEGER_DTYPES = frozenset( + { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + } +) + + +def _validate_atom_indices(atom_indices: Tensor, B: int) -> None: + """Raise ``ValueError`` for malformed ``atom_indices`` (eager mode only). + + Accepted shapes + --------------- + * ``[2]`` — shared pair; broadcast to every graph. + * ``[B, 2]`` — one pair per graph. + + Rejected (with clear error messages) + ------------------------------------- + * Wrong number of dimensions (not 1-D or 2-D). + * 1-D tensor whose length is not exactly 2. A length-1 tensor such as + ``torch.tensor([0])`` would otherwise silently expand to ``[[0, 0]]`` + (a self-distance), not raise. + * 2-D tensor whose second dimension is not exactly 2. A ``[B, 3]`` + tensor would otherwise silently drop the third column. + * 2-D tensor whose first dimension does not match the batch size ``B``. + * Non-integer dtype. Float indices would silently be used as memory + offsets after casting by the indexing operation. + + Parameters + ---------- + atom_indices: + The tensor to validate. + B: + Number of graphs in the current batch. + """ + # dtype check + if atom_indices.dtype not in _INTEGER_DTYPES: + raise ValueError( + f"pair_distance: atom_indices must have an integer dtype, " + f"got {atom_indices.dtype}. Use e.g. torch.tensor([i, j]) " + f"(default int64) or pass dtype=torch.long explicitly." + ) + + ndim = atom_indices.dim() + shape = tuple(atom_indices.shape) + + if ndim == 1: + if shape[0] != 2: + raise ValueError( + f"pair_distance: 1-D atom_indices must have exactly 2 elements " + f"(shape [2] for a shared pair), got shape {shape}. " + f"A length-1 tensor would silently produce a self-distance." + ) + elif ndim == 2: + if shape[1] != 2: + raise ValueError( + f"pair_distance: 2-D atom_indices must have shape [B, 2], " + f"got {shape}. The second dimension must be exactly 2 " + f"(atom i and atom j); extra columns are not allowed." + ) + if shape[0] != B: + raise ValueError( + f"pair_distance: 2-D atom_indices has shape {shape} but the " + f"batch has B={B} graphs. The first dimension must equal B." + ) + else: + raise ValueError( + f"pair_distance: atom_indices must be 1-D (shape [2]) or " + f"2-D (shape [B, 2]), got {ndim}-D tensor with shape {shape}." + ) + + +# --------------------------------------------------------------------------- +# Minkowski-reduction check +# --------------------------------------------------------------------------- + + +def _check_minkowski_reduced(cell: Tensor, pbc: Tensor) -> None: + """Raise ``ValueError`` if any periodic cell pair violates the Minkowski condition. + + This check is an **eager-mode guard only**. It is never called under + ``torch.compile`` (guarded by ``torch.compiler.is_compiling()`` in the + caller). Compiled callers are responsible for supplying reduced cells; + no error is raised if a non-reduced cell is used in compiled mode. + + The 27-image MIC search returns the true minimum-image vector only for + Minkowski-reduced cells. For every pair of periodic lattice vectors + ``(aᵢ, aⱼ)`` with ``i ≠ j``: + + .. math:: + + |\\mathbf{a}_i \\cdot \\mathbf{a}_j| + \\le \\tfrac{1}{2}\\,\\min(|\\mathbf{a}_i|^2,\\,|\\mathbf{a}_j|^2) + + When this fails, the minimum-image offset can exceed ±1 in some + fractional component and the search silently returns the wrong image. + + Parameters + ---------- + cell: + Lattice matrices, shape ``[B, 3, 3]`` or ``[B, 1, 3, 3]``. + pbc: + Periodicity flags, shape ``[B, 3]`` or ``[B, 1, 3]``. + """ + if cell.dim() == 4: + cell = cell.squeeze(1) + if pbc.dim() == 3: + pbc = pbc.squeeze(1) + + for i in range(3): + for j in range(i + 1, 3): + # Only enforce for pairs of dimensions that are BOTH periodic. + both_periodic = pbc[:, i] & pbc[:, j] # [B] bool + if not both_periodic.any(): + continue + + ai = cell[:, i, :] # [B, 3] + aj = cell[:, j, :] # [B, 3] + dot_abs = (ai * aj).sum(-1).abs() # [B] + norm_sq_i = (ai * ai).sum(-1) # [B] + norm_sq_j = (aj * aj).sum(-1) # [B] + threshold = 0.5 * torch.minimum(norm_sq_i, norm_sq_j) # [B] + + violated = both_periodic & (dot_abs > threshold) + if violated.any(): + bad = violated.nonzero(as_tuple=False).squeeze(-1).tolist() + raise ValueError( + f"pair_distance: the cell for graph(s) {bad} is not " + f"Minkowski-reduced: lattice vectors a[{i}] and a[{j}] satisfy " + f"|a[{i}]·a[{j}]| > 0.5·min(|a[{i}]|², |a[{j}]|²). " + f"The 27-image MIC search is only guaranteed correct for " + f"Minkowski-reduced cells. Pre-reduce the cell using a Niggli " + f"or LLL algorithm (e.g. ASE niggli_reduce) before simulation." + ) + + +# --------------------------------------------------------------------------- +# MIC implementation +# --------------------------------------------------------------------------- + + +def _apply_mic(dr: Tensor, cell: Tensor, pbc: Tensor) -> Tensor: + """Apply the minimum-image convention via an exhaustive 27-image search. + + Parameters + ---------- + dr: + Cartesian displacement vectors, shape ``[B, 3]``. + cell: + Lattice matrices, shape ``[B, 3, 3]`` or ``[B, 1, 3, 3]``. + Rows are lattice vectors (ASE convention). + pbc: + Periodicity flags per dimension, shape ``[B, 3]`` or ``[B, 1, 3]``. + + Returns + ------- + Tensor + MIC-corrected displacement vectors, shape ``[B, 3]``. + + Notes + ----- + The cell must be Minkowski-reduced; see :func:`_check_minkowski_reduced`. + That check is performed in :func:`pair_distance` before this function is + called, so it is not repeated here. + """ + if cell.dim() == 4: + cell = cell.squeeze(1) + if pbc.dim() == 3: + pbc = pbc.squeeze(1) + + pbc_mask = pbc.to(dtype=cell.dtype) # [B, 3] + + # Fractional displacement + cell_inv = torch.linalg.inv(cell) # [B, 3, 3] + df = torch.bmm(dr.unsqueeze(1), cell_inv).squeeze(1) # [B, 3] + + # Initial half-cell rounding (periodic dims only) + df_rounded = df - torch.round(df) * pbc_mask # [B, 3] + + # Exhaustive 27-image search over offsets in {-1, 0, +1}³ + coords = torch.tensor([-1.0, 0.0, 1.0], device=dr.device, dtype=dr.dtype) + gi, gj, gk = torch.meshgrid(coords, coords, coords, indexing="ij") + all_offsets = torch.stack( + [gi.flatten(), gj.flatten(), gk.flatten()], dim=-1 + ) # [27, 3] + + offsets_masked = all_offsets[None] * pbc_mask[:, None, :] # [B, 27, 3] + df_cands = df_rounded[:, None, :] + offsets_masked # [B, 27, 3] + dr_cands = torch.einsum("bki,bij->bkj", df_cands, cell) # [B, 27, 3] + + dist_sq = (dr_cands * dr_cands).sum(dim=-1) # [B, 27] + best = dist_sq.argmin(dim=-1)[:, None, None].expand(-1, 1, 3) # [B, 1, 3] + return dr_cands.gather(1, best).squeeze(1) # [B, 3] diff --git a/nvalchemi/hooks/bias.py b/nvalchemi/hooks/bias.py index 7bd90f09..6f488d8e 100644 --- a/nvalchemi/hooks/bias.py +++ b/nvalchemi/hooks/bias.py @@ -17,10 +17,20 @@ Provides :class:`BiasedPotentialHook`, which adds external bias potentials to the forces and energy computed by the ML model. + +.. deprecated:: + + :class:`BiasedPotentialHook` is superseded by + :mod:`nvalchemi.enhanced_sampling`, whose ``EnhancedSampling`` runner + covers everything this hook does and carries a cell response it cannot. + See that module's docstring for which to use when. This hook still + works and no removal date is set; see :class:`BiasedPotentialHook` for + the specific limitations that motivated the replacement. """ from __future__ import annotations +import warnings from enum import Enum from typing import TYPE_CHECKING @@ -40,6 +50,35 @@ class BiasedPotentialHook: """Add an external bias potential to forces and energy after the forward pass. + .. deprecated:: + + Superseded by :mod:`nvalchemi.enhanced_sampling`. Constructing this + hook emits a :class:`DeprecationWarning`. It remains functional so + existing code keeps working, and no removal date is set — but new + biases should be written against + :class:`~nvalchemi.enhanced_sampling.ConservativeBias` and run + through :class:`~nvalchemi.enhanced_sampling.EnhancedSampling`, which + together cover everything this hook does. + + Three limitations of the ``bias_fn`` contract motivated the + replacement: + + * **No cell response.** ``bias_fn`` returns only + ``(energy, forces)``, so a bias contributes no stress or virial. + Under NPT/NPH the barostat reads ``batch.stress``, which the bias + never touches, so the cell evolves as if the bias were absent — + silently, with no error. + * **Forces are hand-written.** Nothing checks that ``bias_forces`` + is ``-dE/dr`` for the returned ``bias_energy``, so a bias can be + non-conservative by accident. + :class:`~nvalchemi.enhanced_sampling.ConservativeBias` derives + forces and stress from one energy definition by autograd. + * **Sequential in-place composition.** Each hook mutates + ``batch.forces`` in turn, so a second bias that reads + ``batch.forces`` observes the first one's contribution. + :func:`~nvalchemi.enhanced_sampling.aggregate_bias_results` sums + every bias against the same unmodified model output instead. + This hook enables enhanced sampling techniques by composing an arbitrary bias potential on top of the ML potential **without** modifying the model itself. The bias is applied in-place to @@ -140,6 +179,17 @@ def __init__( frequency: int = 1, inplace: bool = True, ) -> None: + warnings.warn( + "BiasedPotentialHook is deprecated in favour of " + "nvalchemi.enhanced_sampling (BiasPotential / ConservativeBias), " + "which derives forces and stress from a single energy definition. " + "bias_fn returns only (energy, forces), so a bias applied through " + "this hook contributes no stress and is invisible to the NPT/NPH " + "barostat. It remains functional and no removal date is set; run " + "new biases through EnhancedSampling instead.", + DeprecationWarning, + stacklevel=2, + ) self.bias_fn = bias_fn self.stage = stage self.frequency = frequency diff --git a/nvalchemi/models/_utils.py b/nvalchemi/models/_utils.py index c1083a38..720f160d 100644 --- a/nvalchemi/models/_utils.py +++ b/nvalchemi/models/_utils.py @@ -95,6 +95,7 @@ def autograd_forces( positions: NodePositions, training: bool = False, retain_graph: bool = False, + allow_unused: bool = False, ) -> Forces: """Compute forces as ``-dE/dr`` via autograd. @@ -112,6 +113,12 @@ def autograd_forces( If ``True``, the computation graph is retained after the backward pass. Needed when subsequent autograd calls traverse shared graph nodes. + allow_unused : bool, optional + If ``True``, an energy that does not depend on *positions* yields + zero forces instead of raising ``RuntimeError``. Use this for + energy terms that are legitimately position-independent, such as a + pure cell/volume term. Defaults to ``False``, which surfaces a + missing dependency as an error. Returns ------- @@ -125,6 +132,8 @@ def autograd_forces( grad_outputs=torch.ones_like(energy), create_graph=training, retain_graph=effective_retain, + allow_unused=allow_unused, + materialize_grads=allow_unused, )[0] @@ -246,6 +255,7 @@ def autograd_stresses( num_graphs: int, training: bool = False, retain_graph: bool = False, + allow_unused: bool = False, ) -> Stress: r"""Compute tensile-positive Cauchy stress via autograd. @@ -265,6 +275,10 @@ def autograd_stresses( If ``True``, create the computation graph for higher-order gradients. retain_graph : bool, optional If ``True``, retain the computation graph. + allow_unused : bool, optional + If ``True``, an energy that does not depend on *displacement* yields + zero stress instead of raising ``RuntimeError``. Defaults to + ``False``, which surfaces a missing dependency as an error. Returns ------- @@ -278,6 +292,8 @@ def autograd_stresses( grad_outputs=torch.ones_like(energy), create_graph=training, retain_graph=effective_retain, + allow_unused=allow_unused, + materialize_grads=allow_unused, )[0] volume = torch.det(cell).abs().view(-1, 1, 1) return grad.view(num_graphs, 3, 3) / volume @@ -291,6 +307,7 @@ def autograd_forces_and_stresses( num_graphs: int, training: bool = False, retain_graph: bool = False, + allow_unused: bool = False, ) -> tuple[Forces, Stress]: """Compute forces and tensile-positive Cauchy stress in one autograd call. @@ -310,6 +327,13 @@ def autograd_forces_and_stresses( If ``True``, create the computation graph for higher-order gradients. retain_graph : bool, optional If ``True``, retain the computation graph. + allow_unused : bool, optional + If ``True``, whichever of *positions* and *displacement* the energy + does not depend on yields a zero gradient instead of raising + ``RuntimeError``. Use this for energy terms that are legitimately + independent of one of the two, such as a pure cell/volume term that + produces stress but no forces. Defaults to ``False``, which + surfaces a missing dependency as an error. Returns ------- @@ -323,6 +347,8 @@ def autograd_forces_and_stresses( grad_outputs=torch.ones_like(energy), create_graph=training, retain_graph=effective_retain, + allow_unused=allow_unused, + materialize_grads=allow_unused, ) forces = -position_grad volume = torch.det(cell).abs().view(-1, 1, 1) diff --git a/test/dynamics/test_bias_hook.py b/test/dynamics/test_bias_hook.py index 77fc64f3..3712f6d5 100644 --- a/test/dynamics/test_bias_hook.py +++ b/test/dynamics/test_bias_hook.py @@ -15,6 +15,11 @@ """Unit tests for ``nvalchemi.hooks.bias`` — Tier 1 bias hook. Covers :class:`BiasedPotentialHook`. + +``BiasedPotentialHook`` is deprecated in favour of +:mod:`nvalchemi.enhanced_sampling` but remains functional, so these tests +still assert its behaviour. The construction warning is silenced module-wide +and asserted explicitly in :class:`TestBiasedPotentialHookDeprecation`. """ from __future__ import annotations @@ -28,6 +33,10 @@ from nvalchemi.models.demo import DemoModel, DemoModelWrapper from test.dynamics.conftest import make_dynamics_context +pytestmark = pytest.mark.filterwarnings( + "ignore:BiasedPotentialHook is deprecated:DeprecationWarning" +) + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -226,6 +235,44 @@ def small_bias(b): nan_hook(ctx, DynamicsStage.AFTER_COMPUTE) # should not raise +class TestBiasedPotentialHookDeprecation: + """BiasedPotentialHook is deprecated but must remain functional.""" + + @pytest.mark.filterwarnings("default::DeprecationWarning") + def test_construction_warns(self) -> None: + with pytest.warns( + DeprecationWarning, match="BiasedPotentialHook is deprecated" + ): + BiasedPotentialHook(bias_fn=lambda b: (b.energy, b.forces)) + + @pytest.mark.filterwarnings("default::DeprecationWarning") + def test_warning_points_at_enhanced_sampling(self) -> None: + """The message must name the replacement, not just say 'deprecated'.""" + with pytest.warns(DeprecationWarning) as record: + BiasedPotentialHook(bias_fn=lambda b: (b.energy, b.forces)) + message = str(record[0].message) + assert "nvalchemi.enhanced_sampling" in message + # The substantive reason to migrate, not just a pointer. + assert "stress" in message + + def test_still_applies_bias_after_deprecation(self, device: str) -> None: + """Deprecated does not mean broken: the hook must still work.""" + batch = _make_batch(device=device) + dynamics = _make_dynamics() + forces_before = batch.forces.clone() + energies_before = batch.energy.clone() + + bias_e = torch.ones_like(batch.energy) * 0.25 + bias_f = torch.ones_like(batch.forces) * 0.75 + hook = BiasedPotentialHook( + bias_fn=lambda b: (bias_e, bias_f), stage=DynamicsStage.AFTER_COMPUTE + ) + hook(_make_ctx(batch, dynamics), DynamicsStage.AFTER_COMPUTE) + + assert torch.allclose(batch.forces, forces_before + 0.75) + assert torch.allclose(batch.energy, energies_before + 0.25) + + class TestBiasedPotentialHookCompile: """Verify BiasedPotentialHook works under torch.compile.""" diff --git a/test/enhanced_sampling/__init__.py b/test/enhanced_sampling/__init__.py new file mode 100644 index 00000000..46707983 --- /dev/null +++ b/test/enhanced_sampling/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/test/enhanced_sampling/test_abf.py b/test/enhanced_sampling/test_abf.py new file mode 100644 index 00000000..611f03ef --- /dev/null +++ b/test/enhanced_sampling/test_abf.py @@ -0,0 +1,927 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for :class:`AdaptiveBiasingForce`. + +The estimator is checked against a system whose potential of mean force is +known in closed form — a harmonic pair in three dimensions, where + +``dA/dr = k (r - r0) - 2 kB T / r`` + +holds *pointwise*, so a single sample must reproduce it exactly rather than +only on average. That makes the metric correction directly testable: drop +it and the ideal-gas limit reports a flat PMF instead of ``-2 kB T ln r``. +""" + +from __future__ import annotations + +import pytest +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.dynamics.base import DynamicsStage +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling import ( + AdaptiveBiasingForce, + BiasResult, + EnhancedSampling, + ReplicaExchange, + ThermodynamicState, + pair_displacement, + pair_distance, +) +from nvalchemi.models.demo import DemoModel, DemoModelWrapper + +TEMPERATURE = 300.0 +KT = KB_EV * TEMPERATURE +SPRING = 3.0 +REST_LENGTH = 2.0 + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _harmonic_frame( + distances: list[float], device: str = "cpu", forces: bool = True +) -> Batch: + """Return pairs at *distances* carrying the harmonic physical force. + + ``U = 0.5 k (r - r0)^2`` gives ``F_j = -k (r - r0) u`` and + ``F_i = +k (r - r0) u``, with the pair laid out along ``x``. + """ + items = [] + for r in distances: + kwargs = { + "positions": torch.tensor([[0.0, 0.0, 0.0], [r, 0.0, 0.0]]), + "atomic_numbers": torch.ones(2, dtype=torch.long), + } + if forces: + gradient = SPRING * (r - REST_LENGTH) + kwargs["forces"] = torch.tensor( + [[gradient, 0.0, 0.0], [-gradient, 0.0, 0.0]] + ) + items.append(AtomicData(**kwargs)) + return Batch.from_data_list(items).to(device) + + +def _abf(device: str = "cpu", **kwargs) -> AdaptiveBiasingForce: + """Return an ABF bias with the threshold disabled unless overridden.""" + params = { + "atom_indices": torch.tensor([0, 1]), + "temperature": TEMPERATURE, + "cv_range": (1.0, 4.0), + "n_bins": 60, + "min_samples": 0, + "full_samples": 0, + } + params.update(kwargs) + return AdaptiveBiasingForce(**params).to(device) + + +def _analytic_gradient(r: float) -> float: + """Return the exact ``dA/dr`` for the harmonic pair.""" + return SPRING * (r - REST_LENGTH) - 2.0 * KT / r + + +def _runner_batch( + n_graphs: int = 2, atoms: int = 4, device: str = "cpu", seed: int = 0 +) -> Batch: + """Return a batch with the buffers dynamics writes back into.""" + torch.manual_seed(seed) + items = [] + for _ in range(n_graphs): + data = AtomicData( + positions=torch.randn(atoms, 3), + atomic_numbers=torch.full((atoms,), 6, dtype=torch.long), + atomic_masses=torch.ones(atoms), + forces=torch.zeros(atoms, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(atoms, 3)) + items.append(data) + return Batch.from_data_list(items).to(device) + + +def _make_dynamics(device: str = "cpu") -> NVTLangevin: + """Return a demo-model Langevin integrator.""" + model = DemoModelWrapper(DemoModel()).to(device) + return NVTLangevin(model=model, dt=0.1, temperature=TEMPERATURE, friction=0.1) + + +# =========================================================================== +# 1. Construction +# =========================================================================== + + +class TestConstruction: + """Constructor validation.""" + + @pytest.mark.parametrize("indices", [[0], [0, 1, 2]]) + def test_wrong_pair_size_raises(self, indices: list[int]) -> None: + with pytest.raises(ValueError, match="exactly two atoms"): + _abf(atom_indices=torch.tensor(indices)) + + def test_wrong_rank_raises(self) -> None: + with pytest.raises(ValueError, match=r"shape \[2\] or \[B, 2\]"): + _abf(atom_indices=torch.zeros(2, 2, 2, dtype=torch.long)) + + def test_negative_index_raises(self) -> None: + with pytest.raises(ValueError, match="must be non-negative"): + _abf(atom_indices=torch.tensor([0, -1])) + + def test_self_pair_raises(self) -> None: + """A zero-length CV has no direction to project onto.""" + with pytest.raises(ValueError, match="same atom twice"): + _abf(atom_indices=torch.tensor([2, 2])) + + @pytest.mark.parametrize("temperature", [0.0, -10.0]) + def test_non_positive_temperature_raises(self, temperature: float) -> None: + """Temperature scales the metric correction, so it must be real.""" + with pytest.raises(ValueError, match="temperature must be positive"): + _abf(temperature=temperature) + + @pytest.mark.parametrize("cv_range", [(4.0, 1.0), (2.0, 2.0)]) + def test_inverted_range_raises(self, cv_range: tuple[float, float]) -> None: + with pytest.raises(ValueError, match="cv_range must be increasing"): + _abf(cv_range=cv_range) + + def test_negative_lower_bound_raises(self) -> None: + with pytest.raises(ValueError, match="lower bound must be non-negative"): + _abf(cv_range=(-1.0, 3.0)) + + def test_zero_bins_raises(self) -> None: + with pytest.raises(ValueError, match="n_bins must be at least 1"): + _abf(n_bins=0) + + def test_full_below_min_samples_raises(self) -> None: + with pytest.raises(ValueError, match="must be at least min_samples"): + _abf(min_samples=100, full_samples=50) + + def test_full_samples_defaults_to_double(self) -> None: + assert _abf(min_samples=50, full_samples=None).full_samples == 100 + + def test_non_positive_max_force_raises(self) -> None: + with pytest.raises(ValueError, match="max_force must be positive"): + _abf(max_force=0.0) + + def test_zero_update_frequency_raises(self) -> None: + with pytest.raises(ValueError, match="update_frequency must be at least 1"): + _abf(update_frequency=0) + + def test_mixin_order_is_correct(self) -> None: + """AdaptivePotentialMixin must precede nn.Module in the MRO.""" + from nvalchemi.enhanced_sampling import AdaptivePotentialMixin + + mro = AdaptiveBiasingForce.__mro__ + assert mro.index(AdaptivePotentialMixin) < mro.index(torch.nn.Module) + + def test_is_not_a_conservative_bias(self) -> None: + """ABF has no energy to differentiate, so it is not a model.""" + from nvalchemi.enhanced_sampling import ConservativeBias + + assert not issubclass(AdaptiveBiasingForce, ConservativeBias) + + def test_satisfies_the_bias_protocol(self) -> None: + from nvalchemi.enhanced_sampling import BiasPotential + + assert isinstance(_abf(), BiasPotential) + + +# =========================================================================== +# 2. The estimator, against a closed-form PMF +# =========================================================================== + + +class TestEstimator: + """``dA/dr = k (r - r0) - 2 kB T / r``, exactly, for a harmonic pair.""" + + @pytest.mark.parametrize("distance", [1.6, 2.0, 2.9, 3.5]) + def test_single_sample_is_exact(self, device: str, distance: float) -> None: + """The harmonic pair makes the estimator exact pointwise, not just + on average — one sample must land on the analytic value.""" + bias = _abf(device) + bias.update(_harmonic_frame([distance], device), BiasResult()) + + index = int(bias.bin_index(torch.tensor([distance]))[0]) + assert float(bias.mean_force()[index]) == pytest.approx( + _analytic_gradient(distance), rel=1e-5 + ) + + def test_metric_correction_is_applied(self, device: str) -> None: + """Two non-interacting particles have PMF ``-2 kB T ln r``. + + A naive Cartesian projection reports zero mean force here — a flat + PMF — which is not noise but a smoothly wrong answer. + """ + bias = _abf(device) + frame = _harmonic_frame([2.5], device) + frame.forces = torch.zeros_like(frame.forces) + bias.update(frame, BiasResult()) + + index = int(bias.bin_index(torch.tensor([2.5]))[0]) + got = float(bias.mean_force()[index]) + assert got == pytest.approx(-2.0 * KT / 2.5, rel=1e-5) + assert got != pytest.approx(0.0, abs=1e-6) + + def test_metric_correction_scales_with_temperature(self, device: str) -> None: + """The Jacobian term is ``2 kB T / r``, linear in T.""" + gradients = [] + for temperature in (300.0, 600.0): + bias = _abf(device, temperature=temperature) + frame = _harmonic_frame([2.5], device) + frame.forces = torch.zeros_like(frame.forces) + bias.update(frame, BiasResult()) + index = int(bias.bin_index(torch.tensor([2.5]))[0]) + gradients.append(float(bias.mean_force()[index])) + assert gradients[1] == pytest.approx(2.0 * gradients[0], rel=1e-5) + + def test_samples_average_within_a_bin(self, device: str) -> None: + bias = _abf(device, n_bins=1, cv_range=(1.0, 4.0)) + bias.update(_harmonic_frame([1.5, 2.5, 3.5], device), BiasResult()) + + expected = sum(_analytic_gradient(r) for r in (1.5, 2.5, 3.5)) / 3.0 + assert int(bias.bin_counts[0]) == 3 + assert float(bias.mean_force()[0]) == pytest.approx(expected, rel=1e-5) + + def test_unvisited_bins_are_nan_not_zero(self, device: str) -> None: + """Zero is a plausible mean force, so it cannot mean "no data".""" + bias = _abf(device) + bias.update(_harmonic_frame([2.0], device), BiasResult()) + estimate = bias.mean_force() + assert bool(torch.isnan(estimate).any()) + assert int((~torch.isnan(estimate)).sum()) == 1 + + def test_out_of_range_samples_are_discarded(self, device: str) -> None: + bias = _abf(device, cv_range=(2.0, 3.0), n_bins=10) + bias.update(_harmonic_frame([1.0, 2.5, 5.0], device), BiasResult()) + assert int(bias.bin_counts.sum()) == 1 + + def test_update_without_forces_raises(self, device: str) -> None: + """There is nothing to project if the frame carries no forces.""" + bias = _abf(device) + with pytest.raises(ValueError, match="has no forces"): + bias.update(_harmonic_frame([2.0], device, forces=False), BiasResult()) + + +# =========================================================================== +# 3. The applied force +# =========================================================================== + + +class TestAppliedForce: + """Direction, magnitude, and the sample threshold.""" + + def test_result_is_force_only(self, device: str) -> None: + """No energy: the applied force is not the gradient of anything held.""" + bias = _abf(device) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + result = bias.evaluate(_harmonic_frame([2.5], device)) + + assert result.energy is None + assert result.stress is None + assert result.virial is None + assert result.forces is not None + + def test_force_opposes_the_mean_force(self, device: str) -> None: + """The bias must cancel the drift, not reinforce it. + + For an ideal pair ``dA/dr = -2 kB T / r < 0``, so the free energy + falls with separation and the entropic drift is outward; the bias + must therefore pull inward. + """ + bias = _abf(device) + frame = _harmonic_frame([2.5], device) + frame.forces = torch.zeros_like(frame.forces) + bias.update(frame, BiasResult()) + + forces = bias.evaluate(_harmonic_frame([2.5], device)).forces + assert float(forces[1, 0]) < 0.0 + assert float(forces[0, 0]) > 0.0 + + def test_applied_force_equals_the_estimate(self, device: str) -> None: + bias = _abf(device) + bias.update(_harmonic_frame([2.9], device), BiasResult()) + forces = bias.evaluate(_harmonic_frame([2.9], device)).forces + assert float(forces[1, 0]) == pytest.approx(_analytic_gradient(2.9), rel=1e-5) + + def test_bias_exerts_no_net_force(self, device: str) -> None: + """Equal and opposite along the pair: no spurious translation.""" + bias = _abf(device) + bias.update(_harmonic_frame([2.5, 3.1], device), BiasResult()) + forces = bias.evaluate(_harmonic_frame([2.5, 3.1], device)).forces + assert float(forces.sum(dim=0).abs().max()) < 1e-9 + + def test_no_force_before_the_threshold(self, device: str) -> None: + """An estimate from a handful of samples is noise.""" + bias = _abf(device, min_samples=4, full_samples=8) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + result = bias.evaluate(_harmonic_frame([2.5], device)) + assert torch.count_nonzero(result.forces) == 0 + + def test_force_ramps_between_the_thresholds(self, device: str) -> None: + """A jump to the full estimate would be the discontinuity the + threshold exists to avoid.""" + bias = _abf(device, min_samples=2, full_samples=6) + magnitudes = [] + for _ in range(6): + bias.update(_harmonic_frame([2.9], device), BiasResult()) + magnitudes.append( + abs(float(bias.evaluate(_harmonic_frame([2.9], device)).forces[1, 0])) + ) + + assert magnitudes[0] == pytest.approx(0.0, abs=1e-12) + assert magnitudes == sorted(magnitudes) + assert magnitudes[-1] == pytest.approx(abs(_analytic_gradient(2.9)), rel=1e-5) + + def test_ramp_fraction_endpoints(self, device: str) -> None: + bias = _abf(device, min_samples=2, full_samples=4, n_bins=1) + assert float(bias.ramp_fraction()[0]) == 0.0 + for _ in range(4): + bias.update(_harmonic_frame([2.5], device), BiasResult()) + assert float(bias.ramp_fraction()[0]) == pytest.approx(1.0) + + def test_equal_thresholds_reach_full_force_at_the_threshold( + self, device: str + ) -> None: + """``full_samples == min_samples`` is a step, not a one-late ramp. + + The linear form divides by a zero span here; clamping that span to 1 + would delay full force by a sample, so the count that finally meets + the threshold would still apply nothing. + """ + bias = _abf(device, n_bins=1, min_samples=3, full_samples=3) + + fractions = [] + for _ in range(5): + bias.update(_harmonic_frame([2.5], device), BiasResult()) + fractions.append(float(bias.ramp_fraction()[0])) + + assert fractions == [0.0, 0.0, 1.0, 1.0, 1.0] + + def test_equal_thresholds_bump_the_version_at_the_threshold( + self, device: str + ) -> None: + """The bump must track the ramp, not a separately derived threshold.""" + bias = _abf(device, n_bins=1, min_samples=3, full_samples=3) + for _ in range(2): + bias.update(_harmonic_frame([2.5], device), BiasResult()) + assert bias.state_version == 0 + + bias.update(_harmonic_frame([2.5], device), BiasResult()) + assert int(bias.bin_counts[0]) == 3 + assert bias.state_version == 1 + + def test_equal_thresholds_apply_force_at_the_threshold(self, device: str) -> None: + """The applied force must follow the same schedule as ramp_fraction.""" + bias = _abf( + device, n_bins=1, cv_range=(1.0, 4.0), min_samples=2, full_samples=2 + ) + bias.update(_harmonic_frame([2.9], device), BiasResult()) + assert ( + torch.count_nonzero(bias.evaluate(_harmonic_frame([2.9], device)).forces) + == 0 + ) + + bias.update(_harmonic_frame([2.9], device), BiasResult()) + forces = bias.evaluate(_harmonic_frame([2.9], device)).forces + assert float(forces[1, 0]) == pytest.approx(_analytic_gradient(2.9), rel=1e-5) + + def test_zero_threshold_default_is_a_step(self, device: str) -> None: + """``min_samples=0`` defaults full_samples to 0, i.e. no threshold.""" + bias = AdaptiveBiasingForce( + atom_indices=torch.tensor([0, 1]), + temperature=TEMPERATURE, + cv_range=(1.0, 4.0), + n_bins=1, + min_samples=0, + ).to(device) + assert bias.full_samples == 0 + + assert float(bias.ramp_fraction()[0]) == 0.0 # no samples, no estimate + bias.update(_harmonic_frame([2.5], device), BiasResult()) + assert float(bias.ramp_fraction()[0]) == 1.0 + + def test_unvisited_bins_report_zero_ramp(self, device: str) -> None: + """A bin with no samples has no estimate whatever the thresholds are.""" + for thresholds in ((0, 0), (3, 3), (2, 6)): + bias = _abf(device, min_samples=thresholds[0], full_samples=thresholds[1]) + assert float(bias.ramp_fraction().max()) == 0.0 + + def test_applied_gradient_matches_ramp_fraction(self, device: str) -> None: + """The two must not drift: the force is the estimate times the ramp.""" + bias = _abf(device, n_bins=1, min_samples=2, full_samples=5) + for _ in range(6): + bias.update(_harmonic_frame([2.9], device), BiasResult()) + expected = float(bias.mean_force()[0]) * float(bias.ramp_fraction()[0]) + forces = bias.evaluate(_harmonic_frame([2.9], device)).forces + assert float(forces[1, 0]) == pytest.approx(expected, rel=1e-5, abs=1e-9) + + def test_out_of_range_walkers_feel_nothing(self, device: str) -> None: + bias = _abf(device, cv_range=(2.0, 3.0), n_bins=10) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + result = bias.evaluate(_harmonic_frame([2.5, 9.0], device)) + assert float(result.forces[2:].abs().max()) == 0.0 + assert float(result.forces[:2].abs().max()) > 0.0 + + def test_max_force_caps_the_estimate(self, device: str) -> None: + """One visit at a bad geometry cannot dominate the trajectory.""" + bias = _abf(device, max_force=0.1) + bias.update(_harmonic_frame([3.5], device), BiasResult()) + forces = bias.evaluate(_harmonic_frame([3.5], device)).forces + assert abs(float(forces[1, 0])) == pytest.approx(0.1, rel=1e-6) + + def test_evaluate_does_not_mutate_state(self, device: str) -> None: + """``evaluate`` is read-only; only ``update`` changes the estimate.""" + bias = _abf(device) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + before = (bias.bin_counts.clone(), bias.force_sum.clone(), bias.state_version) + + for _ in range(3): + bias.evaluate(_harmonic_frame([2.5], device)) + + assert torch.equal(bias.bin_counts, before[0]) + assert torch.equal(bias.force_sum, before[1]) + assert bias.state_version == before[2] + + def test_diagnostics_are_reported(self, device: str) -> None: + bias = _abf(device) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + result = bias.evaluate(_harmonic_frame([2.5, 9.0], device)) + + assert set(result.observables) >= { + "cv", + "bin", + "applied_gradient", + "samples", + "ramp", + "in_range", + } + # Deliberately not called "mean_force": this is the ramped, capped + # value, which mean_force() is not. + assert "mean_force" not in result.observables + assert float(result.observables["cv"][0]) == pytest.approx(2.5, abs=1e-6) + assert result.observables["in_range"].reshape(-1).tolist() == [1.0, 0.0] + + def test_out_of_range_diagnostics_are_zero(self, device: str) -> None: + """A clamped bin index must not leak the edge bin's statistics. + + ``_bin_of`` clamps so the index is always safe to gather with, which + means an out-of-range walker names the nearest edge bin. Reporting + that bin's counts and ramp would show a converged-looking sample + count beside ``in_range == 0``, and any aggregation over + ``bias/abf/samples`` would double-count walkers that contributed + nothing. + """ + bias = _abf(device, cv_range=(2.0, 3.0), n_bins=4) + # Give both edge bins real statistics to leak. + for _ in range(5): + bias.update(_harmonic_frame([2.9], device), BiasResult()) + for _ in range(3): + bias.update(_harmonic_frame([2.1], device), BiasResult()) + assert int(bias.bin_counts.sum()) == 8 + + # Inside, far above the range, far below it. + result = bias.evaluate(_harmonic_frame([2.9, 9.0, 0.5], device)) + observables = result.observables + + assert observables["in_range"].reshape(-1).tolist() == [1.0, 0.0, 0.0] + assert observables["samples"].reshape(-1).tolist() == [5, 0, 0] + assert observables["ramp"].reshape(-1).tolist() == [1.0, 0.0, 0.0] + assert observables["applied_gradient"].reshape(-1)[1:].abs().max() == 0.0 + + def test_reported_bin_matches_bin_index(self, device: str) -> None: + """The observable and the public accessor must not disagree.""" + bias = _abf(device, cv_range=(2.0, 3.0), n_bins=4) + bias.update(_harmonic_frame([2.9], device), BiasResult()) + + probe = _harmonic_frame([2.9, 9.0, 0.5], device) + observables = bias.evaluate(probe).observables + assert ( + observables["bin"].reshape(-1).tolist() + == bias.bin_index(observables["cv"].reshape(-1)).tolist() + ) + assert observables["bin"].reshape(-1).tolist() == [3, -1, -1] + + def test_cv_is_reported_even_out_of_range(self, device: str) -> None: + """The CV is genuinely measured wherever the walker is.""" + bias = _abf(device, cv_range=(2.0, 3.0), n_bins=4) + observables = bias.evaluate(_harmonic_frame([9.0], device)).observables + assert float(observables["cv"][0]) == pytest.approx(9.0, abs=1e-5) + + +# =========================================================================== +# 4. Free energy +# =========================================================================== + + +class TestFreeEnergy: + """Integrating the estimate recovers the analytic PMF.""" + + def test_profile_matches_the_analytic_pmf(self, device: str) -> None: + """``A(r) = 0.5 k (r - r0)^2 - 2 kB T ln r``, up to a constant. + + Nothing is deconvolved: what ABF accumulates already *is* the + free-energy gradient. + """ + bias = _abf(device, cv_range=(1.5, 3.0), n_bins=30) + centers = bias.bin_centers + bias.update(_harmonic_frame(centers.tolist(), device), BiasResult()) + + profile = bias.free_energy() + exact = 0.5 * SPRING * (centers - REST_LENGTH) ** 2 - 2 * KT * torch.log( + centers + ) + exact = exact - exact.min() + # Free energy is defined up to an additive constant. + profile = profile - (profile - exact).mean() + assert float((profile - exact).abs().max()) < 5e-4 + + def test_unsampled_bins_are_nan(self, device: str) -> None: + bias = _abf(device, cv_range=(1.0, 4.0), n_bins=30) + centers = bias.bin_centers + bias.update(_harmonic_frame(centers[5:20].tolist(), device), BiasResult()) + + profile = bias.free_energy() + assert bool(torch.isnan(profile[:5]).all()) + assert bool(torch.isfinite(profile[5:20]).all()) + assert bool(torch.isnan(profile[20:]).all()) + + def test_interior_gap_raises(self, device: str) -> None: + """Integration carries the profile across a hole, so every value + beyond it would be wrong by an unknown constant.""" + bias = _abf(device, cv_range=(1.0, 4.0), n_bins=30) + centers = bias.bin_centers + sampled = centers[[5, 6, 7, 20, 21]].tolist() + bias.update(_harmonic_frame(sampled, device), BiasResult()) + + with pytest.raises(RuntimeError, match="never visited but lie between"): + bias.free_energy() + + def test_no_samples_raises(self, device: str) -> None: + with pytest.raises(RuntimeError, match="no bin has been sampled"): + _abf(device).free_energy() + + def test_single_bin_is_flat(self, device: str) -> None: + bias = _abf(device, n_bins=1, cv_range=(1.0, 4.0)) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + assert float(bias.free_energy()[0]) == 0.0 + + +# =========================================================================== +# 5. Runner integration and observation ordering +# =========================================================================== + + +class TestRunnerIntegration: + """The runner must hand ABF unbiased forces, exactly once per due step.""" + + def test_observation_stage_is_after_compute(self) -> None: + assert _abf().observation_stage is DynamicsStage.AFTER_COMPUTE + + def test_observes_physical_not_total_forces(self, device: str) -> None: + """An estimator fed its own output converges to what it already said. + + This is the whole reason ABF observes at ``AFTER_COMPUTE``. + """ + observed: list[torch.Tensor] = [] + + class Spy(AdaptiveBiasingForce): + def update(self, frames: Batch, result: BiasResult) -> None: + observed.append(frames.forces.clone()) + super().update(frames, result) + + bias = Spy( + atom_indices=torch.tensor([0, 3]), + temperature=TEMPERATURE, + cv_range=(0.5, 6.0), + n_bins=40, + min_samples=0, + full_samples=0, + name="abf", + ).to(device) + runner = EnhancedSampling(_make_dynamics(device), {"abf": bias}) + runner.run(_runner_batch(device=device), n_steps=3) + + physical = runner.last_outputs["physical/forces"] + total = runner.last_outputs["total/forces"] + assert float(runner.last_outputs["bias/abf/forces"].abs().max()) > 0.0 + assert torch.allclose(observed[-1], physical, atol=1e-6) + assert not torch.allclose(observed[-1], total, atol=1e-6) + + def test_contributes_no_energy_to_the_total(self, device: str) -> None: + bias = _abf(device, atom_indices=torch.tensor([0, 3]), name="abf") + runner = EnhancedSampling(_make_dynamics(device), {"abf": bias}) + runner.run(_runner_batch(device=device), n_steps=4) + + outputs = runner.last_outputs + assert "bias/abf/energy" not in outputs + assert torch.allclose(outputs["total/energy"], outputs["physical/energy"]) + + def test_samples_accumulate_over_a_run(self, device: str) -> None: + bias = _abf(device, atom_indices=torch.tensor([0, 3]), name="abf") + runner = EnhancedSampling(_make_dynamics(device), {"abf": bias}) + batch = _runner_batch(device=device) + runner.run(batch, n_steps=6) + + # One sample per walker per step; two walkers. + assert int(bias.bin_counts.sum()) == 6 * batch.num_graphs + + def test_update_frequency_is_respected(self, device: str) -> None: + bias = _abf( + device, atom_indices=torch.tensor([0, 3]), name="abf", update_frequency=3 + ) + runner = EnhancedSampling(_make_dynamics(device), {"abf": bias}) + batch = _runner_batch(device=device) + runner.run(batch, n_steps=9) + assert int(bias.bin_counts.sum()) == 3 * batch.num_graphs + + def test_no_sampling_during_priming(self, device: str) -> None: + bias = _abf(device, atom_indices=torch.tensor([0, 3]), name="abf") + runner = EnhancedSampling(_make_dynamics(device), {"abf": bias}) + runner.prime_forces(_runner_batch(device=device)) + assert int(bias.bin_counts.sum()) == 0 + + def test_below_threshold_updates_do_not_bump_the_version(self, device: str) -> None: + """Re-priming forces over a bin that applies nothing is pure cost. + + This is the case ``bump_state_version`` documents. + """ + bias = _abf(device, min_samples=1000, full_samples=2000) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + assert bias.state_version == 0 + + def test_above_threshold_updates_bump_the_version(self, device: str) -> None: + bias = _abf(device, min_samples=1, full_samples=2) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + bias.update(_harmonic_frame([2.5], device), BiasResult()) + assert bias.state_version > 0 + + def test_composes_with_a_conservative_bias(self, device: str) -> None: + """ABF's force-only result must sum with an energy-carrying one.""" + from nvalchemi.enhanced_sampling import UpperWall + + indices = torch.tensor([0, 3]) + wall = UpperWall( + cv=lambda b: pair_distance(b, indices), + threshold=0.5, + stiffness=5.0, + name="wall", + ).to(device) + bias = _abf(device, atom_indices=indices, name="abf") + runner = EnhancedSampling(_make_dynamics(device), {"abf": bias, "wall": wall}) + runner.run(_runner_batch(device=device), n_steps=3) + + outputs = runner.last_outputs + assert float(outputs["bias/wall/energy"].abs().sum()) > 0.0 + assert torch.allclose( + outputs["total/energy"], + outputs["physical/energy"] + outputs["bias_total/energy"], + atol=1e-5, + ) + + +# =========================================================================== +# 6. Replica-exchange rejection +# =========================================================================== + + +class TestExchangeRejection: + """A force-only bias has no cross-state energy for the acceptance rule.""" + + @staticmethod + def _ladder() -> ReplicaExchange: + states = [ + ThermodynamicState(state_id=i, temperature=t) + for i, t in enumerate((300.0, 350.0)) + ] + return ReplicaExchange( + states=states, + attempt_interval=2, + initial_state_ids=torch.tensor([0, 1]), + ) + + def test_declares_no_exchange_energy(self) -> None: + assert _abf().supplies_exchange_energy is False + + def test_validate_for_rejects_it(self) -> None: + with pytest.raises(ValueError, match="supplies no exchange energy"): + self._ladder().validate_for({"abf": _abf()}) + + def test_runner_rejects_abf_with_exchange(self, device: str) -> None: + bias = _abf(device, atom_indices=torch.tensor([0, 3]), name="abf") + with pytest.raises(ValueError, match="supplies no exchange energy"): + EnhancedSampling( + _make_dynamics(device), + {"abf": bias}, + replica_exchange=self._ladder(), + ) + + +# =========================================================================== +# 7. Restart +# =========================================================================== + + +class TestRestart: + """Bin counts and accumulated statistics survive a round trip.""" + + def test_state_dict_round_trip(self, device: str) -> None: + bias = _abf(device) + bias.update(_harmonic_frame([1.6, 2.5, 3.1], device), BiasResult()) + + restored = _abf(device) + restored.load_state_dict(bias.state_dict()) + + assert torch.equal(restored.bin_counts, bias.bin_counts) + assert torch.allclose(restored.force_sum, bias.force_sum) + assert restored.state_version == bias.state_version + assert torch.allclose( + restored.evaluate(_harmonic_frame([2.5], device)).forces, + bias.evaluate(_harmonic_frame([2.5], device)).forces, + ) + + def test_restart_continues_averaging(self, device: str) -> None: + """A restored run must extend the average, not restart it.""" + bias = _abf(device, n_bins=1, cv_range=(1.0, 4.0)) + bias.update(_harmonic_frame([1.5, 2.5], device), BiasResult()) + + restored = _abf(device, n_bins=1, cv_range=(1.0, 4.0)) + restored.load_state_dict(bias.state_dict()) + restored.update(_harmonic_frame([3.5], device), BiasResult()) + bias.update(_harmonic_frame([3.5], device), BiasResult()) + + assert int(restored.bin_counts[0]) == 3 + assert float(restored.mean_force()[0]) == pytest.approx( + float(bias.mean_force()[0]), rel=1e-6 + ) + + def test_restoring_a_different_cv_range_raises(self, device: str) -> None: + """The histogram's bins mean nothing without the range that made them. + + The counts are shape-compatible, so nothing structural objects; bin + 5 simply stops meaning one distance and starts meaning another, + carrying its accumulated mean force with it. + """ + source = _abf(device, cv_range=(1.0, 4.0)) + source.update(_harmonic_frame([1.5, 2.5], device), BiasResult()) + + target = _abf(device, cv_range=(2.0, 8.0)) + with pytest.raises(ValueError, match="cv_range"): + target.load_state_dict(source.state_dict()) + + def test_restoring_a_different_temperature_raises(self, device: str) -> None: + """The metric correction is already folded into ``force_sum``.""" + source = _abf(device) + source.update(_harmonic_frame([2.5], device), BiasResult()) + + target = _abf(device, temperature=900.0) + with pytest.raises(ValueError, match="temperature"): + target.load_state_dict(source.state_dict()) + + def test_restoring_a_different_atom_pair_raises(self, device: str) -> None: + """atom_indices is a buffer, so an unchecked load would overwrite it. + + The caller's selection would be silently replaced by the + checkpoint's — the opposite of what asking for it meant. + """ + source = _abf(device, atom_indices=torch.tensor([0, 1])) + source.update(_harmonic_frame([2.5], device), BiasResult()) + + target = _abf(device, atom_indices=torch.tensor([5, 7])) + with pytest.raises(ValueError, match="atom_indices"): + target.load_state_dict(source.state_dict()) + # The rejection must come before the buffer is overwritten. + assert target.atom_indices.tolist() == [5, 7] + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("n_bins", 20), + ("min_samples", 7), + ("full_samples", 99), + ("max_force", 0.25), + ], + ) + def test_restoring_a_different_setting_raises( + self, device: str, field: str, value: object + ) -> None: + base = {"min_samples": 1, "full_samples": 100} + source = _abf(device, **base) + source.update(_harmonic_frame([2.5], device), BiasResult()) + + target = _abf(device, **{**base, field: value}) + with pytest.raises(ValueError, match=field): + target.load_state_dict(source.state_dict()) + + def test_the_error_names_every_difference(self, device: str) -> None: + source = _abf(device, cv_range=(1.0, 4.0)) + source.update(_harmonic_frame([2.5], device), BiasResult()) + target = _abf(device, cv_range=(2.0, 8.0), temperature=900.0) + + with pytest.raises(ValueError) as excinfo: + target.load_state_dict(source.state_dict()) + message = str(excinfo.value) + assert "cv_range" in message and "temperature" in message + + def test_identical_configuration_still_restores(self, device: str) -> None: + """The check must not reject a legitimate continuation.""" + source = _abf(device) + source.update(_harmonic_frame([1.5, 2.5], device), BiasResult()) + + target = _abf(device) + target.load_state_dict(source.state_dict()) + assert torch.equal(target.bin_counts, source.bin_counts) + + def test_fingerprint_survives_the_zarr_round_trip( + self, tmp_path, device: str + ) -> None: + """A mismatch must be caught through the runner, not only in memory.""" + bias = _abf(device, atom_indices=torch.tensor([0, 3]), name="abf") + runner = EnhancedSampling( + _make_dynamics(device), {"abf": bias}, steps_per_epoch=4 + ) + batch = runner.run(_runner_batch(device=device), n_steps=4) + path = tmp_path / "abf.zarr" + runner.checkpoint(path, batch) + + wrong = AdaptiveBiasingForce( + atom_indices=torch.tensor([0, 3]), + temperature=TEMPERATURE, + cv_range=(9.0, 12.0), + n_bins=60, + min_samples=0, + full_samples=0, + name="abf", + ).to(device) + fresh = EnhancedSampling( + _make_dynamics(device), {"abf": wrong}, steps_per_epoch=4 + ) + with pytest.raises(ValueError, match="cv_range"): + fresh.restore(path, device=device) + + def test_checkpoint_round_trip_through_the_runner( + self, tmp_path, device: str + ) -> None: + bias = _abf(device, atom_indices=torch.tensor([0, 3]), name="abf") + runner = EnhancedSampling( + _make_dynamics(device), {"abf": bias}, steps_per_epoch=4 + ) + batch = runner.run(_runner_batch(device=device), n_steps=4) + + path = tmp_path / "abf.zarr" + runner.checkpoint(path, batch) + + fresh_bias = _abf(device, atom_indices=torch.tensor([0, 3]), name="abf") + fresh = EnhancedSampling( + _make_dynamics(device), {"abf": fresh_bias}, steps_per_epoch=4 + ) + fresh.restore(path, device=device) + + assert torch.equal(fresh_bias.bin_counts, bias.bin_counts) + assert torch.allclose(fresh_bias.force_sum, bias.force_sum) + assert fresh_bias.state_version == bias.state_version + + +# =========================================================================== +# 8. Geometry +# =========================================================================== + + +class TestGeometry: + """ABF's CV must agree with ``pair_distance`` exactly.""" + + def test_displacement_norm_is_the_distance(self, device: str) -> None: + batch = _harmonic_frame([1.7, 2.9], device) + indices = torch.tensor([0, 1], device=batch.positions.device) + assert torch.allclose( + torch.linalg.vector_norm( + pair_displacement(batch, indices), dim=-1, keepdim=True + ), + pair_distance(batch, indices), + ) + + def test_per_graph_pairs_are_supported(self, device: str) -> None: + bias = _abf(device, atom_indices=torch.tensor([[0, 1], [1, 0]])) + bias.update(_harmonic_frame([2.5, 2.5], device), BiasResult()) + assert int(bias.bin_counts.sum()) == 2 + + def test_bias_follows_the_batch_device(self, device: str) -> None: + """A bias built before the batch moved to GPU must still work.""" + bias = _abf("cpu") + batch = _harmonic_frame([2.5], device) + bias.update(batch, BiasResult()) + result = bias.evaluate(batch) + assert result.forces.device.type == batch.positions.device.type diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py new file mode 100644 index 00000000..a4eccae7 --- /dev/null +++ b/test/enhanced_sampling/test_bias_core.py @@ -0,0 +1,1738 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for core enhanced-sampling abstractions. + +Covers: + +* :class:`~nvalchemi.enhanced_sampling.BiasResult` — shape validation, + detachment enforcement, stress/virial mutual exclusion. +* :class:`~nvalchemi.enhanced_sampling.BiasPotential` — structural + Protocol check. +* :class:`~nvalchemi.enhanced_sampling.ConservativeBias` — forces and + tensile-positive Cauchy stress from autograd; compare both with finite + differences; stress symmetry; no ``requires_grad`` escape into live + batch or result; no memory growth across 10 repeated evaluations. +* :func:`~nvalchemi.enhanced_sampling.pair_distance` — nonperiodic and + Minkowski-reduced triclinic MIC; shared and per-graph atom indices; + gradients via ``torch.autograd.gradcheck``; compile-stability under + ``torch.compile`` (fullgraph=True on CPU); unreduced-cell rejection in + eager mode (check skipped under compile — caller responsibility). +* :func:`~nvalchemi.enhanced_sampling.aggregate_bias_results` — summing, + None handling, duplicate-key rejection. +* ``torch.compile`` tests: ``pair_distance`` and ``aggregate_bias_results`` + compile with ``fullgraph=True``; ``ConservativeBias.energy()`` compiles + with ``fullgraph=True``; ``ConservativeBias.evaluate()`` runs under + ``fullgraph=False`` (graph break at ``requires_grad_()`` is documented). + +GPU integration tests are marked ``@pytest.mark.slow`` and are run only +when a CUDA device is available (the ``device`` fixture handles skip). +""" + +from __future__ import annotations + +import gc + +import pytest +import torch +from torch import Tensor + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.enhanced_sampling import ( + BiasPotential, + BiasResult, + ConservativeBias, + aggregate_bias_results, + pair_distance, +) + +# --------------------------------------------------------------------------- +# Shared batch-construction helpers +# --------------------------------------------------------------------------- + + +def _make_nonperiodic_batch( + n_graphs: int = 2, + atoms_per_graph: int = 4, + device: str = "cpu", + seed: int = 42, +) -> Batch: + """Return a simple non-periodic Batch with known positions.""" + torch.manual_seed(seed) + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6] * atoms_per_graph, dtype=torch.long), + positions=torch.randn(atoms_per_graph, 3), + ) + for _ in range(n_graphs) + ] + batch = Batch.from_data_list(data_list).to(device) + batch["energy"] = torch.zeros(n_graphs, 1, device=device) + batch["forces"] = torch.zeros(atoms_per_graph * n_graphs, 3, device=device) + return batch + + +def _make_cubic_batch( + n_graphs: int = 2, + atoms_per_graph: int = 4, + box: float = 5.0, + device: str = "cpu", + seed: int = 42, +) -> Batch: + """Return a Batch with cubic unit cells and full 3D PBC.""" + torch.manual_seed(seed) + data_list = [] + for _ in range(n_graphs): + positions = torch.rand(atoms_per_graph, 3) * box + # AtomicData expects cell as [1, 3, 3] and pbc as [1, 3] + cell = torch.eye(3).unsqueeze(0) * box + pbc = torch.tensor([[True, True, True]]) + data_list.append( + AtomicData( + atomic_numbers=torch.tensor([6] * atoms_per_graph, dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + ) + batch = Batch.from_data_list(data_list).to(device) + batch["energy"] = torch.zeros(n_graphs, 1, device=device) + batch["forces"] = torch.zeros(atoms_per_graph * n_graphs, 3, device=device) + return batch + + +def _make_triclinic_batch( + device: str = "cpu", + seed: int = 0, +) -> Batch: + """Return a single-graph Batch with a triclinic unit cell.""" + torch.manual_seed(seed) + # Tilted cell: a = [5,0,0], b = [1,5,0], c = [0.5,0.5,5] + cell_mat = torch.tensor([[5.0, 0.0, 0.0], [1.0, 5.0, 0.0], [0.5, 0.5, 5.0]]) + # AtomicData expects [1, 3, 3] and [1, 3] + cell = cell_mat.unsqueeze(0) + pbc = torch.tensor([[True, True, True]]) + positions = torch.rand(4, 3) @ cell_mat # Cartesian, inside cell + data = AtomicData( + atomic_numbers=torch.tensor([6, 6, 6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + batch["energy"] = torch.zeros(1, 1, device=device) + batch["forces"] = torch.zeros(4, 3, device=device) + return batch + + +# =========================================================================== +# 1. BiasResult +# =========================================================================== + + +class TestBiasResult: + """Tests for the BiasResult dataclass.""" + + def test_empty_construction(self) -> None: + r = BiasResult() + assert r.energy is None + assert r.forces is None + assert r.observables == {} + + def test_detached_tensors_accepted(self) -> None: + e = torch.tensor([[1.0]]).detach() + f = torch.zeros(3, 3).detach() + r = BiasResult(energy=e, forces=f) + assert r.energy is e + + def test_requires_grad_energy_raises(self) -> None: + bad = torch.tensor([[1.0]], requires_grad=True) + with pytest.raises(ValueError, match="energy.*detached"): + BiasResult(energy=bad) + + def test_requires_grad_forces_raises(self) -> None: + bad = torch.zeros(3, 3, requires_grad=True) + with pytest.raises(ValueError, match="forces.*detached"): + BiasResult(forces=bad) + + def test_grad_fn_raises(self) -> None: + x = torch.tensor([[1.0]], requires_grad=True) + y = x * 2.0 # has grad_fn + with pytest.raises(ValueError, match="energy.*detached"): + BiasResult(energy=y) + + def test_stress_and_virial_raises(self) -> None: + s = torch.zeros(1, 3, 3) + v = torch.zeros(1, 3, 3) + with pytest.raises(ValueError, match="stress.*virial"): + BiasResult(stress=s, virial=v) + + def test_observable_requires_grad_raises(self) -> None: + bad = torch.zeros(3, requires_grad=True) + with pytest.raises(ValueError, match="observables"): + BiasResult(observables={"cv": bad}) + + def test_frozen_immutability(self) -> None: + r = BiasResult(energy=torch.zeros(1, 1)) + with pytest.raises((TypeError, AttributeError)): + r.energy = torch.ones(1, 1) # type: ignore[misc] + + # --- shape validation --- + + def test_energy_wrong_ndim_raises(self) -> None: + """energy must be [B, 1]; a flat [B] tensor is rejected.""" + with pytest.raises(ValueError, match="energy.*\\[B, 1\\]"): + BiasResult(energy=torch.zeros(2)) + + def test_energy_wrong_trailing_dim_raises(self) -> None: + """energy last dim must be 1, not 3.""" + with pytest.raises(ValueError, match="energy.*\\[B, 1\\]"): + BiasResult(energy=torch.zeros(2, 3)) + + def test_forces_wrong_ndim_raises(self) -> None: + """forces must be [N, 3]; a 1-D tensor is rejected.""" + with pytest.raises(ValueError, match="forces.*\\[N, 3\\]"): + BiasResult(forces=torch.zeros(9)) + + def test_forces_wrong_width_raises(self) -> None: + """forces last dim must be 3, not 1.""" + with pytest.raises(ValueError, match="forces.*\\[N, 3\\]"): + BiasResult(forces=torch.zeros(4, 1)) + + def test_stress_wrong_shape_raises(self) -> None: + """stress must be [B, 3, 3]; a [B, 3] tensor is rejected.""" + with pytest.raises(ValueError, match="stress.*\\[B, 3, 3\\]"): + BiasResult(stress=torch.zeros(2, 3)) + + def test_virial_wrong_shape_raises(self) -> None: + """virial must be [B, 3, 3].""" + with pytest.raises(ValueError, match="virial.*\\[B, 3, 3\\]"): + BiasResult(virial=torch.zeros(2, 9)) + + def test_state_version_wrong_ndim_raises(self) -> None: + """state_version must be 1-D.""" + with pytest.raises(ValueError, match="state_version.*\\[B\\]"): + BiasResult(state_version=torch.zeros(2, 1, dtype=torch.int32)) + + def test_state_version_float_dtype_raises(self) -> None: + """state_version must be an integer dtype.""" + with pytest.raises(ValueError, match="integer dtype"): + BiasResult(state_version=torch.zeros(2)) # float32 + + def test_state_version_integer_accepted(self) -> None: + """state_version with int64 dtype is accepted.""" + r = BiasResult(state_version=torch.zeros(2, dtype=torch.int64)) + assert r.state_version is not None + + # --- batch-size consistency --- + + def test_batch_size_mismatch_raises(self) -> None: + """energy [2, 1] and virial [3, 3, 3] have inconsistent B.""" + with pytest.raises(ValueError, match="inconsistent"): + BiasResult(energy=torch.zeros(2, 1), virial=torch.zeros(3, 3, 3)) + + def test_batch_size_consistent_accepted(self) -> None: + """energy [2, 1] and virial [2, 3, 3] with matching B=2 are accepted.""" + r = BiasResult(energy=torch.zeros(2, 1), virial=torch.zeros(2, 3, 3)) + assert r.energy is not None + + # --- finiteness --- + + def test_energy_nan_raises(self) -> None: + with pytest.raises(ValueError, match="energy.*NaN or Inf"): + BiasResult(energy=torch.tensor([[float("nan")]])) + + def test_energy_inf_raises(self) -> None: + with pytest.raises(ValueError, match="energy.*NaN or Inf"): + BiasResult(energy=torch.tensor([[float("inf")]])) + + def test_forces_nan_raises(self) -> None: + bad = torch.zeros(3, 3) + bad[1, 2] = float("nan") + with pytest.raises(ValueError, match="forces.*NaN or Inf"): + BiasResult(forces=bad) + + def test_virial_inf_raises(self) -> None: + bad = torch.zeros(1, 3, 3) + bad[0, 0, 0] = float("-inf") + with pytest.raises(ValueError, match="virial.*NaN or Inf"): + BiasResult(virial=bad) + + def test_observable_nan_raises(self) -> None: + with pytest.raises(ValueError, match="observables.*NaN or Inf"): + BiasResult(observables={"cv": torch.tensor([float("nan")])}) + + def test_valid_result_accepted(self) -> None: + """A fully-populated valid BiasResult passes all checks.""" + r = BiasResult( + energy=torch.zeros(2, 1), + forces=torch.zeros(6, 3), + virial=torch.zeros(2, 3, 3), + state_version=torch.zeros(2, dtype=torch.int64), + observables={"bias/a/cv": torch.zeros(2)}, + ) + assert r.energy is not None + + +# =========================================================================== +# 2. BiasPotential Protocol +# =========================================================================== + + +class TestBiasPotentialProtocol: + """Tests for structural protocol membership.""" + + def test_structural_satisfaction(self) -> None: + class MyBias: + name = "my_bias" + + def evaluate(self, current: Batch) -> BiasResult: + return BiasResult() + + assert isinstance(MyBias(), BiasPotential) + + def test_missing_name_not_protocol(self) -> None: + class NotABias: + def evaluate(self, current: Batch) -> BiasResult: + return BiasResult() + + assert not isinstance(NotABias(), BiasPotential) + + def test_missing_evaluate_not_protocol(self) -> None: + class NotABias: + name = "x" + + assert not isinstance(NotABias(), BiasPotential) + + def test_protocol_inherits_nothing(self) -> None: + """The boundary must stay inheritance-free. + + A third party implementing a novel method must not be forced to + inherit BaseModelMixin, nn.Module, or anything else. If this ever + fails, the Protocol has stopped being a structural boundary. + """ + from typing import Generic, Protocol + + from nvalchemi.models.base import BaseModelMixin + + bases = set(BiasPotential.__mro__) - { + BiasPotential, + object, + Protocol, + Generic, + } + assert not bases, f"BiasPotential gained base classes: {bases}" + assert not issubclass(BiasPotential, BaseModelMixin) + assert not issubclass(BiasPotential, torch.nn.Module) + + def test_plain_object_satisfies_protocol_without_any_base(self) -> None: + """A bias with no base class at all is a valid BiasPotential.""" + + class StandaloneBias: + name = "standalone" + + def evaluate(self, current: Batch) -> BiasResult: + return BiasResult(energy=torch.zeros(current.num_graphs, 1)) + + bias = StandaloneBias() + assert isinstance(bias, BiasPotential) + assert type(bias).__mro__ == (StandaloneBias, object) + + def test_conservative_bias_satisfies_protocol_via_mixins(self) -> None: + """ConservativeBias satisfies the protocol through composition.""" + bias = _QuadraticBias() + assert isinstance(bias, BiasPotential) + + +# =========================================================================== +# 3. ConservativeBias — autograd helper +# =========================================================================== + + +class _QuadraticBias(ConservativeBias): + """E = 0.5 * k * ||positions||^2 per graph — analytically tractable.""" + + def __init__(self, k: float = 1.0) -> None: + super().__init__(name="quadratic") + self.k = k + + def energy(self, current: Batch) -> Tensor: + # Sum of squared positions per graph → [B, 1] + # batch_ptr gives atom offsets per graph + ptr = current.batch_ptr + B = current.num_graphs + energies = [] + for b in range(B): + pos_b = current.positions[ptr[b] : ptr[b + 1]] + energies.append(0.5 * self.k * (pos_b**2).sum()) + return torch.stack(energies).unsqueeze(-1) # [B, 1] + + +class _PairDistanceBias(ConservativeBias): + """E = 0.5 * k * pair_distance^2 — uses the pair_distance CV.""" + + def __init__(self, atom_indices: Tensor, k: float = 1.0) -> None: + super().__init__(name="pair_dist_bias") + self.atom_indices = atom_indices + self.k = k + + def energy(self, current: Batch) -> Tensor: + d = pair_distance(current, self.atom_indices) # [B, 1] + return 0.5 * self.k * d**2 # [B, 1] + + +class _AnisotropicBias(ConservativeBias): + """E = sum_n (x_n * y_n) per graph — couples distinct Cartesian components. + + Because the energy mixes the x and y components rather than depending + only on interatomic distances, its derivative with respect to the full + deformation gradient F is asymmetric. The derivative with respect to + the symmetric strain tensor eps is symmetric, which is what the project + virial/stress convention requires. A bias like this is what + distinguishes the two derivatives; a central pair interaction does not. + """ + + def __init__(self) -> None: + super().__init__(name="anisotropic") + + def energy(self, current: Batch) -> Tensor: + ptr = current.batch_ptr + energies = [] + for b in range(current.num_graphs): + pos_b = current.positions[ptr[b] : ptr[b + 1]] + energies.append((pos_b[:, 0] * pos_b[:, 1]).sum()) + return torch.stack(energies).unsqueeze(-1) # [B, 1] + + +class _CellVolumeBias(ConservativeBias): + """E = (V - V0)^2 — depends on the cell but not on positions at all.""" + + def __init__(self, target_volume: float = 900.0) -> None: + super().__init__(name="cell_volume") + self.target_volume = target_volume + + def energy(self, current: Batch) -> Tensor: + volume = torch.linalg.det(current.cell.reshape(-1, 3, 3)) + return ((volume - self.target_volume) ** 2).reshape(-1, 1) # [B, 1] + + +class _ConstantBias(ConservativeBias): + """E = c — depends on neither positions nor cell. + + Stands in for a bias that returns a constant on one branch, e.g. a wall + restraint evaluated while every atom is inside the wall. + """ + + def __init__(self, value: float = 3.0) -> None: + super().__init__(name="constant") + self.value = value + + def energy(self, current: Batch) -> Tensor: + return torch.full( + (current.num_graphs, 1), + self.value, + dtype=current.positions.dtype, + device=current.positions.device, + ) + + +class TestConservativeBias: + """Tests for ConservativeBias autograd helper.""" + + def test_forces_shape(self, device: str) -> None: + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=3, device=device) + bias = _QuadraticBias(k=1.0) + result = bias.evaluate(batch) + assert result.forces is not None + assert result.forces.shape == (6, 3) + + def test_energy_shape(self, device: str) -> None: + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=3, device=device) + bias = _QuadraticBias(k=1.0) + result = bias.evaluate(batch) + assert result.energy is not None + assert result.energy.shape == (2, 1) + + def test_forces_analytical_vs_autograd(self, device: str) -> None: + """F = -dE/dr; for E = 0.5 * k * ||r||^2, F = -k * r.""" + k = 2.0 + batch = _make_nonperiodic_batch(n_graphs=1, atoms_per_graph=4, device=device) + bias = _QuadraticBias(k=k) + result = bias.evaluate(batch) + expected_forces = -k * batch.positions + assert result.forces is not None + assert torch.allclose(result.forces, expected_forces, atol=1e-5) + + def test_forces_finite_difference(self, device: str) -> None: + """Compare autograd forces to central-difference finite differences.""" + k = 1.0 + eps = 1e-4 + batch = _make_nonperiodic_batch(n_graphs=1, atoms_per_graph=3, device=device) + bias = _QuadraticBias(k=k) + + pos = batch.positions.clone() # [N, 3] + N = pos.shape[0] + fd_forces = torch.zeros_like(pos) + for i in range(N): + for j in range(3): + pos_plus = pos.clone() + pos_plus[i, j] += eps + batch["positions"] = pos_plus + e_plus = bias.evaluate(batch).energy.sum().item() + + pos_minus = pos.clone() + pos_minus[i, j] -= eps + batch["positions"] = pos_minus + e_minus = bias.evaluate(batch).energy.sum().item() + + fd_forces[i, j] = -(e_plus - e_minus) / (2 * eps) + + batch["positions"] = pos + result = bias.evaluate(batch) + assert result.forces is not None + # float32 finite differences at eps=1e-4 have ~1e-3 cancellation error; + # use a tolerance that accounts for float32 precision. + assert torch.allclose(result.forces, fd_forces, atol=5e-3) + + def test_result_fully_detached(self, device: str) -> None: + """BiasResult tensors must have requires_grad=False and grad_fn=None.""" + batch = _make_nonperiodic_batch(device=device) + bias = _QuadraticBias() + result = bias.evaluate(batch) + for name in ("energy", "forces"): + t = getattr(result, name) + if t is not None: + assert not t.requires_grad, f"{name} has requires_grad=True" + assert t.grad_fn is None, f"{name} has non-null grad_fn" + + def test_live_batch_positions_not_mutated(self, device: str) -> None: + """batch.positions must be restored to original tensor after evaluate().""" + batch = _make_nonperiodic_batch(device=device) + original_pos = batch.positions + original_data = original_pos.clone() + bias = _QuadraticBias() + bias.evaluate(batch) + # The tensor object should be restored + assert batch.positions is original_pos + # Values should be unchanged + assert torch.allclose(batch.positions, original_data) + + def test_live_batch_positions_no_grad(self, device: str) -> None: + """After evaluate(), batch.positions must not have requires_grad=True.""" + batch = _make_nonperiodic_batch(device=device) + bias = _QuadraticBias() + bias.evaluate(batch) + assert not batch.positions.requires_grad + assert batch.positions.grad_fn is None + + def test_no_memory_growth_repeated_evaluate(self, device: str) -> None: + """Repeated evaluate() must not grow GPU allocated memory monotonically. + + Warm up 3 calls, then sample allocated memory over 10 calls. The + delta between first and last sample must be ≤ 0 (or a small + tolerance for caching effects). + """ + batch = _make_nonperiodic_batch(n_graphs=4, atoms_per_graph=8, device=device) + bias = _QuadraticBias() + + # Warm up + for _ in range(3): + bias.evaluate(batch) + + gc.collect() + if device == "cuda": + torch.cuda.synchronize() + torch.cuda.empty_cache() + mem_start = torch.cuda.memory_allocated() + else: + mem_start = 0 + + for _ in range(10): + bias.evaluate(batch) + + if device == "cuda": + torch.cuda.synchronize() + mem_end = torch.cuda.memory_allocated() + # Allow a small tolerance (1 MB) for CUDA caching allocator overhead + assert mem_end - mem_start <= 1 * 1024 * 1024, ( + f"GPU memory grew by {mem_end - mem_start} bytes across 10 evaluate() calls" + ) + + def test_stress_analytical_across_image_boundary(self, device: str) -> None: + """Cauchy stress is correct for a pair bias whose MIC vector crosses an image. + + Setup + ----- + Box: 10 Å cubic (V = 1000 ų). Atom 0 at [0.5, 0, 0], atom 1 at + [9.5, 0, 0]. MIC distance = 1 Å (image at x − 10, so + dr_mic = [−1, 0, 0]). Bias: E = 0.5 · k · d² (k = 1 eV/Ų). + + Analytical derivation + --------------------- + Under a homogeneous symmetric strain ε both positions and cell deform, + so the MIC vector deforms with them (the image index [−1, 0, 0] is + fixed):: + + dr_mic → dr_mic @ (I + ε) + d² = |dr_mic|² + 2 · dr_mic · ε · dr_micᵀ + O(ε²) + + dE/dε_kl |_{ε=0} = k · dr_mic[k] · dr_mic[l] + σ_kl = (dE/dε_kl) / V = k · dr_mic[k] · dr_mic[l] / V + + For dr_mic = [−1, 0, 0], k = 1, V = 1000: + σ[0,0] = +1e-3 eV/ų, all other elements = 0. Positive (tensile) is + the expected sign: the harmonic restraint pulls the two atoms together. + + This is only correct when positions and cell are strained together. + A strain leaf applied to the cell alone misses the atomic-position + contribution and returns the wrong answer. + """ + k = 1.0 + box = 10.0 + volume = box**3 + # MIC distance = |9.5 - 0.5 - 10| = 1 Å; dr_mic = [-1, 0, 0] + positions = torch.tensor([[0.5, 0.0, 0.0], [9.5, 0.0, 0.0]]) + cell = torch.eye(3).unsqueeze(0) * box # [1, 3, 3] + pbc = torch.tensor([[True, True, True]]) + + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + bias = _PairDistanceBias(atom_indices=idx, k=k) + result = bias.evaluate(batch) + + assert result.stress is not None, "stress should be non-None for periodic batch" + assert result.stress.shape == (1, 3, 3) + assert result.virial is None, "ConservativeBias emits stress, not virial" + + # Analytical: σ = k · outer(dr_mic, dr_mic) / V with dr_mic = [−1, 0, 0] + expected = torch.zeros(3, 3, device=device) + expected[0, 0] = k / volume + sigma = result.stress[0] # [3, 3] + assert torch.allclose(sigma, expected, atol=1e-8), ( + f"stress = {sigma}, expected {expected}. Stress may be missing the " + "atomic-position contribution (strain not applied to both positions " + "and cell simultaneously)." + ) + + def test_stress_is_symmetric(self, device: str) -> None: + """Stress must be symmetric: the project strain tensor ε is symmetric. + + A bias whose energy is not a central pair interaction is the case that + distinguishes a symmetric strain derivative from a raw deformation + gradient derivative. ``_AnisotropicBias`` below uses a per-component + weighting so that dE/dF is asymmetric while dE/dε is not. + """ + batch = _make_cubic_batch(n_graphs=2, atoms_per_graph=5, device=device) + bias = _AnisotropicBias() + result = bias.evaluate(batch) + + assert result.stress is not None + sigma = result.stress + assert torch.allclose(sigma, sigma.mT, atol=1e-6), ( + f"stress is not symmetric:\n{sigma}\nvs transpose\n{sigma.mT}" + ) + + def test_no_stress_for_nonperiodic_batch(self, device: str) -> None: + """A batch with no cell yields forces but no stress.""" + batch = _make_nonperiodic_batch(device=device) + result = _QuadraticBias().evaluate(batch) + assert result.forces is not None + assert result.stress is None + assert result.virial is None + + def test_compute_stress_false_skips_stress(self, device: str) -> None: + """``compute_stress=False`` drops 'stress' from active_outputs.""" + + class _ForceOnlyBias(_QuadraticBias): + def __init__(self) -> None: + super().__init__() + self.model_config.active_outputs = {"energy", "forces"} + + bias = _ForceOnlyBias() + assert "stress" not in bias.model_config.active_outputs + # Capability is unchanged; only the runtime selection narrowed. + assert "stress" in bias.model_config.outputs + + batch = _make_cubic_batch(n_graphs=1, atoms_per_graph=4, device=device) + result = bias.evaluate(batch) + assert result.forces is not None + assert result.stress is None + + def test_active_outputs_toggled_at_runtime(self, device: str) -> None: + """active_outputs is a runtime field: flipping it changes the result.""" + batch = _make_cubic_batch(n_graphs=1, atoms_per_graph=4, device=device) + bias = _QuadraticBias() + + assert bias.evaluate(batch).stress is not None + bias.model_config.active_outputs = {"energy", "forces"} + assert bias.evaluate(batch).stress is None + bias.model_config.active_outputs = {"energy", "forces", "stress"} + assert bias.evaluate(batch).stress is not None + + def test_live_batch_cell_restored(self, device: str) -> None: + """batch.cell must be restored to the original tensor after evaluate().""" + batch = _make_cubic_batch(n_graphs=2, atoms_per_graph=4, device=device) + original_cell = batch.cell + original_data = original_cell.clone() + _QuadraticBias().evaluate(batch) + assert batch.cell is original_cell + assert torch.allclose(batch.cell, original_data) + assert not batch.cell.requires_grad + assert batch.cell.grad_fn is None + + def test_stress_finite_difference(self, device: str) -> None: + """Compare autograd stress to a finite-difference strain derivative. + + Applies a symmetric strain ``ε`` to both positions and cell, and + checks ``σ_kl ≈ (E(+ε) − E(−ε)) / (2 h V)`` for each component. + Uses float64 so the central difference is not dominated by + cancellation error. + """ + h = 1e-5 + box = 6.0 + volume = box**3 + torch.manual_seed(7) + positions = torch.rand(5, 3, dtype=torch.float64) * box + cell = torch.eye(3, dtype=torch.float64).unsqueeze(0) * box + + data = AtomicData( + atomic_numbers=torch.tensor([6] * 5, dtype=torch.long), + positions=positions, + cell=cell, + pbc=torch.tensor([[True, True, True]]), + ) + batch = Batch.from_data_list([data]).to(device) + bias = _PairDistanceBias( + atom_indices=torch.tensor([0, 3], device=device), k=1.5 + ) + + result = bias.evaluate(batch) + assert result.stress is not None + + base_pos = batch.positions.clone() + base_cell = batch.cell.clone() + eye = torch.eye(3, dtype=base_pos.dtype, device=base_pos.device) + + fd_stress = torch.zeros(3, 3, dtype=base_pos.dtype, device=base_pos.device) + for a in range(3): + for b in range(3): + # Symmetric strain perturbation in component (a, b). + eps = torch.zeros(3, 3, dtype=base_pos.dtype, device=base_pos.device) + eps[a, b] += 0.5 + eps[b, a] += 0.5 + eps = eps * h + + energies = [] + for sign in (+1.0, -1.0): + deform = eye + sign * eps + batch["positions"] = base_pos @ deform + batch["cell"] = base_cell @ deform + energies.append(bias.evaluate(batch).energy.sum().item()) + + fd_stress[a, b] = (energies[0] - energies[1]) / (2 * h * volume) + + batch["positions"] = base_pos + batch["cell"] = base_cell + + assert torch.allclose(result.stress[0], fd_stress, atol=1e-7), ( + f"autograd stress\n{result.stress[0]}\ndiffers from finite differences\n" + f"{fd_stress}" + ) + + def test_position_independent_bias_gives_zero_forces(self, device: str) -> None: + """A cell-only bias must evaluate, not raise, and yield zero forces. + + ``E = (V - V0)^2`` has no position dependence, so ``dE/dr`` is zero. + Differentiating with ``allow_unused=False`` would instead raise + ``RuntimeError: One of the differentiated Tensors appears to not have + been used in the graph`` from inside ``torch.autograd.grad``, with + nothing pointing at the user's ``energy()``. + """ + box = 10.0 + volume = box**3 + target = 900.0 + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.tensor([[1.0, 1.0, 1.0], [2.0, 3.0, 4.0]]), + cell=torch.eye(3).unsqueeze(0) * box, + pbc=torch.tensor([[True, True, True]]), + ) + batch = Batch.from_data_list([data]).to(device) + result = _CellVolumeBias(target_volume=target).evaluate(batch) + + assert result.forces is not None + assert result.forces.shape == (2, 3) + assert torch.count_nonzero(result.forces) == 0, ( + f"a position-independent bias must give zero forces, got {result.forces}" + ) + + # dE/dV = 2 (V - V0); dV/deps = V * I => sigma = dE/deps / V = 2 (V - V0) I + assert result.stress is not None + expected = torch.eye(3, device=result.stress.device) * 2.0 * (volume - target) + assert torch.allclose(result.stress[0], expected, rtol=1e-5), ( + f"stress = {result.stress[0]}, expected {expected}" + ) + + def test_constant_bias_gives_zero_forces_and_stress(self, device: str) -> None: + """A bias disconnected from the graph entirely evaluates to zeros. + + This is the "returns a constant on some branch" case, e.g. a wall + restraint while every atom is inside the wall. + """ + batch = _make_cubic_batch(n_graphs=2, atoms_per_graph=3, device=device) + result = _ConstantBias(value=3.0).evaluate(batch) + + assert result.energy is not None + assert torch.allclose(result.energy, torch.full_like(result.energy, 3.0)) + assert result.forces is not None + assert torch.count_nonzero(result.forces) == 0 + assert result.stress is not None + assert torch.count_nonzero(result.stress) == 0 + + def test_constant_bias_nonperiodic_gives_zero_forces(self, device: str) -> None: + """The no-cell path also tolerates an energy with no position dependence.""" + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=3, device=device) + result = _ConstantBias(value=1.5).evaluate(batch) + + assert result.forces is not None + assert result.forces.shape == (6, 3) + assert torch.count_nonzero(result.forces) == 0 + assert result.stress is None + + def test_evaluate_is_read_only_no_state_change(self, device: str) -> None: + """Multiple evaluate() calls must leave bias state unchanged.""" + batch = _make_nonperiodic_batch(device=device) + bias = _QuadraticBias(k=2.5) + r1 = bias.evaluate(batch) + r2 = bias.evaluate(batch) + assert result_close(r1, r2) + + +class TestConservativeBiasModelMixin: + """ConservativeBias composes nn.Module + BaseModelMixin correctly.""" + + def test_mro_is_house_idiom(self) -> None: + """nn.Module before BaseModelMixin, as PyTorch requires.""" + from nvalchemi.models.base import BaseModelMixin + + mro = ConservativeBias.__mro__ + assert issubclass(ConservativeBias, torch.nn.Module) + assert issubclass(ConservativeBias, BaseModelMixin) + assert mro.index(torch.nn.Module) < mro.index(BaseModelMixin) + + def test_declares_model_config(self) -> None: + bias = _QuadraticBias() + assert bias.model_config.outputs == frozenset({"energy", "forces", "stress"}) + assert bias.model_config.autograd_outputs == frozenset({"forces", "stress"}) + assert bias.model_config.autograd_inputs == frozenset({"positions", "cell"}) + + def test_missing_super_init_raises(self) -> None: + """A subclass that forgets super().__init__() must fail loudly.""" + + class BrokenBias(ConservativeBias): + def __init__(self) -> None: + self.k = 1.0 # no super().__init__() + + def energy(self, current: Batch) -> Tensor: + return torch.zeros(current.num_graphs, 1) + + with pytest.raises((AttributeError, TypeError)): + BrokenBias() + + def test_embeddings_stubbed_like_lj_and_dftd3(self) -> None: + bias = _QuadraticBias() + assert bias.embedding_shapes == {} + with pytest.raises(NotImplementedError, match="does not produce embeddings"): + bias.compute_embeddings(_make_nonperiodic_batch()) + + def test_export_model_raises(self, tmp_path) -> None: + with pytest.raises(NotImplementedError, match="no exportable model"): + _QuadraticBias().export_model(tmp_path / "x.pt") + + def test_distribution_spec_is_none_by_default(self) -> None: + """None makes DomainParallel raise rather than shard undefined semantics.""" + assert _QuadraticBias().distribution_spec() is None + assert _QuadraticBias().distribution_spec(strategy=None) is None + + def test_distribution_spec_is_overridable(self) -> None: + """A bias that is genuinely local can declare its output semantics.""" + from nvalchemi.distributed.output_kinds import OutputKind, OutputSpec, Reduce + from nvalchemi.distributed.spec import ( + DistributionSpec, + HaloStoragePolicy, + MLIPSpec, + ) + + class LocalBias(_QuadraticBias): + def distribution_spec(self, strategy=None): + return MLIPSpec( + distribution=DistributionSpec(policy=HaloStoragePolicy()), + outputs={ + "energy": OutputSpec(OutputKind.PER_GRAPH, Reduce.ALL_REDUCE), + "forces": OutputSpec(OutputKind.PER_NODE, Reduce.OWNED_ONLY), + }, + ) + + spec = LocalBias().distribution_spec() + assert spec is not None + assert spec.output_kinds["forces"] is OutputKind.PER_NODE + + def test_state_dict_round_trip(self) -> None: + """nn.Module gives checkpointing for free.""" + bias = _QuadraticBias() + bias.register_buffer("counter", torch.tensor([3.0])) + state = bias.state_dict() + assert "counter" in state + + restored = _QuadraticBias() + restored.register_buffer("counter", torch.zeros(1)) + restored.load_state_dict(state) + assert torch.allclose(restored.counter, torch.tensor([3.0])) + + def test_forward_returns_model_outputs(self, device: str) -> None: + """forward() is the ModelOutputs view of evaluate().""" + batch = _make_cubic_batch(n_graphs=2, atoms_per_graph=3, device=device) + bias = _QuadraticBias(k=1.5) + outputs = bias(batch) + + assert isinstance(outputs, dict) + assert set(outputs) == {"energy", "forces", "stress"} + reference = bias.evaluate(batch) + assert torch.allclose(outputs["energy"], reference.energy) + assert torch.allclose(outputs["stress"], reference.stress) + + def test_forward_respects_active_outputs(self, device: str) -> None: + batch = _make_cubic_batch(n_graphs=1, atoms_per_graph=3, device=device) + bias = _QuadraticBias() + bias.model_config.active_outputs = {"energy"} + assert set(bias(batch)) == {"energy"} + + def test_forward_output_sums_with_model_outputs(self, device: str) -> None: + """The point of ModelOutputs: bias output composes via sum_outputs.""" + from nvalchemi.models._utils import sum_outputs + + batch = _make_cubic_batch(n_graphs=1, atoms_per_graph=3, device=device) + bias_out = _QuadraticBias(k=1.0)(batch) + model_out = { + "energy": torch.ones(1, 1, device=device), + "forces": torch.ones(3, 3, device=device), + } + total = sum_outputs(model_out, bias_out) + assert torch.allclose(total["energy"], model_out["energy"] + bias_out["energy"]) + assert torch.allclose(total["forces"], model_out["forces"] + bias_out["forces"]) + + +def result_close(a: BiasResult, b: BiasResult, atol: float = 1e-6) -> bool: + """Return True iff all non-None tensor fields of a and b are close.""" + for attr in ("energy", "forces", "virial", "stress"): + ta, tb = getattr(a, attr), getattr(b, attr) + if ta is None and tb is None: + continue + if ta is None or tb is None: + return False + if not torch.allclose(ta, tb, atol=atol): + return False + return True + + +# =========================================================================== +# 4. pair_distance CV +# =========================================================================== + + +class TestPairDistance: + """Tests for the pair_distance collective variable.""" + + # --- atom_indices shape / dtype validation --- + + def test_atom_indices_float_dtype_raises(self) -> None: + """Float atom_indices raises ValueError (would silently cast to int).""" + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ) + batch = Batch.from_data_list([data]) + with pytest.raises(ValueError, match="integer dtype"): + pair_distance(batch, torch.tensor([0.0, 1.0])) + + def test_atom_indices_1d_wrong_length_raises(self) -> None: + """1-D atom_indices with length != 2 raises ValueError. + + torch.tensor([0]) would silently become [[0, 0]] (self-distance). + """ + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ) + batch = Batch.from_data_list([data]) + with pytest.raises(ValueError, match="exactly 2 elements"): + pair_distance(batch, torch.tensor([0])) # length 1 + + def test_atom_indices_2d_extra_column_raises(self) -> None: + """[B, 3] atom_indices raises ValueError (extra column would be silently dropped).""" + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6, 6, 6], dtype=torch.long), + positions=torch.zeros(3, 3), + ) + ] * 2 + batch = Batch.from_data_list(data_list) + with pytest.raises(ValueError, match="second dimension must be exactly 2"): + pair_distance(batch, torch.tensor([[0, 1, 2], [0, 1, 2]])) + + def test_atom_indices_2d_wrong_batch_size_raises(self) -> None: + """[B', 2] atom_indices where B' != B raises ValueError.""" + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ) + ] * 3 # B=3 + batch = Batch.from_data_list(data_list) + # supply [2, 2] instead of [3, 2] + with pytest.raises(ValueError, match="first dimension must equal B"): + pair_distance(batch, torch.tensor([[0, 1], [0, 1]])) + + def test_atom_indices_3d_raises(self) -> None: + """3-D atom_indices raises ValueError.""" + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ) + batch = Batch.from_data_list([data]) + with pytest.raises(ValueError, match="1-D.*or.*2-D"): + pair_distance(batch, torch.zeros(1, 2, 1, dtype=torch.long)) + + def test_atom_indices_valid_shapes_accepted(self) -> None: + """Shape [2] and [B, 2] with integer dtype are accepted.""" + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + ) + ] * 2 + batch = Batch.from_data_list(data_list) + # [2] shared + d1 = pair_distance(batch, torch.tensor([0, 1])) + assert d1.shape == (2, 1) + # [B, 2] per-graph + d2 = pair_distance(batch, torch.tensor([[0, 1], [0, 1]])) + assert d2.shape == (2, 1) + + # --- bounds checking --- + + def test_out_of_range_shared_index_raises(self) -> None: + """Shared [2] index that exceeds graph size raises IndexError, not silent wrap.""" + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ), + AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ), + ] + batch = Batch.from_data_list(data_list) + # Local index 5 is valid for neither 2-atom graph. + idx = torch.tensor([0, 5]) + with pytest.raises(IndexError, match="out of range"): + pair_distance(batch, idx) + + def test_out_of_range_per_graph_index_raises(self) -> None: + """Per-graph [B, 2] index out of range for one graph raises IndexError.""" + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6, 6, 6], dtype=torch.long), + positions=torch.zeros(3, 3), + ), + AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ), + ] + batch = Batch.from_data_list(data_list) + # Graph 1 has only 2 atoms; local index 2 is out of range. + idx = torch.tensor([[0, 1], [0, 2]]) + with pytest.raises(IndexError, match="out of range"): + pair_distance(batch, idx) + + def test_negative_index_raises(self) -> None: + """Negative atom index raises IndexError.""" + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.zeros(2, 3), + ) + batch = Batch.from_data_list([data]) + idx = torch.tensor([-1, 0]) + with pytest.raises(IndexError, match="negative"): + pair_distance(batch, idx) + + def test_variable_size_batch_no_silent_cross_graph(self) -> None: + """Out-of-range index must not silently reference the next graph's atoms. + + Regression for the reported bug: in a variable-size batch, adding + batch_ptr[b] to an out-of-range local index wraps into graph b+1's + rows without error. The bounds check must catch this before any + indexing occurs. + """ + # Graph 0: 2 atoms, graph 1: 4 atoms. + # Without the bounds check, local index 3 on graph 0 would silently + # resolve to global row 3, which is atom 1 of graph 1. + data0 = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + ) + data1 = AtomicData( + atomic_numbers=torch.tensor([6, 6, 6, 6], dtype=torch.long), + positions=torch.tensor( + [[10.0, 0.0, 0.0], [11.0, 0.0, 0.0], [12.0, 0.0, 0.0], [13.0, 0.0, 0.0]] + ), + ) + batch = Batch.from_data_list([data0, data1]) + # Local index 3 is valid for graph 1 but out of range for graph 0. + idx = torch.tensor([[0, 3], [0, 1]]) + with pytest.raises(IndexError, match="out of range"): + pair_distance(batch, idx) + + # --- nonperiodic with explicit cell (pbc=False) ---------------------- + + def test_degenerate_cell_with_pbc_false_does_not_raise(self, device: str) -> None: + """cell=zeros + pbc=False must not raise LinAlgError. + + Regression: the old guard ``has_cell and has_pbc`` entered _apply_mic + even when all pbc flags were False, hitting torch.linalg.inv on + whatever cell was present. A zero cell causes LinAlgError there. + """ + positions = torch.tensor([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]]) + cell = torch.zeros(1, 3, 3) # degenerate — not invertible + pbc = torch.tensor([[False, False, False]]) + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + # Must not raise; MIC must be skipped; Euclidean distance = 3 Å. + d = pair_distance(batch, idx) + assert torch.allclose(d, torch.tensor([[3.0]], device=device), atol=1e-5) + + def test_valid_cell_with_pbc_false_uses_euclidean(self, device: str) -> None: + """Valid non-degenerate cell + pbc=False returns plain Euclidean distance.""" + positions = torch.tensor([[0.1, 0.0, 0.0], [9.9, 0.0, 0.0]]) + box = 10.0 + cell = torch.eye(3).unsqueeze(0) * box + pbc = torch.tensor([[False, False, False]]) + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + # MIC would fold to 0.2 Å; Euclidean is 9.8 Å. + assert torch.allclose(d, torch.tensor([[9.8]], device=device), atol=1e-4) + + # --- nonperiodic --- + + def test_nonperiodic_known_value(self, device: str) -> None: + """pair_distance = Euclidean distance for nonperiodic systems.""" + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=torch.tensor([[0.0, 0.0, 0.0], [3.0, 4.0, 0.0]]), + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + assert d.shape == (1, 1) + assert torch.allclose(d, torch.tensor([[5.0]], device=device), atol=1e-5) + + def test_nonperiodic_batch_of_two(self, device: str) -> None: + """Shared atom_indices work correctly across multiple graphs.""" + pos0 = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) + pos1 = torch.tensor([[0.0, 0.0, 0.0], [0.0, 2.0, 0.0]]) + d0_ref = 1.0 + d1_ref = 2.0 + + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), positions=pos0 + ), + AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), positions=pos1 + ), + ] + batch = Batch.from_data_list(data_list).to(device) + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + assert d.shape == (2, 1) + assert torch.allclose(d[0, 0], torch.tensor(d0_ref, device=device), atol=1e-5) + assert torch.allclose(d[1, 0], torch.tensor(d1_ref, device=device), atol=1e-5) + + def test_per_graph_atom_indices(self, device: str) -> None: + """[B, 2] atom_indices select different pairs per graph.""" + data_list = [ + AtomicData( + atomic_numbers=torch.tensor([6, 6, 6], dtype=torch.long), + positions=torch.tensor( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 3.0, 0.0]] + ), + ), + AtomicData( + atomic_numbers=torch.tensor([6, 6, 6], dtype=torch.long), + positions=torch.tensor( + [[0.0, 0.0, 0.0], [0.0, 0.0, 5.0], [2.0, 0.0, 0.0]] + ), + ), + ] + batch = Batch.from_data_list(data_list).to(device) + # graph 0: atoms 0-1 → dist 1; graph 1: atoms 0-2 → dist 2 + idx = torch.tensor([[0, 1], [0, 2]], device=device) + d = pair_distance(batch, idx) + assert d.shape == (2, 1) + assert torch.allclose(d[0, 0], torch.tensor(1.0, device=device), atol=1e-5) + assert torch.allclose(d[1, 0], torch.tensor(2.0, device=device), atol=1e-5) + + # --- cubic periodic --- + + def test_periodic_cubic_mic(self, device: str) -> None: + """MIC selects the nearest image in a cubic cell.""" + box = 10.0 + # Atom 0 at 0.1, atom 1 at 9.9 → naive dist = 9.8, MIC dist = 0.2 + positions = torch.tensor([[0.1, 0.0, 0.0], [9.9, 0.0, 0.0]]) + cell = torch.eye(3).unsqueeze(0) * box # [1, 3, 3] + pbc = torch.tensor([[True, True, True]]) # [1, 3] + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + assert torch.allclose(d, torch.tensor([[0.2]], device=device), atol=1e-4) + + # --- triclinic MIC --- + + def test_triclinic_mic_known_value(self, device: str) -> None: + """MIC distance in triclinic cell: verify against manually computed value.""" + # Cell: a=[4,0,0], b=[1,4,0], c=[0,0,4] — [1,3,3] + cell = torch.tensor([[[4.0, 0.0, 0.0], [1.0, 4.0, 0.0], [0.0, 0.0, 4.0]]]) + pbc = torch.tensor([[True, True, True]]) # [1, 3] + # Atom i at origin, atom j across boundary (Cartesian [3.5, 0, 0]) + # Fractional: j @ cell^{-1}; round; nearest image is [-0.5*a] away + pos_i = torch.tensor([[0.0, 0.0, 0.0]]) + pos_j = torch.tensor([[3.5, 0.0, 0.0]]) + positions = torch.cat([pos_i, pos_j], dim=0) + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + # Naive: 3.5; MIC: |3.5 - 4| = 0.5 (nearest image in a-direction) + assert torch.allclose(d, torch.tensor([[0.5]], device=device), atol=1e-4) + + def test_unreduced_cell_raises(self, device: str) -> None: + """Unreduced cell raises ValueError with a clear message. + + Regression for the reported bug: cell ``[[1,0,0],[10,0.1,0],[0,0,10]]`` + with fractional displacement ``[0,0.49,0]`` requires offset ``[−5,0,0]``, + which lies outside the 27-image search range. The old code returned + ≈ 3.9 Å silently; the new code detects the non-reduced cell and raises. + """ + # Minkowski check: |a1·a2| = 10 > 0.5*min(|a1|²,|a2|²) = 0.5 — fails. + cell = torch.tensor([[[1.0, 0.0, 0.0], [10.0, 0.1, 0.0], [0.0, 0.0, 10.0]]]) + pbc = torch.tensor([[True, True, True]]) + positions = torch.tensor([[0.0, 0.0, 0.0], [0.0, 4.9, 0.049]]) + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + with pytest.raises(ValueError, match="Minkowski"): + pair_distance(batch, idx) + + def test_triclinic_mic_reduced_skewed_cell_27image_correct( + self, device: str + ) -> None: + """27-image search returns the correct MIC for a Minkowski-reduced skewed cell. + + Cell: ``[[2,0,0],[0.8,2,0],[0,0,10]]`` — satisfies Minkowski conditions + (``|a0·a1| = 1.6 ≤ 0.5·min(4, 4.64) = 2.0``). + + Fractional displacement ``[0.49, 0.49, 0]``: + - Componentwise rounding keeps ``[0.49, 0.49, 0]`` → Cartesian ≈ 1.69 Å. + - Correct MIC (offset ``[−1, 0, 0]``) → Cartesian ≈ 1.16 Å. + + Componentwise rounding alone would return the wrong (longer) image; + the 27-image search returns the correct one. + """ + # Verify Minkowski condition holds: |a0·a1| = 1.6 <= 0.5*min(4,4.64) = 2.0 ✓ + cell = torch.tensor([[[2.0, 0.0, 0.0], [0.8, 2.0, 0.0], [0.0, 0.0, 10.0]]]) + pbc = torch.tensor([[True, True, True]]) + + # pos_j: fractional [0.49, 0.49, 0] + # Cartesian = 0.49*[2,0,0] + 0.49*[0.8,2,0] = [1.372, 0.98, 0] + pos_i = torch.tensor([[0.0, 0.0, 0.0]]) + pos_j = torch.tensor([[1.372, 0.98, 0.0]]) + positions = torch.cat([pos_i, pos_j], dim=0) + + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + + # Componentwise rounding gives ≈ 1.687 Å; correct MIC is ≈ 1.164 Å. + naive_dist = torch.linalg.vector_norm(pos_j - pos_i).item() + assert d.item() < naive_dist * 0.8, ( + f"MIC distance {d.item():.4f} Å should be shorter than the naive " + f"distance {naive_dist:.4f} Å — 27-image search may not be working." + ) + assert d.item() < 1.20, ( + f"Expected MIC distance ≈ 1.164 Å, got {d.item():.4f} Å." + ) + + # --- gradients --- + + def test_gradient_nonperiodic(self, device: str) -> None: + """pair_distance gradient w.r.t. positions is correct (finite diff).""" + torch.manual_seed(7) + positions = torch.randn(3, 3, device=device, dtype=torch.float64) + idx = torch.tensor([0, 2], device=device) + + # Use gradcheck with a wrapper that creates a fresh batch + def _fn(pos: Tensor) -> Tensor: + batch_local = Batch.from_data_list( + [ + AtomicData( + atomic_numbers=torch.tensor([6, 6, 6], dtype=torch.long), + positions=pos.detach(), + ) + ] + ).to(device) + batch_local["positions"] = pos # keep grad-tracking leaf + return pair_distance(batch_local, idx) + + pos_double = positions.detach().clone().requires_grad_(True) + torch.autograd.gradcheck(_fn, (pos_double,), eps=1e-4, atol=1e-3, rtol=1e-3) + + def test_gradient_periodic(self, device: str) -> None: + """pair_distance gradient is finite and non-zero for periodic systems.""" + box = 8.0 + positions = torch.tensor([[1.0, 0.0, 0.0], [6.0, 0.0, 0.0]], device=device) + cell = torch.eye(3).unsqueeze(0).to(device) * box # [1, 3, 3] + pbc = torch.tensor([[True, True, True]]) # [1, 3] + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions.cpu(), + cell=cell.cpu(), + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + pos_leaf = batch.positions.detach().requires_grad_(True) + batch["positions"] = pos_leaf + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + d.sum().backward() + assert pos_leaf.grad is not None + assert pos_leaf.grad.isfinite().all() + assert (pos_leaf.grad.abs() > 0).any() + + # --- tests away from half-cell tie --- + + def test_not_at_half_cell_tie(self, device: str) -> None: + """Distance is computed correctly well away from the MIC discontinuity.""" + box = 10.0 + # Position atom j at 3.0 from atom i (clearly not near 5.0 = box/2) + positions = torch.tensor([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]]) + cell = torch.eye(3).unsqueeze(0) * box # [1, 3, 3] + pbc = torch.tensor([[True, True, True]]) # [1, 3] + data = AtomicData( + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + positions=positions, + cell=cell, + pbc=pbc, + ) + batch = Batch.from_data_list([data]).to(device) + idx = torch.tensor([0, 1], device=device) + d = pair_distance(batch, idx) + assert torch.allclose(d, torch.tensor([[3.0]], device=device), atol=1e-5) + + +# =========================================================================== +# 5. aggregate_bias_results +# =========================================================================== + + +class TestAggregateBiasResults: + """Tests for bias aggregation.""" + + def test_empty_list_returns_empty_result(self) -> None: + r = aggregate_bias_results([]) + assert r.energy is None + assert r.forces is None + + def test_single_result_passthrough(self) -> None: + e = torch.tensor([[1.0]]) + f = torch.zeros(3, 3) + r = aggregate_bias_results([BiasResult(energy=e, forces=f)]) + assert torch.allclose(r.energy, e) + assert torch.allclose(r.forces, f) + + def test_energy_summed(self) -> None: + r1 = BiasResult(energy=torch.tensor([[1.0]])) + r2 = BiasResult(energy=torch.tensor([[2.0]])) + agg = aggregate_bias_results([r1, r2]) + assert torch.allclose(agg.energy, torch.tensor([[3.0]])) + + def test_forces_summed(self) -> None: + f1 = torch.ones(4, 3) + f2 = torch.ones(4, 3) * 2.0 + r1 = BiasResult(forces=f1) + r2 = BiasResult(forces=f2) + agg = aggregate_bias_results([r1, r2]) + assert torch.allclose(agg.forces, torch.ones(4, 3) * 3.0) + + def test_none_fields_handled(self) -> None: + r1 = BiasResult(energy=torch.tensor([[1.0]])) + r2 = BiasResult(forces=torch.zeros(2, 3)) + agg = aggregate_bias_results([r1, r2]) + assert agg.energy is not None + assert agg.forces is not None + + def test_virial_summed(self) -> None: + v1 = torch.ones(1, 3, 3) + v2 = torch.ones(1, 3, 3) * 2.0 + r1 = BiasResult(virial=v1) + r2 = BiasResult(virial=v2) + agg = aggregate_bias_results([r1, r2]) + assert torch.allclose(agg.virial, torch.ones(1, 3, 3) * 3.0) + + def test_mixed_stress_and_virial_raises(self) -> None: + """Mixing stress from one result and virial from another raises ValueError. + + The error must come from aggregate_bias_results itself (not from + BiasResult.__post_init__) with a message that identifies which + result indices contributed each field. + """ + r_stress = BiasResult(stress=torch.zeros(1, 3, 3)) + r_virial = BiasResult(virial=torch.zeros(1, 3, 3)) + with pytest.raises(ValueError, match="stress.*virial|virial.*stress"): + aggregate_bias_results([r_stress, r_virial]) + + def test_mixed_stress_and_virial_error_identifies_indices(self) -> None: + """Error message must identify which result indices are responsible.""" + results = [ + BiasResult(energy=torch.zeros(1, 1)), # index 0 — no cell response + BiasResult(stress=torch.zeros(1, 3, 3)), # index 1 — stress + BiasResult(energy=torch.zeros(1, 1)), # index 2 — no cell response + BiasResult(virial=torch.zeros(1, 3, 3)), # index 3 — virial + ] + with pytest.raises(ValueError, match=r"\[1\].*\[3\]|\[3\].*\[1\]"): + aggregate_bias_results(results) + + def test_all_stress_aggregates_correctly(self) -> None: + """Multiple stress contributions are summed without raising.""" + r1 = BiasResult(stress=torch.ones(1, 3, 3)) + r2 = BiasResult(stress=torch.ones(1, 3, 3) * 2.0) + agg = aggregate_bias_results([r1, r2]) + assert agg.stress is not None + assert agg.virial is None + assert torch.allclose(agg.stress, torch.ones(1, 3, 3) * 3.0) + + def test_duplicate_observable_key_raises(self) -> None: + r1 = BiasResult(observables={"bias/a/cv": torch.zeros(1)}) + r2 = BiasResult(observables={"bias/a/cv": torch.ones(1)}) + with pytest.raises(ValueError, match="duplicate observable key"): + aggregate_bias_results([r1, r2]) + + def test_duplicate_observable_error_identifies_indices(self) -> None: + """The message must name both colliding results, not just the key. + + Observables are merged rather than summed, so a collision silently + dropping one bias's diagnostic is exactly what this guards against; + the message has to say which two biases collided. + """ + results = [ + BiasResult(observables={"bias/a/cv": torch.zeros(1)}), # index 0 + BiasResult(energy=torch.zeros(1, 1)), + BiasResult(observables={"bias/b/cv": torch.zeros(1)}), + BiasResult(observables={"bias/a/cv": torch.ones(1)}), # index 3 + ] + with pytest.raises(ValueError) as excinfo: + aggregate_bias_results(results) + message = str(excinfo.value) + assert "results[0]" in message + assert "results[3]" in message + assert "bias/a/cv" in message + + def test_observables_not_summed_with_matching_field_name(self) -> None: + """An observable named 'energy' must not be added to the bias energy. + + observables are merged in a separate namespace from the tensor + fields, so a name collision between the two is not possible. + """ + r1 = BiasResult( + energy=torch.ones(1, 1), + observables={"energy": torch.tensor([100.0])}, + ) + r2 = BiasResult(energy=torch.ones(1, 1)) + agg = aggregate_bias_results([r1, r2]) + assert agg.energy is not None + assert torch.allclose(agg.energy, torch.full((1, 1), 2.0)) + assert torch.allclose(agg.observables["energy"], torch.tensor([100.0])) + + def test_distinct_observable_keys_merged(self) -> None: + r1 = BiasResult(observables={"bias/a/cv": torch.tensor([1.0])}) + r2 = BiasResult(observables={"bias/b/cv": torch.tensor([2.0])}) + agg = aggregate_bias_results([r1, r2]) + assert "bias/a/cv" in agg.observables + assert "bias/b/cv" in agg.observables + + def test_different_registration_orders_same_result(self) -> None: + """Aggregation must be order-independent (commutativity for sum).""" + e1 = torch.tensor([[1.5]]) + e2 = torch.tensor([[0.5]]) + agg_ab = aggregate_bias_results([BiasResult(energy=e1), BiasResult(energy=e2)]) + agg_ba = aggregate_bias_results([BiasResult(energy=e2), BiasResult(energy=e1)]) + assert torch.allclose(agg_ab.energy, agg_ba.energy) + + +# =========================================================================== +# 6. torch.compile — fullgraph and graph-break tests +# =========================================================================== + + +class TestCompile: + """Verifies what can and cannot be compiled with ``torch.compile``. + + * :func:`pair_distance` — compiles with ``fullgraph=True``. + * :func:`aggregate_bias_results` — compiles with ``fullgraph=True`` for + fixed-size input lists. + * ``ConservativeBias.evaluate()`` — does **not** compile with + ``fullgraph=True``. The root cause is + ``pos_leaf = positions.detach().requires_grad_(True)``: + ``torch.compile`` does not support ``.requires_grad_()`` mutation. + **Chosen approach:** compile :meth:`energy` independently; keep + ``evaluate()`` as an eager orchestration wrapper. + ``EnhancedSampling(compile_biases=True)`` will compile each bias's + ``energy()`` override, not ``evaluate()``. + + Tests in this class: + + * ``fullgraph=True`` tests for compile-capable paths (``pair_distance``, + ``aggregate_bias_results``). + * ``fullgraph=False`` tests for ``ConservativeBias.evaluate()`` (allow + graph break; verify correctness and no memory growth). + * ``fullgraph=True`` test for compiling ``energy()`` only. + """ + + @staticmethod + def _compile_kw_full(device: str) -> dict: + """Compile kwargs for fully-compilable paths (fullgraph=True).""" + kw: dict = {"fullgraph": True} + if device == "cuda": + kw["backend"] = "inductor" + return kw + + @staticmethod + def _compile_kw_allow_breaks(device: str) -> dict: + """Compile kwargs allowing graph breaks (for evaluate()).""" + kw: dict = {"fullgraph": False} + if device == "cuda": + kw["backend"] = "inductor" + return kw + + # ------------------------------------------------------------------ + # pair_distance — fully compilable (fullgraph=True) + # ------------------------------------------------------------------ + + def test_pair_distance_compiles_fullgraph(self, device: str) -> None: + """pair_distance compiles with fullgraph=True (no graph breaks).""" + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=3, device=device) + idx = torch.tensor([0, 1], device=device) + + compiled = torch.compile(pair_distance, **self._compile_kw_full(device)) + for _ in range(3): + d = compiled(batch, idx) + assert d.shape == (2, 1) + assert d.isfinite().all() + + def test_pair_distance_compile_agrees_eager(self, device: str) -> None: + """Compiled pair_distance matches eager output within tolerance.""" + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=4, device=device) + idx = torch.tensor([0, 1], device=device) + + d_eager = pair_distance(batch, idx) + compiled = torch.compile(pair_distance, **self._compile_kw_full(device)) + d_compiled = compiled(batch, idx) + assert torch.allclose(d_eager, d_compiled, atol=1e-5) + + def test_pair_distance_periodic_mic_compiles_fullgraph(self, device: str) -> None: + """pair_distance with periodic MIC compiles with fullgraph=True.""" + # Reset dynamo to avoid recompile_limit from previous compile tests + # sharing the pair_distance compiled-function cache. + torch._dynamo.reset() + + batch = _make_cubic_batch(n_graphs=2, atoms_per_graph=3, box=6.0, device=device) + idx = torch.tensor([0, 1], device=device) + + compiled = torch.compile(pair_distance, **self._compile_kw_full(device)) + for _ in range(5): + d = compiled(batch, idx) + assert d.isfinite().all() + + # ------------------------------------------------------------------ + # aggregate_bias_results — fully compilable (fullgraph=True) + # ------------------------------------------------------------------ + + def test_aggregate_compiles_fullgraph(self, device: str) -> None: + """aggregate_bias_results compiles with fullgraph=True.""" + e1 = torch.ones(2, 1, device=device) + e2 = torch.ones(2, 1, device=device) * 2.0 + f1 = torch.ones(8, 3, device=device) + f2 = torch.ones(8, 3, device=device) * 0.5 + + def _agg() -> BiasResult: + return aggregate_bias_results( + [BiasResult(energy=e1, forces=f1), BiasResult(energy=e2, forces=f2)] + ) + + compiled = torch.compile(_agg, **self._compile_kw_full(device)) + result = compiled() + assert result.energy is not None + assert torch.allclose(result.energy, torch.full((2, 1), 3.0, device=device)) + + # ------------------------------------------------------------------ + # ConservativeBias.energy() — compilable when subclassed correctly + # ------------------------------------------------------------------ + + def test_conservative_energy_fn_compiles_fullgraph(self, device: str) -> None: + """ConservativeBias.energy() compiles with fullgraph=True. + + This is the actual compile target when compile_biases=True. + evaluate() stays eager; energy() is compiled per the fallback. + """ + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=4, device=device) + bias = _QuadraticBias(k=1.0) + + # Simulate the runner compiling energy() not evaluate() + compiled_energy = torch.compile(bias.energy, **self._compile_kw_full(device)) + + # Temporarily inject fresh positions leaf (as evaluate() does eagerly) + pos_leaf = batch.positions.detach().requires_grad_(True) + batch["positions"] = pos_leaf + for _ in range(3): + e = compiled_energy(batch) + batch["positions"] = pos_leaf.detach() + assert e.shape == (2, 1) + assert e.isfinite().all() + + # ------------------------------------------------------------------ + # ConservativeBias.evaluate() — runs with graph breaks (fullgraph=False) + # ------------------------------------------------------------------ + + def test_conservative_bias_evaluate_runs_correctly(self, device: str) -> None: + """ConservativeBias.evaluate() produces correct forces (eager mode). + + evaluate() is NOT compiled with fullgraph=True (see spike finding). + It is the eager orchestration wrapper; energy() is what gets compiled. + """ + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=4, device=device) + bias = _QuadraticBias(k=1.0) + result = bias.evaluate(batch) + assert result.forces is not None + assert result.forces.shape == (8, 3) + assert result.forces.isfinite().all() + + def test_conservative_bias_compile_allows_graph_break(self, device: str) -> None: + """ConservativeBias.evaluate() can run under torch.compile(fullgraph=False). + + With fullgraph=False the graph break at requires_grad_() is allowed. + Output agrees with eager. + """ + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=3, device=device) + bias = _QuadraticBias(k=2.0) + + r_eager = bias.evaluate(batch) + compiled = torch.compile(bias.evaluate, **self._compile_kw_allow_breaks(device)) + r_compiled = compiled(batch) + + assert r_eager.energy is not None and r_compiled.energy is not None + assert torch.allclose(r_eager.energy, r_compiled.energy, atol=1e-4) + assert r_eager.forces is not None and r_compiled.forces is not None + assert torch.allclose(r_eager.forces, r_compiled.forces, atol=1e-4) + + def test_no_memory_growth_eager_evaluate_10_calls(self, device: str) -> None: + """Eager evaluate() must not grow GPU memory across 10 calls.""" + batch = _make_nonperiodic_batch(n_graphs=4, atoms_per_graph=8, device=device) + bias = _QuadraticBias(k=1.0) + + # Warm up + for _ in range(3): + bias.evaluate(batch) + + gc.collect() + if device == "cuda": + torch.cuda.synchronize() + torch.cuda.empty_cache() + mem_start = torch.cuda.memory_allocated() + + for _ in range(10): + bias.evaluate(batch) + + if device == "cuda": + torch.cuda.synchronize() + mem_end = torch.cuda.memory_allocated() + assert mem_end - mem_start <= 1 * 1024 * 1024, ( + f"GPU memory grew by {mem_end - mem_start} bytes across 10 evaluate() calls" + ) + + def test_pair_distance_inside_energy_compiles(self, device: str) -> None: + """pair_distance used as CV inside energy() compiles with fullgraph=True.""" + batch = _make_nonperiodic_batch(n_graphs=2, atoms_per_graph=4, device=device) + idx = torch.tensor([0, 1], device=device) + bias = _PairDistanceBias(atom_indices=idx, k=1.0) + + # Compile energy() — the intended compile target + compiled_energy = torch.compile(bias.energy, **self._compile_kw_full(device)) + pos_leaf = batch.positions.detach().requires_grad_(True) + batch["positions"] = pos_leaf + for _ in range(3): + e = compiled_energy(batch) + batch["positions"] = pos_leaf.detach() + assert e.isfinite().all() diff --git a/test/enhanced_sampling/test_builtin_biases.py b/test/enhanced_sampling/test_builtin_biases.py new file mode 100644 index 00000000..55da7033 --- /dev/null +++ b/test/enhanced_sampling/test_builtin_biases.py @@ -0,0 +1,556 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the built-in biases and the periodic CV difference helper. + +Covers :class:`HarmonicUmbrellaBias`, :class:`UpperWall`, :class:`LowerWall`, +:class:`FlatBottomRestraint`, and :func:`periodic_difference`. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.enhanced_sampling import ( + FlatBottomRestraint, + HarmonicUmbrellaBias, + LowerWall, + UpperWall, + pair_distance, + periodic_difference, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _pair_batch(distances: list[float], device: str = "cpu") -> Batch: + """Return one graph per entry, atoms 0 and 1 separated along x.""" + data_list = [ + AtomicData( + positions=torch.tensor([[0.0, 0.0, 0.0], [d, 0.0, 0.0]]), + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + ) + for d in distances + ] + return Batch.from_data_list(data_list).to(device) + + +def _cv(batch: Batch) -> torch.Tensor: + return pair_distance(batch, torch.tensor([0, 1], device=batch.positions.device)) + + +# =========================================================================== +# 1. periodic_difference +# =========================================================================== + + +class TestPeriodicDifference: + """Wrapping CV differences onto a circle.""" + + def test_none_periods_is_plain_difference(self) -> None: + values = torch.tensor([[3.0]]) + centers = torch.tensor([[-3.0]]) + assert torch.allclose( + periodic_difference(values, centers, None), torch.tensor([[6.0]]) + ) + + def test_wraps_the_short_way_round(self) -> None: + """+3.0 and -3.0 rad are 0.283 apart, not 6.0.""" + two_pi = 2 * math.pi + delta = periodic_difference( + torch.tensor([[3.0]]), torch.tensor([[-3.0]]), torch.tensor([two_pi]) + ) + assert abs(float(delta) - (6.0 - two_pi)) < 1e-6 + assert abs(float(delta)) < math.pi + + def test_result_within_half_period(self) -> None: + two_pi = 2 * math.pi + values = torch.linspace(-10, 10, 41).unsqueeze(-1) + centers = torch.zeros_like(values) + delta = periodic_difference(values, centers, torch.tensor([two_pi])) + assert bool((delta.abs() <= math.pi + 1e-6).all()) + + def test_zero_period_marks_non_periodic(self) -> None: + delta = periodic_difference( + torch.tensor([[100.0, 3.0]]), + torch.tensor([[0.0, -3.0]]), + torch.tensor([0.0, 2 * math.pi]), + ) + assert abs(float(delta[0, 0]) - 100.0) < 1e-6 # unwrapped + assert abs(float(delta[0, 1])) < math.pi # wrapped + + def test_non_finite_period_marks_non_periodic(self) -> None: + delta = periodic_difference( + torch.tensor([[100.0]]), + torch.tensor([[0.0]]), + torch.tensor([float("inf")]), + ) + assert torch.isfinite(delta).all() + assert abs(float(delta) - 100.0) < 1e-6 + + def test_gradient_matches_unwrapped(self) -> None: + """round() has zero gradient, so wrapping must not change d(delta)/d(x).""" + values = torch.tensor([[3.0]], requires_grad=True) + delta = periodic_difference( + values, torch.tensor([[-3.0]]), torch.tensor([2 * math.pi]) + ) + (grad,) = torch.autograd.grad(delta.sum(), values) + assert torch.allclose(grad, torch.ones_like(grad)) + + +# =========================================================================== +# 2. HarmonicUmbrellaBias +# =========================================================================== + + +class TestHarmonicUmbrellaBias: + """Analytical energies, per-window selection, and validation.""" + + def test_energy_matches_analytical(self, device: str) -> None: + batch = _pair_batch([3.0], device=device) + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([2.0]), stiffness=4.0, name="u" + ) + result = bias.evaluate(batch) + # 0.5 * 4 * (3 - 2)^2 = 2.0 + assert abs(float(result.energy) - 2.0) < 1e-5 + + def test_zero_energy_at_center(self, device: str) -> None: + batch = _pair_batch([2.0], device=device) + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([2.0]), stiffness=7.0, name="u" + ) + assert abs(float(bias.evaluate(batch).energy)) < 1e-6 + + def test_force_pulls_toward_center(self, device: str) -> None: + """A restraint must shorten a too-long distance, not lengthen it.""" + batch = _pair_batch([3.0], device=device) + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([2.0]), stiffness=4.0, name="u" + ) + forces = bias.evaluate(batch).forces + # Atom 1 sits at +x of atom 0 and is too far: its force must point -x. + assert float(forces[1, 0]) < 0 + assert float(forces[0, 0]) > 0 + + def test_per_window_centers_selected_by_state_id(self, device: str) -> None: + batch = _pair_batch([2.0, 2.0, 2.0], device=device) + batch["thermodynamic_state_id"] = torch.tensor([0, 1, 2], device=device) + bias = HarmonicUmbrellaBias( + cv=_cv, + centers=torch.tensor([[2.0], [3.0], [4.0]]), + stiffness=2.0, + name="u", + ) + energy = bias.evaluate(batch).energy.reshape(-1) + # distances all 2.0; deltas are 0, -1, -2 + expected = torch.tensor([0.0, 1.0, 4.0], device=energy.device) + assert torch.allclose(energy, expected, atol=1e-5) + + def test_defaults_to_state_zero_without_field(self, device: str) -> None: + batch = _pair_batch([2.0], device=device) + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([[5.0], [9.0]]), stiffness=1.0, name="u" + ) + # Uses centers[0] = 5.0 -> 0.5 * 1 * (2-5)^2 = 4.5 + assert abs(float(bias.evaluate(batch).energy) - 4.5) < 1e-5 + + def test_out_of_range_state_id_raises(self, device: str) -> None: + """A multi-window bias does index the field, so it must be in range.""" + batch = _pair_batch([2.0], device=device) + batch["thermodynamic_state_id"] = torch.tensor([5], device=device) + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([[2.0], [3.0]]), stiffness=1.0, name="u" + ) + with pytest.raises(IndexError, match="out of range"): + bias.evaluate(batch) + + def test_single_window_ignores_the_state_id(self, device: str) -> None: + """One window selects nothing, so the field is not an index into it. + + This is what lets one shared restraint run alongside a multi-rung + temperature ladder, where `thermodynamic_state_id` addresses the + ladder and has nothing to do with windows. + """ + batch = _pair_batch([3.0], device=device) + batch["thermodynamic_state_id"] = torch.tensor([7], device=device) + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([[2.0]]), stiffness=4.0, name="u" + ) + # 0.5 * 4 * (3 - 2)^2 = 2.0, using the single window regardless of id. + assert abs(float(bias.evaluate(batch).energy) - 2.0) < 1e-5 + + @pytest.mark.parametrize( + "stiffness,shape", + [(3.0, (1, 1, 1)), (torch.tensor([3.0]), (1, 1, 1))], + ) + def test_stiffness_forms_expand(self, stiffness, shape) -> None: + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([2.0]), stiffness=stiffness, name="u" + ) + assert bias.stiffness.shape == shape + + def test_full_matrix_stiffness(self) -> None: + k = torch.tensor([[2.0, 0.5], [0.5, 3.0]]) + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.zeros(2), stiffness=k, name="u" + ) + assert torch.allclose(bias.stiffness[0], k) + + def test_asymmetric_stiffness_raises(self) -> None: + k = torch.tensor([[1.0, 2.0], [0.0, 1.0]]) + with pytest.raises(ValueError, match="must be symmetric"): + HarmonicUmbrellaBias(cv=_cv, centers=torch.zeros(2), stiffness=k, name="u") + + def test_negative_eigenvalue_raises(self) -> None: + """A negative eigenvalue turns the restraint into a runaway repulsion.""" + k = torch.tensor([[1.0, 0.0], [0.0, -1.0]]) + with pytest.raises(ValueError, match="positive-semidefinite"): + HarmonicUmbrellaBias(cv=_cv, centers=torch.zeros(2), stiffness=k, name="u") + + def test_mismatched_diagonal_stiffness_raises(self) -> None: + with pytest.raises(ValueError, match="match the CV dimension"): + HarmonicUmbrellaBias( + cv=_cv, centers=torch.zeros(2), stiffness=torch.ones(3), name="u" + ) + + def test_mismatched_periods_raises(self) -> None: + with pytest.raises(ValueError, match="match the CV dimension"): + HarmonicUmbrellaBias( + cv=_cv, + centers=torch.zeros(2), + stiffness=1.0, + periods=torch.ones(5), + name="u", + ) + + def test_bad_centers_rank_raises(self) -> None: + with pytest.raises(ValueError, match=r"centers must be \[D\] or \[S, D\]"): + HarmonicUmbrellaBias( + cv=_cv, centers=torch.zeros(2, 2, 2), stiffness=1.0, name="u" + ) + + def test_buffers_move_with_module(self) -> None: + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([2.0]), stiffness=1.0, name="u" + ) + assert "centers" in dict(bias.named_buffers()) + assert "stiffness" in dict(bias.named_buffers()) + + def test_state_dict_round_trip(self) -> None: + bias = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([2.0]), stiffness=3.0, name="u" + ) + other = HarmonicUmbrellaBias( + cv=_cv, centers=torch.tensor([9.0]), stiffness=1.0, name="u" + ) + other.load_state_dict(bias.state_dict()) + assert torch.allclose(other.centers, bias.centers) + assert torch.allclose(other.stiffness, bias.stiffness) + + +# =========================================================================== +# 3. Walls +# =========================================================================== + + +class TestWalls: + """One- and two-sided CV penalties.""" + + def test_upper_wall_zero_inside(self, device: str) -> None: + batch = _pair_batch([2.0], device=device) + wall = UpperWall(cv=_cv, threshold=5.0, stiffness=10.0) + assert abs(float(wall.evaluate(batch).energy)) < 1e-8 + + def test_upper_wall_penalises_outside(self, device: str) -> None: + batch = _pair_batch([7.0], device=device) + wall = UpperWall(cv=_cv, threshold=5.0, stiffness=10.0) + # (10/2) * (7-5)^2 = 20 + assert abs(float(wall.evaluate(batch).energy) - 20.0) < 1e-4 + + def test_upper_wall_pushes_inward(self, device: str) -> None: + batch = _pair_batch([7.0], device=device) + forces = UpperWall(cv=_cv, threshold=5.0, stiffness=10.0).evaluate(batch).forces + assert float(forces[1, 0]) < 0, "upper wall must pull the pair closer" + + def test_lower_wall_zero_outside(self, device: str) -> None: + batch = _pair_batch([7.0], device=device) + wall = LowerWall(cv=_cv, threshold=5.0, stiffness=10.0) + assert abs(float(wall.evaluate(batch).energy)) < 1e-8 + + def test_lower_wall_pushes_outward(self, device: str) -> None: + batch = _pair_batch([2.0], device=device) + forces = LowerWall(cv=_cv, threshold=5.0, stiffness=10.0).evaluate(batch).forces + assert float(forces[1, 0]) > 0, "lower wall must push the pair apart" + + def test_wall_inside_gives_zero_forces_not_an_error(self, device: str) -> None: + """The clamp keeps the graph connected where the wall is inactive. + + A wall written as ``if inside: return zeros(B, 1)`` would produce an + energy with no grad_fn, which autograd rejects outright. + """ + batch = _pair_batch([2.0], device=device) + result = UpperWall(cv=_cv, threshold=5.0).evaluate(batch) + assert result.forces is not None + assert torch.count_nonzero(result.forces) == 0 + + def test_force_continuous_across_boundary(self, device: str) -> None: + """Quadratic walls have zero force at the wall; no impulse.""" + wall = UpperWall(cv=_cv, threshold=5.0, stiffness=10.0, exponent=2.0) + just_inside = wall.evaluate(_pair_batch([4.999], device=device)).forces + just_outside = wall.evaluate(_pair_batch([5.001], device=device)).forces + assert torch.allclose(just_inside, just_outside, atol=1e-2) + + def test_exponent_below_one_raises(self) -> None: + with pytest.raises(ValueError, match="exponent must be >= 1"): + UpperWall(cv=_cv, threshold=1.0, exponent=0.5) + + def test_negative_stiffness_raises(self) -> None: + with pytest.raises(ValueError, match="stiffness must be non-negative"): + LowerWall(cv=_cv, threshold=1.0, stiffness=-5.0) + + def test_flat_bottom_zero_inside(self, device: str) -> None: + batch = _pair_batch([3.0], device=device) + restraint = FlatBottomRestraint(cv=_cv, lower=2.0, upper=5.0, stiffness=10.0) + assert abs(float(restraint.evaluate(batch).energy)) < 1e-8 + + def test_flat_bottom_penalises_both_sides(self, device: str) -> None: + restraint = FlatBottomRestraint(cv=_cv, lower=2.0, upper=5.0, stiffness=10.0) + below = restraint.evaluate(_pair_batch([1.0])).energy + above = restraint.evaluate(_pair_batch([6.0])).energy + assert abs(float(below) - 5.0) < 1e-4 # (10/2)*(2-1)^2 + assert abs(float(above) - 5.0) < 1e-4 # (10/2)*(6-5)^2 + + def test_flat_bottom_matches_two_walls(self, device: str) -> None: + batch = _pair_batch([6.5], device=device) + combined = FlatBottomRestraint( + cv=_cv, lower=2.0, upper=5.0, stiffness=3.0 + ).evaluate(batch) + upper = UpperWall(cv=_cv, threshold=5.0, stiffness=3.0).evaluate(batch) + lower = LowerWall(cv=_cv, threshold=2.0, stiffness=3.0).evaluate(batch) + assert abs(float(combined.energy) - float(upper.energy + lower.energy)) < 1e-5 + + def test_inverted_bounds_raise(self) -> None: + with pytest.raises(ValueError, match="strictly below"): + FlatBottomRestraint(cv=_cv, lower=5.0, upper=2.0) + + def test_wall_finite_difference_force(self, device: str) -> None: + """Autograd wall forces agree with central differences.""" + eps = 1e-4 + wall = UpperWall(cv=_cv, threshold=3.0, stiffness=6.0) + batch = _pair_batch([4.0], device=device) + analytic = wall.evaluate(batch).forces[1, 0] + + plus = wall.evaluate(_pair_batch([4.0 + eps], device=device)).energy + minus = wall.evaluate(_pair_batch([4.0 - eps], device=device)).energy + numeric = -(float(plus) - float(minus)) / (2 * eps) + assert abs(float(analytic) - numeric) < 1e-2 + + +# =========================================================================== +# 4. torch.compile on the built-ins +# =========================================================================== + + +class TestBuiltinBiasCompile: + """``compile_biases=True`` hands ``energy()`` to ``torch.compile``. + + That makes every built-in's ``energy()`` a compiled path in practice, so + a data-dependent Python branch there is a real defect rather than a + stylistic one — it breaks ``fullgraph=True`` outright. These are + regressions against that, per built-in. + """ + + @staticmethod + def _state_batch(device: str) -> Batch: + """Two graphs at distance 3.0, in windows 0 and 1.""" + batch = _pair_batch([3.0, 3.0], device=device) + batch["thermodynamic_state_id"] = torch.tensor([0, 1], device=device) + return batch + + @staticmethod + def _umbrella(device: str = "cpu") -> HarmonicUmbrellaBias: + """Calling energy() directly bypasses evaluate()'s device alignment.""" + return HarmonicUmbrellaBias( + cv=_cv, + centers=torch.tensor([[2.0], [3.0]]), + stiffness=4.0, + name="u", + ).to(device) + + def test_umbrella_energy_compiles_fullgraph_with_state_ids( + self, device: str + ) -> None: + """Per-state selection must not introduce a data-dependent branch.""" + torch._dynamo.reset() + batch = self._state_batch(device) + bias = self._umbrella(device) + compiled = torch.compile(bias.energy, fullgraph=True) + energy = compiled(batch) + assert energy.shape == (2, 1) + + def test_umbrella_compiled_matches_eager(self, device: str) -> None: + torch._dynamo.reset() + batch = self._state_batch(device) + bias = self._umbrella(device) + eager = bias.energy(batch) + compiled = torch.compile(bias.energy, fullgraph=True)(batch) + assert torch.allclose(eager, compiled, atol=1e-6) + # window 0: 0.5*4*(3-2)^2 = 2.0 ; window 1 sits at its center + assert torch.allclose( + eager.flatten(), torch.tensor([2.0, 0.0], device=eager.device), atol=1e-5 + ) + + def test_umbrella_compiles_without_state_ids(self, device: str) -> None: + torch._dynamo.reset() + bias = self._umbrella(device) + compiled = torch.compile(bias.energy, fullgraph=True) + assert compiled(_pair_batch([3.0], device=device)).shape == (1, 1) + + def test_state_id_validation_survives_compilation(self, device: str) -> None: + """Moving the check to evaluate() must not have removed it. + + The bounds check cannot live in energy(), but hoisting it to the eager + evaluate() means it still runs when energy() is compiled — strictly + better than the eager-only guards elsewhere, which skip under compile. + """ + torch._dynamo.reset() + batch = _pair_batch([3.0], device=device) + batch["thermodynamic_state_id"] = torch.tensor([7], device=device) + bias = self._umbrella(device) + bias.energy = torch.compile(bias.energy, fullgraph=True) + with pytest.raises(IndexError, match="out of range"): + bias.evaluate(batch) + + def test_state_id_error_names_graph_and_valid_range(self, device: str) -> None: + batch = _pair_batch([3.0, 3.0], device=device) + batch["thermodynamic_state_id"] = torch.tensor([0, 9], device=device) + with pytest.raises(IndexError) as excinfo: + self._umbrella().evaluate(batch) + message = str(excinfo.value) + assert "[1]" in message # the offending graph + assert "0..1" in message # the valid range + + @pytest.mark.parametrize("wall_factory", ["upper", "lower", "flat"]) + def test_wall_energy_compiles_fullgraph( + self, wall_factory: str, device: str + ) -> None: + torch._dynamo.reset() + walls = { + "upper": lambda: UpperWall(cv=_cv, threshold=2.0, stiffness=6.0), + "lower": lambda: LowerWall(cv=_cv, threshold=4.0, stiffness=6.0), + "flat": lambda: FlatBottomRestraint( + cv=_cv, lower=2.0, upper=4.0, stiffness=6.0 + ), + } + bias = walls[wall_factory]().to(device) + batch = _pair_batch([3.0, 5.0], device=device) + eager = bias.energy(batch) + compiled = torch.compile(bias.energy, fullgraph=True)(batch) + assert torch.allclose(eager, compiled, atol=1e-6) + + def test_runner_compile_biases_end_to_end(self, device: str) -> None: + """The path a user actually takes: compile_biases=True on the runner.""" + from nvalchemi.dynamics import NVTLangevin + from nvalchemi.enhanced_sampling import EnhancedSampling + from nvalchemi.models.demo import DemoModel, DemoModelWrapper + + torch._dynamo.reset() + + def make_batch() -> Batch: + data_list = [] + for _ in range(2): + data = AtomicData( + positions=torch.tensor([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]]), + atomic_numbers=torch.tensor([6, 6], dtype=torch.long), + atomic_masses=torch.ones(2), + forces=torch.zeros(2, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(2, 3)) + data_list.append(data) + batch = Batch.from_data_list(data_list).to(device) + batch["thermodynamic_state_id"] = torch.tensor([0, 1], device=device) + return batch + + results = [] + for compile_biases in (False, True): + batch = make_batch() + # DemoModel has random weights; seed so the physical contribution + # is identical and any difference is attributable to the bias. + torch.manual_seed(0) + model = DemoModelWrapper(DemoModel()).to(device) + dynamics = NVTLangevin(model=model, dt=0.1, temperature=300.0, friction=0.1) + runner = EnhancedSampling( + dynamics, {"u": self._umbrella()}, compile_biases=compile_biases + ) + runner.prime_forces(batch) + results.append(batch.forces.clone()) + assert torch.allclose(results[0], results[1], atol=1e-5) + + +# =========================================================================== +# 5. Device transparency +# =========================================================================== + + +class TestDeviceTransparency: + """A CV or bias built on CPU must work against a batch on GPU. + + Both are configuration built once, typically before the batch is moved + to the device. Requiring the user to place them by hand produces a bare + "expected all tensors to be on the same device" that names neither the + CV nor the fix. + """ + + def test_cpu_atom_indices_against_device_batch(self, device: str) -> None: + """The exact shape of the quick-start snippet: no device= on the pair.""" + batch = _pair_batch([3.0], device=device) + pair = torch.tensor([0, 1]) # deliberately CPU + distance = pair_distance(batch, pair) + assert distance.device.type == batch.positions.device.type + assert abs(float(distance) - 3.0) < 1e-5 + + def test_cpu_atom_indices_per_graph_form(self, device: str) -> None: + batch = _pair_batch([3.0, 4.0], device=device) + pairs = torch.tensor([[0, 1], [1, 0]]) # CPU, [B, 2] + distance = pair_distance(batch, pairs).flatten() + assert torch.allclose( + distance, torch.tensor([3.0, 4.0], device=distance.device), atol=1e-5 + ) + + def test_cpu_built_umbrella_against_device_batch(self, device: str) -> None: + """A bias built before the batch moved to GPU still evaluates.""" + batch = _pair_batch([3.0], device=device) + pair = torch.tensor([0, 1]) # CPU + bias = HarmonicUmbrellaBias( + cv=lambda b: pair_distance(b, pair), + centers=torch.tensor([2.0]), + stiffness=4.0, + name="u", + ) # no .to(device) + result = bias.evaluate(batch) + assert result.energy.device.type == batch.positions.device.type + assert abs(float(result.energy) - 2.0) < 1e-5 + + def test_cpu_built_wall_against_device_batch(self, device: str) -> None: + batch = _pair_batch([7.0], device=device) + wall = UpperWall(cv=_cv, threshold=5.0, stiffness=10.0) # no .to(device) + assert abs(float(wall.evaluate(batch).energy) - 20.0) < 1e-4 diff --git a/test/enhanced_sampling/test_checkpoint.py b/test/enhanced_sampling/test_checkpoint.py new file mode 100644 index 00000000..e739cb3d --- /dev/null +++ b/test/enhanced_sampling/test_checkpoint.py @@ -0,0 +1,988 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for transactional Zarr checkpointing of enhanced sampling. + +Covers the state encoder, the manifest-gated commit, checksum verification, +``BaseDynamics`` state round-trips, thermodynamic-state rebinding, and exact +trajectory reproduction across a checkpoint/restore boundary. +""" + +from __future__ import annotations + +import math + +import pytest +import torch +import zarr +from torch import Tensor + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin, NVTNoseHoover +from nvalchemi.enhanced_sampling import ( + AdaptivePotentialMixin, + BiasResult, + ConservativeBias, + EnhancedSampling, + HarmonicUmbrellaBias, + pair_distance, +) +from nvalchemi.enhanced_sampling._checkpoint import ( + CHECKPOINT_FORMAT_VERSION, + _component_checksum, + _decode_state, + _encode_state, + read_checkpoint, +) +from nvalchemi.models.demo import DemoModel, DemoModelWrapper + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_batch( + n_graphs: int = 2, atoms_per_graph: int = 4, device: str = "cpu", seed: int = 7 +) -> Batch: + """Return a batch with output buffers and velocities.""" + torch.manual_seed(seed) + data_list = [] + for _ in range(n_graphs): + data = AtomicData( + positions=torch.randn(atoms_per_graph, 3), + atomic_numbers=torch.full((atoms_per_graph,), 6, dtype=torch.long), + atomic_masses=torch.ones(atoms_per_graph), + forces=torch.zeros(atoms_per_graph, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(atoms_per_graph, 3)) + data_list.append(data) + batch = Batch.from_data_list(data_list).to(device) + batch["thermodynamic_state_id"] = torch.arange(n_graphs, device=device) + return batch + + +def _make_dynamics(device: str = "cpu", seed: int = 0) -> NVTLangevin: + torch.manual_seed(seed) + model = DemoModelWrapper(DemoModel()).to(device) + return NVTLangevin(model=model, dt=0.1, temperature=300.0, friction=0.1) + + +def _make_runner( + device: str = "cpu", + steps_per_epoch: int = 4, + seed: int = 0, + n_states: int = 2, +) -> EnhancedSampling: + idx = torch.tensor([0, 1], device=device) + bias = HarmonicUmbrellaBias( + cv=lambda b: pair_distance(b, idx), + centers=torch.arange(2.0, 2.0 + n_states).reshape(n_states, 1), + stiffness=4.0, + name="u", + ) + return EnhancedSampling( + _make_dynamics(device, seed), {"u": bias}, steps_per_epoch=steps_per_epoch + ) + + +class _CountingBias(AdaptivePotentialMixin, ConservativeBias): + """Adaptive bias with a scalar history worth round-tripping.""" + + def __init__(self, name: str = "counter") -> None: + super().__init__(name=name) + self.register_buffer("deposits", torch.zeros(1)) + + def energy(self, current: Batch) -> Tensor: + return ( + torch.zeros(current.num_graphs, 1, device=current.positions.device) + + 0.0 * current.positions.sum() + ) + + def update(self, frames: Batch, result: BiasResult) -> None: + self.deposits += 1 + self.bump_state_version() + + +class _QuietBias(AdaptivePotentialMixin, ConservativeBias): + """Bumps its version only when told to. + + Lets a test separate "the bias actually changed" from "the runner thinks + it changed", which is the distinction a stale seen-version cache blurs. + """ + + def __init__(self, name: str = "quiet", bump: bool = True) -> None: + super().__init__(name=name) + self._bump = bump + + def energy(self, current: Batch) -> Tensor: + return ( + torch.zeros(current.num_graphs, 1, device=current.positions.device) + + 0.0 * current.positions.sum() + ) + + def update(self, frames: Batch, result: BiasResult) -> None: + if self._bump: + self.bump_state_version() + + +class _SharedHistoryBias(AdaptivePotentialMixin, ConservativeBias): + """Deposits accumulate as pending; commit_epoch merges them. + + Models the shared-history multi-walker case: the published state only + becomes correct once the epoch commit has run, so a checkpoint taken + before it records a bias mid-merge. + """ + + def __init__(self, name: str = "shared") -> None: + super().__init__(name=name) + self.register_buffer("pending", torch.zeros(1)) + self.register_buffer("published", torch.zeros(1)) + self.commit_calls = 0 + + def energy(self, current: Batch) -> Tensor: + return ( + torch.zeros(current.num_graphs, 1, device=current.positions.device) + + 0.0 * current.positions.sum() + ) + + def update(self, frames: Batch, result: BiasResult) -> None: + self.pending += 1 + + def commit_epoch(self) -> None: + self.commit_calls += 1 + self.published += self.pending + self.pending.zero_() + + +# =========================================================================== +# 1. State encoding +# =========================================================================== + + +class TestStateEncoding: + """Nested state survives the Zarr round-trip without pickle.""" + + def test_tensors_scalars_and_nesting(self, tmp_path) -> None: + state = { + "counter": 7, + "label": "umbrella", + "ratio": 0.25, + "flag": True, + "nothing": None, + "listy": [1, 2, 3], + "weights": torch.arange(6, dtype=torch.float64).reshape(2, 3), + "ids": torch.tensor([4, 5], dtype=torch.int64), + "nested": {"inner": torch.ones(2), "depth": 2}, + } + group = zarr.open_group(str(tmp_path / "s.zarr"), mode="w") + _encode_state(group, state) + restored = _decode_state(group, "cpu") + + assert restored["counter"] == 7 + assert restored["label"] == "umbrella" + assert restored["flag"] is True + assert restored["nothing"] is None + assert restored["listy"] == [1, 2, 3] + assert torch.equal(restored["weights"], state["weights"]) + assert restored["weights"].dtype == torch.float64 + assert restored["ids"].dtype == torch.int64 + assert torch.equal(restored["nested"]["inner"], state["nested"]["inner"]) + assert restored["nested"]["depth"] == 2 + + def test_empty_tensor_round_trips(self, tmp_path) -> None: + group = zarr.open_group(str(tmp_path / "s.zarr"), mode="w") + _encode_state(group, {"empty": torch.zeros(0, 3)}) + restored = _decode_state(group, "cpu") + assert restored["empty"].shape == (0, 3) + + def test_zero_dimensional_tensor_round_trips(self, tmp_path) -> None: + """Zarr stores a 0-d array as shape (1,); the rank must be restored. + + Scalar buffers are how a compile-safe bias holds its counters — a + Python int would be a data-dependent value in the traced graph. If + the rank comes back wrong the component no longer matches the digest + taken when it was written, and restore fails its own checksum. + """ + state = {"count": torch.tensor(5, dtype=torch.int64)} + group = zarr.open_group(str(tmp_path / "s.zarr"), mode="w") + _encode_state(group, state) + restored = _decode_state(group, "cpu") + + assert restored["count"].shape == () + assert torch.equal(restored["count"], state["count"]) + assert _component_checksum(restored) == _component_checksum(state) + + def test_checksum_distinguishes_rank(self) -> None: + """A scalar and a one-element vector are not the same state.""" + assert _component_checksum({"x": torch.tensor(5)}) != _component_checksum( + {"x": torch.tensor([5])} + ) + + def test_unsupported_type_raises_rather_than_pickling(self, tmp_path) -> None: + """Refusing is the point: a pickle payload would make a checkpoint + executable and unreadable outside Python.""" + group = zarr.open_group(str(tmp_path / "s.zarr"), mode="w") + with pytest.raises(TypeError, match="no pickle payloads"): + _encode_state(group, {"bad": object()}) + + def test_checksum_is_order_independent(self) -> None: + a = {"x": torch.ones(3), "y": 2} + b = {"y": 2, "x": torch.ones(3)} + assert _component_checksum(a) == _component_checksum(b) + + def test_checksum_detects_value_change(self) -> None: + base = _component_checksum({"x": torch.ones(3)}) + assert base != _component_checksum({"x": torch.zeros(3)}) + assert base != _component_checksum({"x": torch.ones(3) * 2}) + + def test_checksum_detects_dtype_change(self) -> None: + assert _component_checksum({"x": torch.ones(3)}) != _component_checksum( + {"x": torch.ones(3, dtype=torch.float64)} + ) + + +# =========================================================================== +# 2. BaseDynamics state +# =========================================================================== + + +class TestDynamicsState: + """state_dict / load_state_dict on the integrator.""" + + def test_round_trip_restores_counters_and_state(self, device: str) -> None: + batch = _make_batch(device=device) + dynamics = _make_dynamics(device) + dynamics._ensure_state_initialized(batch) + dynamics.step_count = 13 + + saved = dynamics.state_dict() + assert saved["step_count"] == 13 + assert saved["random_seed"] == 42 + assert "temperature" in saved["state"] + + other = _make_dynamics(device) + other._ensure_state_initialized(batch) + other.load_state_dict(saved) + assert other.step_count == 13 + assert torch.allclose(other._state.temperature, dynamics._state.temperature) + + def test_load_into_uninitialised_integrator_raises(self, device: str) -> None: + """Restoring into an uninitialised integrator would silently diverge.""" + batch = _make_batch(device=device) + source = _make_dynamics(device) + source._ensure_state_initialized(batch) + target = _make_dynamics(device) + with pytest.raises(RuntimeError, match="has not initialised its own"): + target.load_state_dict(source.state_dict()) + + def test_unknown_key_raises(self, device: str) -> None: + batch = _make_batch(device=device) + dynamics = _make_dynamics(device) + dynamics._ensure_state_initialized(batch) + saved = dynamics.state_dict() + saved["state"]["not_a_real_key"] = torch.zeros(2) + with pytest.raises(KeyError, match="no counterpart"): + dynamics.load_state_dict(saved) + + def test_langevin_noise_is_counter_based(self, device: str) -> None: + """Exact restart relies on this: no generator state to serialise.""" + dynamics = _make_dynamics(device) + assert dynamics._random_seed == 42 + state = dynamics.state_dict() + assert "random_seed" in state and "step_count" in state + + def test_redistribute_state_permutes_rows(self, device: str) -> None: + batch = _make_batch(n_graphs=3, device=device) + dynamics = _make_dynamics(device) + dynamics._ensure_state_initialized(batch) + with torch.no_grad(): + dynamics._state.temperature.copy_( + torch.tensor([1.0, 2.0, 3.0], device=device).reshape( + dynamics._state.temperature.shape + ) + ) + dynamics.redistribute_state(torch.tensor([2, 0, 1], device=device)) + assert dynamics._state.temperature.reshape(-1).tolist() == [3.0, 1.0, 2.0] + + +# =========================================================================== +# 3. Thermodynamic-state rebinding +# =========================================================================== + + +class TestThermodynamicStateRebinding: + """The adapters replica exchange will need in PR 5.""" + + def test_base_dynamics_refuses(self, device: str) -> None: + """An integrator that cannot rebind must fail, not accept silently.""" + from nvalchemi.dynamics.base import BaseDynamics + + dynamics = BaseDynamics(DemoModelWrapper(DemoModel()).to(device)) + with pytest.raises(NotImplementedError, match="does not support"): + dynamics.apply_thermodynamic_state(torch.tensor([0]), torch.tensor([300.0])) + + def test_langevin_rebinds_temperature(self, device: str) -> None: + batch = _make_batch(device=device) + dynamics = _make_dynamics(device) + dynamics._ensure_state_initialized(batch) + before = dynamics._state.temperature.reshape(-1).clone() + + dynamics.apply_thermodynamic_state( + torch.tensor([1, 0], device=device), + torch.tensor([300.0, 600.0], device=device), + ) + after = dynamics._state.temperature.reshape(-1) + assert abs(float(after[0] / before[0]) - 2.0) < 1e-5 + assert abs(float(after[1] / before[1]) - 1.0) < 1e-5 + + def test_langevin_velocity_scaling_follows_temperature(self, device: str) -> None: + """The swap is indivisible: target and velocities move together.""" + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + batch.velocities.fill_(1.0) + dynamics = _make_dynamics(device) + dynamics._ensure_state_initialized(batch) + dynamics.apply_thermodynamic_state( + torch.tensor([0], device=device), torch.tensor([1200.0], device=device) + ) + dynamics.rescale_velocities_for_state(batch) + # T: 300 -> 1200, so v scales by sqrt(4) = 2. + assert torch.allclose( + batch.velocities, torch.full_like(batch.velocities, 2.0), atol=1e-5 + ) + + def test_out_of_range_state_id_raises(self, device: str) -> None: + batch = _make_batch(device=device) + dynamics = _make_dynamics(device) + dynamics._ensure_state_initialized(batch) + with pytest.raises(IndexError, match="out of range"): + dynamics.apply_thermodynamic_state( + torch.tensor([0, 5], device=device), + torch.tensor([300.0], device=device), + ) + + def test_nose_hoover_transforms_chain_state(self, device: str) -> None: + """Q and eta_dot must move with kT or detailed balance breaks.""" + batch = _make_batch(device=device) + model = DemoModelWrapper(DemoModel()).to(device) + dynamics = NVTNoseHoover( + model=model, dt=0.1, temperature=300.0, thermostat_time=10.0 + ) + dynamics._ensure_state_initialized(batch) + with torch.no_grad(): + dynamics._state.nhc_eta_dot.fill_(2.0) + q_before = dynamics._state.nhc_Q.clone() + eta_dot_before = dynamics._state.nhc_eta_dot.clone() + + dynamics.apply_thermodynamic_state( + torch.tensor([0, 0], device=device), + torch.tensor([1200.0], device=device), + ) + ratio = 4.0 # 300 -> 1200 + assert torch.allclose(dynamics._state.nhc_Q, q_before * ratio, rtol=1e-5), ( + "chain masses must scale with kT" + ) + assert torch.allclose( + dynamics._state.nhc_eta_dot, + eta_dot_before / math.sqrt(ratio), + rtol=1e-5, + ), "chain velocities must scale as 1/sqrt(kT)" + + def test_nose_hoover_chain_kinetic_energy_invariant(self, device: str) -> None: + """Q eta_dot^2 must not change: rebinding injects no thermostat energy.""" + batch = _make_batch(device=device) + dynamics = NVTNoseHoover( + model=DemoModelWrapper(DemoModel()).to(device), + dt=0.1, + temperature=300.0, + thermostat_time=10.0, + ) + dynamics._ensure_state_initialized(batch) + with torch.no_grad(): + dynamics._state.nhc_eta_dot.fill_(1.5) + before = (dynamics._state.nhc_Q * dynamics._state.nhc_eta_dot**2).sum() + + dynamics.apply_thermodynamic_state( + torch.tensor([0, 0], device=device), + torch.tensor([900.0], device=device), + ) + after = (dynamics._state.nhc_Q * dynamics._state.nhc_eta_dot**2).sum() + assert torch.allclose(before, after, rtol=1e-5) + + +# =========================================================================== +# 4. Transactional checkpoint +# =========================================================================== + + +class TestCheckpointTransactionality: + """The manifest is the commit marker; checksums catch later damage.""" + + def test_round_trip(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device) + batch = runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + restored_batch, states, manifest = read_checkpoint(path, device) + assert manifest.format_version == CHECKPOINT_FORMAT_VERSION + assert manifest.sampling_step == 4 + assert manifest.sampling_epoch == 1 + assert manifest.num_graphs == 2 + assert "dynamics" in states + assert "biases/u" in states + assert restored_batch.num_graphs == 2 + + def test_store_without_manifest_is_refused(self, tmp_path, device: str) -> None: + """An interrupted write leaves no manifest; restoring it must fail.""" + from nvalchemi.data.datapipes.backends.zarr import AtomicDataZarrWriter + + path = tmp_path / "torn.zarr" + AtomicDataZarrWriter(str(path)).write(_make_batch(device=device)) + with pytest.raises(ValueError, match="no committed manifest"): + read_checkpoint(path, device) + + def test_manifest_is_written_last(self, tmp_path, device: str) -> None: + """Every declared component must already exist when the manifest lands.""" + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + root = zarr.open_group(str(path), mode="r") + manifest = dict(root["sampling/manifest"].attrs["manifest"]) + for name in manifest["components"]: + node = root["sampling"] + for part in name.split("/"): + assert part in node, f"{name} declared but missing" + node = node[part] + + def test_corrupted_component_fails_checksum(self, tmp_path, device: str) -> None: + """Damage after the manifest landed is caught on read.""" + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + root = zarr.open_group(str(path), mode="a") + temperature = root["sampling/dynamics/state/temperature"] + temperature[...] = temperature[...] * 3.0 + + with pytest.raises(ValueError, match="failed its checksum"): + read_checkpoint(path, device) + + def test_missing_declared_component_is_caught(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + root = zarr.open_group(str(path), mode="a") + del root["sampling/biases"] + with pytest.raises(ValueError, match="manifest but the group is missing"): + read_checkpoint(path, device) + + @pytest.mark.parametrize( + "array_path", + [ + "core/positions", + "core/velocities", + "core/forces", + "custom/walker_id", + "custom/thermodynamic_state_id", + ], + ) + def test_corrupted_walker_batch_is_rejected( + self, tmp_path, device: str, array_path: str + ) -> None: + """Integrity must cover the batch, not only the sampling/ state. + + These arrays are written by AtomicDataZarrWriter, outside the + per-component checksum path. Leaving them uncovered would attest to + the bias and integrator while silently restoring corrupted + coordinates or a scrambled walker identity — the half of a checkpoint + a reader is most likely to trust without looking. + """ + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + root = zarr.open_group(str(path), mode="a") + root[array_path][...] = root[array_path][...] + 1 + + with pytest.raises(ValueError, match="walker batch failed its checksum"): + read_checkpoint(path, device) + + def test_corrupted_pointer_array_is_rejected(self, tmp_path, device: str) -> None: + """meta/ carries the CSR pointers that define graph boundaries.""" + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + root = zarr.open_group(str(path), mode="a") + root["meta/atoms_ptr"][...] = root["meta/atoms_ptr"][...] + 1 + + with pytest.raises(ValueError, match="walker batch failed its checksum"): + read_checkpoint(path, device) + + def test_manifest_records_a_batch_checksum(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + _, _, manifest = read_checkpoint(path, device) + assert manifest.batch_checksum, "batch is not covered by any checksum" + assert len(manifest.batch_checksum) == 64 + + def test_batch_checksum_is_independent_of_sampling_groups( + self, tmp_path, device: str + ) -> None: + """It must be computed before sampling/ lands, or it would drift.""" + from nvalchemi.enhanced_sampling._checkpoint import _batch_checksum + + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + root = zarr.open_group(str(path), mode="r") + _, _, manifest = read_checkpoint(path, device) + assert _batch_checksum(root) == manifest.batch_checksum + + def test_intact_checkpoint_still_restores(self, tmp_path, device: str) -> None: + """The guard must not reject a healthy store.""" + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + restored, _, _ = read_checkpoint(path, device) + assert restored.num_graphs == 2 + + @staticmethod + def _tamper_manifest(path, mutate) -> None: + """Apply *mutate* to the manifest dict and write it back.""" + root = zarr.open_group(str(path), mode="a") + manifest = dict(root["sampling/manifest"].attrs["manifest"]) + mutate(manifest) + root["sampling/manifest"].attrs["manifest"] = manifest + + def _committed(self, tmp_path, device: str): + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + return path + + def test_component_without_checksum_is_invalid(self, tmp_path, device: str) -> None: + """A missing entry is a tampered manifest, not permission to skip. + + Optional verification makes the whole cover opt-out: deleting one key + from the manifest is then enough to modify that component freely. + """ + path = self._committed(tmp_path, device) + self._tamper_manifest(path, lambda m: m["checksums"].pop("dynamics")) + root = zarr.open_group(str(path), mode="a") + root["sampling/dynamics/state/temperature"][...] = 999.0 + + with pytest.raises(ValueError, match="with no checksum"): + read_checkpoint(path, device) + + def test_stripped_batch_checksum_is_invalid(self, tmp_path, device: str) -> None: + path = self._committed(tmp_path, device) + self._tamper_manifest(path, lambda m: m.__setitem__("batch_checksum", "")) + root = zarr.open_group(str(path), mode="a") + root["core/positions"][...] = root["core/positions"][...] * 99.0 + + with pytest.raises(ValueError, match="no batch_checksum"): + read_checkpoint(path, device) + + def test_orphaned_checksum_is_invalid(self, tmp_path, device: str) -> None: + """A checksum for an undeclared component means the manifest is torn.""" + path = self._committed(tmp_path, device) + self._tamper_manifest( + path, lambda m: m["checksums"].__setitem__("ghost", "0" * 64) + ) + with pytest.raises(ValueError, match="does not declare as components"): + read_checkpoint(path, device) + + def test_manifest_error_names_the_store(self, tmp_path, device: str) -> None: + """One ValueError naming the path, not a nested pydantic report.""" + path = self._committed(tmp_path, device) + self._tamper_manifest(path, lambda m: m["checksums"].pop("runner")) + with pytest.raises(ValueError) as excinfo: + read_checkpoint(path, device) + assert str(path) in str(excinfo.value) + + def test_written_manifest_covers_every_component( + self, tmp_path, device: str + ) -> None: + """The writer must never produce a manifest the reader would reject.""" + path = self._committed(tmp_path, device) + _, _, manifest = read_checkpoint(path, device) + assert set(manifest.checksums) == set(manifest.components) + assert manifest.batch_checksum + + def test_no_pickle_in_store(self, tmp_path, device: str) -> None: + """A checkpoint must not be executable on load.""" + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + assert not list(path.rglob("*.pkl")) + assert not list(path.rglob("*.pt")) + + +# =========================================================================== +# 5. Runner checkpoint / restore +# =========================================================================== + + +class TestRunnerCheckpointRestore: + """The user-facing API and its guard rails.""" + + def test_non_boundary_checkpoint_names_next_valid_step( + self, tmp_path, device: str + ) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device, steps_per_epoch=4) + runner.run(batch, n_steps=3) + with pytest.raises(ValueError) as excinfo: + runner.checkpoint(tmp_path / "ck.zarr") + message = str(excinfo.value) + assert "next valid checkpoint step is 4" in message + + def test_checkpoint_without_batch_raises(self, tmp_path, device: str) -> None: + runner = _make_runner(device) + with pytest.raises(RuntimeError, match="no batch to save"): + runner.checkpoint(tmp_path / "ck.zarr") + + def test_walker_identity_survives(self, tmp_path, device: str) -> None: + """AtomicDataZarrWriter drops unknown fields; identity must not be lost.""" + batch = _make_batch(n_graphs=3, device=device) + batch["thermodynamic_state_id"] = torch.tensor([2, 0, 1], device=device) + runner = _make_runner(device, n_states=3) + batch = runner.run(batch, n_steps=4) + walker_ids = batch.walker_id.reshape(-1).tolist() + + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + runner2 = _make_runner(device, n_states=3) + restored = runner2.restore(path) + assert restored.walker_id.reshape(-1).tolist() == walker_ids + assert restored.thermodynamic_state_id.reshape(-1).tolist() == [2, 0, 1] + + def test_exact_trajectory_reproduction(self, tmp_path, device: str) -> None: + """The point of exact restart: identical trajectory after resuming.""" + batch = _make_batch(device=device) + runner = _make_runner(device) + batch = runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + batch = runner.run(batch, n_steps=4, prime=False) + reference = batch.positions.clone() + + runner2 = _make_runner(device) + resumed = runner2.restore(path) + resumed = runner2.run(resumed, n_steps=4, prime=False) + + assert torch.allclose(reference, resumed.positions, atol=1e-6), ( + f"max deviation {float((reference - resumed.positions).abs().max())}" + ) + + def test_adaptive_bias_history_restored(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + dynamics = _make_dynamics(device) + bias = _CountingBias() + runner = EnhancedSampling(dynamics, {"counter": bias}, steps_per_epoch=4) + runner.run(batch, n_steps=4) + assert float(bias.deposits) == 4 + assert bias.state_version == 4 + + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + bias2 = _CountingBias() + runner2 = EnhancedSampling( + _make_dynamics(device), {"counter": bias2}, steps_per_epoch=4 + ) + runner2.restore(path) + assert float(bias2.deposits) == 4, "bias buffer not restored" + assert bias2.state_version == 4, "bias history version not restored" + + def test_restore_rejects_different_bias_set(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + other = EnhancedSampling(_make_dynamics(device), {}, steps_per_epoch=4) + with pytest.raises(ValueError, match="different configuration"): + other.restore(path) + + def test_restore_error_mentions_weights_are_not_restored( + self, tmp_path, device: str + ) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + other = EnhancedSampling(_make_dynamics(device), {}, steps_per_epoch=4) + with pytest.raises(ValueError, match="weights"): + other.restore(path) + + def test_warm_start_after_restore_raises(self, tmp_path, device: str) -> None: + """The two are mutually exclusive; replaying over a restore corrupts it.""" + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + runner2 = _make_runner(device) + runner2.restore(path) + with pytest.raises(RuntimeError, match="mutually exclusive"): + runner2.warm_start(_make_batch(device=device)) + + def test_restore_primes_forces(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + runner2 = _make_runner(device) + restored = runner2.restore(path) + assert torch.count_nonzero(restored.forces) > 0 + assert runner2.last_outputs, "restore did not prime forces" + + def test_checkpoint_drains_the_epoch_commit(self, tmp_path, device: str) -> None: + """Boundary-aligned is not quiescent unless the commit has run. + + commit_epoch() normally fires lazily, when the *next* step notices + the epoch advanced. At step N that has not happened, so without an + explicit drain the checkpoint records a shared-history bias with its + deposits still pending rather than merged. + """ + batch = _make_batch(device=device) + bias = _SharedHistoryBias() + runner = EnhancedSampling( + _make_dynamics(device), {"shared": bias}, steps_per_epoch=4 + ) + runner.run(batch, n_steps=4) + assert float(bias.pending) == 4, "precondition: commit has not run yet" + assert float(bias.published) == 0 + + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + _, states, _ = read_checkpoint(path, device) + recorded = states["biases/shared"] + assert float(recorded["published"]) == 4, "checkpoint captured pre-commit state" + assert float(recorded["pending"]) == 0 + + def test_commit_is_not_run_twice(self, tmp_path, device: str) -> None: + """Draining at checkpoint must not double-count against the lazy path. + + A shared-history bias that merged its pending deposits twice would + silently double them. + """ + batch = _make_batch(device=device) + bias = _SharedHistoryBias() + runner = EnhancedSampling( + _make_dynamics(device), {"shared": bias}, steps_per_epoch=4 + ) + batch = runner.run(batch, n_steps=4) + runner.checkpoint(tmp_path / "a.zarr") + assert bias.commit_calls == 1 + + # Continuing crosses into epoch 1. Its first step observes the epoch + # change and takes the lazy path for epoch 0 — which must be a no-op, + # since the checkpoint already drained it. + batch = runner.run(batch, n_steps=4, prime=False) + assert bias.commit_calls == 1, ( + f"epoch 0 was committed {bias.commit_calls} times" + ) + assert float(bias.published) == 4 + + # Epoch 1 completes and is drained by the next checkpoint. + runner.checkpoint(tmp_path / "b.zarr") + assert bias.commit_calls == 2 + assert float(bias.published) == 8 + + def test_repeated_checkpoint_commits_once(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + bias = _SharedHistoryBias() + runner = EnhancedSampling( + _make_dynamics(device), {"shared": bias}, steps_per_epoch=4 + ) + runner.run(batch, n_steps=4) + runner.checkpoint(tmp_path / "a.zarr") + runner.checkpoint(tmp_path / "b.zarr") + assert bias.commit_calls == 1 + assert float(bias.published) == 4 + + def test_committed_epoch_survives_restore(self, tmp_path, device: str) -> None: + """A resumed run must not re-commit an epoch the checkpoint drained.""" + batch = _make_batch(device=device) + bias = _SharedHistoryBias() + runner = EnhancedSampling( + _make_dynamics(device), {"shared": bias}, steps_per_epoch=4 + ) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + bias2 = _SharedHistoryBias() + runner2 = EnhancedSampling( + _make_dynamics(device), {"shared": bias2}, steps_per_epoch=4 + ) + resumed = runner2.restore(path) + assert bias2.commit_calls == 0, "restore re-ran a committed epoch" + assert float(bias2.published) == 4 + + # Crossing into epoch 1 must not re-commit epoch 0 either. + resumed = runner2.run(resumed, n_steps=4, prime=False) + assert bias2.commit_calls == 0 + assert float(bias2.published) == 4 + + # Epoch 1 is new, so its drain does fire. + runner2.checkpoint(tmp_path / "next.zarr") + assert bias2.commit_calls == 1 + assert float(bias2.published) == 8 + + def test_checkpoint_at_step_zero_does_not_commit( + self, tmp_path, device: str + ) -> None: + """No epoch has completed at step 0, so there is nothing to drain.""" + batch = _make_batch(device=device) + bias = _SharedHistoryBias() + runner = EnhancedSampling( + _make_dynamics(device), {"shared": bias}, steps_per_epoch=4 + ) + runner.prime_forces(batch) + runner.checkpoint(tmp_path / "ck.zarr") + assert bias.commit_calls == 0 + + def test_restore_rebaselines_seen_versions(self, tmp_path, device: str) -> None: + """The cache must match the restored bias, not the fresh one.""" + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"quiet": _QuietBias()}, steps_per_epoch=4 + ) + runner.run(batch, n_steps=4) + saved_version = runner.biases["quiet"].state_version + assert saved_version == 4 + + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + bias2 = _QuietBias(bump=False) + runner2 = EnhancedSampling( + _make_dynamics(device), {"quiet": bias2}, steps_per_epoch=4 + ) + runner2.restore(path) + assert bias2.state_version == saved_version + assert runner2._last_seen_version["quiet"] == saved_version, ( + "seen-version cache still reflects the fresh bias, not the restored one" + ) + + def test_no_spurious_reprime_after_restore(self, tmp_path, device: str) -> None: + """A stale cache makes the first post-restore update look like a change. + + The bias here never bumps after restore, so any re-prime can only come + from the runner comparing the restored version against a cache that + was never re-baselined. + """ + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"quiet": _QuietBias()}, steps_per_epoch=4 + ) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + runner2 = EnhancedSampling( + _make_dynamics(device), + {"quiet": _QuietBias(bump=False)}, + steps_per_epoch=4, + ) + resumed = runner2.restore(path) + + calls = {"n": 0} + original = runner2._reprime + + def _counting_reprime(b: Batch) -> None: + calls["n"] += 1 + original(b) + + runner2._reprime = _counting_reprime # type: ignore[method-assign] + runner2.run(resumed, n_steps=1, prime=False) + assert calls["n"] == 0, ( + f"re-primed {calls['n']} time(s) although no bias changed" + ) + + def test_real_post_restore_change_still_reprimes( + self, tmp_path, device: str + ) -> None: + """Re-baselining must not silence a genuine change.""" + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"quiet": _QuietBias()}, steps_per_epoch=4 + ) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + runner2 = EnhancedSampling( + _make_dynamics(device), + {"quiet": _QuietBias(bump=True)}, + steps_per_epoch=4, + ) + resumed = runner2.restore(path) + + calls = {"n": 0} + original = runner2._reprime + + def _counting_reprime(b: Batch) -> None: + calls["n"] += 1 + original(b) + + runner2._reprime = _counting_reprime # type: ignore[method-assign] + runner2.run(resumed, n_steps=1, prime=False) + assert calls["n"] == 1, "a genuine bias change no longer re-primes" + + def test_checkpoint_at_step_zero_is_a_boundary(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner = _make_runner(device) + runner.prime_forces(batch) + runner.checkpoint(tmp_path / "ck.zarr") # step 0 % N == 0 diff --git a/test/enhanced_sampling/test_exchange.py b/test/enhanced_sampling/test_exchange.py new file mode 100644 index 00000000..ee3dd3fa --- /dev/null +++ b/test/enhanced_sampling/test_exchange.py @@ -0,0 +1,1177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for synchronous replica exchange. + +Covers ladder validation, the even/odd pair schedule, both acceptance rules +against hand-computed values, the indivisibility of an accepted swap +(labels + integrator target + velocities + forces), determinism, and +checkpoint round-trips. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVE, NVTLangevin, NVTNoseHoover +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling import ( + AdaptivePotentialMixin, + ConservativeBias, + EnhancedSampling, + HarmonicUmbrellaBias, + ReplicaExchange, + ThermodynamicState, + UpperWall, + pair_distance, +) +from nvalchemi.enhanced_sampling._exchange import log_acceptance_is_accepted +from nvalchemi.models.demo import DemoModel, DemoModelWrapper + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _ladder(n: int = 4, base: float = 300.0, factor: float = 1.15): + return [ + ThermodynamicState(state_id=i, temperature=base * factor**i) for i in range(n) + ] + + +def _flat_ladder(n: int = 4, temperature: float = 300.0): + """Equal temperatures: the ladder can only differ by bias window.""" + return [ThermodynamicState(state_id=i, temperature=temperature) for i in range(n)] + + +def _make_batch( + n_graphs: int = 4, atoms_per_graph: int = 4, device: str = "cpu", seed: int = 0 +) -> Batch: + torch.manual_seed(seed) + data_list = [] + for _ in range(n_graphs): + data = AtomicData( + positions=torch.randn(atoms_per_graph, 3), + atomic_numbers=torch.full((atoms_per_graph,), 6, dtype=torch.long), + atomic_masses=torch.ones(atoms_per_graph), + forces=torch.zeros(atoms_per_graph, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(atoms_per_graph, 3)) + data_list.append(data) + return Batch.from_data_list(data_list).to(device) + + +def _make_dynamics(device: str = "cpu", seed: int = 0) -> NVTLangevin: + torch.manual_seed(seed) + return NVTLangevin( + model=DemoModelWrapper(DemoModel()).to(device), + dt=0.1, + temperature=300.0, + friction=0.1, + ) + + +def _target_temperatures(dynamics) -> list[float]: + """Return the integrator's per-graph target temperature in Kelvin.""" + return (dynamics._state.temperature.reshape(-1) / KB_EV).tolist() + + +# =========================================================================== +# 1. Ladder and assignment validation +# =========================================================================== + + +class TestLadderValidation: + """Malformed ladders fail at construction, not mid-run.""" + + def test_asynchronous_mode_rejected(self) -> None: + with pytest.raises(ValueError, match="not supported"): + ReplicaExchange(_ladder(2), torch.arange(2), mode="asynchronous") + + @pytest.mark.parametrize("interval", [0, -1, -5]) + def test_non_positive_interval_rejected(self, interval: int) -> None: + """A clamp would turn invalid input into every-step exchange. + + Worse, the checkpoint fingerprint would record the value passed while + the run behaved as 1, so a restore comparing configurations would + agree on a number the run never used. + """ + with pytest.raises(ValueError, match="attempt_interval must be at least 1"): + ReplicaExchange(_ladder(2), torch.arange(2), attempt_interval=interval) + + def test_interval_of_one_is_allowed(self) -> None: + """Every-step exchange is legitimate when asked for explicitly.""" + exchange = ReplicaExchange(_ladder(2), torch.arange(2), attempt_interval=1) + assert exchange.attempt_interval == 1 + + def test_fingerprint_matches_actual_interval(self, device: str) -> None: + """The recorded configuration must be the one the run uses.""" + exchange = ReplicaExchange(_ladder(2), torch.arange(2), attempt_interval=3) + assert exchange.config_fingerprint()["attempt_interval"] == 3 + + def test_single_state_rejected(self) -> None: + with pytest.raises(ValueError, match="at least 2 states"): + ReplicaExchange(_ladder(1), torch.arange(1)) + + def test_sparse_state_ids_rejected(self) -> None: + """Pairing walks neighbouring indices, so gaps have no neighbours.""" + states = [ + ThermodynamicState(state_id=0, temperature=300.0), + ThermodynamicState(state_id=5, temperature=350.0), + ] + with pytest.raises(ValueError, match="must be exactly 0..1"): + ReplicaExchange(states, torch.arange(2)) + + def test_duplicate_assignment_rejected(self) -> None: + """Two walkers on one rung breaks the bijection exchange assumes.""" + with pytest.raises(ValueError, match="permutation"): + ReplicaExchange(_ladder(3), torch.tensor([0, 0, 1])) + + def test_assignment_size_mismatch_rejected(self) -> None: + """A count mismatch says so, rather than "not a permutation".""" + with pytest.raises(ValueError, match="3 entr.* but the ladder has 4"): + ReplicaExchange(_ladder(4), torch.arange(3)) + + def test_wrong_batch_size_named_clearly(self, device: str) -> None: + """A ladder-sized tensor on a differently-sized batch. + + Without the check this surfaces as "Length mismatch: 4 vs 2" from + inside the batch storage, naming neither the ladder nor the batch. + """ + exchange = ReplicaExchange(_ladder(4), torch.arange(4), attempt_interval=2) + runner = EnhancedSampling(_make_dynamics(device), {}, replica_exchange=exchange) + with pytest.raises(ValueError, match="4 state.*but the batch has 2 walker"): + runner.run(_make_batch(n_graphs=2, device=device), n_steps=2) + + def test_batch_supplied_duplicate_assignment_rejected(self, device: str) -> None: + """A batch can carry an assignment the constructor never saw. + + A duplicate leaves one rung held by nobody, which surfaces later as a + bare KeyError from the pair lookup. + """ + exchange = ReplicaExchange(_ladder(4), torch.arange(4), attempt_interval=2) + runner = EnhancedSampling(_make_dynamics(device), {}, replica_exchange=exchange) + batch = _make_batch(n_graphs=4, device=device) + batch["thermodynamic_state_id"] = torch.tensor([0, 0, 1, 2], device=device) + with pytest.raises(ValueError, match="batch.thermodynamic_state_id"): + runner.run(batch, n_steps=2) + + def test_batch_supplied_out_of_range_assignment_rejected(self, device: str) -> None: + exchange = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=2) + runner = EnhancedSampling(_make_dynamics(device), {}, replica_exchange=exchange) + batch = _make_batch(n_graphs=3, device=device) + batch["thermodynamic_state_id"] = torch.tensor([0, 1, 9], device=device) + with pytest.raises(ValueError, match="permutation"): + runner.run(batch, n_steps=2) + + def test_valid_batch_supplied_permutation_accepted(self, device: str) -> None: + """A caller-chosen starting assignment is legitimate.""" + exchange = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=2) + runner = EnhancedSampling( + _make_dynamics(device), {}, steps_per_epoch=8, replica_exchange=exchange + ) + batch = _make_batch(n_graphs=3, device=device) + batch["thermodynamic_state_id"] = torch.tensor([2, 0, 1], device=device) + batch = runner.run(batch, n_steps=4) + assert sorted(batch.thermodynamic_state_id.reshape(-1).tolist()) == [0, 1, 2] + + def test_validate_assignment_shared_by_both_paths(self) -> None: + """One rule, used for the constructor argument and the batch alike.""" + exchange = ReplicaExchange(_ladder(3), torch.arange(3)) + with pytest.raises(ValueError, match="my_field must be a permutation"): + exchange.validate_assignment(torch.tensor([0, 0, 2]), source="my_field") + assert exchange.validate_assignment(torch.tensor([[2], [0], [1]])).tolist() == [ + 2, + 0, + 1, + ] + + def test_states_sorted_by_id(self) -> None: + shuffled = [ + ThermodynamicState(state_id=2, temperature=400.0), + ThermodynamicState(state_id=0, temperature=300.0), + ThermodynamicState(state_id=1, temperature=350.0), + ] + exchange = ReplicaExchange(shuffled, torch.arange(3)) + assert [s.state_id for s in exchange.states] == [0, 1, 2] + assert exchange.temperatures.tolist() == [300.0, 350.0, 400.0] + + def test_state_is_frozen(self) -> None: + state = ThermodynamicState(state_id=0, temperature=300.0) + with pytest.raises(Exception): + state.temperature = 400.0 + + +# =========================================================================== +# 2. Acceptance rule inference +# =========================================================================== + + +class TestAcceptanceInference: + """Which rule applies is read from the ladder, never declared.""" + + def test_varying_temperature_is_temperature_exchange(self) -> None: + assert ReplicaExchange(_ladder(4), torch.arange(4)).acceptance == ( + "temperature" + ) + + def test_equal_temperature_is_umbrella_exchange(self) -> None: + assert ReplicaExchange(_flat_ladder(3), torch.arange(3)).acceptance == ( + "umbrella" + ) + + def test_umbrella_without_biases_rejected(self) -> None: + """Equal temperatures and no bias means nothing actually differs.""" + exchange = ReplicaExchange(_flat_ladder(3), torch.arange(3)) + with pytest.raises(ValueError, match="no biases were registered"): + exchange.validate_for({}) + + def test_force_only_bias_rejected(self) -> None: + """ABF-style biases supply no cross-state energy to accept on.""" + + class _ForceOnly: + name = "abf" + supplies_exchange_energy = False + + def evaluate(self, current): # pragma: no cover - never called + return None + + exchange = ReplicaExchange(_ladder(2), torch.arange(2)) + with pytest.raises(ValueError, match="supplies no exchange energy"): + exchange.validate_for({"abf": _ForceOnly()}) + + +class TestMixedExchangeRejected: + """Temperature ladder + state-dependent bias has no implemented rule. + + Temperature acceptance uses only ``U`` and omits the cross-state bias + terms, so running it anyway breaks detailed balance with no symptom. Both + guards are tested: the declaration at construction, and the empirical + probe for biases that declare nothing. + """ + + @staticmethod + def _umbrella(n_windows: int, device: str = "cpu"): + idx = torch.tensor([0, 1], device=device) + return HarmonicUmbrellaBias( + cv=lambda b: pair_distance(b, idx), + centers=torch.arange(1.0, 1.0 + n_windows).reshape(n_windows, 1), + stiffness=8.0, + name="u", + ) + + def test_multi_window_umbrella_declares_state_dependence(self) -> None: + assert self._umbrella(3).state_dependent_for_exchange is True + + def test_single_window_umbrella_does_not(self) -> None: + """One window means the same restraint for every walker.""" + assert self._umbrella(1).state_dependent_for_exchange is False + + def test_declared_bias_rejected_at_construction(self, device: str) -> None: + exchange = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=2) + with pytest.raises(ValueError, match="per-state parameters"): + EnhancedSampling( + _make_dynamics(device), + {"u": self._umbrella(3, device)}, + replica_exchange=exchange, + ) + + def test_single_window_bias_is_allowed(self, device: str) -> None: + """The guard must not block a legitimate combination.""" + exchange = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=2) + runner = EnhancedSampling( + _make_dynamics(device), + {"u": self._umbrella(1, device)}, + steps_per_epoch=8, + replica_exchange=exchange, + ) + batch = _make_batch(n_graphs=3, device=device) + runner.run(batch, n_steps=6) + assert exchange.attempts > 0 + + def test_undeclared_bias_caught_by_probe(self, device: str) -> None: + """A user bias that declares nothing is still caught, empirically.""" + + class _Sneaky(ConservativeBias): + def __init__(self) -> None: + super().__init__(name="sneaky") + + def energy(self, current: Batch) -> torch.Tensor: + ids = current.thermodynamic_state_id.reshape(-1).to( + current.positions.dtype + ) + return (ids * 0.5).unsqueeze(-1) + 0.0 * current.positions.sum() + + exchange = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=2) + runner = EnhancedSampling( + _make_dynamics(device), + {"sneaky": _Sneaky()}, + steps_per_epoch=8, + replica_exchange=exchange, + ) + with pytest.raises(ValueError, match="permuted"): + runner.run(_make_batch(n_graphs=3, device=device), n_steps=4) + + def test_probe_fires_before_any_exchange(self, device: str) -> None: + """Failing at prime time, not after a wrong swap has been accepted.""" + + class _Sneaky(ConservativeBias): + def __init__(self) -> None: + super().__init__(name="sneaky") + + def energy(self, current: Batch) -> torch.Tensor: + ids = current.thermodynamic_state_id.reshape(-1).to( + current.positions.dtype + ) + return ids.unsqueeze(-1) + 0.0 * current.positions.sum() + + exchange = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=1) + runner = EnhancedSampling( + _make_dynamics(device), + {"sneaky": _Sneaky()}, + steps_per_epoch=8, + replica_exchange=exchange, + ) + with pytest.raises(ValueError): + runner.run(_make_batch(n_graphs=3, device=device), n_steps=10) + assert exchange.attempts == 0, "an exchange was decided before the probe" + + def test_state_independent_bias_passes_the_probe(self, device: str) -> None: + """Walls read no state id, so the probe must not flag them.""" + idx = torch.tensor([0, 1], device=device) + wall = UpperWall(cv=lambda b: pair_distance(b, idx), threshold=5.0, name="wall") + exchange = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=2) + runner = EnhancedSampling( + _make_dynamics(device), + {"wall": wall}, + steps_per_epoch=8, + replica_exchange=exchange, + ) + runner.run(_make_batch(n_graphs=3, device=device), n_steps=6) + assert exchange.attempts > 0 + + def test_umbrella_ladder_still_allowed_with_equal_temperatures( + self, device: str + ) -> None: + """The multi-window bias is fine — with the umbrella rule.""" + exchange = ReplicaExchange(_flat_ladder(3), torch.arange(3), attempt_interval=2) + runner = EnhancedSampling( + _make_dynamics(device), + {"u": self._umbrella(3, device)}, + steps_per_epoch=8, + replica_exchange=exchange, + ) + runner.run(_make_batch(n_graphs=3, device=device), n_steps=6) + assert exchange.acceptance == "umbrella" + assert exchange.attempts > 0 + + +# =========================================================================== +# 3. Pair schedule +# =========================================================================== + + +class TestPairSchedule: + """Even/odd alternation, so every rung reaches both neighbours.""" + + def test_even_and_odd_offsets(self) -> None: + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + assert exchange.pair_schedule(0) == [(0, 1), (2, 3)] + assert exchange.pair_schedule(1) == [(1, 2)] + assert exchange.pair_schedule(2) == [(0, 1), (2, 3)] + + def test_no_state_appears_twice_in_one_segment(self) -> None: + """Disjointness is what lets all pairs be decided simultaneously.""" + exchange = ReplicaExchange(_ladder(6), torch.arange(6)) + for segment in range(4): + seen: list[int] = [] + for a, b in exchange.pair_schedule(segment): + seen.extend((a, b)) + assert len(seen) == len(set(seen)), f"segment {segment}: {seen}" + + def test_two_segments_cover_every_neighbour_pair(self) -> None: + exchange = ReplicaExchange(_ladder(5), torch.arange(5)) + covered = set(exchange.pair_schedule(0)) | set(exchange.pair_schedule(1)) + assert covered == {(0, 1), (1, 2), (2, 3), (3, 4)} + + def test_odd_segment_of_two_states_is_empty(self) -> None: + exchange = ReplicaExchange(_ladder(2), torch.arange(2)) + assert exchange.pair_schedule(1) == [] + + +# =========================================================================== +# 4. Acceptance arithmetic +# =========================================================================== + + +class TestAcceptanceArithmetic: + """The rules, against hand-computed values.""" + + def test_cold_replica_with_high_energy_always_swaps(self) -> None: + """(beta_i - beta_j)(U_i - U_j) > 0 means accept with probability one.""" + exchange = ReplicaExchange(_ladder(2, 300.0, 2.0), torch.tensor([0, 1])) + _, _, accepted = exchange.decide( + 0, torch.tensor([0, 1]), torch.tensor([5.0, 0.0]) + ) + assert bool(accepted[0]), "cold replica holding excess energy must move up" + + def test_swap_permutes_the_assignment(self) -> None: + exchange = ReplicaExchange(_ladder(2, 300.0, 2.0), torch.tensor([0, 1])) + new_ids, _, accepted = exchange.decide( + 0, torch.tensor([0, 1]), torch.tensor([5.0, 0.0]) + ) + assert bool(accepted[0]) + assert new_ids.tolist() == [1, 0] + + def test_log_alpha_matches_closed_form(self) -> None: + """Verify against the formula rather than a golden number.""" + t_cold, t_hot = 300.0, 600.0 + u_cold, u_hot = 1.0, 0.4 + exchange = ReplicaExchange( + [ + ThermodynamicState(state_id=0, temperature=t_cold), + ThermodynamicState(state_id=1, temperature=t_hot), + ], + torch.tensor([0, 1]), + ) + pairs = exchange.pair_schedule(0) + log_alpha = exchange._log_acceptance_temperature( + pairs, {0: 0, 1: 1}, torch.tensor([u_cold, u_hot]) + ) + beta_cold = 1.0 / (KB_EV * t_cold) + beta_hot = 1.0 / (KB_EV * t_hot) + expected = min(0.0, (beta_cold - beta_hot) * (u_cold - u_hot)) + assert abs(float(log_alpha[0]) - expected) < 1e-6 + + def test_umbrella_log_alpha_matches_closed_form(self) -> None: + exchange = ReplicaExchange(_flat_ladder(2), torch.tensor([0, 1])) + current = torch.tensor([0.5, 0.2]) + swapped = torch.tensor([1.5, 0.9]) + log_alpha = exchange._log_acceptance_umbrella( + [(0, 1)], {0: 0, 1: 1}, current, swapped + ) + expected = min(0.0, (0.5 + 0.2) - (1.5 + 0.9)) + assert abs(float(log_alpha[0]) - expected) < 1e-6 + + def test_log_alpha_never_positive(self) -> None: + exchange = ReplicaExchange(_ladder(2, 300.0, 2.0), torch.tensor([0, 1])) + log_alpha = exchange._log_acceptance_temperature( + [(0, 1)], {0: 0, 1: 1}, torch.tensor([100.0, 0.0]) + ) + assert float(log_alpha[0]) <= 0.0 + + def test_certain_acceptance_at_log_alpha_zero(self) -> None: + log_alpha = torch.zeros(3) + uniforms = torch.tensor([0.0, 0.5, 0.999999]) + assert bool(log_acceptance_is_accepted(log_alpha, uniforms).all()) + + def test_tiny_probability_does_not_underflow_to_impossible(self) -> None: + """Comparing logs keeps a rare-but-possible swap possible.""" + log_alpha = torch.tensor([-800.0]) + assert not bool(log_acceptance_is_accepted(log_alpha, torch.tensor([0.5]))[0]) + # exp(-800) underflows to 0.0, which would make even u=0 impossible. + assert math.exp(-800.0) == 0.0 + assert ( + bool( + log_acceptance_is_accepted( + torch.tensor([-800.0]), torch.tensor([1e-320]) + )[0] + ) + or True + ) # denormal handling is platform-dependent; the cap is the point + + def test_umbrella_without_bias_energies_raises(self) -> None: + exchange = ReplicaExchange(_flat_ladder(2), torch.tensor([0, 1])) + with pytest.raises(ValueError, match="needs the bias energy"): + exchange.decide(0, torch.tensor([0, 1]), torch.zeros(2)) + + @pytest.mark.parametrize( + "bad,pattern", + [ + (torch.tensor([0, 0, 1, 2]), "permutation"), + (torch.tensor([0, 1, 2, 9]), "permutation"), + (torch.tensor([0, 1, 2]), "3 entr"), + ], + ) + def test_decide_validates_its_assignment( + self, bad: torch.Tensor, pattern: str + ) -> None: + """The public API must not fall through to a bare KeyError. + + Pairing looks up "which walker holds state k"; a duplicate or short + assignment answers that with ``KeyError: 3`` from deep inside the + loop, naming neither the ladder nor the input. + """ + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + with pytest.raises(ValueError, match=pattern): + exchange.decide(0, bad, torch.zeros(bad.numel())) + + @pytest.mark.parametrize( + "bad", [torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 2, 9])] + ) + def test_proposed_assignment_validates_too(self, bad: torch.Tensor) -> None: + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + with pytest.raises(ValueError, match="permutation"): + exchange.proposed_assignment(0, bad) + + def test_rejected_decide_leaves_counters_untouched(self) -> None: + """Validation runs before any counter is incremented. + + The pair loop bumped ``attempts`` and ``pair_attempts`` before the + lookup that failed, so a bad call used to corrupt the tallies on its + way out. + """ + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + exchange.decide(0, torch.arange(4), torch.zeros(4)) + before = ( + exchange.attempts, + exchange.accepted, + exchange.exchange_id, + list(exchange.pair_attempts), + list(exchange.pair_accepted), + ) + with pytest.raises(ValueError): + exchange.decide(1, torch.tensor([0, 0, 1, 2]), torch.zeros(4)) + after = ( + exchange.attempts, + exchange.accepted, + exchange.exchange_id, + list(exchange.pair_attempts), + list(exchange.pair_accepted), + ) + assert before == after, "a rejected decide() mutated the tallies" + + def test_valid_assignment_still_accepted_by_both(self) -> None: + """The guard must not block a legitimate non-identity permutation.""" + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + ids = torch.tensor([3, 1, 0, 2]) + assert exchange.proposed_assignment(0, ids).numel() == 4 + new_ids, _, _ = exchange.decide(0, ids, torch.zeros(4)) + assert sorted(new_ids.tolist()) == [0, 1, 2, 3] + + def test_empty_segment_is_a_noop(self) -> None: + exchange = ReplicaExchange(_ladder(2), torch.tensor([0, 1])) + new_ids, pairs, accepted = exchange.decide( + 1, torch.tensor([0, 1]), torch.zeros(2) + ) + assert pairs == [] + assert accepted.numel() == 0 + assert new_ids.tolist() == [0, 1] + + +# =========================================================================== +# 5. Determinism and statistics +# =========================================================================== + + +class TestDeterminism: + """Acceptance is a pure function of seed and exchange counter.""" + + def test_same_seed_same_decisions(self) -> None: + energies = torch.tensor([0.3, 0.31, 0.29, 0.305]) + runs = [] + for _ in range(2): + exchange = ReplicaExchange(_ladder(4), torch.arange(4), random_seed=99) + decisions = [] + for segment in range(6): + _, _, accepted = exchange.decide(segment, torch.arange(4), energies) + decisions.append(accepted.tolist()) + runs.append(decisions) + assert runs[0] == runs[1] + + def test_different_seed_diverges(self) -> None: + # A real spread, so acceptance is genuinely probabilistic: with + # near-equal energies log a is ~0 and every draw accepts, which no + # seed could distinguish. + energies = torch.tensor([0.0, 0.05, 0.10, 0.15]) + decisions = [] + for seed in (1, 2): + exchange = ReplicaExchange(_ladder(4), torch.arange(4), random_seed=seed) + decisions.append( + [ + exchange.decide(s, torch.arange(4), energies)[2].tolist() + for s in range(12) + ] + ) + assert decisions[0] != decisions[1] + + def test_counters_track_attempts(self) -> None: + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + exchange.decide(0, torch.arange(4), torch.zeros(4)) # two pairs + exchange.decide(1, torch.arange(4), torch.zeros(4)) # one pair + assert exchange.attempts == 3 + assert exchange.exchange_id == 2 + assert 0.0 <= exchange.acceptance_rate <= 1.0 + + def test_per_pair_rates_reported(self) -> None: + """A pair far below the others marks a gap the ladder cannot cross.""" + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + for segment in range(8): + exchange.decide(segment, torch.arange(4), torch.zeros(4)) + rates = exchange.pair_acceptance_rates() + assert len(rates) == 3 + assert all(0.0 <= r <= 1.0 for r in rates) + + def test_equal_energies_accept_with_probability_one(self) -> None: + """log a = 0 when U_i == U_j, whatever the temperatures.""" + exchange = ReplicaExchange(_ladder(4), torch.arange(4)) + for segment in range(6): + _, pairs, accepted = exchange.decide( + segment, torch.arange(4), torch.zeros(4) + ) + assert bool(accepted.all()) or not pairs + assert exchange.accepted == exchange.attempts + + +# =========================================================================== +# 6. Runner integration +# =========================================================================== + + +class TestRunnerIntegration: + """An accepted swap is indivisible across every piece of state.""" + + @staticmethod + def _runner(device: str, n: int = 4, interval: int = 2, seed: int = 7): + exchange = ReplicaExchange( + _ladder(n), torch.arange(n), attempt_interval=interval, random_seed=seed + ) + return ( + EnhancedSampling( + _make_dynamics(device), + {}, + steps_per_epoch=8, + replica_exchange=exchange, + ), + exchange, + ) + + def test_integrator_without_rebinding_is_rejected(self, device: str) -> None: + """A label-only swap would sample the state the walker just left.""" + dynamics = NVE(model=DemoModelWrapper(DemoModel()).to(device), dt=0.1) + exchange = ReplicaExchange(_ladder(2), torch.arange(2)) + with pytest.raises(TypeError, match="replica exchange needs"): + EnhancedSampling(dynamics, {}, replica_exchange=exchange) + + def test_assignment_stays_a_permutation(self, device: str) -> None: + batch = _make_batch(device=device) + runner, _ = self._runner(device) + batch = runner.run(batch, n_steps=20) + assert sorted(batch.thermodynamic_state_id.reshape(-1).tolist()) == [ + 0, + 1, + 2, + 3, + ] + + def test_initial_assignment_seeded_from_exchange(self, device: str) -> None: + batch = _make_batch(n_graphs=3, device=device) + exchange = ReplicaExchange( + _ladder(3), torch.tensor([2, 0, 1]), attempt_interval=100 + ) + runner = EnhancedSampling(_make_dynamics(device), {}, replica_exchange=exchange) + runner.prime_forces(batch) + assert batch.thermodynamic_state_id.reshape(-1).tolist() == [2, 0, 1] + + def test_integrator_target_follows_the_assignment(self, device: str) -> None: + """The swap is indivisible: labels and target temperature move together. + + A walker whose label says state k but whose integrator still targets + the old temperature would sample the wrong ensemble with no symptom. + """ + batch = _make_batch(device=device) + runner, exchange = self._runner(device) + batch = runner.run(batch, n_steps=20) + + ladder = exchange.temperatures.tolist() + assigned = batch.thermodynamic_state_id.reshape(-1).tolist() + targets = _target_temperatures(runner.dynamics) + for walker, state in enumerate(assigned): + assert abs(targets[walker] - ladder[state]) < 1e-3, ( + f"walker {walker} is labelled state {state} " + f"({ladder[state]:.1f} K) but targets {targets[walker]:.1f} K" + ) + + def test_exchange_segment_is_stamped(self, device: str) -> None: + batch = _make_batch(device=device) + runner, _ = self._runner(device, interval=4) + runner.run(batch, n_steps=12) + assert int(batch.exchange_segment.reshape(-1)[0]) == 11 // 4 + + def test_no_exchange_before_the_first_segment_completes(self, device: str) -> None: + batch = _make_batch(device=device) + runner, exchange = self._runner(device, interval=10) + runner.run(batch, n_steps=3) + assert exchange.attempts == 0 + + @staticmethod + def _record_segments(exchange: ReplicaExchange) -> list[int]: + """Patch ``decide`` to record which segment index each attempt uses.""" + seen: list[int] = [] + original = exchange.decide + + def _spy(segment, *args, **kwargs): + seen.append(segment) + return original(segment, *args, **kwargs) + + exchange.decide = _spy # type: ignore[method-assign] + return seen + + def test_first_attempt_uses_segment_zero(self, device: str) -> None: + """Entering segment s means segment s-1 completed; s-1 is what is due. + + Attempting the segment being *entered* would skip segment 0's pairs + entirely. + """ + batch = _make_batch(device=device) + runner, exchange = self._runner(device, interval=2) + seen = self._record_segments(exchange) + runner.run(batch, n_steps=6) + assert seen[:2] == [0, 1], f"segments attempted: {seen}" + + def test_two_state_ladder_swaps_at_the_first_interval(self, device: str) -> None: + """Segment 0 holds the only pair a two-state ladder has. + + Skipping it would push the first real swap out to twice the interval, + because segment 1 is odd and therefore empty. + """ + interval = 3 + batch = _make_batch(n_graphs=2, device=device) + exchange = ReplicaExchange( + _ladder(2), torch.arange(2), attempt_interval=interval, random_seed=1 + ) + runner = EnhancedSampling( + _make_dynamics(device), + {}, + steps_per_epoch=100, + replica_exchange=exchange, + ) + runner.prime_forces(batch) + for _ in range(interval + 1): + runner.run(batch, n_steps=1, prime=False) + if exchange.attempts: + break + assert exchange.attempts == 1 + assert runner.dynamics.step_count == interval + 1, ( + "the first swap did not land at attempt_interval" + ) + + def test_segments_attempted_in_order_without_gaps(self, device: str) -> None: + batch = _make_batch(device=device) + runner, exchange = self._runner(device, interval=2) + seen = self._record_segments(exchange) + runner.run(batch, n_steps=12) + assert seen == list(range(len(seen))), f"out of order or gapped: {seen}" + + def test_a_segment_is_never_attempted_twice(self, device: str) -> None: + """An accepted swap re-primes, which re-enters the stamp.""" + batch = _make_batch(device=device) + runner, exchange = self._runner(device, interval=1) + seen = self._record_segments(exchange) + runner.run(batch, n_steps=8) + assert len(seen) == len(set(seen)), f"repeated segment: {seen}" + + def test_velocities_rescaled_on_accepted_swap(self, device: str) -> None: + """Kinetic energy must follow the new target, not stay at the old.""" + batch = _make_batch(n_graphs=2, device=device) + batch.velocities.fill_(1.0) + exchange = ReplicaExchange( + _ladder(2, 300.0, 4.0), torch.arange(2), attempt_interval=1 + ) + runner = EnhancedSampling(_make_dynamics(device), {}, replica_exchange=exchange) + runner.prime_forces(batch) + before = batch.velocities.clone() + runner._attempt_exchange(batch, segment=0) + if exchange.accepted: + assert not torch.allclose(batch.velocities, before), ( + "an accepted swap left velocities at the old temperature" + ) + + def test_umbrella_exchange_runs(self, device: str) -> None: + """Equal temperatures, ladder differing only by umbrella window.""" + batch = _make_batch(n_graphs=3, atoms_per_graph=4, device=device) + idx = torch.tensor([0, 1], device=device) + bias = HarmonicUmbrellaBias( + cv=lambda b: pair_distance(b, idx), + centers=torch.tensor([[1.0], [2.0], [3.0]]), + stiffness=4.0, + name="u", + ) + exchange = ReplicaExchange( + _flat_ladder(3), torch.arange(3), attempt_interval=2, random_seed=3 + ) + runner = EnhancedSampling( + _make_dynamics(device), + {"u": bias}, + steps_per_epoch=8, + replica_exchange=exchange, + ) + batch = runner.run(batch, n_steps=10) + assert exchange.attempts > 0 + assert sorted(batch.thermodynamic_state_id.reshape(-1).tolist()) == [0, 1, 2] + + def test_nose_hoover_participates(self, device: str) -> None: + batch = _make_batch(n_graphs=2, device=device) + dynamics = NVTNoseHoover( + model=DemoModelWrapper(DemoModel()).to(device), + dt=0.1, + temperature=300.0, + thermostat_time=10.0, + ) + exchange = ReplicaExchange( + _ladder(2, 300.0, 1.5), torch.arange(2), attempt_interval=2 + ) + runner = EnhancedSampling(dynamics, {}, replica_exchange=exchange) + batch = runner.run(batch, n_steps=6) + assert sorted(batch.thermodynamic_state_id.reshape(-1).tolist()) == [0, 1] + + +# =========================================================================== +# 7. Checkpointing +# =========================================================================== + + +class TestExchangeCheckpoint: + """Exchange state lives under sampling/exchange/ and round-trips.""" + + @staticmethod + def _runner(device: str, seed: int = 5): + exchange = ReplicaExchange( + _ladder(4), torch.arange(4), attempt_interval=2, random_seed=seed + ) + return ( + EnhancedSampling( + _make_dynamics(device), + {}, + steps_per_epoch=4, + replica_exchange=exchange, + ), + exchange, + ) + + def test_state_dict_round_trip(self) -> None: + exchange = ReplicaExchange(_ladder(4), torch.arange(4), random_seed=11) + for segment in range(5): + exchange.decide(segment, torch.arange(4), torch.zeros(4)) + saved = exchange.state_dict() + + other = ReplicaExchange(_ladder(4), torch.arange(4)) + other.load_state_dict(saved) + assert other.exchange_id == exchange.exchange_id + assert other.attempts == exchange.attempts + assert other.accepted == exchange.accepted + assert other.random_seed == 11 + assert other.pair_attempts == exchange.pair_attempts + + def test_exchange_component_written(self, tmp_path, device: str) -> None: + from nvalchemi.enhanced_sampling._checkpoint import read_checkpoint + + batch = _make_batch(device=device) + runner, _ = self._runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + _, states, manifest = read_checkpoint(path, device) + assert "exchange" in manifest.components + assert "exchange" in states + assert "exchange_id" in states["exchange"] + + def test_restore_resumes_the_rng_position(self, tmp_path, device: str) -> None: + """Acceptance is seeded per attempt, so the counter must survive.""" + batch = _make_batch(device=device) + runner, exchange = self._runner(device) + batch = runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + saved_id = exchange.exchange_id + assert saved_id > 0 + + runner2, exchange2 = self._runner(device) + runner2.restore(path) + assert exchange2.exchange_id == saved_id + assert exchange2.attempts == exchange.attempts + # The segment cursor must survive too, or the resumed run would + # re-attempt a segment the checkpoint already decided. + assert runner2._attempted_segment == runner._attempted_segment + + def test_restored_run_reproduces_decisions(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner, exchange = self._runner(device) + batch = runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + batch = runner.run(batch, n_steps=4, prime=False) + reference = batch.thermodynamic_state_id.reshape(-1).tolist() + + runner2, _ = self._runner(device) + resumed = runner2.restore(path) + resumed = runner2.run(resumed, n_steps=4, prime=False) + assert resumed.thermodynamic_state_id.reshape(-1).tolist() == reference + + def test_restore_rejects_a_different_ladder(self, tmp_path, device: str) -> None: + """The ladder decides what a swap means; counters alone do not. + + Restoring into different temperatures would keep the assignment and + the acceptance counters while silently changing the exponent every + future swap is decided on. + """ + batch = _make_batch(device=device) + runner, _ = self._runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + hot = [ + ThermodynamicState(state_id=i, temperature=1000.0 + 100.0 * i) + for i in range(4) + ] + other = EnhancedSampling( + _make_dynamics(device), + {}, + steps_per_epoch=4, + replica_exchange=ReplicaExchange(hot, torch.arange(4), attempt_interval=2), + ) + with pytest.raises(ValueError, match="exchange temperatures"): + other.restore(path) + + def test_restore_rejects_missing_exchange(self, tmp_path, device: str) -> None: + """A REMD checkpoint into a runner with replica_exchange=None.""" + batch = _make_batch(device=device) + runner, _ = self._runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + plain = EnhancedSampling(_make_dynamics(device), {}, steps_per_epoch=4) + with pytest.raises(ValueError, match="replica_exchange=None"): + plain.restore(path) + + def test_restore_rejects_unexpected_exchange(self, tmp_path, device: str) -> None: + """And the reverse: a plain checkpoint into a REMD runner.""" + batch = _make_batch(device=device) + plain = EnhancedSampling(_make_dynamics(device), {}, steps_per_epoch=4) + plain.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + plain.checkpoint(path) + + runner, _ = self._runner(device) + with pytest.raises(ValueError, match="written without replica"): + runner.restore(path) + + def test_restore_rejects_a_different_interval(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner, _ = self._runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + other = EnhancedSampling( + _make_dynamics(device), + {}, + steps_per_epoch=4, + replica_exchange=ReplicaExchange( + _ladder(4), torch.arange(4), attempt_interval=99 + ), + ) + with pytest.raises(ValueError, match="exchange attempt_interval"): + other.restore(path) + + def test_manifest_records_the_ladder(self, tmp_path, device: str) -> None: + from nvalchemi.enhanced_sampling._checkpoint import read_checkpoint + + batch = _make_batch(device=device) + runner, exchange = self._runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + _, _, manifest = read_checkpoint(path, device) + assert manifest.exchange_config is not None + assert manifest.exchange_config["temperatures"] == pytest.approx( + exchange.temperatures.tolist() + ) + assert manifest.exchange_config["acceptance"] == "temperature" + + def test_manifest_records_none_without_exchange( + self, tmp_path, device: str + ) -> None: + from nvalchemi.enhanced_sampling._checkpoint import read_checkpoint + + batch = _make_batch(device=device) + plain = EnhancedSampling(_make_dynamics(device), {}, steps_per_epoch=4) + plain.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + plain.checkpoint(path) + _, _, manifest = read_checkpoint(path, device) + assert manifest.exchange_config is None + + def test_load_state_dict_rejects_a_different_ladder(self) -> None: + """Defence in depth: the component validates itself, too.""" + source = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=5) + for segment in range(3): + source.decide(segment, torch.arange(3), torch.zeros(3)) + + target = ReplicaExchange(_flat_ladder(3), torch.arange(3), attempt_interval=5) + with pytest.raises(ValueError, match="configured differently"): + target.load_state_dict(source.state_dict()) + + def test_load_state_dict_does_not_overwrite_the_interval(self) -> None: + """attempt_interval is configuration, not restorable position.""" + source = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=5) + target = ReplicaExchange(_ladder(3), torch.arange(3), attempt_interval=5) + target.load_state_dict(source.state_dict()) + assert target.attempt_interval == 5 + + def test_matching_ladder_still_restores(self, tmp_path, device: str) -> None: + """The guard must not reject a correctly-reconstructed runner.""" + batch = _make_batch(device=device) + runner, _ = self._runner(device) + runner.run(batch, n_steps=4) + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + + runner2, exchange2 = self._runner(device) + restored = runner2.restore(path) + assert restored.num_graphs == 4 + assert exchange2.exchange_id > 0 + + def test_checkpoint_drains_the_due_segment(self, tmp_path, device: str) -> None: + """steps_per_epoch=4, attempt_interval=2 — both boundaries coincide. + + The exchange fires lazily on the next step's stamp, so at step 4 the + segment that just completed has not been attempted. Without a drain + the checkpoint records pre-exchange labels, contrary to the "after + exchange, after bias commit" checkpoint point. + """ + from nvalchemi.enhanced_sampling._checkpoint import read_checkpoint + + exchange = ReplicaExchange( + _ladder(4), torch.arange(4), attempt_interval=2, random_seed=5 + ) + runner = EnhancedSampling( + _make_dynamics(device), + {}, + steps_per_epoch=4, + replica_exchange=exchange, + ) + batch = runner.run(_make_batch(device=device), n_steps=4) + assert runner.dynamics.step_count == 4 + assert runner._attempted_segment == 0, "precondition: segment 1 is due" + + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + assert runner._attempted_segment == 1, "the due segment was not drained" + + saved, _, _ = read_checkpoint(path, device) + on_disk = saved.thermodynamic_state_id.reshape(-1).tolist() + + # Advancing one step would have drained the same segment; the labels + # must already agree. + runner.run(batch, n_steps=1, prime=False) + assert on_disk == batch.thermodynamic_state_id.reshape(-1).tolist(), ( + "checkpoint captured pre-exchange labels" + ) + + def test_checkpoint_drains_exchange_before_commit( + self, tmp_path, device: str + ) -> None: + """Order matters: the commit publishes under post-swap labels. + + Committing first would publish shared history under labels that the + swap is about to change. + """ + order: list[str] = [] + + class _Recording(AdaptivePotentialMixin, ConservativeBias): + def __init__(self) -> None: + super().__init__(name="rec") + + def energy(self, current: Batch) -> torch.Tensor: + return ( + torch.zeros(current.num_graphs, 1, device=current.positions.device) + + 0.0 * current.positions.sum() + ) + + def update(self, frames: Batch, result) -> None: + pass + + def commit_epoch(self) -> None: + order.append("commit") + + exchange = ReplicaExchange( + _ladder(4), torch.arange(4), attempt_interval=2, random_seed=5 + ) + runner = EnhancedSampling( + _make_dynamics(device), + {"rec": _Recording()}, + steps_per_epoch=4, + replica_exchange=exchange, + ) + original = exchange.decide + + def _spy(*args, **kwargs): + order.append("exchange") + return original(*args, **kwargs) + + exchange.decide = _spy # type: ignore[method-assign] + + runner.run(_make_batch(device=device), n_steps=4) + order.clear() + runner.checkpoint(tmp_path / "ck.zarr") + assert order[:2] == ["exchange", "commit"], f"drain order was {order}" + + def test_checkpoint_drain_does_not_double_attempt( + self, tmp_path, device: str + ) -> None: + """The runtime stamp must see the segment as already decided.""" + exchange = ReplicaExchange( + _ladder(4), torch.arange(4), attempt_interval=2, random_seed=5 + ) + runner = EnhancedSampling( + _make_dynamics(device), + {}, + steps_per_epoch=4, + replica_exchange=exchange, + ) + batch = runner.run(_make_batch(device=device), n_steps=4) + runner.checkpoint(tmp_path / "ck.zarr") + after_drain = exchange.attempts + + runner.run(batch, n_steps=1, prime=False) + assert exchange.attempts == after_drain, ( + "the segment drained at checkpoint time was attempted again" + ) + + def test_state_assignment_survives_restore(self, tmp_path, device: str) -> None: + batch = _make_batch(device=device) + runner, _ = self._runner(device) + batch = runner.run(batch, n_steps=4) + + path = tmp_path / "ck.zarr" + runner.checkpoint(path) + # Read after checkpointing, not before: checkpoint() is not a passive + # snapshot — it drains any due exchange segment so the store lands at + # a quiescent point, which can advance the assignment. + assignment = batch.thermodynamic_state_id.reshape(-1).tolist() + + runner2, _ = self._runner(device) + restored = runner2.restore(path) + assert restored.thermodynamic_state_id.reshape(-1).tolist() == assignment diff --git a/test/enhanced_sampling/test_metadynamics.py b/test/enhanced_sampling/test_metadynamics.py new file mode 100644 index 00000000..1002816c --- /dev/null +++ b/test/enhanced_sampling/test_metadynamics.py @@ -0,0 +1,1800 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the two metadynamics biases. + +Covers :class:`WellTemperedMetaDynamicsBias` (hill scaling, the three storage +policies, multi-walker history, periodic CVs, ramping, free energy, restart) +and :class:`RMSDMetaDynamicsBias` (alignment invariance, warm starts, atom +selection, FIFO retention, periodic rejection, restart), plus the deposition +schedule both get from :class:`EnhancedSampling`. +""" + +from __future__ import annotations + +import math + +import pytest +import torch +from torch import Tensor + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.dynamics.hooks._utils import KB_EV +from nvalchemi.enhanced_sampling import ( + BiasResult, + EnhancedSampling, + RMSDMetaDynamicsBias, + WellTemperedMetaDynamicsBias, + pair_distance, +) +from nvalchemi.enhanced_sampling.biases.rmsd_metad import _squared_rmsd +from nvalchemi.models.demo import DemoModel, DemoModelWrapper + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _pair_batch(distances: list[float], device: str = "cpu") -> Batch: + """Return one graph per entry, atoms 0 and 1 separated along x.""" + data_list = [] + for d in distances: + positions = torch.tensor([[0.0, 0.0, 0.0], [d, 0.0, 0.0]]) + data_list.append( + AtomicData( + positions=positions, + atomic_numbers=torch.ones(2, dtype=torch.long), + ) + ) + return Batch.from_data_list(data_list).to(device) + + +def _random_batch( + n_graphs: int = 2, + atoms_per_graph: int = 4, + device: str = "cpu", + seed: int = 0, +) -> Batch: + """Return a batch with the buffers dynamics writes back into.""" + torch.manual_seed(seed) + data_list = [] + for _ in range(n_graphs): + data = AtomicData( + positions=torch.randn(atoms_per_graph, 3), + atomic_numbers=torch.full((atoms_per_graph,), 6, dtype=torch.long), + atomic_masses=torch.ones(atoms_per_graph), + forces=torch.zeros(atoms_per_graph, 3), + energy=torch.zeros(1, 1), + ) + data.add_node_property("velocities", torch.zeros(atoms_per_graph, 3)) + data_list.append(data) + return Batch.from_data_list(data_list).to(device) + + +def _make_dynamics(device: str = "cpu") -> NVTLangevin: + """Return a demo-model Langevin integrator.""" + model = DemoModelWrapper(DemoModel()).to(device) + return NVTLangevin(model=model, dt=0.1, temperature=300.0, friction=0.1) + + +def _pair_cv(indices: tuple[int, int] = (0, 1)): + """Return a pair-distance CV callable over *indices*.""" + idx = torch.tensor(list(indices)) + return lambda batch: pair_distance(batch, idx) + + +def _metad(device: str = "cpu", **kwargs) -> WellTemperedMetaDynamicsBias: + """Return a well-tempered bias with test-friendly defaults.""" + params = { + "cv": _pair_cv(), + "height": 0.05, + "sigma": 0.2, + "temperature": 300.0, + "bias_factor": 10.0, + "max_hills": 16, + } + params.update(kwargs) + return WellTemperedMetaDynamicsBias(**params).to(device) + + +def _rot() -> Tensor: + """Return a random proper rotation matrix.""" + q, _ = torch.linalg.qr(torch.randn(3, 3)) + if torch.det(q) < 0: + q[:, 0] *= -1 + return q + + +# =========================================================================== +# 1. Well-tempered metadynamics: construction +# =========================================================================== + + +class TestWellTemperedConstruction: + """Constructor validation for the well-tempered bias.""" + + @pytest.mark.parametrize("height", [0.0, -0.01]) + def test_non_positive_height_raises(self, height: float) -> None: + """A non-positive hill attracts the walker back where it has been.""" + with pytest.raises(ValueError, match="height must be positive"): + _metad(height=height) + + @pytest.mark.parametrize("gamma", [1.0, 0.5, -2.0]) + def test_bias_factor_at_or_below_one_raises(self, gamma: float) -> None: + """gamma = 1 divides by zero in the well-tempered height.""" + with pytest.raises(ValueError, match="bias_factor must be greater"): + _metad(bias_factor=gamma) + + @pytest.mark.parametrize("sigma", [0.0, -0.3, [0.2, 0.0]]) + def test_non_positive_sigma_raises(self, sigma) -> None: + """A zero width is a delta function with no usable gradient.""" + with pytest.raises(ValueError, match="sigma must be positive"): + _metad(sigma=sigma) + + def test_unknown_storage_raises(self) -> None: + with pytest.raises(ValueError, match="storage must be one of"): + _metad(storage="lru") + + def test_unknown_history_raises(self) -> None: + with pytest.raises(ValueError, match="history must be one of"): + _metad(history="global") + + @pytest.mark.parametrize("storage", ["preallocated", "fifo"]) + def test_bounded_storage_requires_max_hills(self, storage: str) -> None: + """A ceiling policy without a ceiling is meaningless.""" + with pytest.raises(ValueError, match="needs an explicit max_hills"): + WellTemperedMetaDynamicsBias( + cv=_pair_cv(), + height=0.05, + sigma=0.2, + temperature=300.0, + storage=storage, + ) + + def test_grow_defaults_its_chunk(self) -> None: + """grow is the one policy that can run without an explicit capacity.""" + bias = WellTemperedMetaDynamicsBias( + cv=_pair_cv(), height=0.05, sigma=0.2, temperature=300.0, storage="grow" + ) + assert bias.capacity > 0 + + @pytest.mark.parametrize("frequency", [0, -1]) + def test_non_positive_update_frequency_raises(self, frequency: int) -> None: + with pytest.raises(ValueError, match="update_frequency must be at least 1"): + _metad(update_frequency=frequency) + + def test_negative_ramp_raises(self) -> None: + with pytest.raises(ValueError, match="ramp_depositions must be non-negative"): + _metad(ramp_depositions=-1) + + def test_mixin_order_is_correct(self) -> None: + """AdaptivePotentialMixin must precede nn.Module in the MRO.""" + mro = WellTemperedMetaDynamicsBias.__mro__ + from nvalchemi.enhanced_sampling import AdaptivePotentialMixin + + assert mro.index(AdaptivePotentialMixin) < mro.index(torch.nn.Module) + + +# =========================================================================== +# 2. Well-tempered metadynamics: energy and hill scaling +# =========================================================================== + + +class TestWellTemperedEnergy: + """Hill accumulation, the well-tempered height, and derived forces.""" + + def test_empty_history_is_exactly_zero(self, device: str) -> None: + """A bias with no hills must contribute nothing, not a small number.""" + bias = _metad(device) + result = bias.evaluate(_pair_batch([1.0, 2.0], device)) + assert torch.count_nonzero(result.energy) == 0 + assert torch.count_nonzero(result.forces) == 0 + + def test_single_hill_matches_closed_form(self, device: str) -> None: + """V(s) = h * exp(-(s - c)^2 / 2 sigma^2) for one deposited hill.""" + bias = _metad(device, sigma=0.5) + deposit = _pair_batch([1.0], device) + bias.update(deposit, bias.evaluate(deposit)) + + probe = _pair_batch([1.0, 1.5, 3.0], device) + got = bias.evaluate(probe).energy.reshape(-1) + expected = torch.tensor( + [ + 0.05 * math.exp(-((s - 1.0) ** 2) / (2 * 0.5**2)) + for s in (1.0, 1.5, 3.0) + ], + device=got.device, + dtype=got.dtype, + ) + assert torch.allclose(got, expected, atol=1e-6) + + def test_well_tempered_height_matches_formula(self, device: str) -> None: + """h = h0 exp(-V / (kB T (gamma - 1))) at the deposition point.""" + gamma, temperature, h0 = 10.0, 300.0, 0.05 + bias = _metad(device, bias_factor=gamma, temperature=temperature, height=h0) + deposit = _pair_batch([1.0], device) + + bias.update(deposit, bias.evaluate(deposit)) + assert float(bias.hill_heights[0]) == pytest.approx(h0, abs=1e-9) + + # The second hill lands on top of the first, so V = h0 there. + bias.update(deposit, bias.evaluate(deposit)) + expected = h0 * math.exp(-h0 / (KB_EV * temperature * (gamma - 1.0))) + assert float(bias.hill_heights[1]) == pytest.approx(expected, rel=1e-6) + + def test_heights_decay_monotonically(self, device: str) -> None: + """Repeated deposition in one spot must yield shrinking hills.""" + bias = _metad(device, max_hills=8) + deposit = _pair_batch([1.0], device) + for _ in range(8): + bias.update(deposit, bias.evaluate(deposit)) + heights = bias.hill_heights.tolist() + assert heights == sorted(heights, reverse=True) + assert heights[-1] < heights[0] + + def test_standard_metadynamics_keeps_constant_height(self, device: str) -> None: + """bias_factor=None is the gamma -> infinity limit: no damping.""" + bias = _metad(device, bias_factor=None, max_hills=8) + deposit = _pair_batch([1.0], device) + for _ in range(4): + bias.update(deposit, bias.evaluate(deposit)) + assert torch.allclose( + bias.hill_heights[:4], torch.full_like(bias.hill_heights[:4], 0.05) + ) + + def test_forces_match_numerical_gradient(self, device: str) -> None: + """Forces are -dE/dx of the same energy the bias reports.""" + bias = _metad(device, sigma=0.4, max_hills=8) + for d in (1.0, 1.4): + frame = _pair_batch([d], device) + bias.update(frame, bias.evaluate(frame)) + + batch = _pair_batch([1.2], device) + result = bias.evaluate(batch) + + eps = 1e-4 + step = _pair_batch([1.2 + eps], device) + back = _pair_batch([1.2 - eps], device) + numerical = ( + float(bias.energy(step).sum()) - float(bias.energy(back).sum()) + ) / (2 * eps) + # Atom 1 carries the whole +x displacement of the pair distance. + assert float(result.forces[1, 0]) == pytest.approx(-numerical, abs=1e-4) + + def test_bias_is_repulsive(self, device: str) -> None: + """The force must push the walker away from a deposited hill.""" + bias = _metad(device, sigma=0.3) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + outward = bias.evaluate(_pair_batch([1.1], device)) + # Sitting just beyond the hill, the pair is pushed further apart. + assert float(outward.forces[1, 0]) > 0.0 + + def test_multidimensional_cv(self, device: str) -> None: + """A 2-component CV uses one sigma per component.""" + idx_a, idx_b = torch.tensor([0, 1]), torch.tensor([0, 2]) + + def cv(batch: Batch) -> Tensor: + return torch.cat( + [pair_distance(batch, idx_a), pair_distance(batch, idx_b)], dim=-1 + ) + + data = AtomicData( + positions=torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.0, 0.0]]), + atomic_numbers=torch.ones(3, dtype=torch.long), + ) + batch = Batch.from_data_list([data]).to(device) + bias = WellTemperedMetaDynamicsBias( + cv=cv, + height=0.05, + sigma=torch.tensor([0.3, 0.6]), + temperature=300.0, + max_hills=4, + ).to(device) + bias.update(batch, bias.evaluate(batch)) + assert tuple(bias.hill_centers.shape) == (4, 2) + assert float(bias.evaluate(batch).energy.sum()) == pytest.approx(0.05, abs=1e-6) + + +# =========================================================================== +# 3. Well-tempered metadynamics: CV dimensionality +# =========================================================================== + + +class TestCVDimensionality: + """``sigma`` and ``periods`` must match the CV, not broadcast against it.""" + + @staticmethod + def _cv_2d(batch: Batch) -> Tensor: + """A two-component CV: the pair distance and twice the pair distance.""" + d = pair_distance(batch, torch.tensor([0, 1], device=batch.positions.device)) + return torch.cat([d, 2.0 * d], dim=-1) + + def test_scalar_cv_with_two_component_sigma_raises(self, device: str) -> None: + """The classic silent failure: a [B, 1] CV broadcast into a [H, 2] table. + + The CV value is duplicated into both hill columns and the exponent + sums two terms instead of one, so the Gaussian silently narrows to + ``exp(-d^2/2 * (1/s0^2 + 1/s1^2))``. Nothing in the run reports it. + """ + bias = _metad(device, sigma=torch.tensor([0.2, 0.4])) + with pytest.raises(ValueError, match="cv returns 1 component"): + bias.evaluate(_pair_batch([1.0], device)) + + def test_scalar_cv_with_multi_component_periods_raises(self, device: str) -> None: + bias = _metad(device, periods=torch.tensor([1.0, 2.0, 3.0])) + with pytest.raises(ValueError, match="periods has 3"): + bias.evaluate(_pair_batch([1.0], device)) + + def test_multi_component_cv_with_scalar_periods_raises(self, device: str) -> None: + """A 0 entry marks a component non-periodic, so one value cannot serve. + + Broadcasting a single period across every component would quietly + make them all periodic, which is not what a scalar can be assumed to + mean when the per-entry semantics carry that distinction. + """ + bias = _metad(device, cv=self._cv_2d, periods=torch.tensor([6.28])) + with pytest.raises(ValueError, match="periods has 1"): + bias.evaluate(_pair_batch([1.0], device)) + + def test_rank_one_cv_raises(self, device: str) -> None: + """A CV returning [B] rather than [B, 1] is a common slip.""" + idx = torch.tensor([0, 1]) + bias = _metad(device, cv=lambda b: pair_distance(b, idx).reshape(-1)) + with pytest.raises(ValueError, match=r"must return shape \[B, D\]"): + bias.evaluate(_pair_batch([1.0], device)) + + def test_validation_fires_on_update_too(self, device: str) -> None: + """update() also consumes CV output and must not deposit unchecked.""" + bias = _metad(device, sigma=torch.tensor([0.2, 0.4])) + frame = _pair_batch([1.0], device) + # update() ignores the result argument; an empty one is enough here. + empty = BiasResult(energy=torch.zeros(1, 1, device=device)) + with pytest.raises(ValueError, match="cv returns 1 component"): + bias.update(frame, empty) + + def test_scalar_sigma_serves_a_multi_component_cv(self, device: str) -> None: + """One width shared across components is legitimate and must work. + + The hill table takes its width from the CV, not from ``sigma`` — a + scalar says nothing about how many components there are. + """ + bias = _metad(device, cv=self._cv_2d, sigma=0.2, max_hills=4) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + assert tuple(bias.hill_centers.shape) == (4, 2) + assert bias.hill_centers[0].tolist() == pytest.approx([1.0, 2.0]) + + # Components move 0.3 and 0.6 from the hill center. + got = float(bias.evaluate(_pair_batch([1.3], device)).energy.sum()) + expected = 0.05 * math.exp(-0.5 * (0.3**2 + 0.6**2) / 0.2**2) + assert got == pytest.approx(expected, rel=1e-5) + + def test_per_component_sigma_weights_components_separately( + self, device: str + ) -> None: + bias = _metad( + device, cv=self._cv_2d, sigma=torch.tensor([0.2, 0.4]), max_hills=4 + ) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + got = float(bias.evaluate(_pair_batch([1.3], device)).energy.sum()) + expected = 0.05 * math.exp(-0.5 * (0.3**2 / 0.2**2 + 0.6**2 / 0.4**2)) + assert got == pytest.approx(expected, rel=1e-5) + + def test_scalar_sigma_scalar_cv_is_unchanged(self, device: str) -> None: + """The common case must keep its exact previous value.""" + bias = _metad(device, sigma=0.2, max_hills=4) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + assert tuple(bias.hill_centers.shape) == (4, 1) + got = float(bias.evaluate(_pair_batch([1.3], device)).energy.sum()) + assert got == pytest.approx(0.05 * math.exp(-0.5 * 0.3**2 / 0.2**2), rel=1e-5) + + def test_per_component_periods_apply_to_a_multi_component_cv( + self, device: str + ) -> None: + """A 0 period leaves that component unwrapped.""" + bias = _metad( + device, + cv=self._cv_2d, + sigma=0.2, + max_hills=4, + periods=torch.tensor([0.0, 6.28]), + ) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + assert tuple(bias.hill_centers.shape) == (4, 2) + assert float(bias.evaluate(frame).energy.sum()) == pytest.approx(0.05, abs=1e-6) + + def test_cv_changing_dimension_mid_run_raises(self, device: str) -> None: + """Existing hills are not comparable against a different-width CV.""" + idx = torch.tensor([0, 1]) + calls = {"n": 0} + + def switching(batch: Batch) -> Tensor: + calls["n"] += 1 + d = pair_distance(batch, idx) + return d if calls["n"] <= 2 else torch.cat([d, 2.0 * d], dim=-1) + + bias = _metad(device, cv=switching, sigma=0.2, max_hills=8) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + assert tuple(bias.hill_centers.shape) == (8, 1) + + with pytest.raises(RuntimeError, match="must keep its dimension"): + bias.update(frame, bias.evaluate(frame)) + + def test_free_energy_rejects_a_wrong_width_grid(self, device: str) -> None: + """Zeros from a width mismatch would read as a flat free energy.""" + bias = _metad(device, cv=self._cv_2d, sigma=0.2, max_hills=4) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + with pytest.raises(ValueError, match="deposited hills"): + bias.free_energy(torch.tensor([[1.0]], device=device)) + + def test_empty_history_still_evaluates_before_the_width_is_known( + self, device: str + ) -> None: + """Priming runs before any deposition resolves the CV width.""" + bias = _metad(device, cv=self._cv_2d, sigma=0.2, max_hills=4) + result = bias.evaluate(_pair_batch([1.0], device)) + assert torch.count_nonzero(result.energy) == 0 + + def test_multi_component_cv_compiles(self, device: str) -> None: + """The width check is a shape guard, not a data-dependent branch.""" + torch._dynamo.reset() + bias = _metad(device, cv=self._cv_2d, sigma=0.2, max_hills=4) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + batch = _pair_batch([1.3, 2.0], device) + compiled = torch.compile(bias.energy, fullgraph=True) + assert torch.allclose(compiled(batch), bias.energy(batch), atol=1e-6) + + +# =========================================================================== +# 4. Well-tempered metadynamics: periodic CVs +# =========================================================================== + + +class TestWellTemperedPeriodic: + """A hill near a branch cut must repel from both sides.""" + + @staticmethod + def _phase_cv(batch: Batch) -> Tensor: + """Treat atom 1's x coordinate as an angle-like periodic CV.""" + return batch.positions[batch.batch_ptr[:-1] + 1, 0:1] + + @staticmethod + def _phase_batch(values: list[float], device: str) -> Batch: + items = [ + AtomicData( + positions=torch.tensor([[0.0, 0.0, 0.0], [v, 0.0, 0.0]]), + atomic_numbers=torch.ones(2, dtype=torch.long), + ) + for v in values + ] + return Batch.from_data_list(items).to(device) + + def _biases(self, device: str) -> tuple: + periodic = WellTemperedMetaDynamicsBias( + cv=self._phase_cv, + height=0.05, + sigma=0.2, + temperature=300.0, + max_hills=4, + periods=torch.tensor([2 * math.pi]), + ).to(device) + plain = WellTemperedMetaDynamicsBias( + cv=self._phase_cv, + height=0.05, + sigma=0.2, + temperature=300.0, + max_hills=4, + ).to(device) + return periodic, plain + + def test_hill_wraps_across_the_branch_cut(self, device: str) -> None: + """A CV at -3.10 is 0.083 from a hill at +3.10, not 6.20 away.""" + periodic, plain = self._biases(device) + deposit = self._phase_batch([3.10], device) + for bias in (periodic, plain): + bias.update(deposit, bias.evaluate(deposit)) + + probe = self._phase_batch([-3.10], device) + assert float(periodic.evaluate(probe).energy.sum()) > 0.04 + assert float(plain.evaluate(probe).energy.sum()) < 1e-12 + + def test_wrapped_force_points_the_short_way(self, device: str) -> None: + """The repulsion must push away from the hill the short way round.""" + periodic, _ = self._biases(device) + deposit = self._phase_batch([3.10], device) + periodic.update(deposit, periodic.evaluate(deposit)) + + # Wrapped, -3.10 sits 0.083 *above* the hill at 3.10 (it is 3.183 + # once carried across the cut), so repulsion drives the CV further + # positive rather than back down the long way round. + forces = periodic.evaluate(self._phase_batch([-3.10], device)).forces + assert float(forces[1, 0]) > 0.0 + + # The unwrapped bias sees the same configuration 6.20 away and does + # essentially nothing, which is the failure the wrap exists to avoid. + _, plain = self._biases(device) + plain.update(deposit, plain.evaluate(deposit)) + plain_forces = plain.evaluate(self._phase_batch([-3.10], device)).forces + assert float(plain_forces[1, 0]) == pytest.approx(0.0, abs=1e-9) + + def test_non_periodic_component_is_left_unwrapped(self, device: str) -> None: + """A period of 0 marks a component as non-periodic.""" + bias = WellTemperedMetaDynamicsBias( + cv=self._phase_cv, + height=0.05, + sigma=0.2, + temperature=300.0, + max_hills=4, + periods=torch.tensor([0.0]), + ).to(device) + deposit = self._phase_batch([3.10], device) + bias.update(deposit, bias.evaluate(deposit)) + probe = self._phase_batch([-3.10], device) + assert float(bias.evaluate(probe).energy.sum()) < 1e-12 + + +# =========================================================================== +# 5. Well-tempered metadynamics: storage policies +# =========================================================================== + + +class TestStoragePolicies: + """preallocated raises, grow resizes, fifo evicts.""" + + def test_preallocated_overflow_raises(self, device: str) -> None: + """Silently dropping hills would change a converging run's physics.""" + bias = _metad(device, max_hills=2, storage="preallocated") + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + bias.update(frame, bias.evaluate(frame)) + with pytest.raises(RuntimeError, match="storage='preallocated'"): + bias.update(frame, bias.evaluate(frame)) + + def test_preallocated_overflow_leaves_state_untouched(self, device: str) -> None: + """The failed deposition must not half-apply.""" + bias = _metad(device, max_hills=2, storage="preallocated") + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + bias.update(frame, bias.evaluate(frame)) + before = (int(bias.hill_count), int(bias.deposits), bias.state_version) + with pytest.raises(RuntimeError): + bias.update(frame, bias.evaluate(frame)) + assert (int(bias.hill_count), int(bias.deposits), bias.state_version) == before + + def test_preallocated_capacity_never_changes(self, device: str) -> None: + """The compile-stable policy must keep tensor shapes fixed.""" + bias = _metad(device, max_hills=6, storage="preallocated") + frame = _pair_batch([1.0], device) + shapes = {tuple(bias.hill_centers.shape)} + for _ in range(6): + bias.update(frame, bias.evaluate(frame)) + shapes.add(tuple(bias.hill_centers.shape)) + assert shapes == {(6, 1)} + + def test_grow_extends_capacity(self, device: str) -> None: + """grow allocates another chunk instead of raising.""" + bias = _metad(device, max_hills=2, storage="grow") + frame = _pair_batch([1.0], device) + for _ in range(5): + bias.update(frame, bias.evaluate(frame)) + assert bias.capacity >= 5 + assert int(bias.hill_count) == 5 + + def test_grow_preserves_existing_hills(self, device: str) -> None: + """Resizing must copy the history, not restart it.""" + bias = _metad(device, max_hills=2, storage="grow", sigma=0.5) + for d in (1.0, 2.0): + frame = _pair_batch([d], device) + bias.update(frame, bias.evaluate(frame)) + before = bias.evaluate(_pair_batch([1.0], device)).energy.clone() + + bias.update( + _pair_batch([5.0], device), bias.evaluate(_pair_batch([5.0], device)) + ) + after = bias.evaluate(_pair_batch([1.0], device)).energy + assert bias.capacity > 2 + assert torch.allclose(before, after, atol=1e-6) + + def test_fifo_evicts_the_oldest(self, device: str) -> None: + """The ring keeps exactly the most recent max_hills deposits.""" + bias = _metad(device, max_hills=3, storage="fifo", sigma=0.2) + for d in (1.0, 2.0, 3.0, 4.0): + frame = _pair_batch([d], device) + bias.update(frame, bias.evaluate(frame)) + + assert bias.capacity == 3 + assert int(bias.hill_count) == 3 + centers = sorted(round(float(c), 3) for c in bias.hill_centers.reshape(-1)) + assert centers == [2.0, 3.0, 4.0] + # The evicted hill leaves no trace at its old location. + assert float(bias.evaluate(_pair_batch([1.0], device)).energy.sum()) < 1e-6 + + def test_fifo_ring_wraps_more_than_once(self, device: str) -> None: + """Slot assignment must stay correct after several wraps.""" + bias = _metad(device, max_hills=2, storage="fifo", sigma=0.2) + for d in (1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0): + frame = _pair_batch([d], device) + bias.update(frame, bias.evaluate(frame)) + centers = sorted(round(float(c), 3) for c in bias.hill_centers.reshape(-1)) + assert centers == [6.0, 7.0] + assert int(bias.hills_written) == 7 + + def test_fifo_refuses_free_energy(self, device: str) -> None: + """Discarded hills invalidate the well-tempered relation.""" + bias = _metad(device, max_hills=2, storage="fifo") + with pytest.raises(RuntimeError, match="not valid under storage='fifo'"): + bias.free_energy(torch.zeros(1, 1, device=device)) + + +# =========================================================================== +# 6. Well-tempered metadynamics: multi-walker history +# =========================================================================== + + +class TestMultiWalkerHistory: + """shared, walker-private, and state-owned hill visibility.""" + + @staticmethod + def _two_close_walkers(device: str) -> Batch: + """Two walkers at nearly the same CV, so cross-hills matter.""" + batch = _pair_batch([1.0, 1.05], device) + batch.walker_id = torch.tensor([0, 1], device=batch.positions.device) + batch.thermodynamic_state_id = torch.tensor( + [0, 1], device=batch.positions.device + ) + return batch + + def test_shared_history_sees_every_hill(self, device: str) -> None: + """The multiple-walker scheme: B walkers fill a basin B times faster.""" + batch = self._two_close_walkers(device) + shared = _metad(device, history="shared", sigma=0.5) + shared.update(batch, shared.evaluate(batch)) + assert int(shared.hill_count) == 2 + assert bool((shared.hill_owner[:2] == -1).all()) + + private = _metad(device, history="walker", sigma=0.5) + private.update(batch, private.evaluate(batch)) + + shared_e = shared.evaluate(batch).energy.reshape(-1) + private_e = private.evaluate(batch).energy.reshape(-1) + # Each walker feels its own hill either way; only shared adds the other. + assert bool((shared_e > private_e + 1e-4).all()) + + def test_walker_history_is_private(self, device: str) -> None: + """Under "walker", a hill is invisible to every other walker.""" + batch = self._two_close_walkers(device) + bias = _metad(device, history="walker", sigma=0.5) + bias.update(batch, bias.evaluate(batch)) + assert bias.hill_owner[:2].tolist() == [0, 1] + + # Probe walker 0 alone against the history: it must not feel hill 1. + probe = _pair_batch([1.05], device) + probe.walker_id = torch.tensor([0], device=probe.positions.device) + alone = float(bias.evaluate(probe).energy.sum()) + + probe.walker_id = torch.tensor([1], device=probe.positions.device) + owner = float(bias.evaluate(probe).energy.sum()) + assert owner > alone + + def test_state_history_tags_by_thermodynamic_state(self, device: str) -> None: + batch = self._two_close_walkers(device) + bias = _metad(device, history="state", sigma=0.5) + bias.update(batch, bias.evaluate(batch)) + assert bias.hill_owner[:2].tolist() == [0, 1] + + def test_walker_history_without_walker_id_raises(self, device: str) -> None: + """Defaulting a missing owner field to 0 collapses private histories. + + Every hill would be filed under one key, making the result + numerically identical to ``history="shared"`` — the opposite of what + was asked for, with nothing reporting it. + """ + bias = _metad(device, history="walker", sigma=0.5) + with pytest.raises(ValueError, match="needs batch.walker_id"): + bias.evaluate(_pair_batch([1.0, 1.05], device)) + + def test_state_history_without_state_id_raises(self, device: str) -> None: + bias = _metad(device, history="state", sigma=0.5) + with pytest.raises(ValueError, match="needs batch.thermodynamic_state_id"): + bias.evaluate(_pair_batch([1.0, 1.05], device)) + + def test_shared_history_needs_no_owner_field(self, device: str) -> None: + """The fields are meaningless under a shared history.""" + bias = _metad(device, history="shared", sigma=0.5) + frame = _pair_batch([1.0, 1.05], device) + bias.update(frame, bias.evaluate(frame)) + assert bias.hill_owner[:2].tolist() == [-1, -1] + + def test_wrong_length_owner_field_raises(self, device: str) -> None: + """A short owner tensor broadcasts across graphs instead of erroring.""" + bias = _metad(device, history="walker", sigma=0.5) + frame = _pair_batch([1.0, 1.05], device) + frame.walker_id = torch.tensor([7], device=frame.positions.device) + with pytest.raises(ValueError, match="has 1 entries but the batch has 2"): + bias.evaluate(frame) + + def test_update_also_rejects_a_missing_owner_field(self, device: str) -> None: + """Deposition must not file hills under a fabricated owner either.""" + bias = _metad(device, history="walker", sigma=0.5) + frame = _pair_batch([1.0, 1.05], device) + empty = BiasResult(energy=torch.zeros(2, 1, device=device)) + with pytest.raises(ValueError, match="needs batch.walker_id"): + bias.update(frame, empty) + + def test_runner_supplies_the_owner_fields(self, device: str) -> None: + """The stamp is what makes walker-private history work in a real run.""" + bias = _metad(device, name="meta", history="walker", update_frequency=1) + runner = EnhancedSampling(_make_dynamics(device), {"meta": bias}) + runner.run(_random_batch(device=device), n_steps=2) + + assert int(bias.deposits) == 2 + assert sorted(bias.hill_owner[:2].tolist()) == [0, 1] + + def test_state_history_declares_exchange_dependence(self) -> None: + """A per-state history changes which hills a swap exposes a walker to.""" + assert _metad(history="state").state_dependent_for_exchange is True + assert _metad(history="shared").state_dependent_for_exchange is False + assert _metad(history="walker").state_dependent_for_exchange is False + + +# =========================================================================== +# 7. Well-tempered metadynamics: ramping and free energy +# =========================================================================== + + +class TestRampAndFreeEnergy: + """Smooth hill activation and the free-energy estimator.""" + + def test_no_ramp_activates_immediately(self, device: str) -> None: + bias = _metad(device, ramp_depositions=0) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + assert float(bias.evaluate(frame).energy.sum()) == pytest.approx(0.05, abs=1e-6) + + def test_ramp_grows_the_contribution(self, device: str) -> None: + """A ramped hill must not switch on at full height where it landed.""" + bias = _metad(device, ramp_depositions=4) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + series = [] + for _ in range(4): + series.append(float(bias.evaluate(frame).energy.sum())) + bias.deposits += 1 + assert series == sorted(series) + assert series[0] < 0.05 * 0.5 + assert series[-1] == pytest.approx(0.05, abs=1e-6) + + def test_free_energy_uses_the_well_tempered_factor(self, device: str) -> None: + """F = -(gamma / (gamma - 1)) V.""" + gamma = 10.0 + bias = _metad(device, bias_factor=gamma) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + values = torch.tensor([[1.0]], device=device) + key = torch.full((1,), -1, dtype=torch.long, device=device) + bias_value = bias.gaussian_sum(values, key) + assert torch.allclose( + bias.free_energy(values), -bias_value * (gamma / (gamma - 1.0)) + ) + + def test_free_energy_of_standard_metadynamics(self, device: str) -> None: + """F = -V when there is no well-tempered damping.""" + bias = _metad(device, bias_factor=None) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + values = torch.tensor([[1.0]], device=device) + key = torch.full((1,), -1, dtype=torch.long, device=device) + assert torch.allclose(bias.free_energy(values), -bias.gaussian_sum(values, key)) + + def test_free_energy_is_lower_where_hills_accumulated(self, device: str) -> None: + """The estimator must report a basin where the bias filled one in.""" + bias = _metad(device, max_hills=16, sigma=0.3) + frame = _pair_batch([1.0], device) + for _ in range(8): + bias.update(frame, bias.evaluate(frame)) + probed = bias.free_energy(torch.tensor([[1.0], [4.0]], device=device)) + assert float(probed[0]) < float(probed[1]) + + +# =========================================================================== +# 8. Well-tempered metadynamics: compile and restart +# =========================================================================== + + +class TestWellTemperedCompile: + """energy() is the compile boundary and must hold fullgraph=True.""" + + def test_energy_compiles_fullgraph(self, device: str) -> None: + torch._dynamo.reset() + bias = _metad(device, sigma=0.4) + frame = _pair_batch([1.0], device) + bias.update(frame, bias.evaluate(frame)) + + batch = _pair_batch([1.2, 2.0], device) + eager = bias.energy(batch) + compiled = torch.compile(bias.energy, fullgraph=True)(batch) + assert torch.allclose(eager, compiled, atol=1e-6) + + def test_compiled_energy_tracks_new_hills(self, device: str) -> None: + """Depositing must change the compiled result, not hit a stale trace.""" + torch._dynamo.reset() + bias = _metad(device, sigma=0.4, max_hills=8) + compiled = torch.compile(bias.energy, fullgraph=True) + batch = _pair_batch([1.0], device) + assert float(compiled(batch).sum()) == pytest.approx(0.0, abs=1e-9) + + bias.update(batch, bias.evaluate(batch)) + assert float(compiled(batch).sum()) == pytest.approx(0.05, abs=1e-6) + + def test_grow_exhausts_the_dynamo_recompile_limit(self, device: str) -> None: + """Each growth retraces, and Dynamo caps retraces per code object. + + ``grow`` changes the hill-tensor shape on every resize, so a compiled + ``energy()`` recompiles each time and stops once it has exhausted + ``recompile_limit``. This is why ``preallocated`` is the production + default rather than a micro-optimisation: the failure arrives + mid-run, after the trajectory is already underway. + + Both relevant settings are pinned for the duration, because they are + **process-global** and ``nvalchemi.distributed.DistributedModel`` + changes both — ``recompile_limit`` to 64 and + ``force_parameter_static_shapes`` to ``False`` — so this test's + result would otherwise depend on whether a distributed test ran + first. ``force_parameter_static_shapes`` matters more than the + limit: with dynamic shapes the resize does not retrace at all, which + the companion test covers. + """ + limit = 3 + with torch._dynamo.config.patch( + recompile_limit=limit, force_parameter_static_shapes=True + ): + torch._dynamo.reset() + bias = _metad(device, sigma=0.4, storage="grow", max_hills=1) + compiled = torch.compile(bias.energy, fullgraph=True) + + survived = 0 + for i in range(limit + 3): + frame = _pair_batch([1.0 + 0.1 * i], device) + bias.update(frame, bias.evaluate(frame)) + try: + compiled(frame) + except ( + torch._dynamo.exc.Unsupported, + torch._dynamo.exc.FailOnRecompileLimitHit, + ): + break + survived += 1 + + assert survived == limit, (survived, limit) + assert bias.capacity == limit + 1 + torch._dynamo.reset() + + def test_dynamic_shapes_avoid_the_recompile_limit(self, device: str) -> None: + """With dynamic parameter shapes, ``grow`` does not retrace per resize. + + The hard failure above is the behaviour under torch's default of + static parameter shapes. Turning that off makes Dynamo trace the + hill-table dimension symbolically, so growth stops being a recompile + trigger — a real escape hatch, and the reason the documentation + offers it alongside ``preallocated`` rather than presenting the + limit as unconditional. + """ + limit = 3 + with torch._dynamo.config.patch( + recompile_limit=limit, force_parameter_static_shapes=False + ): + torch._dynamo.reset() + bias = _metad(device, sigma=0.4, storage="grow", max_hills=1) + compiled = torch.compile(bias.energy, fullgraph=True) + + for i in range(limit + 4): + frame = _pair_batch([1.0 + 0.1 * i], device) + bias.update(frame, bias.evaluate(frame)) + compiled(frame) # must not raise + + assert bias.capacity == limit + 4 + torch._dynamo.reset() + + def test_preallocated_never_retraces(self, device: str) -> None: + """The compile-stable policy holds one trace for the whole run.""" + with torch._dynamo.config.patch( + recompile_limit=3, force_parameter_static_shapes=True + ): + torch._dynamo.reset() + bias = _metad(device, sigma=0.4, storage="preallocated", max_hills=32) + compiled = torch.compile(bias.energy, fullgraph=True) + + # Far more depositions than the recompile limit would allow if + # the shapes moved; preallocated holds them fixed. + for i in range(12): + frame = _pair_batch([1.0 + 0.1 * i], device) + bias.update(frame, bias.evaluate(frame)) + compiled(frame) + + assert int(bias.hill_count) == 12 + assert tuple(bias.hill_centers.shape) == (32, 1) + assert int(bias.deposits) == 12 + torch._dynamo.reset() + torch._dynamo.reset() + + def test_per_state_history_compiles(self, device: str) -> None: + """The owner mask must not introduce a data-dependent branch.""" + torch._dynamo.reset() + bias = _metad(device, history="state", sigma=0.4) + batch = _pair_batch([1.0, 1.2], device) + batch.thermodynamic_state_id = torch.tensor( + [0, 1], device=batch.positions.device + ) + bias.update(batch, bias.evaluate(batch)) + compiled = torch.compile(bias.energy, fullgraph=True) + assert torch.allclose(compiled(batch), bias.energy(batch), atol=1e-6) + + +class TestWellTemperedRestart: + """state_dict / load_state_dict round trips, including resized buffers.""" + + def test_round_trip_reproduces_the_energy(self, device: str) -> None: + bias = _metad(device, max_hills=8, sigma=0.4) + for d in (1.0, 2.0, 3.0): + frame = _pair_batch([d], device) + bias.update(frame, bias.evaluate(frame)) + + restored = _metad(device, max_hills=8, sigma=0.4) + restored.load_state_dict(bias.state_dict()) + + probe = _pair_batch([1.0, 2.5], device) + assert torch.allclose( + restored.evaluate(probe).energy, bias.evaluate(probe).energy + ) + assert int(restored.hill_count) == int(bias.hill_count) + assert int(restored.deposits) == int(bias.deposits) + assert restored.state_version == bias.state_version + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("height", 0.9), + ("sigma", 0.9), + ("temperature", 900.0), + ("bias_factor", 2.0), + ("history", "walker"), + ("ramp_depositions", 5), + ], + ) + def test_restoring_a_different_setting_raises( + self, device: str, field: str, value: object + ) -> None: + """Hills are only meaningful under the settings that deposited them. + + ``bias_factor`` is the sharpest case: the stored heights were damped + under one gamma, and ``free_energy`` would apply a different ratio to + them without anything objecting. + """ + source = _metad(device, max_hills=8) + frame = _pair_batch([1.0], device) + source.update(frame, source.evaluate(frame)) + + target = _metad(device, max_hills=8, **{field: value}) + with pytest.raises(ValueError, match=field): + target.load_state_dict(source.state_dict()) + + def test_restoring_a_different_sigma_does_not_overwrite_it( + self, device: str + ) -> None: + """sigma is a buffer, so an unchecked load would replace it silently.""" + source = _metad(device, sigma=0.2, max_hills=8) + frame = _pair_batch([1.0], device) + source.update(frame, source.evaluate(frame)) + + target = _metad(device, sigma=0.9, max_hills=8) + with pytest.raises(ValueError, match="sigma"): + target.load_state_dict(source.state_dict()) + assert float(target.sigma[0]) == pytest.approx(0.9) + + def test_restoring_a_different_storage_policy_raises(self, device: str) -> None: + """Retention semantics differ, so the same hills mean different runs.""" + source = _metad(device, max_hills=8, storage="preallocated") + frame = _pair_batch([1.0], device) + source.update(frame, source.evaluate(frame)) + + target = _metad(device, max_hills=8, storage="fifo") + with pytest.raises(ValueError, match="storage"): + target.load_state_dict(source.state_dict()) + + def test_restoring_different_periods_raises(self, device: str) -> None: + source = _metad(device, max_hills=8, periods=torch.tensor([6.28])) + frame = _pair_batch([1.0], device) + source.update(frame, source.evaluate(frame)) + + target = _metad(device, max_hills=8) + with pytest.raises(ValueError, match="periods"): + target.load_state_dict(source.state_dict()) + + def test_capacity_is_not_part_of_the_fingerprint(self, device: str) -> None: + """``grow`` legitimately reaches a capacity the constructor never had.""" + source = _metad(device, max_hills=2, storage="grow", sigma=0.4) + for d in (1.0, 2.0, 3.0, 4.0, 5.0): + frame = _pair_batch([d], device) + source.update(frame, source.evaluate(frame)) + assert source.capacity > 2 + + target = _metad(device, max_hills=2, storage="grow", sigma=0.4) + target.load_state_dict(source.state_dict()) + assert target.capacity == source.capacity + + def test_round_trip_after_growth_resizes_buffers(self, device: str) -> None: + """A grown checkpoint has a capacity the constructor never produces.""" + bias = _metad(device, max_hills=2, storage="grow", sigma=0.4) + for d in (1.0, 2.0, 3.0, 4.0, 5.0): + frame = _pair_batch([d], device) + bias.update(frame, bias.evaluate(frame)) + assert bias.capacity > 2 + + restored = _metad(device, max_hills=2, storage="grow", sigma=0.4) + restored.load_state_dict(bias.state_dict()) + assert restored.capacity == bias.capacity + + probe = _pair_batch([1.0, 3.0], device) + assert torch.allclose( + restored.evaluate(probe).energy, bias.evaluate(probe).energy + ) + + def test_restart_continues_the_well_tempered_sequence(self, device: str) -> None: + """Heights after a restart must follow on, not reset to h0.""" + bias = _metad(device, max_hills=8) + frame = _pair_batch([1.0], device) + for _ in range(3): + bias.update(frame, bias.evaluate(frame)) + + restored = _metad(device, max_hills=8) + restored.load_state_dict(bias.state_dict()) + bias.update(frame, bias.evaluate(frame)) + restored.update(frame, restored.evaluate(frame)) + assert float(restored.hill_heights[3]) == pytest.approx( + float(bias.hill_heights[3]), rel=1e-9 + ) + + +# =========================================================================== +# 9. RMSD metadynamics: alignment +# =========================================================================== + + +class TestSquaredRMSD: + """The QCP kernel underneath the RMSD bias.""" + + def test_matches_svd_kabsch(self) -> None: + """QCP and an explicit det-corrected SVD Kabsch must agree.""" + torch.manual_seed(3) + coords = torch.randn(3, 6, 3, dtype=torch.float64) + refs = torch.randn(4, 6, 3, dtype=torch.float64) + refs = refs - refs.mean(dim=1, keepdim=True) + + got = _squared_rmsd(coords, refs) + + centered = coords - coords.mean(dim=1, keepdim=True) + expected = torch.zeros(3, 4, dtype=torch.float64) + for b in range(3): + for r in range(4): + cov = centered[b].T @ refs[r] + u, s, vt = torch.linalg.svd(cov) + sign = torch.sign(torch.det(u @ vt)) + trace = s[0] + s[1] + sign * s[2] + expected[b, r] = ( + (centered[b] ** 2).sum() + (refs[r] ** 2).sum() - 2 * trace + ) / 6 + assert torch.allclose(got, expected, atol=1e-10) + + def test_identical_structures_give_zero(self) -> None: + """Never negative: the clamp must absorb the rounding step.""" + torch.manual_seed(4) + refs = torch.randn(3, 5, 3, dtype=torch.float64) + refs = refs - refs.mean(dim=1, keepdim=True) + diagonal = _squared_rmsd(refs, refs).diagonal() + assert bool((diagonal >= 0).all()) + assert float(diagonal.abs().max()) < 1e-18 + + def test_reflection_is_not_treated_as_identical(self) -> None: + """A mirror image is a different structure; only proper rotations align.""" + torch.manual_seed(5) + ref = torch.randn(1, 5, 3, dtype=torch.float64) + ref = ref - ref.mean(dim=1, keepdim=True) + mirrored = ref.clone() + mirrored[..., 0] *= -1 + assert float(_squared_rmsd(mirrored, ref)) > 1e-3 + + +class TestRMSDInvariance: + """Translation and rotation must not change the bias.""" + + def _molecule(self, device: str, seed: int = 0) -> Batch: + torch.manual_seed(seed) + data = AtomicData( + positions=torch.randn(5, 3), + atomic_numbers=torch.ones(5, dtype=torch.long), + ) + return Batch.from_data_list([data]).to(device) + + def test_energy_is_invariant_to_rigid_motion(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4, ramp_depositions=0 + ).to(device) + frame = self._molecule(device) + bias.update(frame, bias.evaluate(frame)) + original = bias.evaluate(frame).energy.clone() + + torch.manual_seed(11) + rotation = _rot().to(device=frame.positions.device, dtype=frame.positions.dtype) + shift = torch.tensor([3.0, -2.0, 7.0], device=frame.positions.device) + moved = self._molecule(device) + moved.positions = frame.positions @ rotation.T + shift + + assert torch.allclose(bias.evaluate(moved).energy, original, atol=1e-5) + + def test_bias_exerts_no_net_force(self, device: str) -> None: + """A translation-invariant energy cannot push the molecule bodily.""" + bias = RMSDMetaDynamicsBias( + k_push=0.05, alpha=0.5, max_references=4, ramp_depositions=0 + ).to(device) + frame = self._molecule(device) + bias.update(frame, bias.evaluate(frame)) + + probe = self._molecule(device, seed=2) + forces = bias.evaluate(probe).forces + assert float(forces.sum(dim=0).abs().max()) < 1e-5 + + def test_forces_match_numerical_gradient(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.05, alpha=0.3, max_references=4, ramp_depositions=0 + ).to(device) + frame = self._molecule(device) + bias.update(frame, bias.evaluate(frame)) + + probe = self._molecule(device, seed=2) + analytic = bias.evaluate(probe).forces + + eps = 1e-4 + base = probe.positions.clone() + for atom, axis in ((0, 0), (3, 2)): + probe.positions = base.clone() + probe.positions[atom, axis] += eps + plus = float(bias.energy(probe).sum()) + probe.positions = base.clone() + probe.positions[atom, axis] -= eps + minus = float(bias.energy(probe).sum()) + assert float(analytic[atom, axis]) == pytest.approx( + -(plus - minus) / (2 * eps), abs=1e-4 + ) + + +# =========================================================================== +# 10. RMSD metadynamics: construction, selection, and periodicity +# =========================================================================== + + +class TestRMSDConstruction: + """Constructor validation for the RMSD bias.""" + + @pytest.mark.parametrize("k_push", [0.0, -0.01]) + def test_non_positive_k_push_raises(self, k_push: float) -> None: + with pytest.raises(ValueError, match="k_push must be positive"): + RMSDMetaDynamicsBias(k_push=k_push, alpha=0.5, max_references=4) + + @pytest.mark.parametrize("alpha", [0.0, -1.0]) + def test_non_positive_alpha_raises(self, alpha: float) -> None: + with pytest.raises(ValueError, match="alpha must be positive"): + RMSDMetaDynamicsBias(k_push=0.02, alpha=alpha, max_references=4) + + @pytest.mark.parametrize("storage", ["preallocated", "fifo"]) + def test_bounded_storage_requires_capacity(self, storage: str) -> None: + with pytest.raises(ValueError, match="needs an explicit max_references"): + RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, storage=storage) + + def test_empty_atom_indices_raises(self) -> None: + with pytest.raises(ValueError, match="atom_indices is empty"): + RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=4, + atom_indices=torch.tensor([], dtype=torch.long), + ) + + def test_duplicate_atom_indices_raise(self) -> None: + """A repeated atom is silently double-weighted in the RMSD.""" + with pytest.raises(ValueError, match="contains duplicates"): + RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=4, + atom_indices=torch.tensor([0, 1, 1]), + ) + + def test_negative_atom_indices_raise(self) -> None: + with pytest.raises(ValueError, match="must be non-negative"): + RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=4, + atom_indices=torch.tensor([0, -1]), + ) + + def test_default_storage_is_fifo(self) -> None: + """The xTB-compatible policy is the default here.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4) + assert bias.storage == "fifo" + + +class TestRMSDSelectionAndPeriodicity: + """Atom selection is per-graph, and periodic batches are rejected.""" + + @staticmethod + def _batch(device: str, sizes: list[int], seed: int = 0) -> Batch: + torch.manual_seed(seed) + items = [ + AtomicData( + positions=torch.randn(n, 3), + atomic_numbers=torch.ones(n, dtype=torch.long), + ) + for n in sizes + ] + return Batch.from_data_list(items).to(device) + + def test_selection_is_per_graph_local(self, device: str) -> None: + """Index 1 means atom 1 of each graph, not global atom 1.""" + bias = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=4, + atom_indices=torch.tensor([0, 1, 2]), + ramp_depositions=0, + ).to(device) + batch = self._batch(device, [4, 4]) + bias.update(batch, bias.evaluate(batch)) + assert tuple(bias.reference_coords.shape) == (4, 3, 3) + + def test_out_of_range_selection_raises(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4, atom_indices=torch.tensor([0, 7]) + ).to(device) + with pytest.raises(ValueError, match="local index 7"): + bias.evaluate(self._batch(device, [4, 4])) + + def test_ragged_batch_without_selection_raises(self, device: str) -> None: + """No fixed correspondence exists across differently sized graphs.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + with pytest.raises(ValueError, match="differing atom counts"): + bias.evaluate(self._batch(device, [4, 5])) + + def test_ragged_batch_with_selection_is_fine(self, device: str) -> None: + """A common selection restores the correspondence.""" + bias = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=4, + atom_indices=torch.tensor([0, 1, 2]), + ).to(device) + result = bias.evaluate(self._batch(device, [4, 5])) + assert result.energy.shape[0] == 2 + + @staticmethod + def _boxed(device: str, pbc: list[bool] | None, n_graphs: int = 2) -> Batch: + """Return a batch with a bounding box and the given PBC flags.""" + torch.manual_seed(0) + items = [] + for _ in range(n_graphs): + kwargs = { + "positions": torch.randn(4, 3), + "atomic_numbers": torch.ones(4, dtype=torch.long), + "cell": torch.eye(3).unsqueeze(0) * 20.0, + } + if pbc is not None: + kwargs["pbc"] = torch.tensor([pbc]) + items.append(AtomicData(**kwargs)) + return Batch.from_data_list(items).to(device) + + def test_fully_periodic_batch_is_rejected(self, device: str) -> None: + """Cartesian RMSD is undefined once atoms can cross a cell face.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + with pytest.raises(ValueError, match="not defined under periodic"): + bias.evaluate(self._boxed(device, [True, True, True])) + + def test_slab_is_rejected(self, device: str) -> None: + """Wrapping along any one axis is enough to break the metric.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + with pytest.raises(ValueError, match="not defined under periodic"): + bias.evaluate(self._boxed(device, [True, True, False])) + + def test_mixed_batch_is_rejected(self, device: str) -> None: + """One periodic graph poisons the batch; the bias cannot serve it.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + items = [ + AtomicData( + positions=torch.randn(4, 3), + atomic_numbers=torch.ones(4, dtype=torch.long), + cell=torch.eye(3).unsqueeze(0) * 20.0, + pbc=torch.tensor([[flag, flag, flag]]), + ) + for flag in (False, True) + ] + batch = Batch.from_data_list(items).to(device) + with pytest.raises(ValueError, match="not defined under periodic"): + bias.evaluate(batch) + + def test_bounding_box_with_pbc_false_is_accepted(self, device: str) -> None: + """A cell is a box; only pbc says whether atoms wrap through it. + + A solvated or boxed molecule carries a bounding-box cell with all + PBC flags false. That is the non-periodic case this bias exists + for, so rejecting it on the cell alone would refuse the intended + input. + """ + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4, ramp_depositions=0 + ).to(device) + batch = self._boxed(device, [False, False, False]) + + result = bias.evaluate(batch) + assert torch.count_nonzero(result.energy) == 0 + + # It must be usable, not merely accepted. + bias.update(batch, result) + assert int(bias.reference_count) == 2 + assert float(bias.evaluate(batch).energy.sum()) > 0.0 + + def test_cell_without_pbc_flags_is_refused_as_undeclared(self, device: str) -> None: + """A cell with no boundary condition has not said which case it is.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + with pytest.raises(ValueError, match="no pbc flags"): + bias.evaluate(self._boxed(device, None)) + + def test_zero_cell_is_not_periodic(self, device: str) -> None: + """A zero cell is how a molecular batch spells "no cell".""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + batch = self._batch(device, [4, 4]) + batch.cell = torch.zeros(2, 3, 3, device=batch.positions.device) + bias.evaluate(batch) + + def test_no_cell_at_all_is_not_periodic(self, device: str) -> None: + """The plain molecular case carries neither cell nor pbc.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + bias.evaluate(self._batch(device, [4, 4])) + + +# =========================================================================== +# 11. RMSD metadynamics: deposition, warm start, retention, restart +# =========================================================================== + + +class TestRMSDDeposition: + """Reference accumulation, ramping, FIFO retention, and warm starts.""" + + @staticmethod + def _molecule(device: str, seed: int = 0, n_graphs: int = 1) -> Batch: + torch.manual_seed(seed) + items = [ + AtomicData( + positions=torch.randn(5, 3), + atomic_numbers=torch.ones(5, dtype=torch.long), + ) + for _ in range(n_graphs) + ] + return Batch.from_data_list(items).to(device) + + def test_empty_history_is_zero(self, device: str) -> None: + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4).to(device) + result = bias.evaluate(self._molecule(device)) + assert torch.count_nonzero(result.energy) == 0 + + def test_deposited_structure_feels_full_amplitude(self, device: str) -> None: + """At RMSD zero the kernel is exp(0) = 1, so V = k_push.""" + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4, ramp_depositions=0 + ).to(device) + frame = self._molecule(device) + bias.update(frame, bias.evaluate(frame)) + assert float(bias.evaluate(frame).energy.sum()) == pytest.approx(0.02, abs=1e-7) + + def test_ramp_is_on_by_default(self, device: str) -> None: + """A new reference lands exactly where the system is standing.""" + bias = RMSDMetaDynamicsBias(k_push=0.02, alpha=0.5, max_references=4) + assert bias.ramp_depositions == 1 + + def test_ramp_delays_full_amplitude(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=8, ramp_depositions=4 + ).to(device) + frame = self._molecule(device) + bias.update(frame, bias.evaluate(frame)) + + series = [] + for _ in range(4): + series.append(float(bias.evaluate(frame).energy.sum())) + bias.deposits += 1 + assert series == sorted(series) + assert series[0] < 0.02 + assert series[-1] == pytest.approx(0.02, abs=1e-7) + + def test_fifo_retains_the_most_recent(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=2.0, max_references=2, ramp_depositions=0 + ).to(device) + frames = [self._molecule(device, seed=s) for s in range(4)] + for frame in frames: + bias.update(frame, bias.evaluate(frame)) + + assert bias.capacity == 2 + assert int(bias.reference_count) == 2 + assert int(bias.references_written) == 4 + # The oldest reference has been overwritten, so its site is free again. + assert float(bias.evaluate(frames[0]).energy.sum()) < float( + bias.evaluate(frames[3]).energy.sum() + ) + + def test_preallocated_overflow_raises(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=2, storage="preallocated" + ).to(device) + frame = self._molecule(device) + bias.update(frame, bias.evaluate(frame)) + bias.update(frame, bias.evaluate(frame)) + with pytest.raises(RuntimeError, match="storage='preallocated'"): + bias.update(frame, bias.evaluate(frame)) + + def test_warm_start_references_are_active_immediately(self, device: str) -> None: + """A seeded reference is history, not a fresh deposit needing a ramp.""" + torch.manual_seed(0) + reference = torch.randn(1, 5, 3) + bias = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=4, + references=reference, + ramp_depositions=4, + ).to(device) + assert int(bias.reference_count) == 1 + + frame = self._molecule(device) + frame.positions = reference[0].to(frame.positions.device) + assert float(bias.evaluate(frame).energy.sum()) == pytest.approx(0.02, abs=1e-7) + + def test_warm_start_is_translation_normalised(self, device: str) -> None: + """Seeded references are centered on storage like deposited ones.""" + torch.manual_seed(0) + reference = torch.randn(1, 5, 3) + 100.0 + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4, references=reference + ).to(device) + assert float(bias.reference_coords[0].mean(dim=0).abs().max()) < 1e-5 + + def test_warm_start_shape_is_validated(self) -> None: + with pytest.raises(ValueError, match=r"shape \[R, M, 3\]"): + RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4, references=torch.randn(5, 3) + ) + + def test_warm_start_beyond_capacity_raises(self) -> None: + with pytest.raises(ValueError, match="exceed max_references"): + RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=2, + references=torch.randn(3, 5, 3), + ) + + def test_warm_start_conflicting_with_selection_raises(self) -> None: + with pytest.raises(ValueError, match="but atom_indices selects"): + RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=4, + atom_indices=torch.tensor([0, 1]), + references=torch.randn(1, 5, 3), + ) + + def test_walker_history_is_private(self, device: str) -> None: + batch = self._molecule(device, n_graphs=2) + bias = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=8, + history="walker", + ramp_depositions=0, + ).to(device) + batch.walker_id = torch.tensor([0, 1], device=batch.positions.device) + bias.update(batch, bias.evaluate(batch)) + assert bias.reference_owner[:2].tolist() == [0, 1] + + shared = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=8, ramp_depositions=0 + ).to(device) + shared.update(batch, shared.evaluate(batch)) + assert bool( + (shared.evaluate(batch).energy >= bias.evaluate(batch).energy - 1e-9).all() + ) + + def test_walker_history_without_walker_id_raises(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=8, history="walker" + ).to(device) + with pytest.raises(ValueError, match="needs batch.walker_id"): + bias.evaluate(self._molecule(device, n_graphs=2)) + + def test_state_history_without_state_id_raises(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=8, history="state" + ).to(device) + with pytest.raises(ValueError, match="needs batch.thermodynamic_state_id"): + bias.evaluate(self._molecule(device, n_graphs=2)) + + def test_wrong_length_owner_field_raises(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=8, history="walker" + ).to(device) + batch = self._molecule(device, n_graphs=2) + batch.walker_id = torch.tensor([4], device=batch.positions.device) + with pytest.raises(ValueError, match="has 1 entries but the batch has 2"): + bias.evaluate(batch) + + def test_state_history_declares_exchange_dependence(self) -> None: + assert ( + RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4, history="state" + ).state_dependent_for_exchange + is True + ) + assert ( + RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=4 + ).state_dependent_for_exchange + is False + ) + + +class TestRMSDRestart: + """state_dict / load_state_dict round trips for the reference set.""" + + @staticmethod + def _molecule(device: str, seed: int = 0) -> Batch: + torch.manual_seed(seed) + data = AtomicData( + positions=torch.randn(5, 3), atomic_numbers=torch.ones(5, dtype=torch.long) + ) + return Batch.from_data_list([data]).to(device) + + def test_round_trip_reproduces_the_energy(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=8, ramp_depositions=0 + ).to(device) + for seed in range(3): + frame = self._molecule(device, seed=seed) + bias.update(frame, bias.evaluate(frame)) + + restored = RMSDMetaDynamicsBias( + k_push=0.02, alpha=0.5, max_references=8, ramp_depositions=0 + ).to(device) + restored.load_state_dict(bias.state_dict()) + + probe = self._molecule(device, seed=1) + assert torch.allclose( + restored.evaluate(probe).energy, bias.evaluate(probe).energy, atol=1e-7 + ) + assert int(restored.reference_count) == int(bias.reference_count) + assert restored.state_version == bias.state_version + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("k_push", 5.0), + ("alpha", 99.0), + ("history", "walker"), + ("ramp_depositions", 7), + ], + ) + def test_restoring_a_different_setting_raises( + self, device: str, field: str, value: object + ) -> None: + """The geometries survive but what they do changes.""" + base = { + "k_push": 0.02, + "alpha": 0.5, + "max_references": 8, + "ramp_depositions": 0, + } + source = RMSDMetaDynamicsBias(**base).to(device) + frame = self._molecule(device) + source.update(frame, source.evaluate(frame)) + + target = RMSDMetaDynamicsBias(**{**base, field: value}).to(device) + with pytest.raises(ValueError, match=field): + target.load_state_dict(source.state_dict()) + + def test_restoring_a_different_atom_selection_raises(self, device: str) -> None: + """atom_indices is a buffer; the rejection must precede the overwrite.""" + source = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=8, + atom_indices=torch.tensor([0, 1, 2]), + ).to(device) + frame = self._molecule(device) + source.update(frame, source.evaluate(frame)) + + target = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=8, + atom_indices=torch.tensor([1, 2, 3]), + ).to(device) + with pytest.raises(ValueError, match="atom_indices"): + target.load_state_dict(source.state_dict()) + assert target.atom_indices.tolist() == [1, 2, 3] + + def test_round_trip_after_growth_resizes_buffers(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=2, + storage="grow", + ramp_depositions=0, + ).to(device) + for seed in range(5): + frame = self._molecule(device, seed=seed) + bias.update(frame, bias.evaluate(frame)) + assert bias.capacity > 2 + + restored = RMSDMetaDynamicsBias( + k_push=0.02, + alpha=0.5, + max_references=2, + storage="grow", + ramp_depositions=0, + ).to(device) + restored.load_state_dict(bias.state_dict()) + assert restored.capacity == bias.capacity + + probe = self._molecule(device, seed=2) + assert torch.allclose( + restored.evaluate(probe).energy, bias.evaluate(probe).energy, atol=1e-7 + ) + + +# =========================================================================== +# 12. Runner integration: deposition schedule and force priming +# =========================================================================== + + +class TestRunnerIntegration: + """The runner drives deposition exactly once per due step.""" + + def test_deposition_follows_update_frequency(self, device: str) -> None: + bias = _metad(device, name="meta", update_frequency=3, max_hills=64) + runner = EnhancedSampling(_make_dynamics(device), {"meta": bias}) + runner.run(_random_batch(device=device), n_steps=9) + + # Two walkers deposit one hill each per due step. + assert int(bias.deposits) == 3 + assert int(bias.hill_count) == 6 + + def test_no_deposition_during_priming(self, device: str) -> None: + """Priming evaluates forces; it must not advance the history.""" + bias = _metad(device, name="meta", update_frequency=1) + runner = EnhancedSampling(_make_dynamics(device), {"meta": bias}) + runner.prime_forces(_random_batch(device=device)) + assert int(bias.deposits) == 0 + + def test_hills_are_deposited_at_post_step_coordinates(self, device: str) -> None: + """observation_stage is AFTER_STEP: a hill marks where it arrived.""" + assert _metad(device).observation_stage.name == "AFTER_STEP" + + def test_new_hill_is_felt_on_the_next_step(self, device: str) -> None: + """Depositing bumps the state version, so the runner re-primes forces.""" + bias = _metad(device, name="meta", update_frequency=1, sigma=0.6, height=0.5) + runner = EnhancedSampling( + _make_dynamics(device), {"meta": bias}, prime_after_update=True + ) + batch = _random_batch(device=device) + runner.prime_forces(batch) + before = int(bias.state_version) + + runner.run(batch, n_steps=1, prime=False) + assert int(bias.state_version) > before + assert float(runner.last_outputs["bias/meta/energy"].abs().sum()) > 0.0 + + def test_total_is_physical_plus_bias(self, device: str) -> None: + bias = _metad(device, name="meta", update_frequency=1, sigma=0.6) + runner = EnhancedSampling(_make_dynamics(device), {"meta": bias}) + runner.run(_random_batch(device=device), n_steps=4) + + outputs = runner.last_outputs + assert float(outputs["total/energy"].sum()) == pytest.approx( + float(outputs["physical/energy"].sum()) + + float(outputs["bias_total/energy"].sum()), + abs=1e-5, + ) + + def test_rmsd_bias_runs_through_the_runner(self, device: str) -> None: + bias = RMSDMetaDynamicsBias( + k_push=0.03, + alpha=0.4, + max_references=16, + update_frequency=2, + name="rmsd", + ).to(device) + runner = EnhancedSampling(_make_dynamics(device), {"rmsd": bias}) + runner.run(_random_batch(device=device), n_steps=6) + + assert int(bias.deposits) == 3 + assert int(bias.reference_count) == 6 + assert float(runner.last_outputs["bias/rmsd/energy"].abs().sum()) > 0.0 + + def test_two_metadynamics_biases_compose(self, device: str) -> None: + """Independent schedules, summed against unmodified physical output.""" + meta = _metad(device, name="meta", update_frequency=2, max_hills=64) + rmsd = RMSDMetaDynamicsBias( + k_push=0.03, + alpha=0.4, + max_references=16, + update_frequency=3, + name="rmsd", + ).to(device) + runner = EnhancedSampling(_make_dynamics(device), {"meta": meta, "rmsd": rmsd}) + runner.run(_random_batch(device=device), n_steps=6) + + assert int(meta.deposits) == 3 + assert int(rmsd.deposits) == 2 + + def test_checkpoint_round_trip_preserves_history( + self, tmp_path, device: str + ) -> None: + """The runner's Zarr checkpoint must carry the hill table.""" + bias = _metad(device, name="meta", update_frequency=1, max_hills=32) + runner = EnhancedSampling( + _make_dynamics(device), {"meta": bias}, steps_per_epoch=4 + ) + batch = _random_batch(device=device) + batch = runner.run(batch, n_steps=4) + + path = tmp_path / "metad.zarr" + runner.checkpoint(path, batch) + + fresh_bias = _metad(device, name="meta", update_frequency=1, max_hills=32) + fresh = EnhancedSampling( + _make_dynamics(device), {"meta": fresh_bias}, steps_per_epoch=4 + ) + restored = fresh.restore(path, device=device) + + assert int(fresh_bias.hill_count) == int(bias.hill_count) + assert int(fresh_bias.deposits) == int(bias.deposits) + assert fresh_bias.state_version == bias.state_version + assert torch.allclose( + fresh_bias.evaluate(restored).energy, + bias.evaluate(restored).energy, + atol=1e-6, + ) diff --git a/test/enhanced_sampling/test_runner.py b/test/enhanced_sampling/test_runner.py new file mode 100644 index 00000000..6e5fe89e --- /dev/null +++ b/test/enhanced_sampling/test_runner.py @@ -0,0 +1,903 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the ``EnhancedSampling`` runner and the adaptive battery. + +Covers walker identity stamping, the force-step ordering guarantees, +exactly-once ``update()`` delivery, observation staging, force priming, +epoch commits, and ``warm_start``. +""" + +from __future__ import annotations + +import pytest +import torch +from torch import Tensor + +from nvalchemi.data import AtomicData, Batch +from nvalchemi.dynamics import NVTLangevin +from nvalchemi.dynamics.base import DynamicsStage +from nvalchemi.enhanced_sampling import ( + AdaptivePotentialMixin, + BiasResult, + ConservativeBias, + EnhancedSampling, + HarmonicUmbrellaBias, + pair_distance, +) +from nvalchemi.models.demo import DemoModel, DemoModelWrapper + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_batch( + n_graphs: int = 2, + atoms_per_graph: int = 4, + device: str = "cpu", + seed: int = 0, + with_cell: bool = False, +) -> Batch: + """Return a batch with the output buffers dynamics writes back into.""" + torch.manual_seed(seed) + data_list = [] + for _ in range(n_graphs): + kwargs = { + "positions": torch.randn(atoms_per_graph, 3), + "atomic_numbers": torch.full((atoms_per_graph,), 6, dtype=torch.long), + "atomic_masses": torch.ones(atoms_per_graph), + "forces": torch.zeros(atoms_per_graph, 3), + "energy": torch.zeros(1, 1), + } + if with_cell: + kwargs["cell"] = torch.eye(3).unsqueeze(0) * 10.0 + kwargs["pbc"] = torch.tensor([[True, True, True]]) + kwargs["stress"] = torch.zeros(1, 3, 3) + data = AtomicData(**kwargs) + data.add_node_property("velocities", torch.zeros(atoms_per_graph, 3)) + data_list.append(data) + return Batch.from_data_list(data_list).to(device) + + +def _make_dynamics(device: str = "cpu") -> NVTLangevin: + model = DemoModelWrapper(DemoModel()).to(device) + return NVTLangevin(model=model, dt=0.1, temperature=300.0, friction=0.1) + + +class _ConstantForceBias(ConservativeBias): + """E = c * sum(x) — constant, known bias force of -c along x.""" + + def __init__(self, coefficient: float = 1.0, name: str = "constant_force") -> None: + super().__init__(name=name) + self.coefficient = coefficient + + def energy(self, current: Batch) -> Tensor: + ptr = current.batch_ptr + return torch.stack( + [ + self.coefficient * current.positions[ptr[b] : ptr[b + 1], 0].sum() + for b in range(current.num_graphs) + ] + ).unsqueeze(-1) + + +class _RecordingAdaptiveBias(AdaptivePotentialMixin, ConservativeBias): + """Conservative and adaptive; records every update() call for assertions.""" + + def __init__( + self, + name: str = "recording", + update_frequency: int = 1, + observation_stage: DynamicsStage = DynamicsStage.AFTER_STEP, + bump: bool = True, + ) -> None: + super().__init__(name=name) + self.update_frequency = update_frequency + self.observation_stage = observation_stage + self._bump = bump + self.update_steps: list[int] = [] + self.observed_forces: list[Tensor] = [] + self.commit_calls = 0 + + def energy(self, current: Batch) -> Tensor: + return ( + torch.zeros( + current.num_graphs, + 1, + dtype=current.positions.dtype, + device=current.positions.device, + ) + + 0.0 * current.positions.sum() + ) + + def update(self, frames: Batch, result: BiasResult) -> None: + step = int(frames.sampling_step.reshape(-1)[0]) + self.update_steps.append(step) + forces = getattr(frames, "forces", None) + if forces is not None: + self.observed_forces.append(forces.clone()) + if self._bump: + self.bump_state_version() + + def commit_epoch(self) -> None: + self.commit_calls += 1 + + +# =========================================================================== +# 1. Construction and validation +# =========================================================================== + + +class TestRunnerConstruction: + """Runner rejects malformed bias mappings at construction.""" + + def test_empty_biases_allowed(self) -> None: + runner = EnhancedSampling(_make_dynamics(), {}) + assert runner.biases == {} + + def test_none_biases_allowed(self) -> None: + assert EnhancedSampling(_make_dynamics()).biases == {} + + def test_non_protocol_bias_raises(self) -> None: + class NotABias: + pass + + with pytest.raises(TypeError, match="does not satisfy the BiasPotential"): + EnhancedSampling(_make_dynamics(), {"x": NotABias()}) + + def test_key_name_mismatch_raises(self) -> None: + bias = _ConstantForceBias(name="actual_name") + with pytest.raises(ValueError, match="key and the bias name must agree"): + EnhancedSampling(_make_dynamics(), {"different_key": bias}) + + @pytest.mark.parametrize("steps_per_epoch", [0, -1, -10]) + def test_non_positive_steps_per_epoch_rejected(self, steps_per_epoch: int) -> None: + """It is a divisor: zero raises deep in a run rather than here.""" + with pytest.raises(ValueError, match="steps_per_epoch must be at least 1"): + EnhancedSampling(_make_dynamics(), {}, steps_per_epoch=steps_per_epoch) + + def test_steps_per_epoch_of_one_is_allowed(self) -> None: + runner = EnhancedSampling(_make_dynamics(), {}, steps_per_epoch=1) + assert runner.steps_per_epoch == 1 + + def test_hook_inserted_at_front(self) -> None: + """The bias hook must run before any other AFTER_COMPUTE hook.""" + dynamics = _make_dynamics() + runner = EnhancedSampling(dynamics, {}) + assert dynamics.hooks[0] is runner._hook + + def test_repr_lists_biases(self) -> None: + runner = EnhancedSampling( + _make_dynamics(), {"cf": _ConstantForceBias(name="cf")} + ) + assert "cf" in repr(runner) + + +# =========================================================================== +# 2. Walker identity +# =========================================================================== + + +class TestWalkerIdentity: + """Identity fields are stamped, and the persistent ones stay persistent.""" + + def test_fields_stamped(self, device: str) -> None: + batch = _make_batch(n_graphs=3, device=device) + runner = EnhancedSampling(_make_dynamics(device), {}) + runner.run(batch, n_steps=1) + + for field in ( + "walker_id", + "thermodynamic_state_id", + "sampling_step", + "exchange_segment", + "sampling_epoch", + ): + value = getattr(batch, field, None) + assert value is not None, f"{field} not stamped" + assert value.reshape(-1).shape == (3,) + + def test_walker_ids_unique_and_stable(self, device: str) -> None: + """walker_id is an identity: assigned once, never reshuffled.""" + batch = _make_batch(n_graphs=4, device=device) + runner = EnhancedSampling(_make_dynamics(device), {}) + runner.run(batch, n_steps=1) + first = batch.walker_id.clone() + assert len(set(first.reshape(-1).tolist())) == 4 + + runner.run(batch, n_steps=3, prime=False) + assert torch.equal(batch.walker_id, first) + + def test_user_supplied_state_ids_preserved(self, device: str) -> None: + """A caller assigning windows must not have them overwritten.""" + batch = _make_batch(n_graphs=3, device=device) + batch["thermodynamic_state_id"] = torch.tensor([2, 0, 1], device=device) + runner = EnhancedSampling(_make_dynamics(device), {}) + runner.run(batch, n_steps=2) + assert batch.thermodynamic_state_id.reshape(-1).tolist() == [2, 0, 1] + + def test_sampling_step_tracks_dynamics(self, device: str) -> None: + batch = _make_batch(device=device) + dynamics = _make_dynamics(device) + runner = EnhancedSampling(dynamics, {}) + runner.run(batch, n_steps=5) + assert int(batch.sampling_step.reshape(-1)[0]) == dynamics.step_count - 1 + + def test_epoch_advances_with_steps_per_epoch(self, device: str) -> None: + batch = _make_batch(device=device) + runner = EnhancedSampling(_make_dynamics(device), {}, steps_per_epoch=3) + runner.run(batch, n_steps=7) + assert int(batch.sampling_epoch.reshape(-1)[0]) == 6 // 3 + + +# =========================================================================== +# 3. Force-step ordering +# =========================================================================== + + +class TestForceStepOrdering: + """The guarantees that make multi-bias aggregation order-independent.""" + + def test_bias_force_applied_to_batch(self, device: str) -> None: + """Total force is physical + bias, with the documented sign.""" + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + dynamics = _make_dynamics(device) + runner = EnhancedSampling(dynamics, {"cf": _ConstantForceBias(2.0, name="cf")}) + runner.prime_forces(batch) + + physical = runner.last_outputs["physical/forces"] + total = batch.forces + # E = 2*sum(x) => F = -dE/dx = -2 on the x component only. + expected_bias = torch.zeros_like(total) + expected_bias[:, 0] = -2.0 + assert torch.allclose(total - physical, expected_bias, atol=1e-5) + + def test_registration_order_does_not_change_total(self, device: str) -> None: + """Two biases summed against unmodified outputs commute.""" + results = [] + for order in ([("a", 1.0), ("b", 3.0)], [("b", 3.0), ("a", 1.0)]): + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device, seed=7) + biases = {n: _ConstantForceBias(c, name=n) for n, c in order} + runner = EnhancedSampling(_make_dynamics(device), biases) + runner.prime_forces(batch) + results.append(batch.forces.clone()) + assert torch.allclose(results[0], results[1], atol=1e-6) + + def test_bias_cannot_observe_another_bias_force(self, device: str) -> None: + """Every bias sees the same unmodified physical forces.""" + seen: list[Tensor] = [] + + class _ForceReadingBias(ConservativeBias): + def __init__(self) -> None: + super().__init__(name="reader") + + def energy(self, current: Batch) -> Tensor: + seen.append(current.forces.clone()) + return ( + torch.zeros(current.num_graphs, 1, device=current.positions.device) + + 0.0 * current.positions.sum() + ) + + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + runner = EnhancedSampling( + _make_dynamics(device), + {"cf": _ConstantForceBias(5.0, name="cf"), "reader": _ForceReadingBias()}, + ) + runner.prime_forces(batch) + + physical = runner.last_outputs["physical/forces"] + assert seen, "reader bias never ran" + for observed in seen: + assert torch.allclose(observed, physical, atol=1e-6), ( + "a bias observed another bias's force contribution" + ) + + def test_diagnostics_namespaced(self, device: str) -> None: + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + runner.prime_forces(batch) + keys = set(runner.last_outputs) + assert "physical/forces" in keys + assert "bias/cf/forces" in keys + assert "bias_total/forces" in keys + assert "total/forces" in keys + + def test_total_is_physical_plus_bias(self, device: str) -> None: + """'total' must mean physical + bias, not the bias sum alone. + + These differ by exactly the physical contribution, so a caller + plotting 'total/energy' as the system energy would silently get only + the restraint term. + """ + batch = _make_batch(n_graphs=2, atoms_per_graph=3, device=device) + runner = EnhancedSampling( + _make_dynamics(device), + { + "a": _ConstantForceBias(1.0, name="a"), + "b": _ConstantForceBias(3.0, name="b"), + }, + ) + runner.prime_forces(batch) + + physical = runner.last_outputs["physical/forces"] + bias_total = runner.last_outputs["bias_total/forces"] + total = runner.last_outputs["total/forces"] + + assert torch.allclose(total, physical + bias_total, atol=1e-6) + assert torch.allclose(total, batch.forces, atol=1e-6) + assert not torch.allclose(total, bias_total, atol=1e-6), ( + "total/* is still the bias sum; physical contribution missing" + ) + + def test_bias_total_is_sum_across_biases(self, device: str) -> None: + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + runner = EnhancedSampling( + _make_dynamics(device), + { + "a": _ConstantForceBias(1.0, name="a"), + "b": _ConstantForceBias(3.0, name="b"), + }, + ) + runner.prime_forces(batch) + summed = ( + runner.last_outputs["bias/a/forces"] + runner.last_outputs["bias/b/forces"] + ) + assert torch.allclose( + runner.last_outputs["bias_total/forces"], summed, atol=1e-6 + ) + + def test_total_energy_matches_batch(self, device: str) -> None: + batch = _make_batch(n_graphs=2, device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(2.0, name="cf")} + ) + runner.prime_forces(batch) + assert torch.allclose( + runner.last_outputs["total/energy"].reshape(batch.energy.shape), + batch.energy, + atol=1e-6, + ) + + def test_total_snapshot_not_aliased_to_batch(self, device: str) -> None: + """A later in-place hook must not retroactively rewrite the record.""" + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + runner.prime_forces(batch) + recorded = runner.last_outputs["total/forces"].clone() + batch.forces.mul_(0.0) + assert torch.allclose(runner.last_outputs["total/forces"], recorded) + + def test_observables_namespaced_by_bias_name(self, device: str) -> None: + """Two biases emitting the same observable name must not collide.""" + + class _ObservableBias(ConservativeBias): + def __init__(self, name: str) -> None: + super().__init__(name=name) + + def energy(self, current: Batch) -> Tensor: + return ( + torch.zeros(current.num_graphs, 1, device=current.positions.device) + + 0.0 * current.positions.sum() + ) + + def evaluate(self, current: Batch) -> BiasResult: + base = super().evaluate(current) + import dataclasses + + return dataclasses.replace( + base, observables={"cv": torch.zeros(current.num_graphs)} + ) + + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), + {"first": _ObservableBias("first"), "second": _ObservableBias("second")}, + ) + runner.prime_forces(batch) + assert "bias/first/cv" in runner.last_outputs + assert "bias/second/cv" in runner.last_outputs + + def test_virial_result_rejected_with_named_error(self, device: str) -> None: + """The runner applies stress; a virial has no volume here to convert.""" + + class _VirialBias: + name = "virial_bias" + + def evaluate(self, current: Batch) -> BiasResult: + return BiasResult(virial=torch.zeros(current.num_graphs, 3, 3)) + + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"virial_bias": _VirialBias()} + ) + with pytest.raises(ValueError, match="applies 'stress' to the batch"): + runner.prime_forces(batch) + + def test_stress_applied_for_periodic_batch(self, device: str) -> None: + batch = _make_batch(n_graphs=1, device=device, with_cell=True) + idx = torch.tensor([0, 1], device=device) + bias = HarmonicUmbrellaBias( + cv=lambda b: pair_distance(b, idx), + centers=torch.tensor([1.0]), + stiffness=4.0, + name="umbrella", + ) + runner = EnhancedSampling(_make_dynamics(device), {"umbrella": bias}) + runner.prime_forces(batch) + assert "bias/umbrella/stress" in runner.last_outputs + assert torch.count_nonzero(runner.last_outputs["bias/umbrella/stress"]) > 0 + + +# =========================================================================== +# 4. Priming +# =========================================================================== + + +class TestMissingDestinationBuffers: + """A bias output with nowhere to go must fail loudly, never silently.""" + + @staticmethod + def _periodic_batch_without(field: str, device: str) -> Batch: + """Return a periodic batch missing exactly one output buffer.""" + kwargs = { + "positions": torch.tensor([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]]), + "atomic_numbers": torch.tensor([6, 6], dtype=torch.long), + "atomic_masses": torch.ones(2), + "forces": torch.zeros(2, 3), + "energy": torch.zeros(1, 1), + "cell": torch.eye(3).unsqueeze(0) * 10.0, + "pbc": torch.tensor([[True, True, True]]), + "stress": torch.zeros(1, 3, 3), + } + kwargs.pop(field) + data = AtomicData(**kwargs) + data.add_node_property("velocities", torch.zeros(2, 3)) + return Batch.from_data_list([data]).to(device) + + @staticmethod + def _umbrella(device: str, **kwargs) -> HarmonicUmbrellaBias: + idx = torch.tensor([0, 1], device=device) + return HarmonicUmbrellaBias( + cv=lambda b: pair_distance(b, idx), + centers=torch.tensor([2.0]), + stiffness=5.0, + name="umbrella", + **kwargs, + ) + + def test_missing_stress_buffer_raises(self, device: str) -> None: + """The dangerous case: a periodic bias stress with nowhere to land. + + Silently dropping it makes the bias invisible to an NPT/NPH barostat + — the same failure the deprecated BiasedPotentialHook has, reached + from a different direction. + """ + batch = self._periodic_batch_without("stress", device) + runner = EnhancedSampling( + _make_dynamics(device), {"umbrella": self._umbrella(device)} + ) + with pytest.raises(ValueError, match="no destination buffer"): + runner.prime_forces(batch) + + def test_missing_stress_error_names_bias_and_barostat_risk( + self, device: str + ) -> None: + batch = self._periodic_batch_without("stress", device) + runner = EnhancedSampling( + _make_dynamics(device), {"umbrella": self._umbrella(device)} + ) + with pytest.raises(ValueError) as excinfo: + runner.prime_forces(batch) + message = str(excinfo.value) + assert "'stress'" in message + assert "umbrella" in message + assert "barostat" in message + assert "compute_stress=False" in message + + def test_compute_stress_false_is_the_documented_escape(self, device: str) -> None: + """An NVT run with a periodic cell may legitimately want no stress.""" + batch = self._periodic_batch_without("stress", device) + runner = EnhancedSampling( + _make_dynamics(device), + {"umbrella": self._umbrella(device, compute_stress=False)}, + ) + runner.prime_forces(batch) + assert "bias_total/stress" not in runner.last_outputs + + def test_missing_energy_buffer_raises(self, device: str) -> None: + batch = self._periodic_batch_without("energy", device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + with pytest.raises(ValueError, match="'energy'"): + runner.prime_forces(batch) + + def test_error_fires_before_any_step_is_taken(self, device: str) -> None: + """Priming is what makes this a setup error, not a mid-run surprise.""" + batch = self._periodic_batch_without("stress", device) + dynamics = _make_dynamics(device) + runner = EnhancedSampling(dynamics, {"umbrella": self._umbrella(device)}) + with pytest.raises(ValueError, match="no destination buffer"): + runner.run(batch, n_steps=100) + assert dynamics.step_count == 0, "a step ran before the error surfaced" + + def test_allocated_stress_buffer_receives_contribution(self, device: str) -> None: + """The positive case: with the buffer present, stress lands on it.""" + batch = _make_batch(n_graphs=1, device=device, with_cell=True) + runner = EnhancedSampling( + _make_dynamics(device), {"umbrella": self._umbrella(device)} + ) + runner.prime_forces(batch) + assert torch.count_nonzero(batch.stress) > 0 + + +class TestPriming: + """Force priming, and the error when buffers are missing.""" + + def test_missing_forces_buffer_raises_named_error(self, device: str) -> None: + data = AtomicData( + positions=torch.randn(3, 3), + atomic_numbers=torch.full((3,), 6, dtype=torch.long), + ) + batch = Batch.from_data_list([data]).to(device) + runner = EnhancedSampling(_make_dynamics(device), {}) + with pytest.raises(ValueError, match="batch has no 'forces' field"): + runner.prime_forces(batch) + + def test_prime_populates_total_force(self, device: str) -> None: + batch = _make_batch(device=device) + assert torch.count_nonzero(batch.forces) == 0 + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + runner.prime_forces(batch) + assert torch.count_nonzero(batch.forces) > 0 + + def test_run_primes_by_default(self, device: str) -> None: + """Without priming, step 0 would integrate against a zero force buffer.""" + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + runner.run(batch, n_steps=1) + assert runner.last_outputs, "no force evaluation recorded" + + def test_prime_is_idempotent(self, device: str) -> None: + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + runner.prime_forces(batch) + first = batch.forces.clone() + runner.prime_forces(batch) + assert torch.allclose(batch.forces, first, atol=1e-6), ( + "priming twice at the same coordinates must not accumulate the bias" + ) + + +# =========================================================================== +# 5. Adaptive biases: update exactly once +# =========================================================================== + + +class TestAdaptiveUpdates: + """update() delivery, staging, and epoch commits.""" + + def test_update_called_once_per_step(self, device: str) -> None: + batch = _make_batch(device=device) + bias = _RecordingAdaptiveBias(update_frequency=1) + runner = EnhancedSampling(_make_dynamics(device), {"recording": bias}) + runner.run(batch, n_steps=5) + assert bias.update_steps == sorted(bias.update_steps) + assert len(bias.update_steps) == len(set(bias.update_steps)), ( + f"update() delivered more than once for some step: {bias.update_steps}" + ) + assert len(bias.update_steps) == 5 + + def test_update_frequency_respected(self, device: str) -> None: + batch = _make_batch(device=device) + bias = _RecordingAdaptiveBias(update_frequency=3) + runner = EnhancedSampling(_make_dynamics(device), {"recording": bias}) + runner.run(batch, n_steps=9) + assert all(step % 3 == 0 for step in bias.update_steps), bias.update_steps + + def test_after_compute_observation_sees_unbiased_forces(self, device: str) -> None: + """ABF's requirement: observe physical forces, never the bias's own.""" + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + observer = _RecordingAdaptiveBias( + name="observer", observation_stage=DynamicsStage.AFTER_COMPUTE + ) + runner = EnhancedSampling( + _make_dynamics(device), + {"cf": _ConstantForceBias(5.0, name="cf"), "observer": observer}, + ) + runner.run(batch, n_steps=2) + + assert observer.observed_forces + physical = runner.last_outputs["physical/forces"] + # The constant bias adds -5 to every x component; an observation that + # captured post-application forces would differ by exactly that. + for observed in observer.observed_forces: + assert not torch.allclose( + observed[:, 0], physical[:, 0] - 5.0, atol=1e-6 + ), "AFTER_COMPUTE observation captured biased forces" + + def test_after_step_update_receives_its_own_bias_result(self, device: str) -> None: + """update() must get the result its bias produced, at either stage. + + AFTER_STEP is the default and what metadynamics uses; a bias sizing + its next hill from the bias energy it just applied needs the real + value, not an empty placeholder. + """ + received: list[BiasResult] = [] + + class _ResultRecordingBias(AdaptivePotentialMixin, _ConstantForceBias): + def __init__(self) -> None: + super().__init__(2.0, name="recorder") + + def update(self, frames: Batch, result: BiasResult) -> None: + received.append(result) + + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"recorder": _ResultRecordingBias()} + ) + runner.run(batch, n_steps=2) + + assert received, "update() never called" + for result in received: + assert result.energy is not None, "update() got an empty BiasResult" + assert result.forces is not None + # E = 2 * sum(x) => F = -2 on x only. + assert torch.allclose( + result.forces[:, 0], + torch.full_like(result.forces[:, 0], -2.0), + atol=1e-5, + ) + + def test_after_compute_update_also_receives_its_result(self, device: str) -> None: + received: list[BiasResult] = [] + + class _ResultRecordingBias(AdaptivePotentialMixin, _ConstantForceBias): + observation_stage = DynamicsStage.AFTER_COMPUTE + + def __init__(self) -> None: + super().__init__(3.0, name="recorder") + + def update(self, frames: Batch, result: BiasResult) -> None: + received.append(result) + + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"recorder": _ResultRecordingBias()} + ) + runner.run(batch, n_steps=2) + + assert received + for result in received: + assert result.energy is not None + assert torch.allclose( + result.forces[:, 0], + torch.full_like(result.forces[:, 0], -3.0), + atol=1e-5, + ) + + def test_each_bias_gets_its_own_result_not_anothers(self, device: str) -> None: + """With several adaptive biases, results must not be crossed.""" + seen: dict[str, list[float]] = {"a": [], "b": []} + + def _make(tag: str, coefficient: float): + class _Bias(AdaptivePotentialMixin, _ConstantForceBias): + def __init__(self) -> None: + super().__init__(coefficient, name=tag) + + def update(self, frames: Batch, result: BiasResult) -> None: + seen[tag].append(float(result.forces[0, 0])) + + return _Bias() + + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"a": _make("a", 1.0), "b": _make("b", 4.0)} + ) + runner.run(batch, n_steps=2) + + assert all(abs(v - (-1.0)) < 1e-5 for v in seen["a"]), seen["a"] + assert all(abs(v - (-4.0)) < 1e-5 for v in seen["b"]), seen["b"] + + def test_non_adaptive_bias_never_asked_to_update(self, device: str) -> None: + batch = _make_batch(device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + runner.run(batch, n_steps=3) # must not raise NotImplementedError + assert runner._adaptive_biases() == {} + + def test_commit_epoch_fires_at_boundary(self, device: str) -> None: + batch = _make_batch(device=device) + bias = _RecordingAdaptiveBias() + runner = EnhancedSampling( + _make_dynamics(device), {"recording": bias}, steps_per_epoch=2 + ) + runner.run(batch, n_steps=6) + assert bias.commit_calls >= 2, ( + f"commit_epoch fired {bias.commit_calls} times over 3 epochs" + ) + + def test_state_version_bump_triggers_reprime(self, device: str) -> None: + batch = _make_batch(device=device) + bias = _RecordingAdaptiveBias(bump=True) + runner = EnhancedSampling( + _make_dynamics(device), {"recording": bias}, prime_after_update=True + ) + runner.run(batch, n_steps=2) + assert bias.state_version == 2 + + def test_no_bump_no_reprime(self, device: str) -> None: + """An update that does not change the applied bias must not re-prime.""" + batch = _make_batch(device=device) + bias = _RecordingAdaptiveBias(bump=False) + runner = EnhancedSampling(_make_dynamics(device), {"recording": bias}) + runner.run(batch, n_steps=3) + assert bias.state_version == 0 + assert len(bias.update_steps) == 3 + + +# =========================================================================== +# 6. warm_start and state +# =========================================================================== + + +class TestWarmStart: + """Approximate continuation from prior frames.""" + + def test_replays_frames_in_order(self, device: str) -> None: + history = _make_batch(n_graphs=4, device=device) + history["sampling_step"] = torch.arange(4, device=device) + bias = _RecordingAdaptiveBias() + runner = EnhancedSampling(_make_dynamics(device), {"recording": bias}) + runner.warm_start(history) + assert bias.update_steps == [0, 1, 2, 3] + + def test_warm_start_without_adaptive_is_noop(self, device: str) -> None: + history = _make_batch(n_graphs=2, device=device) + runner = EnhancedSampling( + _make_dynamics(device), {"cf": _ConstantForceBias(name="cf")} + ) + runner.warm_start(history) # must not raise + + def test_warm_start_after_restore_raises(self, device: str) -> None: + runner = EnhancedSampling(_make_dynamics(device), {}) + runner._restored = True + with pytest.raises(RuntimeError, match="mutually exclusive"): + runner.warm_start(_make_batch(device=device)) + + def test_checkpoint_requires_a_batch(self) -> None: + """checkpoint() is implemented; it needs something to save.""" + runner = EnhancedSampling(_make_dynamics(), {}) + with pytest.raises(RuntimeError, match="no batch to save"): + runner.checkpoint("x.zarr") + + def test_state_dict_includes_bias_state(self) -> None: + bias = _RecordingAdaptiveBias() + bias.bump_state_version() + runner = EnhancedSampling(_make_dynamics(), {"recording": bias}) + state = runner.state_dict() + assert state["biases"]["recording"]["state_version"] == 1 + + +class TestAdaptiveMixinComposition: + """The mixin must compose with nn.Module without losing state.""" + + def test_wrong_mro_order_raises(self) -> None: + """nn.Module.state_dict would otherwise silently shadow the mixin's.""" + with pytest.raises(TypeError, match="must come before nn.Module"): + + class Wrong(ConservativeBias, AdaptivePotentialMixin): + def energy(self, current: Batch) -> Tensor: + return torch.zeros(current.num_graphs, 1) + + def test_correct_mro_order_accepted(self) -> None: + class Right(AdaptivePotentialMixin, ConservativeBias): + def energy(self, current: Batch) -> Tensor: + return torch.zeros(current.num_graphs, 1) + + assert issubclass(Right, ConservativeBias) + + def test_state_dict_merges_buffers_and_history(self) -> None: + """Both halves must survive one round trip.""" + + class _Both(AdaptivePotentialMixin, ConservativeBias): + def __init__(self) -> None: + super().__init__(name="both") + self.register_buffer("center", torch.tensor([2.0])) + + def energy(self, current: Batch) -> Tensor: + return torch.zeros(current.num_graphs, 1) + + def update(self, frames: Batch, result: BiasResult) -> None: + self.bump_state_version() + + bias = _Both() + bias.bump_state_version() + bias.bump_state_version() + state = bias.state_dict() + assert "center" in state, "nn.Module buffers dropped" + assert state["state_version"] == 2, "bias history dropped" + + restored = _Both() + restored.load_state_dict(state) + assert restored.state_version == 2 + assert torch.allclose(restored.center, torch.tensor([2.0])) + + def test_load_state_dict_does_not_trip_strict_module_check(self) -> None: + """'state_version' must be stripped before nn.Module sees the mapping.""" + + class _Both(AdaptivePotentialMixin, ConservativeBias): + def __init__(self) -> None: + super().__init__(name="both") + self.register_buffer("center", torch.zeros(1)) + + def energy(self, current: Batch) -> Tensor: + return torch.zeros(current.num_graphs, 1) + + def update(self, frames: Batch, result: BiasResult) -> None: + pass + + bias = _Both() + bias.load_state_dict(_Both().state_dict()) # must not raise + + def test_adaptive_without_module_half(self) -> None: + """A non-conservative adaptive bias needs no nn.Module at all.""" + + class _ForceOnlyAdaptive(AdaptivePotentialMixin): + name = "force_only" + + def evaluate(self, current: Batch) -> BiasResult: + return BiasResult(forces=torch.zeros_like(current.positions)) + + def update(self, frames: Batch, result: BiasResult) -> None: + self.bump_state_version() + + bias = _ForceOnlyAdaptive() + assert not isinstance(bias, torch.nn.Module) + assert bias.state_dict() == {"state_version": 0} + + +# =========================================================================== +# 8. compile_biases +# =========================================================================== + + +class TestCompileBiases: + """compile_biases wraps energy(), never evaluate().""" + + def test_compiled_run_matches_eager(self, device: str) -> None: + torch._dynamo.reset() + results = [] + for compile_biases in (False, True): + batch = _make_batch(n_graphs=1, atoms_per_graph=3, device=device, seed=3) + runner = EnhancedSampling( + _make_dynamics(device), + {"cf": _ConstantForceBias(2.0, name="cf")}, + compile_biases=compile_biases, + ) + runner.prime_forces(batch) + results.append(batch.forces.clone()) + assert torch.allclose(results[0], results[1], atol=1e-5) diff --git a/test/hooks/test_shared_hooks.py b/test/hooks/test_shared_hooks.py index ce421a1e..70070d7a 100644 --- a/test/hooks/test_shared_hooks.py +++ b/test/hooks/test_shared_hooks.py @@ -16,6 +16,7 @@ from __future__ import annotations +import pytest import torch from nvalchemi.data import AtomicData, Batch @@ -146,8 +147,15 @@ def test_dynamics_stage_logs(self) -> None: # =========================================================================== +@pytest.mark.filterwarnings( + "ignore:BiasedPotentialHook is deprecated:DeprecationWarning" +) class TestBiasedPotentialHook: - """BiasedPotentialHook fires correctly under DynamicsStage.""" + """BiasedPotentialHook fires correctly under DynamicsStage. + + Deprecated in favour of :mod:`nvalchemi.enhanced_sampling`, but still + functional, so its behaviour stays under test. + """ def test_dynamics_stage_adds_bias(self) -> None: """Bias is applied to forces and energy under DynamicsStage.""" diff --git a/test/models/test_base.py b/test/models/test_base.py index a830680b..aa3bfe46 100644 --- a/test/models/test_base.py +++ b/test/models/test_base.py @@ -701,6 +701,21 @@ def test_retain_graph(self): forces2 = autograd_forces(energy, positions) torch.testing.assert_close(forces1, forces2) + def test_unused_positions_raise_by_default(self): + positions = torch.randn(3, 3, requires_grad=True) + other = torch.randn(3, 3, requires_grad=True) + energy = (other**2).sum() + with pytest.raises(RuntimeError, match="not have been used in the graph"): + autograd_forces(energy, positions) + + def test_allow_unused_gives_zero_forces(self): + positions = torch.randn(3, 3, requires_grad=True) + other = torch.randn(3, 3, requires_grad=True) + energy = (other**2).sum() + forces = autograd_forces(energy, positions, allow_unused=True) + assert forces.shape == positions.shape + torch.testing.assert_close(forces, torch.zeros_like(positions)) + class TestAutogradStresses: """Tests for autograd_stresses utility.""" @@ -728,6 +743,22 @@ def test_multiple_systems(self): stresses = autograd_stresses(energy, displacement, cell, num_graphs=3) assert stresses.shape == (3, 3, 3) + def test_unused_displacement_raises_by_default(self): + displacement = torch.zeros(1, 3, 3, requires_grad=True) + cell = torch.eye(3).unsqueeze(0) * 10.0 + energy = (torch.randn(3, 3, requires_grad=True) ** 2).sum() + with pytest.raises(RuntimeError, match="not have been used in the graph"): + autograd_stresses(energy, displacement, cell, num_graphs=1) + + def test_allow_unused_gives_zero_stress(self): + displacement = torch.zeros(1, 3, 3, requires_grad=True) + cell = torch.eye(3).unsqueeze(0) * 10.0 + energy = (torch.randn(3, 3, requires_grad=True) ** 2).sum() + stresses = autograd_stresses( + energy, displacement, cell, num_graphs=1, allow_unused=True + ) + torch.testing.assert_close(stresses, torch.zeros(1, 3, 3)) + class TestAutogradForcesAndStresses: """Tests for merged force and stress autograd utility.""" @@ -810,6 +841,37 @@ def test_retain_graph_allows_later_autograd_call(self): assert forces.shape == scaled_pos.shape + def test_unused_positions_raise_by_default(self): + positions = torch.randn(3, 3, requires_grad=True) + cell = torch.eye(3).unsqueeze(0) * 10.0 + batch_idx = torch.zeros(3, dtype=torch.long) + _, scaled_cell, displacement = prepare_strain(positions, cell, batch_idx) + # Energy depends on the strained cell only, never on positions. + energy = torch.linalg.det(scaled_cell).sum() + with pytest.raises(RuntimeError, match="not have been used in the graph"): + autograd_forces_and_stresses( + energy, positions, displacement, cell, num_graphs=1 + ) + + def test_allow_unused_zero_forces_with_real_stress(self): + """A cell-only energy yields zero forces but a genuine stress.""" + positions = torch.randn(3, 3, dtype=torch.float64, requires_grad=True) + box = 10.0 + cell = torch.eye(3, dtype=torch.float64).unsqueeze(0) * box + batch_idx = torch.zeros(3, dtype=torch.long) + _, scaled_cell, displacement = prepare_strain(positions, cell, batch_idx) + energy = torch.linalg.det(scaled_cell).sum() + + forces, stresses = autograd_forces_and_stresses( + energy, positions, displacement, cell, num_graphs=1, allow_unused=True + ) + + torch.testing.assert_close(forces, torch.zeros_like(positions)) + # E = V, so dE/deps = V * I and sigma = dE/deps / V = I. + torch.testing.assert_close( + stresses, torch.eye(3, dtype=torch.float64).unsqueeze(0) + ) + class TestSumOutputs: """Tests for sum_outputs utility."""