On-the-fly training using ASE

This is a quick introduction of how to set up our ASE-OTF interface to train a force field. We will train a force field model for diamond. To run the on-the-fly training, we will need to

  1. Create a supercell with ASE Atoms object
  2. Set up FLARE ASE calculator, including the kernel functions, hyperparameters, cutoffs for Gaussian process, and mapping parameters (if Mapped Gaussian Process is used)
  3. Set up DFT ASE calculator. Here we will give an example of Quantum Espresso
  4. Set up on-the-fly training with ASE MD engine

Please make sure you are using the LATEST FLARE code in our master branch.

Step 1: Set up supercell with ASE

Here we create a 2x1x1 supercell with lattice constant 3.855, and randomly perturb the positions of the atoms, so that they will start MD with non-zero forces.

[1]:
import numpy as np
from ase import units
from ase.spacegroup import crystal
from ase.build import bulk

np.random.seed(12345)

a = 3.52678
super_cell = bulk('C', 'diamond', a=a, cubic=True)

Step 2: Set up FLARE calculator

Now let’s set up our Gaussian process model in the same way as introduced before

[2]:
from flare.gp import GaussianProcess
from flare.utils.parameter_helper import ParameterHelper

# set up GP hyperparameters
kernels = ['twobody', 'threebody'] # use 2+3 body kernel
parameters = {'cutoff_twobody': 5.0,
              'cutoff_threebody': 3.5}
pm = ParameterHelper(
    kernels = kernels,
    random = True,
    parameters=parameters
)

hm = pm.as_dict()
hyps = hm['hyps']
cut = hm['cutoffs']
print('hyps', hyps)

gp_model = GaussianProcess(
    kernels = kernels,
    component = 'sc', # single-component. For multi-comp, use 'mc'
    hyps = hyps,
    cutoffs = cut,
    hyp_labels = ['sig2','ls2','sig3','ls3','noise'],
    opt_algorithm = 'L-BFGS-B',
    n_cpus = 1
)
twobody0 cutoff is not define. it's going to use the universal cutoff.
threebody0 cutoff is not define. it's going to use the universal cutoff.
hyps [0.92961609 0.31637555 0.18391881 0.20456028 0.05      ]

Optional

If you want to use Mapped Gaussian Process (MGP), then set up MGP as follows

[3]:
from flare.mgp import MappedGaussianProcess

grid_params = {'twobody':   {'grid_num': [64]},
               'threebody': {'grid_num': [16, 16, 16]}}

mgp_model = MappedGaussianProcess(grid_params,
                                  unique_species = [6],
                                  n_cpus = 1,
                                  map_force = False,
                                  mean_only = False)

Now let’s set up FLARE’s ASE calculator. If you want to use MGP model, then set use_mapping = True and mgp_model = mgp_model below.

[4]:
from flare.ase.calculator import FLARE_Calculator

flare_calculator = FLARE_Calculator(gp_model,
                                    par = True,
                                    mgp_model = None,
                                    use_mapping = False)

super_cell.set_calculator(flare_calculator)

Step 3: Set up DFT calculator

For DFT calculator, you can use any calculator provided by ASE, e.g. Quantum Espresso (QE), VASP, etc.

For a quick illustration of our interface, we use the Lennard-Jones (LJ) potential as an example.

[5]:
from ase.calculators.lj import LennardJones
lj_calc = LennardJones()

Optional: alternatively, set up Quantum Espresso calculator

We also give the code below for setting up the ASE quantum espresso calculator, following the instruction on ASE website.

First, we need to set up our environment variable ASE_ESPRESSO_COMMAND to our QE executable, so that ASE can find this calculator. Then set up our input parameters of QE and create an ASE calculator

[6]:
import os
from ase.calculators.espresso import Espresso

# ---------------- set up executable ----------------
label = 'C'
input_file = label+'.pwi'
output_file = label+'.pwo'
no_cpus = 32
npool = 32
pw_loc = 'path/to/pw.x'

# serial
os.environ['ASE_ESPRESSO_COMMAND'] = f'{pw_loc} < {input_file} > {output_file}'

## parallel qe using mpirun
# os.environ['ASE_ESPRESSO_COMMAND'] = f'mpirun -np {no_cpus} {pw_loc} -npool {npool} < {input_file} > {output_file}'

## parallel qe using srun (for slurm system)
# os.environ['ASE_ESPRESSO_COMMAND'] = 'srun -n {no_cpus} --mpi=pmi2 {pw_loc} -npool {npool} < {input_file} > {output_file}'


# -------------- set up input parameters --------------
input_data = {'control':   {'prefix': label,
                            'pseudo_dir': './',
                            'outdir': './out',
                            'calculation': 'scf'},
              'system':    {'ibrav': 0,
                            'ecutwfc': 60,
                            'ecutrho': 360},
              'electrons': {'conv_thr': 1.0e-9,
                            'electron_maxstep': 100,
                            'mixing_beta': 0.7}}

# ----------------  pseudo-potentials -----------------
ion_pseudo = {'C': 'C.pz-rrkjus.UPF'}

# -------------- create ASE calculator ----------------
dft_calc = Espresso(pseudopotentials=ion_pseudo, label=label,
                    tstress=True, tprnfor=True, nosym=True,
                    input_data=input_data, kpts=(8, 8, 8))

Step 4: Set up On-The-Fly MD engine

Finally, our OTF is compatible with 5 MD engines that ASE supports: VelocityVerlet, NVTBerendsen, NPTBerendsen, NPT and Langevin. We can choose any of them, and set up the parameters based on ASE requirements. After everything is set up, we can run the on-the-fly training.

Note: Currently, only VelocityVerlet is tested on real system, NPT may have issue with pressure and stress.

Set up ASE_OTF training engine: 1. Initialize the velocities of atoms as 500K 2. Set up MD arguments as a dictionary based on ASE MD parameters. For VelocityVerlet, we don’t need to set up extra parameters.

E.g. for NVTBerendsen, we can set `md_kwargs = {'temperature': 500, 'taut': 0.5e3 * units.fs}`
[7]:
from ase import units
from ase.md.velocitydistribution import (MaxwellBoltzmannDistribution,
                                         Stationary, ZeroRotation)

temperature = 500
MaxwellBoltzmannDistribution(super_cell, temperature * units.kB)
Stationary(super_cell)  # zero linear momentum
ZeroRotation(super_cell)  # zero angular momentum

md_engine = 'VelocityVerlet'
md_kwargs = {}
  1. Set up parameters for On-The-Fly (OTF) training. The descriptions of the parameters are in ASE OTF module.
  2. Set up the ASE_OTF training engine, and run
  3. Check otf.out after the training is done.
[8]:
from flare.ase.otf import ASE_OTF

otf_params = {'init_atoms': [0, 1, 2, 3],
              'output_name': 'otf',
              'std_tolerance_factor': 2,
              'max_atoms_added' : 4,
              'freeze_hyps': 10}

test_otf = ASE_OTF(super_cell,
                   timestep = 1 * units.fs,
                   number_of_steps = 3,
                   dft_calc = lj_calc,
                   md_engine = md_engine,
                   md_kwargs = md_kwargs,
                   **otf_params)

test_otf.run()
INFO:otflog:2020-06-10 14:53:00.574573
INFO:otflog:
GaussianProcess Object
Number of cpu cores: 1
Kernel: ['twobody', 'threebody']
Training points: 0
Cutoffs: {'twobody': 5.0, 'threebody': 3.5}
Model Likelihood: None
Number of hyperparameters: 5
Hyperparameters_array: [0.92961609 0.31637555 0.18391881 0.20456028 0.05      ]
Hyperparameters:
sig2: 0.9296160928171479
ls2: 0.3163755545817859
sig3: 0.18391881167709445
ls3: 0.2045602785530397
noise: 0.05
Hyps_mask train_noise: True
Hyps_mask nspecie: 1
Hyps_mask twobody_start: 0
Hyps_mask ntwobody: 1
Hyps_mask threebody_start: 2
Hyps_mask nthreebody: 1
Hyps_mask kernels: ['twobody', 'threebody']
Hyps_mask cutoffs: {'twobody': 5.0, 'threebody': 3.5}

uncertainty tolerance: 2 times noise hyperparameter
timestep (ps): 0.09822694788464063
number of frames: 3
number of atoms: 8
system species: {'C'}
periodic cell:
[[3.52678 0.      0.     ]
 [0.      3.52678 0.     ]
 [0.      0.      3.52678]]

previous positions (A):
C        0.0000    0.0000    0.0000
C        0.8817    0.8817    0.8817
C        0.0000    1.7634    1.7634
C        0.8817    2.6451    2.6451
C        1.7634    0.0000    1.7634
C        2.6451    0.8817    2.6451
C        1.7634    1.7634    0.0000
C        2.6451    2.6451    0.8817
--------------------------------------------------------------------------------

INFO:otflog:
Calling DFT...

INFO:otflog:DFT run complete.
INFO:otflog:number of DFT calls: 1
INFO:otflog:wall time from start: 0.03 s
INFO:otflog:Adding atom [0, 1, 2, 3] to the training set.
INFO:otflog:Uncertainty: [0. 0. 0.]
INFO:otflog:
GP hyperparameters:
INFO:otflog:Hyp0 : sig2 = 0.9296
INFO:otflog:Hyp1 : ls2 = 0.3164
INFO:otflog:Hyp2 : sig3 = 0.1839
INFO:otflog:Hyp3 : ls3 = 0.2046
INFO:otflog:Hyp4 : noise = 0.0010
INFO:otflog:likelihood: 71.8658
INFO:otflog:likelihood gradient: [-1.79487637e-06 -6.53005436e-06 -8.57610391e-06 -1.76345959e-05
 -1.19999904e+04]
INFO:otflog:wall time from start: 7.23 s
INFO:otflog:
*-Frame: 0
Simulation Time: 0.0 ps
El            Position (A)                    DFT Force (ev/A)                  Std. Dev (ev/A)                  Velocities (A/ps)
C        0.0000    0.0000    0.0000        0.0000    0.0000    0.0000        0.0000    0.0000    0.0000        0.0417    0.0382    0.0035
C        0.8817    0.8817    0.8817        0.0000    0.0000    0.0000        0.0000    0.0000    0.0000       -0.0135    0.0222    0.0320
C        0.0000    1.7634    1.7634        0.0000    0.0000   -0.0000        0.0000    0.0000    0.0000        0.0204   -0.0705    0.0194
C        0.8817    2.6451    2.6451       -0.0000   -0.0000   -0.0000        0.0000    0.0000    0.0000       -0.0013    0.0346    0.0278
C        1.7634    0.0000    1.7634        0.0000    0.0000   -0.0000        0.0000    0.0000    0.0000       -0.0676   -0.0129    0.0242
C        2.6451    0.8817    2.6451       -0.0000   -0.0000   -0.0000        0.0000    0.0000    0.0000       -0.0027   -0.0121   -0.0340
C        1.7634    1.7634    0.0000        0.0000   -0.0000    0.0000        0.0000    0.0000    0.0000        0.0658   -0.0278   -0.0497
C        2.6451    2.6451    0.8817       -0.0000   -0.0000    0.0000        0.0000    0.0000    0.0000       -0.0427    0.0283   -0.0231

temperature: 199.00 K
kinetic energy: 0.180060 eV

INFO:otflog:wall time from start: 7.23 s
INFO:otflog:--------------------------------------------------------------------------------
-Frame: 1
Simulation Time: 0.0982 ps
El            Position (A)                    GP Force (ev/A)                   Std. Dev (ev/A)                  Velocities (A/ps)
C        0.0082    0.0075    0.0007       -0.0000   -0.0000   -0.0000       16.1332   16.9877    6.7169        0.0417    0.0382    0.0035
C        0.8790    0.8861    0.8880        0.0000   -0.0000   -0.0000       15.3645    9.5079   17.2434       -0.0135    0.0222    0.0320
C        0.0040    1.7495    1.7672       -0.0000    0.0000   -0.0000       10.8208   37.7660    9.7448        0.0204   -0.0705    0.0194
C        0.8814    2.6519    2.6505       -0.0000   -0.0000    0.0000        9.5129   20.9930    4.7440       -0.0013    0.0346    0.0278
C        1.7501   -0.0025    1.7682        0.0000    0.0000   -0.0000       23.4305   10.7241   12.5271       -0.0676   -0.0129    0.0242
C        2.6446    0.8793    2.6384        0.0000   -0.0000    0.0000        3.1513   13.4187    7.3315       -0.0027   -0.0121   -0.0340
C        1.7763    1.7579   -0.0098       -0.0000    0.0000    0.0000       20.5584    9.0554   17.2228        0.0658   -0.0278   -0.0497
C        2.6367    2.6506    0.8772        0.0000   -0.0000    0.0000       17.8540    7.7304   12.6641       -0.0427    0.0283   -0.0231

temperature: 199.00 K
kinetic energy: 0.180060 eV

INFO:otflog:wall time from start: 9.65 s
INFO:otflog:
Calling DFT...

INFO:otflog:DFT run complete.
INFO:otflog:number of DFT calls: 2
INFO:otflog:wall time from start: 9.67 s
INFO:otflog:
*-Frame: 1
Simulation Time: 0.0982 ps
El            Position (A)                    DFT Force (ev/A)                  Std. Dev (ev/A)                  Velocities (A/ps)
C        0.0164    0.0150    0.0013        0.0372   -0.0001   -0.0205       16.1332   16.9877    6.7169        0.0835    0.0764    0.0069
C        0.8763    0.8904    0.8943       -0.0669   -0.0327    0.0578       15.3645    9.5079   17.2434       -0.0274    0.0442    0.0641
C        0.0080    1.7356    1.7710        0.0322   -0.1194    0.0093       10.8208   37.7660    9.7448        0.0409   -0.1414    0.0388
C        0.8812    2.6587    2.6560        0.0353    0.0737   -0.0165        9.5129   20.9930    4.7440       -0.0025    0.0695    0.0555
C        1.7368   -0.0051    1.7729       -0.0330    0.0001    0.0250       23.4305   10.7241   12.5271       -0.1353   -0.0259    0.0486
C        2.6440    0.8770    2.6317       -0.0014    0.0643    0.0114        3.1513   13.4187    7.3315       -0.0053   -0.0238   -0.0680
C        1.7893    1.7525   -0.0195        0.0479    0.0129   -0.0207       20.5584    9.0554   17.2228        0.1317   -0.0555   -0.0994
C        2.6283    2.6562    0.8726       -0.0513    0.0013   -0.0459       17.8540    7.7304   12.6641       -0.0856    0.0565   -0.0465

temperature: 798.57 K
kinetic energy: 0.722563 eV

INFO:otflog:wall time from start: 9.68 s
INFO:otflog:mean absolute error: 0.0340 eV/A
INFO:otflog:mean absolute dft component: 0.0340 eV/A
INFO:otflog:Adding atom [6, 3, 4, 2] to the training set.
INFO:otflog:Uncertainty: [20.55839508  9.05540846 17.22284583]
INFO:otflog:
GP hyperparameters:
INFO:otflog:Hyp0 : sig2 = 0.7038
INFO:otflog:Hyp1 : ls2 = 2.0405
INFO:otflog:Hyp2 : sig3 = 0.0000
INFO:otflog:Hyp3 : ls3 = 9.6547
INFO:otflog:Hyp4 : noise = 0.0010
INFO:otflog:likelihood: 122.4930
INFO:otflog:likelihood gradient: [-1.34065483e+00 -1.79554908e-01 -4.94110742e-02  1.54534584e-10
 -1.82026091e+04]
INFO:otflog:wall time from start: 30.46 s
INFO:otflog:--------------------------------------------------------------------------------
-Frame: 2
Simulation Time: 0.196 ps
El            Position (A)                    GP Force (ev/A)                   Std. Dev (ev/A)                  Velocities (A/ps)
C        0.0247    0.0225    0.0020        0.0748   -0.0003   -0.0400        0.0008    0.0015    0.0008        0.0000    0.0000    0.0000
C        0.8735    0.8947    0.9007       -0.1357   -0.0645    0.1173        0.0014    0.0015    0.0010        0.0000    0.0000    0.0000
C        0.0121    1.7215    1.7748        0.0632   -0.2385    0.0151        0.0010    0.0015    0.0016        0.0000    0.0000    0.0000
C        0.8810    2.6657    2.6614        0.0692    0.1497   -0.0328        0.0011    0.0013    0.0010        0.0000    0.0000    0.0000
C        1.7235   -0.0076    1.7778       -0.0661   -0.0006    0.0515        0.0015    0.0013    0.0008        0.0000    0.0000    0.0000
C        2.6435    0.8748    2.6251       -0.0019    0.1297    0.0235        0.0009    0.0017    0.0010        0.0000    0.0000    0.0000
C        1.8023    1.7471   -0.0293        0.0980    0.0221   -0.0451        0.0015    0.0018    0.0012        0.0000    0.0000    0.0000
C        2.6197    2.6617    0.8679       -0.1015    0.0024   -0.0895        0.0013    0.0012    0.0012        0.0000    0.0000    0.0000

temperature: 0.00 K
kinetic energy: 0.000000 eV

INFO:otflog:wall time from start: 33.38 s
INFO:otflog:--------------------
INFO:otflog:Run complete.