deploy.graphs

Graphs for exporting potential and force models.

Graphs and Neighbor Lists

class NeighborList

Abstract class for neighbor list graphs.

static create_from_args(r_cutoff, nbr_order, position, local_mask, valid_mask, newton_pair, *args)

Build a model graph from engine-provided neighbor arrays.

local_mask selects atoms owned by the current rank, while valid_mask also includes its real ghost atoms. newton_pair is fixed for one exported model variant.

Return type:

Tuple[NeighborList, Dict]

abstractmethod static set_properties(proto, *, include_pair_type=False, newton_pair=True)

Assigns the graph type to the protobuf message.

class SimpleSparseNeighborList(senders, receivers, max_edges, pair_type=None)

Represents a precomputed neighbor list as a sparse graph.

A simulation engine such as LAMMPS supplies the graph instead of asking chemtrain to infer neighbors from positions. Before model evaluation, chemtrain removes edges beyond the model cutoff and ghost-only edges that cannot affect owned forces. Pruning reduces both computation and data transfer.

Variables:

• senders (jax.jaxlib._jax.Array) – The sender indices of the edges.

• receivers (jax.jaxlib._jax.Array) – The receiver indices of the edges.

• max_edges (jax.jaxlib._jax.Array) – Internally computed Boolean mask selecting relevant edges. The corresponding exported input is only a shape carrier: its length sets the static pruning capacity and its values are ignored.

• pair_type (jax.jaxlib._jax.Array | None) – Optional topology category aligned with senders and receivers. Invalid and padding edges have category zero.

static create_from_args(r_cutoff, nbr_order, position, local_mask, valid_mask, newton_pair, *args)

Build a model graph from engine-provided neighbor arrays.

local_mask selects atoms owned by the current rank, while valid_mask also includes its real ghost atoms. newton_pair is fixed for one exported model variant.

Return type:

Tuple[SimpleSparseNeighborList, NeighborListStatistics]

static set_properties(proto, *, include_pair_type=False, newton_pair=True)

Assigns the graph type to the protobuf message.

class SimpleDenseNeighborList(nbrs, max_edges, max_triplets, pair_type=None)

Represents a precomputed neighbor list as a dense graph.

A simulation engine such as LAMMPS supplies one row per central atom instead of asking chemtrain to infer neighbors from positions. Before model evaluation, chemtrain removes entries beyond the model cutoff and ghost-only entries that cannot affect owned forces. Pruning reduces both computation and data transfer.

Variables:

• nbrs (jax.jaxlib._jax.Array) – Dense matrix of receiver indices, one row per sender.

• max_edges (jax.jaxlib._jax.Array) – Internally computed Boolean edge mask. The corresponding exported input only carries its static capacity. Values are ignored.

• max_triplets (jax.jaxlib._jax.Array) – Internally computed Boolean triplet mask. The corresponding exported input only carries its static capacity. Values are ignored.

• pair_type (jax.jaxlib._jax.Array | None) – Optional topology category with the same shape as nbrs. Invalid and padding entries have category zero.

static create_from_args(r_cutoff, nbr_order, position, local_mask, valid_mask, newton_pair, *args)

Build a model graph from engine-provided neighbor arrays.

local_mask selects atoms owned by the current rank, while valid_mask also includes its real ghost atoms. newton_pair is fixed for one exported model variant.

Return type:

Tuple[SimpleDenseNeighborList, NeighborListStatistics]

static set_properties(proto, *, include_pair_type=False, newton_pair=True)

Assigns the graph type to the protobuf message.

Neighbor-List Statistics

Neighbor-list statistics describe capacities used during graph construction, pruning, and executable compilation.

class NeighborListStatistics

Capacity statistics for supported engine-provided neighbor lists.

max_neighbors is the retained directed-edge count. For sparse graphs, overlong is the valid edge count before dependency pruning. For dense graphs, overlong is the retained ordered-triplet count. The field names are retained for model-format compatibility.

Utility Functions

prune_neighbor_list(list, local, max_edges, nbr_order, half_list=False)

Prunes a sparse graph to dependencies of owned atoms.

Starting from local, the function follows both edge directions for the requested graph depth. Edges incident to the reachable atoms are retained. A half list is expanded first so model code receives the same directed graph representation in both Newton modes.

Parameters:

• list – Sparse neighbor list to prune.

• local – Mask specifying the local atoms.

• max_edges – Maximum number of edges in the pruned list.

• nbr_order (int) – Maximum order of neighbors required for the force computation.

• half_list (bool) – If True, the engine supplied only one directed edge for each physical pair. Pruning adds the reverse edge so exported models see the same full directed graph in both Newton modes.

Returns:

The pruned neighbor list and its number of valid edges.

prune_neighbor_list_dense(list, local, nbr_order)

Prunes a dense graph to dependencies of owned atoms.

Parameters:

• list – Dense neighbor list to prune.

• local – Mask specifying the local atoms.

• nbr_order (int) – Maximum order of neighbors required for the force computation.

Returns:

The pruned neighbor list, number of valid edges, and number of triplets formed by the valid edges.