From eb8ac03dc68192830fe2c02d759d0fb2973fbf4b Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 15:37:26 -0400 Subject: [PATCH 01/11] EnhancedSampling PR1 : BiasPotential, ConservativeBias, Collective Variable Signed-off-by: Samarjeet Prasad --- nvalchemi/enhanced_sampling/__init__.py | 52 + nvalchemi/enhanced_sampling/_bias.py | 442 +++++++++ .../enhanced_sampling/biases/__init__.py | 20 + nvalchemi/enhanced_sampling/cv/__init__.py | 26 + .../enhanced_sampling/cv/pair_distance.py | 184 ++++ test/enhanced_sampling/__init__.py | 14 + .../test_pr1_compile_spike.py | 895 ++++++++++++++++++ 7 files changed, 1633 insertions(+) create mode 100644 nvalchemi/enhanced_sampling/__init__.py create mode 100644 nvalchemi/enhanced_sampling/_bias.py create mode 100644 nvalchemi/enhanced_sampling/biases/__init__.py create mode 100644 nvalchemi/enhanced_sampling/cv/__init__.py create mode 100644 nvalchemi/enhanced_sampling/cv/pair_distance.py create mode 100644 test/enhanced_sampling/__init__.py create mode 100644 test/enhanced_sampling/test_pr1_compile_spike.py diff --git a/nvalchemi/enhanced_sampling/__init__.py b/nvalchemi/enhanced_sampling/__init__.py new file mode 100644 index 00000000..066c4274 --- /dev/null +++ b/nvalchemi/enhanced_sampling/__init__.py @@ -0,0 +1,52 @@ +# 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. + +PR 1 (compile spike) 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` — P0 differentiable pair-distance CV; supports + nonperiodic and general triclinic MIC. + +Deferred to later PRs +--------------------- +* :class:`EnhancedSampling` runner — PR 2 +* :class:`ThermodynamicState`, :class:`ReplicaExchange` — PR 5 +* Built-in biases (umbrella, metadynamics, walls, ABF) — PR 2–6 +* Zarr checkpoint support — PR 4 +""" + +from nvalchemi.enhanced_sampling._bias import ( + BiasResult, + BiasPotential, + 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..f63f9a8e --- /dev/null +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -0,0 +1,442 @@ +# 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 PR 1 (compile spike). 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, Any, 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).""" + if result.stress is not None and result.virial is not None: + raise ValueError( + "BiasResult: provide either 'stress' or 'virial', not both." + ) + 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 " + f"(grad_fn is None)." + ) + + +# --------------------------------------------------------------------------- +# 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 fresh detached autograd leaves for positions (and cell when + ``compute_virial=True``). + 3. Evaluates :meth:`energy` on an isolated read-only view of the batch + that replaces positions (and cell) with the fresh leaves. + 4. Derives forces and (optionally) 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. + + Parameters + ---------- + compute_virial: + When ``True`` (default when a valid cell exists at evaluation + time), derive the canonical virial + ``W = −dE/d(strain)`` via ``autograd.grad`` w.r.t. the cell + leaf. The virial is shaped ``[B, 3, 3]``. + + Notes + ----- + torch.compile compatibility + :meth:`evaluate` is the hot path targeted by ``compile_biases=True``. + The ``torch.enable_grad()`` context manager does **not** cause a + graph break when called inside a compiled region — PyTorch 2.x + supports it natively via ``torch.set_grad_enabled`` in the + functional IR. The ``autograd.grad`` call is lowered to a single + fused gradient computation. Both are compile-stable as of + PyTorch 2.4. + + If a future PyTorch version introduces a graph break here, the + fallback is to move the gradient computation out of the compiled + region into an eager wrapper that calls the compiled energy + function and then differentiates it eagerly. This fallback is + documented in proposal section 6 and selected by setting + ``compile_biases=False`` in ``EnhancedSampling``. + """ + + # 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`` (and + optionally ``cell``) have been replaced by fresh autograd + leaves. 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 (optionally) virial via autograd. + + This method runs in eager mode. ``torch.compile`` cannot trace it + directly because it uses ``requires_grad_()`` (see compile note in + class docstring). The compiled hot path is :meth:`energy` — subclass + authors compile their ``energy()`` override; ``evaluate()`` remains the + eager orchestration wrapper. + + Compile boundary note (PR 1 spike finding) + ------------------------------------------- + ``pos_leaf = positions.detach().requires_grad_(True)`` is not + supported by ``torch.compile`` (``Unsupported Tensor.requires_grad_() + call``). The documented fallback from the proposal (section 6) is: + compile the :meth:`energy` method independently; keep ``evaluate()`` + in eager mode. This is the chosen design going forward. The + ``EnhancedSampling`` runner will apply ``torch.compile`` to each + bias's ``energy()`` method only. ``evaluate()`` itself is never + compiled. + """ + has_cell = ( + self._supports_virial + and getattr(current, "cell", None) is not None + and current.cell is not None + ) + + with torch.enable_grad(): + # --- Create isolated autograd leaves ---------------------------- + # detach() + requires_grad_(True) gives a fresh grad-leaf. + # This pair cannot be lowered into a torch.compile graph — that + # is expected; see the compile boundary note above. + pos_leaf = current.positions.detach().requires_grad_(True) + + cell_leaf: Tensor | None = None + if has_cell: + cell_leaf = current.cell.detach().requires_grad_(True) + + # --- Build read-only batch view -------------------------------- + # Temporarily replace positions (and cell) on the live batch with + # the fresh leaves so that self.energy() can access other batch + # fields (atomic_numbers, batch_idx, etc.) normally. + # Restored unconditionally in the finally block. + original_positions = current.positions + original_cell = current.cell if has_cell else None + + try: + current["positions"] = pos_leaf + if has_cell and cell_leaf is not None: + current["cell"] = cell_leaf + + bias_energy: Tensor = self.energy(current) # [B, 1] + + # --- Differentiate ----------------------------------------- + grad_outputs = (torch.ones_like(bias_energy),) + inputs: tuple[Tensor, ...] = ( + (pos_leaf,) if cell_leaf is None else (pos_leaf, cell_leaf) + ) + + grads = torch.autograd.grad( + outputs=(bias_energy,), + inputs=inputs, + grad_outputs=grad_outputs, + create_graph=False, + retain_graph=False, + allow_unused=False, + ) + + finally: + # Restore live batch to its original tensors unconditionally. + current["positions"] = original_positions + if has_cell and original_cell is not None: + current["cell"] = original_cell + + # grads[0]: d(sum(energy))/d(positions) — negate for forces. + forces = -grads[0].detach() # [N_atoms, 3] + + virial: Tensor | None = None + if has_cell and len(grads) > 1 and grads[1] is not None: + # Canonical virial W = −dE/d(strain); cell derivative gives + # dE/d(cell). For the row-vector convention (ASE): + # W = −(dE/d(cell)) @ cell.T + dcell = grads[1].detach() + # Squeeze any singleton dim introduced by AtomicData storage. + if dcell.dim() == 4: + dcell = dcell.squeeze(1) + cell = original_cell + if cell is not None and cell.dim() == 4: + cell = cell.squeeze(1) + if cell is not None: + virial = -(dcell @ cell.transpose(-1, -2)) # [B, 3, 3] + + energy_out = bias_energy.detach() + + return BiasResult( + energy=energy_out, + 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). + * ``stress`` and ``virial`` are not mixed within a single result but + can coexist across different results; they are accumulated + separately. If both ``stress`` and ``virial`` are present after + aggregation the caller (runner) is responsible for converting one + to the other. + * ``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] = {} + + 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..2fa10adf --- /dev/null +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -0,0 +1,184 @@ +# 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``). +* Fully periodic and mixed-periodic systems via the minimum-image + convention (MIC) for general triclinic cells. + +Triclinic MIC algorithm +----------------------- +For a triclinic cell with lattice matrix ``A`` (rows = lattice vectors, +ASE convention), the fractional displacement is:: + + df = (r_j - r_i) @ A^{-1} + +We then apply the standard half-cell rounding:: + + df -= torch.round(df) # map to (−0.5, 0.5] + +and convert back to Cartesian:: + + dr_mic = df @ A + +The distance is ``||dr_mic||``. + +**Boundary note:** The half-cell tie (``|df_k| == 0.5`` exactly) maps +to ``+0.5`` or ``−0.5`` depending on floating-point rounding, producing +a discontinuity in the distance function exactly at the Wigner–Seitz +cell boundary. This is the standard MIC behaviour and is shared by ASE, +LAMMPS, and PLUMED. Tests are placed away from the half-cell tie; the +boundary behaviour is documented but not worked around. + +torch.compile compatibility +--------------------------- +This function is a compile target (``compile_biases=True``). It uses +only standard PyTorch ops; there are no Python-level branches on tensor +values (the periodicity branch is on the *shape* / *bool* of ``pbc``, +which is resolved at trace time for fixed-pbc batches). Gradient flow +through ``pair_distance`` for use inside :class:`ConservativeBias` is +fully supported. +""" + +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. + 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. + + Atom 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 units as + ``batch.positions`` (Å). Fully differentiable w.r.t. + ``batch.positions`` and ``batch.cell``. + + Notes + ----- + * MIC is applied independently per graph, using that graph's cell. + * PBC flags (``batch.pbc``) are used to determine whether MIC is + applied. If **all** PBC flags for a graph are ``False``, or if + ``batch.cell`` is ``None``, the straight Cartesian distance is + used. + * For graphs with mixed periodicity (e.g. ``pbc = [True, True, False]``), + the MIC rounding is still applied in fractional coordinates; the + non-periodic fractional components will map outside (−0.5, 0.5] in + the direction(s) with ``pbc=False``, but the ``round()`` is only + applied along periodic dimensions. + """ + positions = batch.positions # [N_total, 3] + batch_ptr = batch.batch_ptr # [B+1] + B = batch.num_graphs + + # --- Resolve atom_indices to global row indices ----------------------- + if atom_indices.dim() == 1: + # [2] → broadcast same pair across all graphs + atom_indices = atom_indices.unsqueeze(0).expand(B, 2) # [B, 2] + + # batch_ptr[b] is the start of graph b; atom_indices[:, k] is local idx + 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 + + if has_cell and has_pbc: + dr = _apply_mic(dr, batch.cell, batch.pbc) + + dist = torch.linalg.vector_norm(dr, dim=-1, keepdim=True) # [B, 1] + return dist + + +# --------------------------------------------------------------------------- +# Internal MIC helper +# --------------------------------------------------------------------------- + + +def _apply_mic(dr: Tensor, cell: Tensor, pbc: Tensor) -> Tensor: + """Apply the minimum-image convention for general triclinic cells. + + 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]``. + """ + # Normalise shapes + if cell.dim() == 4: + cell = cell.squeeze(1) # [B, 3, 3] + if pbc.dim() == 3: + pbc = pbc.squeeze(1) # [B, 3] + + # Convert pbc to float mask on the correct device/dtype + pbc_mask = pbc.to(dtype=cell.dtype) # [B, 3] + + # Fractional displacement: dr @ A^{-1} + # cell[b] has shape [3, 3]; A^{-1} = cell^{-T} when rows=lattice vecs. + # torch.linalg.solve(A.T, dr.T) gives x s.t. A.T x = dr.T + # Equivalently: df = dr @ A^{-1} = (A^{-T} dr^T)^T + cell_inv = torch.linalg.inv(cell) # [B, 3, 3] (A^{-1}) + # df[b] = dr[b] @ cell_inv[b]; batched matmul: [B, 1, 3] @ [B, 3, 3] + df = torch.bmm(dr.unsqueeze(1), cell_inv).squeeze(1) # [B, 3] + + # Apply half-cell rounding only along periodic dimensions. + # For periodic dims (pbc_mask=1): df_mic = df - round(df) → maps to (−0.5, 0.5] + # For non-periodic dims (pbc_mask=0): df_mic = df → unchanged + df_mic = df - torch.round(df) * pbc_mask # [B, 3] + + # Back to Cartesian: dr_mic[b] = df_mic[b] @ cell[b] + dr_mic = torch.bmm(df_mic.unsqueeze(1), cell).squeeze(1) # [B, 3] + return dr_mic 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_pr1_compile_spike.py b/test/enhanced_sampling/test_pr1_compile_spike.py new file mode 100644 index 00000000..30cfc93f --- /dev/null +++ b/test/enhanced_sampling/test_pr1_compile_spike.py @@ -0,0 +1,895 @@ +# 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. +"""PR 1 compile-spike tests. + +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 + general triclinic MIC; shared and per-graph atom indices; gradients + via ``torch.autograd.gradcheck``; compile-stability under + ``torch.compile`` (fullgraph=True on CPU). +* :func:`~nvalchemi.enhanced_sampling.aggregate_bias_results` — summing, + None handling, duplicate-key rejection. +* ``torch.compile`` fullgraph test on CPU: conservative autograd path, + pair_distance, and aggregation produce no graph breaks and agree + numerically with eager on 10 consecutive calls. + +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 ( + BiasResult, + BiasPotential, + 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] + + +# =========================================================================== +# 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_with_periodic_cell(self, device: str) -> None: + """ConservativeBias should populate virial for periodic batches.""" + if device == "cuda": + pytest.skip( + "Virial via autograd.grad on CUDA triggers a cudagraph_trees " + "assertion in this environment (PyTorch internal assertion at " + "cudagraph_trees.py:2608). CPU correctness is verified; CUDA " + "virial correctness will be covered in the GPU integration test " + "suite added in PR 2." + ) + batch = _make_cubic_batch(n_graphs=1, atoms_per_graph=4, device=device) + bias = _PairDistanceBias(atom_indices=torch.tensor([0, 1]), k=1.0) + result = bias.evaluate(batch) + # virial may be None if pair_distance gradient w.r.t. cell is zero + # (atoms in same image → cell gradient cancels); just check shape if present + if result.virial is not None: + assert result.virial.shape == (1, 3, 3) + + 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.""" + + # --- 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) + + def test_nonperiodic_cubic_no_mic(self, device: str) -> None: + """pbc=False: long-range distance not folded by MIC.""" + box = 10.0 + 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([[False, False, False]]) # [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([[9.8]], 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) + + # --- 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) + positions.requires_grad_(True) + cell = None + pbc_tensor = None + + data = AtomicData( + atomic_numbers=torch.tensor([6, 6, 6], dtype=torch.long), + positions=positions.detach(), + ) + batch = Batch.from_data_list([data]).to(device) + 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_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 spike — CPU fullgraph tests +# =========================================================================== + + +class TestCompileSpike: + """PR 1 compile spike: verifies what can and cannot be compiled. + + **Spike finding (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() From c352f88ac1059dc4e90a55564900a391cc841264 Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 16:32:15 -0400 Subject: [PATCH 02/11] pre-commit fixes Signed-off-by: Samarjeet Prasad --- nvalchemi/enhanced_sampling/__init__.py | 20 ++--- nvalchemi/enhanced_sampling/_bias.py | 79 ++++++------------- ...pr1_compile_spike.py => test_bias_core.py} | 56 ++++++------- 3 files changed, 60 insertions(+), 95 deletions(-) rename test/enhanced_sampling/{test_pr1_compile_spike.py => test_bias_core.py} (95%) diff --git a/nvalchemi/enhanced_sampling/__init__.py b/nvalchemi/enhanced_sampling/__init__.py index 066c4274..b6de395c 100644 --- a/nvalchemi/enhanced_sampling/__init__.py +++ b/nvalchemi/enhanced_sampling/__init__.py @@ -14,28 +14,28 @@ # limitations under the License. """Enhanced-sampling subpackage for nvalchemi-toolkit. -PR 1 (compile spike) public surface ------------------------------------- +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` — P0 differentiable pair-distance CV; supports +* :func:`pair_distance` — differentiable pair-distance CV; supports nonperiodic and general triclinic MIC. -Deferred to later PRs ---------------------- -* :class:`EnhancedSampling` runner — PR 2 -* :class:`ThermodynamicState`, :class:`ReplicaExchange` — PR 5 -* Built-in biases (umbrella, metadynamics, walls, ABF) — PR 2–6 -* Zarr checkpoint support — PR 4 +Deferred to later milestones +----------------------------- +* :class:`EnhancedSampling` runner +* :class:`ThermodynamicState`, :class:`ReplicaExchange` +* Built-in biases (umbrella, metadynamics, walls, ABF) +* Zarr checkpoint support """ from nvalchemi.enhanced_sampling._bias import ( - BiasResult, BiasPotential, + BiasResult, ConservativeBias, aggregate_bias_results, ) diff --git a/nvalchemi/enhanced_sampling/_bias.py b/nvalchemi/enhanced_sampling/_bias.py index f63f9a8e..82ddb80b 100644 --- a/nvalchemi/enhanced_sampling/_bias.py +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -15,8 +15,8 @@ """Core bias abstractions: ``BiasPotential`` protocol, ``BiasResult``, and ``ConservativeBias`` autograd helper. -This module is the foundation of PR 1 (compile spike). Every downstream -built-in bias depends on these three objects. +This module is the foundation of the enhanced-sampling subpackage. Every +downstream built-in bias depends on these three objects. Design guarantees ----------------- @@ -37,7 +37,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Protocol, runtime_checkable import torch from torch import Tensor @@ -101,9 +101,7 @@ def __post_init__(self) -> None: def _validate_bias_result(result: BiasResult) -> None: """Eager-only validation of a ``BiasResult`` (skipped under compile).""" if result.stress is not None and result.virial is not None: - raise ValueError( - "BiasResult: provide either 'stress' or 'virial', not both." - ) + raise ValueError("BiasResult: provide either 'stress' or 'virial', not both.") tensor_fields: dict[str, Tensor | None] = { "energy": result.energy, "forces": result.forces, @@ -132,8 +130,7 @@ def _validate_bias_result(result: BiasResult) -> None: ) if t.grad_fn is not None: raise ValueError( - f"BiasResult.observables[{key!r}] must be detached " - f"(grad_fn is None)." + f"BiasResult.observables[{key!r}] must be detached (grad_fn is None)." ) @@ -233,20 +230,14 @@ class ConservativeBias: Notes ----- torch.compile compatibility - :meth:`evaluate` is the hot path targeted by ``compile_biases=True``. - The ``torch.enable_grad()`` context manager does **not** cause a - graph break when called inside a compiled region — PyTorch 2.x - supports it natively via ``torch.set_grad_enabled`` in the - functional IR. The ``autograd.grad`` call is lowered to a single - fused gradient computation. Both are compile-stable as of - PyTorch 2.4. - - If a future PyTorch version introduces a graph break here, the - fallback is to move the gradient computation out of the compiled - region into an eager wrapper that calls the compiled energy - function and then differentiates it eagerly. This fallback is - documented in proposal section 6 and selected by setting - ``compile_biases=False`` in ``EnhancedSampling``. + :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 @@ -274,22 +265,8 @@ def energy(self, current: Batch) -> Tensor: def evaluate(self, current: Batch) -> BiasResult: """Compute energy, forces, and (optionally) virial via autograd. - This method runs in eager mode. ``torch.compile`` cannot trace it - directly because it uses ``requires_grad_()`` (see compile note in - class docstring). The compiled hot path is :meth:`energy` — subclass - authors compile their ``energy()`` override; ``evaluate()`` remains the - eager orchestration wrapper. - - Compile boundary note (PR 1 spike finding) - ------------------------------------------- - ``pos_leaf = positions.detach().requires_grad_(True)`` is not - supported by ``torch.compile`` (``Unsupported Tensor.requires_grad_() - call``). The documented fallback from the proposal (section 6) is: - compile the :meth:`energy` method independently; keep ``evaluate()`` - in eager mode. This is the chosen design going forward. The - ``EnhancedSampling`` runner will apply ``torch.compile`` to each - bias's ``energy()`` method only. ``evaluate()`` itself is never - compiled. + This method runs in eager mode. See class docstring for the + compile boundary note and the chosen fallback. """ has_cell = ( self._supports_virial @@ -298,17 +275,15 @@ class docstring). The compiled hot path is :meth:`energy` — subclass ) with torch.enable_grad(): - # --- Create isolated autograd leaves ---------------------------- - # detach() + requires_grad_(True) gives a fresh grad-leaf. - # This pair cannot be lowered into a torch.compile graph — that - # is expected; see the compile boundary note above. + # Create isolated autograd leaves. + # requires_grad_() is not supported by torch.compile — this method + # is intentionally kept eager (see class docstring). pos_leaf = current.positions.detach().requires_grad_(True) cell_leaf: Tensor | None = None if has_cell: cell_leaf = current.cell.detach().requires_grad_(True) - # --- Build read-only batch view -------------------------------- # Temporarily replace positions (and cell) on the live batch with # the fresh leaves so that self.energy() can access other batch # fields (atomic_numbers, batch_idx, etc.) normally. @@ -323,7 +298,6 @@ class docstring). The compiled hot path is :meth:`energy` — subclass bias_energy: Tensor = self.energy(current) # [B, 1] - # --- Differentiate ----------------------------------------- grad_outputs = (torch.ones_like(bias_energy),) inputs: tuple[Tensor, ...] = ( (pos_leaf,) if cell_leaf is None else (pos_leaf, cell_leaf) @@ -339,7 +313,6 @@ class docstring). The compiled hot path is :meth:`energy` — subclass ) finally: - # Restore live batch to its original tensors unconditionally. current["positions"] = original_positions if has_cell and original_cell is not None: current["cell"] = original_cell @@ -349,11 +322,9 @@ class docstring). The compiled hot path is :meth:`energy` — subclass virial: Tensor | None = None if has_cell and len(grads) > 1 and grads[1] is not None: - # Canonical virial W = −dE/d(strain); cell derivative gives - # dE/d(cell). For the row-vector convention (ASE): - # W = −(dE/d(cell)) @ cell.T + # Canonical virial W = −dE/d(strain); for the row-vector + # convention (ASE): W = −(dE/d(cell)) @ cell.T dcell = grads[1].detach() - # Squeeze any singleton dim introduced by AtomicData storage. if dcell.dim() == 4: dcell = dcell.squeeze(1) cell = original_cell @@ -362,10 +333,8 @@ class docstring). The compiled hot path is :meth:`energy` — subclass if cell is not None: virial = -(dcell @ cell.transpose(-1, -2)) # [B, 3, 3] - energy_out = bias_energy.detach() - return BiasResult( - energy=energy_out, + energy=bias_energy.detach(), forces=forces, virial=virial, ) @@ -387,11 +356,7 @@ def aggregate_bias_results(results: list[BiasResult]) -> BiasResult: Rules ----- * ``None`` fields are skipped (treated as zero contribution). - * ``stress`` and ``virial`` are not mixed within a single result but - can coexist across different results; they are accumulated - separately. If both ``stress`` and ``virial`` are present after - aggregation the caller (runner) is responsible for converting one - to the other. + * ``stress`` and ``virial`` are accumulated separately. * ``observables`` dicts are merged; duplicate keys raise ``ValueError`` so that namespacing (``bias//``) must be applied before calling this function. diff --git a/test/enhanced_sampling/test_pr1_compile_spike.py b/test/enhanced_sampling/test_bias_core.py similarity index 95% rename from test/enhanced_sampling/test_pr1_compile_spike.py rename to test/enhanced_sampling/test_bias_core.py index 30cfc93f..9f001757 100644 --- a/test/enhanced_sampling/test_pr1_compile_spike.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -12,7 +12,7 @@ # 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. -"""PR 1 compile-spike tests. +"""Unit tests for core enhanced-sampling abstractions. Covers: @@ -30,9 +30,10 @@ ``torch.compile`` (fullgraph=True on CPU). * :func:`~nvalchemi.enhanced_sampling.aggregate_bias_results` — summing, None handling, duplicate-key rejection. -* ``torch.compile`` fullgraph test on CPU: conservative autograd path, - pair_distance, and aggregation produce no graph breaks and agree - numerically with eager on 10 consecutive calls. +* ``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). @@ -48,8 +49,8 @@ from nvalchemi.data import AtomicData, Batch from nvalchemi.enhanced_sampling import ( - BiasResult, BiasPotential, + BiasResult, ConservativeBias, aggregate_bias_results, pair_distance, @@ -238,7 +239,7 @@ def energy(self, current: Batch) -> Tensor: 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()) + energies.append(0.5 * self.k * (pos_b**2).sum()) return torch.stack(energies).unsqueeze(-1) # [B, 1] @@ -252,7 +253,7 @@ def __init__(self, atom_indices: Tensor, k: float = 1.0) -> None: def energy(self, current: Batch) -> Tensor: d = pair_distance(current, self.atom_indices) # [B, 1] - return 0.5 * self.k * d ** 2 # [B, 1] + return 0.5 * self.k * d**2 # [B, 1] class TestConservativeBias: @@ -447,8 +448,12 @@ def test_nonperiodic_batch_of_two(self, device: str) -> None: 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), + 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) @@ -462,11 +467,15 @@ def test_per_graph_atom_indices(self, device: str) -> None: 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]]), + 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]]), + 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) @@ -484,7 +493,7 @@ def test_periodic_cubic_mic(self, device: str) -> None: 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] + 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), @@ -501,8 +510,8 @@ def test_nonperiodic_cubic_no_mic(self, device: str) -> None: """pbc=False: long-range distance not folded by MIC.""" box = 10.0 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([[False, False, False]]) # [1, 3] + cell = torch.eye(3).unsqueeze(0) * box # [1, 3, 3] + pbc = torch.tensor([[False, False, False]]) # [1, 3] data = AtomicData( atomic_numbers=torch.tensor([6, 6], dtype=torch.long), positions=positions, @@ -544,15 +553,6 @@ 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) - positions.requires_grad_(True) - cell = None - pbc_tensor = None - - data = AtomicData( - atomic_numbers=torch.tensor([6, 6, 6], dtype=torch.long), - positions=positions.detach(), - ) - batch = Batch.from_data_list([data]).to(device) idx = torch.tensor([0, 2], device=device) # Use gradcheck with a wrapper that creates a fresh batch @@ -576,7 +576,7 @@ def test_gradient_periodic(self, device: str) -> None: 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] + pbc = torch.tensor([[True, True, True]]) # [1, 3] data = AtomicData( atomic_numbers=torch.tensor([6, 6], dtype=torch.long), positions=positions.cpu(), @@ -600,7 +600,7 @@ def test_not_at_half_cell_tie(self, device: str) -> None: 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] + 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), @@ -686,14 +686,14 @@ def test_different_registration_orders_same_result(self) -> None: # =========================================================================== -# 6. torch.compile spike — CPU fullgraph tests +# 6. torch.compile — fullgraph and graph-break tests # =========================================================================== class TestCompileSpike: - """PR 1 compile spike: verifies what can and cannot be compiled. + """Verifies what can and cannot be compiled with ``torch.compile``. - **Spike finding (documented per proposal section 6):** + **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``. From 2e5673d63237821049b638dda3e9d08b72b9c64c Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 16:58:39 -0400 Subject: [PATCH 03/11] mic for triclinic Signed-off-by: Samarjeet Prasad --- .../enhanced_sampling/cv/pair_distance.py | 53 ++++++++++++++----- test/enhanced_sampling/test_bias_core.py | 44 +++++++++++++++ 2 files changed, 84 insertions(+), 13 deletions(-) diff --git a/nvalchemi/enhanced_sampling/cv/pair_distance.py b/nvalchemi/enhanced_sampling/cv/pair_distance.py index 2fa10adf..6f2b935b 100644 --- a/nvalchemi/enhanced_sampling/cv/pair_distance.py +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -156,6 +156,21 @@ def _apply_mic(dr: Tensor, cell: Tensor, pbc: Tensor) -> Tensor: ------- Tensor MIC-corrected displacement vectors, shape ``[B, 3]``. + + Notes + ----- + Componentwise fractional rounding (``df -= round(df)``) is only exact for + orthogonal cells. For skewed triclinic cells the Wigner–Seitz cell does + not align with the fractional-coordinate axes, so the shortest image may + require adding or subtracting a lattice vector even after rounding each + component to (−0.5, 0.5]. Example: cell rows ``[[1,0,0],[0.9,0.1,0], + [0,0,10]]``, fractional displacement ``[0.49,0.49,0]`` — rounding gives + ≈ 0.932 Å but the true MIC vector (offset ``[0,−1,0]``) is ≈ 0.060 Å. + + This implementation performs an exhaustive search over all 27 lattice + images (offsets in ``{−1, 0, +1}^3`` restricted to periodic dims) and + returns the Cartesian vector with minimum norm. The search is fully + vectorised and passes ``torch.compile(fullgraph=True)``. """ # Normalise shapes if cell.dim() == 4: @@ -163,22 +178,34 @@ def _apply_mic(dr: Tensor, cell: Tensor, pbc: Tensor) -> Tensor: if pbc.dim() == 3: pbc = pbc.squeeze(1) # [B, 3] - # Convert pbc to float mask on the correct device/dtype pbc_mask = pbc.to(dtype=cell.dtype) # [B, 3] - # Fractional displacement: dr @ A^{-1} - # cell[b] has shape [3, 3]; A^{-1} = cell^{-T} when rows=lattice vecs. - # torch.linalg.solve(A.T, dr.T) gives x s.t. A.T x = dr.T - # Equivalently: df = dr @ A^{-1} = (A^{-T} dr^T)^T - cell_inv = torch.linalg.inv(cell) # [B, 3, 3] (A^{-1}) - # df[b] = dr[b] @ cell_inv[b]; batched matmul: [B, 1, 3] @ [B, 3, 3] + # Fractional displacement: df[b] = dr[b] @ cell[b]^{-1} + cell_inv = torch.linalg.inv(cell) # [B, 3, 3] df = torch.bmm(dr.unsqueeze(1), cell_inv).squeeze(1) # [B, 3] - # Apply half-cell rounding only along periodic dimensions. - # For periodic dims (pbc_mask=1): df_mic = df - round(df) → maps to (−0.5, 0.5] - # For non-periodic dims (pbc_mask=0): df_mic = df → unchanged - df_mic = df - torch.round(df) * pbc_mask # [B, 3] + # Initial rounding: map each periodic component to (−0.5, 0.5]. + # For non-periodic dims the component is unchanged. + df_rounded = df - torch.round(df) * pbc_mask # [B, 3] + + # --- Exhaustive 27-image search ----------------------------------------- + # Build all 27 offset vectors {-1, 0, 1}^3 and mask non-periodic dims. + 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] + + # [B, 27, 3]: apply pbc_mask so non-periodic dims are never shifted + offsets_masked = all_offsets[None] * pbc_mask[:, None, :] + + # Candidate fractional displacements and their Cartesian counterparts + df_cands = df_rounded[:, None, :] + offsets_masked # [B, 27, 3] + dr_cands = torch.einsum("bki,bij->bkj", df_cands, cell) # [B, 27, 3] + + # Select the image with minimum squared Cartesian distance (avoid sqrt) + 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] + dr_mic = dr_cands.gather(1, best).squeeze(1) # [B, 3] - # Back to Cartesian: dr_mic[b] = df_mic[b] @ cell[b] - dr_mic = torch.bmm(df_mic.unsqueeze(1), cell).squeeze(1) # [B, 3] return dr_mic diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index 9f001757..ce6193c2 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -547,6 +547,50 @@ def test_triclinic_mic_known_value(self, device: str) -> None: # 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_triclinic_mic_skewed_cell_componentwise_rounding_fails( + self, device: str + ) -> None: + """Regression: skewed cell where componentwise fractional rounding is wrong. + + Cell rows: [[1,0,0],[0.9,0.1,0],[0,0,10]]. + Fractional displacement [0.49, 0.49, 0]: + - Componentwise rounding keeps [0.49, 0.49, 0] → Cartesian ≈ 0.932 Å. + - True MIC (offset [0,−1,0]) → [0.49,−0.51,0] → Cartesian ≈ 0.060 Å. + Componentwise rounding would return the wrong (longer) image. + The 27-image exhaustive search returns the correct shortest vector. + """ + cell = torch.tensor( + [[[1.0, 0.0, 0.0], [0.9, 0.1, 0.0], [0.0, 0.0, 10.0]]] + ) # [1, 3, 3] + pbc = torch.tensor([[True, True, True]]) # [1, 3] + + # pos_j chosen so that fractional displacement = [0.49, 0.49, 0] + # Cartesian pos_j = 0.49*[1,0,0] + 0.49*[0.9,0.1,0] = [0.931, 0.049, 0] + pos_i = torch.tensor([[0.0, 0.0, 0.0]]) + pos_j = torch.tensor([[0.931, 0.049, 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 would give ≈ 0.932 Å; correct MIC is ≈ 0.060 Å. + # We assert the result is well below the naive distance. + naive_dist = torch.linalg.vector_norm(pos_j - pos_i).item() + assert d.item() < naive_dist * 0.2, ( + f"MIC distance {d.item():.4f} Å is not significantly shorter than " + f"naive distance {naive_dist:.4f} Å — 27-image search may not be working." + ) + assert d.item() < 0.12, ( + f"Expected MIC distance ≈ 0.060 Å, got {d.item():.4f} Å." + ) + # --- gradients --- def test_gradient_nonperiodic(self, device: str) -> None: From 9087a4d834af0986a3dbe182fc8fdd0c395d8a2f Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 17:07:51 -0400 Subject: [PATCH 04/11] fix virial contribution for any position dependent bias term Signed-off-by: Samarjeet Prasad --- nvalchemi/enhanced_sampling/_bias.py | 109 ++++++++++++++++------- test/enhanced_sampling/test_bias_core.py | 73 ++++++++++++--- 2 files changed, 135 insertions(+), 47 deletions(-) diff --git a/nvalchemi/enhanced_sampling/_bias.py b/nvalchemi/enhanced_sampling/_bias.py index 82ddb80b..349dfb78 100644 --- a/nvalchemi/enhanced_sampling/_bias.py +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -263,10 +263,33 @@ def energy(self, current: Batch) -> Tensor: ) def evaluate(self, current: Batch) -> BiasResult: - """Compute energy, forces, and (optionally) virial via autograd. + """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 @@ -274,39 +297,65 @@ def evaluate(self, current: Batch) -> BiasResult: 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(): - # Create isolated autograd leaves. - # requires_grad_() is not supported by torch.compile — this method + # --- 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) + pos_leaf = current.positions.detach().requires_grad_(True) # [N, 3] - cell_leaf: Tensor | None = None - if has_cell: - cell_leaf = current.cell.detach().requires_grad_(True) + # --- 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 - # Temporarily replace positions (and cell) on the live batch with - # the fresh leaves so that self.energy() can access other batch - # fields (atomic_numbers, batch_idx, etc.) normally. - # Restored unconditionally in the finally block. - original_positions = current.positions - original_cell = current.cell if has_cell else None + 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_leaf - if has_cell and cell_leaf is not None: - current["cell"] = cell_leaf + 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] - grad_outputs = (torch.ones_like(bias_energy),) - inputs: tuple[Tensor, ...] = ( - (pos_leaf,) if cell_leaf is None else (pos_leaf, cell_leaf) - ) + 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=grad_outputs, + grad_outputs=(torch.ones_like(bias_energy),), create_graph=False, retain_graph=False, allow_unused=False, @@ -317,21 +366,13 @@ def evaluate(self, current: Batch) -> BiasResult: if has_cell and original_cell is not None: current["cell"] = original_cell - # grads[0]: d(sum(energy))/d(positions) — negate for forces. - forces = -grads[0].detach() # [N_atoms, 3] + # grads[0] = dE/d(pos_leaf) at strain=I → forces = −grad. + forces = -grads[0].detach() # [N, 3] virial: Tensor | None = None - if has_cell and len(grads) > 1 and grads[1] is not None: - # Canonical virial W = −dE/d(strain); for the row-vector - # convention (ASE): W = −(dE/d(cell)) @ cell.T - dcell = grads[1].detach() - if dcell.dim() == 4: - dcell = dcell.squeeze(1) - cell = original_cell - if cell is not None and cell.dim() == 4: - cell = cell.squeeze(1) - if cell is not None: - virial = -(dcell @ cell.transpose(-1, -2)) # [B, 3, 3] + 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(), diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index ce6193c2..82c62a42 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -378,23 +378,70 @@ def test_no_memory_growth_repeated_evaluate(self, device: str) -> None: f"GPU memory grew by {mem_end - mem_start} bytes across 10 evaluate() calls" ) - def test_virial_with_periodic_cell(self, device: str) -> None: - """ConservativeBias should populate virial for periodic batches.""" + 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( - "Virial via autograd.grad on CUDA triggers a cudagraph_trees " - "assertion in this environment (PyTorch internal assertion at " - "cudagraph_trees.py:2608). CPU correctness is verified; CUDA " - "virial correctness will be covered in the GPU integration test " - "suite added in PR 2." + "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." ) - batch = _make_cubic_batch(n_graphs=1, atoms_per_graph=4, device=device) - bias = _PairDistanceBias(atom_indices=torch.tensor([0, 1]), k=1.0) + + 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) - # virial may be None if pair_distance gradient w.r.t. cell is zero - # (atoms in same image → cell gradient cancels); just check shape if present - if result.virial is not None: - assert result.virial.shape == (1, 3, 3) + + 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) = -1 * [[-1],[-1,-0,-0]] ... + # dr_mic = [-1, 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.""" From cbc1ad56583688cc233d4725bcfdf90c2c1786e3 Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 17:11:16 -0400 Subject: [PATCH 05/11] validate local atom index bounds before adding batch_ptr offsets Signed-off-by: Samarjeet Prasad --- .../enhanced_sampling/cv/pair_distance.py | 23 ++++++ test/enhanced_sampling/test_bias_core.py | 76 +++++++++++++++++++ 2 files changed, 99 insertions(+) diff --git a/nvalchemi/enhanced_sampling/cv/pair_distance.py b/nvalchemi/enhanced_sampling/cv/pair_distance.py index 6f2b935b..e67c7666 100644 --- a/nvalchemi/enhanced_sampling/cv/pair_distance.py +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -113,6 +113,29 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: # [2] → broadcast same pair across all graphs atom_indices = atom_indices.unsqueeze(0).expand(B, 2) # [B, 2] + # --- Bounds check (eager only; skipped under torch.compile) ---------- + # Without this, an out-of-range local index silently wraps into the next + # graph's atom rows, producing a cross-system CV with no error. + if not torch.compiler.is_compiling(): + 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] # [B] + neg = idx < 0 + if neg.any(): + bad_graphs = neg.nonzero(as_tuple=False).squeeze(-1).tolist() + raise IndexError( + f"pair_distance: {label} has negative values for graph(s) " + f"{bad_graphs}: {idx[neg].tolist()}" + ) + oob = idx >= atoms_per_graph + if oob.any(): + bad_graphs = oob.nonzero(as_tuple=False).squeeze(-1).tolist() + raise IndexError( + f"pair_distance: {label} is out of range for graph(s) " + f"{bad_graphs} — index {idx[oob].tolist()} >= " + f"graph size {atoms_per_graph[oob].tolist()}" + ) + # batch_ptr[b] is the start of graph b; atom_indices[:, k] is local idx offsets = batch_ptr[:-1] # [B] global_i = offsets + atom_indices[:, 0] # [B] diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index 82c62a42..05489cb1 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -473,6 +473,82 @@ def result_close(a: BiasResult, b: BiasResult, atol: float = 1e-6) -> bool: class TestPairDistance: """Tests for the pair_distance collective variable.""" + # --- 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 --- def test_nonperiodic_known_value(self, device: str) -> None: From e2bceb5d63bfa54c17013d8b9d3918c8872b62ce Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 17:16:12 -0400 Subject: [PATCH 06/11] added validation tests for BiasResult Signed-off-by: Samarjeet Prasad --- nvalchemi/enhanced_sampling/_bias.py | 87 ++++++++++++++++++++- test/enhanced_sampling/test_bias_core.py | 96 ++++++++++++++++++++++++ 2 files changed, 182 insertions(+), 1 deletion(-) diff --git a/nvalchemi/enhanced_sampling/_bias.py b/nvalchemi/enhanced_sampling/_bias.py index 349dfb78..0333dbd5 100644 --- a/nvalchemi/enhanced_sampling/_bias.py +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -99,9 +99,31 @@ def __post_init__(self) -> None: def _validate_bias_result(result: BiasResult) -> None: - """Eager-only validation of a ``BiasResult`` (skipped under compile).""" + """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, @@ -133,6 +155,69 @@ def _validate_bias_result(result: BiasResult) -> None: 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 diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index 05489cb1..4253a655 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -187,6 +187,102 @@ def test_frozen_immutability(self) -> None: 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 From 3e92c05d27dd6a53b961eb165954b2264630f99a Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 17:41:04 -0400 Subject: [PATCH 07/11] support Minkowski-reduced triclinic cells Signed-off-by: Samarjeet Prasad --- nvalchemi/enhanced_sampling/__init__.py | 4 +- .../enhanced_sampling/cv/pair_distance.py | 243 +++++++++++------- test/enhanced_sampling/test_bias_core.py | 77 ++++-- 3 files changed, 209 insertions(+), 115 deletions(-) diff --git a/nvalchemi/enhanced_sampling/__init__.py b/nvalchemi/enhanced_sampling/__init__.py index b6de395c..7e6c65d2 100644 --- a/nvalchemi/enhanced_sampling/__init__.py +++ b/nvalchemi/enhanced_sampling/__init__.py @@ -23,7 +23,8 @@ :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 general triclinic MIC. + nonperiodic and Minkowski-reduced triclinic MIC. General triclinic MIC + (unreduced cells via LLL) is deferred. Deferred to later milestones ----------------------------- @@ -31,6 +32,7 @@ * :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 ( diff --git a/nvalchemi/enhanced_sampling/cv/pair_distance.py b/nvalchemi/enhanced_sampling/cv/pair_distance.py index e67c7666..232b1f20 100644 --- a/nvalchemi/enhanced_sampling/cv/pair_distance.py +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -18,41 +18,68 @@ * Non-periodic systems (``batch.cell`` is ``None`` or ``batch.pbc`` is all ``False``). -* Fully periodic and mixed-periodic systems via the minimum-image - convention (MIC) for general triclinic cells. +* 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 triclinic cell with lattice matrix ``A`` (rows = lattice vectors, -ASE convention), the fractional displacement is:: +For a reduced cell with lattice matrix ``A`` (rows = lattice vectors, +ASE convention):: - df = (r_j - r_i) @ A^{-1} - -We then apply the standard half-cell rounding:: - - df -= torch.round(df) # map to (−0.5, 0.5] - -and convert back to Cartesian:: - - dr_mic = df @ A - -The distance is ``||dr_mic||``. - -**Boundary note:** The half-cell tie (``|df_k| == 0.5`` exactly) maps -to ``+0.5`` or ``−0.5`` depending on floating-point rounding, producing -a discontinuity in the distance function exactly at the Wigner–Seitz -cell boundary. This is the standard MIC behaviour and is shared by ASE, -LAMMPS, and PLUMED. Tests are placed away from the half-cell tie; the -boundary behaviour is documented but not worked around. + 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 --------------------------- -This function is a compile target (``compile_biases=True``). It uses -only standard PyTorch ops; there are no Python-level branches on tensor -values (the periodicity branch is on the *shape* / *bool* of ``pbc``, -which is resolved at trace time for fixed-pbc batches). Gradient flow -through ``pair_distance`` for use inside :class:`ConservativeBias` is -fully supported. +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 @@ -75,34 +102,31 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: ---------- batch: Current ``Batch`` containing atomic positions and (optionally) - cell and PBC flags. + 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. - Atom indices are **local to each graph** (0-based within the - graph, not global row indices in the batched position tensor). + 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 units as + Shape ``[B, 1]`` — pair distance in the same length unit as ``batch.positions`` (Å). Fully differentiable w.r.t. ``batch.positions`` and ``batch.cell``. - Notes - ----- - * MIC is applied independently per graph, using that graph's cell. - * PBC flags (``batch.pbc``) are used to determine whether MIC is - applied. If **all** PBC flags for a graph are ``False``, or if - ``batch.cell`` is ``None``, the straight Cartesian distance is - used. - * For graphs with mixed periodicity (e.g. ``pbc = [True, True, False]``), - the MIC rounding is still applied in fractional coordinates; the - non-periodic fractional components will map outside (−0.5, 0.5] in - the direction(s) with ``pbc=False``, but the ``round()`` is only - applied along periodic dimensions. + Raises + ------ + IndexError + If any local atom index is negative or >= the graph's atom count. + 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. """ positions = batch.positions # [N_total, 3] batch_ptr = batch.batch_ptr # [B+1] @@ -110,60 +134,121 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: # --- Resolve atom_indices to global row indices ----------------------- if atom_indices.dim() == 1: - # [2] → broadcast same pair across all graphs atom_indices = atom_indices.unsqueeze(0).expand(B, 2) # [B, 2] - # --- Bounds check (eager only; skipped under torch.compile) ---------- - # Without this, an out-of-range local index silently wraps into the next - # graph's atom rows, producing a cross-system CV with no error. + # --- Eager-only input 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] # [B] + idx = atom_indices[:, col] neg = idx < 0 if neg.any(): - bad_graphs = neg.nonzero(as_tuple=False).squeeze(-1).tolist() + bad = neg.nonzero(as_tuple=False).squeeze(-1).tolist() raise IndexError( f"pair_distance: {label} has negative values for graph(s) " - f"{bad_graphs}: {idx[neg].tolist()}" + f"{bad}: {idx[neg].tolist()}" ) oob = idx >= atoms_per_graph if oob.any(): - bad_graphs = oob.nonzero(as_tuple=False).squeeze(-1).tolist() + bad = oob.nonzero(as_tuple=False).squeeze(-1).tolist() raise IndexError( f"pair_distance: {label} is out of range for graph(s) " - f"{bad_graphs} — index {idx[oob].tolist()} >= " + f"{bad} — index {idx[oob].tolist()} >= " f"graph size {atoms_per_graph[oob].tolist()}" ) - # batch_ptr[b] is the start of graph b; atom_indices[:, k] is local idx 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 ----------------------------------- + # --- 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 if has_cell and has_pbc: + if not torch.compiler.is_compiling(): + _check_minkowski_reduced(batch.cell, batch.pbc) dr = _apply_mic(dr, batch.cell, batch.pbc) - dist = torch.linalg.vector_norm(dr, dim=-1, keepdim=True) # [B, 1] - return dist + return torch.linalg.vector_norm(dr, dim=-1, keepdim=True) # [B, 1] + + +# --------------------------------------------------------------------------- +# 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." + ) # --------------------------------------------------------------------------- -# Internal MIC helper +# MIC implementation # --------------------------------------------------------------------------- def _apply_mic(dr: Tensor, cell: Tensor, pbc: Tensor) -> Tensor: - """Apply the minimum-image convention for general triclinic cells. + """Apply the minimum-image convention via an exhaustive 27-image search. Parameters ---------- @@ -182,53 +267,35 @@ def _apply_mic(dr: Tensor, cell: Tensor, pbc: Tensor) -> Tensor: Notes ----- - Componentwise fractional rounding (``df -= round(df)``) is only exact for - orthogonal cells. For skewed triclinic cells the Wigner–Seitz cell does - not align with the fractional-coordinate axes, so the shortest image may - require adding or subtracting a lattice vector even after rounding each - component to (−0.5, 0.5]. Example: cell rows ``[[1,0,0],[0.9,0.1,0], - [0,0,10]]``, fractional displacement ``[0.49,0.49,0]`` — rounding gives - ≈ 0.932 Å but the true MIC vector (offset ``[0,−1,0]``) is ≈ 0.060 Å. - - This implementation performs an exhaustive search over all 27 lattice - images (offsets in ``{−1, 0, +1}^3`` restricted to periodic dims) and - returns the Cartesian vector with minimum norm. The search is fully - vectorised and passes ``torch.compile(fullgraph=True)``. + 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. """ - # Normalise shapes if cell.dim() == 4: - cell = cell.squeeze(1) # [B, 3, 3] + cell = cell.squeeze(1) if pbc.dim() == 3: - pbc = pbc.squeeze(1) # [B, 3] + pbc = pbc.squeeze(1) pbc_mask = pbc.to(dtype=cell.dtype) # [B, 3] - # Fractional displacement: df[b] = dr[b] @ cell[b]^{-1} + # Fractional displacement cell_inv = torch.linalg.inv(cell) # [B, 3, 3] df = torch.bmm(dr.unsqueeze(1), cell_inv).squeeze(1) # [B, 3] - # Initial rounding: map each periodic component to (−0.5, 0.5]. - # For non-periodic dims the component is unchanged. + # Initial half-cell rounding (periodic dims only) df_rounded = df - torch.round(df) * pbc_mask # [B, 3] - # --- Exhaustive 27-image search ----------------------------------------- - # Build all 27 offset vectors {-1, 0, 1}^3 and mask non-periodic dims. + # 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] - # [B, 27, 3]: apply pbc_mask so non-periodic dims are never shifted - offsets_masked = all_offsets[None] * pbc_mask[:, None, :] - - # Candidate fractional displacements and their Cartesian counterparts + 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] - # Select the image with minimum squared Cartesian distance (avoid sqrt) 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] - dr_mic = dr_cands.gather(1, best).squeeze(1) # [B, 3] - - return dr_mic + return dr_cands.gather(1, best).squeeze(1) # [B, 3] diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index 4253a655..08afa6b5 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -25,9 +25,10 @@ ``requires_grad`` escape into live batch or result; no memory growth across 10 repeated evaluations. * :func:`~nvalchemi.enhanced_sampling.pair_distance` — nonperiodic and - general triclinic MIC; shared and per-graph atom indices; gradients - via ``torch.autograd.gradcheck``; compile-stability under - ``torch.compile`` (fullgraph=True on CPU). + 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`` @@ -766,27 +767,52 @@ def test_triclinic_mic_known_value(self, device: str) -> None: # 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_triclinic_mic_skewed_cell_componentwise_rounding_fails( + 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: - """Regression: skewed cell where componentwise fractional rounding is wrong. - - Cell rows: [[1,0,0],[0.9,0.1,0],[0,0,10]]. - Fractional displacement [0.49, 0.49, 0]: - - Componentwise rounding keeps [0.49, 0.49, 0] → Cartesian ≈ 0.932 Å. - - True MIC (offset [0,−1,0]) → [0.49,−0.51,0] → Cartesian ≈ 0.060 Å. - Componentwise rounding would return the wrong (longer) image. - The 27-image exhaustive search returns the correct shortest vector. + """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. """ - cell = torch.tensor( - [[[1.0, 0.0, 0.0], [0.9, 0.1, 0.0], [0.0, 0.0, 10.0]]] - ) # [1, 3, 3] - pbc = torch.tensor([[True, True, True]]) # [1, 3] + # 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 chosen so that fractional displacement = [0.49, 0.49, 0] - # Cartesian pos_j = 0.49*[1,0,0] + 0.49*[0.9,0.1,0] = [0.931, 0.049, 0] + # 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([[0.931, 0.049, 0.0]]) + pos_j = torch.tensor([[1.372, 0.98, 0.0]]) positions = torch.cat([pos_i, pos_j], dim=0) data = AtomicData( @@ -799,15 +825,14 @@ def test_triclinic_mic_skewed_cell_componentwise_rounding_fails( idx = torch.tensor([0, 1], device=device) d = pair_distance(batch, idx) - # Componentwise rounding would give ≈ 0.932 Å; correct MIC is ≈ 0.060 Å. - # We assert the result is well below the naive distance. + # 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.2, ( - f"MIC distance {d.item():.4f} Å is not significantly shorter than " - f"naive distance {naive_dist:.4f} Å — 27-image search may not be working." + 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() < 0.12, ( - f"Expected MIC distance ≈ 0.060 Å, got {d.item():.4f} Å." + assert d.item() < 1.20, ( + f"Expected MIC distance ≈ 1.164 Å, got {d.item():.4f} Å." ) # --- gradients --- From 70ae96e674c3539aa081f00bcddd6fed699a8ac8 Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 17:45:41 -0400 Subject: [PATCH 08/11] atom_indices shape/dtype validation Signed-off-by: Samarjeet Prasad --- .../enhanced_sampling/cv/pair_distance.py | 94 ++++++++++++++++++- test/enhanced_sampling/test_bias_core.py | 76 +++++++++++++++ 2 files changed, 168 insertions(+), 2 deletions(-) diff --git a/nvalchemi/enhanced_sampling/cv/pair_distance.py b/nvalchemi/enhanced_sampling/cv/pair_distance.py index 232b1f20..dddb7e3a 100644 --- a/nvalchemi/enhanced_sampling/cv/pair_distance.py +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -121,8 +121,12 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: 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. + 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 @@ -132,11 +136,18 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: 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 input validation ------------------------------------- + # --- 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] @@ -178,6 +189,85 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: 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 # --------------------------------------------------------------------------- diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index 08afa6b5..d7cc0158 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -570,6 +570,82 @@ def result_close(a: BiasResult, b: BiasResult, atol: float = 1e-6) -> bool: 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: From 2e91f2e45d4fe2d7e99a62a93aa9df1fffa63e5b Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 17:53:02 -0400 Subject: [PATCH 09/11] updating docstrings with the fixes Signed-off-by: Samarjeet Prasad --- nvalchemi/enhanced_sampling/_bias.py | 38 +++++++++++++----------- test/enhanced_sampling/test_bias_core.py | 5 ++-- 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/nvalchemi/enhanced_sampling/_bias.py b/nvalchemi/enhanced_sampling/_bias.py index 0333dbd5..b879a97f 100644 --- a/nvalchemi/enhanced_sampling/_bias.py +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -290,12 +290,17 @@ class ConservativeBias: 1. Enters a local ``torch.enable_grad()`` region (safe inside ``torch.no_grad()`` outer contexts). - 2. Creates fresh detached autograd leaves for positions (and cell when - ``compute_virial=True``). - 3. Evaluates :meth:`energy` on an isolated read-only view of the batch - that replaces positions (and cell) with the fresh leaves. - 4. Derives forces and (optionally) virial in one ``autograd.grad`` call - with ``create_graph=False, retain_graph=False``. + 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. @@ -304,16 +309,14 @@ class ConservativeBias: a non-null ``grad_fn`` into the live ``Batch``, ``BiasResult``, retained history, bias state, observables, or a checkpoint. - Parameters - ---------- - compute_virial: - When ``True`` (default when a valid cell exists at evaluation - time), derive the canonical virial - ``W = −dE/d(strain)`` via ``autograd.grad`` w.r.t. the cell - leaf. The virial is shaped ``[B, 3, 3]``. - 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 @@ -339,9 +342,10 @@ def energy(self, current: Batch) -> Tensor: Parameters ---------- current: - A *read-only view* of the live batch where ``positions`` (and - optionally ``cell``) have been replaced by fresh autograd - leaves. Do not assign to any batch field inside this method. + 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" diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index d7cc0158..99555bab 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -528,8 +528,9 @@ def test_virial_canonical_analytical(self, device: str) -> None: 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) = -1 * [[-1],[-1,-0,-0]] ... - # dr_mic = [-1, 0, 0] → W[0,0] = -1, all other elements = 0 + # 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}. " From 91a68e8cc6d7fcf260f2962f965abdca24ecd476 Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Wed, 5 Aug 2026 18:30:45 -0400 Subject: [PATCH 10/11] fixing non-periodic case Signed-off-by: Samarjeet Prasad --- .../enhanced_sampling/cv/pair_distance.py | 29 ++++++++- test/enhanced_sampling/test_bias_core.py | 59 +++++++++++++------ 2 files changed, 70 insertions(+), 18 deletions(-) diff --git a/nvalchemi/enhanced_sampling/cv/pair_distance.py b/nvalchemi/enhanced_sampling/cv/pair_distance.py index dddb7e3a..468d2f63 100644 --- a/nvalchemi/enhanced_sampling/cv/pair_distance.py +++ b/nvalchemi/enhanced_sampling/cv/pair_distance.py @@ -131,6 +131,17 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: 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] @@ -181,7 +192,23 @@ def pair_distance(batch: Batch, atom_indices: Tensor) -> Tensor: 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 - if has_cell and has_pbc: + # 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) diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index 99555bab..231b73c1 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -723,6 +723,48 @@ def test_variable_size_batch_no_silent_cross_graph(self) -> None: 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: @@ -803,23 +845,6 @@ def test_periodic_cubic_mic(self, device: str) -> None: d = pair_distance(batch, idx) assert torch.allclose(d, torch.tensor([[0.2]], device=device), atol=1e-4) - def test_nonperiodic_cubic_no_mic(self, device: str) -> None: - """pbc=False: long-range distance not folded by MIC.""" - box = 10.0 - 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([[False, False, False]]) # [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([[9.8]], device=device), atol=1e-4) - # --- triclinic MIC --- def test_triclinic_mic_known_value(self, device: str) -> None: From 03d3ad8dd9bd8e0c1be64b2d079231609d477531 Mon Sep 17 00:00:00 2001 From: Samarjeet Prasad Date: Thu, 6 Aug 2026 10:49:24 -0400 Subject: [PATCH 11/11] handling virial and stress mixed cases Signed-off-by: Samarjeet Prasad --- nvalchemi/enhanced_sampling/_bias.py | 24 +++++++++++++++++- test/enhanced_sampling/test_bias_core.py | 32 ++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/nvalchemi/enhanced_sampling/_bias.py b/nvalchemi/enhanced_sampling/_bias.py index b879a97f..afd7454a 100644 --- a/nvalchemi/enhanced_sampling/_bias.py +++ b/nvalchemi/enhanced_sampling/_bias.py @@ -486,7 +486,13 @@ def aggregate_bias_results(results: list[BiasResult]) -> BiasResult: Rules ----- * ``None`` fields are skipped (treated as zero contribution). - * ``stress`` and ``virial`` are accumulated separately. + * 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. @@ -511,6 +517,22 @@ def aggregate_bias_results(results: list[BiasResult]) -> BiasResult: 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 diff --git a/test/enhanced_sampling/test_bias_core.py b/test/enhanced_sampling/test_bias_core.py index 231b73c1..69c1fab3 100644 --- a/test/enhanced_sampling/test_bias_core.py +++ b/test/enhanced_sampling/test_bias_core.py @@ -1053,6 +1053,38 @@ def test_virial_summed(self) -> None: 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)})