xnn.dnn.models.physnet#

PhysNet (Unke & Meuwly 2019): message-passing HDNN with explicit physics.

A faithful, self-contained PyTorch translation of the original TensorFlow implementation (MMunibas/PhysNet, neural_network/NeuralNetwork.py and layers/) on the xnn abstractions: it subclasses InteratomicPotential directly (PhysNet is a message-passing high-dimensional NN – it has no hand-crafted descriptor, so DescriptorPotential does not apply) and gets forces/stress from the shared ForceStressOutput. Given the same weights it reproduces the original TensorFlow graph to machine precision (tests/test_physnet.py and the block-by-block notebook).

Architecture (paper eqs 3-15, J. Chem. Theory Comput. 15, 3678, 2019):

  • nuclear charges are embedded into F-vectors (a 95-row table indexed by Z directly – all elements up to Pu, no species list needed; eq 3);

  • distances are expanded in K radial basis functions g_k(r) = phi(r) exp(-beta_k (exp(-r) - mu_k)^2) with learnable centers and widths (softplus-reparametrized for positivity) and the smooth cutoff phi (eqs 7-8);

  • num_blocks modules refine the features: an interaction layer computes the message v from gated features and the distance-based attention mask G g(r_ij) (eqs 5-6), followed by pre-activation residual blocks (eq 4);

  • every module feeds an output block whose zero-initialized linear head predicts per-atom energy and partial-charge contributions; module outputs are summed and scaled/shifted per element (eqs 9-10);

  • predicted charges are corrected to the exact total charge (eq 14) and enter a damped/switched Coulomb term (eqs 12-13 – the code form: shielded 1/sqrt(r^2+1) below sr_cut/2, smoothstep-switched to 1/r, and force-shifted at lr_cutoff when one is set);

  • Grimme D3(BJ) dispersion (d3, an independent implementation verified against the upstream TF module, tables included) with optionally learnable s6/s8/a1/a2 completes the total energy (eq 12).

Upstream conventions preserved: shifted-softplus activation, semi-orthogonal Glorot weight init with zero biases, zero-initialized k2f/output heads, per-element scale/shift tables of length 95, kehalf Coulomb constant in eV*Angstrom units, and the non-hierarchicality penalty returned as "nh_loss". Dropout (upstream keep_prob, default 1.0 = off) is not implemented.

The neighbor-list radius (self.cutoff) is lr_cutoff when set, otherwise sr_cut: radial-basis features vanish identically beyond sr_cut because of the phi envelope, so feeding the longer-range edge list to the interaction blocks is mathematically identical to upstream’s separate short-range index list. Without lr_cutoff upstream evaluates electrostatics/dispersion over all pairs; in xnn the pair list is the graph’s, so set lr_cutoff (or a large cutoff) to capture long-range terms explicitly.

Functions

semi_orthogonal_glorot_weights(n_in, n_out)

Random (semi-)orthogonal weights rescaled to Glorot variance.

softplus_inverse(x)

Return y such that softplus(y) = x.

Classes

PhysNet([cutoff, lr_cutoff, n_features, ...])

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