diff --git a/nvalchemi/enhanced_sampling/__init__.py b/nvalchemi/enhanced_sampling/__init__.py new file mode 100644 index 00000000..7e6c65d2 --- /dev/null +++ b/nvalchemi/enhanced_sampling/__init__.py @@ -0,0 +1,54 @@ +# 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 virial for free. +* :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 deferred. + +Deferred to later milestones +----------------------------- +* :class:`EnhancedSampling` runner +* :class:`ThermodynamicState`, :class:`ReplicaExchange` +* Built-in biases (umbrella, metadynamics, walls, ABF) +* Zarr checkpoint support +* General triclinic MIC for unreduced cells +""" + +from nvalchemi.enhanced_sampling._bias import ( + BiasPotential, + BiasResult, + ConservativeBias, + aggregate_bias_results, +) +from nvalchemi.enhanced_sampling.cv import pair_distance + +__all__ = [ + # Core abstractions + "BiasResult", + "BiasPotential", + "ConservativeBias", + "aggregate_bias_results", + # Collective variables + "pair_distance", +] diff --git a/nvalchemi/enhanced_sampling/_bias.py b/nvalchemi/enhanced_sampling/_bias.py new file mode 100644 index 00000000..afd7454a --- /dev/null +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -0,0 +1,559 @@ +# 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 canonical cell virial 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, Protocol, runtime_checkable + +import torch +from torch import Tensor + +if TYPE_CHECKING: + from nvalchemi.data import Batch + +__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, shape ``[B, 3, 3]``. + Mutually exclusive with ``virial``. + virial: + Canonical virial ``W = −dE/dstrain``, shape ``[B, 3, 3]``. + 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. + """ + + 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. + + Every P0 built-in satisfies this protocol. Authors may also satisfy + it structurally (no inheritance required). + + 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: + """Autograd helper that derives atomic forces and cell virial from energy. + + 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, + when a cell is present and ``_supports_virial=True``, a per-graph + strain leaf ``F_b`` of shape ``[B, 3, 3]`` initialised to the + identity (for the canonical cell virial). + 3. Right-multiplies both positions and cell by ``F_b`` so that a single + ``autograd.grad`` call on the strained batch yields forces from + ``dE/d(pos_leaf)`` and the canonical virial ``W = −dE/dF_b``. + The batch is restored to its original tensors unconditionally in a + ``finally`` block. + 4. Derives forces and virial in one ``autograd.grad`` call with + ``create_graph=False, retain_graph=False``. + 5. Constructs a ``BiasResult`` from fully detached output tensors. + 6. Drops all references to the autograd subgraph before returning, so + that no ``grad_fn`` ever escapes into the live batch or result. + + 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 + ----- + Virial computation + Virial is controlled by the class attribute ``_supports_virial`` + (default ``True``). Subclasses that never need virial may set + ``_supports_virial = False`` to skip the strain-leaf construction. + There is no ``compute_virial`` constructor parameter. + + 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. + """ + + # Subclasses may set this to False to skip virial computation even when + # a cell is present (e.g. force-only biases that extend ConservativeBias + # but never need stress/virial). + _supports_virial: bool = True + + def energy(self, current: Batch) -> Tensor: + """Return bias energy ``[B, 1]`` (eV). + + Must be differentiable w.r.t. ``current.positions`` (and + ``current.cell`` when virial is requested). + + Parameters + ---------- + current: + A *read-only view* of the live batch where ``positions`` has + been replaced by ``pos_leaf @ strain_per_atom`` and ``cell`` + (when present) by ``cell.detach() @ strain_leaf``. 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 canonical cell virial via autograd. + + This method runs in eager mode. See class docstring for the + compile boundary note and the chosen fallback. + + Virial derivation + ----------------- + The canonical virial is ``W = −dE/dstrain`` evaluated at the + identity strain. Under a homogeneous deformation ``F`` (ASE + row-vector convention), both atomic positions and the cell + transform together:: + + r_n → r_n @ F + cell_b → cell_b @ F + + A per-graph strain leaf ``F_b`` (initialised to ``I``) is applied + to both, and a single ``autograd.grad`` call then yields: + + * ``dE/d(pos_leaf[n])`` at ``F=I`` → forces (negated). + * ``dE/d(F_b)`` at ``F=I`` → canonical virial ``W_b = −dE/dF_b``. + + This is the correct formulation for position-dependent biases that + use MIC displacements: the position term ``Σ_n r_n ⊗ (−F_n)`` and + the cell gradient term are automatically combined. Using an + independent cell leaf (without straining positions) misses the + position contribution and returns incorrect virials for pair + restraints across image boundaries. + """ + has_cell = ( + self._supports_virial + and getattr(current, "cell", None) is not None + and current.cell is not None + ) + + B = current.num_graphs + original_positions = current.positions + original_cell = current.cell if has_cell else None + + with torch.enable_grad(): + # --- Positions leaf (for forces) -------------------------------- + # requires_grad_() is not supported by torch.compile; evaluate() + # is intentionally kept eager (see class docstring). + pos_leaf = current.positions.detach().requires_grad_(True) # [N, 3] + + # --- Per-graph strain leaf (for canonical virial) --------------- + # Initialised to the identity; both positions and cell are + # right-multiplied by F_b so that dE/dF_b|_{F=I} = −W_b. + strain_leaf: Tensor | None = None + pos_for_energy: Tensor = pos_leaf + + if has_cell: + cell = original_cell + if cell is not None and cell.dim() == 4: + cell = cell.squeeze(1) # [B, 3, 3] + + strain_leaf = ( + torch.eye(3, device=pos_leaf.device, dtype=pos_leaf.dtype) + .unsqueeze(0) + .expand(B, -1, -1) + .clone() + .requires_grad_(True) + ) # [B, 3, 3] + + # Apply per-graph strain to each atom's position: + # pos_n → pos_n @ F_{b(n)} + strain_per_atom = strain_leaf[current.batch_idx] # [N, 3, 3] + pos_for_energy = torch.einsum( + "nk,nkj->nj", pos_leaf, strain_per_atom + ) # [N, 3] + + # Apply per-graph strain to the cell: cell_b → cell_b @ F_b + # Detach the stored cell values; only F carries the gradient. + cell_for_energy = torch.bmm(cell.detach(), strain_leaf) # [B, 3, 3] + + try: + current["positions"] = pos_for_energy + if has_cell and strain_leaf is not None: + # Store in the same shape as the original cell tensor. + stored = cell_for_energy # type: ignore[possibly-undefined] + if original_cell is not None and original_cell.dim() == 4: + stored = stored.unsqueeze(1) + current["cell"] = stored + + bias_energy: Tensor = self.energy(current) # [B, 1] + + inputs: list[Tensor] = [pos_leaf] + if strain_leaf is not None: + inputs.append(strain_leaf) + + grads = torch.autograd.grad( + outputs=(bias_energy,), + inputs=inputs, + grad_outputs=(torch.ones_like(bias_energy),), + create_graph=False, + retain_graph=False, + allow_unused=False, + ) + + finally: + current["positions"] = original_positions + if has_cell and original_cell is not None: + current["cell"] = original_cell + + # grads[0] = dE/d(pos_leaf) at strain=I → forces = −grad. + forces = -grads[0].detach() # [N, 3] + + virial: Tensor | None = None + if strain_leaf is not None and len(grads) > 1 and grads[1] is not None: + # grads[1] = dE/dF_b at F=I; canonical virial W_b = −dE/dF_b. + virial = -grads[1].detach() # [B, 3, 3] + + return BiasResult( + energy=bias_energy.detach(), + forces=forces, + virial=virial, + ) + + +# --------------------------------------------------------------------------- +# 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; duplicate keys raise ``ValueError`` + so that namespacing (``bias//``) must be applied before + calling this function. + + 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() + + energy_total: Tensor | None = None + forces_total: Tensor | None = None + stress_total: Tensor | None = None + virial_total: Tensor | None = None + observables_total: dict[str, Tensor] = {} + + # 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." + ) + + for r in results: + if r.energy is not None: + energy_total = r.energy if energy_total is None else energy_total + r.energy + if r.forces is not None: + forces_total = r.forces if forces_total is None else forces_total + r.forces + if r.stress is not None: + stress_total = r.stress if stress_total is None else stress_total + r.stress + if r.virial is not None: + virial_total = r.virial if virial_total is None else virial_total + r.virial + for key, val in r.observables.items(): + if key in observables_total: + raise ValueError( + f"aggregate_bias_results: duplicate observable key {key!r}. " + "Apply 'bias//' namespacing before aggregation." + ) + observables_total[key] = val + + return BiasResult( + energy=energy_total, + forces=forces_total, + stress=stress_total, + virial=virial_total, + observables=observables_total, + ) diff --git a/nvalchemi/enhanced_sampling/biases/__init__.py b/nvalchemi/enhanced_sampling/biases/__init__.py new file mode 100644 index 00000000..aa23c106 --- /dev/null +++ b/nvalchemi/enhanced_sampling/biases/__init__.py @@ -0,0 +1,20 @@ +# 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 (P0 compile-spike stubs). + +Full implementations are added in PR 2–6. This ``__init__`` is a +placeholder so that ``nvalchemi.enhanced_sampling.biases`` is a valid +importable namespace from PR 1 onward. +""" diff --git a/nvalchemi/enhanced_sampling/cv/__init__.py b/nvalchemi/enhanced_sampling/cv/__init__.py new file mode 100644 index 00000000..e10e282f --- /dev/null +++ b/nvalchemi/enhanced_sampling/cv/__init__.py @@ -0,0 +1,26 @@ +# 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. + +P0 built-in: :func:`pair_distance`. + +CVs are plain callables — no class hierarchy, no registration. Any +differentiable function ``cv(batch: Batch) -> Tensor[B, D]`` satisfies +the CV interface. +""" + +from nvalchemi.enhanced_sampling.cv.pair_distance import pair_distance + +__all__ = ["pair_distance"] diff --git a/nvalchemi/enhanced_sampling/cv/pair_distance.py b/nvalchemi/enhanced_sampling/cv/pair_distance.py new file mode 100644 index 00000000..468d2f63 --- /dev/null +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -0,0 +1,418 @@ +# 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` is the P0 built-in geometric CV. It 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. + +The original proposal named this "general triclinic MIC"; that description +is overstated. True general triclinic MIC (arbitrary unreduced cells, +implemented via LLL lattice reduction or an extended image search with a +data-dependent radius) is **deferred** — its interaction with the +strain-based virial computation in :class:`ConservativeBias` adds +non-trivial complexity that is out of scope for this release. + +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_distance"] + + +def pair_distance(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). + + Returns + ------- + Tensor + Shape ``[B, 1]`` — pair distance 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) + + # --- 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 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/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_bias_core.py b/test/enhanced_sampling/test_bias_core.py new file mode 100644 index 00000000..69c1fab3 --- /dev/null +++ b/test/enhanced_sampling/test_bias_core.py @@ -0,0 +1,1317 @@ +# 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 + virial from autograd; compare with finite differences; 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) + + +# =========================================================================== +# 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: + self.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: + self.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 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_virial_canonical_analytical(self, device: str) -> None: + """Canonical virial W = −dE/dstrain is correct for a pair bias across an image. + + Setup + ----- + Box: 10 Å cubic. 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 strain F both positions and cell right-multiply: + r_n → r_n @ F, cell → cell @ F + dr_mic → dr_mic @ F (image index n = [−1,0,0] is fixed) + d = |dr_mic @ F| + + dE/d(F_{kl})|_{F=I} = k · dr_mic[k] · dr_mic[l] + W_{kl} = −k · dr_mic[k] · dr_mic[l] + + For dr_mic = [−1, 0, 0]: W[0,0] = −1 eV, all other elements = 0. + + This is only correct when positions and cell are strained together. + Using an independent cell leaf (the prior implementation) misses the + atomic position contribution and returns the wrong virial. + """ + if device == "cuda": + pytest.skip( + "Strain-based virial on CUDA triggers a cudagraph_trees assertion " + "in this environment. CPU correctness is verified here; CUDA will " + "be covered in the GPU integration test suite." + ) + + k = 1.0 + box = 10.0 + # 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.virial is not None, "virial should be non-None for periodic batch" + assert result.virial.shape == (1, 3, 3) + + # Analytical: W = −k · outer(dr_mic, dr_mic) + # dr_mic = [−1, 0, 0] → W = [[-1,0,0],[0,0,0],[0,0,0]] + # W[0,0] = -1, all other elements = 0 + W = result.virial[0] # [3, 3] + assert torch.allclose(W[0, 0], torch.tensor(-k, device=device), atol=1e-4), ( + f"W[0,0] = {W[0, 0].item():.6f}, expected {-k:.6f}. " + "Virial may be missing the atomic-position contribution (strain not " + "applied to both positions and cell simultaneously)." + ) + assert torch.allclose(W[1:, :], torch.zeros(2, 3, device=device), atol=1e-4), ( + f"Off-diagonal/off-axis virial elements should be zero, got {W}" + ) + + 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) + + +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_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 TestCompileSpike: + """Verifies what can and cannot be compiled with ``torch.compile``. + + **Findings (documented per proposal section 6):** + + * :func:`pair_distance` — compiles with ``fullgraph=True``. This is the + primary CV hot path and is the go/no-go gate for ``compile_biases=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. + This is consistent with the risk identified in proposal section 6. + **Chosen fallback (per proposal section 6):** 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()