diff --git a/deepmd/_vendors/ndtensorflow/_namespace.py b/deepmd/_vendors/ndtensorflow/_namespace.py index e5ee9bd6db..f981e22599 100644 --- a/deepmd/_vendors/ndtensorflow/_namespace.py +++ b/deepmd/_vendors/ndtensorflow/_namespace.py @@ -1650,7 +1650,11 @@ def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array: axis = 0 axis = _normalize_axis(axis, tensor.shape.rank) indices_ = tf.cast(_unwrap(indices), tf.int64) - dim = tensor.shape[axis] + # ``TensorShape.__getitem__`` returns either ``None`` or + # ``Dimension(None)`` depending on TensorFlow's global v2-shape setting. + # ``as_list`` normalizes both representations, which keeps dynamic-axis + # indexing valid when another test or caller enables legacy TensorShape. + dim = tensor.shape.as_list()[axis] dim = ( tf.shape(tensor, out_type=tf.int64)[axis] if dim is None @@ -1664,7 +1668,7 @@ def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array: tensor = _unwrap(x) indices_ = tf.cast(_unwrap(indices), tf.int64) axis = _normalize_axis(axis, tensor.shape.rank) - dim = tensor.shape[axis] + dim = tensor.shape.as_list()[axis] dim = ( tf.shape(tensor, out_type=tf.int64)[axis] if dim is None diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index d730b43669..2316ab0545 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -1,14 +1,18 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Default all-pairs neighbor-list builder (historical deepmd behavior).""" +"""Default neighbor-list builder with dense and cell-list search paths.""" from typing import ( TYPE_CHECKING, + Any, ) import array_api_compat from deepmd.dpmodel.array_api import ( Array, + xp_hint_dynamic_size, + xp_scatter_sum, + xp_take_along_axis, ) from deepmd.dpmodel.utils.neighbor_list import ( EdgeNeighborList, @@ -18,6 +22,7 @@ NeighborList, ) from .nlist import ( + apply_pair_exclusion_nlist, build_neighbor_list, extend_coord_with_ghosts, ) @@ -31,11 +36,642 @@ ) +# Dense broadcasting has a low fixed cost, but its backend scaling differs enough +# that one shared threshold causes regressions. These conservative thresholds +# are the later crossover measured at rcut=3 and rcut=6 on an AMD EPYC 7K62 CPU +# and RTX 5090 GPU. ``nloc`` counts local atoms before periodic ghost extension. +_NUMPY_CPU_PERIODIC_CELL_LIST_THRESHOLD = 32 +_NUMPY_CPU_NONPERIODIC_CELL_LIST_THRESHOLD = 256 +_TORCH_CPU_PERIODIC_CELL_LIST_THRESHOLD = 32 +_TORCH_CPU_NONPERIODIC_CELL_LIST_THRESHOLD = 2048 +_JAX_CPU_PERIODIC_CELL_LIST_THRESHOLD = 32 +_JAX_CPU_NONPERIODIC_CELL_LIST_THRESHOLD = 256 +_TF_CPU_PERIODIC_CELL_LIST_THRESHOLD = 512 +_TF_CPU_NONPERIODIC_CELL_LIST_THRESHOLD = 4096 +_TORCH_CUDA_PERIODIC_CELL_LIST_THRESHOLD = 1024 +_TORCH_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD = 8192 +_JAX_CUDA_PERIODIC_CELL_LIST_THRESHOLD = 2048 +_JAX_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD = 12288 +_TF_CUDA_PERIODIC_CELL_LIST_THRESHOLD = 1024 +_TF_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD = 8192 + +# Conservative measured points where compacting the extended periodic shell +# becomes cheaper than sorting every image on each CPU backend. +_NUMPY_CPU_PERIODIC_COMPACTION_THRESHOLD = 64 +_TORCH_CPU_PERIODIC_COMPACTION_THRESHOLD = 256 +_JAX_CPU_PERIODIC_COMPACTION_THRESHOLD = 1024 +_TF_CPU_PERIODIC_COMPACTION_THRESHOLD = 512 + +# At rcut=6 and density 0.05 on RTX 5090, periodic compaction crosses over +# between 131072 and 262144 atoms. Use the later point conservatively: below it +# the added mask/nonzero/gather kernels cost more than the smaller key sort saves. +_TORCH_CUDA_PERIODIC_COMPACTION_THRESHOLD = 262144 + +# A Cartesian cell width of ``rcut`` guarantees that two atoms within the cutoff +# can differ by at most one cell in each direction. Keeping the offsets as Python +# data avoids recreating the Cartesian product with backend-specific meshgrid APIs. +_NEIGHBOR_CELL_OFFSETS = tuple( + (ii, jj, kk) for ii in (-1, 0, 1) for jj in (-1, 0, 1) for kk in (-1, 0, 1) +) + +# Bound row padding to four times the compact edge count; more skewed candidate +# distributions keep the compact representation to avoid excessive memory use. +_PADDED_CANDIDATE_OVERHEAD_LIMIT = 4 +_INT32_MAX = 2**31 - 1 + + +def _index_iota(value: Array) -> Array: + """Return ``[0, ..., value.shape[0] - 1]`` without an avoidable scan. + + Eager backends expose the data-dependent length as a concrete Python + integer, so ``arange`` avoids allocating an all-ones temporary and running a + cumulative sum over the full candidate stream. Traced backends can carry a + symbolic length that ``arange`` can not consume portably; retain the scan + construction there to preserve TensorFlow and torch.export compatibility. + """ + xp = array_api_compat.array_namespace(value) + length = value.shape[0] + if isinstance(length, int): + return xp.arange( + length, + dtype=xp.int64, + device=array_api_compat.device(value), + ) + return xp.cumulative_sum(xp.ones_like(value)) - 1 + + +def _supports_cell_list(coord: Array, nloc: Any, *, periodic: bool) -> bool: + """Whether automatic dispatch should use a backend-safe cell-list path. + + All supported namespaces use the compact dynamic-candidate implementation. + JAX neighbor-list construction runs eagerly outside the compiled model step, + so its data-dependent ``repeat`` and ``nonzero`` result lengths are valid. A + symbolic ``nloc`` keeps the dense path because the threshold decision itself + must remain a Python-level, shape-only choice while tracing. + """ + if not isinstance(nloc, int): + return False + xp = array_api_compat.array_namespace(coord) + is_numpy = array_api_compat.is_numpy_array(coord) + is_jax = array_api_compat.is_jax_array(coord) + is_torch = array_api_compat.is_torch_array(coord) + if is_jax: + import jax + + # The compact candidate axis is deliberately eager. Preserve the dense + # fallback for callers that do trace the entire builder, where JAX/XLA + # requires ``repeat`` and ``nonzero`` result lengths to be static. + if isinstance(coord, jax.core.Tracer): + return False + is_ndtensorflow = getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow" + if not (is_numpy or is_jax or is_torch or is_ndtensorflow): + return False + device = array_api_compat.device(coord) + if is_torch and getattr(device, "type", None) not in ("cpu", "cuda"): + # Cell-list primitives are only validated on PyTorch CPU/CUDA. Treat + # MPS, XPU, and future device types conservatively instead of assuming + # that searchsorted/repeat/nonzero have complete backend coverage. + return False + if is_jax and getattr(device, "platform", None) not in ("cpu", "gpu"): + # TPU and future JAX platforms need their own measurements and operation + # coverage before inheriting either the CPU or CUDA crossover. + return False + device_name = str(device).upper() + tf_cpu = is_ndtensorflow and ( + device_name.startswith("CPU") or "/DEVICE:CPU:" in device_name + ) + tf_gpu = is_ndtensorflow and ( + device_name.startswith("GPU") or "/DEVICE:GPU:" in device_name + ) + tf_unplaced = is_ndtensorflow and not device_name + if is_ndtensorflow and not (tf_cpu or tf_gpu or tf_unplaced): + # TensorFlow eager arrays identify CPU/GPU placement directly, while + # symbolic tensors may remain unplaced until graph execution. Reject + # named accelerators that have not been validated by this path. + return False + is_cuda = ( + getattr(device, "type", None) == "cuda" + or getattr(device, "platform", None) == "gpu" + or tf_gpu + ) + if is_cuda: + if is_jax: + threshold = ( + _JAX_CUDA_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _JAX_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD + ) + elif is_ndtensorflow: + threshold = ( + _TF_CUDA_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _TF_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD + ) + else: + threshold = ( + _TORCH_CUDA_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _TORCH_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD + ) + elif is_jax: + threshold = ( + _JAX_CPU_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _JAX_CPU_NONPERIODIC_CELL_LIST_THRESHOLD + ) + elif tf_unplaced: + # An unplaced TensorFlow graph may execute on either CPU or GPU. The + # later measured crossover avoids selecting the cell list too early on + # whichever validated device the runtime eventually chooses. + threshold = max( + _TF_CPU_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _TF_CPU_NONPERIODIC_CELL_LIST_THRESHOLD, + _TF_CUDA_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _TF_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD, + ) + elif is_ndtensorflow: + threshold = ( + _TF_CPU_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _TF_CPU_NONPERIODIC_CELL_LIST_THRESHOLD + ) + elif is_torch: + threshold = ( + _TORCH_CPU_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _TORCH_CPU_NONPERIODIC_CELL_LIST_THRESHOLD + ) + else: + threshold = ( + _NUMPY_CPU_PERIODIC_CELL_LIST_THRESHOLD + if periodic + else _NUMPY_CPU_NONPERIODIC_CELL_LIST_THRESHOLD + ) + return nloc >= threshold + + +def _select_nearest_padded( + center: Array, + neighbor_ext: Array, + distance: Array, + ncenters: int, + nsel: int, +) -> Array | None: + """Select neighbors by sorting independent, padded center rows. + + The compact candidate stream is already grouped by center. Padding those + groups into rows makes the center key implicit, so two row-wise stable sorts + implement ``(distance, ext_index)`` ordering instead of three global sorts + plus a final scatter. Highly imbalanced rows can waste substantial memory; + return ``None`` in that case so the compact global-sort path remains usable. + + This helper requires eager, concrete candidate counts because the maximum + row width controls an allocation. Callers must retain the compact path for + traced namespaces with symbolic data-dependent dimensions. + + ``center`` must be a non-decreasing stream. The cell-list caller derives it + from repeated ordered query IDs and preserves that ordering while filtering; + this grouping is what makes the arithmetic row rank collision-free. + """ + xp = array_api_compat.array_namespace(center, neighbor_ext, distance) + device = array_api_compat.device(center) + ones = xp.ones_like(center) + count_per_center = xp_scatter_sum( + xp.zeros((ncenters,), dtype=xp.int64, device=device), 0, center, ones + ) + max_candidates = int(xp.max(count_per_center)) + if max_candidates == 0: + return xp.full((ncenters, nsel), -1, dtype=xp.int64, device=device) + + edge_count = center.shape[0] + padded_size = ncenters * max_candidates + if padded_size > _PADDED_CANDIDATE_OVERHEAD_LIMIT * max(edge_count, 1): + return None + + center_start = xp.cumulative_sum(count_per_center) - count_per_center + edge_iota = _index_iota(center) + rank = edge_iota - xp.take(center_start, center, axis=0) + slot = center * max_candidates + rank + + neighbor_rows = xp_scatter_sum( + xp.zeros((padded_size,), dtype=xp.int64, device=device), + 0, + slot, + neighbor_ext + 1, + ) + neighbor_rows = xp.reshape(neighbor_rows - 1, (ncenters, max_candidates)) + distance_rows = xp_scatter_sum( + xp.zeros((padded_size,), dtype=distance.dtype, device=device), + 0, + slot, + distance, + ) + distance_rows = xp.reshape(distance_rows, (ncenters, max_candidates)) + distance_rows = xp.where( + neighbor_rows >= 0, + distance_rows, + xp.full_like(distance_rows, float("inf")), + ) + + # Stable sorting the secondary key first preserves the dense builder's + # extended-index tie break for equal distances. + order = xp.argsort(neighbor_rows, axis=1, stable=True) + neighbor_rows = xp_take_along_axis(neighbor_rows, order, axis=1) + distance_rows = xp_take_along_axis(distance_rows, order, axis=1) + order = xp.argsort(distance_rows, axis=1, stable=True) + neighbor_rows = xp_take_along_axis(neighbor_rows, order, axis=1) + + selected_width = min(nsel, max_candidates) + nlist = neighbor_rows[:, :selected_width] + if selected_width < nsel: + nlist = xp.concat( + ( + nlist, + xp.full( + (ncenters, nsel - selected_width), + -1, + dtype=xp.int64, + device=device, + ), + ), + axis=1, + ) + return nlist + + +def _supports_padded_selection(coord: Array) -> bool: + """Whether row padding can use eager, data-dependent Python dimensions.""" + if array_api_compat.is_numpy_array(coord): + return True + if array_api_compat.is_torch_array(coord): + import torch + + # torch.export/compile must keep the compact path: converting the maximum + # candidate count to ``int`` would specialize an unbacked symbolic value. + # Keep unmeasured accelerator implementations on the compact path too. + device = array_api_compat.device(coord) + is_compiling = getattr(getattr(torch, "compiler", None), "is_compiling", None) + # Older or reduced PyTorch builds may not expose torch.compiler. Without + # a reliable tracing-state query, keep the dynamic compact path instead + # of risking a data-dependent Python allocation during compilation. + # Eager CUDA pays one synchronization for the concrete row width, but + # measured RTX 5090 builds still favored the bounded row sorts. Keep + # that validated path while compilation and unmeasured devices fall back. + return ( + device.type in ("cpu", "cuda") + and is_compiling is not None + and not is_compiling() + ) + if array_api_compat.is_jax_array(coord): + import jax + + # Neighbor-list construction normally runs eagerly before the compiled + # model step; a tracer still needs the static compact fallback. On JAX + # accelerators, materializing the row width on the host costs more than + # the saved sort work, so retain compact device sorting. + device = array_api_compat.device(coord) + return not isinstance(coord, jax.core.Tracer) and device.platform == "cpu" + return False + + +def _build_neighbor_list_cell( + coord: Array, + atype: Array, + nloc: int, + rcut: float, + nsel: int, + pair_excl: "PairExcludeMask | None" = None, +) -> Array: + """Build a fixed-width neighbor list from Cartesian spatial cells. + + Extended coordinates already contain all required periodic images, so the + search itself is non-periodic: atoms are assigned to axis-aligned cells of + width ``rcut`` and each local center examines only its 27 adjacent cells. + Cell members are represented by sorted integer keys; ``searchsorted`` finds + the member ranges and an array-valued ``repeat`` expands only real candidate + pairs. At constant density this uses O(N) candidate memory and O(N log N) + work, instead of the dense O(N**2) distance matrix. Eager NumPy, eager + PyTorch on validated CPU/CUDA devices, and eager JAX on CPU select neighbors + with two bounded row-wise sorts. TensorFlow, traced PyTorch/JAX, and JAX on + accelerators retain compact global sorting to avoid a host-dependent shape. + + The final stable lexicographic ordering is ``(center, distance, ext_index)``. + It matches the dense builder's nearest-neighbor contract, including selecting + before applying pair exclusions (excluded entries leave holes, without + backfilling farther neighbors). + """ + xp = array_api_compat.array_namespace(coord, atype) + device = array_api_compat.device(coord) + nframes, nall = atype.shape + coord = xp.reshape(coord, (nframes, nall, 3)) + + if nloc == 0: + return xp.full((nframes, 0, nsel), -1, dtype=xp.int64, device=device) + + real_mask = atype >= 0 + coord_flat = xp.reshape(coord, (nframes * nall, 3)) + # NumPy and eager PyTorch gain from halving the candidate-index traffic. + # JAX and TensorFlow retain int64 internals because their measured compact + # selection/graph paths did not recover the added conversion cost. + use_int32_internal_indices = ( + ( + array_api_compat.is_numpy_array(coord) + or array_api_compat.is_torch_array(coord) + ) + and isinstance(nframes, int) + and isinstance(nall, int) + and nframes * nall <= _INT32_MAX + ) + internal_index_dtype = xp.int32 if use_int32_internal_indices else xp.int64 + is_unplaced_tensorflow = getattr( + xp, "__name__", "" + ) == "deepmd._vendors.ndtensorflow" and not str(device) + is_accelerator = ( + getattr(device, "type", None) == "cuda" + or getattr(device, "platform", None) == "gpu" + or "GPU" in str(device).upper() + # A traced TensorFlow tensor is commonly unplaced even when the graph + # later executes on GPU, where small/medium compaction regresses. + or is_unplaced_tensorflow + ) + compact_large_torch_cuda = ( + array_api_compat.is_torch_array(coord) + and getattr(device, "type", None) == "cuda" + and rcut <= 6.0 + and nloc >= _TORCH_CUDA_PERIODIC_COMPACTION_THRESHOLD + ) + compact_cpu = False + if not is_accelerator: + if array_api_compat.is_numpy_array(coord): + compact_cpu = nloc >= _NUMPY_CPU_PERIODIC_COMPACTION_THRESHOLD + elif array_api_compat.is_torch_array(coord): + compact_cpu = nloc >= _TORCH_CPU_PERIODIC_COMPACTION_THRESHOLD + elif array_api_compat.is_jax_array(coord): + compact_cpu = nloc >= _JAX_CPU_PERIODIC_COMPACTION_THRESHOLD + elif getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow": + compact_cpu = nloc >= _TF_CPU_PERIODIC_COMPACTION_THRESHOLD + # The extra mask/nonzero/gather launches outweigh the smaller key sort on + # JAX/TF accelerators and smaller PyTorch CUDA systems. CPU backends benefit + # immediately; eager PyTorch CUDA crosses over only at much larger sizes. + compact_periodic_images = ( + isinstance(nall, int) + and nall > 2 * nloc + and (compact_cpu or compact_large_torch_cuda) + ) + + if compact_periodic_images: + # Periodic extension normally contains every atom in every adjacent box. + # An image outside the local Cartesian bounds expanded by ``rcut`` can + # not neighbor any local center, so omit it from cell construction while + # retaining its original extended index. This reduces the usual 27*N + # periodic sort to the central atoms plus a thin boundary shell. + local_real_mask = real_mask[:, :nloc] + local_coord = coord[:, :nloc, :] + local_min = xp.min( + xp.where( + local_real_mask[..., None], + local_coord, + xp.full_like(local_coord, float("inf")), + ), + axis=1, + keepdims=True, + ) + local_max = xp.max( + xp.where( + local_real_mask[..., None], + local_coord, + xp.full_like(local_coord, float("-inf")), + ), + axis=1, + keepdims=True, + ) + has_local_real = xp.any(local_real_mask, axis=1, keepdims=True) + zero_origin = xp.zeros_like(local_min) + origin = xp.where(has_local_real[..., None], local_min - rcut, zero_origin) + upper = xp.where(has_local_real[..., None], local_max + rcut, zero_origin) + in_search_bounds = xp.all((coord >= origin) & (coord <= upper), axis=-1) + search_mask = real_mask & in_search_bounds + + # Keep one pinned placeholder for an all-virtual frame so reductions over + # the compact stream remain defined without a data-dependent Python branch. + first_in_frame = ( + xp.arange(nall, dtype=internal_index_dtype, device=device)[None, :] == 0 + ) + search_mask = search_mask | (first_in_frame & ~has_local_real) + (flat_selected,) = xp.nonzero(xp.reshape(search_mask, (-1,))) + flat_selected = xp.reshape(flat_selected, (-1,)) + if not isinstance(flat_selected.shape[0], int): + xp_hint_dynamic_size(flat_selected) + selected_frame = flat_selected // nall + selected_ext_index = flat_selected - selected_frame * nall + if use_int32_internal_indices: + selected_ext_index = xp.astype(selected_ext_index, xp.int32) + selected_coord = xp.take(coord_flat, flat_selected, axis=0) + selected_atype = xp.take(xp.reshape(atype, (-1,)), flat_selected, axis=0) + selected_origin = xp.take( + xp.reshape(origin, (nframes, 3)), selected_frame, axis=0 + ) + coord_for_binning = xp.where( + selected_atype[:, None] >= 0, selected_coord, selected_origin + ) + cell = xp.astype( + xp.floor((coord_for_binning - selected_origin) / rcut), xp.int64 + ) + center_coord_for_binning = xp.where( + local_real_mask[..., None], local_coord, origin + ) + center_cell = xp.astype( + xp.floor((center_coord_for_binning - origin) / rcut), xp.int64 + ) + dims = xp.max(cell, axis=0) + 1 + cells_per_frame = dims[0] * dims[1] * dims[2] + frame_base = xp.arange(nframes, dtype=xp.int64, device=device) * cells_per_frame + flat_key = cell[:, 0] + dims[0] * (cell[:, 1] + dims[1] * cell[:, 2]) + flat_key = flat_key + xp.take(frame_base, selected_frame, axis=0) + sentinel = cells_per_frame * nframes + flat_key = xp.where(selected_atype >= 0, flat_key, sentinel) + order = xp.argsort(flat_key, stable=True) + sorted_key = xp.take(flat_key, order, axis=0) + sorted_ext_index = xp.take(selected_ext_index, order, axis=0) + else: + # Virtual atoms can carry arbitrary placeholder coordinates. Pin them to + # the per-frame real-atom origin so they do not expand or overflow the grid. + real_coord = xp.where( + real_mask[..., None], + coord, + xp.full_like(coord, float("inf")), + ) + origin = xp.min(real_coord, axis=1, keepdims=True) + has_real = xp.any(real_mask, axis=1, keepdims=True) + origin = xp.where(has_real[..., None], origin, xp.zeros_like(origin)) + coord_for_binning = xp.where(real_mask[..., None], coord, origin) + cell = xp.astype(xp.floor((coord_for_binning - origin) / rcut), xp.int64) + center_cell = cell[:, :nloc, :] + dims = xp.max(xp.reshape(cell, (-1, 3)), axis=0) + 1 + cells_per_frame = dims[0] * dims[1] * dims[2] + frame_base = xp.arange(nframes, dtype=xp.int64, device=device) * cells_per_frame + cell_key = cell[..., 0] + dims[0] * (cell[..., 1] + dims[1] * cell[..., 2]) + cell_key = cell_key + frame_base[:, None] + sentinel = cells_per_frame * nframes + cell_key = xp.where(atype >= 0, cell_key, sentinel) + flat_key = xp.reshape(cell_key, (-1,)) + local_ext_index = xp.broadcast_to( + xp.arange(nall, dtype=internal_index_dtype, device=device)[None, :], + (nframes, nall), + ) + order = xp.argsort(flat_key, stable=True) + sorted_key = xp.take(flat_key, order, axis=0) + sorted_ext_index = xp.take(xp.reshape(local_ext_index, (-1,)), order, axis=0) + + # Query the 27 cells surrounding every local center. Out-of-grid queries and + # virtual centers get key -1; searchsorted then returns an empty interval. + # Keep this fixed 27x3 constant graph-local. Caching backend arrays would + # retain device storage and, for TensorFlow, could reuse a tensor from the + # wrong graph; the small materialization cost is included in the thresholds. + offsets = xp.asarray(_NEIGHBOR_CELL_OFFSETS, dtype=xp.int64, device=device) + query_cell = center_cell[:, :, None, :] + offsets[None, None, :, :] + in_bounds = xp.all( + (query_cell >= 0) & (query_cell < dims[None, None, None, :]), axis=-1 + ) + in_bounds = in_bounds & (atype[:, :nloc, None] >= 0) + query_key = query_cell[..., 0] + dims[0] * ( + query_cell[..., 1] + dims[1] * query_cell[..., 2] + ) + query_key = query_key + frame_base[:, None, None] + query_key = xp.where(in_bounds, query_key, xp.full_like(query_key, -1)) + query_key = xp.reshape(query_key, (-1,)) + + starts = xp.astype(xp.searchsorted(sorted_key, query_key, side="left"), xp.int64) + ends = xp.astype(xp.searchsorted(sorted_key, query_key, side="right"), xp.int64) + counts = ends - starts + + # Expand each cell query into its actual members. The candidate length is + # data-dependent; traced namespaces retain the cumulative-sum iota fallback + # because ``arange`` can not consume an unbacked symbolic size portably. + # Query IDs only range over ``27 * ncenters``. Use 32-bit values when that + # static range permits, then widen the much smaller filtered center stream + # before selection; this cuts several full-candidate index temporaries in half. + query_count = query_key.shape[0] + use_int32_query_ids = use_int32_internal_indices + query_index_dtype = ( + xp.int32 + if use_int32_query_ids + and isinstance(query_count, int) + and query_count <= _INT32_MAX + else xp.int64 + ) + query_ids = xp.repeat( + xp.arange(query_count, dtype=query_index_dtype, device=device), counts + ) + if not isinstance(query_ids.shape[0], int): + xp_hint_dynamic_size(query_ids) + query_output_start = xp.cumulative_sum(counts) - counts + candidate_count = query_ids.shape[0] + use_int32_candidate_positions = ( + use_int32_query_ids + and isinstance(candidate_count, int) + and candidate_count <= _INT32_MAX + and isinstance(nall, int) + and nframes * nall <= _INT32_MAX + ) + if use_int32_candidate_positions: + query_output_start = xp.astype(query_output_start, xp.int32) + candidate_iota = xp.arange(candidate_count, dtype=xp.int32, device=device) + candidate_starts = xp.astype(starts, xp.int32) + else: + candidate_iota = _index_iota(query_ids) + candidate_starts = starts + member_offset = candidate_iota - xp.take(query_output_start, query_ids, axis=0) + sorted_position = xp.take(candidate_starts, query_ids, axis=0) + member_offset + neighbor_ext = xp.take(sorted_ext_index, sorted_position, axis=0) + + center = query_ids // len(_NEIGHBOR_CELL_OFFSETS) + frame = center // nloc + center_local = center - frame * nloc + center_coord = xp.take(coord_flat, frame * nall + center_local, axis=0) + neighbor_coord = xp.take(coord_flat, frame * nall + neighbor_ext, axis=0) + diff = neighbor_coord - center_coord + # Use the same norm as the historical dense implementation for both cutoff + # and ordering. Sorting squared distances can distinguish values that round + # to the same norm, changing the stable order of symmetry-equivalent images. + distance = xp.linalg.vector_norm(diff, axis=-1) + keep_mask = (neighbor_ext != center_local) & (distance <= rcut) + (keep,) = xp.nonzero(keep_mask) + keep = xp.reshape(keep, (-1,)) + if not isinstance(keep.shape[0], int): + xp_hint_dynamic_size(keep) + center = xp.take(center, keep, axis=0) + neighbor_ext = xp.take(neighbor_ext, keep, axis=0) + if use_int32_internal_indices: + center = xp.astype(center, xp.int64) + neighbor_ext = xp.astype(neighbor_ext, xp.int64) + distance = xp.take(distance, keep, axis=0) + + ncenters = nframes * nloc + if _supports_padded_selection(coord): + # repeat(arange(query_count), counts) is non-decreasing; integer division + # and the ordered ``keep`` gather above preserve that property for center. + nlist = _select_nearest_padded(center, neighbor_ext, distance, ncenters, nsel) + if nlist is not None: + nlist = xp.reshape(nlist, (nframes, nloc, nsel)) + return apply_pair_exclusion_nlist(nlist, atype, pair_excl) + + # Stable sorts from the least- to most-significant key produce the desired + # lexicographic order while using only standard Array API operations. + order = xp.argsort(neighbor_ext, stable=True) + order = xp.take( + order, + xp.argsort(xp.take(distance, order, axis=0), stable=True), + axis=0, + ) + order = xp.take( + order, + xp.argsort(xp.take(center, order, axis=0), stable=True), + axis=0, + ) + center = xp.take(center, order, axis=0) + neighbor_ext = xp.take(neighbor_ext, order, axis=0) + + ones = xp.ones_like(center) + count_per_center = xp_scatter_sum( + xp.zeros((ncenters,), dtype=xp.int64, device=device), 0, center, ones + ) + center_start = xp.cumulative_sum(count_per_center) - count_per_center + edge_iota = _index_iota(center) + rank = edge_iota - xp.take(center_start, center, axis=0) + (selected,) = xp.nonzero(rank < nsel) + selected = xp.reshape(selected, (-1,)) + if not isinstance(selected.shape[0], int): + xp_hint_dynamic_size(selected) + selected_center = xp.take(center, selected, axis=0) + selected_rank = xp.take(rank, selected, axis=0) + selected_neighbor = xp.take(neighbor_ext, selected, axis=0) + + # Each (center, rank) slot is unique, so scatter-add into a zero array and + # store ``neighbor + 1``; subtracting one afterwards creates -1 padding. + slot = selected_center * nsel + selected_rank + nlist = xp_scatter_sum( + xp.zeros((ncenters * nsel,), dtype=xp.int64, device=device), + 0, + slot, + selected_neighbor + 1, + ) + nlist = xp.reshape(nlist - 1, (nframes, nloc, nsel)) + return apply_pair_exclusion_nlist(nlist, atype, pair_excl) + + class DefaultNeighborList(NeighborList): - """All-pairs builder: replicate the cell into periodic images and rank by - distance (:func:`~deepmd.dpmodel.utils.nlist.extend_coord_with_ghosts` + - :func:`~deepmd.dpmodel.utils.nlist.build_neighbor_list`). This is the - default when no strategy is supplied, so results are unchanged. + """Adaptive builder using dense search for small systems and spatial cells + for larger systems. + + Both paths first replicate the cell into periodic images with + :func:`~deepmd.dpmodel.utils.nlist.extend_coord_with_ghosts`. Small systems + retain the historical all-pairs distance matrix for its lower constant + overhead. Larger supported arrays use a Cartesian cell list whose candidate + count is linear at constant density. JAX builds this data-dependent list + eagerly before the fixed-shape model computation enters ``jit``. """ def build( @@ -67,8 +703,9 @@ def build( Must be ``"extended"`` (the only mode this builder supports). pair_excl : PairExcludeMask or None, optional When provided, excluded type pairs are erased from the returned - neighbor list immediately after the geometric search by - :func:`~deepmd.dpmodel.utils.nlist.build_neighbor_list`. + neighbor list after the dense or cell-list geometric search selects + neighbors. Excluded entries leave holes without backfilling farther + atoms. Returns ------- @@ -92,16 +729,27 @@ def build( extended_coord, extended_atype, mapping = extend_coord_with_ghosts( coord_normalized, atype, box, rcut ) - # types are distinguished in the lower interface, so keep them merged here; - # pair_excl is forwarded so exclusion is applied at build time. - nlist = build_neighbor_list( - extended_coord, - extended_atype, - nloc, - rcut, - sel, - distinguish_types=False, - pair_excl=pair_excl, - ) + # Types are distinguished in the lower interface, so keep them merged + # here. The dense path remains faster for small systems; the spatial + # path avoids the nloc*nall distance matrix for larger supported arrays. + if _supports_cell_list(coord, nloc, periodic=box is not None): + nlist = _build_neighbor_list_cell( + extended_coord, + extended_atype, + nloc, + rcut, + sum(sel), + pair_excl=pair_excl, + ) + else: + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + sel, + distinguish_types=False, + pair_excl=pair_excl, + ) extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) return extended_coord, extended_atype, nlist, mapping diff --git a/source/tests/common/dpmodel/test_default_neighbor_list.py b/source/tests/common/dpmodel/test_default_neighbor_list.py new file mode 100644 index 0000000000..ffaa8d4afa --- /dev/null +++ b/source/tests/common/dpmodel/test_default_neighbor_list.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the adaptive Array API default neighbor-list builder.""" + +from typing import ( + Any, +) + +import numpy as np +import pytest + +import deepmd.dpmodel.utils.default_neighbor_list as default_nlist +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) + + +def _force_search(monkeypatch: pytest.MonkeyPatch, *, cell: bool) -> None: + """Force one public search path without changing the builder API.""" + threshold = 0 if cell else 10**9 + for name in ( + "_NUMPY_CPU_PERIODIC_CELL_LIST_THRESHOLD", + "_NUMPY_CPU_NONPERIODIC_CELL_LIST_THRESHOLD", + ): + monkeypatch.setattr(default_nlist, name, threshold) + + +def _random_system( + *, nframes: int = 2, nloc: int = 64, dtype: Any = np.float64 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Create reproducible triclinic frames without distance degeneracies.""" + rng = np.random.default_rng(20260721) + box = np.asarray( + [ + [[9.0, 0.0, 0.0], [0.8, 8.5, 0.0], [0.3, 0.6, 9.5]], + [[8.5, 0.0, 0.0], [-0.5, 9.2, 0.0], [0.4, -0.2, 8.8]], + ][:nframes], + dtype=dtype, + ) + fractional = rng.random((nframes, nloc, 3), dtype=dtype) + coord = np.matmul(fractional, box) + atype = rng.integers(0, 2, size=(nframes, nloc), dtype=np.int64) + return coord, atype, box + + +def _build( + monkeypatch: pytest.MonkeyPatch, + *, + cell: bool, + coord: Any, + atype: Any, + box: Any, + rcut: float, + sel: list[int], + pair_excl: PairExcludeMask | None = None, +) -> tuple[Any, Any, Any, Any]: + """Build with a forced dense or cell-list implementation.""" + _force_search(monkeypatch, cell=cell) + return default_nlist.DefaultNeighborList().build( + coord, atype, box, rcut, sel, pair_excl=pair_excl + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("periodic", [False, True]) +def test_cell_list_matches_dense( + monkeypatch: pytest.MonkeyPatch, periodic: bool, dtype: Any +) -> None: + """The sparse candidate search preserves the dense quartet exactly.""" + coord, atype, box = _random_system(dtype=dtype) + box_arg = box if periodic else None + dense = _build( + monkeypatch, + cell=False, + coord=coord, + atype=atype, + box=box_arg, + rcut=3.2, + sel=[24, 24], + ) + cell = _build( + monkeypatch, + cell=True, + coord=coord, + atype=atype, + box=box_arg, + rcut=3.2, + sel=[24, 24], + ) + for dense_value, cell_value in zip(dense, cell, strict=True): + np.testing.assert_allclose(dense_value, cell_value) + + +def test_cell_list_exact_cutoff_virtual_and_exclusion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Boundary neighbors, virtual atoms, and exclusion holes match dense search.""" + coord = np.asarray( + [ + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [2.0, 0.0, 0.0], + [1.0e12, -1.0e12, 1.0e12], + ] + ], + dtype=np.float64, + ) + atype = np.asarray([[0, 1, 0, -1]], dtype=np.int64) + pair_excl = PairExcludeMask(2, [(0, 1)]) + dense = _build( + monkeypatch, + cell=False, + coord=coord, + atype=atype, + box=None, + rcut=1.0, + sel=[4, 4], + pair_excl=pair_excl, + ) + cell = _build( + monkeypatch, + cell=True, + coord=coord, + atype=atype, + box=None, + rcut=1.0, + sel=[4, 4], + pair_excl=pair_excl, + ) + for dense_value, cell_value in zip(dense, cell, strict=True): + np.testing.assert_array_equal(dense_value, cell_value) + + +def test_cell_list_preserves_equal_distance_image_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Stable ordering of symmetry-equivalent periodic images matches dense search.""" + coord = np.asarray([[[0.1, 0.2, 0.3], [1.1, 1.2, 1.3]]], dtype=np.float64) + atype = np.asarray([[0, 0]], dtype=np.int64) + box = 2.0 * np.eye(3, dtype=np.float64)[None, :, :] + dense = _build( + monkeypatch, + cell=False, + coord=coord, + atype=atype, + box=box, + rcut=2.1, + sel=[100], + ) + cell = _build( + monkeypatch, + cell=True, + coord=coord, + atype=atype, + box=box, + rcut=2.1, + sel=[100], + ) + for dense_value, cell_value in zip(dense, cell, strict=True): + np.testing.assert_array_equal(dense_value, cell_value) + + +def test_padded_selection_orders_rows_and_pads() -> None: + """Row-wise selection preserves distance/index order and empty slots.""" + center = np.asarray([0, 0, 0, 1, 1], dtype=np.int64) + neighbor = np.asarray([3, 2, 1, 4, 0], dtype=np.int64) + distance = np.asarray([1.0, 1.0, 0.5, 2.0, 1.0], dtype=np.float64) + result = default_nlist._select_nearest_padded( + center, neighbor, distance, ncenters=3, nsel=4 + ) + np.testing.assert_array_equal( + result, + np.asarray( + [[1, 2, 3, -1], [0, 4, -1, -1], [-1, -1, -1, -1]], + dtype=np.int64, + ), + ) + + +def test_padded_selection_rejects_excessive_imbalance() -> None: + """A single wide row falls back before allocating a mostly empty matrix.""" + center = np.asarray([0] * 9 + [1], dtype=np.int64) + neighbor = np.arange(10, dtype=np.int64) + distance = np.arange(10, dtype=np.float64) + assert ( + default_nlist._select_nearest_padded( + center, neighbor, distance, ncenters=10, nsel=4 + ) + is None + ) + + +def test_automatic_cpu_thresholds() -> None: + """Measured CPU crossovers keep small systems on the dense fast path.""" + assert not default_nlist._supports_cell_list( + np.zeros((1, 31, 3)), + 31, + periodic=True, + ) + assert default_nlist._supports_cell_list( + np.zeros((1, 32, 3)), + 32, + periodic=True, + ) + assert not default_nlist._supports_cell_list( + np.zeros((1, 255, 3)), + 255, + periodic=False, + ) + assert default_nlist._supports_cell_list( + np.zeros((1, 256, 3)), + 256, + periodic=False, + ) diff --git a/source/tests/consistent/test_default_neighbor_list.py b/source/tests/consistent/test_default_neighbor_list.py new file mode 100644 index 0000000000..aa631e2b84 --- /dev/null +++ b/source/tests/consistent/test_default_neighbor_list.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-gated consistency tests for the default neighbor-list builder.""" + +import unittest +import unittest.mock as mock +from typing import ( + Any, +) + +import array_api_compat +import numpy as np + +import deepmd.dpmodel.utils.default_neighbor_list as default_nlist +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) + +from .common import ( + INSTALLED_ARRAY_API_STRICT, + INSTALLED_JAX, + INSTALLED_PT, + INSTALLED_TF2, +) + +if INSTALLED_ARRAY_API_STRICT: + import array_api_strict as strict + +if INSTALLED_PT: + import torch + +if INSTALLED_JAX: + import jax + import jax.numpy as jnp + +if INSTALLED_TF2: + import tensorflow as tf + + from deepmd._vendors import ndtensorflow as ndtf + + +def _random_system( + *, nloc: int, dtype: Any = np.float64 +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Create one reproducible triclinic frame for backend comparisons.""" + rng = np.random.default_rng(20260721) + box = np.asarray( + [[[9.0, 0.0, 0.0], [0.8, 8.5, 0.0], [0.3, 0.6, 9.5]]], + dtype=dtype, + ) + coord = np.matmul(rng.random((1, nloc, 3), dtype=dtype), box) + atype = rng.integers(0, 2, size=(1, nloc), dtype=np.int64) + return coord, atype, box + + +class _CellListBackendMixin: + """Run one shared geometric/exclusion case through each Array API backend.""" + + def _backend_devices(self) -> list[Any]: + """Return validated devices, including an accelerator when available.""" + return [None] + + def _backend_arrays( + self, coord: np.ndarray, atype: np.ndarray, device: Any + ) -> tuple[Any, Any]: + """Convert the shared NumPy case to one backend device under test.""" + raise NotImplementedError + + def _to_numpy(self, value: Any) -> np.ndarray: + """Synchronize and convert one backend result for exact comparison.""" + return np.asarray(value) + + def test_backend_cell_list_matches_dense(self) -> None: + """All backends preserve common periodic and nonperiodic references.""" + nloc = 24 + rcut = 3.2 + nsel = 32 + coord, atype, box = _random_system(nloc=nloc, dtype=np.float32) + pair_excl = PairExcludeMask(2, [(0, 1)]) + for periodic in (False, True): + if periodic: + search_coord, search_atype, _ = default_nlist.extend_coord_with_ghosts( + coord, + atype, + box, + rcut, + ) + else: + search_coord = coord.reshape(1, -1) + search_atype = atype + reference = default_nlist.build_neighbor_list( + search_coord, + search_atype, + nloc, + rcut, + [nsel], + distinguish_types=False, + pair_excl=pair_excl, + ) + for device in self._backend_devices(): + with self.subTest(periodic=periodic, device=str(device)): + coord_backend, atype_backend = self._backend_arrays( + search_coord, search_atype, device + ) + # Exercise the periodic boundary-shell compaction route on + # every backend/device where it is supported. Production + # thresholds are benchmark-derived and otherwise exceed the + # deliberately small cross-backend fixture. + with mock.patch.multiple( + default_nlist, + _NUMPY_CPU_PERIODIC_COMPACTION_THRESHOLD=0, + _TORCH_CPU_PERIODIC_COMPACTION_THRESHOLD=0, + _JAX_CPU_PERIODIC_COMPACTION_THRESHOLD=0, + _TF_CPU_PERIODIC_COMPACTION_THRESHOLD=0, + _TORCH_CUDA_PERIODIC_COMPACTION_THRESHOLD=0, + ): + result = default_nlist._build_neighbor_list_cell( + coord_backend, + atype_backend, + nloc, + rcut, + nsel, + pair_excl=pair_excl, + ) + self.assertTrue(array_api_compat.is_array_api_obj(result)) + np.testing.assert_array_equal(self._to_numpy(result), reference) + + def test_backend_virtual_outlier_does_not_expand_grid(self) -> None: + """Virtual placeholder coordinates do not affect real-atom cells.""" + coord = np.asarray( + [ + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [2.0, 0.0, 0.0], + [1.0e12, -1.0e12, 1.0e12], + ], + [ + [1.0e12, 0.0, 0.0], + [0.0, -1.0e12, 0.0], + [0.0, 0.0, 1.0e12], + [-1.0e12, 0.0, 0.0], + ], + ], + dtype=np.float32, + ) + atype = np.asarray([[0, 1, 0, -1], [-1, -1, -1, -1]], dtype=np.int64) + reference = default_nlist.build_neighbor_list( + coord, + atype, + nloc=4, + rcut=1.0, + sel=[8], + distinguish_types=False, + ) + for device in self._backend_devices(): + with self.subTest(device=str(device)): + coord_backend, atype_backend = self._backend_arrays( + coord, atype, device + ) + result = default_nlist._build_neighbor_list_cell( + coord_backend, + atype_backend, + nloc=4, + rcut=1.0, + nsel=8, + ) + np.testing.assert_array_equal(self._to_numpy(result), reference) + + +@unittest.skipUnless(INSTALLED_ARRAY_API_STRICT, "array_api_strict is not installed") +class TestArrayAPIStrictDefaultNeighborList(_CellListBackendMixin, unittest.TestCase): + """Validate strict Array API execution and conservative public dispatch.""" + + def _backend_arrays( + self, coord: np.ndarray, atype: np.ndarray, device: Any + ) -> tuple[Any, Any]: + return strict.asarray(coord), strict.asarray(atype) + + def test_unknown_array_namespace_uses_dense_path(self) -> None: + """Unbenchmarked Array API namespaces do not inherit another threshold.""" + coord = strict.zeros((1, 1, 3), dtype=strict.float32) + self.assertFalse( + default_nlist._supports_cell_list(coord, 10**6, periodic=False) + ) + + +@unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") +class TestTorchDefaultNeighborList(_CellListBackendMixin, unittest.TestCase): + """Validate PyTorch gradients, dispatch thresholds, and supported devices.""" + + def _backend_devices(self) -> list[Any]: + devices = [torch.device("cpu")] + if torch.cuda.is_available(): + devices.append(torch.device("cuda")) + return devices + + def _backend_arrays( + self, coord: np.ndarray, atype: np.ndarray, device: Any + ) -> tuple[Any, Any]: + return torch.as_tensor(coord, device=device), torch.as_tensor( + atype, device=device + ) + + def _to_numpy(self, value: Any) -> np.ndarray: + return value.detach().cpu().numpy() + + def test_torch_namespace_and_gradient(self) -> None: + coord, atype, box = _random_system(nloc=32) + coord_t = torch.tensor(coord, dtype=torch.float64, requires_grad=True) + atype_t = torch.tensor(atype, dtype=torch.int64) + box_t = torch.tensor(box, dtype=torch.float64) + + with mock.patch.multiple( + default_nlist, + _TORCH_CPU_PERIODIC_CELL_LIST_THRESHOLD=0, + ): + result = default_nlist.DefaultNeighborList().build( + coord_t, atype_t, box_t, 3.2, [24, 24] + ) + self.assertTrue(all(isinstance(value, torch.Tensor) for value in result)) + result[0].sum().backward() + self.assertIsNotNone(coord_t.grad) + + def test_torch_cpu_threshold(self) -> None: + coord = torch.zeros((1, 1, 3), dtype=torch.float64) + self.assertFalse(default_nlist._supports_cell_list(coord, 2047, periodic=False)) + self.assertTrue(default_nlist._supports_cell_list(coord, 2048, periodic=False)) + + def test_torch_unvalidated_device_uses_dense_path(self) -> None: + """Non-CPU/CUDA torch devices retain the conservative dense path.""" + coord = torch.zeros((1, 1, 3), dtype=torch.float64) + with mock.patch.object( + default_nlist.array_api_compat, + "device", + return_value=mock.Mock(type="mps"), + ): + self.assertFalse( + default_nlist._supports_cell_list(coord, 10**6, periodic=False) + ) + + def test_torch_without_compiler_api_uses_compact_selection(self) -> None: + """PyTorch variants without compiler state avoid dynamic row padding.""" + coord = torch.zeros((1, 1, 3), dtype=torch.float64) + with mock.patch.object(torch, "compiler", None): + self.assertFalse(default_nlist._supports_padded_selection(coord)) + + +@unittest.skipUnless(INSTALLED_JAX, "JAX is not installed") +class TestJAXDefaultNeighborList(_CellListBackendMixin, unittest.TestCase): + """Validate eager JAX construction and traced/unknown-device fallbacks.""" + + def _backend_devices(self) -> list[Any]: + devices = [jax.devices("cpu")[0]] + try: + gpu_devices = jax.devices("gpu") + except RuntimeError: + gpu_devices = [] + if gpu_devices: + devices.append(gpu_devices[0]) + return devices + + def _backend_arrays( + self, coord: np.ndarray, atype: np.ndarray, device: Any + ) -> tuple[Any, Any]: + return jax.device_put(coord, device), jax.device_put(atype, device) + + def test_jax_dispatch_and_tracing(self) -> None: + coord, atype, _ = _random_system(nloc=256, dtype=np.float32) + coord_jax, atype_jax = self._backend_arrays(coord, atype, jax.devices("cpu")[0]) + + result = default_nlist.DefaultNeighborList().build( + coord_jax, + atype_jax, + None, + 3.2, + [32], + )[2] + result.block_until_ready() + self.assertEqual(result.shape, (1, 256, 32)) + self.assertFalse( + default_nlist._supports_cell_list(coord_jax, 255, periodic=False) + ) + self.assertTrue( + default_nlist._supports_cell_list(coord_jax, 256, periodic=False) + ) + self.assertTrue(default_nlist._supports_padded_selection(coord_jax)) + traced_support = jax.jit( + lambda cc: jnp.asarray( + default_nlist._supports_cell_list(cc, 256, periodic=False) + ) + )(coord_jax) + self.assertFalse(bool(np.asarray(traced_support))) + traced_padded_support = jax.jit( + lambda cc: jnp.asarray(default_nlist._supports_padded_selection(cc)) + )(coord_jax) + self.assertFalse(bool(np.asarray(traced_padded_support))) + + def test_jax_unvalidated_device_uses_dense_path(self) -> None: + coord, _, _ = _random_system(nloc=1, dtype=np.float32) + coord_jax, _ = self._backend_arrays( + coord, + np.zeros((1, 1), dtype=np.int64), + jax.devices("cpu")[0], + ) + with mock.patch.object( + default_nlist.array_api_compat, + "device", + return_value=mock.Mock(platform="tpu"), + ): + self.assertFalse( + default_nlist._supports_cell_list(coord_jax, 10**6, periodic=False) + ) + + +@unittest.skipUnless(INSTALLED_TF2, "TF2 backend is not installed") +class TestTF2DefaultNeighborList(_CellListBackendMixin, unittest.TestCase): + """Validate the Array API neighbor search in the opt-in TF2 test job.""" + + def _backend_devices(self) -> list[Any]: + devices = ["/CPU:0"] + if tf.config.list_logical_devices("GPU"): + devices.append("/GPU:0") + return devices + + def _backend_arrays( + self, coord: np.ndarray, atype: np.ndarray, device: Any + ) -> tuple[Any, Any]: + with tf.device(device): + return ndtf.asarray(tf.convert_to_tensor(coord)), ndtf.asarray( + tf.convert_to_tensor(atype) + ) + + def _to_numpy(self, value: Any) -> np.ndarray: + return value.unwrap().numpy() + + def test_tf2_cpu_threshold(self) -> None: + """TensorFlow keeps sub-4096-atom CPU systems on dense search.""" + with tf.device("/CPU:0"): + coord = ndtf.asarray(tf.zeros((1, 1, 3), dtype=tf.float64)) + self.assertFalse(default_nlist._supports_cell_list(coord, 4095, periodic=False)) + self.assertTrue(default_nlist._supports_cell_list(coord, 4096, periodic=False)) + + def test_tf2_unvalidated_device_uses_dense_path(self) -> None: + with tf.device("/CPU:0"): + coord = ndtf.asarray(tf.zeros((1, 1, 3), dtype=tf.float32)) + with mock.patch.object( + default_nlist.array_api_compat, + "device", + return_value="TPU:0", + ): + self.assertFalse( + default_nlist._supports_cell_list(coord, 10**6, periodic=False) + ) + + def test_tf2_unplaced_device_uses_later_threshold(self) -> None: + """An unplaced TF graph is conservative across validated CPU/GPU paths.""" + with tf.device("/CPU:0"): + coord = ndtf.asarray(tf.zeros((1, 1, 3), dtype=tf.float32)) + with ( + mock.patch.object( + default_nlist.array_api_compat, + "device", + return_value="", + ), + mock.patch.multiple( + default_nlist, + _TF_CPU_NONPERIODIC_CELL_LIST_THRESHOLD=1024, + _TF_CUDA_NONPERIODIC_CELL_LIST_THRESHOLD=2048, + ), + ): + self.assertFalse( + default_nlist._supports_cell_list(coord, 2047, periodic=False) + ) + self.assertTrue( + default_nlist._supports_cell_list(coord, 2048, periodic=False) + ) + + def test_tf2_function_periodic_cell_list_matches_dense(self) -> None: + """A dynamic-batch TF graph preserves the dense periodic result.""" + coord, atype, box = _random_system(nloc=24) + + # Force opposite public paths so the comparison does not depend on + # benchmark-derived dispatch thresholds or the available TF2 device. + with mock.patch.multiple( + default_nlist, + _NUMPY_CPU_PERIODIC_CELL_LIST_THRESHOLD=10**9, + _TF_CPU_PERIODIC_CELL_LIST_THRESHOLD=0, + _TF_CUDA_PERIODIC_CELL_LIST_THRESHOLD=0, + ): + dense = default_nlist.DefaultNeighborList().build( + coord, atype, box, 3.2, [24, 24] + ) + + # The TF1 consistency suite can leave legacy Dimension objects + # enabled process-wide. Reproduce that shape representation here, + # then restore the TF2 job's original setting after graph tracing. + v2_shapes_enabled = tf.TensorSpec([None], tf.float64).shape[0] is None + tf.compat.v1.disable_v2_tensorshape() + try: + + @tf.function( + autograph=False, + input_signature=[ + tf.TensorSpec([None, 24, 3], tf.float64), + tf.TensorSpec([None, 24], tf.int64), + tf.TensorSpec([None, 3, 3], tf.float64), + ], + ) + def build_graph(cc: Any, aa: Any, bb: Any) -> tuple[Any, Any, Any, Any]: + result = default_nlist.DefaultNeighborList().build( + ndtf.asarray(cc), + ndtf.asarray(aa), + ndtf.asarray(bb), + 3.2, + [24, 24], + ) + return tuple(value.unwrap() for value in result) + + cell = build_graph(coord, atype, box) + finally: + if v2_shapes_enabled: + tf.compat.v1.enable_v2_tensorshape() + + for dense_value, cell_value in zip(dense, cell, strict=True): + np.testing.assert_allclose(dense_value, cell_value.numpy()) + + +if __name__ == "__main__": + unittest.main()