Improve performance of IO and serialisation (#935)

.gitlab-ci.yml

@@ -9,7 +9,8 @@ stages:
 julia/1.9-n2:
   stage: test
   variables:
-    SCHEDULER_PARAMETERS: -N 1 -n 1 -c 16 --gres=gpu:a100:1 --qos=devel -p dgx -t 00:30:00 -A hpc-prf-dftkjl
+    # SCHEDULER_PARAMETERS: -N 1 -n 1 -c 16 --gres=gpu:a100:1 --qos=devel -p dgx -t 00:30:00 -A hpc-prf-dftkjl
+    SCHEDULER_PARAMETERS: -N 1 -n 1 -c 16 --gres=gpu:a100:1 --qos=devel -p dgx -t 01:30:00 -A hpc-prf-dftkjl
     JULIA_DEPOT_PATH: /scratch/hpc-prf-dftkjl/dftkjl01/.julia-ci
     JULIA_NUM_THREADS: 1  # GPU and multi-threading not yet compatible
   coverage: '/\(\d+.\d+\%\) covered/'
@@ -22,9 +23,12 @@ julia/1.9-n2:
     - module load lang/JuliaHPC/1.9.3-foss-2022a-CUDA-11.7.0
     - julia --project=. -e '
         using Pkg; Pkg.instantiate();
+        using DFTK;
+        using Preferences;
+        set_preferences!(DFTK, "precompile_workload" => false; force=true);
         Pkg.add("CUDA");
         using CUDA;
-        CUDA.set_runtime_version!(v"11.7"; local_toolkit=true)
+        CUDA.set_runtime_version!(v"11.7"; local_toolkit=true);
       '
     - julia --color=yes --project=. -e '
         using Pkg;

Project.toml

@@ -1,7 +1,7 @@
 name = "DFTK"
 uuid = "acf6eb54-70d9-11e9-0013-234b7a5f5337"
 authors = ["Michael F. Herbst <[email protected]>", "Antoine Levitt <[email protected]>"]
-version = "0.6.15"
+version = "0.6.16"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -27,7 +27,6 @@ LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
-OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
 PkgVersion = "eebad327-c553-4316-9ea0-9fa01ccd7688"
 Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
@@ -55,8 +54,8 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9"
-wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
+wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 
 [extensions]
 DFTKCUDAExt = "CUDA"
@@ -74,34 +73,33 @@ AbstractFFTs = "1"
 Artifacts = "1"
 AtomsBase = "0.3.1"
 Brillouin = "0.5.14"
-ChainRulesCore = "1.15"
 CUDA = "5"
+ChainRulesCore = "1.15"
 Dates = "1"
 DftFunctionals = "0.2"
 DocStringExtensions = "0.9"
 FFTW = "1.5"
 ForwardDiff = "0.10"
-GenericLinearAlgebra = "0.3"
 GPUArraysCore = "0.1"
+GenericLinearAlgebra = "0.3"
 Interpolations = "0.14, 0.15"
 IntervalArithmetic = "0.20"
-IterativeSolvers = "0.9"
 IterTools = "1"
+IterativeSolvers = "0.9"
 JLD2 = "0.4"
 JSON3 = "1"
 LazyArtifacts = "1.3"
 Libxc = "0.3.17"
+LineSearches = "7"
 LinearAlgebra = "1"
 LinearMaps = "3"
-LineSearches = "7"
 LoopVectorization = "0.12"
 MPI = "0.20.13"
 Markdown = "1"
 Optim = "1"
-OrderedCollections = "1"
 PeriodicTable = "1"
-Plots = "1"
 PkgVersion = "0.3"
+Plots = "1"
 Polynomials = "3, 4"
 PrecompileTools = "1"
 Preferences = "1"
@@ -119,9 +117,9 @@ TimerOutputs = "0.5.12"
 Unitful = "1"
 UnitfulAtomic = "1"
 Wannier = "0.3.2"
-wannier90_jll = "3.1"
 WriteVTK = "1"
 julia = "1.9"
+wannier90_jll = "3.1"
 
 [extras]
 ASEconvert = "3da9722f-58c2-4165-81be-b4d7253e8fd2"

docs/src/tricks/scf_checkpoints.jl

@@ -10,22 +10,18 @@
 # by [`self_consistent_field`](@ref) on disk or retrieve it back into memory,
 # respectively. For this purpose DFTK uses the
 # [JLD2.jl](https://github.com/JuliaIO/JLD2.jl) file format and Julia package.
-# For the moment this process is considered an experimental feature and
-# has a number of caveats, see the warnings below.
 #
-# !!! warning "Saving `scfres` is experimental"
-#     The [`load_scfres`](@ref) and [`save_scfres`](@ref) pair of functions
-#     are experimental features. This means:
+# !!! note "Availability of `load_scfres`, `save_scfres` and checkpointing"
+#     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 interface of these functions
-#       as well as the format in which the data is stored on disk can
-#       change incompatibly in the future. At this point we make no promises ...
-#     - JLD2 is not yet completely matured
-#       and it is recommended to only use it for short-term storage
-#       and **not** to archive scientific results.
-#     - If you are using the functions to transfer data between different
-#       machines ensure that you use the **same version of Julia, JLD2 and DFTK**
-#       for saving and loading data.
+# !!! warning "DFTK data formats are not yet fully matured"
+#     The data format in which DFTK saves data as well as the general interface
+#     of the [`load_scfres`](@ref) and [`save_scfres`](@ref) 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
@@ -55,7 +51,7 @@ save_scfres("scfres.jld2", scfres);
 #-
 scfres.energies
 
-# The `scfres.jld2` file could now be transfered to a different computer,
+# 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:
 
@@ -69,6 +65,19 @@ loaded.energies
 # 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](@ref) on JLD2 files.
 
 # ## Checkpointing of SCF calculations
 # A related feature, which is very useful especially for longer calculations with DFTK
@@ -77,42 +86,48 @@ loaded.energies
 # to overrunning the walltime limit one does not need to start from scratch,
 # but can continue the calculation from the last checkpoint.
 #
-# To enable automatic checkpointing in DFTK one needs to pass the `ScfSaveCheckpoints`
-# callback to [`self_consistent_field`](@ref), for example:
+# The easiest way to enable checkpointing is to use the [`kwargs_scf_checkpoints`](@ref)
+# function, which does two things. (1) It sets up checkpointing using the
+# [`ScfSaveCheckpoints`](@ref) 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`](@ref) appropriately, e.g. by using the density
+# or orbitals from the checkpoint file. For example:
 
-callback = DFTK.ScfSaveCheckpoints()
-scfres = self_consistent_field(basis;
-                               ρ=guess_density(basis, magnetic_moments),
-                               tol=1e-2, callback);
+checkpointargs = kwargs_scf_checkpoints(basis; ρ=guess_density(basis, magnetic_moments))
+scfres = self_consistent_field(basis; tol=1e-2, checkpointargs...)
 
-# Notice that using this callback makes the SCF go silent since the passed
-# callback parameter overwrites the default value (namely `DefaultScfCallback()`)
-# which exactly gives the familiar printing of the SCF convergence.
-# If you want to have both (printing and checkpointing) you need to chain
-# both callbacks:
-callback = DFTK.ScfDefaultCallback() ∘ DFTK.ScfSaveCheckpoints(; keep=true)
-scfres = self_consistent_field(basis;
-                               ρ=guess_density(basis, magnetic_moments),
-                               tol=1e-2, callback);
+# Notice that the `ρ` argument is now passed to kwargs_scf_checkpoints instead.
+# If we run in the same folder the SCF again (here using a tighter tolerance),
+# the calculation just continues.
 
-# For more details on using callbacks with DFTK's `self_consistent_field` function
-# see [Monitoring self-consistent field calculations](@ref).
+checkpointargs = kwargs_scf_checkpoints(basis; ρ=guess_density(basis, magnetic_moments))
+scfres = self_consistent_field(basis; tol=1e-3, checkpointargs...)
 
-# By default checkpoint is saved in the file `dftk_scf_checkpoint.jld2`, which is
-# deleted automatically once the SCF completes successfully. If one wants to keep
-# the file one needs to specify `keep=true` as has been done in the ultimate SCF
-# for demonstration purposes: now we can continue the previous calculation
-# from the last checkpoint as if the SCF had been aborted.
-# For this one just loads the checkpoint with [`load_scfres`](@ref):
+# 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`](@ref). 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-3);
+scfres   = self_consistent_field(oldstate.basis, ρ=oldstate.ρ, ψ=oldstate.ψ, tol=1e-4);
 
-# !!! note "Availability of `load_scfres`, `save_scfres` and `ScfSaveCheckpoints`"
-#     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`).
+# 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](@ref).
 
 # (Cleanup files generated by this notebook)
 rm("dftk_scf_checkpoint.jld2")

examples/convergence_study.jl

@@ -32,7 +32,7 @@ function run_scf(; a=5.0, Ecut, nkpt, tol)
     model  = model_LDA(lattice, atoms, position; temperature=1e-2)
     basis  = PlaneWaveBasis(model; Ecut, kgrid=(nkpt, nkpt, nkpt))
     println("nkpt = $nkpt Ecut = $Ecut")
-    self_consistent_field(basis; is_converged=DFTK.ScfConvergenceEnergy(tol))
+    self_consistent_field(basis; is_converged=ScfConvergenceEnergy(tol))
 end;
 
 # Moreover we define some parameters. To make the calculations run fast for the

examples/density_methods.jl

@@ -29,10 +29,8 @@ function silicon_scf(guess_method)
     model = model_LDA(lattice, atoms, positions)
     basis = PlaneWaveBasis(model; Ecut=12, kgrid=[4, 4, 4])
 
-    ρguess = guess_density(basis, guess_method)
-
-    is_converged = DFTK.ScfConvergenceEnergy(1e-10)
-    self_consistent_field(basis; is_converged, ρ=ρguess)
+    self_consistent_field(basis; ρ=guess_density(basis, guess_method),
+                          is_converged=ScfConvergenceEnergy(1e-10))
 end;
 
 results = Dict();

examples/geometry_optimization.jl

@@ -38,7 +38,7 @@ function compute_scfres(x)
     if isnothing(ρ)
         ρ = guess_density(basis)
     end
-    is_converged = DFTK.ScfConvergenceForce(tol / 10)
+    is_converged = ScfConvergenceForce(tol / 10)
     scfres = self_consistent_field(basis; ψ, ρ, is_converged, callback=identity)
     ψ = scfres.ψ
     ρ = scfres.ρ

examples/input_output.jl

@@ -47,6 +47,13 @@ basis = PlaneWaveBasis(model; Ecut=10, kgrid=(2, 2, 2))
 ρ0 = guess_density(basis, system)
 scfres = self_consistent_field(basis, ρ=ρ0);
 
+# !!! warning "DFTK data formats are not yet fully matured"
+#     The data format in which DFTK saves data as well as the general interface
+#     of the [`load_scfres`](@ref) and [`save_scfres`](@ref) 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.
+
 # ## Writing VTK files for visualization
 # For visualizing the density or the Kohn-Sham orbitals DFTK supports storing
 # the result of an SCF calculations in the form of VTK files.
@@ -97,10 +104,15 @@ save_bands("iron_afm_bands.json", bands)
 # stored on disk in form of an [JLD2.jl](https://github.com/JuliaIO/JLD2.jl) file.
 # This file can be read from other Julia scripts
 # as well as other external codes supporting the HDF5 file format
-# (since the JLD2 format is based on HDF5).
+# (since the JLD2 format is based on HDF5). This includes notably `h5py`
+# to read DFTK output from python.
 
 using JLD2
-save_scfres("iron_afm.jld2", scfres);
+save_scfres("iron_afm.jld2", scfres)
+
+# Saving such JLD2 files supports some options, such as `save_ψ=false`, which avoids saving
+# the Bloch waves (much faster and smaller files). Notice that JLD2 files can also be used
+# with [`save_bands`](@ref).
 
 # Since such JLD2 can also be read by DFTK to start or continue a calculation,
 # these can also be used for checkpointing or for transferring results

examples/publications/2020_silicon_scf_convergence.jl

@@ -31,7 +31,7 @@ function my_callback(info)
 end
 my_isconverged = info -> norm(info.ρout - info.ρin) < tol
 opts = (; callback=my_callback, is_converged=my_isconverged, maxiter, tol,
-        determine_diagtol=info -> diagtol)
+        diagtolalg=AdaptiveDiagtol(; diagtol_max=diagtol))
 
 global errs = []
 global gaps = []

examples/scf_callbacks.jl

@@ -21,22 +21,21 @@ basis  = PlaneWaveBasis(model; Ecut=5, kgrid=[3, 3, 3]);
 
 # DFTK already defines a few callback functions for standard
 # tasks. One example is the usual convergence table,
-# which is defined in the callback `ScfDefaultCallback`.
-# Another example is `ScfPlotTrace`, which records the total
-# energy at each iteration and uses it to plot the convergence
-# of the SCF graphically once it is converged.
-# For details and other callbacks
-# see [`src/scf/scf_callbacks.jl`](https://dftk.org/blob/master/src/scf/scf_callbacks.jl).
-#
-# !!! note "Callbacks are not exported"
-#     Callbacks are not exported from the DFTK namespace as of now,
-#     so you will need to use them, e.g., as `DFTK.ScfDefaultCallback`
-#     and `DFTK.ScfPlotTrace`.
-
+# which is defined in the callback [`ScfDefaultCallback`](@ref).
+# Another example is [`ScfSaveCheckpoints`](@ref), which stores the state
+# of an SCF at each iterations to allow resuming from a failed
+# calculation at a later point.
+# See [Saving SCF results on disk and SCF checkpoints](@ref) for details
+# how to use checkpointing with DFTK.
 
 # In this example we define a custom callback, which plots
 # the change in density at each SCF iteration after the SCF
-# has finished. For this we first define the empty plot canvas
+# has finished. This example is a bit artificial, since the norms
+# of all density differences is available as `scfres.history_Δρ`
+# after the SCF has finished and could be directly plotted, but
+# the following nicely illustrates the use of callbacks in DFTK.
+
+# To enable plotting we first define the empty canvas
 # and an empty container for all the density differences:
 using Plots
 p = plot(; yaxis=:log)
@@ -60,12 +59,14 @@ function plot_callback(info)
     end
     info
 end
-callback = DFTK.ScfDefaultCallback() ∘ plot_callback;
+callback = ScfDefaultCallback() ∘ plot_callback;
 
 # Notice that for constructing the `callback` function we chained the `plot_callback`
-# (which does the plotting) with the `ScfDefaultCallback`, such that when using
-# the `plot_callback` function with `self_consistent_field` we still get the usual
-# convergence table printed. We run the SCF with this callback …
+# (which does the plotting) with the `ScfDefaultCallback`. The latter is the function
+# responsible for printing the usual convergence table. Therefore if we simply did
+# `callback=plot_callback` the SCF would go silent. The chaining of both callbacks
+# (`plot_callback` for plotting and `ScfDefaultCallback()` for the convergence table)
+# makes sure both features are enabled. We run the SCF with the chained callback …
 scfres = self_consistent_field(basis; tol=1e-5, callback);
 
 # … and show the plot