Custom solvers

In this example, we show how to define custom solvers. Our system will again be silicon, because we are not very imaginative

using DFTK
using LinearAlgebra
using PseudoPotentialData
using AtomsBuilder

# We take very (very) crude parameters
pseudopotentials = PseudoFamily("dojo.nc.sr.lda.v0_4_1.standard.upf")
model = model_DFT(bulk(:Si); functionals=LDA(), pseudopotentials)
basis = PlaneWaveBasis(model; Ecut=5, kgrid=[1, 1, 1]);

We define our custom fix-point solver: simply a damped fixed-point

function my_fp_solver(f, x0, info0; maxiter)
    mixing_factor = .7
    x = x0
    info = info0
    for n = 1:maxiter
        fx, info = f(x, info)
        if info.converged || info.timedout
            break
        end
        x = x + mixing_factor * (fx - x)
    end
    (; fixpoint=x, info)
end;

Note that the fixpoint map f operates on an auxiliary variable info for state bookkeeping. Early termination criteria are flagged from inside the function f using boolean flags info.converged and info.timedout. For control over these criteria, see the is_converged and maxtime keyword arguments of self_consistent_field.

Our eigenvalue solver just forms the dense matrix and diagonalizes it explicitly (this only works for very small systems)

function my_eig_solver(A, X0; maxiter, tol, kwargs...)
    n = size(X0, 2)
    A = Array(A)
    E = eigen(A)
    λ = E.values[1:n]
    X = E.vectors[:, 1:n]
    (; λ, X, residual_norms=[], n_iter=0, converged=true, n_matvec=0)
end;

Finally we also define our custom mixing scheme. It will be a mixture of simple mixing (for the first 2 steps) and than default to Kerker mixing. In the mixing interface δF is $(ρ_\text{out} - ρ_\text{in})$, i.e. the difference in density between two subsequent SCF steps and the mix function returns $δρ$, which is added to $ρ_\text{in}$ to yield $ρ_\text{next}$, the density for the next SCF step.

struct MyMixing
    n_simple  # Number of iterations for simple mixing
end
MyMixing() = MyMixing(2)

function DFTK.mix_density(mixing::MyMixing, basis, δF; n_iter, kwargs...)
    if n_iter <= mixing.n_simple
        return δF  # Simple mixing -> Do not modify update at all
    else
        # Use the default KerkerMixing from DFTK
        DFTK.mix_density(KerkerMixing(), basis, δF; kwargs...)
    end
end

That's it! Now we just run the SCF with these solvers

scfres = self_consistent_field(basis;
                               tol=1e-4,
                               solver=my_fp_solver,
                               eigensolver=my_eig_solver,
                               mixing=MyMixing());

n     Energy            log10(ΔE)   log10(Δρ)   Diag   Δtime
---   ---------------   ---------   ---------   ----   ------
  1   -7.824595558215                   -0.55    0.0    271ms
  2   -7.852678862895       -1.55       -0.92    0.0    334ms
  3   -7.857032699882       -2.36       -1.38    0.0   84.1ms
  4   -7.857258210502       -3.65       -1.70    0.0   82.9ms
  5   -7.857310659971       -4.28       -2.00    0.0   83.4ms
  6   -7.857323136747       -4.90       -2.30    0.0   71.9ms
  7   -7.857326229041       -5.51       -2.59    0.0   72.9ms
  8   -7.857327035177       -6.09       -2.88    0.0    189ms
  9   -7.857327256939       -6.65       -3.15    0.0   71.5ms
 10   -7.857327321198       -7.19       -3.41    0.0   70.4ms
 11   -7.857327340706       -7.71       -3.67    0.0   81.8ms
 12   -7.857327346865       -8.21       -3.92    0.0    113ms
 13   -7.857327348871       -8.70       -4.16    0.0   84.7ms

Note that the default convergence criterion is the difference in density. When this gets below tol, the fixed-point solver terminates. You can also customize this with the is_converged keyword argument to self_consistent_field, as shown below.

Customizing the convergence criterion

Here is an example of a defining a custom convergence criterion and specifying it using the is_converged callback keyword to self_consistent_field.

function my_convergence_criterion(info)
    tol = 1e-10
    length(info.history_Etot) < 2 && return false
    ΔE = (info.history_Etot[end-1] - info.history_Etot[end])
    ΔE < tol
end

scfres2 = self_consistent_field(basis;
                                solver=my_fp_solver,
                                is_converged=my_convergence_criterion,
                                eigensolver=my_eig_solver,
                                mixing=MyMixing());

n     Energy            log10(ΔE)   log10(Δρ)   Diag   Δtime
---   ---------------   ---------   ---------   ----   ------
  1   -7.824595558215                   -0.55    0.0    155ms
  2   -7.852678862895       -1.55       -0.92    0.0    152ms
  3   -7.857032699882       -2.36       -1.38    0.0    181ms
  4   -7.857258210502       -3.65       -1.70    0.0   72.8ms
  5   -7.857310659971       -4.28       -2.00    0.0   82.7ms
  6   -7.857323136747       -4.90       -2.30    0.0   83.6ms
  7   -7.857326229041       -5.51       -2.59    0.0   87.3ms
  8   -7.857327035177       -6.09       -2.88    0.0   85.0ms
  9   -7.857327256939       -6.65       -3.15    0.0    136ms
 10   -7.857327321198       -7.19       -3.41    0.0   71.2ms
 11   -7.857327340706       -7.71       -3.67    0.0   79.4ms
 12   -7.857327346865       -8.21       -3.92    0.0   82.7ms
 13   -7.857327348871       -8.70       -4.16    0.0   84.8ms
 14   -7.857327349541       -9.17       -4.40    0.0    115ms
 15   -7.857327349769       -9.64       -4.64    0.0    100ms
 16   -7.857327349848      -10.10       -4.87    0.0   80.6ms