deploy.graphs

deploy.graphs#

Graphs for exporting potential and force models.

Graphs / Neighbor lists#

class SimpleSparseNeighborList(senders, receivers, max_edges)[source]#

Simple neighbor list representation using precomputed neighbor list.

This neighbor list is a sparse representation of a graph. It does not infer the neighbor list from the positions but acts as an interface between the precomputed neighbor list, e.g., from LAMMPS, and the exported model.

Nevertheless, this class increases the efficiency of the exported model while reducing necessary data transfer by pruning the neighbor list. Therefore, the class filters out all edges that are longer than the specified model cutoff distance. Moreover, it prunes all edges between ghost atoms (atoms not in the local domain) that are not relevant for a correct force computation.

Variables:

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

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

  • max_edges (jax.Array) – The maximum number of relevant edges in the neighbor list.

static create_from_args(r_cutoff, nbr_order, position, species, ghost_mask, valid_mask, newton, *args)[source]#

Creates the neighbor list from inputs to the exported function.

Return type:

Tuple[SimpleSparseNeighborList, NeighborListStatistics]

static set_properties(proto)[source]#

Assigns the graph type to the protobuf message.

class SimpleDenseNeighborList(nbrs, max_edges, max_triplets)[source]#

Simple dense neighbor list representation using precomputed neighbor list.

This neighbor list is a semi-sparse representation of a graph. It does not infer the neighbor list from the positions but acts as an interface between the precomputed neighbor list, e.g., from LAMMPS, and the exported model.

This class increases the efficiency of the exported model while reducing necessary data transfer by pruning the neighbor list. Therefore, the class filters out all edges that are longer than the specified model cutoff distance. Moreover, it prunes all edges between ghost atoms (atoms not in the local domain) that are not relevant for a correct force computation.

Variables:

  • senders – The sender indices of the edges.

  • receivers – The receiver indices of the edges.

  • max_edges (jax.Array) – The maximum number of relevant edges in the neighbor list.

static create_from_args(r_cutoff, nbr_order, position, species, ghost_mask, valid_mask, newton, *args)[source]#

Creates the neighbor list from inputs to the exported function.

Return type:

Tuple[SimpleDenseNeighborList, NeighborListStatistics]

static set_properties(proto)[source]#

Assigns the graph type to the protobuf message.

Neighbor List Statistics#

Neighborlist statistic are necessary to correctly allocate any buffer necessary in neighbor list communications.

ListStatistics#

alias of Dict

class NeighborListStatistics[source]#

Statistics for the SimpleSparseNeighborList.

Utility Functions#

compute_cell_list(position, id_buffer, cutoff, mask=None, eps=0.001)[source]#

Assigns particle IDs into a 3D grid.

This implementation follows the JAX, M.D. implementation, but aims to support building a cell list by only using shape information from the input arguments.

Parameters:

  • position – The position of the atom.

  • id_buffer – Determines the dimensions of the grid and the cell capacities. Shape (nx, ny, nz, c) correponds to the numbers of cells in x,y,z dimensions and the maximum capacity per cell c.

  • cutoff – Cutoff to check the dimensions of the cells. If the cell dimensions are smaller than the cutoff, increases the box size to enlarge the cells. Has the downside that cells will get fuller than usual, but will still yield correct neighbor list results.

  • mask – Specifies whether particles should be ignored (mask = 0)

  • eps – Tolerance increasing the box and cells to avoid wrong classification

Returns:

Returns a tuple with updated particle ids per grid and a dataclass containing statistics of the build.

prune_neighbor_list(list, local, max_edges, nbr_order, half_list=False)[source]#

Prunes the neighbor list by removing edges irrelevant to local atoms.

For simplicity, a neighbor list might be built for all atoms within a, e.g., rectangular domain. However, this list can contain atoms that are not relevant for the force computation of local atoms. Therefore, this function prunes the neighbor list by removing all edges that are not relevant for the local atoms. For example, given a simple lennard-jones potential, the neighbor list should only contain atoms that are first-order neighbors to any local atoms.

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 neighbor list is a half list. This means that an edge from i to j implies an edge from j to i.

Returns:

Returns the pruned neighbor list and the number of valid edges.

prune_neighbor_list_dense(list, local, nbr_order)[source]#

Prunes a dense neighbor list.

Parameters:

  • list – Sparse neighbor list to prune.

  • local – Mask specifying the local atoms.

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

Returns:

Returns the pruned neighbor list, the number of valid edges, and the number of triplets from the valid edges.

Contents