Monitoring self-consistent field calculations

The self_consistent_field function takes as the callback keyword argument one function to be called after each iteration. This function gets passed the complete internal state of the SCF solver and can thus be used both to monitor and debug the iterations as well as to quickly patch it with additional functionality.

This example discusses a few aspects of the callback function taking again our favourite silicon example.

We setup silicon in an LDA model using AtomsBuilder to build a bulk silicon lattice, (see AtomsBase integration for more details.

using DFTK
using AtomsBuilder

system = attach_psp(bulk(:Si); Si="hgh/pbe/si-q4")
model  = model_DFT(system; functionals=LDA())
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 ScfSaveCheckpoints, 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 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. 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)
density_differences = Float64[];

The callback function itself gets passed a named tuple similar to the one returned by self_consistent_field, which contains the input and output density of the SCF step as ρin and ρout. Since the callback gets called both during the SCF iterations as well as after convergence just before self_consistent_field finishes we can both collect the data and initiate the plotting in one function.

using LinearAlgebra

function plot_callback(info)
    if info.stage == :finalize
        plot!(p, density_differences, label="|ρout - ρin|", markershape=:x)
    else
        push!(density_differences, norm(info.ρout - info.ρin))
    end
    info
end
callback = ScfDefaultCallback() ∘ plot_callback;

Notice that for constructing the callback function we chained the plot_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);

n     Energy            log10(ΔE)   log10(Δρ)   α      Diag   Δtime
---   ---------------   ---------   ---------   ----   ----   ------
  1   -7.775285893720                   -0.70   0.80    4.5    134ms
  2   -7.779119140560       -2.42       -1.51   0.80    1.0    146ms
  3   -7.779316193659       -3.71       -2.60   0.80    1.2   22.2ms
  4   -7.779349857228       -4.47       -2.80   0.80    2.8   28.2ms
  5   -7.779350276150       -6.38       -2.85   0.80    1.0   21.8ms
  6   -7.779350842311       -6.25       -4.45   0.80    1.0   21.7ms
  7   -7.779350856051       -7.86       -4.47   0.80    2.5   27.8ms
  8   -7.779350856122      -10.15       -5.17   0.80    1.0   22.1ms

… and show the plot

p

The info object passed to the callback contains not just the densities but also the complete Bloch wave (in ψ), the occupation, band eigenvalues and so on. See src/scf/self_consistent_field.jl for all currently available keys.