Manual: Add sub-directories to methods section

This commit is contained in:
Ole Schütt 2023-12-10 16:31:38 +01:00 committed by Ole Schütt
parent 85135fe7db
commit 668feff676
42 changed files with 26 additions and 35 deletions

View file

@ -0,0 +1,76 @@
# Nequip and Allegro
This
[Colab tutorial](https://colab.research.google.com/github/gabriele16/cp2k/blob/nequip-cp2k-colab/colab/allegro-cp2k-tutorial.ipynb)
illustrates how to train an equivariant neural network interatomic potential for bulk water using
the Allegro framework. You will learn how to train a model, deploy it in production, and run
molecular dynamics simulations in CP2K. The training and inference will be carried out on the GPU
provided by the Colab environment.
Allegro is designed for constructing highly accurate and scalable interatomic potentials for
molecular dynamics simulations. The methodology is described in detail in this paper
([](#Musaelian2023)). An open-source package that implements
[Allegro](https://github.com/mir-group/allegro), built on the
[Nequip framework](https://github.com/mir-group/nequip) was developed by the Allegro and NequIP
authors, A. Musaelian, S. Batzner, A. Johansson, L. Sun, C. J. Owen, M. Kornbluth, B. Kozinsky.
## Input Section
Inference in CP2K is performed through the
[NEQUIP](#CP2K_INPUT.FORCE_EVAL.MM.FORCEFIELD.NONBONDED.NEQUIP) and
[ALLEGRO](#CP2K_INPUT.FORCE_EVAL.MM.FORCEFIELD.NONBONDED.ALLEGRO) sections. As an example, the
relevant section for Allegro (or similarly for NequIP) is:
```none
&ALLEGRO
ATOMS Si
PARM_FILE_NAME Allegro/si-deployed.pth
UNIT_COORDS angstrom
UNIT_ENERGY eV
UNIT_FORCES eV*angstrom^-1
&END ALLEGRO
```
where the `si-deployed.pth` refers to the PyTorch model that was deployed using the Allegro
framework, and the `UNIT` tags refer to the units of the coordinates, energy and forces of the model
itself. An example for the full input file can be found in the
[Colab tutorial](https://colab.research.google.com/github/gabriele16/cp2k/blob/nequip-cp2k-colab/colab/allegro-cp2k-tutorial.ipynb)
and on the regtests, see
[Allegro_si_MD.inp](https://github.com/cp2k/cp2k/blob/master/tests/Fist/regtest-allegro/Allegro_si_MD.inp)
### Input details
The tag [ATOMS](#CP2K_INPUT.FORCE_EVAL.MM.FORCEFIELD.NONBONDED.ALLEGRO.ATOMS) expects a list of
elements/kinds in a way and order that is consistent with the YAML file of NequIP and Allegro. If
this is not done unphysical results will be obtained. Additionally, the atomic coordinates in the
[COORD](#CP2K_INPUT.FORCE_EVAL.SUBSYS.COORD) or [TOPOLOGY](#CP2K_INPUT.FORCE_EVAL.SUBSYS.TOPOLOGY)
section have to be provided in a way that is consistent with the YAML file. If this is not done
unphysical results will be obtained. Spotting such issues is quite straightforward as the energy is
significantly wrong. For example, by inverting the order of one of the elements in the test
`regtest-nequip/NequIP_water.inp`, the error with respect to the reference value is of the order of
1 eV. Additionally, running MD leads rapidly to highly unstable simulations.
## Compiling CP2K with LibTorch
Running with NequIP or Allegro requires compiling CP2K with the libtorch library. For the CP2K
binaries running on CPUs installing the toolchain using the flag `--with-libtorch` is enough. To
benefit from (often significant) GPU acceleration, the precompiled Libtorch library for CUDA can be
obtained at [https://pytorch.org](https://pytorch.org), for example for CUDA 11.8:
```shell
wget https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcu118.zip
```
After extracting the libtorch CUDA binaries, the toolchain script `./install_cp2k_toolchain.sh` can
be run providing the appropriate path with the flag `--with-libtorch=<path-to-libtorch-cuda>`.
## Further Resources
For additional references on NequIP, Allegro and equivariant neural networks (e3nn) see:
- Allegro paper [](#Musaelian2023) and code <https://github.com/mir-group/allegro>
- NequIP paper [](#Batzner2022) and code <https://github.com/mir-group/nequip>
- A Tutorial on LAMMPS by the NequIP/Allegro authors is found at the Colab notebook
[here](https://colab.research.google.com/drive/1yq2UwnET4loJYg_Fptt9kpklVaZvoHnq)
- For an introduction to e3nn see [e3nn.org](https://e3nn.org) and
[doi:10.5281/zenodo.7430260](https://dx.doi.org/10.5281/zenodo.7430260)

View file

@ -0,0 +1,11 @@
# Machine Learning
```{toctree}
---
titlesonly:
maxdepth: 1
---
allegro
Neural Network Potentials <https://www.cp2k.org/tools:aml>
pao-ml
```

View file

@ -0,0 +1,198 @@
# PAO-ML
PAO-ML stands for Polarized Atomic Orbitals from Machine Learning. It uses machine learning to
generate geometry adopted small basis sets. It also provides exact ionic forces. The scheme can
serve as an almost drop-in replacement for conventional basis sets to speedup otherwise standard DFT
calculations. The method is similar to semi-empirical models based on minimal basis sets, but offers
improved accuracy and quasi-automatic parameterization. However, the method is still in an early
stage - so use with caution. For more information see: [](#Schuett2018).
## Step 1: Obtain training structures
The PAO-ML scheme takes a set of training structures as input. For each of these structures, the
variational PAO basis is determined via an explicit optimization. The training structures should be
much smaller than the target system, but large enough to contain all the *motifs* of the larger
system. For liquids a good way to obtain structures is to run an MD of a smaller box.
## Step 2: Calculate reference data in primary basis
Choose a primary basis set, e.g. `DZVP-MOLOPT-GTH` and perform a full
[LS_SCF](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF) optimization. You should also enable
[RESTART_WRITE](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.RESTART_WRITE) to save the final density matrix.
It can be used to speed up the next step significantly.
## Step 3: Optimize PAO basis for training structures
Choose a [PAO_BASIS_SIZE](#CP2K_INPUT.FORCE_EVAL.SUBSYS.KIND.PAO_BASIS_SIZE) for each atomic kind.
Good results can already be optained with a minimal basis sets. Slightly larger-than-minimal PAO
basis sets can significantly increase the accuracy. However, they are also tougher to optimize and
machine learn.
Most of the PAO settings are in the [PAO](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO) sections:
```
&PAO
EPS_PAO 1.0E-7 ! convergence threshold of PAO optimization
MAX_PAO 10000 ! minimal PAO basis usually converge withing 2000 steps.
MAX_CYCLES 500 ! tunning parameter for PAO optimization scheme
MIXING 0.5 ! tunning parameter for PAO optimization scheme
PREOPT_DM_FILE primay_basis.dm ! restart DM from primary basis for great speedup
LINPOT_REGULARIZATION_DELTA 1E-6 !!!! Critical parameter for accuracy vs learnability trade-off !!!!
LINPOT_REGULARIZATION_STRENGTH 1E-3 ! rather insensitive parameter, 1e-3 works usually
REGULARIZATION 1.0E-3 ! rather insensitive parameter, 1e-3 works usually
PRECONDITION YES ! not important, don't touch
LINPOT_PRECONDITION_DELTA 0.01 ! not important, don't touch
LINPOT_INITGUESS_DELTA 1E+10 ! not important, don't touch
&PRINT
&RESTART
BACKUP_COPIES 1 ! write restart files, just in case
&END RESTART
&END PRINT
&END PAO
```
Settings for individual atomic kinds are in the [KIND](#CP2K_INPUT.FORCE_EVAL.SUBSYS.KIND) section:
```
&KIND H
PAO_BASIS_SIZE 1 ! set this to at least the minimal basis size
&PAO_POTENTIAL
MAXL 4 ! 4 works usually
BETA 2.0 ! 2 work usually, but is worth exploring in case of accuracy or learnability issues.
&END PAO_POTENTIAL
&END KIND
```
### Tuning the PAO Optimization
Finding the optimal PAO basis poses an intricate minimization problem, because the rotation matrix U
and the Kohn-Sham matrix H have to be optimized in a self-consistent manner. In order to speedup the
optimization, the Kohn-Sham matrix is only updated occasionally while most time is spend on
optimizing U. This alternating scheme is controlled by two input parameters:
- The frequency with which H is recalculated is determined by
[MAX_CYCLES](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.MAX_CYCLES).
- Overshooting during the U optimization is damped via
[MIXING](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.MIXING).
The progress of the PAO optimization can be tracked from lines that start with `PAO| step`. The
columns have the following meaning:
```
step-num energy conv-crit. step-length time
PAO| step 1121 -186.164843303 0.227E-06 0.120E+01 1.440
```
- The step number counts the number of energy evaluation, ie. the number of U matrices probed. It
can increase with different intervals, when the
[ADAPTive](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.LINE_SEARCH.METHOD) line-search method is used.
When the step number reaches [MAX_PAO](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.MAX_PAO) then the
optimization is terminated prematurely.
- The energy is the quantity that is optimized. It contains **only the first order term** of the
total energy, ie. $Tr\[HP\]$, but shares the same variational minima. It furthermore contains the
contributions from the various regularization terms.
- The convergence criterion is the norm of the gradient normalized by system size. It is compared
against [EPS_PAO](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.EPS_PAO) to decided if the PAO
optimization has converged. The overall optimization is terminated if this convergence criterion
is reached within two steps after updating the Kohn-Sham matrix.
- The step length is the outcome of the line search. It should be of order 1. If it starts to behave
erratically towards the end of the optimization, this indicates that further optimization is
hindered by numerical accuracy e.g. from
[EPS_FILTER](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.EPS_FILTER) or
[EPS_SCF](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.EPS_SCF).
- The time is the time spend on this optimization step in seconds. This number can varry accordingly
to the number of performed lines search steps.
## Step 4: Optimize machine learning hyper-parameters
For the simulation of larger systems the PAO-ML scheme infers new PAO basis sets from the training
data. For this two heuristics are employed: A
[descriptor](<https://en.wikipedia.org/wiki/Feature_(machine_learning)>) and an inference algorithm.
Currently, only one simple descriptor and
[Gaussian processes](https://en.wikipedia.org/wiki/Gaussian_process) are implemented. However, this
part offers great opportunities for future research.
In order to obtain good results from the learning machinery a small number of so-called
[hyperparameters](https://en.wikipedia.org/wiki/Hyperparameter) have to be carefully tuned for each
application. For the current implementation this includes the
[GP_SCALE](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.MACHINE_LEARNING.GP_SCALE) and the descriptor's
[BETA](#CP2K_INPUT.FORCE_EVAL.SUBSYS.KIND.PAO_DESCRIPTOR.BETA) and
[SCREENING](#CP2K_INPUT.FORCE_EVAL.SUBSYS.KIND.PAO_DESCRIPTOR.SCREENING).
For the optimization of the hyper-parameter exists no gradient, hence one has to use a
derivative-free method like the one by [Powell](https://en.wikipedia.org/wiki/Powell%27s_method). A
versatile implementation is e.g. the
[scriptmini](https://github.com/cp2k/cp2k/tree/master/tools/scriptmini) tool. A good optimization
criterion is the variance of the energy difference wrt. the primary basis across the training set.
Alternatively, atomic forces could be compared. Despite the missing gradients, this optimization is
rather quick because it only performs calculations in the small PAO basis set.
## Step 5: Run simulation with PAO-ML
Most of the PAO-ML settings are in the
[PAO/MACHINE_LEARNING](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.MACHINE_LEARNING) sections:
```
&PAO
MAX_PAO 0 ! use PAO basis as predicted by ML, required for correct forces
PENALTY_STRENGTH 0.0 ! disable penalty, required for correct forces
&MACHINE_LEARNING
GP_SCALE 0.46 !!! critical tuning parameter - depends also on descriptor settings !!!
GP_NOISE_VAR 0.0001 ! insensitive parameter
METHOD GAUSSIAN_PROCESS ! only implemented method - opportunity for future research
DESCRIPTOR OVERLAP ! only implemented method - opportunity for future research
PRIOR MEAN ! try once ZERO - makes usually no difference
TOLERANCE 1000.0 ! disable check for max variance of GP prediction
&TRAINING_SET
../training/Frame0000/calc_pao_ref-1_0.pao
../training/Frame0100/calc_pao_ref-1_0.pao
../training/Frame0200/calc_pao_ref-1_0.pao
! add more ...
&END TRAINING_SET
&END MACHINE_LEARNING
&END PAO
```
Settings for individual atomic kinds are again in the [KIND](#CP2K_INPUT.FORCE_EVAL.SUBSYS.KIND)
section:
```
&KIND H
PAO_BASIS_SIZE 1 ! use same settings as for training
&PAO_POTENTIAL
MAXL 4 ! use same settings as for training
BETA 2.0 ! use same settings as for training
&END PAO_POTENTIAL
&PAO_DESCRIPTOR
BETA 0.16 !!! important ML hyper-parameter !!!
SCREENING 0.66 !!! important ML hyper-parameter !!!
WEIGHT 1.0 ! usually not needed when BETA and SCREENING are choose properly
&END PAO_DESCRIPTOR
&END KIND
```
## Debugging accuracy vs learnability trade-off
When optimizing the PAO reference data in Step 3 one has to make a trade-off between accuracy and
learnability. Good learnability means that similar structures leads to similar PAO parameters. In
other words the PAO parameters should depend smoothly on the atomic positions. In general, the
settings presented above should yield good results. However, if problems arise in the later machine
learning steps, this might be the culprit.
Unfortunately, there is not yet a simple way to assess learnability. One way to investigate is to
create a set of structures along a reaction coordinate, e.g. a dimer dissociation. One can then plot
the numbers from the `Xblock` in the `.pao` files vs. the reaction coordinate.
The most critical parameters for learnability are
[LINPOT_REGULARIZATION_DELTA](#CP2K_INPUT.FORCE_EVAL.DFT.LS_SCF.PAO.LINPOT_REGULARIZATION_DELTA) and
the potential's [BETA](#CP2K_INPUT.FORCE_EVAL.SUBSYS.KIND.PAO_POTENTIAL.BETA).