xnn.dnn.models.hdnnp.HDNNP#

class xnn.dnn.models.hdnnp.HDNNP(species, featurizer, generation=2, hidden=(15, 15), activation='tanh', atomic_energies=None, charge_featurizer=None, charge_hidden=None, charge_activation=None, constrain_charges=True, gaussian_widths=None, screening=None, hardness='element', hardness_init=10.0, hardness_activation='linear', hardness_hidden=None, hardness_activation_nn=None, charge_neuron=True, ewald_accuracy=1e-10, coulomb_constant=14.39964547846902, energy_scale=1.0, chi_scale=1.0, hardness_scale=1.0)[source]#

Bases: DescriptorPotential

High-dimensional neural network potential, generations 2, 3 and 4.

Parameters:

  • species (sequence of int) – The elements (atomic numbers).

  • featurizer (AtomCenteredSymmetryFunctions) – Descriptor of the short-range networks.

  • generation (int, optional) – 2 (default), 3 or 4.

  • hidden (sequence of int or dict, optional) – Hidden-layer widths of the short-range networks, by default (15, 15).

  • activation (str, sequence or dict, optional) – Activations of the short-range networks (see _ElementNetworks), by default "tanh" with a linear output.

  • atomic_energies (sequence of float, optional) – Per-element energy added to every atom (the free-atom energies), aligned with species.

  • charge_featurizer (AtomCenteredSymmetryFunctions, optional) – Descriptor of the charge (3G) or electronegativity (4G) networks; the short-range descriptor by default.

  • charge_hidden (optional) – Architecture of the charge or electronegativity networks; those of the short-range networks by default.

  • charge_activation (optional) – Architecture of the charge or electronegativity networks; those of the short-range networks by default.

  • constrain_charges (bool, optional) – 3G: shift the network charges uniformly to the total charge (default, RuNNer; the “scaled” charges of Ko et al. 2021); False keeps the raw network outputs (“unscaled”).

  • gaussian_widths (sequence of float or dict, optional) – Width sigma of the Gaussian charge of every element (zero for point charges); default_gaussian_widths() by default.

  • screening (dict, optional) – Cutoff function (kind, r_cut, r_inner, exponent as in cutoff_function()) that screens the electrostatics inside its radius; None (default) for none.

  • hardness (str, optional) – 4G: "element" (default, one trainable constant per element) or "network" (element networks on the charge descriptor).

  • hardness_init (float or sequence, optional) – Initial element hardness parameters p_Z, by default 10.0.

  • hardness_activation (str, optional) – J = hardness_scale * activation(p_Z) of the element hardness, by default "linear".

  • hardness_hidden (optional) – Architecture of the hardness networks (hardness="network").

  • hardness_activation_nn (optional) – Architecture of the hardness networks (hardness="network").

  • charge_neuron (bool, optional) – 4G: feed the charge to the short-range networks, by default True.

  • ewald_accuracy (float, optional) – Truncation of the Ewald sums of periodic cells, by default 1e-10.

  • coulomb_constant (float, optional) – e^2 / (4 pi epsilon_0) in the model’s units, by default eV Angstrom.

  • energy_scale (float, optional) – Factors on the short-range, electronegativity and hardness outputs, by default 1.

  • chi_scale (float, optional) – Factors on the short-range, electronegativity and hardness outputs, by default 1.

  • hardness_scale (float, optional) – Factors on the short-range, electronegativity and hardness outputs, by default 1.

Variables:

  • generation (int) – The generation.

  • element_nets (_ElementNetworks) – Short-range networks.

  • charge_nets (_ElementNetworks or None) – Charge (3G) or electronegativity (4G) networks.

  • hardness_nets (_ElementNetworks or None) – Hardness networks (4G, hardness="network").

  • hardness_raw (torch.nn.Parameter or None) – Element hardness parameters p_Z (4G, hardness="element").

coulomb_matrices(data, dtype)[source]#

(atom indices, A) of every structure, A the Gaussian-charge Coulomb matrix in the model’s energy unit per e^2.

Parameters:
Return type:

list[tuple[Tensor, Tensor]]

charges_and_parameters(data, desc, desc_q, matrices=None)[source]#

The charges of a 3G or 4G model, with the network outputs that set them.

matrices are the coulomb_matrices() of data (computed when not given).

Parameters:
Return type:

dict[str, Tensor]

electrostatic_energy(data, charges, matrices=None)[source]#

Per-atom shares of the (screened) electrostatic energy of charges, (N,).

matrices are the coulomb_matrices() of data (computed when not given).

Parameters:
Return type:

Tensor

charge_input(data, charges)[source]#

The scaled charge input (q - shift_Z) factor_Z of the 4G short-range networks.

Parameters:
Return type:

Tensor

forward(data)[source]#

Energies (and for 3G/4G charges and electrostatics) of a batch.

Parameters:

data (AtomicGraph) – Batched atomic graph.

Returns:

node_energy, energy, energy_short (per structure), node_energy_short (the atomic energies of the short-range networks) and node_features (the short-range descriptor); 3G and 4G add charges, energy_elec and dipole (for molecules; the charges times the positions), 3G charges_raw, 4G electronegativities, hardness and chemical_potential.

Return type:

dict[str, Tensor]

fit_scaling(graphs, mode=None, charge_mode=None)[source]#

Fit the symmetry-function scaling of every featurizer to training structures.

Parameters:

  • graphs (iterable of AtomicGraph) – Training structures; iterated once per featurizer.

  • mode (str, optional) – Scaling modes of the short-range and the charge descriptors.

  • charge_mode (str, optional) – Scaling modes of the short-range and the charge descriptors.

Return type:

None

classmethod from_config(cfg)[source]#

Build an HDNNP (or, for generation: 1, an NNP1G) from a config.

cfg.extra keys: species, generation (2), cutoffs and symmetry_functions (the specs of AtomCenteredSymmetryFunctions; a generic set from default_symmetry_functions() with cfg.cutoff when absent, or radial-only functions from the older etas / rs keys), scaling, scale_range, charge_cutoffs / charge_symmetry_functions / charge_scaling, hidden, activation, charge_hidden, charge_activation, atomic_energies, gaussian_widths, screening, hardness, hardness_init, hardness_activation, hardness_hidden, charge_neuron, constrain_charges, ewald_accuracy; runner (a RuNNer model directory) loads that model instead. 1G: see NNP1G.from_config().

Return type:

HDNNP | NNP1G