deepmd.dpmodel.utils.neighbor_graph.ase_builder

deepmd.dpmodel.utils.neighbor_graph.ase_builder#

Carry-all NeighborGraph builder backed by ASE’s O(N) cell list (optional dep).

build_neighbor_graph_ase is a carry-all search backend: it uses ASE’s neighbor_list("ijS", ...) to enumerate EVERY neighbor within rcut (no sel cutoff), then routes the resulting sparse (i, j, S) edge list through neighbor_graph_from_ijs() so edge_vec is recomputed differentiably from coord/box – ASE’s own distance vectors are intentionally NOT used, to keep the geometry convention and autograd leaf consistent with every other builder. ASE is an OPTIONAL dependency, imported lazily inside the function.

Functions#

build_neighbor_graph_ase(...)

Build a CARRY-ALL NeighborGraph using ASE's O(N) cell-list search.

Module Contents#

deepmd.dpmodel.utils.neighbor_graph.ase_builder.build_neighbor_graph_ase(coord: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, box: deepmd.dpmodel.array_api.Array | None, rcut: float, layout: deepmd.dpmodel.utils.neighbor_graph.graph.GraphLayout | None = None, *, with_csr: bool = False, canonicalize: bool = False, pair_excl: deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask | None = None, compact: bool = False) → deepmd.dpmodel.utils.neighbor_graph.graph.NeighborGraph[source]#

Build a CARRY-ALL NeighborGraph using ASE’s O(N) cell-list search.

Per frame, ASE neighbor_list("ijS", atoms, rcut) returns center i, neighbor j and periodic shift S such that the neighbor image sits at positions[j] + S @ cell. These map directly to the graph convention (src=neighbor=j, dst=center=i), and the edge list is fed to neighbor_graph_from_ijs() which recomputes edge_vec from coord/box (ASE’s distance vectors are discarded for convention + differentiability consistency).

Parameters:

coord: (nf, nloc, 3) local coordinates.
atype: (nf, nloc) local atom types; type < 0 marks a virtual atom, excluded as center and neighbor (the search itself is type-blind).
box: (nf, 3, 3) simulation cell, or None for non-periodic.
rcut: cutoff radius.
layout: edge-axis length policy; None => dynamic (torch) with min_edges guards.
with_csr: Whether to construct destination/source CSR views for a consumer that requires edge-grouped reductions.
canonicalize: Whether to reorder every edge field into destination-major form. Implies with_csr=True.
pair_excl: Optional PairExcludeMask for model-level pair_exclude_types. When given, apply_pair_exclusion() is applied after the geometric search. None (default) leaves all geometrically valid edges present.
compact: Passed to apply_pair_exclusion(); see that function for details. Ignored when pair_excl is None.

Returns:

graph: The carry-all NeighborGraph over the LOCAL atoms (n_node = nloc per frame), with edge_vec recomputed differentiably from coord/box.

Raises:

ImportError: if the optional ase package is not installed.