Saving SCF results on disk and SCF checkpoints
For longer DFT calculations it is pretty standard to run them on a cluster in advance and to perform postprocessing (band structure calculation, plotting of density, etc.) at a later point and potentially on a different machine.
To support such workflows DFTK offers the two functions save_scfres
and load_scfres
, which allow to save the data structure returned by self_consistent_field
on disk or retrieve it back into memory, respectively. For this purpose DFTK uses the JLD2.jl file format and Julia package.
As JLD2 is an optional dependency of DFTK these three functions are only available once one has both imported DFTK and JLD2 (using DFTK
and using JLD2
).
The data format in which DFTK saves data as well as the general interface of the load_scfres
and save_scfres
pair of functions are not yet fully matured. If you use the functions or the produced files expect that you need to adapt your routines in the future even with patch version bumps.
To illustrate the use of the functions in practice we will compute the total energy of the O₂ molecule at PBE level. To get the triplet ground state we use a collinear spin polarisation (see Collinear spin and magnetic systems for details) and a bit of temperature to ease convergence:
using DFTK
using LinearAlgebra
using JLD2
d = 2.079 # oxygen-oxygen bondlength
a = 9.0 # size of the simulation box
lattice = a * I(3)
O = ElementPsp(:O; psp=load_psp("hgh/pbe/O-q6.hgh"))
atoms = [O, O]
positions = d / 2a * [[0, 0, 1], [0, 0, -1]]
magnetic_moments = [1., 1.]
Ecut = 10 # Far too small to be converged
model = model_DFT(lattice, atoms, positions;
functionals=PBE(),
temperature=0.02, smearing=Smearing.Gaussian(),
magnetic_moments)
basis = PlaneWaveBasis(model; Ecut, kgrid=[1, 1, 1])
scfres = self_consistent_field(basis, tol=1e-2, ρ=guess_density(basis, magnetic_moments))
save_scfres("scfres.jld2", scfres);
n Energy log10(ΔE) log10(Δρ) Magnet Diag Δtime
--- --------------- --------- --------- ------ ---- ------
1 -27.64381002978 -0.13 0.001 6.5 101ms
2 -28.92294814681 0.11 -0.82 0.673 2.0 82.2ms
3 -28.93098300057 -2.10 -1.14 1.173 2.5 88.3ms
4 -28.93765514520 -2.18 -1.18 1.766 2.0 75.5ms
5 -28.93952159106 -2.73 -1.46 1.998 1.5 75.7ms
6 -28.93960009118 -4.11 -2.04 1.979 1.0 82.6ms
scfres.energies
Energy breakdown (in Ha):
Kinetic 16.7695503
AtomicLocal -58.4928425
AtomicNonlocal 4.7106323
Ewald -4.8994689
PspCorrection 0.0044178
Hartree 19.3598575
Xc -6.3906715
Entropy -0.0010751
total -28.939600091180
The scfres.jld2
file could now be transferred to a different computer, Where one could fire up a REPL to inspect the results of the above calculation:
using DFTK
using JLD2
loaded = load_scfres("scfres.jld2")
propertynames(loaded)
(:α, :history_Δρ, :converged, :occupation, :occupation_threshold, :algorithm, :basis, :runtime_ns, :n_iter, :n_matvec, :history_Etot, :εF, :energies, :ρ, :timedout, :n_bands_converge, :eigenvalues, :ψ, :ham)
loaded.energies
Energy breakdown (in Ha):
Kinetic 16.7695503
AtomicLocal -58.4928425
AtomicNonlocal 4.7106323
Ewald -4.8994689
PspCorrection 0.0044178
Hartree 19.3598575
Xc -6.3906715
Entropy -0.0010751
total -28.939600091180
Since the loaded data contains exactly the same data as the scfres
returned by the SCF calculation one could use it to plot a band structure, e.g. plot_bandstructure(load_scfres("scfres.jld2"))
directly from the stored data.
Notice that both load_scfres
and save_scfres
work by transferring all data to/from the master process, which performs the IO operations without parallelisation. Since this can become slow, both functions support optional arguments to speed up the processing. An overview:
save_scfres("scfres.jld2", scfres; save_ψ=false)
avoids saving the Bloch wave, which is usually faster and saves storage space.load_scfres("scfres.jld2", basis)
avoids reconstructing the basis from the file, but uses the passed basis instead. This save the time of constructing the basis twice and allows to specify parallelisation options (via the passed basis). Usually this is useful for continuing a calculation on a supercomputer or cluster.
See also the discussion on Input and output formats on JLD2 files.
Checkpointing of SCF calculations
A related feature, which is very useful especially for longer calculations with DFTK is automatic checkpointing, where the state of the SCF is periodically written to disk. The advantage is that in case the calculation errors or gets aborted due to overrunning the walltime limit one does not need to start from scratch, but can continue the calculation from the last checkpoint.
The easiest way to enable checkpointing is to use the kwargs_scf_checkpoints
function, which does two things. (1) It sets up checkpointing using the ScfSaveCheckpoints
callback and (2) if a checkpoint file is detected, the stored density is used to continue the calculation instead of the usual atomic-orbital based guess. In practice this is done by modifying the keyword arguments passed to # self_consistent_field
appropriately, e.g. by using the density or orbitals from the checkpoint file. For example:
checkpointargs = kwargs_scf_checkpoints(basis; ρ=guess_density(basis, magnetic_moments))
scfres = self_consistent_field(basis; tol=1e-2, checkpointargs...);
n Energy log10(ΔE) log10(Δρ) Magnet α Diag Δtime
--- --------------- --------- --------- ------ ---- ---- ------
1 -27.64648759697 -0.13 0.001 0.80 6.0 100ms
2 -28.92283635441 0.11 -0.82 0.676 0.80 2.0 90.5ms
3 -28.93100459035 -2.09 -1.14 1.176 0.80 2.0 90.4ms
4 -28.93766211589 -2.18 -1.19 1.768 0.80 2.0 83.3ms
5 -28.93952655679 -2.73 -1.42 1.998 0.80 2.0 87.4ms
6 -28.93959689844 -4.15 -1.97 1.977 0.80 1.0 77.9ms
7 -28.93961187812 -4.82 -2.82 1.986 0.80 1.0 93.5ms
Notice that the ρ
argument is now passed to kwargsscfcheckpoints instead. If we run in the same folder the SCF again (here using a tighter tolerance), the calculation just continues.
checkpointargs = kwargs_scf_checkpoints(basis; ρ=guess_density(basis, magnetic_moments))
scfres = self_consistent_field(basis; tol=1e-3, checkpointargs...);
n Energy log10(ΔE) log10(Δρ) Magnet α Diag Δtime
--- --------------- --------- --------- ------ ---- ---- ------
1 -28.93961280357 -3.34 1.986 0.80 11.5 159ms
Since only the density is stored in a checkpoint (and not the Bloch waves), the first step needs a slightly elevated number of diagonalizations. Notice, that reconstructing the checkpointargs
in this second call is important as the checkpointargs
now contain different data, such that the SCF continues from the checkpoint. By default checkpoint is saved in the file dftk_scf_checkpoint.jld2
, which can be changed using the filename
keyword argument of kwargs_scf_checkpoints
. Note that the file is not deleted by DFTK, so it is your responsibility to clean it up. Further note that warnings or errors will arise if you try to use a checkpoint, which is incompatible with your calculation.
We can also inspect the checkpoint file manually using the load_scfres
function and use it manually to continue the calculation:
oldstate = load_scfres("dftk_scf_checkpoint.jld2")
scfres = self_consistent_field(oldstate.basis, ρ=oldstate.ρ, ψ=oldstate.ψ, tol=1e-4);
n Energy log10(ΔE) log10(Δρ) Magnet Diag Δtime
--- --------------- --------- --------- ------ ---- ------
1 -28.93876950346 -2.31 1.985 5.0 93.2ms
2 -28.93948686823 -3.14 -2.70 1.985 1.0 69.7ms
3 -28.93961089988 -3.91 -2.57 1.985 4.0 98.3ms
4 -28.93961185299 -6.02 -2.69 1.985 1.0 79.9ms
5 -28.93961269372 -6.08 -2.88 1.985 1.0 68.6ms
6 -28.93961275931 -7.18 -2.88 1.985 1.0 68.1ms
7 -28.93961285361 -7.03 -2.91 1.985 1.0 68.1ms
8 -28.93961308913 -6.63 -3.15 1.985 1.0 70.1ms
9 -28.93961312750 -7.42 -3.28 1.985 1.0 68.1ms
10 -28.93961314419 -7.78 -3.39 1.985 1.5 72.4ms
11 -28.93961316002 -7.80 -3.63 1.985 1.0 84.2ms
12 -28.93961316911 -8.04 -3.91 1.985 1.5 75.6ms
13 -28.93961317165 -8.59 -4.24 1.985 2.0 82.0ms
Some details on what happens under the hood in this mechanism: When using the kwargs_scf_checkpoints
function, the ScfSaveCheckpoints
callback is employed during the SCF, which causes the density to be stored to the JLD2 file in every iteration. When reading the file, the kwargs_scf_checkpoints
transparently patches away the ψ
and ρ
keyword arguments and replaces them by the data obtained from the file. For more details on using callbacks with DFTK's self_consistent_field
function see Monitoring self-consistent field calculations.
(Cleanup files generated by this notebook)
rm("dftk_scf_checkpoint.jld2")
rm("scfres.jld2")