Utils

Shift and Scale

graph_pes.utils.shift_and_scale.guess_per_element_mean_and_var(
per_structure_quantity,
batch,
min_variance=0.01,
)[source]

Guess the per-element mean (\(\mu_Z\)) and variance (\(\sigma_Z^2\)) of a per-structure quantity using ridge regression under the following assumptions:

  1. the per-structure property, \(P\), is a summation over local properties of its components atoms: \(P = \sum_{i=1}^{N} p_{Z_i}\).

  2. the per-atom properties, \(p_{Z_i}\), are independent and identically distributed (i.i.d.) for each atom of type \(Z_i\) according to a normal distribution: \(p_{Z_i} \sim \mathcal{N}(\mu_{Z_i}, \sigma_{Z_i}^2)\).

Parameters:

  • per_structure_quantity (Tensor) – The per-structure quantity to guess the per-element mean and variance of.

  • batch (AtomicGraph) – The batch of graphs to use for guessing the per-element mean and variance.

Returns:

  • means – A dictionary mapping atomic numbers to per-element means.

  • variances – A dictionary mapping atomic numbers to per-element variances.

Return type:

tuple[dict[int, float], dict[int, float]]

Sampling

class graph_pes.utils.sampling.T

alias of TypeVar(‘T’)

class graph_pes.utils.sampling.SequenceSampler(collection, indices=None)[source]

A class that wraps a Sequence of T objects and provides methods for sampling from it without the need to manipulate or access the underlying data.

This is useful for e.g. sub-sampling a collection where individual indexing operations are expensive, such as a database on disk, or when indexing involves some form of pre-processing.

Parameters:

  • collection (Sequence[T]) – The collection to wrap.

  • indices (Sequence[int] | None) – The indices of the elements to include in the collection. If None, all elements are included.

Useful Types and Aliases

class graph_pes.atomic_graph.PropertyKey

Type alias for Literal["energy", "forces", "stress", "virial", "local_energies"].