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:

type_map (Union[Sequence, Array, ndarray])
metric (TypeContributionLstsq)

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

Compute the metric.

Parameters:: regularization (float)

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]#

Computes the fit using the Moore-Penrose pseudo-inverse approach.

This naturally handles rank-deficient cases (like fixed 50:50 ratios) by producing the minimum-norm solution, which assigns equal contributions to indistinguishable types.

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']], 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']], 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])
scales (Union[Array, ndarray, bool, number, bool, int, float, complex])

Return type:

IndexedRescale