Tutorial 5: using machine learning to predict the screening parameters of water molecules

In this tutorial, we will train a machine-learning model to predict the screening parameters of water molecules directly from their orbital densities. To generate a trajectory with 20 different atomic configurations, we run a python script that applies random noise to the atomic positions of a water molecule. The resulting atomic positions are saved in a xyz file and are visualized below

The 20 snapshots generated with perturb_positions.py

Our goal in this tutorial is to perform Koopmans calculations on each of these 20 snapshots using a machine learning model to predict the screening parameters instead of calculating them ab initio.

Running a machine learning workflow

To predict the screening parameters with the machine learning model we must first train the model. In the following we will use the first five snapshots for training and then use the trained machine learning model to predict the screening parameters for the remaining 15 snapshots.

The input file for the machine learning workflow

The input file for performing this task can be downloaded here.

First, we have to specify that we want to perform Koopmans calculations on a whole trajectory of snapshots by setting the "task" keyword in the "workflow" block:

24    "workflow": {
25        "task": "trajectory",
26        "functional": "ki",

For this task, we don’t provide the "atomic_positions" directly to the input file since we don’t want to perform a Koopmans calculation on a single snapshot but on many snapshots. Instead, we provide an xyz file containing all the atomic positions of each snapshot that we would like to simulate

21    "atoms": {
22        "atomic_positions": {
23            "units": "angstrom",
24            "snapshots": "snapshots.xyz"

Finally, we have to provide a ml block with keywords specific to the machine learning model

13    "ml": {
14        "use_ml": true,
15        "n_max": 6,
16        "l_max": 6,
17        "r_min": 1.0,
18        "r_max": 4.0,
19        "number_of_training_snapshots": 5
20    },

To predict the screening parameters from the orbital densities, we have to translate the orbital densities into input vectors for the machine learning model. To do so, we decompose the orbital densities into radial basis functions \(g_{nl}(r)\) and angular basis functions \(Y_{ml}(\theta,\phi)\). This decomposition has the following four hyperparameters that we provided in the input file:

  • \(n_{max}\) determines the number of radial basis functions

  • \(l_{max}\) determines the number of angular basis functions

  • \(r_{min}\) determines the smallest cutoff radius for the radial basis functions

  • \(r_{max}\) determines the largest cutoff radius for the radial basis functions

In anticipation that the machine learning model will be most useful in extended systems (liquids or solids), we apply periodic boundary conditions and use maximally localized Wannier functions as our variational orbitals (despite the fact that our toy water model is not, in fact, a periodic system).

The output file for the machine learning workflow

Running this calculation, the output will show that we compute the screening parameters of the first five snapshots ab initio and add the results to our training data

49  Orbital 1
50  ---------
51   Running calc_alpha/orbital_1/dft_n-1... done
52    
53    predicted screening parameter:  1.00000
54    calculated screening parameter: 0.57472
55    absolute error:                 0.42528
56    
57    Adding this orbital to the training data

Then we use the trained model to predict the screening parameters of the remaining snapshots

619  Orbital 1
620  ---------
621    Predicting the screening parameter with the ML model

Using the script plot_5a.py we can plot predicted ionization potentials of the water molecule across the last 10 snapshots. Of course, they don’t necessarily correspond to anything physical because these configurations have been randomly generated. But in real applications the snapshots will correspond to something physical and the resulting ionization potentials will be meaningful.

The predicted ionization potential for 10 last snapshots

Here there is no way of telling if the model is correct – it has provided us with some screening parameters and we have to trust it. If we want to check if a machine learning model is working properly what we need to do is a convergence analysis with respect to the number of training data. This will be the goal of the following section.

Running a convergence analysis

The input file for the convergence analysis

The corresponding input file differs from the previous input file only in the "task" keyword

1    "workflow": {
2        "task": "convergence_ml",
3        "functional": "ki",

and the "number_of_training_snapshots"

 1    "ml": {
 2        "use_ml": true,
 3        "n_max": 6,
 4        "l_max": 6,
 5        "r_min": 1.0,
 6        "r_max": 4.0,
 7        "number_of_training_snapshots": 10,
 8        "quantities_of_interest": ["alphas", "evs"]
 9    },
10    "atoms": {
11        "atomic_positions": {
12            "units": "angstrom",

For the convergence_ml task, setting "number_of_training_snapshots": 10 means that we will perform the convergence analysis with respect to 1, 2, … , and 10 training snapshots and use the remaining snapshots (in this case snapshots 11 to 20) for testing.

The "quantities_of_interest" is the list of parameters with respect to which we would like to perform the convergence analysis. In addition to performing it only with respect to the screening parameters "alphas", we also perform it with respect to the eigenvalues ("evs"). The latter requires an additional final calculation for each snapshot and therefore takes slightly longer to run.

The output file for the convergence analysis

You should see that the workflow first computes the screening parameters ab-initio for the last 10 snapshots.

18 Obtaining ab-initio results for the last 10 snapshot(s)
19 =======================================================

Next, snapshot 1 is added to the training data.

732 Adding snapshot 1 to the training data
733 ======================================

After having trained the machine learning model on the orbitals of the first snapshot we use the trained model to predict the screening parameters of the last 10 snapshots and compare our results to the results from the ab initio computation.

849 Testing on the last 10 snapshot(s)
850 ==================================

Next, we add snapshot 2 to the training data.

1552 Adding snapshot 2 to the training data
1553 ======================================

Our model is now trained on the orbitals of 2 snapshots. We use this model again to predict the screening parameters of the last 10 snapshots and compare the results to the ab initio calculation. We repeat this procedure until we have added all 10 snapshots to the training data. Then we can have a look at the convergence of the mean absolute error of the predicted screening parameters:

Convergence of the MAE of the predicted screening parameters

and the convergence of the mean absolute error of the predicted orbital energies:

Convergence of the MAE of the predicted eigenenergies

(You can find these plots in the convergence_analysis/final_results/ subdirectory.) We can see that we converged to a reasonable accuracy after about 5 training snapshots (which corresponds to 20 occupied and 10 empty orbitals).

We can now also check (plot_5b.py) that the predicted ionization potentials match with the ionization potentials obtained from the ab-initio computation of the screening parameters:

The predicted and the omputed ionization potential for 10 last snapshots