temp

src/DFTK.jl

@@ -61,6 +61,7 @@ include("SymOp.jl")
 export Smearing
 export Model
 export PlaneWaveBasis
+export Psi, Psik
 export compute_fft_size
 export G_vectors, G_vectors_cart, r_vectors, r_vectors_cart
 export Gplusk_vectors, Gplusk_vectors_cart
@@ -76,6 +77,7 @@ include("Smearing.jl")
 include("Model.jl")
 include("structure.jl")
 include("PlaneWaveBasis.jl")
+include("Psi.jl")
 include("fft.jl")
 include("orbitals.jl")
 include("show.jl")
@@ -242,23 +244,25 @@ function __init__()
 end
 
 # Precompilation block with a basic workflow
-@setup_workload begin
-    # very artificial silicon ground state example
-    a = 10.26
-    lattice = a / 2 * [[0 1 1.];
-                       [1 0 1.];
-                       [1 1 0.]]
-    Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
-    atoms     = [Si, Si]
-    positions = [ones(3)/8, -ones(3)/8]
-    magnetic_moments = [2, -2]
+if isnothing(get(ENV, "DFTK_NO_PRECOMPILATION", nothing))
+    @setup_workload begin
+        # very artificial silicon ground state example
+        a = 10.26
+        lattice = a / 2 * [[0 1 1.];
+                           [1 0 1.];
+                           [1 1 0.]]
+        Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
+        atoms     = [Si, Si]
+        positions = [ones(3)/8, -ones(3)/8]
+        magnetic_moments = [2, -2]
 
-    @compile_workload begin
-        model = model_LDA(lattice, atoms, positions;
-                          magnetic_moments, temperature=0.1, spin_polarization=:collinear)
-        basis = PlaneWaveBasis(model; Ecut=5, kgrid=[2, 2, 2])
-        ρ0 = guess_density(basis, magnetic_moments)
-        scfres = self_consistent_field(basis; ρ=ρ0, tol=1e-2, maxiter=3, callback=identity)
+        @compile_workload begin
+            model = model_LDA(lattice, atoms, positions;
+                              magnetic_moments, temperature=0.1, spin_polarization=:collinear)
+            basis = PlaneWaveBasis(model; Ecut=5, kgrid=[2, 2, 2])
+            ρ0 = guess_density(basis, magnetic_moments)
+            scfres = self_consistent_field(basis; ρ=ρ0, tol=1e-2, maxiter=3, callback=identity)
+        end
     end
 end
 end # module DFTK

src/Model.jl

@@ -26,6 +26,9 @@ struct Model{T <: Real, VT <: Real}
     n_electrons::Union{Int, Nothing}
     εF::Union{T, Nothing}
 
+    # Option for multicomponents computations. TODO: merge this with spins' handling.
+    n_components::Int
+
     # spin_polarization values:
     #     :none       No spin polarization, αα and ββ density identical,
     #                 αβ and βα blocks zero, 1 spin component treated explicitly
@@ -109,6 +112,7 @@ function Model(lattice::AbstractMatrix{T},
                spin_polarization=default_spin_polarization(magnetic_moments),
                symmetries=default_symmetries(lattice, atoms, positions, magnetic_moments,
                                              spin_polarization, terms),
+               n_components=1,
                ) where {T <: Real}
     # Validate εF and n_electrons
     if !isnothing(εF)  # fixed Fermi level
@@ -178,7 +182,8 @@ function Model(lattice::AbstractMatrix{T},
     Model{T,value_type(T)}(model_name,
                            lattice, recip_lattice, n_dim, inv_lattice, inv_recip_lattice,
                            unit_cell_volume, recip_cell_volume,
-                           n_electrons, εF, spin_polarization, n_spin, temperature, smearing,
+                           n_electrons, εF, n_components, spin_polarization, n_spin,
+                           temperature, smearing,
                            atoms, positions, atom_groups, terms, symmetries)
 end
 function Model(lattice::AbstractMatrix{<:Integer}, atoms::Vector{<:Element},

src/Psi.jl

@@ -0,0 +1,25 @@
+struct TempPsik2{T <: Real}
+    basis::PlaneWaveBasis{T}
+    kpoint::Kpoint{T}
+    ψk::AbstractArray{Complex{T}}
+    #ψk::AbstractArray3{T}
+    n_components::Int
+    n_Gk::Int
+    n_bands::Int
+end
+Psik = TempPsik2
+
+# Base.@kwdef struct Psi{T <: Real}
+#     basis::PlaneWaveBasis{T}
+#     ψ::Vector{Psik{T}}
+# end
+# Base.length(ψ::Psi) = length(ψ.ψ)
+function Psi(basis, ψ)
+    if length(size(ψ[1])) == 2
+        @warn "Temp Hack"
+        ψ_n = [reshape(ψk, 1, size(ψk)...) for ψk in ψ]
+    else
+        ψ_n = ψ
+    end
+    [Psik(basis, kpt, ψ_n[ik], size(ψ_n[ik], 1), size(ψ_n[ik], 2), size(ψ_n[ik], 3)) for (ik, kpt) in enumerate(basis.kpoints)]
+end

src/common/ortho.jl

@@ -1,9 +1,14 @@
 @timing function ortho_qr(φk::ArrayType) where {ArrayType <: AbstractArray}
-    x = convert(ArrayType, qr(φk).Q)
+    lead_dim = prod(size(φk)[1:2])
+    Q = qr(reshape(φk, lead_dim, :)).Q
+    # TODO: Does not work on the CI, but Ok locally…
+    #x = convert(ArrayType, reshape(Q, size(φk, 1), size(φk, 2), size(Q, 2)))
+    T = eltype(φk)
+    x = convert(ArrayType, reshape(convert(Matrix{T}, Q), size(φk, 1), size(φk, 2), size(φk, 3)))
     if size(x) == size(φk)
         return x
     else
         # Sometimes QR (but funnily not always) CUDA messes up the size here
-        return x[:, 1:size(φk, 2)]
+        return x[:, :, 1:size(φk, 3)]
     end
 end

src/densities.jl

@@ -12,6 +12,7 @@ using an optional `occupation_threshold`. By default all occupation numbers are
 """
 @views @timing function compute_density(basis::PlaneWaveBasis{T}, ψ, occupation;
                                         occupation_threshold=zero(T)) where {T}
+    @assert basis.model.n_components == 1
     S = promote_type(T, real(eltype(ψ[1])))
     # occupation should be on the CPU as we are going to be doing scalar indexing.
     occupation = [to_cpu(oc) for oc in occupation]
@@ -38,7 +39,7 @@ using an optional `occupation_threshold`. By default all occupation numbers are
                 ψnk_real = ψnk_real_chunklocal[ichunk]
                 kpt = basis.kpoints[ik]
 
-                ifft!(ψnk_real, basis, kpt, ψ[ik][:, n])
+                ifft!(ψnk_real, basis, kpt, ψ[ik][1, :, n])
                 ρ_loc[:, :, :, kpt.spin] .+= (occupation[ik][n] .* basis.kweights[ik]
                                               .* abs2.(ψnk_real))
 
@@ -65,21 +66,22 @@ end
                                    occupation, δoccupation=zero.(occupation);
                                    occupation_threshold=zero(T)) where {T}
     ForwardDiff.derivative(zero(T)) do ε
-        ψ_ε   = [ψk   .+ ε .* δψk   for (ψk,   δψk)   in zip(ψ, δψ)]
-        occ_ε = [occk .+ ε .* δocck for (occk, δocck) in zip(occupation, δoccupation)]
+        ψ_ε   = [ψk.ψk .+ ε .* δψk   for (ψk,   δψk)   in zip(ψ, δψ)]
+        occ_ε = [occk  .+ ε .* δocck for (occk, δocck) in zip(occupation, δoccupation)]
         compute_density(basis, ψ_ε, occ_ε; occupation_threshold)
     end
 end
 
 @views @timing function compute_kinetic_energy_density(basis::PlaneWaveBasis, ψ, occupation)
+    @assert basis.model.n_components == 1
     T = promote_type(eltype(basis), real(eltype(ψ[1])))
     τ = similar(ψ[1], T, (basis.fft_size..., basis.model.n_spin_components))
     τ .= 0
     dαψnk_real = zeros(complex(eltype(basis)), basis.fft_size)
     for (ik, kpt) in enumerate(basis.kpoints)
         G_plus_k = [[Gk[α] for Gk in Gplusk_vectors_cart(basis, kpt)] for α in 1:3]
-        for n = 1:size(ψ[ik], 2), α = 1:3
-            ifft!(dαψnk_real, basis, kpt, im .* G_plus_k[α] .* ψ[ik][:, n])
+        for n = 1:size(ψ[ik], 3), α = 1:3
+            ifft!(dαψnk_real, basis, kpt, im .* G_plus_k[α] .* ψ[ik][1, :, n])
             @. τ[:, :, :, kpt.spin] += occupation[ik][n] * basis.kweights[ik] / 2 * abs2(dαψnk_real)
         end
     end

src/eigen/diag.jl

@@ -13,49 +13,54 @@ function diagonalize_all_kblocks(eigensolver, ham::Hamiltonian, nev_per_kpoint::
                                  prec_type=PreconditionerTPA, interpolate_kpoints=true,
                                  tol=1e-6, miniter=1, maxiter=100, n_conv_check=nothing,
                                  show_progress=false)
-    kpoints = ham.basis.kpoints
+    basis   = ham.basis
+    kpoints = basis.kpoints
     results = Vector{Any}(undef, length(kpoints))
 
     progress = nothing
     if show_progress
         progress = Progress(length(kpoints), desc="Diagonalising Hamiltonian kblocks: ")
     end
+    n_components = basis.model.n_components
     for (ik, kpt) in enumerate(kpoints)
-        n_Gk = length(G_vectors(ham.basis, kpt))
+        n_Gk = length(G_vectors(basis, kpt))
         if n_Gk < nev_per_kpoint
             error("The size of the plane wave basis is $n_Gk, and you are asking for " *
                   "$nev_per_kpoint eigenvalues. Increase Ecut.")
         end
         # Get ψguessk
         if !isnothing(ψguess)
-            if n_Gk != size(ψguess[ik], 1)
-                error("Mismatch in dimension between guess ($(size(ψguess, 1)) and " *
-                      "Hamiltonian $n_Gk")
+            if n_Gk != ψguess[ik].n_Gk
+                error("Mismatch in dimension between guess ($(ψguess[ik].n_Gk)) and " *
+                      "Hamiltonian ($n_Gk)")
             end
-            nev_guess = size(ψguess[ik], 2)
+            nev_guess = ψguess[ik].n_bands
             if nev_guess > nev_per_kpoint
                 ψguessk = ψguess[ik][:, 1:nev_per_kpoint]
             elseif nev_guess == nev_per_kpoint
                 ψguessk = ψguess[ik]
             else
-                X0 = similar(ψguess[ik], n_Gk, nev_per_kpoint)
-                X0[:, 1:nev_guess] = ψguess[ik]
-                X0[:, nev_guess+1:end] = randn(eltype(X0), n_Gk, nev_per_kpoint - nev_guess)
-                ψguessk = ortho_qr(X0)
+                X0 = similar(ψguess[ik].ψk, n_components, n_Gk, nev_per_kpoint)
+                X0[:, :, 1:nev_guess] = ψguess[ik].ψk
+                X0[:, :, nev_guess+1:end] = randn(eltype(X0), n_components, n_Gk,
+                                                              nev_per_kpoint - nev_guess)
+                ψguessk = Psik(basis, kpt, ortho_qr(X0), n_components, n_Gk, nev_per_kpoint)
             end
         elseif interpolate_kpoints && ik > 1
             # use information from previous k-point
-            ψguessk = interpolate_kpoint(results[ik - 1].X, ham.basis, kpoints[ik - 1],
-                                         ham.basis, kpoints[ik])
+            ψguessk = interpolate_kpoint(results[ik - 1].X, basis, kpoints[ik - 1],
+                                         basis, kpoints[ik])
         else
-            ψguessk = random_orbitals(ham.basis, kpt, nev_per_kpoint)
+            ψguessk = random_orbitals(basis, kpt, nev_per_kpoint)
         end
-        @assert size(ψguessk) == (n_Gk, nev_per_kpoint)
+        @assert size(ψguessk.ψk) == (n_components, n_Gk, nev_per_kpoint)
 
         prec = nothing
-        !isnothing(prec_type) && (prec = prec_type(ham.basis, kpt))
-        results[ik] = eigensolver(ham.blocks[ik], ψguessk;
-                                  prec, tol, miniter, maxiter, n_conv_check)
+        !isnothing(prec_type) && (prec = prec_type(basis, kpt))
+        ψk_reshape = reshape(ψguessk.ψk, n_components * n_Gk, :)
+        res = eigensolver(ham.blocks[ik], ψk_reshape;
+                          prec, tol, miniter, maxiter, n_conv_check)
+        results[ik] = merge(res, (; X=reshape(res.X, size(ψguessk.ψk)...)))
 
         # Update progress bar if desired
         !isnothing(progress) && next!(progress)
@@ -78,7 +83,7 @@ Tuple returned by `diagonalize_all_kblocks`.
 """
 function select_eigenpairs_all_kblocks(eigres, range)
     merge(eigres, (λ=[λk[range] for λk in eigres.λ],
-                   X=[Xk[:, range] for Xk in eigres.X],
+                   X=[Xk[:, :, range] for Xk in eigres.X],
                    residual_norms=[resk[range] for resk in eigres.residual_norms]))
 end
 

src/eigen/preconditioners.jl

@@ -42,7 +42,7 @@ function PreconditionerTPA(basis::PlaneWaveBasis{T}, kpt::Kpoint; default_shift=
     PreconditionerTPA{T}(basis, kpt, kin, nothing, default_shift)
 end
 
-@views function ldiv!(Y, P::PreconditionerTPA, R)
+@views function ldiv!(Y, P::PreconditionerTPA, R::AbstractVecOrMat)
     if P.mean_kin === nothing
         ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
     else
@@ -52,6 +52,18 @@ end
     end
     Y
 end
+@views function ldiv!(Y, P::PreconditionerTPA, R::AbstractArray3)
+    if P.mean_kin === nothing
+        ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R)
+    else
+        Threads.@threads for n = 1:size(Y, 3)
+            for σ in 1:size(Y, 1)
+                Y[σ, :, n] .= P.mean_kin[n] ./ (P.mean_kin[n] .+ P.kin) .* R[σ, :, n]
+            end
+        end
+    end
+    Y
+end
 ldiv!(P::PreconditionerTPA, R) = ldiv!(R, P, R)
 (Base.:\)(P::PreconditionerTPA, R) = ldiv!(P, copy(R))
 

src/interpolation.jl

@@ -83,24 +83,28 @@ end
 Interpolate some data from one ``k``-point to another. The interpolation is fast, but not
 necessarily exact. Intended only to construct guesses for iterative solvers.
 """
-function interpolate_kpoint(data_in::AbstractVecOrMat,
+function interpolate_kpoint(data_in::AbstractArray3,
                             basis_in::PlaneWaveBasis,  kpoint_in::Kpoint,
                             basis_out::PlaneWaveBasis, kpoint_out::Kpoint)
     # TODO merge with transfer_blochwave_kpt
     if kpoint_in == kpoint_out
         return copy(data_in)
     end
-    @assert length(G_vectors(basis_in, kpoint_in)) == size(data_in, 1)
+    n_components = basis_in.model.n_components
+    @assert n_components == basis_out.model.n_components == size(data_in, 1)
+    @assert length(G_vectors(basis_in, kpoint_in)) == size(data_in, 2)
 
-    n_bands  = size(data_in, 2)
+    n_bands  = size(data_in, 3)
     n_Gk_out = length(G_vectors(basis_out, kpoint_out))
-    data_out = similar(data_in, n_Gk_out, n_bands) .= 0
+    data_out = similar(data_in, n_components, n_Gk_out, n_bands) .= 0
     # TODO: use a map, or this will not be GPU compatible (scalar indexing)
-    for iin in 1:size(data_in, 1)
-        idx_fft = kpoint_in.mapping[iin]
-        idx_fft in keys(kpoint_out.mapping_inv) || continue
-        iout = kpoint_out.mapping_inv[idx_fft]
-        data_out[iout, :] = data_in[iin, :]
+    for nc in 1:n_components
+        for iin in 1:size(data_in, 2)
+            idx_fft = kpoint_in.mapping[iin]
+            idx_fft in keys(kpoint_out.mapping_inv) || continue
+            iout = kpoint_out.mapping_inv[idx_fft]
+            data_out[nc, iout, :] = data_in[nc, iin, :]
+        end
     end
-    ortho_qr(data_out)  # Re-orthogonalize and renormalize
+    Psik(basis_out, kpoint_out, ortho_qr(data_out), n_components, n_Gk_out, size(data_out, 3))
 end

src/orbitals.jl

@@ -6,18 +6,18 @@ using Random  # Used to have a generic API for CPU and GPU computations alike: s
 # others when using temperature. It is set to 0.0 by default, to treat with insulators.
 function select_occupied_orbitals(basis, ψ, occupation; threshold=0.0)
     N = [something(findlast(x -> x > threshold, occk), 0) for occk in occupation]
-    selected_ψ   = [@view ψk[:, 1:N[ik]] for (ik, ψk)   in enumerate(ψ)]
-    selected_occ = [      occk[1:N[ik]]  for (ik, occk) in enumerate(occupation)]
+    selected_ψ   = [@view ψk.ψk[:, :, 1:N[ik]] for (ik, ψk)   in enumerate(ψ)]
+    selected_occ = [      occk[1:N[ik]]        for (ik, occk) in enumerate(occupation)]
 
     # if we have an insulator, sanity check that the orbitals we kept are the
     # occupied ones
     if iszero(threshold)
         model   = basis.model
         n_spin  = model.n_spin_components
         n_bands = div(model.n_electrons, n_spin * filled_occupation(model), RoundUp)
-        @assert n_bands == size(selected_ψ[1], 2)
+        @assert n_bands == size(selected_ψ[1], 3)
     end
-    (; ψ=selected_ψ, occupation=selected_occ)
+    (; ψ=Psi(basis, selected_ψ), occupation=selected_occ)
 end
 
 # Packing routines used in direct_minimization and newton algorithms.
@@ -54,12 +54,14 @@ end
 unpack_ψ(x, sizes_ψ) = deepcopy(unsafe_unpack_ψ(x, sizes_ψ))
 
 function random_orbitals(basis::PlaneWaveBasis{T}, kpt::Kpoint, howmany) where {T}
+    n_components = basis.model.n_components
+    n_Gk = length(G_vectors(basis, kpt))
     @static if VERSION < v"1.7"  # TaskLocalRNG not yet available.
-        orbitals = randn(Complex{T}, length(G_vectors(basis, kpt)), howmany)
+        orbitals = randn(Complex{T}, n_components, n_Gk, howmany)
         orbitals = to_device(basis.architecture, orbitals)
     else
-        orbitals = similar(basis.G_vectors, Complex{T}, length(G_vectors(basis, kpt)), howmany)
+        orbitals = similar(basis.G_vectors, Complex{T}, n_components, n_Gk, howmany)
         randn!(TaskLocalRNG(), orbitals)  # use the RNG on the device if we're using a GPU
     end
-    ortho_qr(orbitals)
+    Psik(basis, kpt, ortho_qr(orbitals), n_components, n_Gk, howmany)
 end

src/postprocess/dos.jl

@@ -56,7 +56,8 @@ function compute_ldos(ε, basis::PlaneWaveBasis{T}, eigenvalues, ψ;
     # Use compute_density routine to compute LDOS, using just the modified
     # weights (as "occupations") at each k-point. Note, that this automatically puts in the
     # required symmetrization with respect to kpoints and BZ symmetry
-    compute_density(basis, ψ, weights; occupation_threshold=weight_threshold)
+    ψ_arr = [ψk.ψk for ψk in ψ]
+    compute_density(basis, ψ_arr, weights; occupation_threshold=weight_threshold)
 end
 function compute_ldos(scfres::NamedTuple; ε=scfres.εF, kwargs...)
     compute_ldos(ε, scfres.basis, scfres.eigenvalues, scfres.ψ; kwargs...)