deepmd.pt.utils.nv_nlist

deepmd.pt.utils.nv_nlist#

Toolkit-Ops (nvalchemiops) neighbor-list strategy.

A NeighborList implementation that builds the extended representation (extended_coord, extended_atype, nlist, mapping) using the device-resident neighbor-list kernels in nvalchemiops.

Toolkit-Ops returns a dense [total_atoms, max_neighbors] neighbor matrix over the flattened batch. The matrix is converted to the DeePMD extended-atom contract by materializing each unique ghost (frame, src_local, shift) once; the candidate list is then distance-sorted and truncated to sum(sel) so the returned neighbor count is fixed. The search is non-differentiable and runs on detached coordinates, while extended_coord is rebuilt from the input coordinates so force and virial gradients propagate unchanged.

Attributes#

`NV_CELL_LIST_THRESHOLD`
`NV_NONPERIODIC_CELL_LIST_THRESHOLD`
`NV_CPU_CELL_LIST_THRESHOLD`
`log`
`_truncate_to_sel_compiled`

Classes#

NvNeighborList

Neighbor-list strategy using the nvalchemiops kernels.

Functions#

`_suppress_native_stderr`(→ collections.abc.Iterator[None])	Redirect the process `stderr` file descriptor to `os.devnull`.
`is_nv_available`(→ bool)	Whether the `nvalchemiops` Toolkit-Ops neighbor list is importable.
`choose_nv_nlist_method`(→ str)	Choose the Toolkit-Ops neighbor method for a homogeneous batch.
`_input_device_context`(→ collections.abc.Iterator[None])	Run third-party kernels with both default and current devices pinned.
`_grow_search_capacity`(→ int)	Increase Toolkit-Ops capacity by 1.25x, rounded up.
`_truncate_to_sel`(→ torch.Tensor)	Distance-sort the candidate neighbor list and keep the nearest `nsel`
`_matrix_to_extended_inputs`(→ tuple[torch.Tensor, ...)	Convert Toolkit-Ops matrix output to compact extended inputs.

Module Contents#

deepmd.pt.utils.nv_nlist.NV_CELL_LIST_THRESHOLD = 1024[source]#

deepmd.pt.utils.nv_nlist.NV_NONPERIODIC_CELL_LIST_THRESHOLD = 4096[source]#

deepmd.pt.utils.nv_nlist.NV_CPU_CELL_LIST_THRESHOLD = 128[source]#

deepmd.pt.utils.nv_nlist.log[source]#

deepmd.pt.utils.nv_nlist._suppress_native_stderr() → collections.abc.Iterator[None][source]#

Redirect the process stderr file descriptor to os.devnull.

nvalchemiops initializes NVIDIA Warp on first import, which probes for a CUDA driver and prints a native Warp CUDA error 100 line straight to the stderr fd on CPU-only hosts. That line bypasses Python logging, so the only way to mute it is at the descriptor level around the triggering import.

deepmd.pt.utils.nv_nlist.is_nv_available() → bool[source]#: Whether the nvalchemiops Toolkit-Ops neighbor list is importable.

deepmd.pt.utils.nv_nlist.choose_nv_nlist_method(nloc: int, *, periodic: bool = True, device: torch.device | None = None) → str[source]#

Choose the Toolkit-Ops neighbor method for a homogeneous batch.

Parameters:

nloc: Number of local atoms per frame.
periodic: Whether the batch is periodic.
device: Target device. CPU uses a lower cell-list threshold than CUDA because the batch_naive method does not parallelize well there.

Returns:

str: Toolkit-Ops method name.

deepmd.pt.utils.nv_nlist._input_device_context(device: torch.device) → collections.abc.Iterator[None][source]#: Run third-party kernels with both default and current devices pinned.

class deepmd.pt.utils.nv_nlist.NvNeighborList[source]#

Bases: deepmd.dpmodel.utils.neighbor_list.NeighborList

Neighbor-list strategy using the nvalchemiops kernels.

Implements the NeighborList interface on torch tensors; the search runs on the device of the input coordinates. Periodic inputs materialize shifted ghost atoms; non-periodic inputs keep only local atoms.

build(coord: Any, atype: Any, box: Any, rcut: float, sel: list[int], return_mode: str = 'extended') → tuple[Any, Any, Any, Any] | deepmd.dpmodel.utils.neighbor_list.EdgeNeighborList[source]#

Build the extended system and neighbor list.

See deepmd.dpmodel.utils.neighbor_list.NeighborList.build(). The returned nlist is distance-sorted and truncated to sum(sel).

deepmd.pt.utils.nv_nlist._grow_search_capacity(capacity: int) → int[source]#: Increase Toolkit-Ops capacity by 1.25x, rounded up.

deepmd.pt.utils.nv_nlist._truncate_to_sel(extended_coord: torch.Tensor, nlist: torch.Tensor, nsel: int, rcut: float) → torch.Tensor[source]#

Distance-sort the candidate neighbor list and keep the nearest nsel within rcut, padding with -1 when fewer neighbors exist.

The Toolkit-Ops search capacity may exceed sum(sel) on dense systems; this fixes the returned neighbor count at nsel.

The output is the integer nlist; extended_coord is only read to rank candidates and is returned unchanged by the caller. The routine is therefore non-differentiable and runs under no_grad so it never participates in the autograd graph (forward, backward, or the second-order pass used to train forces), which also avoids retaining the distance temporaries for backward.

deepmd.pt.utils.nv_nlist._truncate_to_sel_compiled[source]#

deepmd.pt.utils.nv_nlist._matrix_to_extended_inputs(*, coord: torch.Tensor, atype: torch.Tensor, cell: torch.Tensor | None, nloc: int, neighbor_matrix: torch.Tensor, num_neighbors: torch.Tensor, shifts: torch.Tensor) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor][source]#

Convert Toolkit-Ops matrix output to compact extended inputs.

Toolkit-Ops returns neighbors as a dense matrix over flattened atoms: neighbor_matrix[dst_global, slot] = src_global and shifts[dst_global, slot] = (sx, sy, sz). Here dst_global and src_global are indices in the concatenated nf * nloc input.

DeePMD lower paths use a different contract: nlist stores indices into extended_coord. Local atoms occupy [0, nloc) in each frame, while shifted PBC images must be appended as ghost atoms. This conversion builds the minimal ghost set by materializing each unique (frame, src_local, shift) once, then redirects all shifted nlist entries to the corresponding compact ghost index.