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) → 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 function will leave the shifts of the graph empty and will populate the displacement function in the graph object instead. The total charge is read from atoms.info via ChemicalSystem.from_ase_atoms() so that charge-aware models (e.g. with a Coulomb energy head) can apply the partial-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 neighbour list using displacement_fun for vector computation.

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

Returns:

The graph representing the system.