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_PBE(lattice, atoms, positions; 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.64422225221 -0.13 0.001 6.5 90.9ms
2 -28.92284127513 0.11 -0.82 0.680 2.0 107ms
3 -28.93107731104 -2.08 -1.14 1.183 2.5 84.0ms
4 -28.93772350292 -2.18 -1.19 1.772 2.0 71.0ms
5 -28.93944440106 -2.76 -1.23 2.000 2.0 74.8ms
6 -28.93959091426 -3.83 -1.91 1.976 1.0 65.2ms
7 -28.93961206920 -4.67 -3.22 1.985 1.0 66.7msscfres.energiesEnergy breakdown (in Ha):
Kinetic 16.7721017
AtomicLocal -58.4953967
AtomicNonlocal 4.7094715
Ewald -4.8994689
PspCorrection 0.0044178
Hartree 19.3614174
Xc -6.3912849
Entropy -0.0008701
total -28.939612069200The 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, :history_Etot, :εF, :energies, :ρ, :n_bands_converge, :eigenvalues, :ψ, :ham)loaded.energiesEnergy breakdown (in Ha):
Kinetic 16.7721017
AtomicLocal -58.4953967
AtomicNonlocal 4.7094715
Ewald -4.8994689
PspCorrection 0.0044178
Hartree 19.3614174
Xc -6.3912849
Entropy -0.0008701
total -28.939612069200Since 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.64466988209 -0.13 0.001 0.80 6.0 93.5ms
2 -28.92282519043 0.11 -0.82 0.677 0.80 2.0 101ms
3 -28.93106507618 -2.08 -1.14 1.183 0.80 2.0 84.3ms
4 -28.93768919511 -2.18 -1.19 1.770 0.80 2.0 77.4ms
5 -28.93954098676 -2.73 -1.49 1.997 0.80 2.0 80.3ms
6 -28.93959979505 -4.23 -2.02 1.979 0.80 1.0 71.7msNotice 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.93960591658 -2.90 1.985 0.80 8.0 131ms
2 -28.93961161873 -5.24 -3.32 1.985 0.80 1.0 73.5msSince 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.93865429130 -2.26 1.985 5.5 106ms
2 -28.93947986663 -3.08 -2.69 1.985 1.0 66.2ms
3 -28.93961094116 -3.88 -2.53 1.985 3.5 96.5ms
4 -28.93961166181 -6.14 -2.63 1.985 1.0 64.9ms
5 -28.93961262484 -6.02 -2.87 1.985 1.0 78.4ms
6 -28.93961264703 -7.65 -2.82 1.985 1.0 63.2ms
7 -28.93961266362 -7.78 -2.82 1.985 1.0 63.5ms
8 -28.93961292994 -6.57 -2.91 1.985 1.0 64.6ms
9 -28.93961307957 -6.82 -3.06 1.985 1.0 64.1ms
10 -28.93961311699 -7.43 -3.20 1.985 1.0 64.2ms
11 -28.93961314635 -7.53 -3.35 1.985 1.0 88.4ms
12 -28.93961316395 -7.75 -3.63 1.985 1.5 70.8ms
13 -28.93961316774 -8.42 -4.09 1.985 2.0 76.5msSome 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")