Representation

Representation

In molecular modeling and machine learning-based potentials, selecting an appropriate representation of atomic-scale systems is crucial. The representation must capture both the geometric and physical properties of the system, enabling efficient interpolation across configurations. Additionally, it must respect the inherent symmetries of the problem, such as the rotational, translational, and permutational invariances that arise from the fundamental group O(3). These symmetries ensure that physically equivalent configurations produce identical predictions, a requirement for consistent statistical learning of energy and other molecular properties.

Behler-Parrinello Symmetry Functions

The Behler-Parrinello symmetry functions provide a framework for encoding the local atomic environment, capturing both radial and angular correlations around a central atom. These functions are widely used in neural network potentials due to their ability to represent atomic interactions while respecting key symmetries.

Radial symmetry functions

The radial symmetry function, often denoted as $G_2$, captures pairwise correlations between a central atom and its neighbors within a cutoff radius $r_c$. It is defined as:

$$ G_{2, s}^i = \sum_{j} e^{-\mu_{s}( R_{ij} - R_{s})^2} f_c(R_{ij}), $$

where $R_{ij}$ is the distance between atoms $i$ and $j$, and $\mu$ and $R_s$ are parameters that control the width and position of the Gaussian. The cutoff function $f_c(R_{ij})$ ensures that interactions smoothly decay to zero as the distance approaches $r_c$, thereby limiting the computational expense and enforcing locality in the interactions.

Angular basis functions

Angular symmetry functions, such as $G_3$, capture three-body correlations by considering the relative angles between triplets of atoms in the local environment. The general form of the angular function is:

$$ G_{3}^{i} = \sum_j \sum_{k\ne j} (1+\lambda cos\theta_{ijk}) e^{-\mu_{s}(R_{ij}^{2}+R_{ik}^{2}+R_{jk}^{2})} f_{c}(R_{ij})f_{c}(R_{ik}) $$

where $\theta_{ijk}$ is the bond angle formed by atoms $i$, $j$, and $k$, and $\lambda$ and $\mu$ are tunable parameters that control the strength of angular correlations. This function encodes higher-order interactions, complementing the radial function in fully describing the local atomic environment.

These symmetry functions have been successfully employed in several neural network potentials, including SchNet, PaiNN, PhysNet and AIMNet2, to accurately model atomic interactions while preserving rotational, translational, and permutational invariance.

For further reading: Imbalzano, Giulio, Andrea Anelli, Sinja Klees, J. Örg Behler, and Michele Ceriotti. 2018. “Automatic Selection of Atomic Fingerprints and Reference Configurations for Machine-Learning Potentials.” arXiv. http://arxiv.org/abs/1804.02150.

Fourier-Bessel Representations

To capture electronic structure more explicitly, Fourier-Bessel functions have emerged as a promising alternative. These functions form an invariant basis set for representing the wavefunction of a system, making them well-suited for modeling properties derived from quantum mechanical calculations like density functional theory (DFT). The spherical Bessel functions are part of the solutions to the Schrödinger equation under a zero potential $V(r)$, satisfying the Helmholtz equation:

$$ \nabla^2 + k^{2}\psi(d) = 0. $$

with boundary conditions $\psi(r_c) = 0$. This ensures that the basis functions inherently respect the boundary conditions imposed by the finite spatial extent of the system. The solution of this equation in spherical coordinates leads to a separation of variables, with the solutions expressed in terms of spherical Bessel functions and spherical harmonics.

General Formulation

The general form of the wave function in this representation, using spherical Bessel functions of the first and second kind $j_l(kd)$, respectivly $y_l(kd)$, and spherical harmonics $Y_l^m(\alpha, \phi)$ is

$$ \psi(d,\alpha, \phi) \sum_{l=0}^{\inf} \sum_{m=-l}^{l} (a_{l,m}j_{l}(kd) + b_{l,m}y_{l}(kd)) Y_{l}^{m}(\alpha, \phi) $$

where:

• $d$ is the radial distance
• $\alpha$ and $\phi$ are the angular coordinates in spherical coordinates (polar and azimuthal angles, respectively)
• $l$ is the orbital/angular momentum quantum number
• $m$ is the magnetic quantum number, ranging from -l to l
• $k$ is the wavenumber, related to the spatial frequency or energy of the system
• $a_{l,m}$ and $b_{l,m}$ are coefficients that scale the spherical Bessel functions and are typically determined through boundary conditions or fitting to data
• $j_l(kd)$ and $y_l(kd)$ is the spherical Bessel function of the first and second kind, representing the radial dependence.
• $Y_l^m(\alpha, \phi)$ are the spherical harmonics, which describe the angular part of the wave function.

Simplification to 2D Spherical Fourier-Bessel Basis

To simplify the expression to a form dependent only on the radial distance $d$ and a single angular coordinate $\alpha$, we set the following:

• $b_{l,m}=0$ to zero
• $m=0$ to map to 2D space

This leads to the following 2D spherical Fourier-Bessel basis, where the function depends only on $d$ and $\alpha$

$$ \tilde{a}_{SBF, l,n}(d, \alpha) = \sqrt{\frac{2}{c^3j^2_{l+1}(z_{l,n})}} j_{l}(\frac{z_{l,n}}{c}\cdot d) Y_{l}^{0}(\alpha) $$

where:

• $z_{l,n}$ is the n-th root of the spherical Bessel function, discretizing the wavenumber based on the boundary conditions
• $c$ is a normalization constant related to the system size (cutoff radius)
• $j_{l}$ and $j_{l+1}$ are spherical Bessel functions
• $Y_l^0$ is the spherical harmonic for $m=0$ ​

​Fourier-Bessel representations are particularly effective at capturing spatial oscillations of quantities like the electron density. Since these basis functions naturally arise from the Schrödinger equation, they provide a framework that maintains the physical properties of the system, such as orthogonality and completeness, while ensuring that the wavefunction obeys boundary conditions.

This representation is used in DimNet and DimNet++.

Different implementations of Radial symmetry functions in modelforge

The PhysNet/Tensornet radial basis function (RBF) is fundamentally different from the Ani and SchNet RBFs due to its use of an exponential transformation of the distances. This transformation compresses the input distances, resulting in non-uniform spacing of the peaks when plotted against the original distances.

Because of the exponential transformation, the centers of the Gaussian functions are not equally spaced in the original distance domain. Instead, they are equally spaced in the transformed domain (after applying the exponential function). This means:

• in the original distance domain, the peaks will appear closer together at smaller distances and further apart at larger distances (i.e., peaks are compressed towards smaller distances)
• compared to ANI and SchNet RBFs, which use linear transformations, the PhysNet/Tensornet RBF will naturally show peaks that do not align or have the same spacing.

The following plot shows the RBFs of the three different implementations (SchNet, PhysNet and ANI), each using 100 radial basis functions for

distances = torch.tensor([[0.5], [1.0], [1.5]], dtype=torch.float32) / 10
number_of_radial_basis_functions = 100
max_distance = 2.0  / 10
min_distance = 0.0