fix for components

JuliaMolSim · epolack · Feb 8, 2024 · Dec 11, 2023 · Feb 20, 2024 · Jun 13, 2024
commit 35f72a4ee223e004c5d2328892c9a4c1ca908102
diff --git a/benchmark/regression/testcases.jl b/benchmark/regression/testcases.jl
@@ -39,8 +39,7 @@ end
 
         n_electrons = testcase.n_electrons
         n_bands = div(n_electrons, 2, RoundUp)
-        ψ = [Matrix(qr(randn(ComplexF64, length(G_vectors(basis, kpt)), n_bands)).Q)
-             for kpt in basis.kpoints]
+        ψ = [DFTK.random_orbitals(basis, kpt, n_bands) for kpt in basis.kpoints]
         filled_occ = DFTK.filled_occupation(model)
         occupation = [filled_occ * rand(n_bands) for _ = 1:length(basis.kpoints)]
         occ_scaling = n_electrons / sum(sum(occupation))

diff --git a/ext/DFTKJLD2Ext.jl b/ext/DFTKJLD2Ext.jl
                            "($(basis.model.n_spin_components))")
            end
            if jld["n_components"] != basis.model.n_components
                consistent_kpts = false
                errormsg = ("Mismatch in number of spin components between file " *
                            "($(jld["n_components"])) and supplied basis " *
                            "($(basis.model.n_components))")
            end
diff --git a/src/Psi.jl b/src/Psi.jl
 """
 Reciprocal to [`to_composite_σG`](@ref).
 """
 from_composite_σG(basis::PlaneWaveBasis, ψ::AbstractVector) =
    from_composite_σG.(basis, basis.kpoints, ψ)

 @timing function enforce_phase!(basis::PlaneWaveBasis, kpt::Kpoint, ψk)
    map(eachcol(to_composite_σG(basis, kpt, ψk))) do ψkn
        ϕ = angle(ψkn[end])
        if abs(ϕ) > sqrt(eps(ϕ))
            ε = sign(real(ψkn[end] / cis(ϕ)))
            ψkn ./= ε*cis(ϕ)
        end
    end
    ψk
 end
 """Ensure that the wave function is real."""
 enforce_phase!(basis, ψ::AbstractVector) = enforce_phase!.(basis, basis.kpoints, ψ)
diff --git a/src/common/ortho.jl b/src/common/ortho.jl
 @views @timing function ortho_qr(basis, kpt, φk::AbstractArray{T}) where {T}
    x = convert(Matrix{T}, qr(to_composite_σG(basis, kpt, φk)).Q)
    # Sometimes QR (but funnily not always) CUDA messes up the size here.
    from_composite_σG(basis, kpt, x)[:, :, 1:size(φk, 3)]
diff --git a/src/densities.jl b/src/densities.jl
    symmetrize_ρ(basis, δρ; do_lowpass=false)
 end

 @views @timing function compute_kinetic_energy_density(basis::PlaneWaveBasis{TT}, ψ,
                                                       occupation) where {TT}
    @assert basis.model.n_components == 1
    T = promote_type(TT, real(eltype(ψ[1])))
diff --git a/src/eigen/diag.jl b/src/eigen/diag.jl
        # Get ψguessk
        if !isnothing(ψguess)
            if n_Gk != size(ψguess[ik], 2)
                error("Mismatch in dimension between guess ($(size(ψguess[ik], 2))) and " *
                      "Hamiltonian ($n_Gk)")
            end
            nev_guess = size(ψguess[ik], 3)
diff --git a/src/eigen/preconditioners.jl b/src/eigen/preconditioners.jl
            for σ = 1:size(Y, 1)
                Y[σ, :, n] .= P.mean_kin[n] ./ (P.mean_kin[n] .+ P.kin) .* R[σ, :, n]
            end
        end
    else
        # Fallback for LOBPCG as it works on matrices.
        R_σG = from_composite_σG(P.basis, P.kpoint, R)
diff --git a/src/fft.jl b/src/fft.jl
    normalize && (f_real .*= basis.ifft_normalization)
    f_real
 end
 function ifft!(f_real::AbstractArray4, basis::PlaneWaveBasis, f_fourier::AbstractMatrix)
    mul!(f_real, basis.opBFFT_mc, f_fourier)
    f_real .*= basis.ifft_normalization
 end
 function ifft!(f_real::AbstractArray4, basis::PlaneWaveBasis, kpt::Kpoint,
               f_fourier::AbstractMatrix; normalize=true)
    normalize && (f_fourier .*= basis.fft_normalization)
    f_fourier
 end
 function fft!(f_fourier::AbstractArray4, basis::PlaneWaveBasis, f_real::AbstractArray4)
    if eltype(f_real) <: Real
        f_real = complex.(f_real)
    end
    mul!(f_fourier, basis.opFFT_mc, f_real)
    f_fourier .*= basis.fft_normalization
 end
 function fft!(f_fourier::AbstractMatrix, basis::PlaneWaveBasis, kpt::Kpoint,
              f_real::AbstractArray4; normalize=true)
    opBFFT = inv(opFFT).p
    (; ipFFT, opFFT, ipBFFT, opBFFT)
 end
 function build_fft_plans!(tmp::AbstractArray{Complex{T}};
                          region=1:ndims(tmp)) where {T<:Union{Float32,Float64}}
    ipFFT = AbstractFFTs.plan_fft!(tmp, region)
    opFFT = AbstractFFTs.plan_fft(tmp, region)
    # backwards-FFT by inverting and stripping off normalizations
    ipBFFT = inv(ipFFT).p
    opBFFT = inv(opFFT).p
    (; ipFFT, opFFT, ipBFFT, opBFFT)
 end

 # TODO Some grid sizes are broken in the generic FFT implementation
diff --git a/src/response/hessian.jl b/src/response/hessian.jl
    - `δVind`: Change in potential induced by `δρ` (the term needed on top of `δHextψ`
      to get `δHψ`).
 """
 @timing function solve_ΩplusK_split(ham::Hamiltonian, ρ::AbstractArray{T}, ψ, occupation, εF, 
                                     eigenvalues, rhs; tol=1e-8, tol_sternheimer=tol/10, 
                                     verbose=false, occupation_threshold, q=zero(Vec3{real(T)}), 
                                     kwargs...) where {T}
diff --git a/src/terms/Hamiltonian.jl b/src/terms/Hamiltonian.jl

 # Serves as a wrapper around the specialized multiplication functions for LOBPCG as it
 # expects to work on vectors or matrices. Explicitely create a 3-tensor to prevent ambiguity.
 @timing function LinearAlgebra.mul!(Hψ_σG::AbstractVecOrMat, H::HamiltonianBlock,
                                    ψ_σG::AbstractVecOrMat)
    Hψ = unfold_from_composite_σG(H.basis, H.kpoint, Hψ_σG)
    ψ = unfold_from_composite_σG(H.basis, H.kpoint, ψ_σG)

 # Loop through bands, IFFT to get ψ in real space, loop through terms, FFT and accumulate into Hψ
 # For the common DftHamiltonianBlock there is an optimized version below
 @views @timing "Hamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray3,
                                                                        H::GenericHamiltonianBlock,
                                                                        ψ::AbstractArray3)
    basis = H.basis
 end

 # Fast version, specialized on DFT models. Minimizes the number of FFTs and allocations
 @views @timing "DftHamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray3,
                                                                           H::DftHamiltonianBlock,
                                                                           ψ::AbstractArray3)
    n_bands = size(ψ, 3)
diff --git a/src/terms/operators.jl b/src/terms/operators.jl
                # G_vectors(basis)[ind_ΔG] = G - G'
                ind_ΔG = index_G_vectors(op.basis, G - G′)
                if isnothing(ind_ΔG)
                    error("For full matrix construction, the FFT size must be " *
                          "large enough so that Hamiltonian applications are exact")
                end
                H[σ, i, σ, j] = pot_fourier[ind_ΔG] / sqrt(op.basis.model.unit_cell_volume)
	"($(basis.model.n_spin_components))")
	end
	if jld["n_components"] != basis.model.n_components
	consistent_kpts = false
	errormsg = ("Mismatch in number of spin components between file " *
Check warning on line 94 in ext/DFTKJLD2Ext.jl Codecov / codecov/patch ext/DFTKJLD2Ext.jl#L93-L94 `Added lines #L93 - L94 were not covered by tests`
	"($(jld["n_components"])) and supplied basis " *
	"($(basis.model.n_components))")
	end
	"""
	Reciprocal to [`to_composite_σG`](@ref).
	"""
	from_composite_σG(basis::PlaneWaveBasis, ψ::AbstractVector) =
Check warning on line 35 in src/Psi.jl Codecov / codecov/patch src/Psi.jl#L35 `Added line #L35 was not covered by tests`
	from_composite_σG.(basis, basis.kpoints, ψ)

	@timing function enforce_phase!(basis::PlaneWaveBasis, kpt::Kpoint, ψk)
	map(eachcol(to_composite_σG(basis, kpt, ψk))) do ψkn
	ϕ = angle(ψkn[end])
	if abs(ϕ) > sqrt(eps(ϕ))
	ε = sign(real(ψkn[end] / cis(ϕ)))
	ψkn ./= ε*cis(ϕ)
Check warning on line 43 in src/Psi.jl Codecov / codecov/patch src/Psi.jl#L38-L43 `Added lines #L38 - L43 were not covered by tests`
	end
	end
	ψk
Check warning on line 46 in src/Psi.jl Codecov / codecov/patch src/Psi.jl#L46 `Added line #L46 was not covered by tests`
	end
	"""Ensure that the wave function is real."""
	enforce_phase!(basis, ψ::AbstractVector) = enforce_phase!.(basis, basis.kpoints, ψ)
Check warning on line 49 in src/Psi.jl Codecov / codecov/patch src/Psi.jl#L49 `Added line #L49 was not covered by tests`
	@views @timing function ortho_qr(basis, kpt, φk::AbstractArray{T}) where {T}
Check warning on line 1 in src/common/ortho.jl Codecov / codecov/patch src/common/ortho.jl#L1 `Added line #L1 was not covered by tests`
	x = convert(Matrix{T}, qr(to_composite_σG(basis, kpt, φk)).Q)
	# Sometimes QR (but funnily not always) CUDA messes up the size here.
	from_composite_σG(basis, kpt, x)[:, :, 1:size(φk, 3)]
	symmetrize_ρ(basis, δρ; do_lowpass=false)
	end

	@views @timing function compute_kinetic_energy_density(basis::PlaneWaveBasis{TT}, ψ,
Check warning on line 115 in src/densities.jl Codecov / codecov/patch src/densities.jl#L115 `Added line #L115 was not covered by tests`
	occupation) where {TT}
	@assert basis.model.n_components == 1
	T = promote_type(TT, real(eltype(ψ[1])))
	# Get ψguessk
	if !isnothing(ψguess)
	if n_Gk != size(ψguess[ik], 2)
	error("Mismatch in dimension between guess ($(size(ψguess[ik], 2))) and " *
Check warning on line 27 in src/eigen/diag.jl Codecov / codecov/patch src/eigen/diag.jl#L27 `Added line #L27 was not covered by tests`
	"Hamiltonian ($n_Gk)")
	end
	nev_guess = size(ψguess[ik], 3)
	for σ = 1:size(Y, 1)
	Y[σ, :, n] .= P.mean_kin[n] ./ (P.mean_kin[n] .+ P.kin) .* R[σ, :, n]
	end
	end
Check warning on line 57 in src/eigen/preconditioners.jl Codecov / codecov/patch src/eigen/preconditioners.jl#L57 `Added line #L57 was not covered by tests`
	else
	# Fallback for LOBPCG as it works on matrices.
	R_σG = from_composite_σG(P.basis, P.kpoint, R)
	normalize && (f_real .*= basis.ifft_normalization)
	f_real
	end
	function ifft!(f_real::AbstractArray4, basis::PlaneWaveBasis, f_fourier::AbstractMatrix)
	mul!(f_real, basis.opBFFT_mc, f_fourier)
	f_real .*= basis.ifft_normalization
Check warning on line 42 in src/fft.jl Codecov / codecov/patch src/fft.jl#L40-L42 `Added lines #L40 - L42 were not covered by tests`
	end
	function ifft!(f_real::AbstractArray4, basis::PlaneWaveBasis, kpt::Kpoint,
	f_fourier::AbstractMatrix; normalize=true)
	normalize && (f_fourier .*= basis.fft_normalization)
	f_fourier
	end
	function fft!(f_fourier::AbstractArray4, basis::PlaneWaveBasis, f_real::AbstractArray4)
	if eltype(f_real) <: Real
	f_real = complex.(f_real)
Check warning on line 119 in src/fft.jl Codecov / codecov/patch src/fft.jl#L117-L119 `Added lines #L117 - L119 were not covered by tests`
	end
	mul!(f_fourier, basis.opFFT_mc, f_real)
	f_fourier .*= basis.fft_normalization
Check warning on line 122 in src/fft.jl Codecov / codecov/patch src/fft.jl#L121-L122 `Added lines #L121 - L122 were not covered by tests`
	end
	function fft!(f_fourier::AbstractMatrix, basis::PlaneWaveBasis, kpt::Kpoint,
	f_real::AbstractArray4; normalize=true)
	opBFFT = inv(opFFT).p
	(; ipFFT, opFFT, ipBFFT, opBFFT)
	end
	function build_fft_plans!(tmp::AbstractArray{Complex{T}};
Check warning on line 332 in src/fft.jl Codecov / codecov/patch src/fft.jl#L332 `Added line #L332 was not covered by tests`
	region=1:ndims(tmp)) where {T<:Union{Float32,Float64}}
	ipFFT = AbstractFFTs.plan_fft!(tmp, region)
	opFFT = AbstractFFTs.plan_fft(tmp, region)
Check warning on line 335 in src/fft.jl Codecov / codecov/patch src/fft.jl#L334-L335 `Added lines #L334 - L335 were not covered by tests`
	# backwards-FFT by inverting and stripping off normalizations
	ipBFFT = inv(ipFFT).p
	opBFFT = inv(opFFT).p
	(; ipFFT, opFFT, ipBFFT, opBFFT)
Check warning on line 339 in src/fft.jl Codecov / codecov/patch src/fft.jl#L337-L339 `Added lines #L337 - L339 were not covered by tests`
	end

	# TODO Some grid sizes are broken in the generic FFT implementation
	- `δVind`: Change in potential induced by `δρ` (the term needed on top of `δHextψ`
	to get `δHψ`).
	"""
	@timing function solve_ΩplusK_split(ham::Hamiltonian, ρ::AbstractArray{T}, ψ, occupation, εF,
Check warning on line 149 in src/response/hessian.jl Codecov / codecov/patch src/response/hessian.jl#L149 `Added line #L149 was not covered by tests`
	eigenvalues, rhs; tol=1e-8, tol_sternheimer=tol/10,
	verbose=false, occupation_threshold, q=zero(Vec3{real(T)}),
	kwargs...) where {T}

	# Serves as a wrapper around the specialized multiplication functions for LOBPCG as it
	# expects to work on vectors or matrices. Explicitely create a 3-tensor to prevent ambiguity.
	@timing function LinearAlgebra.mul!(Hψ_σG::AbstractVecOrMat, H::HamiltonianBlock,
Check warning on line 98 in src/terms/Hamiltonian.jl Codecov / codecov/patch src/terms/Hamiltonian.jl#L98 `Added line #L98 was not covered by tests`
	ψ_σG::AbstractVecOrMat)
	Hψ = unfold_from_composite_σG(H.basis, H.kpoint, Hψ_σG)
	ψ = unfold_from_composite_σG(H.basis, H.kpoint, ψ_σG)

	# Loop through bands, IFFT to get ψ in real space, loop through terms, FFT and accumulate into Hψ
	# For the common DftHamiltonianBlock there is an optimized version below
	@views @timing "Hamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray3,
Check warning on line 108 in src/terms/Hamiltonian.jl Codecov / codecov/patch src/terms/Hamiltonian.jl#L108 `Added line #L108 was not covered by tests`
	H::GenericHamiltonianBlock,
	ψ::AbstractArray3)
	basis = H.basis
	end

	# Fast version, specialized on DFT models. Minimizes the number of FFTs and allocations
	@views @timing "DftHamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray3,
Check warning on line 146 in src/terms/Hamiltonian.jl Codecov / codecov/patch src/terms/Hamiltonian.jl#L146 `Added line #L146 was not covered by tests`
	H::DftHamiltonianBlock,
	ψ::AbstractArray3)
	n_bands = size(ψ, 3)
	# G_vectors(basis)[ind_ΔG] = G - G'
	ind_ΔG = index_G_vectors(op.basis, G - G′)
	if isnothing(ind_ΔG)
	error("For full matrix construction, the FFT size must be " *
Check warning on line 86 in src/terms/operators.jl Codecov / codecov/patch src/terms/operators.jl#L86 `Added line #L86 was not covered by tests`
	"large enough so that Hamiltonian applications are exact")
	end
	H[σ, i, σ, j] = pot_fourier[ind_ΔG] / sqrt(op.basis.model.unit_cell_volume)