Create graph from ASE atoms for JAX-MD

mlip.simulation.jax_md.graph_creation.create_graph_from_atoms_and_edges(atoms: Atoms, senders: Array, receivers: Array, displacement_fun: Callable[[Array, Array], Array], cell_to_box_fun: Callable[[Array], Array] | None = None, senders_long_range: Array | None = None, receivers_long_range: Array | None = None, ordering: GraphOrdering | str = GraphOrdering.SENDER) → Graph

Creates a graph from an ase.Atoms object and a list of edges.

This is the graph creation function used in the JAX-MD simulation engine. This method is only called once at the start of the simulation, then an update function (e.g. update_graph_in_simulation_step) is used at each step.

This function leaves the edge shifts of the graph empty and populates the displacement function instead. The total charge is read from atoms.info so that charge-aware models can apply a charge correction.

Parameters:

  • atoms – The ase.Atoms object of the system.

  • senders – The sender indexes of the edges for the graph.

  • receivers – The receiver indexes of the edges for the graph.

  • displacement_fun – Function that takes in two position vectors and returns the displacement vector between them.

  • cell_to_box_fun – A function that takes in a cell and outputs a minimum representation for use in the displ_fun.

  • senders_long_range – Optional sender indexes of the long-range edges. If provided together with receivers_long_range, the resulting graph will carry a long-range neighbor list using displacement_fun for vector computation.

  • receivers_long_range – Optional receiver indexes of the long-range edges.

  • ordering – Edge ordering to be enforced on the output graph. This is applied via graph.sort_edges(ordering) after the initial graph is constructed.

Returns:

The graph representing the system.