examples & tests, missing operators

examples/custom_potential.jl

@@ -70,7 +70,8 @@ tot_local_pot = DFTK.total_local_potential(scfres.ham)[:, 1, 1]; # use only dime
 # Extract other quantities before plotting them
 ρ = scfres.ρ[:, 1, 1, 1]        # converged density, first spin component
 ψ_fourier = scfres.ψ[1][:, 1]   # first k-point, all G components, first eigenvector
-ψ = ifft(basis, basis.kpoints[1], ψ_fourier)[:, 1, 1]
+ψ_t = ifft(basis, basis.kpoints[1], ψ_fourier)[:, 1, 1]
+ψ = reinterpret(reshape, eltype(eltype(ψ_t)), ψ_t)
 ψ /= (ψ[div(end, 2)] / abs(ψ[div(end, 2)]));
 
 using Plots

examples/error_estimates_forces.jl

@@ -124,7 +124,10 @@ function apply_metric(φ, P, δφ, A::Function)
         Aδφk = similar(δφk)
         φk = φ[ik]
         for n = 1:size(δφk,2)
-            Aδφk[:,n] = A(φk, P[ik], δφk[:,n], n)
+            φk_t = reinterpret(reshape, eltype(φk[1]), φk)
+            δφk_t = reinterpret(reshape, eltype(δφk[1,n]), δφk)
+            Aδφk_t = reinterpret(reshape, eltype(Aδφk[1,n]), Aδφk)
+            Aδφk_t[:,n] = A(φk_t, P[ik], δφk_t[:,n], n)
         end
         Aδφk
     end

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,8 +182,9 @@ 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,
-                           atoms, positions, atom_groups, terms, symmetries)
+                           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},
                positions::Vector{<:AbstractVector}; kwargs...)

src/PlaneWaveBasis.jl

@@ -25,6 +25,7 @@ struct Kpoint{T <: Real, GT <: AbstractVector{Vec3{Int}}}
     #                             # G_vectors(basis)[i] == G_vectors(basis, kpt)[mapping_inv[i]]
     G_vectors::GT                 # Wave vectors in integer coordinates (vector of Vec3{Int})
     #                             # ({G, 1/2 |k+G|^2 ≤ Ecut})
+    n_dof::Int                    # Number of components × number of `G_vectors(basis, kpt)`
 end
 
 @doc raw"""
@@ -75,6 +76,7 @@ struct PlaneWaveBasis{T,
     # "cubic" basis in reciprocal and real space, on which potentials and densities are stored
     G_vectors::T_G_vectors
     r_vectors::T_r_vectors
+    n_dof::Int  # number of components × number of `G_vectors`
 
     ## MPI-local information of the kpoints this processor treats
     # Irreducible kpoints. In the case of collinear spin,
@@ -142,9 +144,10 @@ Base.eltype(::PlaneWaveBasis{T}) where {T} = T
         Gvecs_k = to_device(architecture, Gvecs_k)
 
         mapping_inv = Dict(ifull => iball for (iball, ifull) in enumerate(mapping))
+        n_dof_k = length(Gvecs_k) * model.n_components
         for iσ = 1:model.n_spin_components
             push!(kpoints_per_spin[iσ],
-                  Kpoint(iσ, k, mapping, mapping_inv, Gvecs_k))
+                  Kpoint(iσ, k, mapping, mapping_inv, Gvecs_k, n_dof_k))
         end
     end
     vcat(kpoints_per_spin...)  # put all spin up first, then all spin down
@@ -220,7 +223,7 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Number, fft_size, variational,
     # by the same process
     n_kpt   = length(kcoords_global)
     n_procs = mpi_nprocs(comm_kpts)
-    
+
     # Custom reduction operators for MPI are currently not working on aarch64, so
     # fallbacks are defined in common/mpi.jl. For them to work, there cannot be more
     # than 1 MPI process.
@@ -279,13 +282,15 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Number, fft_size, variational,
     r_vectors = to_device(architecture, r_vectors)
     terms = Vector{Any}(undef, length(model.term_types))  # Dummy terms array, filled below
 
+    n_dof = length(Gs) * model.n_components
+
     basis = PlaneWaveBasis{T, value_type(T), typeof(Gs), typeof(r_vectors),
                            typeof(kpoints[1].G_vectors)}(
         model, fft_size, dvol,
         Ecut, variational,
         opFFT, ipFFT, opBFFT, ipBFFT,
         fft_normalization, ifft_normalization,
-        Gs, r_vectors,
+        Gs, r_vectors, n_dof,
         kpoints, kweights_thisproc, kgrid, kshift,
         kcoords_global, kweights_global, comm_kpts, krange_thisproc, krange_allprocs,
         architecture, symmetries, symmetries_respect_rgrid, terms)

src/common/ortho.jl

@@ -1,5 +1,19 @@
 @timing function ortho_qr(φk::ArrayType) where {ArrayType <: AbstractArray}
-    x = convert(ArrayType, qr(φk).Q)
+    TA = eltype(φk)
+    T = eltype(TA)
+    φk_matrix = reinterpret(T, φk)
+    # Matrix(…) for precompilation error in nightly CI.
+    # Failed to precompile DFTK [acf6eb54-70d9-11e9-0013-234b7a5f5337] to "/home/runner/.julia/compiled/v1.10/DFTK/jl_6TSfLT".
+    # ERROR: LoadError: MethodError: no method matching reinterpret(::typeof(reshape), ::Type{StaticArraysCore.SVector{1, ComplexF64}}, ::LinearAlgebra.QRCompactWYQ{ComplexF64, Matrix{ComplexF64}, Matrix{ComplexF64}})
+    #
+    # Closest candidates are:
+    #   reinterpret(::typeof(reshape), ::Type{T}, ::Base.ReinterpretArray{T, N, S, A, true} where {T, N, S, A<:(AbstractArray{S})}) where T
+    #    @ Base reinterpretarray.jl:160
+    #   reinterpret(::typeof(reshape), ::Type{T}, ::AbstractArray{T}) where T
+    #    @ Base reinterpretarray.jl:114
+    #   reinterpret(::typeof(reshape), ::Type{T}, ::A) where {T, S, A<:(AbstractArray{S})}
+    #    @ Base reinterpretarray.jl:82
+    x = reinterpret(reshape, TA, Matrix(qr(φk_matrix).Q))
     if size(x) == size(φk)
         return x
     else

src/densities.jl

@@ -12,7 +12,8 @@ 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}
-    S = promote_type(T, real(eltype(ψ[1])))
+    @assert basis.model.n_components == 1
+    S = promote_type(T, real(eltype(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]
 
@@ -29,7 +30,9 @@ using an optional `occupation_threshold`. By default all occupation numbers are
         ρ_chunklocal = map(1:Threads.nthreads()) do i
             zeros_like(basis.G_vectors, S, basis.fft_size..., basis.model.n_spin_components)
         end
-        ψnk_real_chunklocal = [zeros_like(basis.G_vectors, complex(S), basis.fft_size...)
+        ψnk_real_chunklocal = [zeros_like(basis.G_vectors,
+                                          SVector{basis.model.n_components, complex(S)},
+                                          basis.fft_size...)
                                for _ = 1:Threads.nthreads()]
 
         @sync for (ichunk, chunk) in enumerate(Iterators.partition(ik_n, chunk_length))
@@ -39,8 +42,12 @@ using an optional `occupation_threshold`. By default all occupation numbers are
                 kpt = basis.kpoints[ik]
 
                 ifft!(ψnk_real, basis, kpt, ψ[ik][:, n])
-                ρ_loc[:, :, :, kpt.spin] .+= (occupation[ik][n] .* basis.kweights[ik]
-                                              .* abs2.(ψnk_real))
+                ψnk_real_matrix = reshape(reinterpret(reshape, complex(S), ψnk_real),
+                                          basis.model.n_components, basis.fft_size...)
+                for σ in 1:basis.model.n_components
+                    ρ_loc[:, :, :, kpt.spin] .+= (occupation[ik][n] .* basis.kweights[ik]
+                                                  .* abs2.(ψnk_real_matrix[σ, :, :, :]))
+                end
 
                 synchronize_device(basis.architecture)
             end
@@ -72,8 +79,9 @@ end
 end
 
 @views @timing function compute_kinetic_energy_density(basis::PlaneWaveBasis, ψ, occupation)
-    T = promote_type(eltype(basis), real(eltype(ψ[1])))
-    τ = similar(ψ[1], T, (basis.fft_size..., basis.model.n_spin_components))
+    @assert basis.model.n_components == 1
+    T = promote_type(eltype(basis), real(eltype(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)

src/eigen/diag_lobpcg_hyper.jl

@@ -6,7 +6,8 @@ function lobpcg_hyper(A, X0; maxiter=100, prec=nothing,
     prec === nothing && (prec = I)
 
     @assert !largest "Only seeking the smallest eigenpairs is implemented."
-    result = LOBPCG(A, X0, I, prec, tol, maxiter; n_conv_check, kwargs...)
+    X0_matrix = reinterpret(reshape, eltype(eltype(X0)), X0)
+    result = LOBPCG(A, X0_matrix, I, prec, tol, maxiter; n_conv_check, kwargs...)
 
     n_conv_check === nothing && (n_conv_check = size(X0, 2))
     converged = maximum(result.residual_norms[1:n_conv_check]) < tol

src/external/wannier90.jl

@@ -199,13 +199,14 @@ The  Fourier coefficient of ``g^{per}_n`` at any G
 is given by the value of the Fourier transform of ``g_n`` in G.
 """
 function compute_Ak_gaussian_guess(basis::PlaneWaveBasis, ψk, kpt, centers, n_bands)
+    @assert basis.model.n_components == 1
     n_wannier = length(centers)
     # TODO This function should be improved in performance
 
     # associate a center with the fourier transform of the corresponding gaussian
     fourier_gn(center, qs) = [exp( 2π*(-im*dot(q, center) - dot(q, q) / 4) ) for q in qs]
     qs = vec(map(G -> G .+ kpt.coordinate, G_vectors(basis)))  # all q = k+G in reduced coordinates
-    Ak = zeros(eltype(ψk), (n_bands, n_wannier))
+    Ak = zeros(eltype(ψk), n_bands, n_wannier)
 
     # Compute Ak
     for n in 1:n_wannier
@@ -216,14 +217,15 @@ function compute_Ak_gaussian_guess(basis::PlaneWaveBasis, ψk, kpt, centers, n_b
         # Compute overlap
         for m in 1:n_bands
             # TODO Check the ordering of m and n here!
-            Ak[m, n] = dot(ψk[:, m], coeffs_gn_per)
+            Ak[m, n] = [dot(ψk[:, m], coeffs_gn_per)]
         end
     end
     Ak
 end
 
 
 @timing function write_w90_amn(fileprefix::String, basis::PlaneWaveBasis, ψ; n_bands, centers)
+    @assert basis.model.n_components == 1
     open(fileprefix * ".amn", "w") do fp
         println(fp, "Generated by DFTK at $(now())")
         println(fp, "$n_bands   $(length(basis.kpoints))  $(length(centers))")
@@ -233,7 +235,7 @@ end
             for n in 1:size(Ak, 2)
                 for m in 1:size(Ak, 1)
                     @printf(fp, "%3i %3i %3i  %22.18f %22.18f \n",
-                            m, n, ik, real(Ak[m, n]), imag(Ak[m, n]))
+                            m, n, ik, real(Ak[m, n][1]), imag(Ak[m, n][1]))
                 end
             end
         end
@@ -279,7 +281,8 @@ default_wannier_centres(n_wannier) = [rand(1, 3) for _ in 1:n_wannier]
 
     # Writing the unk files is expensive (requires FFTs), so only do if needed.
     if wannier_plot
-        write_w90_unk(fileprefix, basis, ψ; n_bands, spin=1)
+        ψ_t = [reinterpret(reshape, eltype(eltype(ψ[1])), ψk) for ψk in ψ]
+        write_w90_unk(fileprefix, basis, ψ_t; n_bands, spin=1)
     end
 
     # Run Wannierisation procedure

src/fft.jl

@@ -26,18 +26,26 @@ function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, f_fourier::Abstrac
 end
 function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis,
                  kpt::Kpoint, f_fourier::AbstractVector; normalize=true)
+    @assert basis.model.n_components == 1
     @assert length(f_fourier) == length(kpt.mapping)
     @assert size(f_real) == basis.fft_size
 
     # Pad the input data
-    fill!(f_real, 0)
-    f_real[kpt.mapping] = f_fourier
+    f_real_mat = reinterpret(reshape, eltype(eltype(f_real)), f_real)
+    fill!(f_real_mat, 0)
+    f_real_mat[kpt.mapping] = f_fourier
 
     # Perform an IFFT
-    mul!(f_real, basis.ipBFFT, f_real)
-    normalize && (f_real .*= basis.ifft_normalization)
+    mul!(f_real_mat, basis.ipBFFT, f_real_mat)
+    normalize && (f_real_mat .*= basis.ifft_normalization)
     f_real
 end
+function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, kpt::Kpoint,
+               f_fourier::AbstractVector{SV}; normalize=true) where {SV <: AbstractVector}
+    @assert basis.model.n_components == 1
+    f_fourier_mat = reinterpret(reshape, eltype(eltype(f_fourier)), f_fourier)
+    ifft!(f_real, basis, kpt, f_fourier_mat)
+end
 
 """
     ifft(basis::PlaneWaveBasis, [kpt::Kpoint, ] f_fourier)
@@ -81,15 +89,18 @@ function fft!(f_fourier::AbstractArray3, basis::PlaneWaveBasis, f_real::Abstract
 end
 function fft!(f_fourier::AbstractVector, basis::PlaneWaveBasis,
                  kpt::Kpoint, f_real::AbstractArray3; normalize=true)
+    @assert basis.model.n_components == 1
     @assert size(f_real) == basis.fft_size
     @assert length(f_fourier) == length(kpt.mapping)
 
     # FFT
-    mul!(f_real, basis.ipFFT, f_real)
+    f_real_mat = reinterpret(reshape, eltype(eltype(f_real)), f_real)
+    f_fourier_mat = reinterpret(reshape, eltype(eltype(f_fourier)), f_fourier)
+    mul!(f_real_mat, basis.ipFFT, f_real_mat)
 
     # Truncate
-    f_fourier .= view(f_real, kpt.mapping)
-    normalize && (f_fourier .*= basis.fft_normalization)
+    f_fourier_mat .= view(f_real_mat, kpt.mapping)
+    normalize && (f_fourier_mat .*= basis.fft_normalization)
     f_fourier
 end
 

src/orbitals.jl

@@ -31,8 +31,8 @@ end
 # reinterpret function from julia.
 # /!\ pack_ψ does not share memory while unpack_ψ does
 
-reinterpret_real(x) = reinterpret(real(eltype(x)), x)
-reinterpret_complex(x) = reinterpret(Complex{eltype(x)}, x)
+reinterpret_real(x) = reinterpret(real(eltype(eltype(x))), x)
+reinterpret_complex(x) = reinterpret(Complex{eltype(eltype(x))}, x)
 
 function pack_ψ(ψ)
     # TODO as an optimization, do that lazily? See LazyArrays
@@ -55,10 +55,12 @@ unpack_ψ(x, sizes_ψ) = deepcopy(unsafe_unpack_ψ(x, sizes_ψ))
 
 function random_orbitals(basis::PlaneWaveBasis{T}, kpt::Kpoint, howmany) where {T}
     @static if VERSION < v"1.7"  # TaskLocalRNG not yet available.
-        orbitals = randn(Complex{T}, length(G_vectors(basis, kpt)), howmany)
+        orbitals = randn(SVector{basis.model.n_components, Complex{T}},
+                         length(G_vectors(basis, kpt)), 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, SVector{basis.model.n_components, Complex{T}},
+                           length(G_vectors(basis, kpt)), howmany)
         randn!(TaskLocalRNG(), orbitals)  # use the RNG on the device if we're using a GPU
     end
     ortho_qr(orbitals)

src/response/cg.jl

@@ -62,6 +62,12 @@ function cg!(x::AbstractVector{T}, A::LinearMap{T}, b::AbstractVector{T};
     callback(info)
     info
 end
+function cg!(x::AbstractVector{SV}, A::LinearMap{T}, b::AbstractVector{SV};
+             kwargs...) where {T, SV <: AbstractVector}
+    x_matrix = reinterpret(reshape, eltype(eltype(x)), x)
+    b_matrix = reinterpret(reshape, eltype(eltype(b)), b)
+    cg!(x_matrix, A, b_matrix; kwargs...)
+end
 cg!(x::AbstractVector, A::AbstractMatrix, b::AbstractVector; kwargs...) = cg!(x, LinearMap(A), b; kwargs...)
 cg(A::LinearMap, b::AbstractVector; kwargs...) = cg!(zero(b), A, b; kwargs...)
 cg(A::AbstractMatrix, b::AbstractVector; kwargs...) = cg(LinearMap(A), b; kwargs...)

src/response/chi0.jl

@@ -190,6 +190,29 @@ function sternheimer_solver(Hk, ψk, ε, rhs;
 
     δψkn = ψk_extra * αkn + δψknᴿ
 end
+function  sternheimer_solver(Hk, ψk::AbstractArray{TV}, ε, rhs;
+                            callback=identity, cg_callback=identity,
+                            ψk_extra=nothing, εk_extra=zeros(0),
+                            Hψk_extra=nothing, tol=1e-9) where {TV <: AbstractVector}
+    Tψ = eltype(eltype(ψk))
+    ψk_mat = reinterpret(reshape, Tψ, ψk)
+    if isnothing(ψk_extra)
+        ψk_extra_mat = zeros(size(ψk,1), 0)
+    else
+        ψk_extra_mat = reinterpret(reshape, Tψ, ψk_extra)
+    end
+    if isnothing(Hψk_extra)
+        Hψk_extra_mat = zeros(size(ψk,1), 0)
+    else
+        Hψk_extra_mat = reinterpret(reshape, Tψ, Hψk_extra)
+    end
+    rhs_mat = reinterpret(reshape, Tψ, rhs)
+    δψkn_mat = sternheimer_solver(Hk, ψk_mat, ε, rhs_mat; callback, cg_callback,
+                                  # TODO: fix this
+                                  #ψk_extra=ψk_extra_mat, εk_extra, Hψk_extra=Hψk_extra_mat,
+                                  tol)
+    reinterpret(reshape, SVector{Hk.basis.model.n_components, Tψ}, δψkn_mat)
+end
 
 # Apply the four-point polarizability operator χ0_4P = -Ω^-1
 # Returns (δψ, δocc, δεF) corresponding to a change in *total* Hamiltonian δH