learn.property_prediction.per_species_results

per_species_results(species, per_atom_quantities, species_idxs)

Sorts per-atom results by species and returns a per-species mean.

Only real (non-masked) particles should be input.

Parameters:

• species – An array storing for each atom the corresponding species.

• per_atom_quantities – An array with the same shape as species, storing per-atom quantities to be evaluated per-species.

• species_idxs – A (species,) array storing species-types for evaluation, e.g. jnp.unique(species).

Returns:

A (species,) array of per-species quantities.