xnn.dnn.models.physnet.PhysNet#

class xnn.dnn.models.physnet.PhysNet(cutoff=10.0, lr_cutoff=None, n_features=128, n_rbf=64, num_blocks=5, num_residual_atomic=2, num_residual_interaction=3, num_residual_output=1, use_electrostatics=True, use_dispersion=True, s6=None, s8=None, a1=None, a2=None, d3_references='2010', energy_shift=0.0, energy_scale=1.0, charge_shift=0.0, charge_scale=1.0, species=None, atomic_energies=None, atomic_scales=None)[source]#

Bases: InteratomicPotential

Faithful PhysNet (Unke & Meuwly 2019): energies, forces, and charges.

See the module docstring for the architecture walk-through. All options are read from ModelConfig.extra (see from_config()); upstream train.py spellings are translated by xnn.common.config.translate.

Parameters:

cutoff (float, optional) – Short-range cutoff sr_cut of the neural-network interactions and the radial basis, by default 10.0 (the paper’s value).
lr_cutoff (float or None, optional) – Long-range cutoff for the electrostatic/dispersion terms (upstream lr_cut); the Coulomb term is force-shifted so energy and forces vanish smoothly there. None (default) evaluates the long-range terms un-damped on the graph’s edge list.
n_features (int, optional) – Feature-space width F, by default 128.
n_rbf (int, optional) – Number of radial basis functions K, by default 64.
num_blocks (int, optional) – Number of stacked module blocks, by default 5 (the paper; the upstream code default is 3).
num_residual_atomic (int, optional) – Residual blocks for atom-wise refinements, by default 2.
num_residual_interaction (int, optional) – Residual blocks refining the proto-message, by default 3 (the paper; the upstream code default is 2).
num_residual_output (int, optional) – Residual blocks in the output blocks, by default 1.
use_electrostatics (bool, optional) – Add the switched/shielded Coulomb energy of the predicted partial charges (paper eqs 12-13), by default True.
use_dispersion (bool, optional) – Add Grimme D3(BJ) dispersion, by default True.
s6 (float or None, optional) – D3(BJ) parameters. None (default) makes them learnable (softplus-reparametrized, initialized to the HF values), a number fixes them.
s8 (float or None, optional) – D3(BJ) parameters. None (default) makes them learnable (softplus-reparametrized, initialized to the HF values), a number fixes them.
a1 (float or None, optional) – D3(BJ) parameters. None (default) makes them learnable (softplus-reparametrized, initialized to the HF values), a number fixes them.
a2 (float or None, optional) – D3(BJ) parameters. None (default) makes them learnable (softplus-reparametrized, initialized to the HF values), a number fixes them.
d3_references (str, optional) – D3 reference systems: "2010" (default; Grimme’s original tables, as in upstream PhysNet) or "2024" (the current simple-dftd3 references, which re-parametrize Fr-Pu). Identical for Z <= 86; see xnn.common.models.d3.legacy_c6_table().
energy_shift (float, optional) – Initial value of the per-element energy shift/scale tables (upstream Eshift/Escale), by default 0 and 1.
energy_scale (float, optional) – Initial value of the per-element energy shift/scale tables (upstream Eshift/Escale), by default 0 and 1.
charge_shift (float, optional) – Initial value of the per-element charge shift/scale tables, by default 0 and 1.
charge_scale (float, optional) – Initial value of the per-element charge shift/scale tables, by default 0 and 1.
species (list of int or None, optional) – Only used to interpret atomic_energies/atomic_scales; the model itself handles all elements up to Z = 94.
atomic_energies (array-like or None, optional) – Per-species reference energies loaded into Eshift (aligned with species), like upstream’s dataset-regression initialization.
atomic_scales (array-like or None, optional) – Per-species initial Escale values (aligned with species).

Notes

forward additionally returns "charges" (corrected partial charges, summing exactly to the total charge – 0 unless the graph carries a total_charge attribute), "dipole" (eq 15) and "nh_loss" (the non-hierarchicality penalty, paper eq 18/19) for training-loop use.

property s6: Tensor#: Effective D3 s6 coefficient.

property s8: Tensor#: Effective D3 s8 coefficient.

property a1: Tensor#: Effective D3 a1 coefficient.

property a2: Tensor#: Effective D3 a2 coefficient.

atomic_properties(atomic_numbers, edge_index, edge_vec)[source]#

Scaled atomic energies/charges before the long-range terms.

Parameters:

atomic_numbers (Tensor) – Per-atom atomic numbers, shape (N,).
edge_index (Tensor) – Edge index (2, E); row 0 is the neighbor j, row 1 the center i (upstream idx_j/idx_i).
edge_vec (Tensor) – Edge vectors pos[i] - pos[j], shape (E, 3).

Returns:

Per-atom energies (N,), raw (uncorrected) per-atom charges (N,), edge distances (E,), the scalar non-hierarchicality penalty, and the final per-atom feature vectors (N, n_features).

Return type:

tuple of Tensor

scaled_charges(Qa, batch, num_graphs, total_charge=None)[source]#

Correct the raw charges to the exact total charge (paper eq 14).

Parameters:

Qa (Tensor)
batch (Tensor)
num_graphs (int)
total_charge (Tensor | None)

Return type:

Tensor

electrostatic_energy_per_atom(Dij, Qa, idx_i, idx_j)[source]#

Switched, shielded Coulomb energy per atom (paper eqs 12-13).

At short range the divergent 1/r is traded for the bounded 1/sqrt(r^2 + 1); _switch() blends the two so the kernel is smooth everywhere. When lr_cut is set, both kernels are force-shifted (value and slope zero at the cutoff) and pairs beyond it are dropped.

Parameters:

Dij (Tensor)
Qa (Tensor)
idx_i (Tensor)
idx_j (Tensor)

Return type:

Tensor

dispersion_energy_per_atom(atomic_numbers, Dij, idx_i, idx_j)[source]#

Grimme D3(BJ) dispersion per atom, in eV (paper eq 12).

Parameters:

atomic_numbers (Tensor)
Dij (Tensor)
idx_i (Tensor)
idx_j (Tensor)

Return type:

Tensor

forward(data)[source]#

Predict energies, corrected charges, and the dipole for a graph.

Parameters:: data (xnn.common.data.AtomicGraph) – The input atomic graph. If it carries a total_charge attribute (per-structure tensor), charges are corrected to it; otherwise neutral structures are assumed.
Returns:: "node_energy" (N,) and "energy" (B,) as for every xnn model, plus "charges" (N,) (corrected partial charges), "dipole" (B, 3) (paper eq 15) and "nh_loss" (scalar regularization term).
Return type:: dict of str to torch.Tensor

classmethod from_config(cfg)[source]#

Construct a PhysNet from a core model config.

Core fields: cfg.cutoff -> sr_cut, cfg.n_features -> F, cfg.n_rbf -> K, cfg.n_interactions -> num_blocks. Everything else is read from cfg.extra; upstream train.py spellings are translated by xnn.common.config.translate.

Parameters:: cfg (xnn.common.config.schema.ModelConfig) – The core model config.
Returns:: The instantiated model.
Return type:: PhysNet