Refactor densities

docs/src/advanced/data_structures.md

@@ -149,7 +149,7 @@ Again the function is normalised:
 
 ```@example data_structures
 ψreal = G_to_r(basis, basis.kpoints[1], ψtest)
-sum(abs2.(ψreal)) * model.unit_cell_volume / prod(basis.fft_size)
+sum(abs2.(ψreal)) * basis.dvol
 ```
 
 The list of ``k`` points of the basis can be obtained with `basis.kpoints`.
@@ -197,21 +197,18 @@ Wavefunctions are stored in an array `scfres.ψ` as `ψ[ik][iG, iband]` where
 `ik` is the index of the kpoint (in `basis.kpoints`), `iG` is the
 index of the plane wave (in `G_vectors(basis.kpoints[ik])`) and
 `iband` is the index of the band.
-Densities are usually stored in a
-special type, `RealFourierArray`, from which the representation in
-real and reciprocal space can be accessed using `ρ.real` and
-`ρ.fourier` respectively.
+Densities are stored in real space, as a 4-dimensional array (the third being the spin component).
 
 ```@example data_structures
 using Plots  # hide
 rvecs = collect(r_vectors(basis))[:, 1, 1]  # slice along the x axis
 x = [r[1] for r in rvecs]                   # only keep the x coordinate
-plot(x, scfres.ρ.real[:, 1, 1], label="", xlabel="x", ylabel="ρ", marker=2)
+plot(x, scfres.ρ[:, 1, 1, 1], label="", xlabel="x", ylabel="ρ", marker=2)
 ```
 
 ```@example data_structures
 G_energies = [sum(abs2.(model.recip_lattice * G)) ./ 2 for G in G_vectors(basis)][:]
-scatter(G_energies, abs.(scfres.ρ.fourier[:]);
+scatter(G_energies, abs.(r_to_G(basis, scfres.ρ)[:]);
         yscale=:log10, ylims=(1e-12, 1), label="", xlabel="Energy", ylabel="|ρ|^2")
 ```
 

docs/src/advanced/symmetries.md

@@ -99,7 +99,7 @@ which will result in an almost identical convergence history.
 We can also explicitly verify both methods to yield the same density:
 ```@example symmetries
 using LinearAlgebra  # hide
-(norm(scfres_sym.ρ.real - scfres_nosym.ρ.real),
+(norm(scfres_sym.ρ - scfres_nosym.ρ),
  norm(values(scfres_sym.energies) .- values(scfres_nosym.energies)))
 ```
 

docs/src/guide/input_output.jl

@@ -44,10 +44,10 @@ atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="pbe")) => pos
 
 # Finally we run the calculation.
 
-model  = model_LDA(lattice, atoms, magnetic_moments=magnetic_moments, temperature=0.01)
-basis  = PlaneWaveBasis(model, 10; kgrid=(2, 2, 2))
-ρspin  = guess_spin_density(basis, magnetic_moments)
-scfres = self_consistent_field(basis, tol=1e-4, ρspin=ρspin, mixing=KerkerMixing());
+model = model_LDA(lattice, atoms, magnetic_moments=magnetic_moments, temperature=0.01)
+basis = PlaneWaveBasis(model, 10; kgrid=(2, 2, 2))
+ρ0 = guess_density(basis, magnetic_moments)
+scfres = self_consistent_field(basis, tol=1e-4, ρ=ρ0, mixing=KerkerMixing());
 
 # !!! note "DFTK and ASE"
 #     DFTK also supports using ASE to setup a calculation

examples/anyons.jl

@@ -31,4 +31,4 @@ E = scfres.energies.total
 s = 2
 E11 = π/2 * (2(s+1)/s)^((s+2)/s) * (s/(s+2))^(2(s+1)/s) * E^((s+2)/s) / β
 println("e(1,1) / (2π)= ", E11 / (2π))
-display(heatmap(scfres.ρ.real[:, :, 1], c=:blues))
+display(heatmap(scfres.ρ[:, :, 1, 1], c=:blues))

examples/arbitrary_floattype.jl

@@ -47,7 +47,7 @@ scfres.energies
 #-
 eltype(scfres.energies.total)
 #-
-eltype(scfres.ρ.real)
+eltype(scfres.ρ)
 
 #
 # !!! note "Generic linear algebra routines"

examples/cohen_bergstresser.jl

@@ -30,6 +30,5 @@ using Plots
 
 n_bands = 8
 ρ0 = guess_density(basis)  # Just dummy, has no meaning in this model
-ρspin0 = nothing
-p = plot_bandstructure(basis, ρ0, ρspin0, n_bands, εF=εF, kline_density=10)
+p = plot_bandstructure(basis, ρ0, n_bands, εF=εF, kline_density=10)
 ylims!(p, (-5, 6))

examples/collinear_magnetism.jl

@@ -56,10 +56,10 @@ magnetic_moments = [Fe => [4, ]];
 # We repeat the calculation using the same model as before. DFTK now detects
 # the non-zero moment and switches to a collinear calculation.
 
-model  = model_LDA(lattice, atoms, magnetic_moments=magnetic_moments, temperature=0.01)
-basis  = PlaneWaveBasis(model, Ecut; kgrid=kgrid)
-ρspin  = guess_spin_density(basis, magnetic_moments)
-scfres = self_consistent_field(basis, tol=1e-6, ρspin=ρspin, mixing=KerkerMixing());
+model = model_LDA(lattice, atoms, magnetic_moments=magnetic_moments, temperature=0.01)
+basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid)
+ρ0 = guess_density(basis, magnetic_moments)
+scfres = self_consistent_field(basis, tol=1e-6; ρ=ρ0, mixing=KerkerMixing());
 #-
 scfres.energies
 

examples/custom_potential.jl

@@ -60,8 +60,8 @@ model = Model(lattice; atoms=atoms, n_electrons=n_electrons, terms=terms,
 # a starting density and we choose to start from a zero density.
 Ecut = 500
 basis = PlaneWaveBasis(model, Ecut, kgrid=(1, 1, 1))
-ρ = zeros(complex(eltype(basis)), basis.fft_size)
-scfres = self_consistent_field(basis, tol=1e-8, ρ=from_fourier(basis, ρ))
+ρ = zeros(eltype(basis), basis.fft_size..., 1)
+scfres = self_consistent_field(basis, tol=1e-8, ρ=ρ)
 scfres.energies
 # Computing the forces can then be done as usual:
 hcat(compute_forces(scfres)...)
@@ -70,7 +70,7 @@ hcat(compute_forces(scfres)...)
 tot_local_pot = DFTK.total_local_potential(scfres.ham)[:, 1, 1]; # use only dimension 1
 
 # Extract other quantities before plotting them
-ρ = real(scfres.ρ.real)[:, 1, 1]  # converged density
+ρ = scfres.ρ[:, 1, 1, 1]  # converged density, first spin component
 ψ_fourier = scfres.ψ[1][:, 1];    # first kpoint, all G components, first eigenvector
 ψ = G_to_r(basis, basis.kpoints[1], ψ_fourier)[:, 1, 1]
 ψ /= (ψ[div(end, 2)] / abs(ψ[div(end, 2)]));

examples/custom_solvers.jl

@@ -54,13 +54,13 @@ struct MyMixing
 end
 MyMixing() = MyMixing(0.7)
 
-function DFTK.mix(mixing::MyMixing, basis, δF::RealFourierArray, δF_spin=nothing; n_iter, kwargs...)
+function DFTK.mix(mixing::MyMixing, basis, δF; n_iter, kwargs...)
     if n_iter <= 2
-        ## Just do simple mixing on total density and spin density (if it exists)
-        (mixing.α * δF, isnothing(δF_spin) ? nothing : mixing.α * δF_spin)
+        ## Just do simple mixing
+        mixing.α * δF
     else
         ## Use the KerkerMixing from DFTK
-        DFTK.mix(KerkerMixing(α=mixing.α), basis, δF, δF_spin; kwargs...)
+        DFTK.mix(KerkerMixing(α=mixing.α), basis, δF; kwargs...)
     end
 end
 

examples/dielectric.jl

@@ -22,44 +22,10 @@ scfres = self_consistent_field(basis, tol=1e-14)
 
 # Apply ε† = 1 - χ0 (vc + fxc)
 function eps_fun(dρ)
-    dρ = reshape(dρ, size(scfres.ρ.real))
-    dρ = from_real(basis, dρ)
     dv = apply_kernel(basis, dρ; ρ=scfres.ρ)
-    χdv = apply_χ0(scfres.ham, scfres.ψ, scfres.εF, scfres.eigenvalues, dv...)[1]
-    vec((dρ - χdv).real)
+    χdv = apply_χ0(scfres.ham, scfres.ψ, scfres.εF, scfres.eigenvalues, dv)
+    dρ - χdv
 end
 
-# A straightfoward Arnoldi eigensolver that diagonalizes the matrix at each step
-# This is more efficient than Arpack when `f` is very expensive
-println("Starting Arnoldi ...")
-function arnoldi(f, x0; howmany=5, tol=1e-4, maxiter=30, n_print=howmany)
-    for (V, B, r, nr, b) in ArnoldiIterator(f, x0)
-        # A * V = V * B + r * b'
-        V = hcat(V...)
-        AV = V*B + r*b'
-
-        ew, ev = eigen(B, sortby=real)
-        Vr = V*ev
-        AVr = AV*ev
-        R = AVr - Vr * Diagonal(ew)
-
-        N = size(V, 2)
-        normr = [norm(r) for r in eachcol(R)]
-
-        println("#--- $N ---#")
-        println("idcs      evals     residnorms")
-        inds = unique(append!(collect(1:min(N, n_print)), max(1, N-n_print):N))
-        for i in inds
-            @printf "% 3i  %10.6g  %10.6g\n" i real(ew[i]) normr[i]
-        end
-        any(imag.(ew[inds]) .> 1e-5) && println("Warn: Suppressed imaginary part.")
-        println()
-
-        is_converged = (N ≥ howmany && all(normr[1:howmany] .< tol)
-                        && all(normr[end-howmany:end] .< tol))
-        if is_converged || (N ≥ maxiter)
-            return (λ=ew, X=Vr, AX=AVr, residual_norms=normr)
-        end
-    end
-end
-arnoldi(eps_fun, vec(randn(size(scfres.ρ.real))); howmany=5)
+# eager diagonalizes the subspace matrix at each iteration
+eigsolve(eps_fun, randn(size(scfres.ρ)), 5, :LM; eager=true, verbosity=3)

examples/geometry_optimization.jl

@@ -35,7 +35,7 @@ function compute_scfres(x)
     model = model_LDA(lattice, atoms)
     basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid)
     global ψ, ρ
-    if ρ ===  nothing
+    if ρ === nothing
         ρ = guess_density(basis)
     end
     scfres = self_consistent_field(basis; ψ=ψ, ρ=ρ,

examples/gross_pitaevskii.jl

@@ -61,7 +61,7 @@ scfres.energies
 # We use the opportunity to explore some of DFTK internals.
 #
 # Extract the converged density and the obtained wave function:
-ρ = real(scfres.ρ.real)[:, 1, 1]  # converged density
+ρ = real(scfres.ρ)[:, 1, 1, 1]  # converged density, first spin component
 ψ_fourier = scfres.ψ[1][:, 1];    # first kpoint, all G components, first eigenvector
 
 # Transform the wave function to real space and fix the phase:
@@ -75,7 +75,7 @@ dx = a / N  # real-space grid spacing
 @assert sum(abs2.(ψ)) * dx ≈ 1.0
 
 # The density is simply built from ψ:
-norm(scfres.ρ.real - abs2.(ψ))
+norm(scfres.ρ - abs2.(ψ))
 
 # We summarize the ground state in a nice plot:
 using Plots

examples/gross_pitaevskii_2D.jl

@@ -38,4 +38,4 @@ model = Model(lattice; n_electrons=n_electrons,
               terms=terms, spin_polarization=:spinless)  # "spinless electrons"
 basis = PlaneWaveBasis(model, Ecut, kgrid=(1, 1, 1))
 scfres = direct_minimization(basis, tol=1e-5)  # Reduce tol for production
-heatmap(scfres.ρ.real[:, :, 1], c=:blues)
+heatmap(scfres.ρ[:, :, 1, 1], c=:blues)

examples/polarizability.jl

@@ -29,10 +29,9 @@ Ecut = 30
 tol = 1e-8
 
 ## dipole moment of a given density (assuming the current geometry)
-function dipole(ρ)
-    basis = ρ.basis
+function dipole(basis, ρ)
     rr = [a * (r[1] - 1/2) for r in r_vectors(basis)]
-    d = sum(rr .* ρ.real) * basis.model.unit_cell_volume / prod(basis.fft_size)
+    d = sum(rr .* ρ) * basis.dvol
 end;
 
 # ## Polarizability by finite differences
@@ -41,15 +40,15 @@ end;
 model = model_LDA(lattice, atoms; symmetries=false)
 basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid)
 res = self_consistent_field(basis, tol=tol)
-μref = dipole(res.ρ)
+μref = dipole(basis, res.ρ)
 
 # Then in a small uniform field:
 ε = .01
 model_ε = model_LDA(lattice, atoms; extra_terms=[ExternalFromReal(r -> -ε * (r[1] - a/2))],
                     symmetries=false)
 basis_ε = PlaneWaveBasis(model_ε, Ecut; kgrid=kgrid)
 res_ε = self_consistent_field(basis_ε, tol=tol)
-με = dipole(res_ε.ρ)
+με = dipole(basis_ε, res_ε.ρ)
 
 #-
 polarizability = (με - μref) / ε
@@ -81,30 +80,26 @@ println("Polarizability :   $polarizability")
 
 using KrylovKit
 
-## KrylovKit cannot deal with the density as a 3D array, so we need `vec` to vectorize
-## and `devec` to "unvectorize"
-devec(arr) = from_real(basis, reshape(arr, size(res.ρ.real)))
-
-## Apply (1- χ0 K) to a vectorized dρ
+## Apply (1- χ0 K)
 function dielectric_operator(dρ)
-    dρ = devec(dρ)
     dv = apply_kernel(basis, dρ; ρ=res.ρ)
-    χ0dv = apply_χ0(res.ham, res.ψ, res.εF, res.eigenvalues, dv...)[1]
-    vec((dρ - χ0dv).real)
+    χ0dv = apply_χ0(res.ham, res.ψ, res.εF, res.eigenvalues, dv)
+    dρ - χ0dv
 end
 
 ## dVext is the potential from a uniform field interacting with the dielectric dipole
 ## of the density.
-dVext = from_real(basis, [-a * (r[1] - 1/2) for r in r_vectors(basis)])
+dVext = [-a * (r[1] - 1/2) for r in r_vectors(basis)]
+dVext = cat(dVext; dims=4)
 
 ## Apply χ0 once to get non-interacting dipole
-dρ_nointeract = apply_χ0(res.ham, res.ψ, res.εF, res.eigenvalues, dVext)[1]
+dρ_nointeract = apply_χ0(res.ham, res.ψ, res.εF, res.eigenvalues, dVext)
 
 ## Solve Dyson equation to get interacting dipole
-dρ = devec(linsolve(dielectric_operator, vec(dρ_nointeract.real), verbosity=3)[1])
+dρ = linsolve(dielectric_operator, dρ_nointeract, verbosity=3)[1]
 
-println("Non-interacting polarizability: $(dipole(dρ_nointeract))")
-println("Interacting polarizability:     $(dipole(dρ))")
+println("Non-interacting polarizability: $(dipole(basis, dρ_nointeract))")
+println("Interacting polarizability:     $(dipole(basis, dρ))")
 
 # As expected, the interacting polarizability matches the finite difference
 # result. The non-interacting polarizability is higher.

examples/scf_callbacks.jl

@@ -63,7 +63,7 @@ function plot_callback(info)
     if info.stage == :finalize
         plot!(p, density_differences, label="|ρout - ρin|", markershape=:x)
     else
-        push!(density_differences, norm(info.ρout.real - info.ρin.real))
+        push!(density_differences, norm(info.ρout - info.ρin))
     end
     info
 end

examples/scf_checkpoints.jl

@@ -49,8 +49,7 @@ model = model_PBE(lattice, atoms, temperature=0.02, smearing=smearing=Smearing.G
                   magnetic_moments=magnetic_moments)
 basis = PlaneWaveBasis(model, Ecut; kgrid=[1, 1, 1])
 
-ρspin  = guess_spin_density(basis, magnetic_moments)
-scfres = self_consistent_field(basis, tol=1e-2, ρspin=ρspin)
+scfres = self_consistent_field(basis, tol=1e-2, ρ=guess_density(basis, magnetic_moments))
 save_scfres("scfres.jld2", scfres);
 #-
 scfres.energies
@@ -81,15 +80,17 @@ loaded.energies
 # callback to [`self_consistent_field`](@ref), for example:
 
 callback = DFTK.ScfSaveCheckpoints()
-scfres = self_consistent_field(basis, tol=1e-2, ρspin=ρspin, callback=callback);
+scfres = self_consistent_field(basis; ρ=guess_density(basis, magnetic_moments),
+                               tol=1e-2, callback=callback);
 
 # 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, tol=1e-2, ρspin=ρspin, callback=callback);
+scfres = self_consistent_field(basis; ρ=guess_density(basis, magnetic_moments),
+                               tol=1e-2, callback=callback);
 
 # For more details on using callbacks with DFTK's `self_consistent_field` function
 # see [Monitoring self-consistent field calculations](@ref).
@@ -102,7 +103,7 @@ scfres = self_consistent_field(basis, tol=1e-2, ρspin=ρspin, callback=callback
 # For this one just loads the checkpoint with [`load_scfres`](@ref):
 
 oldstate = load_scfres("dftk_scf_checkpoint.jld2")
-scfres   = self_consistent_field(oldstate.basis, ρ=oldstate.ρ, ρspin=oldstate.ρspin,
+scfres   = self_consistent_field(oldstate.basis, ρ=oldstate.ρ,
                                  ψ=oldstate.ψ, tol=1e-3);
 
 # !!! note "Availability of `load_scfres`, `save_scfres` and `ScfSaveCheckpoints`"

examples/silicon_scf_convergence.jl

@@ -23,11 +23,11 @@ function my_callback(info)
     info.stage == :finalize && return
     global resids
     info.n_iter == 1 && (resids = [])
-    err = norm(info.ρout.fourier - info.ρin.fourier)
+    err = norm(info.ρout - info.ρin)
     println(info.n_iter, " ", err)
     push!(resids, err)
 end
-my_isconverged = info -> norm(info.ρout.fourier - info.ρin.fourier) < tol
+my_isconverged = info -> norm(info.ρout - info.ρin) < tol
 opts = (callback=my_callback, is_converged=my_isconverged, maxiter=maxiter, tol=tol,
         determine_diagtol=info -> diagtol)
 

src/DFTK.jl

@@ -58,11 +58,6 @@ include("Smearing.jl")
 include("Model.jl")
 include("PlaneWaveBasis.jl")
 
-export RealFourierArray
-export from_real
-export from_fourier
-include("RealFourierArray.jl")
-
 export Energies
 include("energies.jl")
 
@@ -88,8 +83,11 @@ export apply_kernel
 export compute_kernel
 include("terms/terms.jl")
 
-export compute_density
 include("occupation.jl")
+export compute_density
+export total_density
+export spin_density
+export ρ_from_total_and_spin
 include("densities.jl")
 include("interpolation.jl")
 
@@ -130,7 +128,7 @@ export kgrid_size_from_minimal_spacing
 include("symmetry.jl")
 include("bzmesh.jl")
 
-export guess_density, guess_spin_density
+export guess_density
 export load_psp
 export list_psp
 include("guess_density.jl")