Batched inference

mlip.inference.batched_inference.run_batched_inference(structures: list[Atoms], force_field: ForceField, batch_size: int = 16, max_n_node: int | None = None, max_n_edge: int | None = None) → list[Prediction]

Runs a batched inference on given structures.

Computes energies, forces, and if available with the given force field, stress tensors. Result will be returned as a list of Prediction objects, one for each input structure.

Parameters:

• structures – The structures to batch and then compute predictions for.

• force_field – The force field object to compute the predictions with.

• batch_size – The batch size. Default is 16.

• max_n_node – This value will be multiplied with the batch size to determine the maximum number of nodes we allow in a batch. Note that a batch will always contain max_n_node * batch_size nodes, as the remaining ones are filled up with dummy nodes. The default is None which means an optimal number is automatically computed for the dataset.

• max_n_edge – This value will be multiplied with the batch size to determine the maximum number of edges we allow in a batch. Note that a batch will always contain max_n_edge * batch_size edges, as the remaining ones are filled up with dummy edges. The default is None which means an optimal number is automatically computed for the dataset.

Returns:

A list of predictions for each structure. These dataclasses will hold a float for energy, a numpy array for forces of shape (num_atoms, 3), and optionally one for stress of shape (3, 3).