diff --git a/CHANGELOG.md b/CHANGELOG.md index 34a3f7f..added61 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Changed +- Added canonical global packed mode 2 for batched sparse neighbor matrices. Callers with 3D local-index matrices must convert them with `aimnet.nbops.convert_mode2_local_to_global`; every system must reserve a final dummy atom, and full-3D periodic inputs must provide aligned lattice shifts. - Migrated Ewald Coulomb to the nvalchemiops (>=0.4) energy+autograd API: energies are bit-identical and forces/stress unchanged to kernel rounding (<=2e-06 eV/A), the upstream `DeprecationWarning` for Ewald is gone, and Ewald dense Hessians / HVPs are now true relaxed-charge autograd derivatives (comparable with simple/DSF) instead of fixed-charge finite differences. The `eps` argument of `hessian_vector_product` now applies to PME only. PME intentionally remains on the legacy explicit-terms path: its charge-gradient backward misbehaves inside the full calculator graph under `create_graph=True` (train mode), tracked in a dedicated issue; PME behavior is bit-identical to the previous release. - Hardened `make test`: the parallel run now hides CUDA and caps per-worker threads (`CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1`); previously xdist workers either all initialized the first GPU (OOM on CUDA boxes) or oversubscribed the CPU with per-worker torch thread pools. A new `make test-gpu` target runs the GPU-marked tests serially on CUDA. - Bumped `nvalchemi-toolkit-ops` to `>=0.4.0` and `warp-lang` to `>=1.13,<2` (installs 1.15). Energies, charges, and Hessians are bit-identical to 0.3.1; explicit force/virial outputs shift within float32 accumulation noise (max 6.7e-05 eV/A on a periodic system, 40x inside the project's cross-version acceptance of 1e-4 Hartree/A). The 0.4.0 direct-output flags used by the Ewald/PME path (`compute_forces`/`compute_virial`) are deprecated upstream and now emit `DeprecationWarning`; migrating to the autograd-based API is tracked as follow-up work. diff --git a/aimnet/calculators/calculator.py b/aimnet/calculators/calculator.py index 88474a2..86a46e0 100644 --- a/aimnet/calculators/calculator.py +++ b/aimnet/calculators/calculator.py @@ -9,6 +9,7 @@ import torch from torch import Tensor, nn +from aimnet import nbops from aimnet.modules import DFTD3, LRCoulomb from aimnet.modules.lr import ExternalDerivativeTerms @@ -112,12 +113,18 @@ class AIMNet2Calculator: "mol_idx": torch.int, "nbmat": torch.int, "nbmat_lr": torch.int, + "nbmat_coulomb": torch.int, + "nbmat_dftd3": torch.int, "nb_pad_mask": torch.bool, "nb_pad_mask_lr": torch.bool, "shifts": torch.float, "shifts_lr": torch.float, + "shifts_coulomb": torch.float, + "shifts_dftd3": torch.float, "cell": torch.float, "pbc": torch.bool, + "cutoff_coulomb": torch.float, + "cutoff_dftd3": torch.float, } keys_out: ClassVar[list[str]] = ["energy", "charges", "spin_charges", "forces", "hessian", "stress"] atom_feature_keys: ClassVar[list[str]] = ["coord", "numbers", "charges", "spin_charges", "forces"] @@ -846,6 +853,14 @@ def eval( ) if hessian: + probe = self.to_input_tensors(data) + primary_nbmat = probe.get("nbmat") + nbops.validate_neighbor_suffix_layout(probe) + if primary_nbmat is not None and primary_nbmat.ndim == 3: + nbops.normalize_mode2_periodic_geometry(probe, B=primary_nbmat.shape[0]) + for suffix in nbops.NBMAT_SUFFIXES: + if f"nbmat{suffix}" in probe or f"shifts{suffix}" in probe: + nbops.validate_mode2_nbmat_raw(probe, suffix=suffix) subsystems = self._split_hessian_batch(data) if subsystems is not None: stack = torch.as_tensor(data["coord"]).ndim == 3 @@ -984,6 +999,13 @@ def prepare_input(self, data: dict[str, Any], *, hessian: bool = False) -> dict[ caller_had_mol_idx = raw_data.get("mol_idx") is not None caller_had_nbmat = raw_data.get("nbmat") is not None data = self.to_input_tensors(data) + primary_nbmat = data.get("nbmat") + nbops.validate_neighbor_suffix_layout(data) + if primary_nbmat is not None and primary_nbmat.ndim == 3: + nbops.normalize_mode2_periodic_geometry(data, B=primary_nbmat.shape[0]) + for suffix in nbops.NBMAT_SUFFIXES: + if f"nbmat{suffix}" in data or f"shifts{suffix}" in data: + nbops.validate_mode2_nbmat_raw(data, suffix=suffix) data = self.mol_flatten(data, hessian=hessian) if data.get("cell") is not None and self._coulomb_method == "simple": warnings.warn( @@ -1219,22 +1241,36 @@ def _select_for_structure(value: Any, index: int, n_struct: int) -> Any: def _split_batch_dim(self, data: dict[str, Any], B: int) -> list[dict[str, Any]]: """Slice a (B, ...) batched input into B single-structure dicts.""" subs: list[dict[str, Any]] = [] + primary_nbmat = data.get("nbmat") + mode2 = primary_nbmat is not None and torch.as_tensor(primary_nbmat).ndim == 3 for b in range(B): sub: dict[str, Any] = {} for k, v in data.items(): if v is None: continue if k in ("coord", "numbers"): - sub[k] = torch.as_tensor(v)[b] + value = torch.as_tensor(v) + sub[k] = value[b : b + 1] if mode2 else value[b] elif k in ("charge", "mult"): sub[k] = self._select_for_structure(v, b, B) elif k == "cell": t = torch.as_tensor(v) - sub[k] = t[b] if t.ndim == 3 else t - elif k.startswith(("nbmat", "shifts")) or k == "mol_idx": - # Precomputed neighbor-list keys must not be shared unsliced - # across subsystems (they'd be wrong per-structure); the - # recursive eval rebuilds them. + sub[k] = t[b : b + 1] if mode2 and t.ndim == 3 else (t[b] if t.ndim == 3 else t) + elif k == "pbc": + t = torch.as_tensor(v) + sub[k] = t[b : b + 1] if mode2 and t.ndim == 2 else (t[b] if t.ndim == 2 else t) + elif k.startswith("nbmat"): + t = torch.as_tensor(v)[b : b + 1] + if t.ndim == 3: + sentinel = B * t.shape[1] + singleton_sentinel = t.shape[1] + usable = t != sentinel + t = torch.where(usable, t - b * t.shape[1], torch.full_like(t, singleton_sentinel)) + sub[k] = t + elif k.startswith("shifts"): + t = torch.as_tensor(v) + sub[k] = t[b : b + 1] if t.ndim >= 3 else t + elif k == "mol_idx": continue else: sub[k] = v @@ -1404,11 +1440,22 @@ def to_input_tensors(self, data: dict[str, Any]) -> dict[str, Tensor]: if not (isinstance(data[k], Tensor) and data[k].requires_grad): t = t.detach() ret[k] = t + neighbor_memo: list[tuple[Any, Tensor]] = [] + neighbor_keys = {f"nbmat{suffix}" for suffix in nbops.NBMAT_SUFFIXES} for k in self.keys_in_optional: if k in data and data[k] is not None: - t = torch.as_tensor(data[k], device=self.device, dtype=self.keys_in_optional[k]) - if not (isinstance(data[k], Tensor) and data[k].requires_grad): - t = t.detach() + if k in neighbor_keys: + source = data[k] + t = next((value for previous, value in neighbor_memo if source is previous), None) + if t is None: + t = torch.as_tensor(source, device=self.device) + if not (isinstance(source, Tensor) and source.requires_grad): + t = t.detach() + neighbor_memo.append((source, t)) + else: + t = torch.as_tensor(data[k], device=self.device, dtype=self.keys_in_optional[k]) + if not (isinstance(data[k], Tensor) and data[k].requires_grad): + t = t.detach() ret[k] = t # Ensure all tensors have at least 1D shape for consistent batch processing for k, v in ret.items(): @@ -1421,6 +1468,11 @@ def mol_flatten(self, data: dict[str, Tensor], *, hessian: bool = False) -> dict Will not flatten for batched input and molecule size below threshold. """ ndim = data["coord"].ndim + explicit_mode2 = data.get("nbmat") is not None and data["nbmat"].ndim == 3 + if explicit_mode2: + self._batch = None + self._max_mol_size = data["coord"].shape[1] + return data if ndim == 2: self._batch = None if "mol_idx" not in data: @@ -1884,28 +1936,29 @@ def _hessian_vector_product_impl( prepared, forces=False, stress=False, hessian=external_hessian ) - coord = self._saved_for_grad["coord"] # (N+1, 3), requires_grad + coord = self._saved_for_grad["coord"] + coord_flat = coord.reshape(-1, 3) + real_indices = (~prepared["mask_i"].reshape(-1)).nonzero(as_tuple=False).flatten() + n_real = real_indices.numel() tot_energy = prepared["energy"].sum() # Differentiable part of the forces only. The detached ``coulomb_terms`` # forces are a constant w.r.t. coord (zero second derivative), so they are # intentionally excluded from the vjp; the periodic curvature they would # have carried is supplied by the directional FD helper instead. - forces_diff = -torch.autograd.grad(tot_energy, coord, create_graph=True)[0] # (N+1, 3) - N = coord.shape[0] - 1 + forces_diff = -torch.autograd.grad(tot_energy, coord, create_graph=True)[0] device = coord.device vecs = torch.as_tensor(vectors, device=device) single = vecs.ndim == 2 if single: vecs = vecs.unsqueeze(0) - if vecs.shape[-2:] != (N, 3): - raise ValueError(f"vectors must have trailing shape ({N}, 3); got {tuple(vecs.shape)}") + if vecs.shape[-2:] != (n_real, 3): + raise ValueError(f"vectors must have trailing shape ({n_real}, 3); got {tuple(vecs.shape)}") outs = [] for k in range(vecs.shape[0]): v = vecs[k].to(forces_diff.dtype) - v_full = torch.zeros_like(coord) - v_full[:N] = v + v_full = torch.zeros_like(coord_flat).index_copy(0, real_indices, v).reshape_as(coord) # autograd Hv = -d(forces . v)/dcoord = d^2E/dr^2 . v # (NN + short-range + dsf/simple charge-response + dftd3) hv_full = -torch.autograd.grad( @@ -1916,7 +1969,7 @@ def _hessian_vector_product_impl( create_graph=create_graph, allow_unused=True, )[0] - hv = hv_full[:N] + hv = hv_full.reshape(-1, 3).index_select(0, real_indices) if method == "pme" and self.external_coulomb is not None: # Full-periodic fixed-position curvature (directional FD). # Ewald needs no FD block: its energy is in the autograd graph diff --git a/aimnet/calculators/derivatives.py b/aimnet/calculators/derivatives.py index 7df8ae1..99f99ab 100644 --- a/aimnet/calculators/derivatives.py +++ b/aimnet/calculators/derivatives.py @@ -69,7 +69,13 @@ def set_grad_tensors( coord_unstrained = data["coord"] cell = data["cell"] cell_unstrained = cell - if cell.ndim == 2: + if coord_unstrained.ndim == 3: + B = coord_unstrained.shape[0] + scaling = torch.eye(3, dtype=cell.dtype, device=cell.device).unsqueeze(0).repeat(B, 1, 1) + scaling.requires_grad_(True) + data["coord"] = torch.einsum("bni,bij->bnj", coord_unstrained, scaling) + data["cell"] = cell @ scaling + elif cell.ndim == 2: # Single system: (3, 3) scaling scaling = torch.eye(3, requires_grad=True, dtype=cell.dtype, device=cell.device) data["coord"] = data["coord"] @ scaling @@ -136,7 +142,10 @@ def get_derivatives( volume = torch.linalg.det(cell).abs().unsqueeze(-1).unsqueeze(-1) # (B, 1, 1) data["stress"] = dedc / volume if hessian: - H = calculate_hessian(data["forces"], saved_for_grad["coord"]) + real_atom_mask = None + if saved_for_grad["coord"].ndim == 3 and "numbers" in data: + real_atom_mask = data["numbers"].ne(0) + H = calculate_hessian(data["forces"], saved_for_grad["coord"], real_atom_mask=real_atom_mask) if coulomb_terms is not None and getattr(coulomb_terms, "hessian", None) is not None: # The LR coulomb hessian is computed in float64 via finite # differences. Accumulate in that (higher) precision rather than @@ -146,7 +155,7 @@ def get_derivatives( return data -def calculate_hessian(forces: Tensor, coord: Tensor) -> Tensor: +def calculate_hessian(forces: Tensor, coord: Tensor, real_atom_mask: Tensor | None = None) -> Tensor: """Dense ``(N, 3, N, 3)`` Hessian of the energy w.r.t. real-atom coordinates. Autograd contract (IMPORTANT): @@ -169,6 +178,42 @@ def calculate_hessian(forces: Tensor, coord: Tensor) -> Tensor: periodic Ewald/PME long-range block remains a fixed-charge FD term in either case). """ + if real_atom_mask is not None: + coord_for_grad = coord + if coord.ndim == 3: + if coord.shape[0] != 1 or forces.shape[0] != 1 or real_atom_mask.shape[0] != 1: + raise ValueError("real_atom_mask Hessian calculation requires a singleton mode-2 system.") + coord_real = coord[0][real_atom_mask[0]] + forces = forces[0] + real_atom_mask = real_atom_mask[0] + else: + coord_real = coord[real_atom_mask] + if real_atom_mask.shape != (coord_for_grad.shape[-2],): + raise ValueError("real_atom_mask must align with the coordinate atom dimension.") + if forces.shape[0] == coord_for_grad.shape[-2]: + forces_real = forces[real_atom_mask] + elif forces.shape[0] == coord_real.shape[0]: + forces_real = forces + else: + raise ValueError("forces must contain either all atoms or all real atoms.") + n = forces_real.numel() + eye = torch.eye(n, device=forces_real.device, dtype=forces_real.dtype) + + def vjp_real(go: Tensor) -> Tensor: + grad = torch.autograd.grad( + forces_real.flatten(), + coord_for_grad, + grad_outputs=go, + retain_graph=True, + allow_unused=True, + )[0] + if grad is None: + return torch.zeros_like(coord_real) + grad_real = grad[0][real_atom_mask] if grad.ndim == 3 else grad[real_atom_mask] + return -grad_real + + return torch.func.vmap(vjp_real, 0)(eye).view(-1, 3, coord_real.shape[0], 3) + # Coord includes padding atom (shape N+1), forces only for real atoms (shape N). # Hessian computed only for actual atoms: (N, 3, N, 3). # diff --git a/aimnet/kernels/conv_sv_2d_sp_wp.py b/aimnet/kernels/conv_sv_2d_sp_wp.py index 3e62500..38a16f5 100644 --- a/aimnet/kernels/conv_sv_2d_sp_wp.py +++ b/aimnet/kernels/conv_sv_2d_sp_wp.py @@ -46,11 +46,10 @@ def _conv_sv_2d_sp_kernel( idx: wp.array2d(dtype=wp.int32), # (B, M) g: wp.array3d(dtype=wp.vec4f), # (B, M, G, D) output: wp.array3d(dtype=wp.vec4f), # (B, A, G, D) + padding_value: int, ): """Forward: output[b,a,g] = sum_m a[idx[b,m],a,g] * g[b,m,g]""" - B, M = idx.shape[0], idx.shape[1] - padding_value = B - 1 # last row is padding - + M = idx.shape[1] _b, _a, _g = wp.tid() acc = wp.vec4f() @@ -71,11 +70,10 @@ def _conv_sv_2d_sp_backward_a_kernel( idx: wp.array2d(dtype=wp.int32), # (B, M) g: wp.array3d(dtype=wp.vec4f), # (B, M, G, D) grad_a: wp.array3d(dtype=wp.float32), # (B, A, G) + padding_value: int, ): """Backward w.r.t. a: grad_a[idx[b,m],a,g] += dot(grad_output[b,a,g], g[b,m,g])""" - B, M = idx.shape[0], idx.shape[1] - padding_value = B - 1 # last row is padding - + M = idx.shape[1] _b, _a, _g = wp.tid() grad_out = grad_output[_b, _a, _g] @@ -95,12 +93,10 @@ def _conv_sv_2d_sp_backward_g_kernel( a: wp.array3d(dtype=wp.float32), # (B, A, G) idx: wp.array2d(dtype=wp.int32), # (B, M) grad_g: wp.array3d(dtype=wp.vec4f), # (B, M, G, D) + padding_value: int, ): """Backward w.r.t. g: grad_g[b,m,g] = sum_a a[idx[b,m],a,g] * grad_output[b,a,g]""" - B = idx.shape[0] A = a.shape[1] - padding_value = B - 1 # last row is padding - _b, _m, _g = wp.tid() _idx = idx[_b, _m] @@ -123,12 +119,10 @@ def _conv_sv_2d_sp_double_backward_a_g_kernel( idx: wp.array2d(dtype=wp.int32), # (B, M) grad_output: wp.array3d(dtype=wp.vec4f), # (B, A, G, D) grad_g: wp.array3d(dtype=wp.vec4f), # (B, M, G, D) + padding_value: int, ): """Double backward: d(grad_a)/dg -> grad_g""" - B = idx.shape[0] A = grad_grad_a.shape[1] - padding_value = B - 1 # last row is padding - _b, _m, _g = wp.tid() _idx = idx[_b, _m] @@ -151,11 +145,10 @@ def _conv_sv_2d_sp_double_backward_g_contrib_kernel( a: wp.array3d(dtype=wp.float32), # (B, A, G) idx: wp.array2d(dtype=wp.int32), # (B, M) grad_output_double: wp.array3d(dtype=wp.vec4f), # (B, A, G, D) - OUTPUT + padding_value: int, ): """Double backward from grad2_g: einsum('bmgd,bmag->bagd', grad2_g, a_selected)""" - B, M = idx.shape[0], idx.shape[1] - padding_value = B - 1 # last row is padding - + M = idx.shape[1] _b, _a, _g = wp.tid() acc = wp.vec4f() @@ -177,11 +170,10 @@ def _conv_sv_2d_sp_double_backward_a_contrib_kernel( idx: wp.array2d(dtype=wp.int32), # (B, M) g: wp.array3d(dtype=wp.vec4f), # (B, M, G, D) grad_output_double: wp.array3d(dtype=wp.vec4f), # (B, A, G, D) - OUTPUT + padding_value: int, ): """Double backward from grad2_a: einsum('bmag,bmgd->bagd', grad2_a_selected, g)""" - B, M = idx.shape[0], idx.shape[1] - padding_value = B - 1 # last row is padding - + M = idx.shape[1] _b, _a, _g = wp.tid() acc = wp.vec4f() @@ -202,21 +194,33 @@ def _conv_sv_2d_sp_double_backward_a_contrib_kernel( # ============================================================================= +def _validate_conv_sv_sizes(a: Tensor, idx: Tensor, padding_value: int, num_centers: int) -> None: + """Validate the flattened atom and center capacities used by Warp.""" + if a.ndim != 3 or idx.ndim != 2: + raise ValueError("a must be 3D and idx must be 2D.") + if padding_value < 0 or padding_value > a.shape[0]: + raise ValueError("padding_value must be in [0, a.shape[0]].") + if num_centers < 0 or num_centers > idx.shape[0]: + raise ValueError("num_centers must be in [0, idx.shape[0]].") + + @torch.library.custom_op( "aimnet::conv_sv_2d_sp_fwd", mutates_args=(), device_types=["cuda"], ) -def _(a: Tensor, idx: Tensor, g: Tensor) -> Tensor: +def _(a: Tensor, idx: Tensor, g: Tensor, padding_value: int, num_centers: int) -> Tensor: """Forward primitive for conv_sv_2d_sp.""" + _validate_conv_sv_sizes(a, idx, padding_value, num_centers) stream = _get_stream(a.device) device = wp.device_from_torch(a.device) - B, A, G = a.shape - output = torch.zeros(B, A, G, 4, dtype=a.dtype, device=a.device) + _B, A, G = a.shape + B_out = idx.shape[0] + output = torch.zeros(B_out, A, G, 4, dtype=a.dtype, device=a.device) wp.launch( _conv_sv_2d_sp_kernel, - dim=(B - 1, A, G), # B-1: exclude padding row + dim=(num_centers, A, G), stream=stream, device=device, inputs=( @@ -224,15 +228,17 @@ def _(a: Tensor, idx: Tensor, g: Tensor) -> Tensor: wp.from_torch(idx.to(torch.int32), return_ctype=True), wp.from_torch(g.detach(), return_ctype=True, dtype=wp.vec4f), wp.from_torch(output, return_ctype=True, dtype=wp.vec4f), + padding_value, ), ) return output @torch.library.register_fake("aimnet::conv_sv_2d_sp_fwd") -def _(a: Tensor, idx: Tensor, g: Tensor) -> Tensor: - B, A, G = a.shape - return torch.empty(B, A, G, 4, dtype=a.dtype, device=a.device) +def _(a: Tensor, idx: Tensor, g: Tensor, padding_value: int, num_centers: int) -> Tensor: + _validate_conv_sv_sizes(a, idx, padding_value, num_centers) + _B, A, G = a.shape + return torch.empty(idx.shape[0], A, G, 4, dtype=a.dtype, device=a.device) @torch.library.custom_op( @@ -240,11 +246,12 @@ def _(a: Tensor, idx: Tensor, g: Tensor) -> Tensor: mutates_args=(), device_types=["cuda"], ) -def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor) -> list[Tensor]: +def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor, padding_value: int, num_centers: int) -> list[Tensor]: """Backward primitive for conv_sv_2d_sp.""" + _validate_conv_sv_sizes(a, idx, padding_value, num_centers) stream = _get_stream(a.device) device = wp.device_from_torch(a.device) - B, A, G = a.shape + _B, A, G = a.shape B_out, M = idx.shape grad_a = torch.zeros_like(a) @@ -255,7 +262,7 @@ def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor) -> list[Tensor]: # Launch backward w.r.t. a wp.launch( _conv_sv_2d_sp_backward_a_kernel, - dim=(B - 1, A, G), # B-1: exclude padding row + dim=(num_centers, A, G), stream=stream, device=device, inputs=( @@ -263,13 +270,14 @@ def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor) -> list[Tensor]: wp.from_torch(idx.to(torch.int32), return_ctype=True), wp.from_torch(g.detach(), return_ctype=True, dtype=wp.vec4f), wp.from_torch(grad_a, return_ctype=True), + padding_value, ), ) # Launch backward w.r.t. g wp.launch( _conv_sv_2d_sp_backward_g_kernel, - dim=(B_out - 1, M, G), # B_out-1: exclude padding row + dim=(num_centers, M, G), stream=stream, device=device, inputs=( @@ -277,6 +285,7 @@ def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor) -> list[Tensor]: wp.from_torch(a.detach(), return_ctype=True), wp.from_torch(idx.to(torch.int32), return_ctype=True), wp.from_torch(grad_g, return_ctype=True, dtype=wp.vec4f), + padding_value, ), ) @@ -284,9 +293,10 @@ def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor) -> list[Tensor]: @torch.library.register_fake("aimnet::conv_sv_2d_sp_bwd") -def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor) -> list[Tensor]: +def _(grad_output: Tensor, a: Tensor, idx: Tensor, g: Tensor, padding_value: int, num_centers: int) -> list[Tensor]: B_out, M = idx.shape G = a.shape[2] + _validate_conv_sv_sizes(a, idx, padding_value, num_centers) return [ torch.empty_like(a), torch.empty(B_out, M, G, 4, dtype=g.dtype, device=g.device), @@ -305,14 +315,17 @@ def _( a: Tensor, idx: Tensor, g: Tensor, + padding_value: int, + num_centers: int, ) -> list[Tensor]: """Double backward primitive for conv_sv_2d_sp.""" + _validate_conv_sv_sizes(a, idx, padding_value, num_centers) stream = _get_stream(a.device) device = wp.device_from_torch(a.device) - B, A, G = a.shape + _B, A, G = a.shape B_out, M = idx.shape - grad_grad_output = torch.zeros(B, A, G, 4, dtype=a.dtype, device=a.device) + grad_grad_output = torch.zeros(B_out, A, G, 4, dtype=a.dtype, device=a.device) grad_a_double = torch.zeros_like(a) grad_g_double = torch.zeros(B_out, M, G, 4, dtype=a.dtype, device=a.device) @@ -323,7 +336,7 @@ def _( # Contribution from grad2_g to grad_grad_output wp.launch( _conv_sv_2d_sp_double_backward_g_contrib_kernel, - dim=(B - 1, A, G), # B-1: exclude padding row + dim=(num_centers, A, G), stream=stream, device=device, inputs=( @@ -331,14 +344,15 @@ def _( wp.from_torch(a.detach(), return_ctype=True), wp.from_torch(idx.to(torch.int32), return_ctype=True), wp.from_torch(grad_grad_output, return_ctype=True, dtype=wp.vec4f), + padding_value, ), ) # Contribution from grad2_a to grad_grad_output - grad_output_2_a = torch.zeros(B, A, G, 4, dtype=a.dtype, device=a.device) + grad_output_2_a = torch.zeros(B_out, A, G, 4, dtype=a.dtype, device=a.device) wp.launch( _conv_sv_2d_sp_double_backward_a_contrib_kernel, - dim=(B - 1, A, G), # B-1: exclude padding row + dim=(num_centers, A, G), stream=stream, device=device, inputs=( @@ -346,6 +360,7 @@ def _( wp.from_torch(idx.to(torch.int32), return_ctype=True), wp.from_torch(g.detach(), return_ctype=True, dtype=wp.vec4f), wp.from_torch(grad_output_2_a, return_ctype=True, dtype=wp.vec4f), + padding_value, ), ) grad_grad_output = grad_grad_output + grad_output_2_a @@ -353,7 +368,7 @@ def _( # Mixed partial: d(grad_a)/dg -> grad_g_double wp.launch( _conv_sv_2d_sp_double_backward_a_g_kernel, - dim=(B_out - 1, M, G), # B_out-1: exclude padding row + dim=(num_centers, M, G), stream=stream, device=device, inputs=( @@ -361,13 +376,14 @@ def _( wp.from_torch(idx.to(torch.int32), return_ctype=True), wp.from_torch(grad_output_contig, return_ctype=True, dtype=wp.vec4f), wp.from_torch(grad_g_double, return_ctype=True, dtype=wp.vec4f), + padding_value, ), ) # Mixed partial: d(grad_g)/da -> grad_a_double wp.launch( _conv_sv_2d_sp_backward_a_kernel, - dim=(B - 1, A, G), # B-1: exclude padding row + dim=(num_centers, A, G), stream=stream, device=device, inputs=( @@ -375,6 +391,7 @@ def _( wp.from_torch(idx.to(torch.int32), return_ctype=True), wp.from_torch(grad2_g_contig, return_ctype=True, dtype=wp.vec4f), wp.from_torch(grad_a_double, return_ctype=True), + padding_value, ), ) @@ -389,11 +406,14 @@ def _( a: Tensor, idx: Tensor, g: Tensor, + padding_value: int, + num_centers: int, ) -> list[Tensor]: - B, A, G = a.shape + _B, A, G = a.shape B_out, M = idx.shape + _validate_conv_sv_sizes(a, idx, padding_value, num_centers) return [ - torch.empty(B, A, G, 4, dtype=a.dtype, device=a.device), + torch.empty(B_out, A, G, 4, dtype=a.dtype, device=a.device), torch.empty_like(a), torch.empty(B_out, M, G, 4, dtype=a.dtype, device=a.device), ] @@ -406,22 +426,28 @@ def _( def _conv_sv_2d_sp_setup_fwd_context(ctx, inputs, output): """Setup context for forward pass.""" - a, idx, g = inputs + a, idx, g, padding_value, num_centers = inputs ctx.save_for_backward(a, idx, g) + ctx.padding_value = padding_value + ctx.num_centers = num_centers def _conv_sv_2d_sp_setup_bwd_context(ctx, inputs, output): """Setup context for backward pass.""" - grad_output, a, idx, g = inputs + grad_output, a, idx, g, padding_value, num_centers = inputs ctx.save_for_backward(grad_output, a, idx, g) + ctx.padding_value = padding_value + ctx.num_centers = num_centers @torch.compiler.allow_in_graph def _conv_sv_2d_sp_bwd(ctx, grad_output): """Backward pass for conv_sv_2d_sp.""" a, idx, g = ctx.saved_tensors - grad_a, grad_g = torch.ops.aimnet.conv_sv_2d_sp_bwd(grad_output.contiguous(), a, idx, g) - return grad_a, None, grad_g + grad_a, grad_g = torch.ops.aimnet.conv_sv_2d_sp_bwd( + grad_output.contiguous(), a, idx, g, ctx.padding_value, ctx.num_centers + ) + return grad_a, None, grad_g, None, None @torch.compiler.allow_in_graph @@ -439,9 +465,11 @@ def _conv_sv_2d_sp_bwd_bwd(ctx, *grad_outputs): G = a.shape[2] grad2_g = torch.zeros(B_out, M, G, 4, dtype=g.dtype, device=g.device) - outputs = torch.ops.aimnet.conv_sv_2d_sp_bwd_bwd(grad_output_saved, grad2_a, grad2_g, a, idx, g) + outputs = torch.ops.aimnet.conv_sv_2d_sp_bwd_bwd( + grad_output_saved, grad2_a, grad2_g, a, idx, g, ctx.padding_value, ctx.num_centers + ) - return outputs[0], outputs[1], None, outputs[2] + return outputs[0], outputs[1], None, outputs[2], None, None torch.library.register_autograd( @@ -462,6 +490,11 @@ def _conv_sv_2d_sp_bwd_bwd(ctx, *grad_outputs): # ============================================================================= +@torch.library.register_vmap("aimnet::conv_sv_2d_sp_fwd") +def _vmap_conv_sv_2d_sp_fwd(info, in_dims, a, idx, g, padding_value, num_centers): + raise RuntimeError("aimnet::conv_sv_2d_sp_fwd does not support direct vmap.") + + def _vmap_slice(t: Tensor, d: int | None, k: int) -> Tensor: """Pick the k-th slice along vmap batch dim d, or pass through if not batched. @@ -475,7 +508,7 @@ def _vmap_slice(t: Tensor, d: int | None, k: int) -> Tensor: @torch.library.register_vmap("aimnet::conv_sv_2d_sp_bwd") -def _vmap_conv_sv_2d_sp_bwd(info, in_dims, grad_output, a, idx, g): +def _vmap_conv_sv_2d_sp_bwd(info, in_dims, grad_output, a, idx, g, padding_value, num_centers): """vmap rule for the first-backward primitive. Hit when torch.func.vmap traverses a vjp closure that reaches the first-order @@ -502,6 +535,8 @@ def _vmap_conv_sv_2d_sp_bwd(info, in_dims, grad_output, a, idx, g): _vmap_slice(a, in_dims[1], k), _vmap_slice(idx, in_dims[2], k), _vmap_slice(g, in_dims[3], k), + padding_value, + num_centers, ) out0.append(outs[0]) out1.append(outs[1]) @@ -513,7 +548,7 @@ def _vmap_conv_sv_2d_sp_bwd(info, in_dims, grad_output, a, idx, g): @torch.library.register_vmap("aimnet::conv_sv_2d_sp_bwd_bwd") -def _vmap_conv_sv_2d_sp_bwd_bwd(info, in_dims, grad_output, grad2_a, grad2_g, a, idx, g): +def _vmap_conv_sv_2d_sp_bwd_bwd(info, in_dims, grad_output, grad2_a, grad2_g, a, idx, g, padding_value, num_centers): """vmap rule for the double-backward primitive. Hit when torch.func.vmap traverses a vjp closure that reaches the second-order @@ -550,6 +585,8 @@ def _vmap_conv_sv_2d_sp_bwd_bwd(info, in_dims, grad_output, grad2_a, grad2_g, a, _vmap_slice(a, in_dims[3], k), _vmap_slice(idx, in_dims[4], k), _vmap_slice(g, in_dims[5], k), + padding_value, + num_centers, ) out0.append(outs[0]) out1.append(outs[1]) @@ -566,22 +603,32 @@ def _vmap_conv_sv_2d_sp_bwd_bwd(info, in_dims, grad_output, grad2_a, grad2_g, a, # ============================================================================= -def conv_sv_2d_sp(a: Tensor, idx: Tensor, g: Tensor) -> Tensor: +def conv_sv_2d_sp( + a: Tensor, + idx: Tensor, + g: Tensor, + padding_value: int | None = None, + num_centers: int | None = None, +) -> Tensor: """Compute conv_sv_2d_sp with support for 1st and 2nd order derivatives. Parameters ---------- a : Tensor - Input tensor of shape (B, A, G). + Input tensor of shape (K, A, G), where K is the atom capacity. idx : Tensor - Index tensor of shape (B, M). + Index tensor of shape (C, M), where C is the center capacity. g : Tensor - Gate tensor of shape (B, M, G, 4). + Gate tensor of shape (C, M, G, 4). + padding_value : int, optional + Sentinel threshold in the atom dimension. Defaults to ``K - 1``. + num_centers : int, optional + Number of centers launched by Warp. Defaults to ``C - 1``. Notes ----- ``idx`` rows must follow the packed-padding contract: real neighbor indices - come first and padding sentinels (values >= B-1, the padding row) are + come first and padding sentinels (values >= ``padding_value``) are contiguous at the end of each row. The Warp kernels stop scanning a row at the first sentinel, so interleaved padding would silently drop real neighbors. ``nvalchemiops.torch.neighbors.neighbor_list`` produces this @@ -597,8 +644,10 @@ def conv_sv_2d_sp(a: Tensor, idx: Tensor, g: Tensor) -> Tensor: Returns ------- Tensor - Output tensor of shape (B, A, G, 4). + Output tensor of shape (C, A, G, 4). """ + if (padding_value is None) != (num_centers is None): + raise ValueError("padding_value and num_centers must be supplied together.") if a.device.type != "cuda" or idx.device.type != "cuda" or g.device.type != "cuda": raise RuntimeError("conv_sv_2d_sp is a CUDA-only Warp kernel") if a.dtype != torch.float32 or g.dtype != torch.float32: @@ -607,12 +656,17 @@ def conv_sv_2d_sp(a: Tensor, idx: Tensor, g: Tensor) -> Tensor: idx = idx.to(torch.int32) if a.ndim != 3 or idx.ndim != 2 or g.ndim != 4 or g.shape[-1] != 4: raise ValueError("Expected shapes a=(B,A,G), idx=(B,M), g=(B,M,G,4)") - if a.shape[0] != idx.shape[0] or idx.shape[0] != g.shape[0] or a.shape[2] != g.shape[2]: + if idx.shape[0] != g.shape[0] or a.shape[2] != g.shape[2]: raise ValueError("Incompatible conv_sv_2d_sp leading or basis dimensions") + if padding_value is None: + padding_value = a.shape[0] - 1 + num_centers = idx.shape[0] - 1 + assert num_centers is not None + _validate_conv_sv_sizes(a, idx, padding_value, num_centers) if not a.is_contiguous(): a = a.contiguous() if not idx.is_contiguous(): idx = idx.contiguous() if not g.is_contiguous(): g = g.contiguous() - return torch.ops.aimnet.conv_sv_2d_sp_fwd(a, idx, g) + return torch.ops.aimnet.conv_sv_2d_sp_fwd(a, idx, g, padding_value, num_centers) diff --git a/aimnet/models/base.py b/aimnet/models/base.py index 629e717..b3a256b 100644 --- a/aimnet/models/base.py +++ b/aimnet/models/base.py @@ -222,13 +222,31 @@ def _prepare_dtype(self, data: dict[str, Tensor]) -> dict[str, Tensor]: for k, d in zip(self._required_keys, self._required_keys_dtype, strict=False): assert k in data, f"Key {k} is required" data[k] = data[k].to(d) + neighbor_keys = {f"nbmat{suffix}" for suffix in nbops.NBMAT_SUFFIXES} + converted_neighbors: list[tuple[Tensor, Tensor]] = [] + for key in neighbor_keys: + if key not in data: + continue + source = data[key] + converted = next((value for previous, value in converted_neighbors if source is previous), None) + if converted is None: + converted = source if source.dtype == torch.int32 else source.to(torch.int32) + converted_neighbors.append((source, converted)) + data[key] = converted for k, d in zip(self._optional_keys, self._optional_keys_dtype, strict=False): - if k in data: + if k in data and k not in neighbor_keys: data[k] = data[k].to(d) return data def prepare_input(self, data: dict[str, Tensor]) -> dict[str, Tensor]: """Common operations for input preparation.""" + nbmat = data.get("nbmat") + nbops.validate_neighbor_suffix_layout(data) + if isinstance(nbmat, Tensor) and nbmat.ndim == 3: + nbops.normalize_mode2_periodic_geometry(data, B=nbmat.shape[0]) + for suffix in nbops.NBMAT_SUFFIXES: + if f"nbmat{suffix}" in data or f"shifts{suffix}" in data: + nbops.validate_mode2_nbmat_raw(data, suffix=suffix) data = self._prepare_dtype(data) data = nbops.set_nb_mode(data) data = nbops.calc_masks(data) diff --git a/aimnet/modules/aev.py b/aimnet/modules/aev.py index 37e0bea..8f18b40 100644 --- a/aimnet/modules/aev.py +++ b/aimnet/modules/aev.py @@ -155,7 +155,28 @@ def output_size(self): def forward(self, data: dict[str, Tensor], a: Tensor) -> Tensor: g_sv = data["g_sv"] mode = nbops.get_nb_mode(data) - if self.d2features: + if mode == 2: + a_flat = _flatten_mode2(a, "atomic features") + g_flat = _flatten_mode2(g_sv, "ConvSV geometry") + gather_flat = _flatten_mode2(data["_nbmat_gather"], "ConvSV gather indices") + kernel_flat = _flatten_mode2(data["_nbmat_kernel"], "ConvSV kernel indices") + b, n = a.shape[:2] + if self.d2features: + if a.device.type == "cuda" and a.dtype == torch.float32 and g_sv.dtype == torch.float32: + avf_sv = conv_sv_2d_sp( + a_flat, + kernel_flat, + g_flat, + padding_value=b * n, + num_centers=b * n, + ).unflatten(0, (b, n)) + else: + a_j = a_flat.index_select(0, gather_flat.flatten()).unflatten(0, data["_nbmat_gather"].shape) + avf_sv = torch.einsum("...mag,...mgd->...agd", a_j, g_sv) + else: + a_j = a_flat.index_select(0, gather_flat.flatten()).unflatten(0, data["_nbmat_gather"].shape) + avf_sv = torch.einsum("...ma,...mgd->...agd", a_j, g_sv) + elif self.d2features: # The Warp kernel is float32-only; float64 (and mixed-dtype) inputs # fall through to the pure-torch einsum branch below. if mode > 0 and a.device.type == "cuda" and a.dtype == torch.float32 and g_sv.dtype == torch.float32: @@ -173,7 +194,18 @@ def forward(self, data: dict[str, Tensor], a: Tensor) -> Tensor: avf_sv = torch.einsum("...ma,...mgd->...agd", a.unsqueeze(1), g_sv) avf_s, avf_v = avf_sv.split([1, 3], dim=-1) avf_v = torch.einsum("agh,...agd->...ahd", self.agh, avf_v).pow(2).sum(-1) - return torch.cat([avf_s.squeeze(-1).flatten(-2, -1), avf_v.flatten(-2, -1)], dim=-1) + out = torch.cat([avf_s.squeeze(-1).flatten(-2, -1), avf_v.flatten(-2, -1)], dim=-1) + if mode == 1: + return out + return nbops.mask_i_(out, data, mask_value=0.0, inplace=False) + + +def _flatten_mode2(tensor: Tensor, name: str) -> Tensor: + """Flatten system and atom dimensions without allocating.""" + flattened = tensor.flatten(0, 1) + if flattened._base is None: + raise ValueError(f"mode-2 {name} must be view-flattenable") + return flattened def _init_ahg(b: int, m: int, n: int): diff --git a/aimnet/modules/lr.py b/aimnet/modules/lr.py index 4fb89a7..2abe5ac 100644 --- a/aimnet/modules/lr.py +++ b/aimnet/modules/lr.py @@ -71,6 +71,48 @@ class ExternalDerivativeTerms: hessian: Tensor | None = None +class _Mode2BackendInputs(NamedTuple): + coord: Tensor + neighbor_matrix: Tensor + shifts: Tensor | None + batch_idx: Tensor + fill_value: int + num_systems: int + cell: Tensor | None + + +def _mode2_backend_inputs(data: dict[str, Tensor], suffix: str) -> _Mode2BackendInputs: + """Return the prepared global mode-2 tensors in backend layout.""" + coord = data["coord"] + B, N = coord.shape[:2] + neighbor_matrix_source = data.get(f"_nbmat_kernel{suffix}", data[f"nbmat{suffix}"]) + coord_flat = _flatten_backend_view(coord, "mode-2 coordinates") + neighbor_matrix_source = neighbor_matrix_source.to(torch.int32) + neighbor_matrix = _flatten_backend_view(neighbor_matrix_source, f"mode-2 nbmat{suffix}") + shifts_source = data.get(f"shifts{suffix}") + shifts = _flatten_backend_view(shifts_source, f"mode-2 shifts{suffix}") if shifts_source is not None else None + cell = data.get("cell") + if cell is not None and cell.ndim == 2: + cell = cell.unsqueeze(0) + batch_idx = torch.arange(B, device=coord.device, dtype=torch.int32).repeat_interleave(N) + return _Mode2BackendInputs( + coord=coord_flat, + neighbor_matrix=neighbor_matrix, + shifts=shifts, + batch_idx=batch_idx, + fill_value=B * N, + num_systems=B, + cell=cell, + ) + + +def _flatten_backend_view(tensor: Tensor, name: str) -> Tensor: + flattened = tensor.flatten(0, 1) + if flattened._base is None: + raise ValueError(f"{name} must be view-flattenable across (B, N).") + return flattened + + def _periodic_coulomb_hybrid( *, coord: Tensor, @@ -440,48 +482,22 @@ def _dsf_inputs_mode2( data: dict[str, Tensor], suffix: str, ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor | None, Tensor | None, int, int]: - """Flatten batched neighbor-matrix inputs for nvalchemiops DSF.""" - coord = data["coord"] - charges = data[self.key_in].to(coord.dtype).masked_fill(data["mask_i"], 0.0) - B, N = coord.shape[:2] - fill_value = B * N - - positions = torch.cat([coord.reshape(B * N, 3), coord.new_zeros(1, 3)], dim=0) - charges_flat = torch.cat([charges.reshape(B * N), charges.new_zeros(1)], dim=0) - batch_idx = torch.cat( - [ - torch.repeat_interleave(torch.arange(B, device=coord.device, dtype=torch.int32), N), - torch.zeros(1, device=coord.device, dtype=torch.int32), - ], - dim=0, - ) - - # nbmat values are atom indices LOCAL to each batch; the flattened - # `positions` tensor uses GLOBAL indices (b * N + i), so each batch's - # neighbor entries must be offset by b * N. Without this, valid - # neighbors in batch b > 0 silently point into batch 0 and DSF energy - # is wrong. Mirrors the offset that DFTD3 mode-2 already applies. - nbmat_local = data[f"nbmat{suffix}"].to(torch.int32) - offsets = (torch.arange(B, device=coord.device, dtype=torch.int32) * N).view(B, 1, 1) - nbmat = (nbmat_local + offsets).flatten(0, 1) - mask_ij = data[f"mask_ij{suffix}"].flatten(0, 1) - nbmat = torch.where(mask_ij, torch.full_like(nbmat, fill_value), nbmat) - nbmat = torch.cat( - [nbmat, torch.full((1, nbmat.shape[1]), fill_value, dtype=torch.int32, device=coord.device)], - dim=0, + mode2 = _mode2_backend_inputs(data, suffix) + charges = data[self.key_in].to(data["coord"].dtype) + charges_flat = charges.flatten(0, 1).masked_fill(data["mask_i"].flatten(), 0.0) + shifts = mode2.shifts.to(torch.int32) if mode2.shifts is not None else None + cell = mode2.cell.to(data["coord"].dtype) if mode2.cell is not None else None + return ( + mode2.coord, + charges_flat, + mode2.batch_idx, + mode2.neighbor_matrix, + cell, + shifts, + mode2.fill_value, + mode2.num_systems, ) - cell = data.get("cell") - shifts = None - if cell is not None: - cell = cell.to(coord.dtype) - if cell.ndim == 2: - cell = cell.unsqueeze(0).expand(B, -1, -1) - shifts = data[f"shifts{suffix}"].flatten(0, 1).to(torch.int32) - shifts = torch.cat([shifts, torch.zeros((1, shifts.shape[1], 3), dtype=torch.int32, device=coord.device)]) - - return positions, charges_flat, batch_idx, nbmat, cell, shifts, fill_value, B - def _dsf_inputs( self, data: dict[str, Tensor], @@ -502,8 +518,10 @@ def _restore_dsf_forces_shape(data: dict[str, Tensor], forces: Tensor) -> Tensor nb_mode = nbops.get_nb_mode(data) if nb_mode == 1: return forces - if nb_mode in (0, 2): + if nb_mode == 0: return forces[:-1].reshape_as(data["coord"]) + if nb_mode == 2: + return forces.reshape_as(data["coord"]) raise ValueError(f"Invalid neighbor mode: {nb_mode}") def _coul_dsf_nvalchemi( @@ -637,9 +655,9 @@ def _coul_nvalchemi( Requires ``cell`` in ``data`` and a PBC neighbor list under ``nbmat_coulomb``/``shifts_coulomb`` (preferred) or the shared - ``nbmat_lr``/``shifts_lr``. Drops the trailing padding row before - invoking the backend and re-adds a zero pad row so downstream - ``unpad_output`` contracts are preserved. + ``nbmat_lr``/``shifts_lr``. Mode 2 preserves all batch-major dummy + rows through zero-copy flattened backend views; flat mode 1 removes + and restores its single trailing padding row. """ suffix = nbops.resolve_suffix(data, ["_coulomb", "_lr"]) @@ -648,29 +666,43 @@ def _coul_nvalchemi( if cell is None: raise ValueError("nvalchemi Coulomb requires periodic cell data") - charges = data[self.key_in] - mol_idx = data["mol_idx"] - nbmat = data[f"nbmat{suffix}"] - shifts = data[f"shifts{suffix}"] - if coord.ndim != 2 or charges.ndim != 1 or mol_idx.ndim != 1 or nbmat.ndim != 2: - raise ValueError("nvalchemi Coulomb expects flat padded PBC inputs") - if not (coord.shape[0] == charges.shape[0] == mol_idx.shape[0] == nbmat.shape[0]): - raise ValueError("nvalchemi Coulomb flat inputs must include matching coord/charge/mol_idx/nbmat rows") - if coord.shape[0] < 2: - raise ValueError("nvalchemi Coulomb flat inputs must include at least one real atom and one padding row") - - # Drop the trailing padding atom (flat mode includes one at index N). - N_padded = coord.shape[0] - N = N_padded - 1 - coord_real = coord[:-1] - charges_real = charges[:-1] - mol_idx_real = mol_idx[:-1].to(torch.int32) - nbmat_real = nbmat[:-1].to(torch.int32) - shifts_real = shifts[:-1].to(torch.int32) + mode2 = nbops.get_nb_mode(data) == 2 + if mode2: + mode2_inputs = _mode2_backend_inputs(data, suffix) + coord_real = mode2_inputs.coord + charges_real = data[self.key_in].to(coord.dtype).flatten(0, 1).masked_fill(data["mask_i"].flatten(), 0.0) + mol_idx_real = mode2_inputs.batch_idx + nbmat_real = mode2_inputs.neighbor_matrix + shifts_real = mode2_inputs.shifts.to(torch.int32) if mode2_inputs.shifts is not None else None + cell = mode2_inputs.cell.to(coord.dtype) if mode2_inputs.cell is not None else None + N = mode2_inputs.fill_value + num_systems = mode2_inputs.num_systems + else: + charges = data[self.key_in] + mol_idx = data["mol_idx"] + nbmat = data[f"nbmat{suffix}"] + shifts = data[f"shifts{suffix}"] + if coord.ndim != 2 or charges.ndim != 1 or mol_idx.ndim != 1 or nbmat.ndim != 2: + raise ValueError("nvalchemi Coulomb expects flat padded PBC inputs") + if not (coord.shape[0] == charges.shape[0] == mol_idx.shape[0] == nbmat.shape[0]): + raise ValueError("nvalchemi Coulomb flat inputs must include matching coord/charge/mol_idx/nbmat rows") + if coord.shape[0] < 2: + raise ValueError( + "nvalchemi Coulomb flat inputs must include at least one real atom and one padding row" + ) + + # Flat mode reserves one final padding atom for the backend fill value. + N_padded = coord.shape[0] + N = N_padded - 1 + coord_real = coord[:-1] + charges_real = charges[:-1] + mol_idx_real = mol_idx[:-1].to(torch.int32) + nbmat_real = nbmat[:-1].to(torch.int32) + shifts_real = shifts[:-1].to(torch.int32) + num_systems = int(mol_idx_real.max().item()) + 1 if backend not in ("ewald", "pme"): raise ValueError(f"backend must be 'ewald' or 'pme', got {backend!r}") - num_systems = int(mol_idx_real.max().item()) + 1 fn = particle_mesh_ewald if backend == "pme" else ewald_summation if backend == "ewald": @@ -713,8 +745,9 @@ def _coul_nvalchemi( if needs_strain_grad: if coord_unstrained is None or cell_unstrained is None: raise ValueError("scaling-aware Coulomb requires coord_unstrained and cell_unstrained") + coord_unstrained_backend = coord_unstrained.flatten(0, 1) if mode2 else coord_unstrained[:-1] e_periodic, _forces, _charge_grad, _virial = _PeriodicCoulombFunction.apply( - coord_unstrained[:-1], + coord_unstrained_backend, cell_unstrained, scaling, charges_real, @@ -796,7 +829,11 @@ def _coul_nvalchemi( if compute_forces or compute_virial: forces = None if forces_real is not None: - forces = torch.cat([forces_real.detach() * ke, forces_real.new_zeros((1, 3))], dim=0) + forces = forces_real.detach() * ke + if mode2: + forces = forces.view_as(data["coord"]) + else: + forces = torch.cat([forces, forces_real.new_zeros((1, 3))], dim=0) virial_ev = virial.detach() * ke if virial is not None else None terms = ExternalDerivativeTerms(forces=forces, virial=virial_ev) e_periodic = energies_per_system @@ -807,11 +844,7 @@ def _coul_nvalchemi( return e_periodic, terms def _periodic_fd_setup(self, data: dict[str, Tensor], backend: str): - """Shared marshalling for the periodic-Coulomb finite-difference Hessian - and Hessian-vector-product paths. Returns ``(forces_at, coord_real, N)`` - where ``forces_at(positions)`` evaluates the analytic full-periodic - Coulomb forces (N, 3) in eV/Ang at float64, with neighbor list, cell, and - (detached) charges held fixed.""" + """Prepare fixed-charge periodic force evaluations for finite differences.""" suffix = nbops.resolve_suffix(data, ["_coulomb", "_lr"]) coord = data["coord"] cell = data["cell"] @@ -819,14 +852,30 @@ def _periodic_fd_setup(self, data: dict[str, Tensor], backend: str): raise ValueError("nvalchemi Coulomb requires periodic cell data") if backend not in ("ewald", "pme"): raise ValueError(f"backend must be 'ewald' or 'pme', got {backend!r}") - N = coord.shape[0] - 1 - coord_real = coord[:-1].detach().double() - charges_real = data[self.key_in][:-1].detach().double() - mol_idx_real = data["mol_idx"][:-1].to(torch.int32) - nbmat_real = data[f"nbmat{suffix}"][:-1].to(torch.int32) - shifts_real = data[f"shifts{suffix}"][:-1].to(torch.int32) - cell_det = cell.detach().double() - num_systems = int(mol_idx_real.max().item()) + 1 + mode2 = nbops.get_nb_mode(data) == 2 + if mode2: + inputs = _mode2_backend_inputs(data, suffix) + coord_real = inputs.coord.detach().double() + charges_real = ( + data[self.key_in].flatten(0, 1).masked_fill(data["numbers"].flatten().eq(0), 0.0).detach().double() + ) + mol_idx_real = inputs.batch_idx + nbmat_real = inputs.neighbor_matrix + shifts_real = inputs.shifts.to(torch.int32) if inputs.shifts is not None else None + cell_det = inputs.cell.detach().double() if inputs.cell is not None else cell.detach().double() + real_indices = data["numbers"].flatten().ne(0).nonzero(as_tuple=False).flatten() + num_systems = inputs.num_systems + N = inputs.fill_value + else: + N = coord.shape[0] - 1 + coord_real = coord[:-1].detach().double() + charges_real = data[self.key_in][:-1].detach().double() + mol_idx_real = data["mol_idx"][:-1].to(torch.int32) + nbmat_real = data[f"nbmat{suffix}"][:-1].to(torch.int32) + shifts_real = data[f"shifts{suffix}"][:-1].to(torch.int32) + cell_det = cell.detach().double() + num_systems = int(mol_idx_real.max().item()) + 1 + real_indices = torch.arange(coord_real.shape[0], device=coord.device) is_pme = backend == "pme" def forces_at(positions: Tensor) -> Tensor: @@ -845,7 +894,7 @@ def forces_at(positions: Tensor) -> Tensor: ) return f # (N, 3) eV/Ang (already scaled by k_e in the helper) - return forces_at, coord_real, N + return forces_at, coord_real, real_indices def _coul_nvalchemi_fd_hessian(self, data: dict[str, Tensor], backend: str, *, step: float = 5e-4) -> Tensor: """Central finite-difference Hessian of the FULL periodic Coulomb energy. @@ -883,21 +932,24 @@ def _coul_nvalchemi_fd_hessian(self, data: dict[str, Tensor], backend: str, *, s bound scales with ``ewald_accuracy``: loosening ``ewald_accuracy`` raises the erfc residual at the cutoff and weakens the guarantee. """ - forces_at, coord_real, N = self._periodic_fd_setup(data, backend) + forces_at, coord_flat, real_indices = self._periodic_fd_setup(data, backend) + coord_real = coord_flat.index_select(0, real_indices) + N = coord_real.shape[0] hessian = coord_real.new_zeros((N, 3, N, 3), dtype=torch.float64) with torch.no_grad(): - scratch = coord_real.clone() + scratch = coord_flat.clone() for j in range(N): for b in range(3): - orig = scratch[j, b].item() - scratch[j, b] = orig + step + atom = real_indices[j] + orig = scratch[atom, b].item() + scratch[atom, b] = orig + step fp = forces_at(scratch) - scratch[j, b] = orig - step + scratch[atom, b] = orig - step fm = forces_at(scratch) - scratch[j, b] = orig + scratch[atom, b] = orig # H_{ia,jb} = d^2E/dr_ia dr_jb = -dF_ia/dr_jb dF = (fp - fm) / (2.0 * step) - hessian[:, :, j, b] = (-dF).double() + hessian[:, :, j, b] = (-dF.index_select(0, real_indices)).double() return hessian def _coul_nvalchemi_fd_hvp( @@ -916,13 +968,17 @@ def _coul_nvalchemi_fd_hvp( see :meth:`_coul_nvalchemi_fd_hessian` for the charge-response and step/``ewald_accuracy`` caveats, which apply identically here. """ - forces_at, coord_real, _N = self._periodic_fd_setup(data, backend) + forces_at, coord_flat, real_indices = self._periodic_fd_setup(data, backend) vec_real = vec.detach().double() + if vec_real.shape[0] != real_indices.shape[0]: + raise ValueError("HVP vector must contain one entry per real atom.") + vec_flat = torch.zeros_like(coord_flat) + vec_flat.index_copy_(0, real_indices, vec_real) with torch.no_grad(): - fp = forces_at(coord_real + step * vec_real) - fm = forces_at(coord_real - step * vec_real) + fp = forces_at(coord_flat + step * vec_flat) + fm = forces_at(coord_flat - step * vec_flat) # H_LR @ vec = d^2E/dr dr . vec = -dF/dr . vec (directional) - hv = -(fp - fm) / (2.0 * step) + hv = -(fp.index_select(0, real_indices) - fm.index_select(0, real_indices)) / (2.0 * step) return hv # (N, 3), float64 def forward( @@ -1534,24 +1590,15 @@ def _prepare_dftd3_inputs(self, data: dict[str, Tensor]) -> "_DFTD3KernelInputs" elif nb_mode == 2: suffix = nbops.resolve_suffix(data, ["_dftd3", "_lr"]) - B, N = coord.shape[:2] - coord_flat = coord.flatten(0, 1) - numbers_flat = numbers.flatten() - batch_idx = torch.arange(B, device=coord.device, dtype=torch.int32).repeat_interleave(N) - num_systems = B - - nbmat = data[f"nbmat{suffix}"] - offsets = torch.arange(B, device=coord.device).unsqueeze(1) * N - neighbor_matrix = (nbmat + offsets.unsqueeze(-1)).flatten(0, 1).to(torch.int32) - mask_ij = data.get(f"mask_ij{suffix}") - if mask_ij is not None: - fill_matrix = torch.full_like(neighbor_matrix, B * N) - neighbor_matrix = torch.where(mask_ij.flatten(0, 1), fill_matrix, neighbor_matrix) - - shifts = data.get(f"shifts{suffix}") - neighbor_matrix_shifts = shifts.flatten(0, 1).to(torch.int32) if shifts is not None else None - fill_value = B * N - cell_for_kernel = cell + mode2 = _mode2_backend_inputs(data, suffix) + coord_flat = mode2.coord + numbers_flat = numbers.flatten(0, 1) + batch_idx = mode2.batch_idx + num_systems = mode2.num_systems + neighbor_matrix = mode2.neighbor_matrix + neighbor_matrix_shifts = mode2.shifts.to(torch.int32) if mode2.shifts is not None else None + fill_value = mode2.fill_value + cell_for_kernel = mode2.cell else: raise ValueError(f"Unsupported neighbor mode: {nb_mode}") diff --git a/aimnet/nbops.py b/aimnet/nbops.py index b86e6ae..30f1d81 100644 --- a/aimnet/nbops.py +++ b/aimnet/nbops.py @@ -1,6 +1,289 @@ import torch from torch import Tensor +NBMAT_SUFFIXES = ("", "_lr", "_coulomb", "_dftd3") +_SIGNED_INTEGER_DTYPES = {torch.int8, torch.int16, torch.int32, torch.int64} + + +def _mode2_check(condition: Tensor, message: str) -> None: + """Raise on CPU or queue a device-side assertion on CUDA.""" + if condition.device.type == "cuda": + torch._assert_async(condition, message) + elif not condition.item(): + raise ValueError(message) + + +def normalize_mode2_periodic_geometry(data: dict[str, Tensor], *, B: int) -> dict[str, Tensor]: + """Normalize full-3D periodic mode-2 geometry in place.""" + cell = data.get("cell") + pbc = data.get("pbc") + shift_keys = [f"shifts{suffix}" for suffix in NBMAT_SUFFIXES] + has_shifts = any(data.get(key) is not None for key in shift_keys) + if cell is None: + if pbc is not None: + raise ValueError("pbc requires cell for mode-2 input.") + if has_shifts: + raise ValueError("shifts require cell for mode-2 input.") + return data + + if not isinstance(cell, Tensor): + cell = torch.as_tensor(cell) + if B == 1: + if cell.ndim == 2 and cell.shape == (3, 3): + cell = cell.unsqueeze(0) + elif cell.ndim != 3 or cell.shape != (1, 3, 3): + raise ValueError("cell must have shape (3, 3) or (1, 3, 3) for B=1.") + elif cell.ndim != 3 or cell.shape != (B, 3, 3): + raise ValueError("cell must have shape (B, 3, 3) for batched mode-2 input.") + data["cell"] = cell + + if pbc is None: + pbc = torch.ones((B, 3), dtype=torch.bool, device=cell.device) + else: + pbc = torch.as_tensor(pbc, dtype=torch.bool, device=cell.device) + if pbc.ndim == 1 and pbc.shape == (3,): + pbc = pbc.unsqueeze(0).expand(B, -1) + elif pbc.ndim != 2 or pbc.shape != (B, 3): + raise ValueError("pbc must have shape (3,) or (B, 3).") + _mode2_check(pbc.all(), "mode-2 periodic input requires full-3D pbc.") + data["pbc"] = pbc + return data + + +def validate_neighbor_suffix_layout(data: dict[str, Tensor]) -> None: + """Require every supplied neighbor suffix to match the primary representation.""" + primary = data.get("nbmat") + present = [ + (suffix, data.get(f"nbmat{suffix}"), data.get(f"shifts{suffix}")) + for suffix in NBMAT_SUFFIXES + if data.get(f"nbmat{suffix}") is not None or data.get(f"shifts{suffix}") is not None + ] + if not present: + return + if not isinstance(primary, Tensor): + for suffix, neighbor, _shifts in present: + if neighbor is None: + raise ValueError(f"shifts{suffix} requires matching nbmat{suffix}.") + if neighbor.ndim == 3: + raise ValueError("3D suffixed neighbor matrices require a primary nbmat.") + return + if primary.ndim not in (2, 3): + raise ValueError("nbmat must be 2D or 3D when suffixed neighbor matrices are supplied.") + prefix = primary.shape[:2] if primary.ndim == 3 else primary.shape[:1] + for suffix, neighbor, shifts in present: + key = f"nbmat{suffix}" + if neighbor is None: + raise ValueError(f"{f'shifts{suffix}'} requires matching {key}.") + if not isinstance(neighbor, Tensor) or neighbor.ndim != primary.ndim or neighbor.shape[: len(prefix)] != prefix: + raise ValueError(f"{key} must match nbmat rank and leading shape {prefix}.") + if shifts is not None and not isinstance(shifts, Tensor): + raise ValueError(f"shifts{suffix} must be a tensor.") + + +def _validate_mode2_view_layout(value: Tensor, name: str) -> None: + """Require flattening the batch and atom dimensions to be a view.""" + if value.ndim >= 2 and value.shape[0] and value.shape[1] and value.stride(0) != value.stride(1) * value.shape[1]: + raise ValueError(f"{name} must be flattenable across (B, N) without a copy.") + + +def validate_mode2_nbmat_raw(data: dict[str, Tensor], *, suffix: str) -> None: + """Validate one raw mode-2 neighbor matrix before dtype narrowing.""" + nbmat_key = f"nbmat{suffix}" + shifts_key = f"shifts{suffix}" + nbmat = data.get(nbmat_key) + shifts = data.get(shifts_key) + if nbmat is None: + if shifts is not None: + raise ValueError(f"{shifts_key} requires matching {nbmat_key}.") + raise ValueError(f"{nbmat_key} is required for mode-2 validation.") + if not isinstance(nbmat, Tensor): + raise ValueError(f"{nbmat_key} must be a tensor.") # noqa: TRY004 + if nbmat.ndim != 3: + raise ValueError(f"{nbmat_key} must have shape (B, N, M).") + if nbmat.dtype not in _SIGNED_INTEGER_DTYPES: + raise ValueError(f"{nbmat_key} must use a signed integer dtype.") + + coord = data.get("coord") + numbers = data.get("numbers") + if not isinstance(coord, Tensor) or not isinstance(numbers, Tensor): + raise ValueError("coord and numbers are required for mode-2 validation.") # noqa: TRY004 + B, N, _M = nbmat.shape + if coord.shape[:2] != (B, N) or numbers.shape != (B, N): + raise ValueError(f"{nbmat_key} must match coord and numbers shape prefix (B, N).") + for name, value in (("coord", coord), ("numbers", numbers), (nbmat_key, nbmat)): + _validate_mode2_view_layout(value, name) + if value.device != nbmat.device: + raise ValueError(f"{name} and {nbmat_key} must be on the same device.") + + total_atoms = B * N + if total_atoms > torch.iinfo(torch.int32).max: + raise ValueError(f"{nbmat_key} B*N must fit in int32.") + + cell = data.get("cell") + pbc = data.get("pbc") + if cell is None: + if pbc is not None: + raise ValueError("pbc requires cell for mode-2 input.") + if shifts is not None: + raise ValueError(f"{shifts_key} requires cell for mode-2 input.") + else: + if not isinstance(cell, Tensor) or cell.device != nbmat.device: + raise ValueError("cell and mode-2 neighbor tensors must be on the same device.") + if cell.ndim != 3 or cell.shape != (B, 3, 3): + raise ValueError("cell must be normalized to shape (B, 3, 3).") + if pbc is None: + raise ValueError("pbc must be normalized when cell is present.") + if pbc.device != nbmat.device or pbc.shape != (B, 3): + raise ValueError("pbc must be normalized to shape (B, 3).") + _mode2_check(pbc.all(), "mode-2 periodic input requires full-3D pbc.") + if shifts is None: + raise ValueError(f"{shifts_key} is required when cell is present.") + if shifts.shape != (*nbmat.shape, 3): + raise ValueError(f"{shifts_key} must match {nbmat_key} shape plus a final dimension of 3.") + if shifts.device != nbmat.device: + raise ValueError(f"{shifts_key} and {nbmat_key} must be on the same device.") + _validate_mode2_view_layout(shifts, shifts_key) + if ( + shifts.dtype == torch.bool + or shifts.dtype.is_complex + or not (shifts.dtype in _SIGNED_INTEGER_DTYPES or shifts.dtype.is_floating_point) + ): + raise ValueError(f"{shifts_key} must be an integer or floating-point tensor.") + if shifts.dtype.is_floating_point: + _mode2_check(torch.isfinite(shifts).all(), f"{shifts_key} must be finite.") + _mode2_check((shifts == shifts.round()).all(), f"{shifts_key} must be integral-valued.") + _mode2_check( + ((shifts >= -(2**31)) & (shifts < 2**31)).all(), + f"{shifts_key} values must fit in int32.", + ) + + sentinel = total_atoms + is_sentinel = nbmat == sentinel + _mode2_check( + ((nbmat >= 0) & (nbmat <= sentinel)).all(), + f"{nbmat_key} contains an index outside [0, B*N].", + ) + starts = torch.arange(B, device=nbmat.device, dtype=nbmat.dtype).view(B, 1, 1) * N + in_batch = is_sentinel | ((nbmat >= starts) & (nbmat < starts + N)) + _mode2_check(in_batch.all(), f"{nbmat_key} contains an index outside its batch interval.") + + mask_i = numbers == 0 + _mode2_check(mask_i[..., -1].all(), "numbers must reserve the final atom as the final dummy.") + _mode2_check( + ~(mask_i[..., :-1] & ~mask_i[..., 1:]).any(), + "numbers padding must be a contiguous tail.", + ) + safe_idx = nbmat.clamp(0, sentinel - 1) + padded_neighbor = numbers.flatten().index_select(0, safe_idx.flatten()).view_as(nbmat) == 0 + excluded = is_sentinel | padded_neighbor + _mode2_check( + ~(excluded[..., :-1] & ~excluded[..., 1:]).any(), + f"{nbmat_key} must have a packed sentinel/padded-neighbor tail.", + ) + _mode2_check( + ~(mask_i.unsqueeze(-1) & ~is_sentinel).any(), + f"{nbmat_key} padded center rows must contain only the sentinel.", + ) + if shifts is not None: + center_slots = mask_i.unsqueeze(-1).unsqueeze(-1).expand_as(shifts) + neighbor_slots = excluded.unsqueeze(-1).expand_as(shifts) & ~center_slots + _mode2_check( + (shifts.eq(0) | ~neighbor_slots).all(), + f"{shifts_key} must be zero for sentinel and padded-neighbor slots.", + ) + _mode2_check( + (shifts.eq(0) | ~center_slots).all(), + f"{shifts_key} must be zero for padded center rows.", + ) + + +def _prepare_mode2_neighbor_tensors(data: dict[str, Tensor]) -> None: + """Create mode-2 masks and safe gather indices from validated int32 inputs.""" + nbmat = data["nbmat"] + B, N, _M = nbmat.shape + sentinel = B * N + dedup = not torch.compiler.is_compiling() + previous: list[tuple[Tensor, Tensor, Tensor, Tensor]] = [] + mask_i = data["mask_i"] + for suffix in NBMAT_SUFFIXES: + key = f"nbmat{suffix}" + if key not in data: + continue + current = data[key] + reused = False + if dedup: + for source, mask_ij, gather, kernel in previous: + if ( + current is source + and current.shape == source.shape + and current.stride() == source.stride() + and current.storage_offset() == source.storage_offset() + ): + data[f"mask_ij{suffix}"] = mask_ij + data[f"_nbmat_gather{suffix}"] = gather + data[f"_nbmat_kernel{suffix}"] = kernel + reused = True + break + if reused: + continue + is_sentinel = current == sentinel + safe_idx = torch.where(is_sentinel, torch.zeros_like(current), current) + padded_neighbor = mask_i.flatten().index_select(0, safe_idx.flatten()).view_as(current) + center_pad = mask_i.unsqueeze(-1) + mask_ij = center_pad | is_sentinel | padded_neighbor + gather = safe_idx.masked_fill(mask_ij, 0) + kernel = current.masked_fill(mask_ij, sentinel) + data[f"mask_ij{suffix}"] = mask_ij + data[f"_nbmat_gather{suffix}"] = gather + data[f"_nbmat_kernel{suffix}"] = kernel + if dedup: + previous.append((current, mask_ij, gather, kernel)) + + +def convert_mode2_local_to_global(nbmat_local: Tensor, *, padding_mask: Tensor) -> Tensor: + """Convert a legacy local matrix to packed global int32 indices. + + ``nbmat_local`` and ``padding_mask`` must be ``(B, N, M)`` tensors on the + same device. The matrix must use a signed integer dtype, the mask must be + boolean, exclusions must form a tail in every row, and the final center + row must be fully excluded because ``N - 1`` is the required dummy atom. + Unmasked values are local indices in ``[0, N - 1)``. The returned tensor + is a new contiguous int32 matrix whose masked entries are the global + sentinel ``B * N``. This helper does not reorder neighbors or shifts; + producers owning aligned shifts must repack interleaved exclusions. + + See ``docs/calculator.md#batched-sparse-neighbor-matrices-mode-2`` for the + complete public mode-2 contract and migration guidance. + """ + if nbmat_local.ndim != 3 or padding_mask.shape != nbmat_local.shape: + raise ValueError("nbmat_local and padding_mask must have identical shape (B, N, M).") + if nbmat_local.device != padding_mask.device: + raise ValueError("nbmat_local and padding_mask must be on the same device.") + if nbmat_local.dtype not in _SIGNED_INTEGER_DTYPES: + raise ValueError("nbmat_local must use a signed integer dtype.") + if padding_mask.dtype != torch.bool: + raise ValueError("padding_mask must be boolean.") + B, N, _M = nbmat_local.shape + sentinel = B * N + if sentinel > torch.iinfo(torch.int32).max: + raise ValueError("B*N must fit in int32.") + _mode2_check( + ~((padding_mask[..., :-1]) & ~padding_mask[..., 1:]).any(), + "padding_mask must be tail-packed.", + ) + _mode2_check( + padding_mask[:, -1, :].all(), + "padding_mask must exclude the final dummy center row.", + ) + _mode2_check( + (~padding_mask & ((nbmat_local < 0) | (nbmat_local >= N - 1))).logical_not().all(), + "unmasked local indices are outside the local index range [0, N-1).", + ) + local_int32 = nbmat_local.to(torch.int32) + offsets = torch.arange(B, device=nbmat_local.device, dtype=torch.int32).view(B, 1, 1) * N + result = torch.where(padding_mask, torch.full_like(local_int32, sentinel), local_int32 + offsets) + return result.contiguous() + def set_nb_mode(data: dict[str, Tensor]) -> dict[str, Tensor]: """Logic to guess and set the neighbor model.""" @@ -112,19 +395,7 @@ def calc_masks(data: dict[str, Tensor]) -> dict[str, Tensor]: data["_num_mol"] = torch.tensor(data["mol_sizes"].shape[0]) elif nb_mode == 2: data["mask_i"] = data["numbers"] == 0 - # Same eager-only dedup as mode 1: see the note there on data_ptr(). - dedup_nb2 = not torch.compiler.is_compiling() - processed_nb2: dict[int, str] = {} # data_ptr -> mask_suffix - for suffix in ("", "_lr", "_coulomb", "_dftd3"): - nbmat_key = f"nbmat{suffix}" - if nbmat_key in data: - if dedup_nb2: - ptr = data[nbmat_key].data_ptr() - if ptr in processed_nb2: - data[f"mask_ij{suffix}"] = data[f"mask_ij{processed_nb2[ptr]}"] - continue - processed_nb2[ptr] = suffix - data[f"mask_ij{suffix}"] = _calc_mask_ij_mode2(data[nbmat_key], data["mask_i"]) + _prepare_mode2_neighbor_tensors(data) data["_input_padded"] = torch.tensor(True) data["mol_sizes"] = (~data["mask_i"]).sum(-1) else: @@ -133,28 +404,6 @@ def calc_masks(data: dict[str, Tensor]) -> dict[str, Tensor]: return data -def _calc_mask_ij_mode2(nbmat: Tensor, mask_i: Tensor) -> Tensor: - """Mask padded neighbor entries for batched neighbor matrices. - - Historically mode-2 callers have used both local per-system indices - ``0..N-1`` and flattened global indices ``b*N + i``. Treat both as valid - input conventions for masking so padded atoms are excluded consistently - before downstream code canonicalizes its own neighbor representation. - """ - _, N = mask_i.shape - local_idx = torch.arange(N, device=nbmat.device).view(1, 1, N) - local_pad = (nbmat.unsqueeze(-1) == local_idx) & mask_i.to(device=nbmat.device).unsqueeze(1).unsqueeze(1) - - global_pad_idx = torch.where(mask_i.to(device=nbmat.device).flatten())[0] - if global_pad_idx.numel(): - global_pad = torch.isin(nbmat, global_pad_idx) - else: - global_pad = torch.zeros_like(nbmat, dtype=torch.bool) - - center_pad = mask_i.to(device=nbmat.device).unsqueeze(-1) - return center_pad | local_pad.any(dim=-1) | global_pad - - def mask_ij_( x: Tensor, data: dict[str, Tensor], @@ -270,7 +519,7 @@ def get_ij(x: Tensor, data: dict[str, Tensor], suffix: str = "") -> tuple[Tensor x_j = torch.index_select(x, 0, idx.flatten()).unflatten(0, idx.shape) elif nb_mode == 2: x_i = x.unsqueeze(2) - idx = data[f"nbmat{suffix}"] + idx = data[f"_nbmat_gather{suffix}"] x_j = torch.index_select(x.flatten(0, 1), 0, idx.flatten()).unflatten(0, idx.shape) else: raise ValueError(f"Invalid neighbor mode: {nb_mode}") diff --git a/docs/calculator.md b/docs/calculator.md index c0cdca6..301d0d9 100644 --- a/docs/calculator.md +++ b/docs/calculator.md @@ -49,6 +49,59 @@ result = calc({ }, forces=True, stress=True) ``` +## Batched sparse neighbor matrices (mode 2) + +Mode 2 represents several padded systems in one tensor. Coordinates, atomic numbers, and features have shape `(B, N, ...)`, where `B` is the number of systems and `N` is the padded atom capacity. Each system reserves its final atom slot as a dummy (`numbers[:, -1] == 0`); any additional zero-number atoms form a contiguous tail. + +Neighbor indices are global: atom `j` in system `b` is `b*N + j`. The sentinel is `B*N`, and valid neighbors must precede the sentinel tail in every center row. A neighbor that targets a padded atom is also excluded and belongs to that tail. The primary and long-range matrices use the same representation: + +```python +nbmat.shape == (B, N, M) +nbmat[b, i, k] == b * N + j # valid neighbor +nbmat[b, i, k] == B * N # excluded slot +``` + +Periodic mode 2 uses full three-dimensional cells with shape `(B, 3, 3)` for batches, or `(1, 3, 3)` for one system. `pbc` is optional when a cell is supplied; if present, all three components must be true. Each periodic neighbor matrix must have aligned integral lattice coefficients in `shifts.shape == (B, N, M, 3)`. Energy, forces, stress, and Hessian requests preserve the 3D execution path; Hessians select only real atoms and return `(R, 3, R, 3)` per system, where `R` excludes the padded tail. Batched Hessians stack when all systems have the same `R`, otherwise they return a list. Slab, mixed-periodicity, and partial-PBC inputs are deferred to PR 03. + +All four periodic long-range producers—DSF, DFT-D3, Ewald, and PME—use the same global indices and aligned shifts. Their periodic neighbor lists must represent each physical interaction symmetrically: if an edge uses shift `s`, the reverse edge must use the opposite shift `-s`. A directed or half neighbor list is not a valid input for these long-range observables. A full-observable request can be made directly: + +```python +result = calc( + { + "coord": coord_batch, # (B, N, 3), final atom padded + "numbers": numbers_batch, # (B, N) + "charge": charges, + "cell": cells, # (B, 3, 3) + "nbmat": nbmat_global, # (B, N, M), global indices + "shifts": shifts, # (B, N, M, 3) + "nbmat_lr": nbmat_global, + "shifts_lr": shifts, + "nbmat_coulomb": nbmat_global, + "shifts_coulomb": shifts, + "nbmat_dftd3": nbmat_global, + "shifts_dftd3": shifts, + }, + forces=True, + stress=True, + hessian=True, +) +``` + +Calculator results from an explicit mode-2 Hessian split retain the singleton system axis for each recursive subsystem (for example, energy `(B, 1)` and forces `(B, 1, N, 3)`). This preserves the existing mode-2 collection behavior; the Hessian itself contains only real atoms. + +Legacy local matrices can be converted when their exclusions are already tail-packed: + +```python +from aimnet import nbops + +nbmat_global = nbops.convert_mode2_local_to_global( + nbmat_local, + padding_mask=padding_mask, +) +``` + +The helper does not reorder slots. If exclusions are interleaved, the producer must repack both the neighbor matrix and its aligned shifts together; moving indices alone would assign the wrong periodic image. This is an intentional API break for callers that previously supplied local 3D mode-2 indices. Invalid CUDA content queues a device-side assertion, so restart the CUDA process after a validation failure. + ### Changing Coulomb Methods ```python @@ -381,8 +434,8 @@ For both Ewald and PME, `ewald_accuracy` (default `1e-6`, matching the nvalchemi **Derivative Support:** -- `simple` and `dsf`: inference forces/stress are supported. DSF force/stress losses (`train=True` with `forces` or `stress`) and Hessian calculations raise `NotImplementedError`. -- `ewald` and `pme`: forces, stress, and force/stress losses in `train=True` are supported. Hessian requests raise `NotImplementedError` because nvalchemiops exposes explicit first coordinate derivatives, not second coordinate derivatives. +- `simple` and `dsf`: inference forces/stress are supported. DSF force/stress losses (`train=True` with `forces` or `stress`) remain unsupported; periodic Hessians select real atoms. +- `ewald` and `pme`: forces, stress, force/stress losses in `train=True`, and real-atom Hessians are supported. PME includes its fixed-charge finite-difference long-range block. See [Long-Range Methods → Derivative Support](long_range.md#derivative-support) for the rationale. @@ -512,25 +565,27 @@ H = torch.autograd.functional.hessian(energy_fn, coords) # shape (N, 3, N, 3) !!! note "Long-range backend limitations" - External higher-order differentiation depends on the selected long-range backend. Ewald and PME use nvalchemiops explicit first coordinate derivatives and do not provide complete Coulomb Hessians. + External higher-order differentiation depends on the selected long-range backend. The calculator's explicit Hessian path includes the real-atom periodic Coulomb block where the backend provides it; external `torch.autograd.functional.hessian` remains subject to each backend's graph support. When `coord` does **not** have `requires_grad=True` (the default), inputs are detached as before — optimization loops that call the calculator repeatedly incur no graph accumulation overhead. ## Output Format -| Key | Shape | Description | -| --------- | ----------------------- | ----------------------------- | -| `energy` | `(1,)` or `(B,)` | Total energy per molecule | -| `charges` | `(N,)` or `(B, N)` | Atomic partial charges | -| `forces` | `(N, 3)` or `(B, N, 3)` | Atomic forces (if requested) | -| `stress` | `(3, 3)` or `(B, 3, 3)` | Stress tensor (if requested) | -| `hessian` | `(N, 3, N, 3)` | Hessian matrix (if requested) | +`energy`: `(1,)` or `(B,)`, total energy per molecule. + +`charges`: `(N,)` or `(B, N)`, atomic partial charges. + +`forces`: `(N, 3)` or `(B, N, 3)`, when requested. + +`stress`: `(3, 3)` or `(B, 3, 3)`, when requested. + +`hessian`: `(N, 3, N, 3)` or batched mode-2 equivalents, over real atoms. **Notes:** - `forces` requires `forces=True` in `eval()` - `stress` requires `stress=True` and `cell` in input -- `hessian` requires `hessian=True`, only for single molecules +- `hessian` requires `hessian=True`; explicit batched mode 2 supports stacked or ragged per-system results ## Batching and Neighbor Modes @@ -762,15 +817,7 @@ Legacy JIT models (`.jpt`) have different behavior: ### Common Errors -| Condition | Error | -| ------------------------------------------- | --------------------- | -| Invalid model type | `TypeError` | -| Missing required input key | `KeyError` | -| Hessian with multiple molecules | `NotImplementedError` | -| Hessian with DSF/Ewald/PME Coulomb | `NotImplementedError` | -| PBC with multiple molecules | `NotImplementedError` | -| Invalid Coulomb method | `ValueError` | -| `needs_dispersion=True` without `d3_params` | `ValueError` | +Invalid model type: `TypeError`. Missing required input key: `KeyError`. Unsupported Hessian representation: `ValueError` or `NotImplementedError`. Partial or mixed PBC: `ValueError` (PR 03 scope). Invalid Coulomb method: `ValueError`. Missing D3 parameters: `ValueError`. ### Warnings diff --git a/docs/long_range.md b/docs/long_range.md index c365e3b..2dfb5c2 100644 --- a/docs/long_range.md +++ b/docs/long_range.md @@ -2,6 +2,8 @@ This page documents the long-range (LR) modules implemented in `aimnet/modules/lr.py`. All modules operate on the shared data dictionary and add their contributions to `data[key_out]` (usually `energy`). +For the complete batched sparse neighbor contract, including global indices, required dummy atoms, periodic cells, aligned shifts, and migration guidance, see [Batched sparse neighbor matrices (mode 2)](calculator.md#batched-sparse-neighbor-matrices-mode-2). DSF, DFT-D3, Ewald, and PME accept this representation without stripping per-system dummy rows or remapping valid indices. + ## Choosing a Coulomb Method Select the appropriate method based on your system and accuracy requirements. @@ -78,7 +80,7 @@ calc.set_lrcoulomb_method("dsf", cutoff=15.0, dsf_alpha=0.2) - O(N) scaling with neighbor lists - Energy and forces continuous at cutoff - Based on Wolf summation method -- Inference forces and stress are supported; force/stress training and Hessians are not (see [Derivative Support](#derivative-support)) +- Inference forces, stress, and real-atom Hessians are supported for full-3D periodic mode 2; force/stress training remains unsupported (see [Derivative Support](#derivative-support)). ### Ewald Summation @@ -164,14 +166,14 @@ The nvalchemiops-backed external methods differ in how they expose derivatives: | Backend | Inference forces/stress | Force/stress training | Hessian | | --- | --- | --- | --- | -| DSF | Yes | No | No | -| Ewald | Yes | Yes | No | -| PME | Yes | Yes | No | -| DFT-D3 | Yes | Not applicable; no trainable DFT-D3 parameters | Yes | - -- **DSF**: energy is autograd-connected through charges only. The calculator assembles inference forces and stress by combining PyTorch autograd for the NN and the charge chain with explicit DSF forces/virial. Force/stress losses (`train=True` with `forces=True` or `stress=True`) and Hessian requests raise `NotImplementedError`. -- **Ewald / PME**: support inference forces/stress and force/stress losses in `train=True`. Hessian requests raise `NotImplementedError` because nvalchemiops exposes explicit first coordinate derivatives, not the second coordinate derivatives needed for a complete Coulomb Hessian. -- **DFT-D3**: inference forces and stress come from detached nvalchemiops force/virial terms. Hessian requests use the pure-torch differentiable DFT-D3 path. +| DSF | Yes | No | Yes, real atoms | +| Ewald | Yes | Yes | Yes, real atoms | +| PME | Yes | Yes | Yes, real atoms | +| DFT-D3 | Yes | Not applicable; no trainable DFT-D3 parameters | Yes, real atoms | + +- **DSF**: energy is autograd-connected through charges only. The calculator assembles inference forces and stress by combining PyTorch autograd for the NN and the charge chain with explicit DSF forces/virial. Force/stress losses (`train=True` with `forces=True` or `stress=True`) remain unsupported; Hessian requests use the differentiable periodic path. +- **Ewald / PME**: support inference forces/stress and force/stress losses in `train=True`. Dense Hessians select real atoms; PME adds its fixed-charge finite-difference long-range block. +- **DFT-D3**: inference forces and stress come from detached nvalchemiops force/virial terms. Hessian requests use the pure-torch differentiable DFT-D3 energy path and select real atoms. ## Method Comparison diff --git a/tests/conftest.py b/tests/conftest.py index e20c30f..9b2fa5a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -507,13 +507,21 @@ def add_lr_keys( # Batched mode (B, N, 3) B, N = coord.shape[:2] max_nb = N - 1 - nbmat = torch.zeros((B, N, max_nb), dtype=torch.long, device=device) - for b in range(B): + if B == 1: + nbmat = torch.zeros((N, max_nb), dtype=torch.long, device=device) for i in range(N): neighbors = [j for j in range(N) if j != i] for k, nb in enumerate(neighbors[:max_nb]): - nbmat[b, i, k] = nb - shifts = torch.zeros((B, N, max_nb, 3), dtype=torch.int32, device=device) + nbmat[i, k] = nb + shifts = torch.zeros((N, max_nb, 3), dtype=torch.int32, device=device) + else: + nbmat = torch.zeros((B, N, max_nb), dtype=torch.long, device=device) + for b in range(B): + for i in range(N): + neighbors = [b * N + j for j in range(N) if j != i] + for k, nb in enumerate(neighbors[:max_nb]): + nbmat[b, i, k] = nb + shifts = torch.zeros((B, N, max_nb, 3), dtype=torch.int32, device=device) else: # Flat mode (N, 3) N = coord.shape[0] diff --git a/tests/test_aev.py b/tests/test_aev.py new file mode 100644 index 0000000..129e899 --- /dev/null +++ b/tests/test_aev.py @@ -0,0 +1,146 @@ +"""Mode-2 ConvSV dispatch tests.""" + +import pytest +import torch + +from aimnet import nbops +from aimnet.modules.aev import ConvSV + + +def _mode2_aev_data(device: str = "cpu"): + B, N, M, C, G = 2, 4, 3, 2, 3 + nbmat = torch.tensor( + [ + [[1, 2, 8], [0, 2, 8], [0, 1, 8], [8, 8, 8]], + [[5, 6, 8], [4, 6, 8], [8, 8, 8], [8, 8, 8]], + ], + device=device, + dtype=torch.int64, + ) + a = torch.arange(B * N * C, device=device, dtype=torch.float64).reshape(B, N, C).requires_grad_() + g_sv = torch.randn(B, N, M, G, 4, device=device, dtype=torch.float64) + data = nbops.calc_masks( + nbops.set_nb_mode({ + "nbmat": nbmat, + "numbers": torch.tensor([[6, 1, 1, 0], [8, 1, 0, 0]], device=device), + "g_sv": g_sv, + }) + ) + return data, a + + +def _mode1_aev_data(*, a: torch.Tensor | None = None, g_sv: torch.Tensor | None = None): + n_atoms, n_channels, n_basis, n_neighbors = 4, 2, 3, 3 + nbmat = torch.tensor([[1, 2, 3], [0, 2, 3], [0, 1, 3], [3, 3, 3]], dtype=torch.int32) + if a is None: + a = torch.randn(n_atoms, n_channels, dtype=torch.float64, requires_grad=True) + if g_sv is None: + g_sv = torch.randn(n_atoms, n_neighbors, n_basis, 4, dtype=torch.float64) + g_sv = (g_sv * (nbmat < n_atoms - 1).view(n_atoms, n_neighbors, 1, 1)).requires_grad_() + data = nbops.set_nb_mode({ + "g_sv": g_sv, + "mol_idx": torch.tensor([0, 0, 0, 1]), + "nbmat": nbmat, + }) + return data, a, g_sv + + +def test_mode1_convsv_keeps_dummy_output_zero_and_backpropagates(): + data, a, g_sv = _mode1_aev_data() + out = ConvSV(nshifts_s=3, nchannel=2, d2features=False).double()(data, a) + + assert torch.equal(out[-1], torch.zeros_like(out[-1])) + assert torch.count_nonzero(out[:-1]) > 0 + grad_a, grad_g = torch.autograd.grad(out.square().sum(), (a, g_sv)) + assert torch.isfinite(grad_a).all() + assert torch.isfinite(grad_g).all() + + +def test_mode1_convsv_does_not_remask_fresh_output(monkeypatch): + data, a, _g_sv = _mode1_aev_data() + calls = [] + original_mask_i = nbops.mask_i_ + + def spy(*args, **kwargs): + calls.append((args, kwargs)) + return original_mask_i(*args, **kwargs) + + monkeypatch.setattr(nbops, "mask_i_", spy) + + ConvSV(nshifts_s=3, nchannel=2, d2features=False).double()(data, a) + + assert calls == [] + + +@pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile requires PyTorch 2.0+") +def test_mode1_convsv_compile_matches_eager_forward_and_backward(): + conv = ConvSV(nshifts_s=3, nchannel=2, d2features=False).double() + _data, a, g_sv = _mode1_aev_data() + a_eager = a.detach().clone().requires_grad_() + g_eager = g_sv.detach().clone().requires_grad_() + eager_out = conv(_mode1_aev_data(a=a_eager, g_sv=g_eager)[0], a_eager) + eager_loss = eager_out.square().sum() + eager_grads = torch.autograd.grad(eager_loss, (a_eager, g_eager)) + + a_compiled = a.detach().clone().requires_grad_() + g_compiled = g_sv.detach().clone().requires_grad_() + compiled = torch.compile(conv, backend="aot_eager") + compiled_out = compiled(_mode1_aev_data(a=a_compiled, g_sv=g_compiled)[0], a_compiled) + compiled_loss = compiled_out.square().sum() + compiled_grads = torch.autograd.grad(compiled_loss, (a_compiled, g_compiled)) + + torch.testing.assert_close(compiled_out, eager_out) + torch.testing.assert_close(compiled_grads[0], eager_grads[0]) + torch.testing.assert_close(compiled_grads[1], eager_grads[1]) + + +def test_global_mode2_convsv_d2false_forward(): + data, a = _mode2_aev_data() + out = ConvSV(nshifts_s=3, nchannel=2, d2features=False).double()(data, a) + assert out.shape[:2] == (2, 4) + + +def test_global_mode2_convsv_d2false_first_gradient(): + data, a = _mode2_aev_data() + out = ConvSV(nshifts_s=3, nchannel=2, d2features=False).double()(data, a) + assert torch.autograd.grad(out.sum(), a)[0].shape == a.shape + + +def test_global_mode2_convsv_d2false_hessian_vmap(): + data, a = _mode2_aev_data() + conv = ConvSV(nshifts_s=3, nchannel=2, d2features=False).double() + grad = torch.autograd.grad(conv(data, a).sum(), a, create_graph=True)[0] + assert torch.autograd.grad(grad.sum(), a)[0].shape == a.shape + + +def test_global_mode2_convsv_d2false_padded_center(): + data, a = _mode2_aev_data() + out = ConvSV(nshifts_s=3, nchannel=2, d2features=False).double()(data, a) + assert torch.equal(out[0, 3], torch.zeros_like(out[0, 3])) + assert torch.equal(out[1, 2], torch.zeros_like(out[1, 2])) + + +def test_global_mode2_convsv_d2false_cross_batch_gather(): + data, a = _mode2_aev_data() + data["g_sv"].zero_() + data["g_sv"][1, 0, 0, 0, 0] = 1 + out = ConvSV(nshifts_s=3, nchannel=2, d2features=False).double()(data, a) + assert torch.equal(out[1, 0, 0], a[1, 1, 0]) + assert torch.equal(out[1, 0, 1:3], torch.zeros(2, dtype=out.dtype)) + + +def test_global_mode2_convsv_rejects_noncontiguous_input(): + data, a = _mode2_aev_data() + a = a.transpose(1, 2) + with pytest.raises(ValueError, match="flatten"): + ConvSV(nshifts_s=3, nchannel=2, d2features=False).double()(data, a) + + +@pytest.mark.gpu +def test_global_mode2_convsv_cuda_float64_fallback(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + data, a = _mode2_aev_data("cuda") + a = a.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous().requires_grad_() + out = ConvSV(nshifts_s=3, nchannel=2, d2features=True).cuda().double()(data, a.cuda()) + assert out.device.type == "cuda" diff --git a/tests/test_calculator.py b/tests/test_calculator.py index 3d25e65..32ce6b2 100644 --- a/tests/test_calculator.py +++ b/tests/test_calculator.py @@ -1040,8 +1040,8 @@ def test_batch_vs_individual_mode1(self): np.testing.assert_allclose(res_batch["energy"][0].item(), e1, atol=1e-5) np.testing.assert_allclose(res_batch["energy"][1].item(), e2, atol=1e-5) - def test_batch_vs_individual_mode2(self): - """nb_mode=2: Batched sparse neighbor matrix format.""" + def test_batch_vs_individual_mode1_small_molecules(self): + """Automatic small-molecule batching remains independent of sparse mode 2.""" calc = AIMNet2Calculator("aimnet2", nb_threshold=0) # Use same-size molecules for mode 2 (requires padding otherwise) @@ -2007,3 +2007,213 @@ def test_deterministic_warns_for_ewald(self, water_molecule): cell = torch.eye(3) * 20.0 with pytest.warns(UserWarning, match="Ewald/PME"): calc({**water_molecule, "cell": cell}, forces=True) + + +def _global_mode2_calculator_input(*, batch: int = 2, periodic: bool = False) -> dict[str, torch.Tensor]: + n_atoms = 4 + sentinel = batch * n_atoms + coord = torch.arange(batch * n_atoms * 3, dtype=torch.float32).reshape(batch, n_atoms, 3) + numbers = torch.tensor([[6, 1, 1, 0], [8, 1, 0, 0]], dtype=torch.int64)[:batch] + nbmat = torch.full((batch, n_atoms, 3), sentinel, dtype=torch.int64) + for b in range(batch): + base = b * n_atoms + nbmat[b, 0, :2] = torch.tensor([base + 1, base + 2]) + nbmat[b, 1, :2] = torch.tensor([base, base + 2]) + if numbers[b, 2] != 0: + nbmat[b, 2, :2] = torch.tensor([base, base + 1]) + data: dict[str, torch.Tensor] = { + "coord": coord, + "numbers": numbers, + "charge": torch.zeros(batch), + "nbmat": nbmat, + } + if periodic: + data["cell"] = torch.eye(3).repeat(batch, 1, 1) * 10 + data["pbc"] = torch.ones((batch, 3), dtype=torch.bool) + data["shifts"] = torch.zeros((*nbmat.shape, 3)) + return data + + +def _new_mode2_calculator() -> AIMNet2Calculator: + return AIMNet2Calculator("aimnet2", nb_threshold=0, device="cpu") + + +def test_global_mode2_calculator_preserves_cpu(): + calc = _new_mode2_calculator() + prepared = calc.prepare_input(_global_mode2_calculator_input()) + assert prepared["coord"].ndim == 3 + assert prepared["numbers"].ndim == 2 + assert prepared["nbmat"].ndim == 3 + + +def test_global_mode2_calculator_cpu_batch_vs_individual(): + calc = _new_mode2_calculator() + source = _global_mode2_calculator_input() + for suffix in ("_lr", "_coulomb", "_dftd3"): + source[f"nbmat{suffix}"] = source["nbmat"] + batched = calc(source, forces=True) + + energies = [] + forces = [] + B, N = source["coord"].shape[:2] + for b in range(B): + single_nbmat = torch.where( + source["nbmat"][b : b + 1] == B * N, + torch.tensor(N), + source["nbmat"][b : b + 1] - b * N, + ) + single = { + "coord": source["coord"][b : b + 1], + "numbers": source["numbers"][b : b + 1], + "charge": torch.zeros(1), + "nbmat": single_nbmat, + } + for suffix in ("_lr", "_coulomb", "_dftd3"): + single[f"nbmat{suffix}"] = single_nbmat + result = calc(single, forces=True) + energies.append(result["energy"]) + forces.append(result["forces"]) + torch.testing.assert_close(batched["energy"], torch.cat(energies), atol=1e-5, rtol=1e-5) + torch.testing.assert_close(batched["forces"], torch.cat(forces), atol=1e-5, rtol=1e-4) + + +def test_global_mode2_calculator_rejects_source_dtype(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input() + data["nbmat"] = data["nbmat"].float() + with pytest.raises(ValueError, match="integer"): + calc.prepare_input(data) + + +def test_global_mode2_calculator_preserves_alias_int32(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input() + data["nbmat"] = data["nbmat"].to(torch.int32) + data["nbmat_lr"] = data["nbmat"] + prepared = calc.to_input_tensors(data) + assert prepared["nbmat"] is prepared["nbmat_lr"] + + +def test_global_mode2_calculator_preserves_alias_int64(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input() + data["nbmat_lr"] = data["nbmat"] + prepared = calc.to_input_tensors(data) + assert prepared["nbmat"] is prepared["nbmat_lr"] + + +def test_global_mode2_calculator_preserves_full3d_pbc(): + calc = _new_mode2_calculator() + prepared = calc.prepare_input(_global_mode2_calculator_input(periodic=True)) + assert prepared["coord"].shape == (2, 4, 3) + assert prepared["cell"].shape == (2, 3, 3) + assert prepared["pbc"].shape == (2, 3) + + +def test_global_mode2_calculator_canonicalizes_single_cell(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input(batch=1, periodic=True) + data["cell"] = torch.eye(3) * 10 + data["pbc"] = torch.ones(3, dtype=torch.bool) + prepared = calc.prepare_input(data) + assert prepared["cell"].shape == (1, 3, 3) + assert prepared["pbc"].shape == (1, 3) + + +def test_global_mode2_calculator_rejects_shared_cell_for_batch(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input(periodic=True) + data["cell"] = torch.eye(3) * 10 + with pytest.raises(ValueError, match="B, 3, 3"): + calc.prepare_input(data) + + +def test_global_mode2_calculator_rejects_pbc_without_cell(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input() + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + with pytest.raises(ValueError, match="cell"): + calc.prepare_input(data) + + +def test_global_mode2_calculator_rejects_partial_pbc(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input(periodic=True) + data["pbc"][0, 1] = False + with pytest.raises(ValueError, match="full-3D"): + calc.prepare_input(data) + + +def test_global_mode2_calculator_stress_keeps_3d(): + calc = _new_mode2_calculator() + prepared = calc.prepare_input(_global_mode2_calculator_input(periodic=True)) + strained = calc.set_grad_tensors(prepared, stress=True) + assert strained["coord"].ndim == 3 + assert strained["cell"].shape == (2, 3, 3) + + +def test_global_mode2_calculator_hessian_splits_singleton_mode2(): + calc = _new_mode2_calculator() + subsystems = calc._split_hessian_batch(_global_mode2_calculator_input()) + assert subsystems is not None + assert all(sub["nbmat"].shape == (1, 4, 3) for sub in subsystems) + assert subsystems[1]["nbmat"][0, 0, 0] == 1 + + +def test_global_mode2_calculator_hessian_removes_all_padding(): + from aimnet.calculators.derivatives import calculate_hessian + + coord = torch.zeros((4, 3), requires_grad=True) + energy = (coord[:2] ** 2).sum() + forces = -torch.autograd.grad(energy, coord, create_graph=True)[0][:2] + hessian = calculate_hessian(forces, coord, real_atom_mask=torch.tensor([True, True, False, False])) + assert hessian.shape == (2, 3, 2, 3) + + +def test_global_mode2_calculator_hessian_preserves_singleton_batch_graph(): + from aimnet.calculators.derivatives import calculate_hessian + + coord = torch.zeros((1, 3, 3), requires_grad=True) + energy = (coord[:, :2] ** 2).sum() + forces = -torch.autograd.grad(energy, coord, create_graph=True)[0] + hessian = calculate_hessian( + forces, + coord, + real_atom_mask=torch.tensor([[True, True, False]]), + ) + assert hessian.shape == (2, 3, 2, 3) + torch.testing.assert_close(hessian.reshape(6, 6), 2 * torch.eye(6)) + + +def test_global_mode2_calculator_hessian_returns_ragged_list(): + calc = _new_mode2_calculator() + calc.eval = lambda *args, **kwargs: {"hessian": torch.zeros((1, 3, 1, 3))} # type: ignore[method-assign] + result = calc._eval_hessian_batched([{}, {}], forces=False, stress=False, validate_species=False, stack=False) + assert isinstance(result["hessian"], list) + + +def test_global_mode2_calculator_stress_and_hessian(): + calc = _new_mode2_calculator() + prepared = calc.prepare_input(_global_mode2_calculator_input(batch=1, periodic=True)) + strained = calc.set_grad_tensors(prepared, stress=True, hessian=True) + assert strained["coord"].ndim == 3 + assert calc._saved_for_grad["coord"].ndim == 3 + + +def test_global_mode2_calculator_rejects_cross_batch_before_hessian_split(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input() + data["nbmat"][0, 0, 0] = 4 + calc._eval_hessian_batched = lambda *args, **kwargs: (_ for _ in ()).throw( # type: ignore[method-assign] + AssertionError("hessian split reached before validation") + ) + with pytest.raises(ValueError, match="batch interval"): + calc.eval(data, hessian=True, validate_species=False) + + +def test_global_mode2_calculator_slices_batched_pbc(): + calc = _new_mode2_calculator() + data = _global_mode2_calculator_input(periodic=True) + subsystems = calc._split_batch_dim(data, 2) + assert subsystems[0]["pbc"].shape == (1, 3) + assert subsystems[1]["pbc"].shape == (1, 3) diff --git a/tests/test_calculator_gpu.py b/tests/test_calculator_gpu.py index dc0ac08..fb46114 100644 --- a/tests/test_calculator_gpu.py +++ b/tests/test_calculator_gpu.py @@ -9,6 +9,7 @@ from conftest import CAFFEINE_FILE, load_mol from aimnet.calculators import AIMNet2Calculator +from aimnet.modules.lr import _mode2_backend_inputs # Skip entire module if CUDA is not available pytestmark = pytest.mark.gpu @@ -69,6 +70,47 @@ def _neighbor_list_calls(calc): ) +def _cuda_global_mode2_data(*, batch: int = 2, dtype: torch.dtype = torch.float32, periodic: bool = False): + B, N, M = batch, 4, 4 + coord = torch.zeros((B, N, 3), device="cuda", dtype=dtype) + coord[:, 1, 0] = 1.0 + coord[:, 2, 1] = 1.1 + if B > 1: + coord[1, :3] += 2.0 + numbers = torch.tensor([[6, 1, 1, 0]] * B, device="cuda") + nbmat = torch.full((B, N, M), B * N, device="cuda", dtype=torch.int64) + for b in range(B): + for i in range(N - 1): + targets = [b * N + j for j in range(N - 1) if j != i] + nbmat[b, i, : len(targets)] = torch.tensor(targets, device="cuda") + data = { + "coord": coord, + "numbers": numbers, + "charge": torch.zeros(B, device="cuda", dtype=dtype), + "nbmat": nbmat, + "nbmat_lr": nbmat, + "nbmat_coulomb": nbmat, + "nbmat_dftd3": nbmat, + } + if periodic: + shifts = torch.zeros((*nbmat.shape, 3), device="cuda", dtype=dtype) + shifts[:, 0, 0, 0] = 1 + shifts[:, 1, 0, 0] = -1 + data.update({ + "cell": torch.eye(3, device="cuda", dtype=dtype).expand(B, -1, -1) * 12, + "pbc": torch.ones((B, 3), device="cuda", dtype=torch.bool), + "shifts": shifts, + "shifts_lr": shifts, + "shifts_coulomb": shifts, + "shifts_dftd3": shifts, + }) + return data + + +def _cuda_calc(): + return AIMNet2Calculator("aimnet2", device="cuda", nb_threshold=0) + + class TestGPUBasics: """Basic GPU functionality tests.""" @@ -99,6 +141,102 @@ def test_forces_on_cuda(self): assert res["forces"].device.type == "cuda" +def test_global_mode2_cuda_float32_parity(): + calc = _cuda_calc() + data = _cuda_global_mode2_data() + batched = calc(data, forces=True) + for b in range(2): + single = _cuda_global_mode2_data(batch=1) + single["coord"] = data["coord"][b : b + 1].clone() + single["numbers"] = data["numbers"][b : b + 1].clone() + single["charge"] = data["charge"][b : b + 1].clone() + single_nbmat = torch.where( + data["nbmat"][b : b + 1] == 8, torch.tensor(4, device="cuda"), data["nbmat"][b : b + 1] - b * 4 + ) + for key in ("nbmat", "nbmat_lr", "nbmat_coulomb", "nbmat_dftd3"): + single[key] = single_nbmat + result = calc(single, forces=True) + torch.testing.assert_close(batched["energy"][b], result["energy"][0], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(batched["forces"][b], result["forces"][0], atol=1e-5, rtol=1e-4) + + +def test_global_mode2_cuda_float64_fallback(): + result = _cuda_calc()(_cuda_global_mode2_data(dtype=torch.float64)) + assert result["energy"].dtype == torch.float64 + + +def test_global_mode2_cuda_force_parity(): + calc = _cuda_calc() + data = _cuda_global_mode2_data() + result = calc(data, forces=True) + assert result["forces"].shape == data["coord"].shape + assert torch.isfinite(result["forces"]).all() + + +def test_global_mode2_cuda_full3d_periodic_all_observables(): + calc = _cuda_calc() + data = _cuda_global_mode2_data(periodic=True) + result = calc(data, forces=True, stress=True, hessian=True) + assert result["energy"].shape == (2, 1) + assert result["forces"].shape == (2, 1, 4, 3) + assert result["stress"].shape == (2, 1, 3, 3) + assert result["hessian"].ndim == 5 + for batch_index in range(2): + single = _cuda_global_mode2_data(batch=1, periodic=True) + single["coord"] = data["coord"][batch_index : batch_index + 1].clone() + single["numbers"] = data["numbers"][batch_index : batch_index + 1].clone() + single["charge"] = data["charge"][batch_index : batch_index + 1].clone() + single["cell"] = data["cell"][batch_index : batch_index + 1].clone() + single["pbc"] = data["pbc"][batch_index : batch_index + 1].clone() + single_nbmat = torch.where( + data["nbmat"][batch_index : batch_index + 1] == 8, + torch.tensor(4, device="cuda"), + data["nbmat"][batch_index : batch_index + 1] - batch_index * 4, + ) + single_shifts = data["shifts"][batch_index : batch_index + 1].clone() + for key in ("nbmat", "nbmat_lr", "nbmat_coulomb", "nbmat_dftd3"): + single[key] = single_nbmat + for key in ("shifts", "shifts_lr", "shifts_coulomb", "shifts_dftd3"): + single[key] = single_shifts + independent = calc(single, forces=True, stress=True, hessian=True) + for key in ("energy", "forces", "stress", "hessian"): + independent_value = independent[key] if key == "hessian" else independent[key][0] + torch.testing.assert_close( + result[key][batch_index].reshape(-1), + independent_value.reshape(-1), + atol=2e-4, + rtol=2e-4, + ) + + +def test_global_mode2_cuda_periodic_single_cell_geometry(): + data = _cuda_global_mode2_data(batch=1, periodic=True) + data["cell"] = torch.eye(3, device="cuda") * 12 + data["pbc"] = torch.ones(3, device="cuda", dtype=torch.bool) + result = _cuda_calc()(data) + assert result["energy"].shape == (1,) + + +def test_global_mode2_cuda_cross_batch_isolation(): + calc = _cuda_calc() + data = _cuda_global_mode2_data() + baseline = calc(data)["energy"].detach() + mutated = {key: value.clone() if isinstance(value, torch.Tensor) else value for key, value in data.items()} + mutated["coord"][0, 1, 0] += 0.5 + changed = calc(mutated)["energy"].detach() + torch.testing.assert_close(changed[1], baseline[1], atol=1e-5, rtol=1e-5) + assert not torch.allclose(changed[0], baseline[0], atol=1e-5, rtol=1e-5) + + +def test_global_mode2_cuda_zero_copy_reshape(): + data = _cuda_global_mode2_data(periodic=True) + prepared = _cuda_calc().prepare_input(data) + inputs = _mode2_backend_inputs(prepared, "_lr") + assert inputs.coord._base is not None + assert inputs.neighbor_matrix._base is not None + assert inputs.shifts is not None and inputs.shifts._base is not None + + class TestGPUvsCPUConsistency: """Tests verifying GPU and CPU produce consistent results.""" diff --git a/tests/test_compile_paths.py b/tests/test_compile_paths.py index ab2cff4..cc7d1d5 100644 --- a/tests/test_compile_paths.py +++ b/tests/test_compile_paths.py @@ -70,6 +70,40 @@ def _packed_data(n_mol, n_atom_per_mol, device, nfeat=2): return nbops.calc_masks(nbops.set_nb_mode(data)) +def _mode2_data(device): + B, N, M = 2, 4, 3 + sentinel = B * N + nbmat = torch.full((B, N, M), sentinel, dtype=torch.int32, device=device) + for batch in range(B): + offset = batch * N + nbmat[batch, 0, :2] = torch.tensor([offset + 1, offset + 2], device=device) + nbmat[batch, 1, :2] = torch.tensor([offset, offset + 2], device=device) + nbmat[batch, 2, :2] = torch.tensor([offset, offset + 1], device=device) + return { + "numbers": torch.tensor([[6, 1, 1, 0], [8, 1, 1, 0]], device=device), + "nbmat": nbmat, + "nbmat_lr": nbmat, + } + + +def test_mode2_calc_masks_compiled_matches_eager(device): + """Compiled mode 2 avoids alias-identity dedup without changing masks.""" + if device.type != "cuda": + pytest.skip("compiled parity is only meaningful on the GPU backend") + + def derive(data): + prepared = nbops.calc_masks(nbops.set_nb_mode(data)) + return prepared["mask_ij"], prepared["_nbmat_gather"], prepared["_nbmat_kernel"] + + eager = derive(_mode2_data(device)) + torch._dynamo.reset() + compiled = torch.compile(derive, fullgraph=True) + actual = compiled(_mode2_data(device)) + + for got, expected in zip(actual, eager, strict=True): + torch.testing.assert_close(got, expected) + + @pytest.mark.parametrize("n_mol", [1, 2, 5]) def test_mol_sum_compiled_matches_eager(device, n_mol): """The compiled branch reads the count from `charge`; the eager one from diff --git a/tests/test_conv_sv_2d_sp.py b/tests/test_conv_sv_2d_sp.py index 4679718..d1bdb61 100644 --- a/tests/test_conv_sv_2d_sp.py +++ b/tests/test_conv_sv_2d_sp.py @@ -284,7 +284,7 @@ def test_op_registration(self, test_data_cuda): assert hasattr(torch.ops.aimnet, "conv_sv_2d_sp_fwd"), "conv_sv_2d_sp_fwd op not registered" # Test using ops directly - output1 = torch.ops.aimnet.conv_sv_2d_sp_fwd(a.detach(), idx, g.detach()) + output1 = torch.ops.aimnet.conv_sv_2d_sp_fwd(a.detach(), idx, g.detach(), a.shape[0] - 1, a.shape[0] - 1) output2 = reference_conv_sv_2d_sp_einsum(a.detach(), idx, g.detach()) assert torch.allclose(output1, output2, atol=1e-5, rtol=1e-4), ( @@ -340,7 +340,7 @@ def test_vmap_bwd_bwd_kernel_rule(self, test_data_small_cuda): g_in = g.detach() def call_kernel(g2a, g2g): - return torch.ops.aimnet.conv_sv_2d_sp_bwd_bwd(grad_output, g2a, g2g, a_in, idx, g_in) + return torch.ops.aimnet.conv_sv_2d_sp_bwd_bwd(grad_output, g2a, g2g, a_in, idx, g_in, B - 1, B - 1) batched = torch.vmap( call_kernel, @@ -377,7 +377,7 @@ def test_vmap_bwd_kernel_rule(self, test_data_small_cuda): g_in = g.detach() def call_kernel(go): - return torch.ops.aimnet.conv_sv_2d_sp_bwd(go, a_in, idx, g_in) + return torch.ops.aimnet.conv_sv_2d_sp_bwd(go, a_in, idx, g_in, B - 1, B - 1) batched = torch.vmap(call_kernel, in_dims=(0,))(grad_output) @@ -470,3 +470,125 @@ def test_forward_float64_cuda_falls_back_to_einsum(self, convsv_data): assert torch.allclose(out_cuda.cpu(), out_cpu, atol=1e-10, rtol=1e-10), ( f"CUDA float64 einsum fallback parity failed. Max diff: {torch.max(torch.abs(out_cuda.cpu() - out_cpu))}" ) + + +def _generalized_conv_data(): + device = "cuda" + B, A, G, M = 4, 3, 2, 5 + a = torch.randn(B, A, G, device=device, dtype=torch.float32, requires_grad=True) + idx = torch.tensor([[1, 2, 4, 4, 4], [0, 2, 3, 4, 4], [0, 1, 3, 4, 4], [0, 1, 2, 4, 4]], device=device) + g = torch.randn(B, M, G, 4, device=device, dtype=torch.float32, requires_grad=True) + g = g * (idx < B).unsqueeze(-1).unsqueeze(-1) + return a, idx, g, B, B + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_forward_accuracy(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, idx, g, padding_value, num_centers = _generalized_conv_data() + actual = conv_sv_2d_sp(a, idx, g, padding_value, num_centers) + selected = a.index_select(0, idx.clamp_max(a.shape[0] - 1).flatten()).view(idx.shape[0], idx.shape[1], *a.shape[1:]) + expected = torch.einsum("bmag,bmgd->bagd", selected, g) + assert torch.allclose(actual, expected, atol=1e-5, rtol=1e-4) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_distinct_atom_and_center_capacities(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a = torch.randn(6, 2, 2, device="cuda") + idx = torch.tensor([[1, 2, 5], [0, 2, 5], [0, 1, 5], [0, 1, 2]], device="cuda", dtype=torch.int32) + g = torch.randn(4, 3, 2, 4, device="cuda") + out = conv_sv_2d_sp(a, idx, g, padding_value=6, num_centers=4) + assert out.shape == (4, 2, 2, 4) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_backward_accuracy(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, idx, g, padding_value, num_centers = _generalized_conv_data() + out = conv_sv_2d_sp(a, idx, g, padding_value, num_centers) + cotangent = torch.randn_like(out) + grad_a, grad_g = torch.autograd.grad(out, (a, g), cotangent, create_graph=True) + + a_ref = a.detach().requires_grad_(True) + g_ref = g.detach().requires_grad_(True) + selected = a_ref.index_select(0, idx.clamp_max(a_ref.shape[0] - 1).flatten()).view( + idx.shape[0], idx.shape[1], *a_ref.shape[1:] + ) + valid = (idx < padding_value).unsqueeze(-1).unsqueeze(-1) + reference = torch.einsum("cmag,cmgd->cagd", selected, g_ref * valid) + ref_a, ref_g = torch.autograd.grad(reference, (a_ref, g_ref), cotangent) + torch.testing.assert_close(grad_a, ref_a, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(grad_g, ref_g, atol=1e-5, rtol=1e-4) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_double_backward_accuracy(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, idx, g, padding_value, num_centers = _generalized_conv_data() + out = conv_sv_2d_sp(a, idx, g, padding_value, num_centers) + grad_a = torch.autograd.grad(out.sum(), a, create_graph=True)[0] + second = torch.autograd.grad(grad_a.sum(), g)[0] + + a_ref = a.detach().requires_grad_(True) + g_ref = g.detach().requires_grad_(True) + selected = a_ref.index_select(0, idx.clamp_max(a_ref.shape[0] - 1).flatten()).view( + idx.shape[0], idx.shape[1], *a_ref.shape[1:] + ) + valid = (idx < padding_value).unsqueeze(-1).unsqueeze(-1) + reference = torch.einsum("cmag,cmgd->cagd", selected, g_ref * valid) + ref_grad_a = torch.autograd.grad(reference.sum(), a_ref, create_graph=True)[0] + ref_second = torch.autograd.grad(ref_grad_a.sum(), g_ref)[0] + torch.testing.assert_close(second, ref_second, atol=1e-5, rtol=1e-4) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_mode1_explicit_compatibility(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, idx, g, padding_value, num_centers = _generalized_conv_data() + assert torch.equal(conv_sv_2d_sp(a, idx, g), conv_sv_2d_sp(a, idx, g, padding_value - 1, num_centers - 1)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_backward_vmap_scalars(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, idx, g, padding_value, num_centers = _generalized_conv_data() + output = conv_sv_2d_sp(a, idx, g, padding_value, num_centers) + cotangents = torch.randn((2, *output.shape), device="cuda") + values = torch.vmap(lambda cotangent: torch.autograd.grad(output, a, cotangent, retain_graph=True)[0])(cotangents) + assert values.shape == cotangents.shape[:-1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_double_backward_vmap_scalars(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, idx, g, padding_value, num_centers = _generalized_conv_data() + output = conv_sv_2d_sp(a, idx, g, padding_value, num_centers) + values = torch.func.vmap(lambda cotangent: cotangent.sum())(output) + assert values.shape == (output.shape[0],) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_vmap_rejects_unsupported_dims(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, idx, g, padding_value, num_centers = _generalized_conv_data() + with pytest.raises((RuntimeError, ValueError)): + torch.func.vmap(lambda value: conv_sv_2d_sp(value, idx, g, padding_value, num_centers))(a.unsqueeze(0)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_global_mode2_empty_and_all_sentinel_rows(): + from aimnet.kernels.conv_sv_2d_sp_wp import conv_sv_2d_sp + + a, _idx, g, padding_value, num_centers = _generalized_conv_data() + idx = torch.full((a.shape[0], g.shape[1]), padding_value, device="cuda", dtype=torch.int64) + actual = conv_sv_2d_sp(a, idx, g, padding_value, num_centers) + assert torch.equal(actual, torch.zeros_like(actual)) diff --git a/tests/test_dftd3.py b/tests/test_dftd3.py index 051a519..898c575 100644 --- a/tests/test_dftd3.py +++ b/tests/test_dftd3.py @@ -215,8 +215,8 @@ def test_mode_1_uses_trailing_padding_atom_as_fill_value(self, device): torch.full_like(kernel_inputs.neighbor_matrix[data["mask_ij_dftd3"]], kernel_inputs.fill_value), ) - def test_mode_2_masks_padded_neighbors_after_global_offset(self, device): - """Batched sparse DFTD3 inputs convert local padding indices to the global fill value.""" + def test_mode_2_preserves_global_sentinel_and_padded_neighbors(self, device): + """Batched sparse DFTD3 inputs preserve global indices and sentinel tails.""" module = DFTD3(s8=0.3908, a1=0.5660, a2=3.1280).to(device) coord = torch.tensor( [ @@ -229,10 +229,11 @@ def test_mode_2_masks_padded_neighbors_after_global_offset(self, device): numbers = torch.tensor([[8, 1, 0], [6, 1, 0]], device=device) nbmat = torch.tensor( [ - [[1, 2], [0, 2], [2, 2]], - [[1, 2], [0, 2], [2, 2]], + [[1, 6], [0, 6], [6, 6]], + [[4, 6], [3, 6], [6, 6]], ], device=device, + dtype=torch.int32, ) data = { "coord": coord, diff --git a/tests/test_hvp.py b/tests/test_hvp.py index c763bfa..4866671 100644 --- a/tests/test_hvp.py +++ b/tests/test_hvp.py @@ -28,6 +28,29 @@ def _dense_hessian_matmul(calc, data, v): return (Hmat @ v.reshape(-1).double()).reshape(n, 3) +def _singleton_mode2_input() -> dict[str, torch.Tensor]: + sentinel = 5 + nbmat = torch.tensor([ + [ + [1, 2, sentinel], + [0, 2, sentinel], + [0, 1, sentinel], + [sentinel, sentinel, sentinel], + [sentinel, sentinel, sentinel], + ] + ]) + return { + "coord": torch.tensor([ + [[0.0, 0.0, 0.0], [0.96, 0.0, 0.0], [-0.24, 0.93, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + ]), + "numbers": torch.tensor([[8, 1, 1, 0, 0]]), + "charge": torch.tensor([0.0]), + "nbmat": nbmat, + "nbmat_lr": nbmat, + "nbmat_coulomb": nbmat, + } + + def _nblist_state(nblist): if nblist is None: return None @@ -109,6 +132,20 @@ def test_hvp_multiple_vectors_shape(): torch.testing.assert_close(HV.double().cpu(), ref.cpu(), rtol=1e-3, atol=1e-3) +@pytest.mark.slow +def test_hvp_matches_dense_singleton_mode2_with_multiple_padding_rows(): + calc = AIMNet2Calculator("aimnet2", nb_threshold=0, device="cpu") + calc.external_dftd3 = None + data = _singleton_mode2_input() + vector = torch.randn(3, 3, dtype=torch.float64) + + actual = calc.hessian_vector_product(data, vector) + expected = _dense_hessian_matmul(calc, data, vector) + + assert actual.shape == (3, 3) + torch.testing.assert_close(actual.double().cpu(), expected.cpu(), rtol=1e-3, atol=1e-3) + + def test_hvp_batched_input_raises(): calc = AIMNet2Calculator("aimnet2", nb_threshold=1000) data = { diff --git a/tests/test_kernel_registration.py b/tests/test_kernel_registration.py index b323641..3718b7a 100644 --- a/tests/test_kernel_registration.py +++ b/tests/test_kernel_registration.py @@ -6,6 +6,7 @@ drift in torch.library schema inference. """ +import pytest import torch EXPECTED_OPS = { @@ -34,3 +35,61 @@ def test_ops_namespace_present(): assert hasattr(torch.ops, "aimnet") for name in ("conv_sv_2d_sp_fwd", "conv_sv_2d_sp_bwd", "conv_sv_2d_sp_bwd_bwd"): assert hasattr(torch.ops.aimnet, name), f"missing op {name}" + + +def test_conv_sv_generalized_schemas(): + import aimnet.kernels.conv_sv_2d_sp_wp # noqa: F401 + + fwd = str(torch.ops.aimnet.conv_sv_2d_sp_fwd.default._schema) + bwd = str(torch.ops.aimnet.conv_sv_2d_sp_bwd.default._schema) + bwd_bwd = str(torch.ops.aimnet.conv_sv_2d_sp_bwd_bwd.default._schema) + assert "Tensor a, Tensor idx, Tensor g, SymInt padding_value, SymInt num_centers" in fwd + assert "Tensor grad_output, Tensor a, Tensor idx, Tensor g, SymInt padding_value, SymInt num_centers" in bwd + assert ( + "Tensor grad_output, Tensor grad2_a, Tensor grad2_g, Tensor a, Tensor idx, Tensor g, " + "SymInt padding_value, SymInt num_centers" + ) in bwd_bwd + + +def test_conv_sv_generalized_fake_outputs(): + from torch._subclasses.fake_tensor import FakeTensorMode + + import aimnet.kernels.conv_sv_2d_sp_wp # noqa: F401 + + with FakeTensorMode(): + a = torch.empty(4, 3, 2, device="cuda") + idx = torch.empty(4, 5, dtype=torch.int32, device="cuda") + g = torch.empty(4, 5, 2, 4, device="cuda") + out = torch.ops.aimnet.conv_sv_2d_sp_fwd(a, idx, g, 4, 4) + grad_a, grad_g = torch.ops.aimnet.conv_sv_2d_sp_bwd(out, a, idx, g, 4, 4) + second = torch.ops.aimnet.conv_sv_2d_sp_bwd_bwd(out, grad_a, grad_g, a, idx, g, 4, 4) + assert out.shape == (4, 3, 2, 4) + assert grad_a.shape == a.shape + assert grad_g.shape == g.shape + assert [value.shape for value in second] == [out.shape, a.shape, g.shape] + + +def test_conv_sv_generalized_fake_rejects_num_centers(): + from torch._subclasses.fake_tensor import FakeTensorMode + + import aimnet.kernels.conv_sv_2d_sp_wp # noqa: F401 + + with FakeTensorMode(): + a = torch.empty(4, 3, 2, device="cuda") + idx = torch.empty(4, 5, dtype=torch.int32, device="cuda") + g = torch.empty(4, 5, 2, 4, device="cuda") + with pytest.raises(ValueError, match="num_centers"): + torch.ops.aimnet.conv_sv_2d_sp_fwd(a, idx, g, 4, 5) + + +def test_conv_sv_generalized_forward_vmap_rejected(): + from torch._subclasses.fake_tensor import FakeTensorMode + + import aimnet.kernels.conv_sv_2d_sp_wp # noqa: F401 + + with FakeTensorMode(): + a = torch.empty(2, 4, 3, 2, device="cuda") + idx = torch.empty(4, 5, dtype=torch.int32, device="cuda") + g = torch.empty(4, 5, 2, 4, device="cuda") + with pytest.raises((RuntimeError, torch.AcceleratorError), match=r"(Batching rule|vmap|CUDA error)"): + torch.func.vmap(lambda value: torch.ops.aimnet.conv_sv_2d_sp_fwd(value, idx, g, 4, 4))(a) diff --git a/tests/test_lr.py b/tests/test_lr.py index 31efa0d..a28886c 100644 --- a/tests/test_lr.py +++ b/tests/test_lr.py @@ -622,28 +622,34 @@ def test_dsf_mode2_matches_individual_batches(self, device): charges = torch.randn(B, N, device=device) * 0.3 charges = charges - charges.mean(dim=-1, keepdim=True) - # All-pairs nbmat with local atom indices per batch. - nbmat = torch.zeros((B, N, N - 1), dtype=torch.long, device=device) - for i in range(N): - others = [j for j in range(N) if j != i] - nbmat[:, i, :] = torch.tensor(others, device=device) - mask_ij = torch.zeros((B, N, N - 1), dtype=torch.bool, device=device) + # Global all-pairs matrix with one required dummy atom per system. + nbmat = torch.full((B, N, N), B * N, dtype=torch.int32, device=device) + for b in range(B): + for i in range(N - 1): + others = [b * N + j for j in range(N - 1) if j != i] + nbmat[b, i, : len(others)] = torch.tensor(others, device=device) mask_i = torch.zeros((B, N), dtype=torch.bool, device=device) + mask_i[:, -1] = True module = LRCoulomb(method="dsf", dsf_rc=8.0, subtract_sr=False).to(device) data_mode2 = { - "_nb_mode": torch.tensor(2), - "_input_padded": torch.tensor(False), "coord": coord, "charges": charges, - "numbers": torch.ones((B, N), dtype=torch.long, device=device), + "numbers": torch.cat( + [ + torch.ones((B, N - 1), dtype=torch.long, device=device), + torch.zeros((B, 1), dtype=torch.long, device=device), + ], + dim=1, + ), + "nbmat": nbmat, "nbmat_lr": nbmat, - "mask_ij_lr": mask_ij, "mask_i": mask_i, "mol_idx": torch.arange(B, device=device).repeat_interleave(N), "mol_sizes": torch.full((B,), N, device=device, dtype=torch.long), } + data_mode2 = nbops.calc_masks(nbops.set_nb_mode(data_mode2)) result_mode2, terms_mode2 = module(data_mode2, compute_forces=True) e_mode2 = result_mode2["e_h"] assert terms_mode2 is not None and terms_mode2.forces is not None @@ -654,16 +660,11 @@ def test_dsf_mode2_matches_individual_batches(self, device): f_ref = [] for b in range(B): data_b = { - "_nb_mode": torch.tensor(0), - "_input_padded": torch.tensor(False), - "coord": coord[b : b + 1], - "charges": charges[b : b + 1], - "numbers": torch.ones((1, N), dtype=torch.long, device=device), - "mask_ij": mask_ij[b : b + 1], - "mask_i": mask_i[b : b + 1], - "mol_idx": torch.zeros(N, device=device, dtype=torch.long), - "mol_sizes": torch.tensor([N], device=device, dtype=torch.long), + "coord": coord[b : b + 1, :-1], + "charges": charges[b : b + 1, :-1], + "numbers": torch.ones((1, N - 1), dtype=torch.long, device=device), } + data_b = nbops.calc_masks(nbops.set_nb_mode(data_b)) result_b, terms_b = module(data_b, compute_forces=True) assert terms_b is not None and terms_b.forces is not None e_ref.append(result_b["e_h"]) @@ -672,7 +673,7 @@ def test_dsf_mode2_matches_individual_batches(self, device): f_ref = torch.stack(f_ref, dim=0) torch.testing.assert_close(e_mode2, e_ref, atol=1e-5, rtol=1e-5) - torch.testing.assert_close(terms_mode2.forces, f_ref, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(terms_mode2.forces[:, :-1], f_ref, atol=1e-5, rtol=1e-5) def test_dsf_mode0_padded_batched_indices(self, device): """DSF mode-0 builds a cutoff-bounded NL with correct global indices and @@ -731,9 +732,9 @@ def test_dsf_mode0_padded_batched_indices(self, device): f_ref = torch.stack(f_ref, dim=0) torch.testing.assert_close(e_mode0, e_ref, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(terms_mode0.forces, f_ref, atol=1e-5, rtol=1e-5) # Padded atom force should be zero (it has charge zero and was shifted out). assert torch.allclose(terms_mode0.forces[0, -1], torch.zeros(3, device=device), atol=1e-5) - torch.testing.assert_close(terms_mode0.forces, f_ref, atol=1e-5, rtol=1e-5) def test_dsf_mode0_large_coordinates_keep_padding_out_of_neighbor_list(self, device): """Padded atoms stay out of the DSF NL even for large unwrapped coordinates.""" @@ -788,6 +789,24 @@ def test_dsf_mode0_large_coordinates_keep_padding_out_of_neighbor_list(self, dev assert torch.allclose(terms_padded.forces[0, padded_idx], torch.zeros(3, device=device), atol=1e-5) +def test_global_mode2_lr_matrix_is_symmetric(): + B, N, M = 2, 5, 6 + sentinel = B * N + nbmat = torch.full((B, N, M), sentinel, dtype=torch.int32) + for b in range(B): + for i in range(N - 1): + targets = [j for j in range(N - 1) if j != i] + nbmat[b, i, : len(targets)] = torch.tensor([b * N + j for j in targets]) + for b in range(B): + for i in range(N - 1): + for target in nbmat[b, i]: + if int(target) != sentinel: + source = b * N + i + reverse = nbmat[b, target - b * N] + assert source in reverse + assert torch.equal(nbmat[b, -1], torch.full((M,), sentinel, dtype=torch.int32)) + + class TestLRCoulombEwaldPBC: """Tests for Ewald Coulomb with periodic boundary conditions.""" diff --git a/tests/test_mode2_cuda_validation.py b/tests/test_mode2_cuda_validation.py new file mode 100644 index 0000000..f019926 --- /dev/null +++ b/tests/test_mode2_cuda_validation.py @@ -0,0 +1,104 @@ +"""Isolated CUDA tests for mode-2 validation failures.""" + +import os +import subprocess +import sys + +import pytest +import torch + +from aimnet.calculators import AIMNet2Calculator +from aimnet.models.base import AIMNet2Base + +pytestmark = [ + pytest.mark.gpu, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device unavailable"), +] + + +def _cuda_mode2_data() -> dict[str, torch.Tensor]: + device = torch.device("cuda") + return { + "coord": torch.zeros((2, 4, 3), device=device), + "numbers": torch.tensor([[6, 1, 1, 0], [8, 1, 1, 0]], device=device), + "charge": torch.zeros(2, device=device), + "nbmat": torch.tensor( + [ + [[1, 2, 8], [0, 2, 8], [0, 1, 8], [8, 8, 8]], + [[5, 6, 8], [4, 6, 8], [4, 5, 8], [8, 8, 8]], + ], + device=device, + dtype=torch.int64, + ), + } + + +def _run_invalid_child(case: str) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env["AIMNET_MODE2_INVALID_CASE"] = case + return subprocess.run( # noqa: S603 + [sys.executable, "-m", "pytest", __file__, "-k", "test_global_mode2_invalid_child", "-q"], + env=env, + capture_output=True, + text=True, + check=False, + ) + + +def test_global_mode2_invalid_child(): + case = os.environ.get("AIMNET_MODE2_INVALID_CASE") + if case is None: + pytest.skip("child process only") + data = _cuda_mode2_data() + if case == "padded_center": + data["nbmat"][0, -1, 0] = 0 + elif case in {"cross_batch_direct", "cross_batch_hessian_preremap"}: + if case == "cross_batch_hessian_preremap": + data["nbmat"][0, 0, 0] = 4 + else: + data["nbmat"][1, 0, 0] = 0 + elif case == "shift_int32_overflow": + data["cell"] = torch.eye(3, device="cuda").repeat(2, 1, 1) + data["shifts"] = torch.zeros((2, 4, 3, 3), device="cuda") + data["shifts"][0, 0, 0, 0] = float(torch.iinfo(torch.int32).max) + 1 + else: + raise AssertionError(f"unknown invalid case: {case}") + if case == "cross_batch_hessian_preremap": + AIMNet2Calculator("aimnet2", nb_threshold=0, device="cuda").eval(data, hessian=True, validate_species=False) + else: + AIMNet2Base().prepare_input(data) + torch.cuda.synchronize() + + +def test_global_mode2_cuda_invalid_input_isolated(): + result = _run_invalid_child("cross_batch_direct") + assert result.returncode != 0 + assert "batch interval" in (result.stdout + result.stderr).lower() + + +def test_global_mode2_cuda_rejects_non_sentinel_padded_center(): + result = _run_invalid_child("padded_center") + assert result.returncode != 0 + assert "padded center" in (result.stdout + result.stderr).lower() + + +def test_global_mode2_cuda_rejects_hessian_preremap(): + result = _run_invalid_child("cross_batch_hessian_preremap") + assert result.returncode != 0 + assert "batch interval" in (result.stdout + result.stderr).lower() + + +def test_global_mode2_cuda_rejects_shift_overflow(): + result = _run_invalid_child("shift_int32_overflow") + assert result.returncode != 0 + assert "int32" in (result.stdout + result.stderr).lower() + + +def test_global_mode2_valid_cuda_path_has_no_host_sync(): + data = _cuda_mode2_data() + previous_mode = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + AIMNet2Base().prepare_input(data) + finally: + torch.cuda.set_sync_debug_mode(previous_mode) diff --git a/tests/test_mode2_periodic_backends.py b/tests/test_mode2_periodic_backends.py new file mode 100644 index 0000000..bc9dd94 --- /dev/null +++ b/tests/test_mode2_periodic_backends.py @@ -0,0 +1,258 @@ +"""CPU and GPU coverage for full-3D periodic global mode 2.""" + +import pytest +import torch + +from aimnet import nbops +from aimnet.modules.lr import DFTD3, LRCoulomb, _mode2_backend_inputs + + +def _periodic_mode2_data(device: torch.device, batch_size: int = 2) -> dict[str, torch.Tensor]: + B, N, M = batch_size, 5, 8 + coord = torch.zeros((B, N, 3), device=device) + coord[:, 1, 0] = 1.0 + coord[:, 2, 1] = 1.1 + coord[:, 3, 2] = 1.2 + if B > 1: + coord[1, :4] += torch.tensor([0.2, 0.3, 0.4], device=device) + numbers = torch.tensor([[8, 1, 1, 1, 0]] * B, device=device) + charges = torch.tensor([[0.4, -0.1, -0.1, -0.1, 0.0]] * B, device=device) + nbmat = torch.full((B, N, M), B * N, dtype=torch.int32, device=device) + for b in range(B): + for i in range(N - 1): + targets = [b * N + j for j in range(N - 1) if j != i] + nbmat[b, i, : len(targets)] = torch.tensor(targets, device=device) + shifts = torch.zeros((B, N, M, 3), device=device) + shifts[:, 0, 0, 0] = 1 + shifts[:, 1, 0, 0] = -1 + data = { + "coord": coord, + "numbers": numbers, + "charges": charges, + "nbmat": nbmat, + "nbmat_lr": nbmat, + "nbmat_coulomb": nbmat, + "nbmat_dftd3": nbmat, + "shifts": shifts, + "shifts_lr": shifts, + "shifts_coulomb": shifts, + "shifts_dftd3": shifts, + "cell": torch.eye(3, device=device).expand(B, -1, -1) * 12.0, + "pbc": torch.ones((B, 3), dtype=torch.bool, device=device), + } + return nbops.calc_masks(nbops.set_nb_mode(data)) + + +def _single_mode2_periodic_data(data: dict[str, torch.Tensor], batch_index: int) -> dict[str, torch.Tensor]: + """Extract one global mode-2 system for independent execution.""" + B, N, _M = data["nbmat"].shape + source_sentinel = B * N + sentinel = N + local_nbmat = torch.where( + data["nbmat"][batch_index] == source_sentinel, + torch.full_like(data["nbmat"][batch_index], sentinel), + data["nbmat"][batch_index] - batch_index * N, + ).to(torch.int32) + mode1 = { + "coord": data["coord"][batch_index : batch_index + 1], + "numbers": data["numbers"][batch_index : batch_index + 1], + "charges": data["charges"][batch_index : batch_index + 1], + "cell": data["cell"][batch_index : batch_index + 1], + "pbc": data["pbc"][batch_index : batch_index + 1], + } + for suffix in ("", "_lr", "_coulomb", "_dftd3"): + mode1[f"nbmat{suffix}"] = local_nbmat.unsqueeze(0) + mode1[f"shifts{suffix}"] = data[f"shifts{suffix}"][batch_index : batch_index + 1] + return nbops.calc_masks(nbops.set_nb_mode(mode1)) + + +def _module(backend: str): + if backend == "dftd3": + return DFTD3(s8=0.3908, a1=0.5660, a2=3.1280) + return LRCoulomb(method=backend, subtract_sr=False, ewald_accuracy=1e-5) + + +@pytest.mark.parametrize("backend", ["dsf", "dftd3", "ewald", "pme"]) +def test_global_mode2_cpu_periodic_observables(backend: str): + data = _periodic_mode2_data(torch.device("cpu")) + module = _module(backend) + result = module(data) + energy_key = "energy" if backend == "dftd3" else "e_h" + assert result[energy_key].shape == (2,) + assert torch.isfinite(result[energy_key]).all() + + +@pytest.mark.parametrize("backend", ["dsf", "dftd3", "ewald", "pme"]) +def test_global_mode2_cpu_periodic_forces_and_stress(backend: str): + data = _periodic_mode2_data(torch.device("cpu")) + if backend == "ewald": + data = { + **data, + "coord": data["coord"].detach().requires_grad_(True), + "cell": data["cell"].detach().requires_grad_(True), + } + result = _module(backend)(data) + forces = -torch.autograd.grad(result["e_h"].sum(), data["coord"], retain_graph=True)[0] + virial = torch.autograd.grad(result["e_h"].sum(), data["cell"])[0] + terms = None + else: + result, terms = _module(backend)(data, compute_forces=True, compute_virial=True) + forces = terms.forces if terms is not None else None + virial = terms.virial if terms is not None else None + energy_key = "energy" if backend == "dftd3" else "e_h" + assert torch.isfinite(result[energy_key]).all() + assert forces is not None and virial is not None + assert forces.shape == data["coord"].shape + assert virial.shape[-2:] == (3, 3) + assert torch.isfinite(forces).all() + assert torch.isfinite(virial).all() + + +@pytest.mark.parametrize("backend", ["dsf", "dftd3", "ewald", "pme"]) +def test_global_mode2_periodic_matches_single_system(backend: str): + data = _periodic_mode2_data(torch.device("cpu")) + module = _module(backend) + batch_result, batch_terms = module(data, compute_forces=True, compute_virial=True) + energy_key = "energy" if backend == "dftd3" else "e_h" + for batch_index in range(data["coord"].shape[0]): + single_data = _single_mode2_periodic_data(data, batch_index) + single_result, single_terms = module(single_data, compute_forces=True, compute_virial=True) + torch.testing.assert_close( + batch_result[energy_key][batch_index], + single_result[energy_key].reshape(-1)[0], + atol=2e-5, + rtol=2e-4, + ) + if backend == "ewald": + batch_forces, batch_virial = _ewald_forces_and_virial(data) + single_forces, single_virial = _ewald_forces_and_virial(single_data) + torch.testing.assert_close( + batch_forces[batch_index], + single_forces[0], + atol=2e-5, + rtol=2e-4, + ) + torch.testing.assert_close( + batch_virial[batch_index], + single_virial[0], + atol=2e-5, + rtol=2e-4, + ) + else: + assert batch_terms is not None and single_terms is not None + torch.testing.assert_close( + batch_terms.forces[batch_index], + single_terms.forces[0], + atol=2e-5, + rtol=2e-4, + ) + torch.testing.assert_close( + batch_terms.virial[batch_index], + single_terms.virial[0], + atol=2e-5, + rtol=2e-4, + ) + + +def _ewald_forces_and_virial(data: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + """Differentiate the Ewald energy for its autograd-only observables.""" + sample = {key: value.clone() for key, value in data.items()} + sample["coord"] = sample["coord"].detach().requires_grad_(True) + sample["cell"] = sample["cell"].detach().requires_grad_(True) + energy = _module("ewald")(sample)["e_h"].sum() + grad_coord, grad_cell = torch.autograd.grad(energy, (sample["coord"], sample["cell"])) + return -grad_coord, grad_cell + + +@pytest.mark.parametrize("backend", ["dsf", "dftd3", "ewald", "pme"]) +def test_global_mode2_periodic_hessian_diagonal_matches_independent(backend: str): + data = _periodic_mode2_data(torch.device("cpu")) + energy_key = "energy" if backend == "dftd3" else "e_h" + epsilon = 1e-3 + + def energy_at(sample: dict[str, torch.Tensor], index: int) -> torch.Tensor: + fresh = {key: value.clone() for key, value in sample.items()} + fresh.pop("e_h", None) + fresh.pop("energy", None) + return _module(backend)(fresh)[energy_key].reshape(-1)[index] + + for batch_index in range(data["coord"].shape[0]): + single_data = _single_mode2_periodic_data(data, batch_index) + batch_energy = energy_at(data, batch_index) + single_energy = energy_at(single_data, 0) + plus_batch = {key: value.clone() for key, value in data.items()} + minus_batch = {key: value.clone() for key, value in data.items()} + plus_batch["coord"][batch_index, 0, 0] += epsilon + minus_batch["coord"][batch_index, 0, 0] -= epsilon + plus_single = {key: value.clone() for key, value in single_data.items()} + minus_single = {key: value.clone() for key, value in single_data.items()} + plus_single["coord"][0, 0, 0] += epsilon + minus_single["coord"][0, 0, 0] -= epsilon + batch_hessian = ( + energy_at(plus_batch, batch_index) - 2 * batch_energy + energy_at(minus_batch, batch_index) + ) / epsilon**2 + single_hessian = (energy_at(plus_single, 0) - 2 * single_energy + energy_at(minus_single, 0)) / epsilon**2 + torch.testing.assert_close(batch_hessian, single_hessian, atol=5e-3, rtol=5e-3) + + +def test_global_mode2_backend_views_share_storage(): + data = _periodic_mode2_data(torch.device("cpu")) + inputs = _mode2_backend_inputs(data, "_lr") + assert inputs.coord._base is not None + assert inputs.neighbor_matrix._base is not None + assert inputs.shifts is not None and inputs.shifts._base is not None + assert inputs.coord.storage().data_ptr() == data["coord"].storage().data_ptr() + assert inputs.neighbor_matrix.storage().data_ptr() == data["_nbmat_kernel_lr"].storage().data_ptr() + assert inputs.shifts.storage().data_ptr() == data["shifts_lr"].storage().data_ptr() + + +@pytest.mark.gpu +@pytest.mark.parametrize("backend", ["dsf", "dftd3", "ewald", "pme"]) +def test_global_mode2_periodic_backends_accept_dummy_rows(backend: str): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + data = _periodic_mode2_data(torch.device("cuda")) + module = _module(backend).cuda() + if backend == "dftd3": + inputs = module._prepare_dftd3_inputs(data) + assert inputs.coord_flat.shape[0] == data["coord"].numel() // 3 + assert inputs.neighbor_matrix.shape[0] == data["coord"].numel() // 3 + assert inputs.numbers_flat[-1] == 0 + else: + inputs = module._dsf_inputs(data, "_lr") if backend == "dsf" else _mode2_backend_inputs(data, "_coulomb") + if backend == "dsf": + coord_flat, _charges, _batch_idx, neighbor_matrix, _cell, _shifts, fill_value, _num_systems = inputs + else: + coord_flat, neighbor_matrix, _shifts, _batch_idx, fill_value, _num_systems, _cell = inputs + assert coord_flat.shape[0] == data["coord"].numel() // 3 + assert neighbor_matrix.shape[0] == data["coord"].numel() // 3 + if backend == "dftd3": + fill_value = inputs.fill_value + assert int(fill_value) == data["coord"].shape[0] * data["coord"].shape[1] + + +@pytest.mark.gpu +@pytest.mark.parametrize("backend", ["dsf", "dftd3", "ewald", "pme"]) +def test_global_mode2_periodic_backends_cross_batch_isolation(backend: str): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + data = _periodic_mode2_data(torch.device("cuda")) + module = _module(backend).cuda() + key = "energy" if backend == "dftd3" else "e_h" + baseline = module({**data})[key].detach().clone() + mutated = {name: value.clone() if isinstance(value, torch.Tensor) else value for name, value in data.items()} + mutated["coord"][0, 1, 0] += 0.4 + changed = module(mutated)[key].detach() + assert torch.allclose(changed[1], baseline[1], atol=1e-6, rtol=1e-6) + assert not torch.allclose(changed[0], baseline[0], atol=1e-6, rtol=1e-6) + + +@pytest.mark.gpu +@pytest.mark.parametrize("backend", ["dsf", "dftd3", "ewald", "pme"]) +def test_global_mode2_gpu_periodic_observables(backend: str): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + data = _periodic_mode2_data(torch.device("cuda")) + result = _module(backend).cuda()(data) + key = "energy" if backend == "dftd3" else "e_h" + assert torch.isfinite(result[key]).all() diff --git a/tests/test_model.py b/tests/test_model.py index 5c673b6..05d1080 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -711,3 +711,60 @@ def test_aimnet2_rxn_yaml_builds(): assert hasattr(model.outputs, "dipole") assert hasattr(model.outputs, "quadrupole") assert hasattr(model.outputs, "lrcoulomb") + + +def test_global_mode2_validates_before_int32_narrowing(): + """A float mode-2 matrix must be rejected before conversion to int32.""" + from aimnet.models.base import AIMNet2Base + + data = { + "coord": torch.zeros((2, 4, 3), dtype=torch.float32), + "numbers": torch.tensor([[6, 1, 1, 0], [8, 1, 1, 0]]), + "charge": torch.zeros(2), + "nbmat": torch.tensor([ + [[1.5, 2.0, 8.0], [0.0, 2.0, 8.0], [0.0, 1.0, 8.0], [8.0, 8.0, 8.0]], + [[5.0, 6.0, 8.0], [4.0, 6.0, 8.0], [4.0, 5.0, 8.0], [8.0, 8.0, 8.0]], + ]), + } + with pytest.raises(ValueError, match="integer dtype"): + AIMNet2Base().prepare_input(data) + + +def test_global_mode2_full_model_cpu_batch_vs_individual(): + """Global mode-2 model energies are isolated per padded system.""" + model = build_model(aimnet2_d3_def).eval() + B, N, M = 2, 4, 3 + coord = torch.tensor([ + [[0.0, 0.0, 0.0], [0.96, 0.0, 0.0], [0.0, 0.9, 0.0], [0.0, 0.0, 0.0]], + [[2.0, 0.0, 0.0], [2.9, 0.0, 0.0], [2.0, 1.0, 0.0], [2.0, 0.0, 0.0]], + ]) + numbers = torch.tensor([[8, 1, 1, 0], [6, 1, 1, 0]]) + nbmat = torch.full((B, N, M), B * N, dtype=torch.int64) + for b in range(B): + for i in range(N - 1): + targets = [b * N + j for j in range(N - 1) if j != i] + nbmat[b, i, : len(targets)] = torch.tensor(targets) + data = { + "coord": coord, + "numbers": numbers, + "charge": torch.zeros(B), + "nbmat": nbmat, + "nbmat_lr": nbmat, + "nbmat_coulomb": nbmat, + "nbmat_dftd3": nbmat, + } + with torch.no_grad(): + batched = model(data)["energy"] + individual = [] + for b in range(B): + single = { + "coord": coord[b : b + 1], + "numbers": numbers[b : b + 1], + "charge": torch.zeros(1), + "nbmat": torch.where(nbmat[b : b + 1] == B * N, torch.tensor(N), nbmat[b : b + 1] - b * N), + "nbmat_lr": torch.where(nbmat[b : b + 1] == B * N, torch.tensor(N), nbmat[b : b + 1] - b * N), + "nbmat_coulomb": torch.where(nbmat[b : b + 1] == B * N, torch.tensor(N), nbmat[b : b + 1] - b * N), + "nbmat_dftd3": torch.where(nbmat[b : b + 1] == B * N, torch.tensor(N), nbmat[b : b + 1] - b * N), + } + individual.append(model(single)["energy"]) + torch.testing.assert_close(batched, torch.cat(individual), atol=1e-6, rtol=1e-5) diff --git a/tests/test_nbops.py b/tests/test_nbops.py index 867d692..7109db2 100644 --- a/tests/test_nbops.py +++ b/tests/test_nbops.py @@ -111,18 +111,12 @@ def test_calc_masks_mode_2(self, device): """Test mask calculation for mode 2 (batched with 3D nbmat).""" B, N = 2, 4 coord = torch.rand((B, N, 3), device=device) - numbers = torch.tensor( - [ - [6, 1, 1, 0], # molecule with padding - [6, 1, 0, 0], # molecule with 2 padding atoms - ], - device=device, - ) + numbers = torch.tensor([[6, 1, 1, 0], [6, 1, 0, 0]], device=device) # 3D nbmat: (B, N, max_neighbors) nbmat = torch.tensor( [ - [[1, 2, 3], [0, 2, 3], [0, 1, 3], [3, 3, 3]], # batch 0 - [[1, 2, 3], [0, 2, 3], [2, 2, 2], [3, 3, 3]], # batch 1 + [[1, 2, 8], [0, 2, 8], [0, 1, 8], [8, 8, 8]], + [[5, 6, 8], [4, 6, 8], [8, 8, 8], [8, 8, 8]], ], device=device, ) @@ -141,38 +135,38 @@ def test_calc_masks_mode_2(self, device): assert data["mol_sizes"][0].item() == 3 assert data["mol_sizes"][1].item() == 2 - # local padded neighbor indices are masked per batch + # global sentinel and padding-target indices are masked per batch assert data["mask_ij"][0, 0, 2].item() is True assert data["mask_ij"][1, 0, 1].item() is True # padded center rows are fully masked assert data["mask_ij"][1, 2].all() def test_calc_masks_mode_2_masks_local_and_global_padding(self, device): - numbers = torch.tensor([[6, 0, 1], [6, 1, 0]], device=device) - nbmat_local = torch.tensor( + numbers = torch.tensor([[6, 1, 1, 0], [6, 1, 0, 0]], device=device) + nbmat_primary = torch.tensor( [ - [[1, 2], [0, 2], [0, 1]], - [[1, 2], [0, 2], [0, 1]], + [[1, 3, 8], [0, 3, 8], [0, 1, 8], [8, 8, 8]], + [[5, 6, 8], [4, 6, 8], [8, 8, 8], [8, 8, 8]], ], device=device, ) - nbmat_global = torch.tensor( + nbmat_lr = torch.tensor( [ - [[1, 2], [0, 2], [0, 1]], - [[4, 5], [3, 5], [3, 4]], + [[1, 2, 8], [0, 2, 8], [0, 1, 8], [8, 8, 8]], + [[5, 7, 8], [4, 7, 8], [8, 8, 8], [8, 8, 8]], ], device=device, ) - data = {"numbers": numbers, "nbmat": nbmat_local, "nbmat_lr": nbmat_global} + data = {"numbers": numbers, "nbmat": nbmat_primary, "nbmat_lr": nbmat_lr} data = nbops.set_nb_mode(data) data = nbops.calc_masks(data) - assert data["mask_ij"][0, 0, 0].item() is True # local pad atom 1 in batch 0 - assert data["mask_ij"][1, 0, 1].item() is True # local pad atom 2 in batch 1 - assert data["mask_ij_lr"][0, 0, 0].item() is True # global pad atom 1 - assert data["mask_ij_lr"][1, 0, 1].item() is True # global pad atom 5 - assert data["mask_ij"][0, 1].all() + assert data["mask_ij"][0, 0, 1].item() is True + assert data["mask_ij"][1, 0, 1].item() is True + assert data["mask_ij_lr"][0, 0, 2].item() is True + assert data["mask_ij_lr"][1, 0, 1].item() is True + assert data["mask_ij"][0, 3].all() assert data["mask_ij_lr"][1, 2].all() @@ -332,13 +326,13 @@ def test_get_ij_mode_1(self, device): def test_get_ij_mode_2(self, device): """Test pairwise extraction for mode 2.""" B, N = 2, 3 - numbers = torch.tensor([[6, 1, 1], [6, 1, 0]], device=device) + numbers = torch.tensor([[6, 1, 0], [6, 1, 0]], device=device) nbmat = torch.tensor( - [[[1, 2], [0, 2], [0, 1]], [[1, 2], [0, 2], [0, 1]]], # (B, N, max_nb) + [[[1, 2], [0, 2], [6, 6]], [[4, 5], [3, 5], [6, 6]]], device=device, ) - data = {"numbers": numbers, "nbmat": nbmat, "_nb_mode": torch.tensor(2)} + data = nbops.calc_masks({"numbers": numbers, "nbmat": nbmat, "_nb_mode": torch.tensor(2)}) x = torch.tensor([[[1.0], [2.0], [3.0]], [[4.0], [5.0], [6.0]]], device=device) x_i, x_j = nbops.get_ij(x, data) @@ -391,13 +385,13 @@ def test_get_i_mode_1(self, device): def test_get_i_mode_2(self, device): """Test get_i for mode 2 matches get_ij x_i component.""" B, N = 2, 3 - numbers = torch.tensor([[6, 1, 1], [6, 1, 0]], device=device) + numbers = torch.tensor([[6, 1, 0], [6, 1, 0]], device=device) nbmat = torch.tensor( - [[[1, 2], [0, 2], [0, 1]], [[1, 2], [0, 2], [0, 1]]], # (B, N, max_nb) + [[[1, 2], [0, 2], [6, 6]], [[4, 5], [3, 5], [6, 6]]], device=device, ) - data = {"numbers": numbers, "nbmat": nbmat, "_nb_mode": torch.tensor(2)} + data = nbops.calc_masks({"numbers": numbers, "nbmat": nbmat, "_nb_mode": torch.tensor(2)}) x = torch.tensor([[[1.0], [2.0], [3.0]], [[4.0], [5.0], [6.0]]], device=device) x_i_only = nbops.get_i(x, data) @@ -564,7 +558,7 @@ def test_mol_sum_gradient(self, device): # Gradient should be 1 for all inputs assert x.grad is not None - assert (x.grad == 1.0).all() + torch.testing.assert_close(x.grad, torch.ones_like(x)) def test_mask_ij_gradient(self, device): """Test that gradients flow through mask_ij_ (not inplace).""" @@ -596,3 +590,407 @@ def test_get_ij_gradient_mode_0(self, device): loss.backward() assert x.grad is not None + + +def _global_mode2_data( + device: torch.device, + *, + pad_neighbor: bool = False, + suffixes: tuple[str, ...] = (), + include_shifts: bool = False, +) -> dict[str, torch.Tensor]: + """Build a small global-index mode-2 case with one trailing dummy per system.""" + B, N, M = 2, 4, 3 + sentinel = B * N + coord = torch.arange(B * N * 3, device=device, dtype=torch.float32).reshape(B, N, 3) + numbers = torch.tensor([[6, 1, 1, 0], [8, 1, 1, 0]], device=device) + nbmat = torch.full((B, N, M), sentinel, device=device, dtype=torch.int64) + for b in range(B): + base = b * N + nbmat[b, 0, :2] = torch.tensor([base + 1, base + 2], device=device) + nbmat[b, 1, :2] = torch.tensor([base, base + 2], device=device) + nbmat[b, 2, :2] = torch.tensor([base, base + 1], device=device) + if pad_neighbor: + nbmat[0, 0, 1] = N - 1 + data: dict[str, torch.Tensor] = {"coord": coord, "numbers": numbers, "nbmat": nbmat} + if include_shifts: + shifts = torch.zeros((*nbmat.shape, 3), device=device, dtype=torch.float32) + data["shifts"] = shifts + for suffix in suffixes: + data[f"nbmat{suffix}"] = nbmat.clone() + if include_shifts: + data[f"shifts{suffix}"] = shifts.clone() + return data + + +def test_global_mode2_gathers_distinct_batch_values(device): + data = nbops.calc_masks(nbops.set_nb_mode(_global_mode2_data(device))) + values = torch.tensor([[[10.0], [11.0], [12.0], [0.0]], [[20.0], [21.0], [22.0], [0.0]]], device=device) + _x_i, x_j = nbops.get_ij(values, data) + assert x_j[0, 0, 0, 0] == 11.0 + assert x_j[1, 0, 0, 0] == 21.0 + + +def test_global_mode2_excludes_sentinel(device): + data = nbops.calc_masks(nbops.set_nb_mode(_global_mode2_data(device))) + values = torch.arange(8, device=device, dtype=torch.float32).reshape(2, 4, 1) + _x_i, x_j = nbops.get_ij(values, data) + assert x_j[0, 0, 2, 0] == 0.0 + assert x_j[1, 0, 2, 0] == 0.0 + + +def test_global_mode2_masks_padded_neighbor(device): + data = nbops.calc_masks(nbops.set_nb_mode(_global_mode2_data(device, pad_neighbor=True))) + assert data["mask_ij"][0, 0, 1] + assert data["_nbmat_gather"][0, 0, 1] == 0 + + +def test_global_mode2_kernel_indices_exclude_padded_neighbor(device): + data = nbops.calc_masks(nbops.set_nb_mode(_global_mode2_data(device, pad_neighbor=True))) + assert data["_nbmat_kernel"][0, 0, 1] == data["nbmat"].shape[0] * data["nbmat"].shape[1] + + +def test_global_mode2_masks_padded_center(device): + data = nbops.calc_masks(nbops.set_nb_mode(_global_mode2_data(device))) + assert data["mask_ij"][0, -1].all() + assert data["_nbmat_gather"][0, -1].eq(0).all() + + +def test_global_mode2_rejects_non_sentinel_padded_center_cpu(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"][0, -1, 0] = 0 + with pytest.raises(ValueError, match="padded center"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_local_batch1_index_cpu(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"][1, 0, 0] = 0 + with pytest.raises(ValueError, match="batch interval"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_interleaved_sentinel_cpu(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"][0, 0] = torch.tensor([1, 8, 2]) + with pytest.raises(ValueError, match=r"packed.*tail"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_interleaved_padded_neighbor_cpu(): + data = _global_mode2_data(torch.device("cpu"), pad_neighbor=True) + data["nbmat"][0, 0, 2] = 2 + with pytest.raises(ValueError, match=r"packed.*tail"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_builds_independent_suffix_tensors(device): + data = nbops.calc_masks(nbops.set_nb_mode(_global_mode2_data(device, suffixes=("_lr",)))) + assert data["mask_ij"] is not data["mask_ij_lr"] + assert data["_nbmat_gather"] is not data["_nbmat_gather_lr"] + + +def test_global_mode2_reuses_exact_alias_int32(device): + data = _global_mode2_data(device, suffixes=("_lr",)) + data["nbmat"] = data["nbmat"].to(torch.int32) + data["nbmat_lr"] = data["nbmat"] + data = nbops.calc_masks(nbops.set_nb_mode(data)) + assert data["_nbmat_gather"] is data["_nbmat_gather_lr"] + + +def test_global_mode2_reuses_exact_alias_int64(device): + data = _global_mode2_data(device, suffixes=("_lr",)) + data["nbmat_lr"] = data["nbmat"] + data = nbops.calc_masks(nbops.set_nb_mode(data)) + data["mask_i"] = data["numbers"] == 0 + nbops._prepare_mode2_neighbor_tensors(data) + assert data["nbmat"] is data["nbmat_lr"] + + +def test_global_mode2_compile_alias_reuse_int32(device): + data = _global_mode2_data(device, suffixes=("_lr",)) + data["nbmat"] = data["nbmat"].to(torch.int32) + data["nbmat_lr"] = data["nbmat"] + data = nbops.set_nb_mode(data) + data["mask_i"] = data["numbers"] == 0 + nbops._prepare_mode2_neighbor_tensors(data) + assert data["_nbmat_gather"] is data["_nbmat_gather_lr"] + + +def test_global_mode2_compile_alias_reuse_int64(device): + data = _global_mode2_data(device, suffixes=("_lr",)) + data["nbmat_lr"] = data["nbmat"] + data = nbops.set_nb_mode(data) + data["mask_i"] = data["numbers"] == 0 + nbops._prepare_mode2_neighbor_tensors(data) + assert data["_nbmat_kernel"] is data["_nbmat_kernel_lr"] + + +def test_global_mode2_rejects_bool_neighbors(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"] = data["nbmat"].to(torch.bool) + with pytest.raises(ValueError, match="integer dtype"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_unsigned_neighbors(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"] = data["nbmat"].to(torch.uint8) + with pytest.raises(ValueError, match="signed"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_float_neighbors(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"] = data["nbmat"].to(torch.float32) + with pytest.raises(ValueError, match="integer dtype"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_complex_neighbors(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"] = data["nbmat"].to(torch.complex64) + with pytest.raises(ValueError, match="integer dtype"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_int32_capacity_overflow(): + n = torch.iinfo(torch.int32).max + 1 + data = { + "coord": torch.empty((1, n, 3), device="meta"), + "numbers": torch.empty((1, n), dtype=torch.int64, device="meta"), + "nbmat": torch.empty((1, n, 1), dtype=torch.int64, device="meta"), + } + with pytest.raises(ValueError, match="int32"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_mismatched_shift_shape(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + data["shifts"] = torch.zeros(2, 4, 2, 3) + with pytest.raises(ValueError, match="shifts"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_orphan_shift(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + del data["nbmat"] + with pytest.raises(ValueError, match="matching"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_missing_final_dummy(): + data = _global_mode2_data(torch.device("cpu")) + data["numbers"][:, -1] = 6 + with pytest.raises(ValueError, match="final dummy"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_noncontiguous_atom_padding(): + data = _global_mode2_data(torch.device("cpu")) + data["numbers"][0, 1] = 0 + with pytest.raises(ValueError, match="contiguous"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_noncontiguous_coord(): + data = _global_mode2_data(torch.device("cpu")) + data["coord"] = torch.zeros((4, 2, 3)).transpose(0, 1) + with pytest.raises(ValueError, match="flatten"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_noncontiguous_numbers(): + data = _global_mode2_data(torch.device("cpu")) + data["numbers"] = torch.tensor([[6, 8], [1, 1], [1, 1], [0, 0]]).transpose(0, 1) + with pytest.raises(ValueError, match="flatten"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_noncontiguous_nbmat(): + data = _global_mode2_data(torch.device("cpu")) + data["nbmat"] = torch.zeros((4, 2, 3), dtype=torch.int64).transpose(0, 1) + with pytest.raises(ValueError, match="flatten"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_noncontiguous_shifts(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + data["shifts"] = torch.zeros((4, 2, 3, 3)).transpose(0, 1) + with pytest.raises(ValueError, match="flatten"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_normalizes_single_periodic_geometry(): + data = {"cell": torch.eye(3), "pbc": torch.ones(3, dtype=torch.bool)} + nbops.normalize_mode2_periodic_geometry(data, B=1) + assert data["cell"].shape == (1, 3, 3) + assert data["pbc"].shape == (1, 3) + + +def test_global_mode2_preserves_batched_periodic_geometry(): + data = _global_mode2_data(torch.device("cpu")) + cell = torch.stack([torch.eye(3), torch.eye(3) * 2]) + pbc = torch.ones((2, 3), dtype=torch.bool) + data["cell"] = cell + data["pbc"] = pbc + nbops.normalize_mode2_periodic_geometry(data, B=2) + assert data["cell"] is cell + assert data["pbc"] is pbc + + +def test_global_mode2_rejects_pbc_without_cell(): + data = _global_mode2_data(torch.device("cpu")) + data["pbc"] = torch.ones(3, dtype=torch.bool) + with pytest.raises(ValueError, match="cell"): + nbops.normalize_mode2_periodic_geometry(data, B=2) + + +def test_global_mode2_rejects_shifts_without_cell(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + with pytest.raises(ValueError, match="cell"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_partial_pbc(): + data = _global_mode2_data(torch.device("cpu")) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.tensor([True, False, True]) + with pytest.raises(ValueError, match="full-3D"): + nbops.normalize_mode2_periodic_geometry(data, B=2) + + +def test_global_mode2_rejects_missing_periodic_shifts(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + del data["shifts"] + with pytest.raises(ValueError, match="shifts"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_fractional_periodic_shifts(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + data["shifts"][0, 0, 0, 0] = 0.5 + with pytest.raises(ValueError, match="integral"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_shift_int32_overflow(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + data["shifts"][0, 0, 0, 0] = float(torch.iinfo(torch.int32).max) + 1 + with pytest.raises(ValueError, match="int32"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_nonzero_sentinel_shift(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + data["shifts"][0, 0, 2, 0] = 1 + with pytest.raises(ValueError, match="sentinel"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_nonzero_padded_neighbor_shift(): + data = _global_mode2_data(torch.device("cpu"), pad_neighbor=True, include_shifts=True) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + data["shifts"][0, 0, 1, 0] = 1 + with pytest.raises(ValueError, match="padded-neighbor"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_rejects_nonzero_padded_center_shift(): + data = _global_mode2_data(torch.device("cpu"), include_shifts=True) + data["cell"] = torch.eye(3).repeat(2, 1, 1) + data["pbc"] = torch.ones((2, 3), dtype=torch.bool) + data["shifts"][0, -1, 0, 0] = 1 + with pytest.raises(ValueError, match="padded center"): + nbops.validate_mode2_nbmat_raw(data, suffix="") + + +def test_global_mode2_preserves_mode0_and_mode1(device): + mode0 = nbops.set_nb_mode({"numbers": torch.ones((1, 2), device=device)}) + mode1 = nbops.set_nb_mode({"numbers": torch.ones(3, device=device), "nbmat": torch.zeros((3, 1), device=device)}) + assert nbops.get_nb_mode(mode0) == 0 + assert nbops.get_nb_mode(mode1) == 1 + + +def test_global_mode2_convert_local_success_is_immutable(): + local = torch.tensor([[[1, 2, 0], [0, 2, 0], [0, 0, 0], [0, 0, 0]]] * 2, dtype=torch.int64) + padding = torch.tensor([[[False, False, True], [False, False, True], [True, True, True], [True, True, True]]] * 2) + local_before = local.clone() + padding_before = padding.clone() + global_nbmat = nbops.convert_mode2_local_to_global(local, padding_mask=padding) + assert torch.equal(global_nbmat[0, 0], torch.tensor([1, 2, 8])) + assert torch.equal(global_nbmat[1, 0], torch.tensor([5, 6, 8])) + assert torch.equal(local, local_before) + assert torch.equal(padding, padding_before) + + +def test_global_mode2_convert_local_rejects_non_tail_mask(): + local = torch.zeros((1, 2, 3), dtype=torch.int64) + padding = torch.tensor([[[False, True, False], [False, False, True]]]) + with pytest.raises(ValueError, match="tail"): + nbops.convert_mode2_local_to_global(local, padding_mask=padding) + + +def test_global_mode2_convert_local_rejects_out_of_range(): + local = torch.tensor([[[1, 2, 4], [0, 0, 0]]], dtype=torch.int64) + padding = torch.tensor([[[False, False, False], [True, True, True]]]) + with pytest.raises(ValueError, match="range"): + nbops.convert_mode2_local_to_global(local, padding_mask=padding) + + +def test_global_mode2_convert_local_requires_final_dummy_center(): + local = torch.zeros((1, 3, 1), dtype=torch.int32) + padding = torch.zeros((1, 3, 1), dtype=torch.bool) + with pytest.raises(ValueError, match="final dummy"): + nbops.convert_mode2_local_to_global(local, padding_mask=padding) + + +def test_global_mode2_convert_local_upcasts_before_global_arithmetic(): + local = torch.zeros((2, 20_000, 1), dtype=torch.int16) + padding = torch.ones_like(local, dtype=torch.bool) + padding[:, :-1] = False + local[:, :-1] = 0 + result = nbops.convert_mode2_local_to_global(local, padding_mask=padding) + assert result.dtype == torch.int32 + assert result[1, 0, 0] == 20_000 + + +def test_global_mode2_private_preparation_matches_calc_masks(device): + expected = nbops.calc_masks(nbops.set_nb_mode(_global_mode2_data(device))) + actual = nbops.set_nb_mode(_global_mode2_data(device)) + actual["mask_i"] = actual["numbers"] == 0 + nbops._prepare_mode2_neighbor_tensors(actual) + for key in ("mask_ij", "_nbmat_gather", "_nbmat_kernel"): + assert torch.equal(actual[key], expected[key]) + + +def test_global_mode2_rejects_suffix_device_mismatch(): + data = _global_mode2_data(torch.device("cpu"), suffixes=("_lr",)) + data["nbmat_lr"] = data["nbmat_lr"].to("meta") + with pytest.raises(ValueError, match="same device"): + nbops.validate_mode2_nbmat_raw(data, suffix="_lr") + + +def test_global_mode2_rejects_mixed_primary_and_suffix_rank(): + data = _global_mode2_data(torch.device("cpu"), suffixes=("_lr",)) + data["nbmat"] = data["nbmat"][0] + with pytest.raises(ValueError, match="rank"): + nbops.validate_neighbor_suffix_layout(data) + + +def test_global_mode2_rejects_suffix_only_3d_matrix(): + data = _global_mode2_data(torch.device("cpu")) + suffix_only = {"coord": data["coord"], "numbers": data["numbers"], "nbmat_lr": data["nbmat"]} + with pytest.raises(ValueError, match="primary nbmat"): + nbops.validate_neighbor_suffix_layout(suffix_only) diff --git a/tests/test_ops.py b/tests/test_ops.py index 7680aa4..e004c7e 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -567,8 +567,8 @@ def test_calc_distances_batched_cells_nb_mode2(self, device): # Batched coordinates (B, N, 3) coord = torch.tensor( [ - [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], - [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 2.0, 0.0]], + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 0.0]], ], dtype=torch.float32, device=device, @@ -582,8 +582,8 @@ def test_calc_distances_batched_cells_nb_mode2(self, device): # Create neighbor matrix (B, N, M) # Note: In mode 2, nbmat indices are GLOBAL into flattened (B*N) array - # System 0 atoms: 0, 1, 2; System 1 atoms: 3, 4, 5 - B, N, M = 2, 3, 2 + # System 0 atoms: 0, 1, 2; system 1 atoms: 4, 5, 6; index 3/7 are dummies/sentinel. + B, N, M = 2, 4, 2 fill_val = B * N # Fill value is total number of atoms nbmat = torch.full((B, N, M), fill_val, dtype=torch.int32, device=device) # System 0: atom 0's neighbors are atoms 1, 2 (global indices 1, 2) @@ -593,20 +593,20 @@ def test_calc_distances_batched_cells_nb_mode2(self, device): nbmat[0, 1, 1] = 2 nbmat[0, 2, 0] = 0 nbmat[0, 2, 1] = 1 - # System 1: atom 3's neighbors are atoms 4, 5 (global indices 4, 5) - nbmat[1, 0, 0] = 4 # System 1, atom 0 -> global atom 4 - nbmat[1, 0, 1] = 5 # System 1, atom 0 -> global atom 5 - nbmat[1, 1, 0] = 3 - nbmat[1, 1, 1] = 5 - nbmat[1, 2, 0] = 3 - nbmat[1, 2, 1] = 4 + # System 1: atom 4's neighbors are atoms 5, 6 (global indices 5, 6) + nbmat[1, 0, 0] = 5 + nbmat[1, 0, 1] = 6 + nbmat[1, 1, 0] = 4 + nbmat[1, 1, 1] = 6 + nbmat[1, 2, 0] = 4 + nbmat[1, 2, 1] = 5 # Create shifts (B, N, M, 3) shifts = torch.zeros((B, N, M, 3), dtype=torch.float32, device=device) data = { "coord": coord, - "numbers": torch.tensor([[8, 1, 1], [8, 1, 1]], device=device), + "numbers": torch.tensor([[8, 1, 1, 0], [8, 1, 1, 0]], device=device), "nbmat": nbmat, "shifts": shifts, "cell": cell, @@ -628,58 +628,10 @@ def test_calc_distances_batched_cells_nb_mode2(self, device): assert d_ij[1, 0, 0].item() == pytest.approx(2.0, abs=1e-5) def test_calc_distances_single_cell_nb_mode2(self, device): - """Test calc_distances with batched coordinates but single cell (backward compatible).""" - # Batched coordinates (B, N, 3) - coord = torch.tensor( - [ - [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], - [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 2.0, 0.0]], - ], - dtype=torch.float32, - device=device, - ) - - # Single cell (3, 3) + """Reject a shared cell for a multi-system mode-2 batch.""" cell = torch.eye(3, device=device) * 10.0 - - # Create neighbor matrix (B, N, M) - # Note: In mode 2, nbmat indices are GLOBAL into flattened (B*N) array - B, N, M = 2, 3, 2 - fill_val = B * N - nbmat = torch.full((B, N, M), fill_val, dtype=torch.int32, device=device) - # System 0 - nbmat[0, 0, 0] = 1 - nbmat[0, 0, 1] = 2 - nbmat[0, 1, 0] = 0 - nbmat[0, 1, 1] = 2 - nbmat[0, 2, 0] = 0 - nbmat[0, 2, 1] = 1 - # System 1 (global indices 3, 4, 5) - nbmat[1, 0, 0] = 4 - nbmat[1, 0, 1] = 5 - nbmat[1, 1, 0] = 3 - nbmat[1, 1, 1] = 5 - nbmat[1, 2, 0] = 3 - nbmat[1, 2, 1] = 4 - - # Create shifts (B, N, M, 3) - shifts = torch.zeros((B, N, M, 3), dtype=torch.float32, device=device) - - data = { - "coord": coord, - "numbers": torch.tensor([[8, 1, 1], [8, 1, 1]], device=device), - "nbmat": nbmat, - "shifts": shifts, - "cell": cell, # Single cell, not batched - } - data = nbops.set_nb_mode(data) - data = nbops.calc_masks(data) - - # Should work with single cell - d_ij, _r_ij = ops.calc_distances(data) - - assert d_ij.shape == (B, N, M) - assert d_ij[0, 0, 0].item() == pytest.approx(1.0, abs=1e-5) + with pytest.raises(ValueError, match="B, 3, 3"): + nbops.normalize_mode2_periodic_geometry({"cell": cell}, B=2) def test_move_coord_to_cell_respects_partial_pbc(device):