spainn.model

Properties per electronic state

Prediction modules for fundamental properties that stem directly from a single electronic states (\(j\)) and are intrinsic to that state without any interaction.

Energies

class model.Atomwise(n_in: int, n_out: int = 1, n_hidden: int | ~typing.Sequence[int] | None = None, n_layers: int = 2, activation: ~typing.Callable = <function silu>, aggregation_mode: str = 'sum', output_key: str = 'energy', per_atom_output_key: str | None = None)

Predicts atom-wise contributions and accumulates global prediction, e.g., for energies.

If aggregation_mode is None, only the per-atom predictions will be returned.

Parameters:

• n_in (input dimension of representation)

• n_out (output dimension of target property (default: 1), i.e., number of states)

• n_hidden (size of hidden layers.) – If an integer, same number of node is used for all hidden layers resulting in a rectangular network. If None, the number of neurons is divided by two after each layer starting n_in resulting in a pyramidal network.

• n_layers (number of layers.)

• aggregation_mode (one of {sum, avg} (default: sum))

• output_key (the key under which the result will be stored)

• per_atom_output_key (If not None, the key under which the per-atom result will be stored)

Return type:

Dictionary with predicted energies for multiple states.

Examples

The following code demonstrates how to define a prediction module for predicting the energies of two electronic states using the Atomwise module SPaiNN.

>>> pred_energy = spainn.Atomwise(
>>>     n_in=30,
>>>     n_out=2
>>> )

forward(inputs: Dict[str, Tensor]) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Forces

class model.Forces(calc_forces: bool = True, calc_stress: bool = False, energy_key: str = 'energy', force_key: str = 'forces')

Predicts forces and stress as response of the energy prediction w.r.t. the atom positions and strain for multiple electronic states.

The number of electronic states, for which the forces are computed is automatically taken from the shape of the energies. For example, energy inputs of shape \([1,3]\) result in the prediction of forces for three electronic states.

Parameters:

• calc_forces (If True, calculate atomic forces.)

• calc_stress (If True, calculate the stress tensor.)

• energy_key (Key of the energy in results.)

• force_key (Key of the forces in results.)

Return type:

Dictionary with predicted forces for multiple states.

Examples

The following code demonstrates how to define a prediction module for predicting the forces from for multiple electronic states from the first derivative of the respective predicted energies with respect to the atom positions using the Forces module. The number of the electronic states is automatically taken from the shape of the energies.

>>> pred_forces = spainn.Forces()

forward(inputs: Dict[str, Tensor]) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Properties per coupling of states

Prediction modules for properties that emerge from the coupling or interaction of two electronic states.

Nonadiabatic couplings

class model.Nacs(n_in: int, n_out: int, n_hidden: int | ~typing.Sequence[int] | None = None, n_layers: int = 2, activation: ~typing.Callable = <function silu>, nac_key: str = 'nacs', use_vector_repr: bool = False)

Predicts nonadiabatic coupling vector, i.e. local, atomic couplings of electronic states of the same multiplicity.

This requires a representation supplying (equivariant) vector features [1].

References

[1]

Schütt, Unke, Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. ICML 2021, http://proceedings.mlr.press/v139/schutt21a.html

Parameters:

• n_in (input dimension of representation)

• n_out (output dimension of target property, i.e.,) – number of couplings

• n_hidden (size of hidden layers) – If an integer, same number of node is used for all hidden layers resulting in a rectangular network. If None, the number of neurons is divided by two after each layer starting n_in resulting in a pyramidal network.

• n_layers (number of layers)

• activation (activation function)

• nac_key (the key under which the nonadiabatic couplings are stored)

• use_vector_repr (If true, use vector representation to predict) – local, atomic couplings.

Return type:

Dictionary with predicted energies for multiple states.

Examples

The following code demonstrates how to define a prediction module for predicting the nonadiabatic couplings of two electronic states, i.e. one coupling vector using the Nacs module SPaiNN.

>>> pred_nacs = spainn.Nacs(
>>>     n_in=30,
>>>     n_out=1
>>> )

forward(inputs: Dict[str, Tensor]) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Dipoles Moments

class model.Dipoles(n_in: int, n_out: int = 1, n_hidden: int | ~typing.Sequence[int] | None = None, n_layers: int = 2, activation: ~typing.Callable = <function silu>, dipole_key: str = 'dipoles', return_charges: bool = True, charges_key: str = 'partial_charges', use_vector_repr: bool = False)

Predict Dipole Moments from latent partial charges. This requires a representation supplying (equivariant) vector features [2], [3], [4].

References

[2]

Schütt, Unke, Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. ICML 2021, http://proceedings.mlr.press/v139/schutt21a.html

[3]

Veit et al. Predicting molecular dipole moments by combining atomic partial charges and atomic dipoles. The Journal of Chemical Physics 153.2 (2020): 024113.

[4]

Gastegger, Behler, Marquetand. Machine learning molecular dynamics for the simulation of infrared spectra. Chemical science 8.10 (2017): 6924-6935.

Parameters:

• n_in (Input dimension of representation.)

• n_out (Output dimension, i.e., number of dipole moments (couplings).)

• n_hidden (Size of hidden layers.) – If an integer, same number of node is used for all hidden layers resulting in a rectangular network. If None, the number of neurons is divided by two after each layer starting n_in resulting in a pyramidal network.

• n_layers (Number of layers.)

• activation (Activation function.)

• dipole_key (Key under which dipoles are stored.)

• return_charges (If true, return partial charges.)

• charges_key (Key under which partial charges are stored.)

• use_vector_repr (If true, use (equivariant) vector representation to) – dipole momemnts.

Return type:

Dictionary with predicted dipoles for multiple states.

Examples

The following code demonstrates how to define a prediction module for predicting the dipole moments of two electronic states, i.e. one transition dipole and two permanent dipoles using the Dipoles module of SPaiNN.

>>> pred_dipoles = spainn.Dipoles(
>>>     n_in=30,
>>>     n_out=1
>>> )

forward(inputs: Dict[str, Tensor]) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Spin-orbit couplings

class model.Socs(n_in: int, n_out: int = 1, n_hidden: int | ~typing.Sequence[int] | None = None, n_layers: int = 2, activation: ~typing.Callable = <function silu>, socs_key: str = 'socs')

Predicts spin-orbit coupling vectors, i.e., local atomic couplings between electronic states of different multiplicity.

This requires a representation supplying (equivariant) vector features [5].

References

[5]

Schütt, Unke, Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. ICML 2021, http://proceedings.mlr.press/v139/schutt21a.html

Parameters:

• n_in (Input dimension of representation.)

• n_out (Output dimension, i.e., number of spin-orbit couplings.)

• n_hidden (Size of hidden layers.) – If an integer, same number of node is used for all hidden layers resulting in a rectangular network. If None, the number of neurons is divided by two after each layer starting n_in resulting in a pyramidal network.

• n_layers (Number of layers.)

• activation (Activation function.)

• socs_key (Key under which spin-orbit couplings are stored.)

Return type:

Dictionary with predicted dipoles for multiple states.

Examples

The following code demonstrates how to define a prediction module for predicting the spin-orbit couplings of two electronic states, e.g. \(S_0\) and \(T_1\) using the Socs module of SPaiNN.

>>> pred_dipoles = spainn.Dipoles(
>>>     n_in=30,
>>>     n_out=1
>>> )

forward(inputs: Dict[str, Tensor]) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.