xnn.cnn.models.schnet#

SchNet (Schuett et al., NIPS 2017) – continuous-filter convolutional network.

A faithful implementation of the architecture described in the manuscripts

  • K. T. Schuett, P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Mueller, “SchNet: A continuous-filter convolutional neural network for modeling quantum interactions”, NIPS 30 (2017) – the architecture; and

  • K. T. Schuett, F. Arbabzadah, S. Chmiela, K. R. Mueller, A. Tkatchenko, “Quantum-chemical insights from deep tensor neural networks”, Nat. Commun. 8, 13890 (2017) – the DTNN predecessor, for the conventions SchNet inherits (per-atom energy standardization, sum pooling).

It is built directly from the papers’ equations on the xnn abstractions (GaussianRBF, scatter_sum(), shifted_softplus(), …); nothing is taken from the schnetpack code base.

Architecture (NIPS paper section 4, Fig. 2):

  • atoms are embedded by nuclear charge, x^0_i = a_{Z_i} (eq 3);

  • interatomic distances are expanded in Gaussian radial basis functions e_k(r) = exp(-gamma (r - mu_k)^2) with centers every 0.1 Angstrom and gamma = 10 per Angstrom^2 (section “Filter-generating networks”);

  • T interaction blocks (no weight sharing across blocks) refine the atom features through the ResNet-style residual x^{l+1}_i = x^l_i + v^l_i, where the residual is atom-wise -> cfconv -> atom-wise -> shifted softplus -> atom-wise (Fig. 2, middle);

  • the continuous-filter convolution (cfconv, eq 2) gates each neighbor’s features element-wise with a filter generated from the distance, x_i = sum_j x_j o W(r_ij), where the filter-generating network is two dense layers with shifted-softplus activations over the RBF expansion (Fig. 2, right);

  • the readout maps the final features through atom-wise (F -> F/2) -> shifted softplus -> atom-wise (F/2 -> 1) and sum-pools the per-atom energies over each structure (Fig. 2, left), after the DTNN per-atom standardization E_i = E_sigma * E^hat_i + E_mu (DTNN Methods, step 4; SchNet.set_energy_scale_shift()).

The shifted softplus ssp(x) = ln(0.5 e^x + 0.5) is used throughout, which keeps the potential-energy surface smooth (infinitely differentiable), so the autograd forces added by ForceStressOutput are smooth and energy-conserving by construction (paper eqs 1 and 4).

Deviations from the papers, all optional and off by default:

  • cutoff_fn="cosine" multiplies the generated filter by a smooth CosineCutoff envelope so the PES stays smooth when a finite neighbor-list cutoff truncates the graph. The paper itself trains without a cutoff – its RBF grid simply ends at 30 Angstrom, beyond any distance in its molecular datasets – which is what the default (None) reproduces.

  • atom_ref, a learnable per-element reference energy (the xnn convention shared by every model here), initialized to zero so it is inert unless set/trained. It plays the role of a per-element E_mu.

Works for molecules and periodic solids unchanged: periodicity enters only through data.edge_vectors(), which already accounts for cell shifts.

Classes

SchNet([n_features, n_interactions, n_rbf, ...])

SchNet continuous-filter convolutional interatomic potential.
Show Source