Spectroscopy from Machine Learning by Accurately Representing the Atomic Polar Tensor

Vibrational spectroscopy clearly is one of the key techniques to elucidate microscopic structure and dynamics. Without the aid of theoretical approaches, it is however, often difficult to understand such spectra at a microscopic level. Therefore, there is a demand to calculate accurate, and thus predictive, vibrational spectra from ab initio techniques at corresponding experimental thermodynamic conditions. Ab initio molecular dynamics have repeatedly proved to be suitable for this purpose, however, the computational demand is huge, especially if expensive electronic structure methods are used. In the last decades, machine learning of potential energy surfaces has become a popular way to accelerate such simulations. This way the information on the electronic structure is usually lost and calculating IR or Raman spectra with proper selection rules is therefore still difficult. Here, the E(3)-equivariant neural network e3nn is used to fit the atomic polar tensor a posteriori on top of existing molecular dynamics simulations. The introduced approach is general and thus transferable to any system, with or without periodic boundary conditions. As an exemplary case study, the new approach is applied to calculate the IR spectrum of ambient liquid water. It is demonstrated, that the machine learned IR spectrum exactly matches the explicitly calculated one and excellently reproduces the experimentally measured spectrum in shape and intensity. It is foreshadowed that this approach may enable one to calculate accurate spectra at even higher levels of theory, such as hybrid DFT and beyond.

Vibrational spectroscopy, be it IR or Raman spectroscopy, is one of the most important techniques to unveil microscopic properties of matter, such as structure determination [1] or structural dynamics of water [2,3]. It can also be used to time-dependently monitor dynamical processes, e.g. intramolecular couplings in peptides [4], or proton transfer mechanisms [5]. Attenuated total reflection (ATR) IR or sum-frequency generation (SFG) [6] can selectively probe molecules at interfaces, such as metal/liquid water interfaces [7,8] or water/air interfaces [9]. Using THz spectroscopy the chemical environment of molecules can directly be probed and it elucidated solvation dynamics in aqueous solutions from ambient [10,11] to extreme thermodynamic conditions, such as high-pressures [12] and supercritical phases [13].
Because of the broad applicability and power of vibrational spectroscopy there is a demand to calculate accurate vibrational spectra from ab initio techniques which are predictive and also aid to understand experiments at the molecular level. When it comes to spectroscopy of condensed phase systems at finite temperature, ab initio molecular dynamics (AIMD) [14], where the electronic structure is calculated on-the-fly at every timestep, is the prime technique for several reasons. First, the potential energy surface can be accurately represented-clearly depending on the underlying electronic structure theory setup employed to drive the MD. Second, as the electronic structure is available at every time step, charge transfer and polarization effects are naturally included and dipole moments can directly be obtained, e.g. from maximally localized Wannier functions [14] or from the electron density [15]. Third, anharmonic effects are also naturally taken into account which can break the standard normal mode analysis [16] if they are too large. Notably, such large anharmonicities can even be present for single molecules (e.g. peptides) in the gas-phase [17,18], where AIMD simulations are required because the normal mode analysis fails to correctly reproduce the measured spectra.
From a MD trajectory, the frequency dependent Beer-Lambert absorption coefficient of IR spectroscopy can be calculated from the time auto correlation function of the total dipole moment vector M(t) of the simulation box (see e.g. [19]), where β = 1/k B T , k B is the Boltzmann constant, T is the temperature, V is the volume of the periodic simulation box, c is the speed of light in vacuum, 0 is the dielectric constant, and n(ω) is the frequency dependent refractive index. Note that the so-called "quantum correction factor" [20] has already been included. The main disadvantage of this technique is that the MD simulations need to be quite long, also to reduce the statistical noise to a minimum. Clearly, this is a problem for AIMD which can be highly demanding computationally, especially if more expensive techniques, such as hybrid DFT, are used to drive the MD. Over the last decades, machine learning (ML) approaches have been introduced with the aim to accelerate AIMD simulations. Therein, the expensive electronic structure calculations are replaced with a cheaper machine learning model, while retaining the same accuracy [21,22,23,24]. Arguably, ML techniques have repeatedly proved to reliably represent the potential energy surface from explicit electronic structure calculations at a fraction of the cost. One apparent problem of these pioneering ML techniques usually is that only the potential energy surface is trained (which is generally sufficient to run MD simulations), but all information on the electronic structure itself is lost. Therefore, total dipole moments at the quality of the underlying electronic structure theory cannot be obtained along the "MLMD" (machine learning molecular dynamics) simulation. One way to circumvent this issue was to extract single snapshots from the MLMD trajectory and explicitly calculate the electronic structure for those snapshots again, e.g. to calculate polarizability tensors for Raman spectra [25]. However, formally, time correlation functions (as in e.g. Eq. 1 for IR spectra), require the electronic structure at each time step; or at least frequently enough such that all vibrations present in the system are correctly sampled. Note, that the sampling theorem can be employed to determine how frequently time-dependent data needs to be provided [26], however, the fastest vibration needs to be known. For example, in case of liquid water, the fastest vibration is the O−H stretch at roughly 3500 cm −1 . According to the sampling theorem, data needs to be provided at least every roughly 4.5 fs to correctly sample this vibration. This introduces a huge bottleneck for MLMD simulations, if a significant amount of configurations needs to be explicitly recalculated anyways to get exact vibrational spectra.
In recent years, training atomic or molecular properties using ML has been an extremely active field, such that it is impossible to give a complete list of work done. Nevertheless, some key ideas on how vibrational spectra can be obtained from ML approaches are summarized in the following. Partial atomic charges have been introduced in third generation NNPs [24], but with the main purpose to include long-range interactions. Such trained atomic partial charges could potentially also be used as an output parameter to calculate dipole moments along a MLMD trajectory. An apparent problem of partial charges in general is, that they are no physical observables. As such, their magnitude depends on the chosen partitioning scheme employed, e.g. Mulliken [27], Hirshfeld [28], or Bader [29] (incomplete list). It could be shown that different partitioning schemes can yield very different results [30,31,32]. Moreover, choosing an unsuitable scheme for a given problem can lead to wrong molecular dipole moments [32] and can yield unphysical atomic charges [33]. Clearly, these caveats potentially also affect the quality of the IR spectrum calculated from such partial atomic charges.
A complementary approach is to train partial charges such, that molecular dipole moments are reproduced correctly [34,31] or to train the positions of Wannier centers [35,36]. Molecular dipole moment are generally measurable and can thus be validated against experiments. In case of gas-phase systems (without periodic boundary conditions), the total dipole moment vector has recently been trained as a whole to predict IR spectra of protonated water clusters [37] as well as of an ethanol and an aspirin molecule [38]. In condensed systems with periodic boundary conditions (as it is the case here to simulate liquid water), the total dipole moment is, however, multivalued [39] which can possibly lead to ambiguities in the training and prediction process. For small single molecules, also the polarizability tensor has been trained [40,38] which can then be used to calculate Raman spectra of these molecules. Finally, even the full electron density of single molecules has been trained [41,42] as well as transition dipole moments to excited states [43] which allow the calculation of UV/Vis spectra.
In this work, I introduce an alternative idea to most accurately calculate vibrational spectra from MLMD simulations. The herein introduced "atomic polar tensor neural network" (APTNN) trains a strictly atomic property, is transferable to any system, and does not rely on partial charges or any definition of molecules or molecular coordinate systems. Using a Fourier transform identity on Eq. 1, the equation can be rewritten, such that whereṀ(t) is the time derivative of the total dipole moment. In an effort to express the dipolar velocity as a function of atomic velocities, the chain rule can be applied to express the ξ-th where ζ and ξ represent the three cartesian coordinates. In matrix notation, this equation can be rewritten in a more compact way, where v i (t) is the velocity of the i-th nuclei and is the "atomic polar tensor" (APT) [44] of atom i. The atom velocities v i (t) are readily available along any MLMD trajectory. Having a ML model for the APT therefore gives full access to the vibrational spectrum at the level of the underlying electronic structure theory. The APTNN is meant to be trained on existing MD trajectories (AIMD, MLMD). After the training, the APT of each sampled configuration can be predicted and thus the time correlation function can properly be sampled. The APT of atom i is the spatial derivative of the total dipole moment with respect to a displacement of that atom (Eq. 5). It is thus a 3x3 tensor which is invariant with respect to translations, but equivariant with respect to rotations. This means that if the atomic coordinates are translated in space, the corresponding APTs do not change. However, if the set of atomic coordinates rotates in space, the APTs also rotate accordingly. Numerically, the derivative is evaluated from central (or "two-sided") finite differences using a displacement of 0.01 A. Due to the finite differences, 6 additional single point calculations are necessary to calculate the APT of a single atom. Previously, it could already be shown that the chosen displacement is small enough for liquid water [45]. Indeed, I also computed the APT explicitly with a displacement of 0.04 A and did not find any significant difference. It might, however, not be sufficiently small enough for other systems.
Besides, the APT is also an interesting property to assign specific atomic motion to spectral features. Inspecting again Eq. 3 closely, the APT represents nothing else than the definition of the IR selection rule. As such, any velocity spectrum (e.g. the vibrational density of states) can be promoted to a proper IR spectrum by weighting with the respective APTs. This has been utilized in the past to decompose the spectrum into individual contributions, since dissecting atomic velocities is much more intuitive compared to dissecting molecular dipole moments. For example, in case of liquid water, the total spectrum has been decomposed into translational, rotational, and vibrational contributions based on atomic velocities [45]. By reweighting with the APT, the proper IR intensities of these contributions can be obtained. This decomposition requires the APT for each participating atom at each time step which can be enormously expensive or even impossible, see above. Therefore, the instantaneous normal mode (INM) approximation had to be used in that work [45]. This approximation is quite drastic because is typically introduces imaginary frequencies (similar to the standard normal mode analysis) which need to be dealt with in some ad hoc way. Instead, using a ML model for the APT, to be introduced in what follows, it is now possible to compute all APTs at each time step of the MD simulation such that the INM approximation is no longer necessary. Moreover, IR spectra of peptides have been dissected in terms of atomic velocities in the past [46], even using sophisticated decompositions based on graph theory [17]. Also in this context, it was not possible to compute the APT at each time frame of the AIMD simulation because of its high cost. Therefore, the APT was only calculated at some snapshots of the AIMD simulation, under the approximation that it does not change much as a function of time [47]. Notably, it could also be shown that this is also a delicate approximation which can yield spurious signals in the IR spectrum if the APT is not computed frequently enough [47].
An apparent problem when training an APT (see Eq. 5) is that it is an equivariant property as already mentioned, while most machine learning models can only infer invariant properties.
Recently, the e3nn framework has been introduced [48] for PyTorch which can be used to train E(3)-equivariant graph Neural Networks and thus enables one to infer also equivariant properties by a machine learning model. Its introduction also caused a huge boom in the field and numerous works have been published, where equivariant properties have been modeled. For example, the NequIP package was recently developed [49] to model potential energy surfaces using equivariant message passing. Moreover, it was also used recently to train the electron density as a whole for gas-phase molecules [42].
Here, I now use the so-called SimpleNetwork (v2106) from the e3nn toolkit to create an APTNN for liquid ambient water which is capable to predict the APTs of all atoms at a given MD snapshot. The introduced approach is, however, generally transferable to any other system, with and without periodic boundary conditions. As the underlying electronic structure theory, I opt to use the RPBE functional [50] supplemented by D3 dispersion corrections [51]. It could be shown repeatedly in the past, that this electronic structure model reproduces fluid water excellently, even far away from ambient conditions [52,53,13,54]. Previously, liquid ambient water has been simulated [52] and these data (16 trajectories, 20 ps each using a timestep of 1 fs) are used in this work. As a starting point, I randomly selected 10 statistically independent configurations from the available AIMD trajectories and calculated the APT. The calculation setup is exactly the same as before [52] and I refer to that reference for an elaborate description. All electronic structure calculations have been performed using version 8.0 of the CP2k program package [55] and the Quickstep module [56]. Each configuration contains 384 atoms and therefore 2304 single point calculations are required. This is undoubtedly a substantial computational commitment, however, this way 384 APTs are obtained from a single MD snapshot, which is quite a lot of training data which will become apparent in the following.
Having 10 MD snapshots with explicitly calculated APTs for all atoms available, an APTNN is trained by randomly selecting 9 (90 %) configurations for its training set. The remaining configuration is used to validate how well the model generalizes during the training. I use a e3nn SimpleNetwork model containing two message-passing layers and 20 features, up to l max = 2, defining the maximum for the spherical harmonics. Even and odd features were considered, resulting in the configuration 20x0o+20x0e+20x1o+20x1e+20x2o+20x2e. Moreover, I use a radial cutoff of 6 A which has been proved to work well to mimic the potential energy surface of liquid ambient water using High-Dimensional NNPs [57]. The model has been trained using the Adam optimizer [58] implemented in pyTorch [59]. An initial learning rate of 0.01 is used which is automatically reduced by a factor of 0.1, when the loss of the validation data set does not decrease further over the last 10 training epochs. The hyperparameters have been chosen manually based on the validation set performance during the training and have then been fixed. The overall performance of the trained model was finally evaluated on an unrelated test set and on the predicted IR spectrum compared with available reference and experimental data, see below.
The learning curves of a APTNN as a function of the data set size (4, 9, 18 and 27 training configurations) is shown in Fig. 1a. All APTNNs have been trained according to the above described training procedure. The presented mean squared error (MSE) is computed on the respective validation set only. The learning curves clearly show that the learning progress follows the expected behavior as a function of training epochs. Moreover, the training error systematically decreases as a function of provided training data set size as expected.
All data presented in the following is calculated from the APTNN trained using 9 configurations only (solid brown learning curve in Fig. 1a) after 100 epochs (marked by a vertical dashed line). To evaluate this model, I randomly select 10 configurations from the available AIMD simulations of liquid water which have not been included in the training set ("test set"). For those configurations, the APT is explicitly calculated using the electronic structure method as before and compared with the prediction of the APTNN. The direct component-wise comparison is shown in Fig. 1b and the overall component-wise RMSE is 3.49 × 10 −2 e, where e is the electron charge. Overall, a very good correlation between the model and the explicit calculation is found.
As a second quality benchmark, I now use the trained APTNN to predict the IR spectrum of liquid ambient water. In practical terms, I take all the available AIMD trajectories of RPBE-D3 liquid ambient water (16 trajectories, 20 ps each using a timestep of 1 fs, see above) and predict the APT for each atom at each time step. Recall that this is a daunting task for explicit electronic structure calculations. Having all these APTs available, I then calculate the total dipole moment derivative at each time step according to Eq. 4. Finally, the machine learned IR spectrum is obtained by Eq. 2 and presented in Fig. 2. In Fig. 2a it is compared with the reference IR spectrum which has been obtained directly from the explicit AIMD simulations. Note that this reference IR spectrum of RPBE-D3 water has already been published before, see e.g. Ref. [13]. The comparison shows that the machine learned spectrum reproduces the explicitly calculated reference spectrum exactly in shape and intensity. Moreover, in Fig. 2b the machine learned IR spectrum is compared with the experimental spectrum of ambient liquid water [60]. Again, the predicted spectrum almost exactly matches the experimental IR spectrum, without any intensity or frequency scalings. Note, that in Fig. 2b, the frequency dependent refractive index, n(ω) has been calculated using Kramers-Kronig relations [61], to obtain α(ω) from α(ω)n(ω), see Eq. 2, to be comparable with experiment.  Figure 2: In (a) the machine learned IR spectrum of liquid ambient water (red) is compared with the explicit AIMD reference spectrum (black), see text. In (b) the same machine learned IR spectrum is compared with accurate experimental measurements [60] (blue). Here, the IR spectrum has been divided by the frequency dependent refractive index n(ω) (calculated via the Kramers-Kronig relations, see e.g. Ref. [61]) to be comparable with the experiment. Mind the therefore different y scales in (a) and (b). Note that no scaling or shifting of frequency or intensity has been applied.
Remarkably, this accuracy is reached by only considering 10 randomly selected configurations for training, i.e. without even using more sophisticated active learning techniques. It foreshadows that training a APTNN with even more computationally expensive electronic structure theory will be possible; indeed, potential energy surfaces of coupled cluster accuracy have been successfully trained already [62,49,37]. This could for example be achieved using a committee of APTNNs which straightforwardly enables active learning. Committees have recently shown to be an efficient tool to obtain a highly accurate model while the training set size is minimized [63,64,65].
The acceleration gained by the APTNN is significant. To have a fair comparison, a benchmark was conducted on the very same machine (16 Intel(R) Xeon(R) E5-2640 v3 CPUs, 128 cores in total). A single MD step including wavefunction extrapolation took about 8 s using the GGA functional RPBE. Calculating an APT for a single atom requires six single points and therefore amounts to 48 s. Predicting an APT for a single atom using the APTNN however only takes about 8 ms, giving an acceleration of more than three orders of magnitude. Note that the acceleration becomes even larger when more expensive electronic structure methods are used, e.g. hybrid DFT or beyond. In that case, the computational time to calculate a single point increases dramatically, while the time needed to predict a single APT from the APTNN remains constant. Moreover, pyTorch and e3nn are designed to run more efficiently on GPUs than CPUs. A significant additional acceleration is therefore expected when switching to GPUs, however, this has not been tested yet.
In conclusion, a Machine Learning approach ("APTNN") has been introduced to model the atomic polar tensor along molecular dynamics simulations. This tensor allows one to calculate correct IR lineshapes and intensities with high accuracy and a significant acceleration compared to explicit electronic structure calculations. As an exemplary case study, the model was applied to liquid ambient water. It was demonstrated that the atomic polar tensor can indeed be successfully trained and utilized to exactly reproduce the IR spectrum calculated directly from available ab initio molecular dynamics simulations. The machine learned IR spectrum furthermore matches the available experimental IR spectrum. Importantly, training the atomic polar tensor does not require any definition of molecules, reference frames, or partial charges. It is straightforwardly transferable to any system including periodic boundary conditions as well as finite systems.
Notably, the training methodology is also general for any 3x3 tensor and thus the very same approach can also be used to train polarizability tensors which are required to calculate Raman spectra from MD simulations without the usual drastic approximations. Having access to the atomic polar tensor and polarizability tensors will furthermore also allow one to calculate accurate sum-frequency-generation spectra from MD simulations. To compute the IR spectrum of liquid ambient water exactly, a surprisingly small training data set is required which foreshadows that the method can also be used with more expensive electronic structure calculations. At the current stage, the model offers a quick and straightforward way to successfully train and predict IR spectra from molecular dynamics simulations at DFT accuracy, which will be utilized in future applications. A general toolkit to automatically and efficiently train an APTNN on available (machine learned) MD trajectories is currently being developed.