Averaging predictions per atom

metatrain.utils.per_atom.average_by_num_atoms(tensor_map_dict: Dict[str, TensorMap], systems: List[System], per_structure_keys: List[str]) → Dict[str, TensorMap]

Averages a dictionary of TensorMap objects by the number of atoms in each system.

This function averages by the number of atoms in each system. Targets that are present in per_structure_keys will not be averaged.

Parameters:

• tensor_map_dict (Dict[str, TensorMap]) – A dictionary of TensorMap objects.

• systems (List[System]) – The systems used to compute the predictions.

• per_structure_keys (List[str]) – A list of keys whose corresponding TensorMap objects that should not be averaged.

Returns:

The dictionary of averaged TensorMap objects.

Return type:

Dict[str, TensorMap]

metatrain.utils.per_atom.divide_by_num_atoms(tensor_map: TensorMap, num_atoms: Tensor) → TensorMap

Takes the average values per atom of a TensorMap.

Since some quantities might already be per atom (e.g., atomic energies or position gradients), this function only divides a block (or gradient block) by the number of atoms if the block’s samples do not contain the “atom” key. In practice, this guarantees the desired behavior for the majority of the cases, including energies, forces, and virials, where the energies and virials should be divided by the number of atoms, while the forces should not.

Parameters:

• tensor_map (TensorMap) – The input tensor map.

• num_atoms (Tensor) – The number of atoms in each system.

Returns:

A new tensor map with the values divided by the number of atoms.

Return type:

TensorMap