tensorial.gcnn.atomic package#

Submodules#

tensorial.gcnn.atomic.keys module#

Module contents#

tensorial.gcnn.atomic.AllAtomicNumbers#

alias of AtomicNumbers

tensorial.gcnn.atomic.AvgNumNeighbours#

alias of AverageFromFun

class tensorial.gcnn.atomic.AvgNumNeighboursByAtomType(atom_types, type_field='atomic_numbers', state=None)[source]#

Bases: AvgNumNeighboursByType

Parameters:

  • atom_types (Union[Sequence[int], Int[Array, 'n_types'], Int[ndarray, 'n_types']])

  • type_field (str)

  • state (list[Average] | None)

Initializes the AvgNumNeighboursByType class with node types, type field, and optional state.

Parameters:

  • node_types – Sequence of integers representing the types of nodes or an array with shape (n_types,).

  • type_field – String indicating the field name used to identify node types, defaults to “type_id”.

  • state – Optional Averages object containing precomputed state information.

  • atom_types (Union[Sequence[int], Int[Array, 'n_types'], Int[ndarray, 'n_types']])

  • type_field (str)

  • state (list[Average] | None)

Raises:

  • ValueError – If node_types is empty or contains invalid values.

  • TypeError – If type_field is not a string or state is not of type Averages or None.

class tensorial.gcnn.atomic.EnergyContributionLstsq(type_map, metric=None)[source]#

Bases: Metric

Parameters:
compute()[source]#

Compute the metric.

create(graphs, *_)[source]#

Create a new metric instance from data.

Parameters:

graphs (GraphsTuple)

Return type:

EnergyContributionLstsq

empty()[source]#

Create a new empty instance.

By default, this will call the constructor with no arguments, if needed, subclasses can overwrite this with custom behaviour.

Return type:

EnergyContributionLstsq

merge(other)[source]#

Merge the metric with data from another metric instance of the same type.

Parameters:

other (EnergyContributionLstsq)

Return type:

EnergyContributionLstsq

update(graphs, *_)[source]#

Update the metric from new data and return a new instance.

Parameters:

graphs (GraphsTuple)

Return type:

EnergyContributionLstsq

class tensorial.gcnn.atomic.EnergyPerAtomLstsq(*args, state=None, **kwargs)[source]#

Bases: FromFun

Calculate the least squares estimate of the energy per atom

Parameters:

state (Metric[TypeVar(_OutT)])

compute()[source]#

Compute function.

Return type:

Array

static func(graph, *_)[source]#
metric#

alias of LeastSquaresEstimate

tensorial.gcnn.atomic.ForceStd#

alias of Force

tensorial.gcnn.atomic.NumSpecies#

alias of Species

class tensorial.gcnn.atomic.SpeciesTransform(atomic_numbers, field='atomic_numbers', out_field='species')[source]#

Bases: Module

Take an ordered list of species and transform them into an integer corresponding to their position in the list

Parameters:

  • atomic_numbers (Union[Sequence[int], Int[Array, 'numbers'], Int[ndarray, 'numbers']])

  • field (str)

  • out_field (str)

atomic_numbers: Int[Array, 'numbers']#
field: str = 'atomic_numbers'#
out_field: str = 'species'#
class tensorial.gcnn.atomic.TypeContributionLstsq(xtx=None, xty=None)[source]#

Bases: Metric[Array | ndarray]

Online Least Squares Metric.

Uses ‘Sufficient Statistics’ (XtX, Xty) to perform linear regression without storing the entire dataset history.

Parameters:

  • xtx (Float[Array, 'n_types n_types'] | None)

  • xty ('] | None)

compute(regularization=1e-06)[source]#

Compute the metric.

Parameters:

regularization (float)

classmethod create(cls, type_counts, values, mask=None, /)[source]#

Create a new metric instance from data.

Parameters:

  • type_counts (Union[Int[Array, 'batch_size n_types'], Int[ndarray, 'batch_size n_types'], Float[Array, 'batch_size n_types'], Float[ndarray, 'batch_size n_types']])

  • values (Union['], ']])

  • mask (Union[Bool[Array, 'batch_size'], Bool[ndarray, 'batch_size'], None])

Return type:

TypeContributionLstsq

classmethod empty()[source]#

Create a new empty instance.

By default, this will call the constructor with no arguments, if needed, subclasses can overwrite this with custom behaviour.

Return type:

TypeContributionLstsq

property is_empty#
merge(other)[source]#

Merge the metric with data from another metric instance of the same type.

Parameters:

other (TypeContributionLstsq)

Return type:

TypeContributionLstsq

update(type_counts, values, mask=None, /)[source]#

Update the metric from new data and return a new instance.

Parameters:

  • type_counts (Union[Int[Array, 'batch_size n_types'], Int[ndarray, 'batch_size n_types'], Float[Array, 'batch_size n_types'], Float[ndarray, 'batch_size n_types']])

  • values (Union['], ']])

  • mask (Union[Bool[Array, 'batch_size'], Bool[ndarray, 'batch_size'], None])

Return type:

TypeContributionLstsq

xtx: Float[Array, 'n_types n_types'] | None = None#
xty: Float[Array, 'n_types ...'] | None = None#
tensorial.gcnn.atomic.graph_from_ase(ase_atoms, r_max, *, key_mapping=None, atom_include_keys=('numbers',), edge_include_keys=(), global_include_keys=(), cell=None, pbc=None, use_calculator=True, **kwargs)[source]#

Create a jraph Graph from an ase.Atoms object

Parameters:

  • ase_atoms – the Atoms object

  • r_max – the maximum neighbour distance to use when considering two atoms to be neighbours

  • key_mapping

  • atom_include_keys

  • global_include_keys

  • cell – an optional unit cell (otherwise will be taken from ase.cell)

  • pbc – an optional periodic boundary conditions array [bool, bool, bool] (otherwise will be taken from ase.pbc)

  • use_calculator – if True, will try to use an attached calculator get additional properties

  • ase_atoms (Atoms)

  • r_max (Number)

  • key_mapping (dict[str, str] | None)

  • atom_include_keys (Iterable | None)

  • edge_include_keys (Iterable | None)

  • global_include_keys (Iterable | None)

  • cell (Optional[Union[Float[Array, '3 3'], Float[ndarray, '3 3']]])

  • pbc (Union[bool, Union[tuple[bool, bool, bool], Bool[Array, '3'], Bool[ndarray, '3'], Bool[TypedNdArray, '3']], None])

  • use_calculator (bool)

Return type:

GraphsTuple

Returns:

the atomic graph

tensorial.gcnn.atomic.graph_from_pymatgen(pymatgen_structure, r_max, *, key_mapping=None, atom_include_keys=('numbers',), edge_include_keys=(), global_include_keys=(), cell=None, pbc=None, graph_globals=None, **kwargs)[source]#

Create a jraph Graph from a pymatgen SiteCollection object or subclass (e.g. Structure, Molecule)

Note that the special atom key “numbers” is used to retrieve atomic numbers using SiteCollection.atomic_numbers. All other keys are used to retrieve site properties using SiteCollection.site_properties.

Parameters:

  • pymatgen_structure – the SiteCollection object

  • r_max – the maximum neighbour distance to use when considering two atoms to be neighbours

  • key_mapping

  • atom_include_keys

  • global_include_keys

  • cell – an optional unit cell (otherwise will be taken from Structure.lattice.matrix if it exists)

  • pbc – an optional periodic boundary conditions array [bool, bool, bool] (otherwise will be taken from Structure.lattice.pbc if it exists)

  • pymatgen_structure (pymatgen.core.SiteCollection)

  • r_max (Number)

  • key_mapping (dict[str, str] | None)

  • atom_include_keys (Iterable | None)

  • edge_include_keys (Iterable | None)

  • global_include_keys (Iterable | None)

  • cell (Optional[Union[Float[Array, '3 3'], Float[ndarray, '3 3']]])

  • pbc (Union[bool, Union[tuple[bool, bool, bool], Bool[Array, '3'], Bool[ndarray, '3'], Bool[TypedNdArray, '3']], None])

  • graph_globals (dict[str, Union[Array, ndarray]] | None)

Return type:

GraphsTuple

Returns:

the atomic graph

tensorial.gcnn.atomic.per_species_rescale(num_types, field, *, types_field=None, out_field=None, shifts=None, scales=None)[source]#
Parameters:

  • num_types (int)

  • field (str)

  • types_field (str)

  • out_field (str)

  • shifts (Union[Array, ndarray, bool, number, bool, int, float, complex, TypedNdArray])

  • scales (Union[Array, ndarray, bool, number, bool, int, float, complex, TypedNdArray])

Return type:

IndexedRescale