diff --git a/docs/src/developer/data_structures.md b/docs/src/developer/data_structures.md index 74a18bdb43..1497a2a81a 100644 --- a/docs/src/developer/data_structures.md +++ b/docs/src/developer/data_structures.md @@ -140,7 +140,7 @@ For example let us check the normalization of the first eigenfunction at the first ``k``-point in reciprocal space: ```@example data_structures -ψtest = scfres.ψ[1][:, 1] +ψtest = scfres.ψ[1][:, :, 1] sum(abs2.(ψtest)) ``` diff --git a/examples/compare_solvers.jl b/examples/compare_solvers.jl index 472dd69237..22a3f256b7 100644 --- a/examples/compare_solvers.jl +++ b/examples/compare_solvers.jl @@ -39,8 +39,8 @@ scfres_dm = direct_minimization(basis; tol=tol^2); scfres_start = self_consistent_field(basis; tol=0.5); # Remove the virtual orbitals (which Newton cannot treat yet) -ψ = DFTK.select_occupied_orbitals(basis, scfres_start.ψ, scfres_start.occupation).ψ -scfres_newton = newton(basis, ψ; tol); +ψ = DFTK.select_occupied_orbitals(scfres_start.ψ, scfres_start.occupation).ψ +scfres_newton = newton(ψ; tol); # ## Comparison of results diff --git a/examples/custom_potential.jl b/examples/custom_potential.jl index e33c0fe8a7..76c37b8cd4 100644 --- a/examples/custom_potential.jl +++ b/examples/custom_potential.jl @@ -72,8 +72,9 @@ compute_forces(scfres) tot_local_pot = DFTK.total_local_potential(scfres.ham)[:, 1, 1]; # use only dimension 1 # 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 +ρ = scfres.ρ[:, 1, 1, 1] # converged density, first spin component +ψ_fourier = scfres.ψ[1][1, :, 1] # first k-point, first (and only) component, + # all G components, first eigenvector # Transform the wave function to real space and fix the phase: ψ = ifft(basis, basis.kpoints[1], ψ_fourier)[:, 1, 1] diff --git a/examples/error_estimates_forces.jl b/examples/error_estimates_forces.jl index 7e7000ca2b..e2c3f209c2 100644 --- a/examples/error_estimates_forces.jl +++ b/examples/error_estimates_forces.jl @@ -54,7 +54,7 @@ tol = 1e-5; # We compute the reference solution ``P_*`` from which we will compute the # references forces. scfres_ref = self_consistent_field(basis_ref; tol, callback=identity) -ψ_ref = DFTK.select_occupied_orbitals(basis_ref, scfres_ref.ψ, scfres_ref.occupation).ψ; +ψ_ref = DFTK.select_occupied_orbitals(scfres_ref.ψ, scfres_ref.occupation).ψ; # We compute a variational approximation of the reference solution with # smaller `Ecut`. `ψr`, `ρr` and `Er` are the quantities computed with `Ecut` @@ -69,66 +69,67 @@ Ecut = 15 basis = PlaneWaveBasis(model; Ecut, kgrid) scfres = self_consistent_field(basis; tol, callback=identity) ψr = DFTK.transfer_blochwave(scfres.ψ, basis, basis_ref) -ρr = compute_density(basis_ref, ψr, scfres.occupation) -Er, hamr = energy_hamiltonian(basis_ref, ψr, scfres.occupation; ρ=ρr); +ρr = compute_density(ψr, scfres.occupation) +Er, hamr = energy_hamiltonian(ψr, scfres.occupation; ρ=ρr); # We then compute several quantities that we need to evaluate the error bounds. # - Compute the residual ``R(P)``, and remove the virtual orbitals, as required # in [`src/scf/newton.jl`](https://github.com/JuliaMolSim/DFTK.jl/blob/fedc720dab2d194b30d468501acd0f04bd4dd3d6/src/scf/newton.jl#L121). -res = DFTK.compute_projected_gradient(basis_ref, ψr, scfres.occupation) -res, occ = DFTK.select_occupied_orbitals(basis_ref, res, scfres.occupation) -ψr = DFTK.select_occupied_orbitals(basis_ref, ψr, scfres.occupation).ψ; +res = DFTK.compute_projected_gradient(ψr, scfres.occupation) +res, occ = DFTK.select_occupied_orbitals(BlochWaves(ψr.basis, res), scfres.occupation) +ψr = DFTK.select_occupied_orbitals(ψr, scfres.occupation).ψ; # - Compute the error ``P-P_*`` on the associated orbitals ``ϕ-ψ`` after aligning # them: this is done by solving ``\min |ϕ - ψU|`` for ``U`` unitary matrix of # size ``N×N`` (``N`` being the number of electrons) whose solution is # ``U = S(S^*S)^{-1/2}`` where ``S`` is the overlap matrix ``ψ^*ϕ``. -function compute_error(basis, ϕ, ψ) +function compute_error(ϕ, ψ) map(zip(ϕ, ψ)) do (ϕk, ψk) - S = ψk'ϕk + S = ψk[1, :, :]'ϕk[1, :, :] U = S*(S'S)^(-1/2) - ϕk - ψk*U + reshape(ϕk[1, :, :] - ψk[1, :, :]*U, size(ϕk)...) end end -err = compute_error(basis_ref, ψr, ψ_ref); +err = compute_error(ψr, ψ_ref); # - Compute ``{\boldsymbol M}^{-1}R(P)`` with ``{\boldsymbol M}^{-1}`` defined in [^CDKL2021]: P = [PreconditionerTPA(basis_ref, kpt) for kpt in basis_ref.kpoints] map(zip(P, ψr)) do (Pk, ψk) - DFTK.precondprep!(Pk, ψk) + DFTK.precondprep!(Pk, ψk[1, :, :]) end function apply_M(φk, Pk, δφnk, n) + fact = reshape(sqrt.(Pk.mean_kin[n] .+ Pk.kin), size(δφnk)...) DFTK.proj_tangent_kpt!(δφnk, φk) - δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk + δφnk = fact .* δφnk DFTK.proj_tangent_kpt!(δφnk, φk) - δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk + δφnk = fact .* δφnk DFTK.proj_tangent_kpt!(δφnk, φk) end function apply_inv_M(φk, Pk, δφnk, n) DFTK.proj_tangent_kpt!(δφnk, φk) - op(x) = apply_M(φk, Pk, x, n) + op(x) = apply_M(φk, Pk, reshape(x, 1, :), n) function f_ldiv!(x, y) - x .= DFTK.proj_tangent_kpt(y, φk) + x .= DFTK.proj_tangent_kpt(reshape(y, 1, :), φk)[1, :] x ./= (Pk.mean_kin[n] .+ Pk.kin) - DFTK.proj_tangent_kpt!(x, φk) + DFTK.proj_tangent_kpt!(reshape(x, 1, :), φk)[1, :] end - J = LinearMap{eltype(φk)}(op, size(δφnk, 1)) - δφnk = cg(J, δφnk; Pl=DFTK.FunctionPreconditioner(f_ldiv!), + J = LinearMap{eltype(φk)}(op, size(δφnk, 2)) + δφnk = cg(J, δφnk[1, :]; Pl=DFTK.FunctionPreconditioner(f_ldiv!), verbose=false, reltol=0, abstol=1e-15) - DFTK.proj_tangent_kpt!(δφnk, φk) + DFTK.proj_tangent_kpt!(reshape(δφnk, 1, :), φk) end function apply_metric(φ, P, δφ, A::Function) map(enumerate(δφ)) do (ik, δφk) Aδφk = similar(δφk) φk = φ[ik] - for n = 1:size(δφk,2) - Aδφk[:,n] = A(φk, P[ik], δφk[:,n], n) + for n = 1:size(δφk, 3) + Aδφk[:, :, n] = A(φk, P[ik], δφk[:, :, n], n) end Aδφk end end -Mres = apply_metric(ψr, P, res, apply_inv_M); +Mres = apply_metric(ψr.data, P, res, apply_inv_M); # We can now compute the modified residual ``R_{\rm Schur}(P)`` using a Schur # complement to approximate the error on low-frequencies[^CDKL2021]: @@ -148,7 +149,7 @@ Mres = apply_metric(ψr, P, res, apply_inv_M); # - Compute the projection of the residual onto the high and low frequencies: resLF = DFTK.transfer_blochwave(res, basis_ref, basis) -resHF = res - DFTK.transfer_blochwave(resLF, basis, basis_ref); +resHF = denest(res) - denest(DFTK.transfer_blochwave(resLF, basis, basis_ref)); # - Compute ``{\boldsymbol M}^{-1}_{22}R_2(P)``: e2 = apply_metric(ψr, P, resHF, apply_inv_M); @@ -158,19 +159,19 @@ e2 = apply_metric(ψr, P, resHF, apply_inv_M); Λ = map(enumerate(ψr)) do (ik, ψk) Hk = hamr.blocks[ik] Hψk = Hk * ψk - ψk'Hψk + ψk[1, :, :]'Hψk[1, :, :] end ΩpKe2 = DFTK.apply_Ω(e2, ψr, hamr, Λ) .+ DFTK.apply_K(basis_ref, e2, ψr, ρr, occ) ΩpKe2 = DFTK.transfer_blochwave(ΩpKe2, basis_ref, basis) -rhs = resLF - ΩpKe2; +rhs = denest(resLF) - denest(ΩpKe2); # - Solve the Schur system to compute ``R_{\rm Schur}(P)``: this is the most # costly step, but inverting ``\boldsymbol{Ω} + \boldsymbol{K}`` on the small space has # the same cost than the full SCF cycle on the small grid. -(; ψ) = DFTK.select_occupied_orbitals(basis, scfres.ψ, scfres.occupation) -e1 = DFTK.solve_ΩplusK(basis, ψ, rhs, occ; tol).δψ +(; ψ) = DFTK.select_occupied_orbitals(scfres.ψ, scfres.occupation) +e1 = DFTK.solve_ΩplusK(ψ, rhs, occ; tol).δψ e1 = DFTK.transfer_blochwave(e1, basis, basis_ref) -res_schur = e1 + Mres; +res_schur = denest(e1) + Mres; # ## Error estimates @@ -196,8 +197,9 @@ relerror["F(P)"] = compute_relerror(f); # To this end, we use the `ForwardDiff.jl` package to compute ``{\rm d}F(P)`` # using automatic differentiation. function df(basis, occupation, ψ, δψ, ρ) - δρ = DFTK.compute_δρ(basis, ψ, δψ, occupation) - ForwardDiff.derivative(ε -> compute_forces(basis, ψ.+ε.*δψ, occupation; ρ=ρ+ε.*δρ), 0) + δρ = DFTK.compute_δρ(ψ, δψ, occupation) + ForwardDiff.derivative(ε -> compute_forces(BlochWaves(ψ.basis, denest(ψ).+ε.*δψ), + occupation; ρ=ρ+ε.*δρ), 0) end; # - Computation of the forces by a linearization argument if we have access to diff --git a/examples/geometry_optimization.jl b/examples/geometry_optimization.jl index 8f27b6464b..4789402b7d 100644 --- a/examples/geometry_optimization.jl +++ b/examples/geometry_optimization.jl @@ -38,6 +38,11 @@ function compute_scfres(x) if isnothing(ρ) ρ = guess_density(basis) end + if isnothing(ψ) + ψ = BlochWaves(basis) + else + ψ = BlochWaves(basis, denest(ψ)) + end is_converged = DFTK.ScfConvergenceForce(tol / 10) scfres = self_consistent_field(basis; ψ, ρ, is_converged, callback=identity) ψ = scfres.ψ diff --git a/examples/gross_pitaevskii.jl b/examples/gross_pitaevskii.jl index 1f298bb0ab..c3bd43cd4b 100644 --- a/examples/gross_pitaevskii.jl +++ b/examples/gross_pitaevskii.jl @@ -57,8 +57,9 @@ scfres.energies # We use the opportunity to explore some of DFTK internals. # # Extract the converged density and the obtained wave function: -ρ = real(scfres.ρ)[:, 1, 1, 1] # converged density, first spin component -ψ_fourier = scfres.ψ[1][:, 1]; # first k-point, all G components, first eigenvector +ρ = real(scfres.ρ)[:, 1, 1, 1] # converged density, first spin component +ψ_fourier = scfres.ψ[1][1, :, 1] # first k-point, first (and only) component, + # all G components, first eigenvector # Transform the wave function to real space and fix the phase: ψ = ifft(basis, basis.kpoints[1], ψ_fourier)[:, 1, 1] @@ -85,7 +86,7 @@ plot!(p, x, ρ, label="ρ") # of a particular state (ψ, occupation). # The density ρ associated to this state is precomputed # and passed to the routine as an optimization. -E, ham = energy_hamiltonian(basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ) +E, ham = energy_hamiltonian(scfres.ψ, scfres.occupation; ρ=scfres.ρ) @assert E.total == scfres.energies.total # Now the Hamiltonian contains all the blocks corresponding to ``k``-points. Here, we just @@ -93,7 +94,7 @@ E, ham = energy_hamiltonian(basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ) H = ham.blocks[1]; # `H` can be used as a linear operator (efficiently using FFTs), or converted to a dense matrix: -ψ11 = scfres.ψ[1][:, 1] # first k-point, first eigenvector +ψ11 = scfres.ψ[1][1, :, 1] # first k-point, first component, first eigenvector Hmat = Array(H) # This is now just a plain Julia matrix, ## which we can compute and store in this simple 1D example @assert norm(Hmat * ψ11 - H * ψ11) < 1e-10 diff --git a/examples/publications/2022_cazalis.jl b/examples/publications/2022_cazalis.jl index 14d6add458..f3ec6bd86f 100644 --- a/examples/publications/2022_cazalis.jl +++ b/examples/publications/2022_cazalis.jl @@ -9,8 +9,8 @@ using Plots struct Hartree2D end struct Term2DHartree <: DFTK.TermNonlinear end (t::Hartree2D)(basis) = Term2DHartree() -function DFTK.ene_ops(term::Term2DHartree, basis::PlaneWaveBasis{T}, - ψ, occ; ρ, kwargs...) where {T} +function DFTK.ene_ops(term::Term2DHartree, ψ::BlochWaves{T}, occ; ρ, kwargs...) where {T} + basis = ψ.basis ## 2D Fourier transform of 3D Coulomb interaction 1/|x| poisson_green_coeffs = 2T(π) ./ [norm(G) for G in G_vectors_cart(basis)] poisson_green_coeffs[1] = 0 # DC component diff --git a/ext/DFTKGenericLinearAlgebraExt/DFTKGenericLinearAlgebraExt.jl b/ext/DFTKGenericLinearAlgebraExt/DFTKGenericLinearAlgebraExt.jl index 1c16df7438..451afe3168 100644 --- a/ext/DFTKGenericLinearAlgebraExt/DFTKGenericLinearAlgebraExt.jl +++ b/ext/DFTKGenericLinearAlgebraExt/DFTKGenericLinearAlgebraExt.jl @@ -29,21 +29,21 @@ end DFTK.default_primes(::Any) = (2, ) # Generic fallback function, Float32 and Float64 specialization in fft.jl -function DFTK.build_fft_plans!(tmp::AbstractArray{<:Complex}) +function DFTK.build_fft_plans!(tmp::AbstractArray{<:Complex}; region=1:ndims(tmp)) # Note: FourierTransforms has no support for in-place FFTs at the moment # ... also it's extension to multi-dimensional arrays is broken and # the algo only works for some cases @assert all(ispow2, size(tmp)) - # opFFT = AbstractFFTs.plan_fft(tmp) # TODO When multidim works + # opFFT = AbstractFFTs.plan_fft(tmp) # TODO When multidim works # opBFFT = inv(opFFT).p - opFFT = generic_plan_fft(tmp) # Fallback for now + opFFT = generic_plan_fft(tmp) # Fallback for now opBFFT = generic_plan_bfft(tmp) # TODO Can be cut once FourierTransforms supports AbstractFFTs properly ipFFT = DummyInplace{typeof(opFFT)}(opFFT) ipBFFT = DummyInplace{typeof(opBFFT)}(opBFFT) - ipFFT, opFFT, ipBFFT, opBFFT + (; ipFFT, opFFT, ipBFFT, opBFFT) end struct GenericPlan{T} @@ -51,7 +51,7 @@ struct GenericPlan{T} factor::T end -function Base.:*(p::GenericPlan, X::AbstractArray) +function Base.:*(p::GenericPlan, X::AbstractArray{T, 3}) where {T} pl1, pl2, pl3 = p.subplans ret = similar(X) for i = 1:size(X, 1), j = 1:size(X, 2) @@ -86,4 +86,18 @@ function generic_plan_bfft(data::AbstractArray{T, 3}) where {T} plan_bfft(data[1, 1, :])], T(1)) end +# TODO: Let's not bother with multicomponents yet. +function generic_plan_fft(data::AbstractArray{T, 4}) where {T} + @assert size(data, 1) == 1 + generic_plan_fft(data[1, :, :, :]) +end +function generic_plan_bfft(data::AbstractArray{T, 4}) where {T} + @assert size(data, 1) == 1 + generic_plan_bfft(data[1, :, :, :]) +end +function Base.:*(p::GenericPlan, X::AbstractArray{T, 4}) where {T} + @assert size(X, 1) == 1 + reshape(p * X[1, :, :, :], size(X)...) +end + end diff --git a/ext/DFTKIntervalArithmeticExt.jl b/ext/DFTKIntervalArithmeticExt.jl index b3d2bcac43..16fe836e64 100644 --- a/ext/DFTKIntervalArithmeticExt.jl +++ b/ext/DFTKIntervalArithmeticExt.jl @@ -42,8 +42,8 @@ function _is_well_conditioned(A::AbstractArray{<:Interval}; kwargs...) end function symmetry_operations(lattice::AbstractMatrix{<:Interval}, atoms, positions, - magnetic_moments=[]; - tol_symmetry=max(SYMMETRY_TOLERANCE, maximum(radius, lattice))) + magnetic_moments=[]; + tol_symmetry=max(SYMMETRY_TOLERANCE, maximum(radius, lattice))) @assert tol_symmetry < 1e-2 symmetry_operations(IntervalArithmetic.mid.(lattice), atoms, positions, magnetic_moments; tol_symmetry) diff --git a/ext/DFTKJLD2Ext.jl b/ext/DFTKJLD2Ext.jl index 0173dda3ca..724ae1b757 100644 --- a/ext/DFTKJLD2Ext.jl +++ b/ext/DFTKJLD2Ext.jl @@ -34,7 +34,11 @@ function DFTK.load_scfres(jld::JLD2.JLDFile) kpt_properties = (:ψ, :occupation, :eigenvalues) # Need splitting over MPI processes for sym in kpt_properties - scfdict[sym] = jld[string(sym)][basis.krange_thisproc] + if sym == :ψ + scfdict[sym] = nest(basis, jld[string(sym)][basis.krange_thisproc]) + else + scfdict[sym] = jld[string(sym)][basis.krange_thisproc] + end end for sym in jld["__propertynames"] sym in (:ham, :basis, :ρ, :energies) && continue # special @@ -42,9 +46,8 @@ function DFTK.load_scfres(jld::JLD2.JLDFile) scfdict[sym] = jld[string(sym)] end - energies, ham = energy_hamiltonian(basis, scfdict[:ψ], scfdict[:occupation]; - ρ=scfdict[:ρ], - eigenvalues=scfdict[:eigenvalues], + energies, ham = energy_hamiltonian(scfdict[:ψ], scfdict[:occupation]; + ρ=scfdict[:ρ], eigenvalues=scfdict[:eigenvalues], εF=scfdict[:εF]) scfdict[:energies] = energies diff --git a/ext/DFTKWriteVTKExt.jl b/ext/DFTKWriteVTKExt.jl index a4c3dc8152..e59d788856 100644 --- a/ext/DFTKWriteVTKExt.jl +++ b/ext/DFTKWriteVTKExt.jl @@ -20,10 +20,12 @@ function DFTK.save_scfres_master(filename::AbstractString, scfres::NamedTuple, : # Storing the bloch waves if save_ψ for ik = 1:length(basis.kpoints) - for iband = 1:size(scfres.ψ[ik])[2] - ψnk_real = ifft(basis, basis.kpoints[ik], scfres.ψ[ik][:, iband]) - vtkfile["ψ_k$(ik)_band$(iband)_real"] = real.(ψnk_real) - vtkfile["ψ_k$(ik)_band$(iband)_imag"] = imag.(ψnk_real) + for iband = 1:size(scfres.ψ[ik], 3) + ψnk_real = ifft(basis, basis.kpoints[ik], scfres.ψ[ik][:, :, iband]) + for σ = 1:basis.model.n_components + vtkfile["σ$(σ)_ψ_k$(ik)_band$(iband)_real"] = real.(ψnk_real[σ, :, :, :]) + vtkfile["σ$(σ)_ψ_k$(ik)_band$(iband)_imag"] = imag.(ψnk_real[σ, :, :, :]) + end end end end diff --git a/src/DFTK.jl b/src/DFTK.jl index 8c8e9db247..7609de88fb 100644 --- a/src/DFTK.jl +++ b/src/DFTK.jl @@ -72,6 +72,10 @@ export compute_fft_size export G_vectors, G_vectors_cart, r_vectors, r_vectors_cart export Gplusk_vectors, Gplusk_vectors_cart export Kpoint +export BlochWaves, view_component, nest, denest +export blochwave_as_matrix +export blochwave_as_tensor +export blochwaves_as_matrices export ifft export irfft export ifft! @@ -83,6 +87,7 @@ include("Model.jl") include("structure.jl") include("bzmesh.jl") include("PlaneWaveBasis.jl") +include("Psi.jl") include("fft.jl") include("orbitals.jl") include("input_output.jl") diff --git a/src/Model.jl b/src/Model.jl index 10d12ae6d6..2e5c51a43d 100644 --- a/src/Model.jl +++ b/src/Model.jl @@ -26,6 +26,10 @@ struct Model{T <: Real, VT <: Real} n_electrons::Union{Int, Nothing} εF::Union{T, Nothing} + # Option for multicomponents computations. + # TODO: Planed to replace current handling of spins, as it will allow full spin computations. + n_components::Int + # spin_polarization values: # :none No spin polarization, αα and ββ density identical, # αβ and βα blocks zero, 1 spin component treated explicitly @@ -106,7 +110,8 @@ function Model(lattice::AbstractMatrix{T}, terms=[Kinetic()], temperature=zero(T), smearing=temperature > 0 ? Smearing.FermiDirac() : Smearing.None(), - spin_polarization=determine_spin_polarization(magnetic_moments), + n_components=1, + spin_polarization=default_spin_polarization(magnetic_moments, n_components), symmetries=default_symmetries(lattice, atoms, positions, magnetic_moments, spin_polarization, terms), ) where {T <: Real} @@ -178,7 +183,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}, @@ -224,6 +230,7 @@ function Model{T}(model::Model; model.temperature, model.smearing, model.εF, + model.n_components, model.spin_polarization, model.symmetries, # Can be safely disabled: this has been checked for model @@ -249,6 +256,11 @@ normalize_magnetic_moment(mm::AbstractVector)::Vec3{Float64} = mm """Number of valence electrons.""" n_electrons_from_atoms(atoms) = sum(n_elec_valence, atoms; init=0) +function default_spin_polarization(magnetic_moments, n_components) + n_components > 1 && return :spinless + determine_spin_polarization(magnetic_moments) +end + """ Default logic to determine the symmetry operations to be used in the model. """ diff --git a/src/PlaneWaveBasis.jl b/src/PlaneWaveBasis.jl index 0e380b3240..c2505b274c 100644 --- a/src/PlaneWaveBasis.jl +++ b/src/PlaneWaveBasis.jl @@ -72,6 +72,11 @@ struct PlaneWaveBasis{T, ipBFFT fft_normalization::T # fft = fft_normalization * FFT ifft_normalization::T # ifft = ifft_normalization * BFFT + # Multicomponents plans + opFFT_mc + ipFFT_mc + opBFFT_mc + ipBFFT_mc # "cubic" basis in reciprocal and real space, on which potentials and densities are stored G_vectors::T_G_vectors @@ -213,7 +218,11 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Real, fft_size::Tuple{Int, Int, I # Setup FFT plans Gs = to_device(architecture, G_vectors(fft_size)) - (ipFFT, opFFT, ipBFFT, opBFFT) = build_fft_plans!(similar(Gs, Complex{T}, fft_size)) + (; ipFFT, opFFT, ipBFFT, opBFFT) = build_fft_plans!(similar(Gs, Complex{T}, fft_size)) + # strided FFT plans for multicomponents + ipFFT_mc, opFFT_mc, ipBFFT_mc, opBFFT_mc = + build_fft_plans!(similar(Gs, Complex{T}, model.n_components, fft_size...); + region=2:4) # Normalization constants # fft = fft_normalization * FFT @@ -294,6 +303,7 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Real, fft_size::Tuple{Int, Int, I Ecut, variational, opFFT, ipFFT, opBFFT, ipBFFT, fft_normalization, ifft_normalization, + opFFT_mc, ipFFT_mc, opBFFT_mc, ipBFFT_mc, Gs, r_vectors, kpoints, kweights, kgrid, kcoords_global, kweights_global, comm_kpts, krange_thisproc, krange_allprocs, diff --git a/src/Psi.jl b/src/Psi.jl new file mode 100644 index 0000000000..ebc53660e7 --- /dev/null +++ b/src/Psi.jl @@ -0,0 +1,55 @@ +struct BlochWaves{T, Tψ, Basis <: AbstractBasis{T}} + basis::Basis + data::Vector{VT} where {VT <: AbstractArray3{Tψ}} + n_components::Int +end +function BlochWaves(basis::PlaneWaveBasis, ψ::Vector{VT}) where {VT <: AbstractArray3} + BlochWaves(basis, ψ, basis.model.n_components) +end +function BlochWaves(basis::PlaneWaveBasis, ψ::Vector{VT}) where {VT <: AbstractMatrix} + n_components = basis.model.n_components + BlochWaves(basis, [reshape(ψk, n_components, :, size(ψk, 2)) for ψk in ψ], n_components) +end +BlochWaves(basis) = BlochWaves(basis, [Array{eltype(basis), 3}(undef, 0, 0, 0)], 0) + +# Helpers function to directly have access to the `data` field. +Base.getindex(ψ::BlochWaves, indices...) = ψ.data[indices...] +Base.isnothing(ψ::BlochWaves) = iszero(ψ.data) +Base.iterate(ψ::BlochWaves, args...) = Base.iterate(ψ.data, args...) +Base.length(ψ::BlochWaves) = length(ψ.data) +Base.similar(ψ::BlochWaves, args...) = similar(ψ.data, args...) +Base.size(ψ::BlochWaves, args...) = Base.size(ψ.data, args...) + +# T@D@: Do not change flatten the size of array by default (1:1) +# T@D@: replace with iterations with eachslice(ψk; dims=1) + +@doc raw""" + view_component(ψk::AbstractArray3, σ) + +View the ``σ``-th component(s) of the wave function `ψk`. +It returns a 2D matrix if `σ` is an integer or a 3D array if it is a list. +""" +@views view_component(ψk::AbstractArray3, σ) = ψk[σ, :, :] +# Apply the previous function for each k-point. +@views view_component(ψ::BlochWaves, σ) = [view_component(ψk, σ) for ψk in ψ] +""" + denest(ψ::BlochWaves; σ) + +Returns the arrays containing the data from the `BlochWaves` structure `ψ`. +If `σ` is given, we can ask for only some components to be extracted. +""" +@views denest(ψ::BlochWaves; σ=1:ψ.basis.model.n_components) = [view_component(ψk, σ) for ψk in ψ] +@views denest(basis, ψ::Vector; σ=1:basis.model.n_components) = [view_component(ψk, σ) for ψk in ψ] +# Wrapper around the BlochWaves creator to have an inverse to the `denest` function. +nest(basis, ψ::Vector{A}) where {A <: AbstractArray} = BlochWaves(basis, ψ) + +eachcomponent(ψk::AbstractArray3) = eachslice(ψk; dims=1) +eachband(ψk::AbstractArray3) = eachslice(ψk; dims=3) + +@views blochwave_as_tensor(ψk::AbstractMatrix, n_components) = reshape(ψk, n_components, :, size(ψk, 2)) +@views blochwave_as_matrix(ψk::AbstractArray3) = reshape(ψk, :, size(ψk, 3)) +# reduce along component direction +@views blochwave_as_matorvec(ψk::AbstractArray3) = reshape(ψk, :, size(ψk, 3)) +@views blochwave_as_matorvec(ψk::AbstractMatrix) = reshape(ψk, size(ψk, 2)) +# Works for BlochWaves & Vector(AbstractArray3). +@views blochwaves_as_matrices(ψ) = @views [reshape(ψk, :, size(ψk, 3)) for ψk in ψ] diff --git a/src/common/ortho.jl b/src/common/ortho.jl index fc5473ca3b..082ef8fb40 100644 --- a/src/common/ortho.jl +++ b/src/common/ortho.jl @@ -1,9 +1,4 @@ -@timing function ortho_qr(φk::ArrayType) where {ArrayType <: AbstractArray} - x = convert(ArrayType, qr(φk).Q) - 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)] - end +@timing function ortho_qr(φk::AbstractArray{T}) where {T} + x = convert(Matrix{T}, qr(blochwave_as_matrix(φk)).Q) + blochwave_as_tensor(x, size(φk, 1))[:, :, 1:size(φk, 3)] end diff --git a/src/common/types.jl b/src/common/types.jl index c70b8f75fc..ddd20533fb 100644 --- a/src/common/types.jl +++ b/src/common/types.jl @@ -9,3 +9,4 @@ value_type(T) = T const Mat3{T} = SMatrix{3, 3, T, 9} where {T} const Vec3{T} = SVector{3, T} where {T} const AbstractArray3{T} = AbstractArray{T, 3} +const AbstractArray4{T} = AbstractArray{T, 4} diff --git a/src/densities.jl b/src/densities.jl index 7c1b48a59b..b9560a80f7 100644 --- a/src/densities.jl +++ b/src/densities.jl @@ -2,20 +2,23 @@ # ρ[ix,iy,iz,iσ] in real space, where iσ ∈ [1:n_spin_components] """ - compute_density(basis::PlaneWaveBasis, ψ::AbstractVector, occupation::AbstractVector) + compute_density(ψ::BlochWaves, occupation::AbstractVector) -Compute the density for a wave function `ψ` discretized on the plane-wave -grid `basis`, where the individual k-points are occupied according to `occupation`. -`ψ` should be one coefficient matrix per ``k``-point. +Compute the density for a wave function `ψ` discretized on the plane-wave grid `ψ.basis`, +where the individual k-points are occupied according to `occupation`. +`ψ` should contain one coefficient matrix per ``k``-point. It is possible to ask only for occupations higher than a certain level to be computed by using an optional `occupation_threshold`. By default all occupation numbers are considered. """ -@views @timing function compute_density(basis::PlaneWaveBasis{T}, ψ, occupation; - occupation_threshold=zero(T)) where {T} - S = promote_type(T, real(eltype(ψ[1]))) +# TODO: We reduce all components for the density. Will need to be though again when we merge +# the components and the spins. +@views @timing function compute_density(ψ::BlochWaves{T, Tψ}, occupation; + occupation_threshold=zero(T)) where {T, Tψ} + S = promote_type(T, real(Tψ)) # occupation should be on the CPU as we are going to be doing scalar indexing. occupation = [to_cpu(oc) for oc in occupation] + basis = ψ.basis mask_occ = [findall(occnk -> abs(occnk) ≥ occupation_threshold, occk) for occk in occupation] if all(isempty, mask_occ) # No non-zero occupations => return zero density @@ -29,19 +32,22 @@ 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...) - for _ = 1:Threads.nthreads()] + ψσnk_real_chunklocal = [zeros_like(basis.G_vectors, complex(S), + basis.model.n_components, basis.fft_size...) + for _ = 1:Threads.nthreads()] @sync for (ichunk, chunk) in enumerate(Iterators.partition(ik_n, chunk_length)) Threads.@spawn for (ik, n) in chunk # spawn a task per chunk ρ_loc = ρ_chunklocal[ichunk] - ψnk_real = ψnk_real_chunklocal[ichunk] + ψσnk_real = ψσnk_real_chunklocal[ichunk] kpt = basis.kpoints[ik] - ifft!(ψnk_real, basis, kpt, ψ[ik][:, n]) - ρ_loc[:, :, :, kpt.spin] .+= (occupation[ik][n] .* basis.kweights[ik] - .* abs2.(ψnk_real)) + ifft!(ψσnk_real, basis, kpt, ψ[ik][:, :, n]) + for σ = 1:basis.model.n_components + ρ_loc[:, :, :, kpt.spin] .+= (occupation[ik][n] .* basis.kweights[ik] + .* abs2.(ψσnk_real[σ, :, :, :])) + end synchronize_device(basis.architecture) end end @@ -61,25 +67,28 @@ using an optional `occupation_threshold`. By default all occupation numbers are end # Variation in density corresponding to a variation in the orbitals and occupations. -@views @timing function compute_δρ(basis::PlaneWaveBasis{T}, ψ, δψ, - occupation, δoccupation=zero.(occupation); +@views @timing function compute_δρ(ψ::BlochWaves{T}, δψ, 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)] - compute_density(basis, ψ_ε, occ_ε; occupation_threshold) + compute_density(BlochWaves(ψ.basis, ψ_ε), occ_ε; occupation_threshold) 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)) +@views @timing function compute_kinetic_energy_density(ψ::BlochWaves{T, Tψ}, + occupation) where {T, Tψ} + basis = ψ.basis + @assert basis.model.n_components == 1 + TT = promote_type(T, real(Tψ)) + τ = similar(ψ[1], TT, (basis.fft_size..., basis.model.n_spin_components)) τ .= 0 - dαψnk_real = zeros(complex(eltype(basis)), basis.fft_size) + dαψnk_real = zeros(complex(T), basis.fft_size) for (ik, kpt) in enumerate(basis.kpoints) G_plus_k = [[Gk[α] for Gk in Gplusk_vectors_cart(basis, kpt)] for α = 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 diff --git a/src/eigen/diag.jl b/src/eigen/diag.jl index e39c9fbbb9..071e4b2638 100644 --- a/src/eigen/diag.jl +++ b/src/eigen/diag.jl @@ -7,48 +7,53 @@ that really does the work, operating on a single ``k``-Block. `prec_type` should be a function that returns a preconditioner when called as `prec(ham, kpt)` """ function diagonalize_all_kblocks(eigensolver, ham::Hamiltonian, nev_per_kpoint::Int; - ψguess=nothing, - prec_type=PreconditionerTPA, interpolate_kpoints=true, - tol=1e-6, miniter=1, maxiter=100, n_conv_check=nothing) - kpoints = ham.basis.kpoints + ψguess=nothing, prec_type=PreconditionerTPA, + interpolate_kpoints=true, tol=1e-6, miniter=1, maxiter=100, + n_conv_check=nothing) + basis = ham.basis + kpoints = basis.kpoints + n_components = basis.model.n_components results = Vector{Any}(undef, length(kpoints)) 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[ik], 1)) and " * + 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], 2) + nev_guess = size(ψguess[ik], 3) 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) + X0 = similar(ψguess[ik], n_components, n_Gk, nev_per_kpoint) + X0[:, :, 1:nev_guess] = ψguess[ik] + X0[:, :, nev_guess+1:end] = randn(eltype(X0), n_components, n_Gk, + nev_per_kpoint - nev_guess) ψguessk = ortho_qr(X0) 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) == (n_components, n_Gk, nev_per_kpoint) prec = nothing !isnothing(prec_type) && (prec = prec_type(ham[ik])) - results[ik] = eigensolver(ham[ik], ψguessk; - prec, tol, miniter, maxiter, n_conv_check) + res = eigensolver(ham[ik], blochwave_as_matrix(ψguessk); + prec, tol, miniter, maxiter, n_conv_check) + results[ik] = merge(res, (; X=blochwave_as_tensor(res.X, n_components))) + end # Transform results into a nicer datastructure @@ -68,7 +73,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 diff --git a/src/eigen/preconditioners.jl b/src/eigen/preconditioners.jl index 41d3588d9a..84d6b346e6 100644 --- a/src/eigen/preconditioners.jl +++ b/src/eigen/preconditioners.jl @@ -40,16 +40,23 @@ function PreconditionerTPA(basis::PlaneWaveBasis{T}, kpt::Kpoint; default_shift= # it's better to pass a HamiltonianBlock directly and read the computed values. kinetic_term = only(kinetic_term) kin = kinetic_energy(kinetic_term, basis.Ecut, Gplusk_vectors_cart(basis, kpt)) - PreconditionerTPA{T}(basis, kpt, kin, nothing, default_shift) + PreconditionerTPA{T}(basis, kpt, repeat(kin, basis.model.n_components), nothing, default_shift) end function PreconditionerTPA(ham::HamiltonianBlock; kwargs...) PreconditionerTPA(ham.basis, ham.kpoint; kwargs...) end -@views function ldiv!(Y, P::PreconditionerTPA, R) +@views function ldiv!(Y, P::PreconditionerTPA, R::AbstractArray) if P.mean_kin === nothing ldiv!(Y, Diagonal(P.kin .+ P.default_shift), R) + elseif ndims(R) == 3 + Threads.@threads for n = 1:size(Y, 3) + 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. Threads.@threads for n = 1:size(Y, 2) Y[:, n] .= P.mean_kin[n] ./ (P.mean_kin[n] .+ P.kin) .* R[:, n] end diff --git a/src/external/wannier_shared.jl b/src/external/wannier_shared.jl index 0aa90bf2bf..7c0746fc1f 100644 --- a/src/external/wannier_shared.jl +++ b/src/external/wannier_shared.jl @@ -190,13 +190,14 @@ end @timing function write_w90_unk(fileprefix::String, basis, ψ; n_bands, spin=1) + @assert basis.model.n_components == 1 fft_size = basis.fft_size for ik in krange_spin(basis, spin) open(dirname(fileprefix) * (@sprintf "/UNK%05i.%i" ik spin), "w") do fp println(fp, "$(fft_size[1]) $(fft_size[2]) $(fft_size[3]) $ik $n_bands") for iband = 1:n_bands - ψnk_real = ifft(basis, basis.kpoints[ik], @view ψ[ik][:, iband]) + ψnk_real = ifft(basis, basis.kpoints[ik], @view ψ[ik][1, :, iband]) for iz = 1:fft_size[3], iy = 1:fft_size[2], ix = 1:fft_size[1] @printf fp "%25.18f %25.18f\n" real(ψnk_real[ix, iy, iz]) imag(ψnk_real[ix, iy, iz]) end @@ -218,6 +219,7 @@ We use here that: ``u_{n(k + G_{\rm shift})}(r) = e^{-i*\langle G_{\rm shift},r \rangle} u_{nk}``. """ @views function overlap_Mmn_k_kpb(basis::PlaneWaveBasis, ψ, ik, ikpb, G_shift, n_bands) + @assert basis.model.n_components == 1 # Search for common Fourier modes and their resp. indices in Bloch states k and kpb # TODO Check if this can be improved using the G vector mapping in the kpoints k = basis.kpoints[ik] @@ -233,7 +235,7 @@ We use here that: for n = 1:n_bands for m = 1:n_bands # Select the coefficient in right order - Mkb[m, n] = dot(ψ[ik][iGk, m], ψ[ikpb][iGkpb, n]) + Mkb[m, n] = dot(ψ[ik][1, iGk, m], ψ[ikpb][1, iGkpb, n]) end end iszero(Mkb) && return Matrix(I, n_bands, n_bands) @@ -297,19 +299,19 @@ function compute_amn_kpoint(basis::PlaneWaveBasis, kpt, ψk, projections, n_band Ak end - @timing function write_w90_amn( fileprefix::String, basis::PlaneWaveBasis, projections, ψ; n_bands) + @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(projections))") for (ik, (ψk, kpt)) in enumerate(zip(ψ, basis.kpoints)) - Ak = compute_amn_kpoint(basis, kpt, ψk, projections, n_bands) + Ak = compute_amn_kpoint(basis, kpt, ψk[1, :, :], projections, n_bands) for n = 1:size(Ak, 2) for m = 1:size(Ak, 1) @printf(fp, "%3i %3i %3i %22.18f %22.18f \n", diff --git a/src/fft.jl b/src/fft.jl index c644d4c93c..0ce3b20588 100644 --- a/src/fft.jl +++ b/src/fft.jl @@ -24,8 +24,8 @@ function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, f_fourier::Abstrac mul!(f_real, basis.opBFFT, f_fourier) f_real .*= basis.ifft_normalization end -function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, - kpt::Kpoint, f_fourier::AbstractVector; normalize=true) +function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, kpt::Kpoint, + f_fourier::AbstractVector; normalize=true) @assert length(f_fourier) == length(kpt.mapping) @assert size(f_real) == basis.fft_size @@ -38,6 +38,26 @@ function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, normalize && (f_real .*= basis.ifft_normalization) f_real end +# For multicomponents +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) + n_components = basis.model.n_components + @assert length(f_fourier) == n_components*length(kpt.mapping) + @assert size(f_real)[2:end] == basis.fft_size + + # Pad the input data + fill!(f_real, 0) + reshape(f_real, n_components, :)[:, kpt.mapping] = f_fourier + + # Perform an IFFT + mul!(f_real, basis.ipBFFT_mc, f_real) + normalize && (f_real .*= basis.ifft_normalization) + f_real +end """ ifft(basis::PlaneWaveBasis, [kpt::Kpoint, ] f_fourier) @@ -58,6 +78,10 @@ end function ifft(basis::PlaneWaveBasis, kpt::Kpoint, f_fourier::AbstractVector; kwargs...) ifft!(similar(f_fourier, basis.fft_size...), basis, kpt, f_fourier; kwargs...) end +function ifft(basis::PlaneWaveBasis, kpt::Kpoint, f_fourier::AbstractMatrix; kwargs...) + ifft!(similar(f_fourier, basis.model.n_components, basis.fft_size...), + basis, kpt, f_fourier; kwargs...) +end """ Perform a real valued iFFT; see [`ifft`](@ref). Note that this function @@ -92,6 +116,29 @@ function fft!(f_fourier::AbstractVector, basis::PlaneWaveBasis, normalize && (f_fourier .*= basis.fft_normalization) f_fourier end +# For multicomponents +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) + @assert size(f_real)[2:end] == basis.fft_size + @assert length(f_fourier) == basis.model.n_components * length(kpt.mapping) + + # FFT + mul!(f_real, basis.ipFFT_mc, f_real) + + # Truncate + for σ = 1:basis.model.n_components + f_fourier[σ, :] .= @view(f_real[σ, :, :, :][kpt.mapping]) + end + normalize && (f_fourier .*= basis.fft_normalization) + f_fourier +end """ fft(basis::PlaneWaveBasis, [kpt::Kpoint, ] f_real) @@ -263,25 +310,32 @@ end Plan a FFT of type `T` and size `fft_size`, spending some time on finding an optimal algorithm. (Inplace, out-of-place) x (forward, backward) FFT plans are returned. """ -function build_fft_plans!(tmp::Array{Complex{Float64}}) - ipFFT = FFTW.plan_fft!(tmp; flags=FFTW.MEASURE) - opFFT = FFTW.plan_fft(tmp; flags=FFTW.MEASURE) +function build_fft_plans!(tmp::Array{Complex{Float64}}; region=1:ndims(tmp)) + ipFFT = FFTW.plan_fft!(tmp, region; flags=FFTW.MEASURE) + opFFT = FFTW.plan_fft(tmp, region; flags=FFTW.MEASURE) # backwards-FFT by inverting and stripping off normalizations - ipFFT, opFFT, inv(ipFFT).p, inv(opFFT).p + ipBFFT = inv(ipFFT).p + opBFFT = inv(opFFT).p + (; ipFFT, opFFT, ipBFFT, opBFFT) end -function build_fft_plans!(tmp::Array{Complex{Float32}}) +function build_fft_plans!(tmp::Array{Complex{Float32}}; region=1:ndims(tmp)) # For Float32 there are issues with aligned FFTW plans, so we # fall back to unaligned FFTW plans (which are generally discouraged). - ipFFT = FFTW.plan_fft!(tmp; flags=FFTW.MEASURE | FFTW.UNALIGNED) - opFFT = FFTW.plan_fft(tmp; flags=FFTW.MEASURE | FFTW.UNALIGNED) + ipFFT = FFTW.plan_fft!(tmp, region; flags=FFTW.MEASURE | FFTW.UNALIGNED) + opFFT = FFTW.plan_fft(tmp, region; flags=FFTW.MEASURE | FFTW.UNALIGNED) # backwards-FFT by inverting and stripping off normalizations - ipFFT, opFFT, inv(ipFFT).p, inv(opFFT).p + ipBFFT = inv(ipFFT).p + opBFFT = inv(opFFT).p + (; ipFFT, opFFT, ipBFFT, opBFFT) end -function build_fft_plans!(tmp::AbstractArray{Complex{T}}) where {T<:Union{Float32,Float64}} - ipFFT = AbstractFFTs.plan_fft!(tmp) - opFFT = AbstractFFTs.plan_fft(tmp) +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 - ipFFT, opFFT, inv(ipFFT).p, inv(opFFT).p + 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/input_output.jl b/src/input_output.jl index c761de0644..fa2478559b 100644 --- a/src/input_output.jl +++ b/src/input_output.jl @@ -8,6 +8,7 @@ function Base.show(io::IO, model::Model) nD = model.n_dim == 3 ? "" : "$(model.n_dim)D, " print(io, "Model(", model.model_name, ", ", nD, + "n_components = ", model.n_components, ", ", "spin_polarization = :", model.spin_polarization, ")") end @@ -33,6 +34,7 @@ function Base.show(io::IO, ::MIME"text/plain", model::Model) if !isnothing(model.n_electrons) showfieldln(io, "num. electrons", model.n_electrons) end + showfieldln(io, "num. components", model.n_components) showfieldln(io, "spin polarization", model.spin_polarization) showfieldln(io, "temperature", @sprintf "%.5g Ha" model.temperature) if model.temperature > 0 diff --git a/src/interpolation.jl b/src/interpolation.jl index 1a219dc69a..f22e2fad73 100644 --- a/src/interpolation.jl +++ b/src/interpolation.jl @@ -83,24 +83,26 @@ 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 = 1:size(data_in, 1) + for iin = 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[iout, :] = data_in[iin, :] + data_out[:, iout, :] = data_in[:, iin, :] end ortho_qr(data_out) # Re-orthogonalize and renormalize end diff --git a/src/orbitals.jl b/src/orbitals.jl index e4e7e6461f..4ef4bdd418 100644 --- a/src/orbitals.jl +++ b/src/orbitals.jl @@ -4,20 +4,20 @@ using Random # Used to have a generic API for CPU and GPU computations alike: s # virtual states (or states with small occupation level for metals). # threshold is a parameter to distinguish between states we want to keep and the # 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) +function select_occupied_orbitals(ψ, 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[:, :, 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 + ψ = BlochWaves(ψ.basis, selected_ψ) + # If we have an insulator, sanity check that the orbitals we kept are the occupied ones. if iszero(threshold) - model = basis.model + 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 all([n_bands == size(ψk, 3) for ψk in ψ]) end - (; ψ=selected_ψ, occupation=selected_occ) + (; ψ, occupation=selected_occ) end # Packing routines used in direct_minimization and newton algorithms. @@ -54,7 +54,8 @@ end unpack_ψ(x, sizes_ψ) = deepcopy(unsafe_unpack_ψ(x, sizes_ψ)) function random_orbitals(basis::PlaneWaveBasis{T}, kpt::Kpoint, howmany::Integer) where {T} - orbitals = similar(basis.G_vectors, Complex{T}, length(G_vectors(basis, kpt)), howmany) + orbitals = similar(basis.G_vectors, Complex{T}, + basis.model.n_components, length(G_vectors(basis, kpt)), howmany) randn!(TaskLocalRNG(), orbitals) # use the RNG on the device if we're using a GPU ortho_qr(orbitals) end diff --git a/src/postprocess/dos.jl b/src/postprocess/dos.jl index 141dd1379d..18878acbaf 100644 --- a/src/postprocess/dos.jl +++ b/src/postprocess/dos.jl @@ -57,7 +57,7 @@ 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) + compute_density(ψ, weights; occupation_threshold=weight_threshold) end function compute_ldos(scfres::NamedTuple; ε=scfres.εF, kwargs...) compute_ldos(ε, scfres.basis, scfres.eigenvalues, scfres.ψ; kwargs...) diff --git a/src/postprocess/forces.jl b/src/postprocess/forces.jl index f18622c7f0..a54dd74f1b 100644 --- a/src/postprocess/forces.jl +++ b/src/postprocess/forces.jl @@ -5,10 +5,11 @@ lattice vectors. To get cartesian forces use [`compute_forces_cart`](@ref). Returns a list of lists of forces (as SVector{3}) in the same order as the `atoms` and `positions` in the underlying [`Model`](@ref). """ -@timing function compute_forces(basis::PlaneWaveBasis{T}, ψ, occupation; kwargs...) where {T} +@timing function compute_forces(ψ::BlochWaves{T}, occupation; kwargs...) where {T} + basis = ψ.basis # no explicit symmetrization is performed here, it is the # responsability of each term to return symmetric forces - forces_per_term = [compute_forces(term, basis, ψ, occupation; kwargs...) + forces_per_term = [compute_forces(term, ψ, occupation; kwargs...) for term in basis.terms] sum(filter(!isnothing, forces_per_term)) end @@ -19,14 +20,14 @@ Returns a list of lists of forces `[[force for atom in positions] for (element, positions) in atoms]` which has the same structure as the `atoms` object passed to the underlying [`Model`](@ref). """ -function compute_forces_cart(basis::PlaneWaveBasis, ψ, occupation; kwargs...) - forces_reduced = compute_forces(basis, ψ, occupation; kwargs...) - covector_red_to_cart.(basis.model, forces_reduced) +function compute_forces_cart(ψ::BlochWaves, occupation; kwargs...) + forces_reduced = compute_forces(ψ, occupation; kwargs...) + covector_red_to_cart.(ψ.basis.model, forces_reduced) end function compute_forces(scfres) - compute_forces(scfres.basis, scfres.ψ, scfres.occupation; scfres.ρ) + compute_forces(scfres.ψ, scfres.occupation; scfres.ρ) end function compute_forces_cart(scfres) - compute_forces_cart(scfres.basis, scfres.ψ, scfres.occupation; scfres.ρ) + compute_forces_cart(scfres.ψ, scfres.occupation; scfres.ρ) end diff --git a/src/postprocess/stresses.jl b/src/postprocess/stresses.jl index f1c4de4412..ff31994d84 100644 --- a/src/postprocess/stresses.jl +++ b/src/postprocess/stresses.jl @@ -12,8 +12,9 @@ Compute the stresses (= 1/Vol dE/d(M*lattice), taken at M=I) of an obtained SCF basis.kgrid, basis.symmetries_respect_rgrid, basis.use_symmetries_for_kpoint_reduction, basis.comm_kpts, basis.architecture) - ρ = compute_density(new_basis, scfres.ψ, scfres.occupation) - energies = energy_hamiltonian(new_basis, scfres.ψ, scfres.occupation; + ψ = BlochWaves(new_basis, denest(scfres.ψ)) + ρ = compute_density(ψ, scfres.occupation) + energies = energy_hamiltonian(ψ, scfres.occupation; ρ, scfres.eigenvalues, scfres.εF).energies energies.total end diff --git a/src/response/chi0.jl b/src/response/chi0.jl index f60a274795..9d19468710 100644 --- a/src/response/chi0.jl +++ b/src/response/chi0.jl @@ -174,7 +174,7 @@ function sternheimer_solver(Hk, ψk, ε, rhs; function R_ldiv!(x, y) x .= R(precon \ R(y)) end - J = LinearMap{eltype(ψk)}(RAR, size(Hk, 1)) + J = LinearMap{eltype(ψk)}(RAR, size(ψk, 1)) res = cg(J, bb; precon=FunctionPreconditioner(R_ldiv!), tol, proj=R, callback=cg_callback) !res.converged && @warn("Sternheimer CG not converged", res.n_iter, @@ -189,6 +189,20 @@ function sternheimer_solver(Hk, ψk, ε, rhs; δψkn = ψk_extra * αkn + δψknᴿ end +function sternheimer_solver_wrapper(Hk, ψk, ε, rhs; + callback=identity, cg_callback=identity, + ψk_extra=zeros_like(ψk, size(ψk)[1:2]..., 0), + εk_extra=zeros(0), + Hψk_extra=zeros_like(ψk, size(ψk)[1:2]..., 0), + tol=1e-9) + ψk_extra_r = reshape(ψk_extra, prod(size(ψk_extra)[1:2]), :) + Hψk_extra_r = reshape(Hψk_extra, prod(size(Hψk_extra)[1:2]), :) + ψk_r = reshape(ψk, :, size(ψk, 3)) + δψkn = sternheimer_solver(Hk, ψk_r, ε, vec(rhs); callback, cg_callback, + ψk_extra=ψk_extra_r, εk_extra, + Hψk_extra=Hψk_extra_r, tol) + reshape(δψkn, size(ψk)[1:2]...) +end # Apply the four-point polarizability operator χ0_4P = -Ω^-1 # Returns (δψ, δocc, δεF) corresponding to a change in *total* Hamiltonian δH @@ -265,7 +279,7 @@ function compute_δocc!(δocc, basis, ψ, εF::T, ε, δHψ) where {T} D = zero(T) for ik = 1:Nk, (n, εnk) in enumerate(ε[ik]) enred = (εnk - εF) / temperature - δεnk = real(dot(ψ[ik][:, n], δHψ[ik][:, n])) + δεnk = real(dot(ψ[ik][:, :, n], δHψ[ik][:, :, n])) fpnk = filled_occ * Smearing.occupation_derivative(smearing, enred) / temperature δocc[ik][n] = δεnk * fpnk D += fpnk * basis.kweights[ik] @@ -290,7 +304,7 @@ Perform in-place computations of the derivatives of the wave functions by solvin a Sternheimer equation for each `k`-points. It is assumed the passed `δψ` are initialised to zero. """ -function compute_δψ!(δψ, basis, H, ψ, εF, ε, δHψ; ψ_extra=[zeros_like(ψk, size(ψk,1), 0) for ψk in ψ], +function compute_δψ!(δψ, basis, H, ψ, εF, ε, δHψ; ψ_extra=[zeros_like(ψk, size(ψk)[1:2]..., 0) for ψk in ψ], kwargs_sternheimer...) model = basis.model temperature = model.temperature @@ -307,7 +321,12 @@ function compute_δψ!(δψ, basis, H, ψ, εF, ε, δHψ; ψ_extra=[zeros_like( ψk_extra = ψ_extra[ik] Hψk_extra = Hk * ψk_extra - εk_extra = diag(real.(ψk_extra' * Hψk_extra)) + εk_extra = let + lead_dim = prod(size(ψk_extra)[1:2]) + ψk_extra_r = reshape(ψk_extra, lead_dim, :) + Hψk_extra_r = reshape(Hψk_extra, lead_dim, :) + diag(real.(ψk_extra_r' * Hψk_extra_r)) + end for n = 1:length(εk) fnk = filled_occ * Smearing.occupation(smearing, (εk[n]-εF) / temperature) @@ -320,12 +339,13 @@ function compute_δψ!(δψ, basis, H, ψ, εF, ε, δHψ; ψ_extra=[zeros_like( ddiff = Smearing.occupation_divided_difference ratio = filled_occ * ddiff(smearing, εk[m], εk[n], εF, temperature) αmn = compute_αmn(fmk, fnk, ratio) # fnk * αmn + fmk * αnm = ratio - δψk[:, n] .+= ψk[:, m] .* αmn .* dot(ψk[:, m], δHψ[ik][:, n]) + δψk[:, :, n] .+= ψk[:, :, m] .* αmn .* dot(ψk[:, :, m], δHψ[ik][:, :, n]) end # Sternheimer contribution - δψk[:, n] .+= sternheimer_solver(Hk, ψk, εk[n], δHψ[ik][:, n]; ψk_extra, - εk_extra, Hψk_extra, kwargs_sternheimer...) + δψk[:, :, n] .+= sternheimer_solver_wrapper(Hk, ψk, εk[n], δHψ[ik][:, :, n]; + ψk_extra, εk_extra, Hψk_extra, + kwargs_sternheimer...) end end end @@ -344,11 +364,11 @@ end mask_occ = map(occk -> findall(occnk -> abs(occnk) ≥ occ_thresh, occk), occupation) mask_extra = map(occk -> findall(occnk -> abs(occnk) < occ_thresh, occk), occupation) - ψ_occ = [ψ[ik][:, maskk] for (ik, maskk) in enumerate(mask_occ)] - ψ_extra = [ψ[ik][:, maskk] for (ik, maskk) in enumerate(mask_extra)] + ψ_occ = [ψ[ik][:, :, maskk] for (ik, maskk) in enumerate(mask_occ)] + ψ_extra = [ψ[ik][:, :, maskk] for (ik, maskk) in enumerate(mask_extra)] ε_occ = [eigenvalues[ik][maskk] for (ik, maskk) in enumerate(mask_occ)] - δHψ_occ = [δHψ[ik][:, maskk] for (ik, maskk) in enumerate(mask_occ)] + δHψ_occ = [δHψ[ik][:, :, maskk] for (ik, maskk) in enumerate(mask_occ)] # First we compute δoccupation. We only need to do this for the actually occupied # orbitals. So we make a fresh array padded with zeros, but only alter the elements @@ -360,7 +380,7 @@ end # Then we compute δψ, again in-place into a zero-padded array δψ = zero.(ψ) - δψ_occ = [δψ[ik][:, maskk] for (ik, maskk) in enumerate(mask_occ)] + δψ_occ = [δψ[ik][:, :, maskk] for (ik, maskk) in enumerate(mask_occ)] compute_δψ!(δψ_occ, basis, ham.blocks, ψ_occ, εF, ε_occ, δHψ_occ; ψ_extra, kwargs_sternheimer...) @@ -390,9 +410,9 @@ function apply_χ0(ham, ψ, occupation, εF, eigenvalues, δV::AbstractArray{T}; δHψ = [RealSpaceMultiplication(basis, kpt, @views δV[:, :, :, kpt.spin]) * ψ[ik] for (ik, kpt) in enumerate(basis.kpoints)] - (; δψ, δoccupation) = apply_χ0_4P(ham, ψ, occupation, εF, eigenvalues, δHψ; + (; δψ, δoccupation) = apply_χ0_4P(ham, denest(ψ), occupation, εF, eigenvalues, δHψ; occupation_threshold, kwargs_sternheimer...) - δρ = compute_δρ(basis, ψ, δψ, occupation, δoccupation; occupation_threshold) + δρ = compute_δρ(ψ, δψ, occupation, δoccupation; occupation_threshold) δρ * normδV end diff --git a/src/response/hessian.jl b/src/response/hessian.jl index 5dc433225a..6ac5e376b5 100644 --- a/src/response/hessian.jl +++ b/src/response/hessian.jl @@ -26,9 +26,12 @@ Compute the application of Ω defined at ψ to δψ. H is the Hamiltonian comput from ψ and Λ is the set of Rayleigh coefficients ψk' * Hk * ψk at each k-point. """ @timing function apply_Ω(δψ, ψ, H::Hamiltonian, Λ) + n_components = H.basis.model.n_components + @assert n_components == 1 δψ = proj_tangent(δψ, ψ) - Ωδψ = [H.blocks[ik] * δψk - δψk * Λ[ik] for (ik, δψk) in enumerate(δψ)] - proj_tangent!(Ωδψ, ψ) + Ωδψ = [H.blocks[ik] * δψk - δψk * Λ[ik] + for (ik, δψk) in enumerate(blochwaves_as_matrices(δψ))] + proj_tangent!(blochwave_as_tensor.(Ωδψ, n_components), ψ) end """ @@ -37,19 +40,21 @@ end Compute the application of K defined at ψ to δψ. ρ is the density issued from ψ. δψ also generates a δρ, computed with `compute_δρ`. """ +# T@D@ basis redundant; change signature maybe? @views @timing function apply_K(basis::PlaneWaveBasis, δψ, ψ, ρ, occupation) δψ = proj_tangent(δψ, ψ) - δρ = compute_δρ(basis, ψ, δψ, occupation) + δρ = compute_δρ(ψ, δψ, occupation) δV = apply_kernel(basis, δρ; ρ) Kδψ = map(enumerate(ψ)) do (ik, ψk) kpt = basis.kpoints[ik] δVψk = similar(ψk) - for n = 1:size(ψk, 2) - ψnk_real = ifft(basis, kpt, ψk[:, n]) - δVψnk_real = δV[:, :, :, kpt.spin] .* ψnk_real - δVψk[:, n] = fft(basis, kpt, δVψnk_real) + for n = 1:size(ψk, 3) + ψnk_real = ifft(basis, kpt, ψk[:, :, n]) + δVψnk_real = reduce(hcat, δV[:, :, :, kpt.spin] .* ψnk_real[σ, :, :, :] + for σ = 1:size(ψk, 1)) + δVψk[:, :, n] = fft(basis, kpt, δVψnk_real) end δVψk end @@ -58,13 +63,14 @@ Compute the application of K defined at ψ to δψ. ρ is the density issued fro end """ - solve_ΩplusK(basis::PlaneWaveBasis{T}, ψ, res, occupation; + solve_ΩplusK(ψ::BlochWaves{T}, rhs, occupation; tol=1e-10, verbose=false) where {T} Return δψ where (Ω+K) δψ = rhs """ -@timing function solve_ΩplusK(basis::PlaneWaveBasis{T}, ψ, rhs, occupation; - callback=identity, tol=1e-10) where {T} +@timing function solve_ΩplusK(ψ::BlochWaves{T}, rhs, occupation; callback=identity, + tol=1e-10) where {T} + basis = ψ.basis filled_occ = filled_occupation(basis.model) # for now, all orbitals have to be fully occupied -> need to strip them beforehand @assert all(all(occ_k .== filled_occ) for occ_k in occupation) @@ -75,9 +81,10 @@ Return δψ where (Ω+K) δψ = rhs @assert mpi_nprocs() == 1 # Distributed implementation not yet available # compute quantites at the point which define the tangent space - ρ = compute_density(basis, ψ, occupation) - H = energy_hamiltonian(basis, ψ, occupation; ρ).ham + ρ = compute_density(ψ, occupation) + H = energy_hamiltonian(ψ, occupation; ρ).ham + ψ_matrices = blochwaves_as_matrices(ψ) pack(ψ) = reinterpret_real(pack_ψ(ψ)) unpack(x) = unpack_ψ(reinterpret_complex(x), size.(ψ)) unsafe_unpack(x) = unsafe_unpack_ψ(reinterpret_complex(x), size.(ψ)) @@ -89,18 +96,19 @@ Return δψ where (Ω+K) δψ = rhs # preconditioner Pks = [PreconditionerTPA(basis, kpt) for kpt in basis.kpoints] for ik = 1:length(Pks) - precondprep!(Pks[ik], ψ[ik]) + precondprep!(Pks[ik], ψ_matrices[ik]) end function f_ldiv!(x, y) δψ = unpack(y) proj_tangent!(δψ, ψ) Pδψ = [ Pks[ik] \ δψk for (ik, δψk) in enumerate(δψ)] - proj_tangent!(Pδψ, ψ) + # T@D@ revert deepcopy + Pδψ = deepcopy(proj_tangent!(Pδψ, ψ)) x .= pack(Pδψ) end # Rayleigh-coefficients - Λ = [ψk'Hψk for (ψk, Hψk) in zip(ψ, H * ψ)] + Λ = [ψk'Hψk for (ψk, Hψk) in zip(ψ_matrices, H * ψ_matrices)] # mapping of the linear system on the tangent space function ΩpK(x) @@ -114,7 +122,8 @@ Return δψ where (Ω+K) δψ = rhs # solve (Ω+K) δψ = rhs on the tangent space with CG function proj(x) δψ = unpack(x) - proj_tangent!(δψ, ψ) + # T@D@ revert deepcopy + δψ = deepcopy(proj_tangent!(δψ, ψ)) pack(δψ) end res = cg(J, rhs_pack; precon=FunctionPreconditioner(f_ldiv!), proj, tol, @@ -145,11 +154,12 @@ Solve the problem `(Ω+K) δψ = rhs` using a split algorithm, where `rhs` is ty basis = ham.basis @assert size(rhs[1]) == size(ψ[1]) # Assume the same number of bands in ψ and rhs + ψ_array = denest(ψ) # compute δρ0 (ignoring interactions) - δψ0, δoccupation0 = apply_χ0_4P(ham, ψ, occupation, εF, eigenvalues, -rhs; + δψ0, δoccupation0 = apply_χ0_4P(ham, ψ_array, occupation, εF, eigenvalues, -rhs; tol=tol_sternheimer, occupation_threshold, kwargs...) # = -χ04P * rhs - δρ0 = compute_δρ(basis, ψ, δψ0, occupation, δoccupation0; occupation_threshold) + δρ0 = compute_δρ(ψ, δψ0, occupation, δoccupation0; occupation_threshold) # compute total δρ pack(δρ) = vec(δρ) @@ -176,13 +186,13 @@ Solve the problem `(Ω+K) δψ = rhs` using a split algorithm, where `rhs` is ty end # Compute total change in eigenvalues - δeigenvalues = map(ψ, δHψ) do ψk, δHψk - map(eachcol(ψk), eachcol(δHψk)) do ψnk, δHψnk + δeigenvalues = map(ψ_array, δHψ) do ψk, δHψk + map(eachslice(ψk; dims=3), eachslice(δHψk; dims=3)) do ψnk, δHψnk real(dot(ψnk, δHψnk)) # δε_{nk} = <ψnk | δH | ψnk> end end - δψ, δoccupation, δεF = apply_χ0_4P(ham, ψ, occupation, εF, eigenvalues, δHψ; + δψ, δoccupation, δεF = apply_χ0_4P(ham, ψ_array, occupation, εF, eigenvalues, δHψ; occupation_threshold, tol=tol_sternheimer, kwargs...) diff --git a/src/scf/direct_minimization.jl b/src/scf/direct_minimization.jl index 99c9e34454..1ef2c60493 100644 --- a/src/scf/direct_minimization.jl +++ b/src/scf/direct_minimization.jl @@ -63,15 +63,18 @@ Computes the ground state by direct minimization. `kwargs...` are passed to `Optim.Options()`. Note that the resulting ψ are not necessarily eigenvectors of the Hamiltonian. """ -direct_minimization(basis::PlaneWaveBasis; kwargs...) = direct_minimization(basis, nothing; kwargs...) -function direct_minimization(basis::PlaneWaveBasis{T}, ψ0; - prec_type=PreconditionerTPA, maxiter=1_000, +direct_minimization(basis::PlaneWaveBasis; kwargs...) = + direct_minimization(BlochWaves(basis); kwargs...) + +function direct_minimization(ψ0::BlochWaves{T}; prec_type=PreconditionerTPA, maxiter=1_000, optim_solver=Optim.LBFGS, tol=1e-6, kwargs...) where {T} if mpi_nprocs() > 1 # need synchronization in Optim error("Direct minimization with MPI is not supported yet") end + basis = ψ0.basis model = basis.model + @assert model.n_components == 1 @assert iszero(model.temperature) # temperature is not yet supported @assert isnothing(model.εF) # neither are computations with fixed Fermi level filled_occ = filled_occupation(model) @@ -79,16 +82,17 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0; n_bands = div(model.n_electrons, n_spin * filled_occ, RoundUp) Nk = length(basis.kpoints) - if ψ0 === nothing - ψ0 = [random_orbitals(basis, kpt, n_bands) for kpt in basis.kpoints] + if isnothing(ψ0) + ψ0 = BlochWaves(basis, [random_orbitals(basis, kpt, n_bands) for kpt in basis.kpoints]) end + ψ0_matrices = blochwaves_as_matrices(ψ0) occupation = [filled_occ * ones(T, n_bands) for _ = 1:Nk] # we need to copy the reinterpret array here to not raise errors in Optim.jl # TODO raise this issue in Optim.jl pack(ψ) = copy(reinterpret_real(pack_ψ(ψ))) - unpack(x) = unpack_ψ(reinterpret_complex(x), size.(ψ0)) - unsafe_unpack(x) = unsafe_unpack_ψ(reinterpret_complex(x), size.(ψ0)) + unpack(x) = unpack_ψ(reinterpret_complex(x), size.(ψ0_matrices)) + unsafe_unpack(x) = unsafe_unpack_ψ(reinterpret_complex(x), size.(ψ0_matrices)) # this will get updated along the iterations H = nothing @@ -96,10 +100,10 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0; ρ = nothing # computes energies and gradients - function fg!(E, G, ψ) + function fg!(::Any, G, ψ) ψ = unpack(ψ) - ρ = compute_density(basis, ψ, occupation) - energies, H = energy_hamiltonian(basis, ψ, occupation; ρ) + ρ = compute_density(BlochWaves(basis, ψ), occupation) + energies, H = energy_hamiltonian(BlochWaves(basis, ψ), occupation; ρ) # The energy has terms like occ * <ψ|H|ψ>, so the gradient is 2occ Hψ if G !== nothing @@ -123,7 +127,7 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0; f_tol=pop!(kwdict, :f_tol, -1), g_tol=pop!(kwdict, :g_tol, -1), iterations=maxiter, kwdict...) - res = Optim.optimize(Optim.only_fg!(fg!), pack(ψ0), + res = Optim.optimize(Optim.only_fg!(fg!), pack(ψ0_matrices), optim_solver(; P, precondprep=precondprep!, manifold, linesearch=LineSearches.BackTracking()), optim_options) @@ -133,9 +137,9 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0; eigenvalues = [] for ik = 1:Nk Hψk = H.blocks[ik] * ψ[ik] - F = eigen(Hermitian(ψ[ik]'Hψk)) + F = eigen(Hermitian(ψ[ik][1, :, :]'Hψk[1, :, :])) push!(eigenvalues, F.values) - ψ[ik] .= ψ[ik] * F.vectors + ψ[ik][1, :, :] .= ψ[ik][1, :, :] * F.vectors end εF = nothing # does not necessarily make sense here, as the @@ -143,5 +147,6 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0; # We rely on the fact that the last point where fg! was called is the minimizer to # avoid recomputing at ψ - (; ham=H, basis, energies, converged=true, ρ, ψ, eigenvalues, occupation, εF, optim_res=res) + (; ham=H, basis, energies, converged=true, ρ, ψ=BlochWaves(basis, ψ), eigenvalues, + occupation, εF, optim_res=res) end diff --git a/src/scf/nbands_algorithm.jl b/src/scf/nbands_algorithm.jl index 77139c07da..aa2f93caec 100644 --- a/src/scf/nbands_algorithm.jl +++ b/src/scf/nbands_algorithm.jl @@ -59,7 +59,7 @@ function determine_n_bands(bands::AdaptiveBands, occupation::Nothing, eigenvalue if isnothing(ψ) n_bands_compute = bands.n_bands_compute else - n_bands_compute = max(bands.n_bands_compute, maximum(ψk -> size(ψk, 2), ψ)) + n_bands_compute = max(bands.n_bands_compute, maximum(ψk -> size(ψk, 3), ψ)) end # Boost number of bands to converge to have more information around in the next step # and to thus make a better decision on the number of bands we actually care about. @@ -67,7 +67,7 @@ function determine_n_bands(bands::AdaptiveBands, occupation::Nothing, eigenvalue (; n_bands_converge, n_bands_compute) end function determine_n_bands(bands::AdaptiveBands, occupation::AbstractVector, - eigenvalues::AbstractVector, ψ::AbstractVector) + eigenvalues::AbstractVector, ψ::BlochWaves) # TODO Could return different bands per k-Points # Determine number of bands to be actually converged @@ -87,7 +87,7 @@ function determine_n_bands(bands::AdaptiveBands, occupation::AbstractVector, end n_bands_compute = max(bands.n_bands_compute, n_bands_compute_ε, n_bands_converge + 3) if !isnothing(ψ) - n_bands_compute = max(n_bands_compute, maximum(ψk -> size(ψk, 2), ψ)) + n_bands_compute = max(n_bands_compute, maximum(ψk -> size(ψk, 3), ψ)) end (; n_bands_converge, n_bands_compute) end diff --git a/src/scf/newton.jl b/src/scf/newton.jl index f8b1f1e713..3fd59aeef6 100644 --- a/src/scf/newton.jl +++ b/src/scf/newton.jl @@ -49,9 +49,9 @@ using IterativeSolvers # Compute the gradient of the energy, projected on the space tangent to ψ, that # is to say H(ψ)*ψ - ψ*λ where λ is the set of Rayleigh coefficients associated # to the ψ. -function compute_projected_gradient(basis::PlaneWaveBasis, ψ, occupation) - ρ = compute_density(basis, ψ, occupation) - H = energy_hamiltonian(basis, ψ, occupation; ρ).ham +function compute_projected_gradient(ψ, occupation) + ρ = compute_density(ψ, occupation) + H = energy_hamiltonian(ψ, occupation; ρ).ham [proj_tangent_kpt(H.blocks[ik] * ψk, ψk) for (ik, ψk) in enumerate(ψ)] end @@ -59,7 +59,10 @@ end # Projections on the space tangent to ψ function proj_tangent_kpt!(δψk, ψk) # δψk = δψk - ψk * (ψk'δψk) - mul!(δψk, ψk, ψk'δψk, -1, 1) + δψk_matrix = blochwave_as_matorvec(δψk) + ψk_matrix = blochwave_as_matorvec(ψk) + mul!(δψk_matrix, ψk_matrix, ψk_matrix'δψk_matrix, -1, 1) + δψk end proj_tangent_kpt(δψk, ψk) = proj_tangent_kpt!(deepcopy(δψk), ψk) @@ -72,20 +75,20 @@ end """ - newton(basis::PlaneWaveBasis{T}, ψ0; + newton(ψ0::BlochWaves{T}; tol=1e-6, tol_cg=tol / 100, maxiter=20, callback=ScfDefaultCallback(), is_converged=ScfConvergenceDensity(tol)) Newton algorithm. Be careful that the starting point needs to be not too far from the solution. """ -function newton(basis::PlaneWaveBasis{T}, ψ0; - tol=1e-6, tol_cg=tol / 100, maxiter=20, - callback=ScfDefaultCallback(), - is_converged=ScfConvergenceDensity(tol)) where {T} +function newton(ψ0::BlochWaves{T}; tol=1e-6, tol_cg=tol / 100, maxiter=20, + callback=ScfDefaultCallback(), is_converged=ScfConvergenceDensity(tol)) where {T} # setting parameters + basis = ψ0.basis model = basis.model + @assert model.n_components == 1 @assert iszero(model.temperature) # temperature is not yet supported @assert isnothing(model.εF) # neither are computations with fixed Fermi level @@ -93,7 +96,7 @@ function newton(basis::PlaneWaveBasis{T}, ψ0; filled_occ = filled_occupation(model) n_spin = model.n_spin_components n_bands = div(model.n_electrons, n_spin * filled_occ, RoundUp) - @assert n_bands == size(ψ0[1], 2) + @assert all([n_bands == size(ψk, 3) for ψk in ψ0]) # number of kpoints and occupation Nk = length(basis.kpoints) @@ -103,9 +106,9 @@ function newton(basis::PlaneWaveBasis{T}, ψ0; n_iter = 0 # orbitals, densities and energies to be updated along the iterations - ψ = deepcopy(ψ0) - ρ = compute_density(basis, ψ, occupation) - energies, H = energy_hamiltonian(basis, ψ, occupation; ρ) + ψ = BlochWaves(basis, denest(ψ0)) + ρ = compute_density(ψ, occupation) + energies, H = energy_hamiltonian(ψ, occupation; ρ) converged = false # perform iterations @@ -113,13 +116,13 @@ function newton(basis::PlaneWaveBasis{T}, ψ0; n_iter += 1 # compute Newton step and next iteration - res = compute_projected_gradient(basis, ψ, occupation) + res = compute_projected_gradient(ψ, occupation) # solve (Ω+K) δψ = -res so that the Newton step is ψ <- ψ + δψ - δψ = solve_ΩplusK(basis, ψ, -res, occupation; tol=tol_cg).δψ - ψ = [ortho_qr(ψ[ik] + δψ[ik]) for ik = 1:Nk] + δψ = solve_ΩplusK(ψ, -res, occupation; tol=tol_cg).δψ + ψ = BlochWaves(basis, [ortho_qr(ψ[ik] + δψ[ik]) for ik = 1:Nk]) - ρ_next = compute_density(basis, ψ, occupation) - energies, H = energy_hamiltonian(basis, ψ, occupation; ρ=ρ_next) + ρ_next = compute_density(ψ, occupation) + energies, H = energy_hamiltonian(ψ, occupation; ρ=ρ_next) info = (; ham=H, basis, converged, stage=:iterate, ρin=ρ, ρout=ρ_next, n_iter, energies, algorithm="Newton") callback(info) @@ -131,11 +134,12 @@ function newton(basis::PlaneWaveBasis{T}, ψ0; # Rayleigh-Ritz eigenvalues = [] + ψ_matrices = blochwaves_as_matrices(ψ) for ik = 1:Nk Hψk = H.blocks[ik] * ψ[ik] - F = eigen(Hermitian(ψ[ik]'Hψk)) + F = eigen(Hermitian(ψ_matrices[ik]'blochwave_as_matrix(Hψk))) push!(eigenvalues, F.values) - ψ[ik] .= ψ[ik] * F.vectors + ψ_matrices[ik] .= ψ_matrices[ik] * F.vectors end εF = nothing # does not necessarily make sense here, as the diff --git a/src/scf/potential_mixing.jl b/src/scf/potential_mixing.jl index 2892a43590..a9e3f4983b 100644 --- a/src/scf/potential_mixing.jl +++ b/src/scf/potential_mixing.jl @@ -229,7 +229,7 @@ trial_damping(damping::FixedDamping, args...) = damping.α fermialg::AbstractFermiAlgorithm=default_fermialg(basis.model), ρ=guess_density(basis), V=nothing, - ψ=nothing, + ψ::BlochWaves=BlochWaves(basis), tol=1e-6, maxiter=100, eigensolver=lobpcg_hyper, @@ -253,15 +253,16 @@ trial_damping(damping::FixedDamping, args...) = damping.α end # Initial guess for V (if none given) - ham = energy_hamiltonian(basis, nothing, nothing; ρ).ham + ham = energy_hamiltonian(BlochWaves(basis), nothing; ρ).ham isnothing(V) && (V = total_local_potential(ham)) - function EVρ(Vin; diagtol=tol / 10, ψ=nothing, eigenvalues=nothing, occupation=nothing) + function EVρ(Vin; diagtol=tol / 10, ψ=BlochWaves(ham.basis), eigenvalues=nothing, + occupation=nothing) ham_V = hamiltonian_with_total_potential(ham, Vin) res_V = next_density(ham_V, nbandsalg, fermialg; eigensolver, ψ, eigenvalues, occupation, miniter=diag_miniter, tol=diagtol) - new_E, new_ham = energy_hamiltonian(basis, res_V.ψ, res_V.occupation; + new_E, new_ham = energy_hamiltonian(res_V.ψ, res_V.occupation; ρ=res_V.ρout, eigenvalues=res_V.eigenvalues, εF=res_V.εF) (; basis, ham=new_ham, energies=new_E, diff --git a/src/scf/scf_callbacks.jl b/src/scf/scf_callbacks.jl index 68f3db56ad..c18907f6c4 100644 --- a/src/scf/scf_callbacks.jl +++ b/src/scf/scf_callbacks.jl @@ -142,7 +142,7 @@ Flag convergence on the change in cartesian force between two iterations. function ScfConvergenceForce(tolerance) previous_force = nothing function is_converged(info) - force = compute_forces_cart(info.basis, info.ψ, info.occupation; ρ=info.ρout) + force = compute_forces_cart(info.ψ, info.occupation; ρ=info.ρout) error = isnothing(previous_force) ? NaN : norm(previous_force - force) previous_force = force error < tolerance diff --git a/src/scf/scfres.jl b/src/scf/scfres.jl index b9af1725f5..5d3d74491d 100644 --- a/src/scf/scfres.jl +++ b/src/scf/scfres.jl @@ -5,7 +5,11 @@ function gather_kpts_scfres(scfres::NamedTuple) kpt_properties = (:ψ, :occupation, :eigenvalues) scfdict = Dict{Symbol, Any}() for symbol in kpt_properties - scfdict[symbol] = gather_kpts(getproperty(scfres, symbol), scfres.basis) + if symbol == :ψ + scfdict[symbol] = gather_kpts(denest(getproperty(scfres, symbol)), scfres.basis) + else + scfdict[symbol] = gather_kpts(getproperty(scfres, symbol), scfres.basis) + end end scfdict[:basis] = gather_kpts(scfres.basis) diff --git a/src/scf/self_consistent_field.jl b/src/scf/self_consistent_field.jl index 91654ed345..a72ff52361 100644 --- a/src/scf/self_consistent_field.jl +++ b/src/scf/self_consistent_field.jl @@ -13,8 +13,9 @@ data structure to determine and adjust the number of bands to be computed. function next_density(ham::Hamiltonian, nbandsalg::NbandsAlgorithm=AdaptiveBands(ham.basis.model), fermialg::AbstractFermiAlgorithm=default_fermialg(ham.basis.model); - eigensolver=lobpcg_hyper, ψ=nothing, eigenvalues=nothing, - occupation=nothing, kwargs...) + eigensolver=lobpcg_hyper, ψ=BlochWaves(ham.basis), + eigenvalues=nothing, occupation=nothing, kwargs...) + @assert ham.basis == ψ.basis n_bands_converge, n_bands_compute = determine_n_bands(nbandsalg, occupation, eigenvalues, ψ) @@ -22,8 +23,8 @@ function next_density(ham::Hamiltonian, increased_n_bands = true else @assert length(ψ) == length(ham.basis.kpoints) - n_bands_compute = max(n_bands_compute, maximum(ψk -> size(ψk, 2), ψ)) - increased_n_bands = n_bands_compute > size(ψ[1], 2) + n_bands_compute = max(n_bands_compute, maximum(ψk -> size(ψk, 3), ψ)) + increased_n_bands = any([n_bands_compute > size(ψk, 3) for ψk in ψ]) end # TODO Synchronize since right now it is assumed that the same number of bands are @@ -49,9 +50,10 @@ function next_density(ham::Hamiltonian, "`nbandsalg=AdaptiveBands(model; n_bands_converge=$(n_bands_converge + 3)`)") end - ρout = compute_density(ham.basis, eigres.X, occupation; nbandsalg.occupation_threshold) - (; ψ=eigres.X, eigenvalues=eigres.λ, occupation, εF, ρout, diagonalization=eigres, - n_bands_converge, nbandsalg.occupation_threshold) + ψ = BlochWaves(ham.basis, eigres.X) + ρout = compute_density(ψ, occupation; nbandsalg.occupation_threshold) + (; ψ, eigenvalues=eigres.λ, occupation, εF, ρout, + diagonalization=eigres, n_bands_converge, nbandsalg.occupation_threshold) end @@ -86,7 +88,7 @@ Overview of parameters: @timing function self_consistent_field( basis::PlaneWaveBasis{T}; ρ=guess_density(basis), - ψ=nothing, + ψ::BlochWaves=BlochWaves(basis), tol=1e-6, is_converged=ScfConvergenceDensity(tol), maxiter=100, @@ -101,10 +103,8 @@ Overview of parameters: compute_consistent_energies=true, response=ResponseOptions(), # Dummy here, only for AD ) where {T} + @assert basis == ψ.basis # All these variables will get updated by fixpoint_map - if !isnothing(ψ) - @assert length(ψ) == length(basis.kpoints) - end occupation = nothing eigenvalues = nothing ρout = ρ @@ -123,7 +123,7 @@ Overview of parameters: # Note that ρin is not the density of ψ, and the eigenvalues # are not the self-consistent ones, which makes this energy non-variational - energies, ham = energy_hamiltonian(basis, ψ, occupation; ρ=ρin, eigenvalues, εF) + energies, ham = energy_hamiltonian(ψ, occupation; ρ=ρin, eigenvalues, εF) # Diagonalize `ham` to get the new state nextstate = next_density(ham, nbandsalg, fermialg; eigensolver, ψ, eigenvalues, @@ -137,8 +137,7 @@ Overview of parameters: # Compute the energy of the new state if compute_consistent_energies - energies = energy_hamiltonian(basis, ψ, occupation; - ρ=ρout, eigenvalues, εF).energies + energies = energy_hamiltonian(ψ, occupation; ρ=ρout, eigenvalues, εF).energies end info = merge(info, (; energies)) @@ -161,7 +160,7 @@ Overview of parameters: # We do not use the return value of solver but rather the one that got updated by fixpoint_map # ψ is consistent with ρout, so we return that. We also perform a last energy computation # to return a correct variational energy - energies, ham = energy_hamiltonian(basis, ψ, occupation; ρ=ρout, eigenvalues, εF) + energies, ham = energy_hamiltonian(ψ, occupation; ρ=ρout, eigenvalues, εF) # Measure for the accuracy of the SCF # TODO probably should be tracked all the way ... diff --git a/src/supercell.jl b/src/supercell.jl index f2ff6f39f3..4d2bac2bc0 100644 --- a/src/supercell.jl +++ b/src/supercell.jl @@ -72,8 +72,7 @@ function cell_to_supercell(ψ, basis::PlaneWaveBasis{T}, # Ensure that the basis is unfolded. length(basis.kgrid) != length(basis.kpoints) && error("basis must be unfolded") - num_kpG = sum(size.(ψ, 1)) - num_bands = size(ψ[1], 2) + num_kpG = sum([size(ψk, 2) for ψk in ψ]) # Maps k+G vector of initial basis to a G vector of basis_supercell function cell_supercell_mapping(kpt) @@ -84,12 +83,12 @@ function cell_to_supercell(ψ, basis::PlaneWaveBasis{T}, # Transfer all ψ[k] independantly and return the hcat of all blocs ψ_out_blocs = [] for (ik, kpt) in enumerate(basis.kpoints) - ψk_supercell = zeros(complex(T), num_kpG, num_bands) - ψk_supercell[cell_supercell_mapping(kpt), :] .= ψ[ik] + ψk_supercell = zeros(complex(T), basis.model.n_components, num_kpG, size(ψ[ik], 3)) + ψk_supercell[:, cell_supercell_mapping(kpt), :] .= ψ[ik] push!(ψ_out_blocs, ψk_supercell) end # Note that each column is normalize since each ψ[ik][:,n] is. - hcat(ψ_out_blocs...) + BlochWaves(basis_supercell, [cat(ψ_out_blocs...; dims=3)]) end @doc raw""" @@ -110,15 +109,14 @@ function cell_to_supercell(scfres::NamedTuple) # Compute supercell basis, ψ, occupations and ρ basis_supercell = cell_to_supercell(basis) - ψ_supercell = [cell_to_supercell(ψ, basis, basis_supercell)] + ψ_supercell = cell_to_supercell(ψ, basis, basis_supercell) eigs_supercell = [vcat(scfres_unfold.eigenvalues...)] occ_supercell = compute_occupation(basis_supercell, eigs_supercell, scfres.εF).occupation - ρ_supercell = compute_density(basis_supercell, ψ_supercell, occ_supercell; - scfres.occupation_threshold) + ρ_supercell = compute_density(ψ_supercell, occ_supercell; scfres.occupation_threshold) # Supercell Energies - Eham_supercell = energy_hamiltonian(basis_supercell, ψ_supercell, occ_supercell; - ρ=ρ_supercell, eigenvalues=eigs_supercell, scfres.εF) + Eham_supercell = energy_hamiltonian(ψ_supercell, occ_supercell; ρ=ρ_supercell, + eigenvalues=eigs_supercell, scfres.εF) merge(scfres, (; ham=Eham_supercell.ham, basis=basis_supercell, ψ=ψ_supercell, energies=Eham_supercell.energies, ρ=ρ_supercell, diff --git a/src/symmetry.jl b/src/symmetry.jl index 6d01d9ae86..64e6a24a15 100644 --- a/src/symmetry.jl +++ b/src/symmetry.jl @@ -216,7 +216,7 @@ Apply a symmetry operation to eigenvectors `ψk` at a given `kpoint` to obtain a equivalent point in [-0.5, 0.5)^3 and associated eigenvectors (expressed in the basis of the new ``k``-point). """ -function apply_symop(symop::SymOp, basis, kpoint, ψk::AbstractVecOrMat) +function apply_symop(symop::SymOp, basis, kpoint, ψk::AbstractArray) S, τ = symop.S, symop.τ isone(symop) && return kpoint, ψk @@ -248,11 +248,13 @@ function apply_symop(symop::SymOp, basis, kpoint, ψk::AbstractVecOrMat) invS = Mat3{Int}(inv(S)) Gs_full = [G + kshift for G in G_vectors(basis, Skpoint)] ψSk = zero(ψk) - for iband = 1:size(ψk, 2) + for iband = 1:size(ψk, 3) for (ig, G_full) in enumerate(Gs_full) igired = index_G_vectors(basis, kpoint, invS * G_full) @assert igired !== nothing - ψSk[ig, iband] = cis2pi(-dot(G_full, τ)) * ψk[igired, iband] + for σ = 1:basis.model.n_components + ψSk[σ, ig, iband] = cis2pi(-dot(G_full, τ)) * ψk[σ, igired, iband] + end end end @@ -413,6 +415,8 @@ function unfold_array_(basis_irred, basis_unfolded, data, is_ψ) error("Brillouin zone symmetry unfolding not supported with MPI yet") end data_unfolded = similar(data, length(basis_unfolded.kpoints)) + n_components = basis_unfolded.model.n_components + @assert n_components == basis_irred.model.n_components for ik_unfolded = 1:length(basis_unfolded.kpoints) kpt_unfolded = basis_unfolded.kpoints[ik_unfolded] ik_irred, symop = unfold_mapping(basis_irred, kpt_unfolded) @@ -420,15 +424,19 @@ function unfold_array_(basis_irred, basis_unfolded, data, is_ψ) # transform ψ_k from data into ψ_Sk in data_unfolded kunfold_coord = kpt_unfolded.coordinate @assert normalize_kpoint_coordinate(kunfold_coord) ≈ kunfold_coord - _, ψSk = apply_symop(symop, basis_irred, - basis_irred.kpoints[ik_irred], data[ik_irred]) + ψSk = apply_symop(symop, basis_irred, basis_irred.kpoints[ik_irred], + data[ik_irred])[2] data_unfolded[ik_unfolded] = ψSk else # simple copy data_unfolded[ik_unfolded] = data[ik_irred] end end - data_unfolded + if is_ψ + BlochWaves(basis_unfolded, denest(basis_unfolded, data_unfolded)) + else + data_unfolded + end end function unfold_bz(scfres) @@ -436,7 +444,7 @@ function unfold_bz(scfres) ψ = unfold_array_(scfres.basis, basis_unfolded, scfres.ψ, true) eigenvalues = unfold_array_(scfres.basis, basis_unfolded, scfres.eigenvalues, false) occupation = unfold_array_(scfres.basis, basis_unfolded, scfres.occupation, false) - energies, ham = energy_hamiltonian(basis_unfolded, ψ, occupation; + energies, ham = energy_hamiltonian(ψ, occupation; scfres.ρ, eigenvalues, scfres.εF) @assert energies.total ≈ scfres.energies.total new_scfres = (; basis=basis_unfolded, ψ, ham, eigenvalues, occupation) diff --git a/src/terms/Hamiltonian.jl b/src/terms/Hamiltonian.jl index 37f8247c6b..17a18264bd 100644 --- a/src/terms/Hamiltonian.jl +++ b/src/terms/Hamiltonian.jl @@ -88,43 +88,57 @@ function LinearAlgebra.mul!(Hψ, H::Hamiltonian, ψ) Hψ end # need `deepcopy` here to copy the elements of the array of arrays ψ (not just pointers) -Base.:*(H::Hamiltonian, ψ) = mul!(deepcopy(ψ), H, ψ) +Base.:*(H::Hamiltonian, ψ::AbstractArray) = mul!(deepcopy(ψ), H, ψ) + +# Serves as a wrapper around the specialized multiplication functions for LOBPCG as it +# expects to work on matrices. +function LinearAlgebra.mul!(Hψ::AbstractVecOrMat, H::HamiltonianBlock, ψ::AbstractVecOrMat) + n_Gk = length(H.kpoint.G_vectors) + n_bands = size(Hψ, 2) + n_components = H.basis.model.n_components + Hψ_r = reshape(Hψ, n_components, n_Gk, n_bands) + ψ_r = reshape(ψ, n_components, n_Gk, n_bands) + mul!(Hψ_r, H, ψ_r) + Hψ +end # 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ψ::AbstractArray, +@views @timing "Hamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray3, H::GenericHamiltonianBlock, - ψ::AbstractArray) + ψ::AbstractArray3) T = eltype(H.basis) - n_bands = size(ψ, 2) - Hψ_fourier = similar(Hψ[:, 1]) + n_bands = size(ψ, 3) + Hψ_fourier = similar(Hψ[1, :, 1]) ψ_real = similar(ψ, complex(T), H.basis.fft_size...) Hψ_real = similar(Hψ, complex(T), H.basis.fft_size...) # take ψi, IFFT it to ψ_real, apply each term to Hψ_fourier and Hψ_real, and add it to Hψ for iband = 1:n_bands Hψ_real .= 0 Hψ_fourier .= 0 - ifft!(ψ_real, H.basis, H.kpoint, ψ[:, iband]) - for op in H.optimized_operators - @timing "$(nameof(typeof(op)))" begin - apply!((; fourier=Hψ_fourier, real=Hψ_real), - op, - (; fourier=ψ[:, iband], real=ψ_real)) + for σ = 1:H.basis.model.n_components + ifft!(ψ_real, H.basis, H.kpoint, ψ[σ, :, iband]) + for op in H.optimized_operators + @timing "$(nameof(typeof(op)))" begin + apply!((fourier=Hψ_fourier, real=Hψ_real), + op, + (fourier=ψ[σ, :, iband], real=ψ_real)) + end end + Hψ[σ, :, iband] .= Hψ_fourier + fft!(Hψ_fourier, H.basis, H.kpoint, Hψ_real) + Hψ[σ, :, iband] .+= Hψ_fourier end - Hψ[:, iband] .= Hψ_fourier - fft!(Hψ_fourier, H.basis, H.kpoint, Hψ_real) - Hψ[:, iband] .+= Hψ_fourier end Hψ end # Fast version, specialized on DFT models. Minimizes the number of FFTs and allocations -@views @timing "DftHamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray, +@views @timing "DftHamiltonian multiplication" function LinearAlgebra.mul!(Hψ::AbstractArray3, H::DftHamiltonianBlock, - ψ::AbstractArray) - n_bands = size(ψ, 2) + ψ::AbstractArray3) + n_bands = size(ψ, 3) iszero(n_bands) && return Hψ # Nothing to do if ψ empty have_divAgrad = !isnothing(H.divAgrad_op) @@ -139,17 +153,20 @@ end ψ_real = H.scratch.ψ_reals[ichunk] @timeit to "local+kinetic" begin - ifft!(ψ_real, H.basis, H.kpoint, ψ[:, iband]; normalize=false) - ψ_real .*= potential - fft!(Hψ[:, iband], H.basis, H.kpoint, ψ_real; normalize=false) # overwrites ψ_real - Hψ[:, iband] .+= H.fourier_op.multiplier .* ψ[:, iband] + for σ = 1:H.basis.model.n_components + ifft!(ψ_real, H.basis, H.kpoint, ψ[σ, :, iband]; normalize=false) + ψ_real .*= potential + fft!(Hψ[σ, :, iband], H.basis, H.kpoint, ψ_real; normalize=false) # overwrites ψ_real + Hψ[σ, :, iband] .+= H.fourier_op.multiplier .* ψ[σ, :, iband] + end end if have_divAgrad + @assert H.basis.model.n_components == 1 @timeit to "divAgrad" begin - apply!((; fourier=Hψ[:, iband], real=nothing), + apply!((; fourier=Hψ[1, :, iband], real=nothing), H.divAgrad_op, - (; fourier=ψ[:, iband], real=nothing), + (; fourier=ψ[1, :, iband], real=nothing), ψ_real) # ψ_real used as scratch end end @@ -174,13 +191,13 @@ end Hψ end - # Get energies and Hamiltonian # kwargs is additional info that might be useful for the energy terms to precompute # (eg the density ρ) -@timing function energy_hamiltonian(basis::PlaneWaveBasis, ψ, occupation; kwargs...) +@timing function energy_hamiltonian(ψ::BlochWaves, occupation; kwargs...) + basis = ψ.basis # it: index into terms, ik: index into kpoints - @timing "ene_ops" ene_ops_arr = [ene_ops(term, basis, ψ, occupation; kwargs...) + @timing "ene_ops" ene_ops_arr = [ene_ops(term, ψ, occupation; kwargs...) for term in basis.terms] energies = [eh.E for eh in ene_ops_arr] operators = [eh.ops for eh in ene_ops_arr] # operators[it][ik] @@ -205,8 +222,8 @@ end energies = Energies(basis.model.term_types, energies) (; energies, ham) end -function Hamiltonian(basis::PlaneWaveBasis; ψ=nothing, occupation=nothing, kwargs...) - energy_hamiltonian(basis, ψ, occupation; kwargs...).ham +function Hamiltonian(basis::PlaneWaveBasis; ψ=BlochWaves(basis), occupation=nothing, kwargs...) + energy_hamiltonian(ψ, occupation; kwargs...).ham end """ diff --git a/src/terms/anyonic.jl b/src/terms/anyonic.jl index 8a4000936a..0f04cba2ed 100644 --- a/src/terms/anyonic.jl +++ b/src/terms/anyonic.jl @@ -99,8 +99,10 @@ function TermAnyonic(basis::PlaneWaveBasis{T}, hbar, β) where {T} TermAnyonic(hbar, β, ρref, Aref) end -function ene_ops(term::TermAnyonic, basis::PlaneWaveBasis{T}, ψ, occupation; +function ene_ops(term::TermAnyonic, ψ::BlochWaves{T}, occupation; ρ, kwargs...) where {T} + basis = ψ.basis + @assert basis.model.n_components == 1 hbar = term.hbar β = term.β @assert ψ !== nothing # the hamiltonian depends explicitly on ψ @@ -132,7 +134,7 @@ function ene_ops(term::TermAnyonic, basis::PlaneWaveBasis{T}, ψ, occupation; β^2 .* (abs2.(Areal[1]) .+ abs2.(Areal[2])))] # Now compute effective local potential - 2β x^⟂/|x|² ∗ (βAρ + J) - J = compute_current(basis, ψ, occupation) + J = compute_current(basis, [ψk[1, :, :] for ψk in ψ], occupation) eff_current = [hbar .* J[α] .+ β .* ρ .* Areal[α] for α = 1:2] eff_current_fourier = [fft(basis, eff_current[α]) for α = 1:2] @@ -153,11 +155,13 @@ function ene_ops(term::TermAnyonic, basis::PlaneWaveBasis{T}, ψ, occupation; ops_ham = [ops_energy..., RealSpaceMultiplication(basis, basis.kpoints[1], eff_pot_real)] E = zero(T) - for iband = 1:size(ψ[1], 2) - ψnk = @views ψ[1][:, iband] - # TODO optimize this - for op in ops_energy - E += occupation[1][iband] * real(dot(ψnk, op * ψnk)) + for iband = 1:size(ψ[1], 3) + for σ = 1:basis.model.n_components + ψσnk = @views ψ[1][σ, :, iband] + # TODO optimize this + for op in ops_energy + E += occupation[1][iband] * real(dot(ψσnk, op * ψσnk)) + end end end diff --git a/src/terms/entropy.jl b/src/terms/entropy.jl index b4870d545c..b07f7d508e 100644 --- a/src/terms/entropy.jl +++ b/src/terms/entropy.jl @@ -8,8 +8,8 @@ struct Entropy end (::Entropy)(basis) = TermEntropy() struct TermEntropy <: Term end -function ene_ops(term::TermEntropy, basis::PlaneWaveBasis{T}, ψ, occupation; - kwargs...) where {T} +function ene_ops(term::TermEntropy, ψ::BlochWaves{T}, occupation; kwargs...) where {T} + basis = ψ.basis ops = [NoopOperator(basis, kpt) for kpt in basis.kpoints] smearing = basis.model.smearing temperature = basis.model.temperature @@ -25,8 +25,8 @@ function ene_ops(term::TermEntropy, basis::PlaneWaveBasis{T}, ψ, occupation; eigenvalues = kwargs[:eigenvalues] E = zero(T) - for (ik, k) in enumerate(basis.kpoints) - for iband = 1:size(ψ[ik], 2) + for (ik, ψk) in enumerate(ψ) + for iband = 1:size(ψk, 3) E -= (temperature * basis.kweights[ik] * filled_occupation(basis.model) diff --git a/src/terms/ewald.jl b/src/terms/ewald.jl index 3d23d166e8..fc052bf406 100644 --- a/src/terms/ewald.jl +++ b/src/terms/ewald.jl @@ -25,10 +25,10 @@ end TermEwald(energy, forces, η) end -function ene_ops(term::TermEwald, basis::PlaneWaveBasis, ψ, occupation; kwargs...) - (; E=term.energy, ops=[NoopOperator(basis, kpt) for kpt in basis.kpoints]) +function ene_ops(term::TermEwald, ψ::BlochWaves, occupation; kwargs...) + (; E=term.energy, ops=[NoopOperator(ψ.basis, kpt) for kpt in ψ.basis.kpoints]) end -compute_forces(term::TermEwald, ::PlaneWaveBasis, ψ, occupation; kwargs...) = term.forces +compute_forces(term::TermEwald, ::BlochWaves, occupation; kwargs...) = term.forces """ Standard computation of energy and forces. diff --git a/src/terms/hartree.jl b/src/terms/hartree.jl index 0c4cd78c12..f193a02d5c 100644 --- a/src/terms/hartree.jl +++ b/src/terms/hartree.jl @@ -39,8 +39,9 @@ function TermHartree(basis::PlaneWaveBasis{T}, scaling_factor) where {T} TermHartree(T(scaling_factor), T(scaling_factor) .* poisson_green_coeffs) end -@timing "ene_ops: hartree" function ene_ops(term::TermHartree, basis::PlaneWaveBasis{T}, - ψ, occupation; ρ, kwargs...) where {T} +@timing "ene_ops: hartree" function ene_ops(term::TermHartree, ψ::BlochWaves{T}, occupation; + ρ, kwargs...) where {T} + basis = ψ.basis ρtot_fourier = fft(basis, total_density(ρ)) pot_fourier = term.poisson_green_coeffs .* ρtot_fourier pot_real = irfft(basis, pot_fourier) diff --git a/src/terms/kinetic.jl b/src/terms/kinetic.jl index 6ceb2a7f0d..1970c13fcd 100644 --- a/src/terms/kinetic.jl +++ b/src/terms/kinetic.jl @@ -34,8 +34,9 @@ function kinetic_energy(kin::Kinetic, Ecut, q) kinetic_energy(kin.blowup, kin.scaling_factor, Ecut, q) end -@timing "ene_ops: kinetic" function ene_ops(term::TermKinetic, basis::PlaneWaveBasis{T}, - ψ, occupation; kwargs...) where {T} +@timing "ene_ops: kinetic" function ene_ops(term::TermKinetic, ψ::BlochWaves{T}, occupation; + kwargs...) where {T} + basis = ψ.basis ops = [FourierMultiplication(basis, kpoint, term.kinetic_energies[ik]) for (ik, kpoint) in enumerate(basis.kpoints)] if isnothing(ψ) || isnothing(occupation) @@ -45,10 +46,10 @@ end E = zero(T) for (ik, ψk) in enumerate(ψ) - for iband = 1:size(ψk, 2) - ψnk = @views ψk[:, iband] + for iband = 1:size(ψk, 3), σ = 1:size(ψk, 1) + ψσkn = @views ψk[σ, :, iband] E += (basis.kweights[ik] * occupation[ik][iband] - * real(dot(ψnk, Diagonal(term.kinetic_energies[ik]), ψnk))) + * real(dot(ψσkn, Diagonal(term.kinetic_energies[ik]), ψσkn))) end end E = mpi_sum(E, basis.comm_kpts) diff --git a/src/terms/local.jl b/src/terms/local.jl index 3ee08cb40b..e2378e06fe 100644 --- a/src/terms/local.jl +++ b/src/terms/local.jl @@ -6,10 +6,10 @@ # two spin components. abstract type TermLocalPotential <: Term end -@timing "ene_ops: local" function ene_ops(term::TermLocalPotential, - basis::PlaneWaveBasis{T}, ψ, occupation; - kwargs...) where {T} +@timing "ene_ops: local" function ene_ops(term::TermLocalPotential, ψ::BlochWaves{T}, + occupation; kwargs...) where {T} potview(data, spin) = ndims(data) == 4 ? (@view data[:, :, :, spin]) : data + basis = ψ.basis ops = [RealSpaceMultiplication(basis, kpt, potview(term.potential_values, kpt.spin)) for kpt in basis.kpoints] if :ρ in keys(kwargs) @@ -123,18 +123,19 @@ function (::AtomicLocal)(basis::PlaneWaveBasis{T}) where {T} TermAtomicLocal(pot_real) end -@timing "forces: local" function compute_forces(::TermAtomicLocal, basis::PlaneWaveBasis{TT}, - ψ, occupation; ρ, kwargs...) where {TT} - T = promote_type(TT, real(eltype(ψ[1]))) +@timing "forces: local" function compute_forces(::TermAtomicLocal, ψ::BlochWaves{T, Tψ}, + occupation; ρ, kwargs...) where {T, Tψ} + TT = promote_type(T, real(Tψ)) + basis = ψ.basis model = basis.model ρ_fourier = fft(basis, total_density(ρ)) # energy = sum of form_factor(G) * struct_factor(G) * rho(G) # where struct_factor(G) = e^{-i G·r} - forces = [zero(Vec3{T}) for _ = 1:length(model.positions)] + forces = [zero(Vec3{TT}) for _ = 1:length(model.positions)] for group in model.atom_groups element = model.atoms[first(group)] - form_factors = [Complex{T}(local_potential_fourier(element, norm(G))) + form_factors = [Complex{TT}(local_potential_fourier(element, norm(G))) for G in G_vectors_cart(basis)] for idx in group r = model.positions[idx] diff --git a/src/terms/local_nonlinearity.jl b/src/terms/local_nonlinearity.jl index 34998827f1..7c4e665cd8 100644 --- a/src/terms/local_nonlinearity.jl +++ b/src/terms/local_nonlinearity.jl @@ -9,8 +9,9 @@ struct TermLocalNonlinearity{TF} <: TermNonlinear end (L::LocalNonlinearity)(::AbstractBasis) = TermLocalNonlinearity(L.f) -function ene_ops(term::TermLocalNonlinearity, basis::PlaneWaveBasis{T}, ψ, occupation; +function ene_ops(term::TermLocalNonlinearity, ψ::BlochWaves{T}, occupation; ρ, kwargs...) where {T} + basis = ψ.basis fp(ρ) = ForwardDiff.derivative(term.f, ρ) E = sum(fρ -> convert_dual(T, fρ), term.f.(ρ)) * basis.dvol potential = convert_dual.(T, fp.(ρ)) diff --git a/src/terms/magnetic.jl b/src/terms/magnetic.jl index 4b92ed0dc8..4e5e9d7b65 100644 --- a/src/terms/magnetic.jl +++ b/src/terms/magnetic.jl @@ -30,8 +30,8 @@ function TermMagnetic(basis::PlaneWaveBasis{T}, Afunction::Function) where {T} TermMagnetic(Apotential) end -function ene_ops(term::TermMagnetic, basis::PlaneWaveBasis{T}, ψ, occupation; - kwargs...) where {T} +function ene_ops(term::TermMagnetic, ψ::BlochWaves{T}, occupation; kwargs...) where {T} + basis = ψ.basis ops = [MagneticFieldOperator(basis, kpoint, term.Apotential) for (ik, kpoint) in enumerate(basis.kpoints)] if isnothing(ψ) || isnothing(occupation) @@ -39,11 +39,13 @@ function ene_ops(term::TermMagnetic, basis::PlaneWaveBasis{T}, ψ, occupation; end E = zero(T) - for (ik, k) in enumerate(basis.kpoints) - for iband = 1:size(ψ[1], 2) - ψnk = @views ψ[ik][:, iband] - # TODO optimize this - E += basis.kweights[ik] * occupation[ik][iband] * real(dot(ψnk, ops[ik] * ψnk)) + for (ik, ψk) in enumerate(ψ) + for iband = 1:size(ψk, 3) + for σ = 1:basis.model.n_components + ψσnk = @views ψ[ik][σ, :, iband] + # TODO optimize this + E += basis.kweights[ik] * occupation[ik][iband] * real(dot(ψσnk, ops[ik] * ψσnk)) + end end end E = mpi_sum(E, basis.comm_kpts) diff --git a/src/terms/nonlocal.jl b/src/terms/nonlocal.jl index e7dfedf985..7a31f83421 100644 --- a/src/terms/nonlocal.jl +++ b/src/terms/nonlocal.jl @@ -25,28 +25,30 @@ struct TermAtomicNonlocal <: Term ops::Vector{NonlocalOperator} end -@timing "ene_ops: nonlocal" function ene_ops(term::TermAtomicNonlocal, - basis::PlaneWaveBasis{T}, - ψ, occupation; kwargs...) where {T} +@timing "ene_ops: nonlocal" function ene_ops(term::TermAtomicNonlocal, ψ::BlochWaves{T}, + occupation; kwargs...) where {T} + basis = ψ.basis if isnothing(ψ) || isnothing(occupation) return (; E=T(Inf), term.ops) end E = zero(T) for (ik, ψk) in enumerate(ψ) - Pψk = term.ops[ik].P' * ψk # nproj x nband - band_enes = dropdims(sum(real.(conj.(Pψk) .* (term.ops[ik].D * Pψk)), dims=1), dims=1) - E += basis.kweights[ik] * sum(band_enes .* occupation[ik]) + for σ = 1:basis.model.n_components + Pψk = term.ops[ik].P' * ψk[σ, :, :] # nproj x nband + band_enes = dropdims(sum(real.(conj.(Pψk) .* (term.ops[ik].D * Pψk)), dims=1), dims=1) + E += basis.kweights[ik] * sum(band_enes .* occupation[ik]) + end end E = mpi_sum(E, basis.comm_kpts) (; E, term.ops) end -@timing "forces: nonlocal" function compute_forces(::TermAtomicNonlocal, - basis::PlaneWaveBasis{TT}, - ψ, occupation; kwargs...) where {TT} - T = promote_type(TT, real(eltype(ψ[1]))) +@timing "forces: nonlocal" function compute_forces(::TermAtomicNonlocal, ψ::BlochWaves{T, Tψ}, + occupation; kwargs...) where {T, Tψ} + basis = ψ.basis + TT = promote_type(T, real(Tψ)) model = basis.model unit_cell_volume = model.unit_cell_volume psp_groups = [group for group in model.atom_groups @@ -56,11 +58,11 @@ end isempty(psp_groups) && return nothing # energy terms are of the form , where P(G) = form_factor(G) * structure_factor(G) - forces = [zero(Vec3{T}) for _ = 1:length(model.positions)] + forces = [zero(Vec3{TT}) for _ = 1:length(model.positions)] for group in psp_groups element = model.atoms[first(group)] - C = build_projection_coefficients_(T, element.psp) + C = build_projection_coefficients_(TT, element.psp) for (ik, kpt) in enumerate(basis.kpoints) # we compute the forces from the irreductible BZ; they are symmetrized later qs = Gplusk_vectors(basis, kpt) @@ -72,12 +74,14 @@ end P = structure_factors .* form_factors ./ sqrt(unit_cell_volume) forces[idx] += map(1:3) do α - dPdR = [-2T(π)*im*q[α] for q in qs] .* P - ψk = ψ[ik] - dHψk = P * (C * (dPdR' * ψk)) - -sum(occupation[ik][iband] * basis.kweights[ik] * - 2real(dot(ψk[:, iband], dHψk[:, iband])) - for iband=1:size(ψk, 2)) + dPdR = [-2TT(π)*im*q[α] for q in qs] .* P + mapreduce(+, 1:model.n_components) do σ + ψkσ = ψ[ik][σ, :, :] + dHψkσ = P * (C * (dPdR' * ψkσ)) + -sum(occupation[ik][iband] * basis.kweights[ik] * + 2real(dot(ψkσ[:, iband], dHψkσ[:, iband])) + for iband=1:size(ψkσ, 2)) + end end # α end # r end # kpt diff --git a/src/terms/operators.jl b/src/terms/operators.jl index 1602e0db26..78ffcdf900 100644 --- a/src/terms/operators.jl +++ b/src/terms/operators.jl @@ -24,13 +24,23 @@ function LinearAlgebra.mul!(Hψ::AbstractVector, op::RealFourierOperator, ψ::Ab Hψ .= Hψ_fourier .+ fft(op.basis, op.kpoint, Hψ_real) Hψ end -function LinearAlgebra.mul!(Hψ::AbstractMatrix, op::RealFourierOperator, ψ::AbstractMatrix) - @views for i = 1:size(ψ, 2) - mul!(Hψ[:, i], op, ψ[:, i]) +# T@D@: remove this hack +function LinearAlgebra.mul!(Hψ::AbstractArray3, op::RealFourierOperator, ψ::AbstractArray3) + @views for i = 1:size(ψ, 3) + @views for σ = 1:size(ψ, 1) + mul!(Hψ[σ, :, i], op, ψ[σ, :, i]) + end end Hψ end Base.:*(op::RealFourierOperator, ψ) = mul!(similar(ψ), op, ψ) +# Default transformation: from two components tensors to diagonal matrices for each +# reciprocal vector. +function Matrix(op::RealFourierOperator) + n_Gk = length(G_vectors(op.basis, op.kpoint)) + n_components = op.basis.model.n_components + reshape(Array(op), n_components*n_Gk, n_components*n_Gk) +end """ Noop operation: don't do anything. @@ -41,9 +51,10 @@ struct NoopOperator{T <: Real} <: RealFourierOperator kpoint::Kpoint{T} end apply!(Hψ, op::NoopOperator, ψ) = nothing -function Matrix(op::NoopOperator) +function Array(op::NoopOperator) n_Gk = length(G_vectors(op.basis, op.kpoint)) - zeros_like(op.basis.G_vectors, eltype(op.basis), n_Gk, n_Gk) + n_components = op.basis.model.n_components + zeros_like(op.basis.G_vectors, eltype(op.basis), n_components, n_Gk, n_components, n_Gk) end """ @@ -58,20 +69,23 @@ end function apply!(Hψ, op::RealSpaceMultiplication, ψ) Hψ.real .+= op.potential .* ψ.real end -function Matrix(op::RealSpaceMultiplication) +function Array(op::RealSpaceMultiplication) # V(G, G') = = 1/sqrt(Ω) pot_fourier = fft(op.basis, op.potential) n_G = length(G_vectors(op.basis, op.kpoint)) - H = zeros(complex(eltype(op.basis)), n_G, n_G) - for (j, G′) in enumerate(G_vectors(op.basis, op.kpoint)) - for (i, G) in enumerate(G_vectors(op.basis, op.kpoint)) - # 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") + n_components = op.basis.model.n_components + H = zeros(complex(eltype(op.basis)), n_components, n_G, n_components, n_G) + for σ = 1:n_components + for (j, G′) in enumerate(G_vectors(op.basis, op.kpoint)) + for (i, G) in enumerate(G_vectors(op.basis, op.kpoint)) + # 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) end - H[i, j] = pot_fourier[ind_ΔG] / sqrt(op.basis.model.unit_cell_volume) end end H @@ -89,7 +103,12 @@ end function apply!(Hψ, op::FourierMultiplication, ψ) Hψ.fourier .+= op.multiplier .* ψ.fourier end -Matrix(op::FourierMultiplication) = Array(Diagonal(op.multiplier)) +function Array(op::FourierMultiplication) + n_Gk = length(G_vectors(op.basis, op.kpoint)) + n_components = op.basis.model.n_components + D = Diagonal(op.multiplier) + reshape(kron(D, I(n_components)), n_components, n_Gk, n_components, n_Gk) +end """ Nonlocal operator in Fourier space in Kleinman-Bylander format, @@ -104,9 +123,29 @@ struct NonlocalOperator{T <: Real, PT, DT} <: RealFourierOperator D::DT end function apply!(Hψ, op::NonlocalOperator, ψ) - Hψ.fourier .+= op.P * (op.D * (op.P' * ψ.fourier)) + fix_that_apply!(Hψ.fourier, op, ψ.fourier) + Hψ.fourier +end +# T@D@ self-explainatory +function fix_that_apply!(Hψ, op::NonlocalOperator, ψ::AbstractVecOrMat) + Hψ .+= op.P * (op.D * (op.P' * ψ)) + Hψ +end +function fix_that_apply!(Hψ, op::NonlocalOperator, ψ::AbstractArray3) + for σ = 1:op.basis.model.n_components + Hψ[σ, :, :] .+= op.P * (op.D * (op.P' * ψ[σ, :, :])) + end + Hψ +end +function Array(op::NonlocalOperator) + n_Gk = length(G_vectors(op.basis, op.kpoint)) + n_components = op.basis.model.n_components + H = zeros(complex(eltype(op.basis)), n_components, n_Gk, n_components, n_Gk) + for σ = 1:op.basis.model.n_components + H[σ, :, σ, :] = op.P * op.D * op.P' + end + H end -Matrix(op::NonlocalOperator) = op.P * op.D * op.P' """ Magnetic field operator A⋅(-i∇). @@ -126,7 +165,7 @@ function apply!(Hψ, op::MagneticFieldOperator, ψ) Hψ.real .+= op.Apot[α] .* ∂αψ_real end end -# TODO Implement Matrix(op::MagneticFieldOperator) +# TODO Implement Array(op::MagneticFieldOperator) @doc raw""" Nonlocal "divAgrad" operator ``-½ ∇ ⋅ (A ∇)`` where ``A`` is a scalar field on the @@ -150,7 +189,7 @@ function apply!(Hψ, op::DivAgradOperator, ψ, end end -# TODO Implement Matrix(op::DivAgrad) +# TODO Implement Array(op::DivAgrad) # Optimize RFOs by combining terms that can be combined diff --git a/src/terms/pairwise.jl b/src/terms/pairwise.jl index 55f597459c..8198a64895 100644 --- a/src/terms/pairwise.jl +++ b/src/terms/pairwise.jl @@ -33,10 +33,10 @@ struct TermPairwisePotential{TV, Tparams, T} <:Term forces::Vector{Vec3{T}} end -function ene_ops(term::TermPairwisePotential, basis::PlaneWaveBasis, ψ, occupation; kwargs...) - (; E=term.energy, ops=[NoopOperator(basis, kpt) for kpt in basis.kpoints]) +function ene_ops(term::TermPairwisePotential, ψ::BlochWaves, occupation; kwargs...) + (; E=term.energy, ops=[NoopOperator(ψ.basis, kpt) for kpt in ψ.basis.kpoints]) end -compute_forces(term::TermPairwisePotential, ::PlaneWaveBasis, ψ, occ; kwargs...) = term.forces +compute_forces(term::TermPairwisePotential, ::BlochWaves, occ; kwargs...) = term.forces """ diff --git a/src/terms/psp_correction.jl b/src/terms/psp_correction.jl index 198b0771bc..52ff73005d 100644 --- a/src/terms/psp_correction.jl +++ b/src/terms/psp_correction.jl @@ -15,8 +15,8 @@ function TermPspCorrection(basis::PlaneWaveBasis) TermPspCorrection(energy_psp_correction(model)) end -function ene_ops(term::TermPspCorrection, basis::PlaneWaveBasis, ψ, occupation; kwargs...) - (; E=term.energy, ops=[NoopOperator(basis, kpt) for kpt in basis.kpoints]) +function ene_ops(term::TermPspCorrection, ψ::BlochWaves, occupation; kwargs...) + (; E=term.energy, ops=[NoopOperator(ψ.basis, kpt) for kpt in ψ.basis.kpoints]) end """ diff --git a/src/terms/terms.jl b/src/terms/terms.jl index b4bf7e4c52..1343b77348 100644 --- a/src/terms/terms.jl +++ b/src/terms/terms.jl @@ -4,7 +4,7 @@ include("operators.jl") # - A Term is something that, given a state, returns a named tuple (; E, hams) with an energy # and a list of RealFourierOperator (for each kpoint). # - Each term must overload -# `ene_ops(term, basis, ψ, occupation; kwargs...)` +# `ene_ops(term, ψ, occupation; kwargs...)` # -> (; E::Real, ops::Vector{RealFourierOperator}). # - Note that terms are allowed to hold on to references to ψ (eg Fock term), # so ψ should not mutated after ene_ops @@ -26,8 +26,8 @@ abstract type TermNonlinear <: Term end A term with a constant zero energy. """ struct TermNoop <: Term end -function ene_ops(term::TermNoop, basis::PlaneWaveBasis{T}, ψ, occupation; kwargs...) where {T} - (; E=zero(eltype(T)), ops=[NoopOperator(basis, kpt) for kpt in basis.kpoints]) +function ene_ops(term::TermNoop, ψ::BlochWaves{T}, occupation; kwargs...) where {T} + (; E=zero(eltype(T)), ops=[NoopOperator(ψ.basis, kpt) for kpt in ψ.basis.kpoints]) end include("Hamiltonian.jl") @@ -59,9 +59,9 @@ include("anyonic.jl") breaks_symmetries(::Anyonic) = true # forces computes either nothing or an array forces[at][α] (by default no forces) -compute_forces(::Term, ::AbstractBasis, ψ, occupation; kwargs...) = nothing +compute_forces(::Term, ::BlochWaves, occupation; kwargs...) = nothing # dynamical matrix for phonons computations (array dynmat[n_dim, n_atom, n_dim, n_atom]) -compute_dynmat(::Term, ::AbstractBasis, ψ, occupation; kwargs...) = nothing +compute_dynmat(::Term, ::BlochWaves, occupation; kwargs...) = nothing @doc raw""" compute_kernel(basis::PlaneWaveBasis; kwargs...) diff --git a/src/terms/xc.jl b/src/terms/xc.jl index 3a79070d0e..f4581f11ae 100644 --- a/src/terms/xc.jl +++ b/src/terms/xc.jl @@ -55,10 +55,10 @@ struct TermXc{T,CT} <: TermNonlinear where {T,CT} ρcore::CT end -function xc_potential_real(term::TermXc, basis::PlaneWaveBasis{T}, ψ, occupation; - ρ, τ=nothing) where {T} +function xc_potential_real(term::TermXc, ψ::BlochWaves{T}, occupation; ρ, τ=nothing) where {T} @assert !isempty(term.functionals) + basis = ψ.basis model = basis.model n_spin = model.n_spin_components potential_threshold = term.potential_threshold @@ -74,7 +74,7 @@ function xc_potential_real(term::TermXc, basis::PlaneWaveBasis{T}, ψ, occupatio if isnothing(ψ) || isnothing(occupation) τ = zero(ρ) else - τ = compute_kinetic_energy_density(basis, ψ, occupation) + τ = compute_kinetic_energy_density(ψ, occupation) end end @@ -135,11 +135,11 @@ function xc_potential_real(term::TermXc, basis::PlaneWaveBasis{T}, ψ, occupatio (; E, potential, Vτ) end -@views @timing "ene_ops: xc" function ene_ops(term::TermXc, basis::PlaneWaveBasis{T}, - ψ, occupation; ρ, τ=nothing, - kwargs...) where {T} - E, Vxc, Vτ = xc_potential_real(term, basis, ψ, occupation; ρ, τ) +@views @timing "ene_ops: xc" function ene_ops(term::TermXc, ψ::BlochWaves, occupation; + ρ, τ=nothing, kwargs...) + E, Vxc, Vτ = xc_potential_real(term, ψ, occupation; ρ, τ) + basis = ψ.basis ops = map(basis.kpoints) do kpt if !isnothing(Vτ) [RealSpaceMultiplication(basis, kpt, Vxc[:, :, :, kpt.spin]), @@ -151,14 +151,14 @@ end (; E, ops) end -@timing "forces: xc" function compute_forces(term::TermXc, basis::PlaneWaveBasis{T}, - ψ, occupation; ρ, τ=nothing, - kwargs...) where {T} +@timing "forces: xc" function compute_forces(term::TermXc, ψ::BlochWaves{T}, occupation; + ρ, τ=nothing, kwargs...) where {T} # the only non-zero force contribution is from the nlcc core charge # early return if nlcc is disabled / no elements have model core charges isnothing(term.ρcore) && return nothing - Vxc_real = xc_potential_real(term, basis, ψ, occupation; ρ, τ).potential + basis = ψ.basis + Vxc_real = xc_potential_real(term, ψ, occupation; ρ, τ).potential # TODO: the factor of 2 here should be associated with the density, not the potential if basis.model.spin_polarization in (:none, :spinless) Vxc_fourier = fft(basis, Vxc_real[:,:,:,1]) diff --git a/src/transfer.jl b/src/transfer.jl index a9c80d9bac..74689588f7 100644 --- a/src/transfer.jl +++ b/src/transfer.jl @@ -95,13 +95,14 @@ Transfer an array `ψk` defined on basis_in ``k``-point kpt_in to basis_out ``k` function transfer_blochwave_kpt(ψk_in, basis_in::PlaneWaveBasis{T}, kpt_in::Kpoint, basis_out::PlaneWaveBasis{T}, kpt_out::Kpoint) where {T} kpt_in == kpt_out && return copy(ψk_in) - @assert length(G_vectors(basis_in, kpt_in)) == size(ψk_in, 1) + @assert length(G_vectors(basis_in, kpt_in)) == size(ψk_in, 2) idcsk_in, idcsk_out = transfer_mapping(basis_in, kpt_in, basis_out, kpt_out) - n_bands = size(ψk_in, 2) - ψk_out = similar(ψk_in, length(G_vectors(basis_out, kpt_out)), n_bands) + n_bands = size(ψk_in, 3) + n_components = basis_in.model.n_components + ψk_out = similar(ψk_in, n_components, length(G_vectors(basis_out, kpt_out)), n_bands) ψk_out .= 0 - ψk_out[idcsk_out, :] .= ψk_in[idcsk_in, :] + ψk_out[:, idcsk_out, :] .= ψk_in[:, idcsk_in, :] ψk_out end @@ -113,14 +114,15 @@ Beware: `ψk_out` can lose information if the shift `ΔG` is large or if the `G_ differ between `k`-points. """ function transfer_blochwave_kpt(ψk_in, basis::PlaneWaveBasis, kpt_in, kpt_out, ΔG) - ψk_out = zeros(eltype(ψk_in), length(G_vectors(basis, kpt_out)), size(ψk_in, 2)) + ψk_out = zeros(eltype(ψk_in), basis.model.n_components, length(G_vectors(basis, kpt_out)), + size(ψk_in, 3)) for (iG, G) in enumerate(G_vectors(basis, kpt_in)) # e^i(kpt_in + G)r = e^i(kpt_out + G')r, where # kpt_out + G' = kpt_in + G = kpt_out + ΔG + G, and # G' = G + ΔG idx_Gp_in_kpoint = index_G_vectors(basis, kpt_out, G - ΔG) if !isnothing(idx_Gp_in_kpoint) - ψk_out[idx_Gp_in_kpoint, :] = ψk_in[iG, :] + ψk_out[:, idx_Gp_in_kpoint, :] = ψk_in[:, iG, :] end end ψk_out @@ -132,6 +134,7 @@ Transfer Bloch wave between two basis sets. Limited feature set. function transfer_blochwave(ψ_in, basis_in::PlaneWaveBasis{T}, basis_out::PlaneWaveBasis{T}) where {T} @assert basis_in.model.lattice == basis_out.model.lattice + @assert basis_in.model.n_components == basis_out.model.n_components @assert length(basis_in.kpoints) == length(basis_out.kpoints) @assert all(basis_in.kpoints[ik].coordinate == basis_out.kpoints[ik].coordinate for ik = 1:length(basis_in.kpoints)) @@ -147,9 +150,10 @@ function transfer_blochwave(ψ_in, basis_in::PlaneWaveBasis{T}, # It is then of size G_vectors(basis_out.kpoints[ik]) and the transfer can be done with # ψ_out[ik] .= ψ_in[ik][idcs_in[ik], :] - map(enumerate(basis_out.kpoints)) do (ik, kpt_out) - transfer_blochwave_kpt(ψ_in[ik], basis_in, basis_in.kpoints[ik], basis_out, kpt_out) - end + BlochWaves(basis_out, map(enumerate(basis_out.kpoints)) do (ik, kpt_out) + transfer_blochwave_kpt(ψ_in[ik], basis_in, basis_in.kpoints[ik], + basis_out, kpt_out) + end) end @doc raw""" diff --git a/src/workarounds/forwarddiff_rules.jl b/src/workarounds/forwarddiff_rules.jl index e5aed15664..431176db1a 100644 --- a/src/workarounds/forwarddiff_rules.jl +++ b/src/workarounds/forwarddiff_rules.jl @@ -105,12 +105,13 @@ end next_working_fft_size(::Type{<:ForwardDiff.Dual}, size::Int) = size -function build_fft_plans!(tmp::AbstractArray{Complex{T}}) where {T<:ForwardDiff.Dual} - opFFT = AbstractFFTs.plan_fft(tmp) - opBFFT = AbstractFFTs.plan_bfft(tmp) +function build_fft_plans!(tmp::AbstractArray{Complex{T}}; + region=1:ndims(tmp)) where {T<:ForwardDiff.Dual} + opFFT = AbstractFFTs.plan_fft(tmp, region) + opBFFT = AbstractFFTs.plan_bfft(tmp, region) ipFFT = DummyInplace{typeof(opFFT)}(opFFT) ipBFFT = DummyInplace{typeof(opBFFT)}(opBFFT) - ipFFT, opFFT, ipBFFT, opBFFT + (; ipFFT, opFFT, ipBFFT, opBFFT) end # determine symmetry operations only from primal lattice values @@ -149,6 +150,7 @@ function construct_value(model::Model{T}) where {T <: ForwardDiff.Dual} temperature=ForwardDiff.value(model.temperature), model.smearing, εF=ForwardDiff.value(model.εF), + model.n_components, model.spin_polarization, model.symmetries, # Can be safely disabled: this has been checked for basis.model @@ -198,14 +200,14 @@ function self_consistent_field(basis_dual::PlaneWaveBasis{T}; ## Compute external perturbation (contained in ham_dual) and from matvec with bands Hψ_dual = let occupation_dual = [T.(occk) for occk in scfres.occupation] - ψ_dual = [Complex.(T.(real(ψk)), T.(imag(ψk))) for ψk in scfres.ψ] - ρ_dual = compute_density(basis_dual, ψ_dual, occupation_dual) + ψ_dual_values = [Complex.(T.(real(ψk)), T.(imag(ψk))) for ψk in scfres.ψ] + ψ_dual = BlochWaves(basis_dual, ψ_dual_values) + ρ_dual = compute_density(ψ_dual, occupation_dual) εF_dual = T(scfres.εF) # Only needed for entropy term eigenvalues_dual = [T.(εk) for εk in scfres.eigenvalues] - ham_dual = energy_hamiltonian(basis_dual, ψ_dual, occupation_dual; - ρ=ρ_dual, eigenvalues=eigenvalues_dual, - εF=εF_dual).ham - ham_dual * ψ_dual + ham_dual = energy_hamiltonian(ψ_dual, occupation_dual; + ρ=ρ_dual, eigenvalues=eigenvalues_dual, εF=εF_dual).ham + ham_dual * ψ_dual_values end ## Implicit differentiation @@ -217,12 +219,12 @@ function self_consistent_field(basis_dual::PlaneWaveBasis{T}; ## Convert and combine DT = ForwardDiff.Dual{ForwardDiff.tagtype(T)} - ψ = map(scfres.ψ, getfield.(δresults, :δψ)...) do ψk, δψk... + ψ = BlochWaves(basis_dual, map(scfres.ψ, getfield.(δresults, :δψ)...) do ψk, δψk... map(ψk, δψk...) do ψnk, δψnk... Complex(DT(real(ψnk), real.(δψnk)), DT(imag(ψnk), imag.(δψnk))) end - end + end) ρ = map((ρi, δρi...) -> DT(ρi, δρi), scfres.ρ, getfield.(δresults, :δρ)...) eigenvalues = map(scfres.eigenvalues, getfield.(δresults, :δeigenvalues)...) do εk, δεk... map((εnk, δεnk...) -> DT(εnk, δεnk), εk, δεk...) @@ -235,7 +237,7 @@ function self_consistent_field(basis_dual::PlaneWaveBasis{T}; # TODO Could add δresults[α].δVind the dual part of the total local potential in ham_dual # and in this way return a ham that represents also the total change in Hamiltonian - energies, ham = energy_hamiltonian(basis_dual, ψ, occupation; ρ, eigenvalues, εF) + energies, ham = energy_hamiltonian(ψ, occupation; ρ, eigenvalues, εF) # This has to be changed whenever the scfres structure changes (; ham, basis=basis_dual, energies, ρ, eigenvalues, occupation, εF, ψ, diff --git a/test/chi0.jl b/test/chi0.jl index b558696757..dcfad9111f 100644 --- a/test/chi0.jl +++ b/test/chi0.jl @@ -35,7 +35,7 @@ function test_chi0(testcase; symmetries=false, temperature=0, spin_polarization= model_kwargs...) basis = PlaneWaveBasis(model; basis_kwargs...) ρ0 = guess_density(basis, magnetic_moments) - ham0 = energy_hamiltonian(basis, nothing, nothing; ρ=ρ0).ham + ham0 = energy_hamiltonian(BlochWaves(basis), nothing; ρ=ρ0).ham nbandsalg = is_εF_fixed ? FixedBands(; n_bands_converge=6) : AdaptiveBands(model) res = DFTK.next_density(ham0, nbandsalg; tol, eigensolver) scfres = (; ham=ham0, res...) @@ -56,7 +56,7 @@ function test_chi0(testcase; symmetries=false, temperature=0, spin_polarization= model = model_LDA(testcase.lattice, testcase.atoms, testcase.positions; model_kwargs..., extra_terms=[term_builder]) basis = PlaneWaveBasis(model; basis_kwargs...) - ham = energy_hamiltonian(basis, nothing, nothing; ρ=ρ0).ham + ham = energy_hamiltonian(BlochWaves(basis), nothing; ρ=ρ0).ham res = DFTK.next_density(ham, nbandsalg; tol, eigensolver) res.ρout end @@ -71,9 +71,9 @@ function test_chi0(testcase; symmetries=false, temperature=0, spin_polarization= @test norm(diff_findiff - diff_applied_χ0) < testtol # Test apply_χ0 without extra bands - ψ_occ, occ_occ = DFTK.select_occupied_orbitals(basis, scfres.ψ, scfres.occupation; + ψ_occ, occ_occ = DFTK.select_occupied_orbitals(scfres.ψ, scfres.occupation; threshold=scfres.occupation_threshold) - ε_occ = [scfres.eigenvalues[ik][1:size(ψk, 2)] for (ik, ψk) in enumerate(ψ_occ)] + ε_occ = [scfres.eigenvalues[ik][1:size(ψk, 3)] for (ik, ψk) in enumerate(ψ_occ)] diff_applied_χ0_noextra = apply_χ0(scfres.ham, ψ_occ, occ_occ, scfres.εF, ε_occ, δV; scfres.occupation_threshold) diff --git a/test/compute_density.jl b/test/compute_density.jl index 0b01dcb26c..510d660071 100644 --- a/test/compute_density.jl +++ b/test/compute_density.jl @@ -25,35 +25,36 @@ res = diagonalize_all_kblocks(lobpcg_hyper, ham, n_bands; tol) occ, εF = DFTK.compute_occupation(basis, res.λ, FermiBisection()) - ρnew = compute_density(basis, res.X, occ) + ρnew = compute_density(BlochWaves(basis, res.X), occ) for it = 1:n_rounds ham = Hamiltonian(basis; ρ=ρnew) - res = diagonalize_all_kblocks(lobpcg_hyper, ham, n_bands; tol, ψguess=res.X) + ψguess = BlochWaves(basis, res.X) + res = diagonalize_all_kblocks(lobpcg_hyper, ham, n_bands; tol, ψguess) occ, εF = DFTK.compute_occupation(basis, res.λ, FermiBisection()) - ρnew = compute_density(basis, res.X, occ) + ρnew = compute_density(BlochWaves(basis, res.X), occ) end ham, res.X, res.λ, ρnew, occ end function test_orthonormality(basis, ψ; tol=1e-8) - n_k = length(ψ) - n_states = size(ψ[1], 2) + n_states = size(ψ[1], 3) n_fft = prod(basis.fft_size) for (ik, kpt) in enumerate(basis.kpoints) - # Fourier-transform the wave functions to real space - ψk = ψ[ik] - ψk_real = cat((DFTK.ifft(basis, kpt, ψik) for ψik in eachcol(ψk))..., dims=4) + for σ = 1:basis.model.n_components + # Fourier-transform the wave functions to real space + ψk = ψ[ik][σ, :, :] + ψk_real = cat((DFTK.ifft(basis, kpt, ψik) for ψik in eachcol(ψk))..., dims=4) - T = real(eltype(ψk_real)) - ψk_real_mat = reshape(ψk_real, n_fft, n_states) - ψk_real_overlap = adjoint(ψk_real_mat) * ψk_real_mat - nondiag = ψk_real_overlap - I * (n_fft / basis.model.unit_cell_volume) + ψk_real_mat = reshape(ψk_real, n_fft, n_states) + ψk_real_overlap = adjoint(ψk_real_mat) * ψk_real_mat + nondiag = ψk_real_overlap - I * (n_fft / basis.model.unit_cell_volume) - @test maximum(abs.(nondiag)) < tol + @test maximum(abs.(nondiag)) < tol + end end end @@ -100,7 +101,8 @@ # yields an eigenfunction of the Hamiltonian # Also check that the accumulated partial densities are equal # to the returned density. - for (ik, kpt) in enumerate(ham_ir.basis.kpoints) + @assert ham_ir.basis.model.n_components == 1 + for ik in eachindex(ham_ir.basis.kpoints) Hk_ir = ham_ir.blocks[ik] for symop in ham_ir.basis.symmetries Skpoint, ψSk = DFTK.apply_symop(symop, ham_ir.basis, Hk_ir.kpoint, ψ_ir[ik]) @@ -113,14 +115,14 @@ range = 1:length(eigenvalues_ir[ik]) - n_ignore for iband in range - ψnSk = ψSk[:, iband] + ψnSk = ψSk[1, :, iband] residual = norm(Hk_full * ψnSk - eigenvalues_ir[ik][iband] * ψnSk) @test residual < 10tol end # iband end # symop end # k - end # eigenvectors - end # testset + end # eigenvectors + end # testset end test_full_vs_irreducible("silicon (3,2,3)", silicon, [3, 2, 3], Ecut=5, tol=1e-10) diff --git a/test/compute_jacobian_eigen.jl b/test/compute_jacobian_eigen.jl index 3823a6a140..98cbf24d78 100644 --- a/test/compute_jacobian_eigen.jl +++ b/test/compute_jacobian_eigen.jl @@ -7,18 +7,22 @@ using DFTK: apply_K, apply_Ω using DFTK: precondprep!, FunctionPreconditioner using LinearMaps -function eigen_ΩplusK(basis::PlaneWaveBasis{T}, ψ, occupation, numval) where {T} +function eigen_ΩplusK(ψ::BlochWaves{T}, occupation, numval) where {T} + basis = ψ.basis + n_components = basis.model.n_components + @assert n_components == 1 + ψ_matrices = blochwaves_as_matrices(ψ) pack(ψ) = reinterpret_real(pack_ψ(ψ)) unpack(x) = unpack_ψ(reinterpret_complex(x), size.(ψ)) # compute quantites at the point which define the tangent space - ρ = compute_density(basis, ψ, occupation) - H = energy_hamiltonian(basis, ψ, occupation; ρ).ham + ρ = compute_density(ψ, occupation) + H = energy_hamiltonian(ψ, occupation; ρ).ham # preconditioner Pks = [PreconditionerTPA(basis, kpt) for kpt in basis.kpoints] - for (ik, ψk) in enumerate(ψ) + for (ik, ψk) in enumerate(ψ_matrices) precondprep!(Pks[ik], ψk) end function f_ldiv!(x, y) @@ -33,18 +37,17 @@ function eigen_ΩplusK(basis::PlaneWaveBasis{T}, ψ, occupation, numval) where { end # random starting point on the tangent space to avoid eigenvalue 0 - n_bands = size(ψ[1], 2) x0 = map(1:numval) do _ - initial = map(basis.kpoints) do kpt + initial = map(enumerate(basis.kpoints)) do (ik, kpt) n_Gk = length(G_vectors(basis, kpt)) - randn(Complex{eltype(basis)}, n_Gk, n_bands) + randn(Complex{eltype(basis)}, n_components, n_Gk, size(ψ[ik], 3)) end pack(proj_tangent(initial, ψ)) end - x0 = hcat(x0...) + x0 = reduce(hcat, x0) # Rayleigh-coefficients - Λ = [ψk'Hψk for (ψk, Hψk) in zip(ψ, H * ψ)] + Λ = [ψk'Hψk for (ψk, Hψk) in zip(ψ_matrices, H * ψ_matrices)] # mapping of the linear system on the tangent space function ΩpK(x) @@ -73,9 +76,9 @@ end model = model_atomic(testcase.lattice, testcase.atoms, testcase.positions) basis = PlaneWaveBasis(model; Ecut=5, kgrid=[1, 1, 1]) scfres = self_consistent_field(basis; tol=1e-8) - ψ, occupation = select_occupied_orbitals(basis, scfres.ψ, scfres.occupation) + ψ, occupation = select_occupied_orbitals(scfres.ψ, scfres.occupation) - res = eigen_ΩplusK(basis, ψ, occupation, numval) + res = eigen_ΩplusK(ψ, occupation, numval) gap = scfres.eigenvalues[1][5] - scfres.eigenvalues[1][4] # in the linear case, the smallest eigenvalue of Ω is the gap @@ -89,9 +92,9 @@ end model = model_LDA(testcase.lattice, testcase.atoms, testcase.positions) basis = PlaneWaveBasis(model; Ecut=5, kgrid=[1, 1, 1]) scfres = self_consistent_field(basis; tol=1e-8) - ψ, occupation = select_occupied_orbitals(basis, scfres.ψ, scfres.occupation) + ψ, occupation = select_occupied_orbitals(scfres.ψ, scfres.occupation) - res = eigen_ΩplusK(basis, ψ, occupation, numval) + res = eigen_ΩplusK(ψ, occupation, numval) @test res.λ[1] > 1e-3 end end diff --git a/test/energies_guess_density.jl b/test/energies_guess_density.jl index c471a45e33..ca2c13a945 100644 --- a/test/energies_guess_density.jl +++ b/test/energies_guess_density.jl @@ -17,7 +17,7 @@ basis = PlaneWaveBasis(model; Ecut, kgrid, fft_size, kshift) ρ0 = guess_density(basis, ValenceDensityGaussian()) - E, H = energy_hamiltonian(basis, nothing, nothing; ρ=ρ0) + E, H = energy_hamiltonian(BlochWaves(basis), nothing; ρ=ρ0) @test E["Hartree"] ≈ 0.3527293727197568 atol=5e-8 @test E["Xc"] ≈ -2.3033165870558165 atol=5e-8 @@ -26,8 +26,8 @@ res = diagonalize_all_kblocks(lobpcg_hyper, H, n_bands, tol=1e-9) occupation = [[2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0] for i = 1:length(basis.kpoints)] - ρ = compute_density(H.basis, res.X, occupation) - E, H = energy_hamiltonian(basis, res.X, occupation; ρ) + ρ = compute_density(BlochWaves(H.basis, res.X), occupation) + E, H = energy_hamiltonian(BlochWaves(basis, res.X), occupation; ρ) @test E["Kinetic"] ≈ 3.3824289861522194 atol=5e-8 @test E["AtomicLocal"] ≈ -2.4178712046759157 atol=5e-8 @@ -49,7 +49,7 @@ PairwisePotential(V, params)], ) basis = PlaneWaveBasis(model; Ecut, kgrid, fft_size, kshift) - E, H = energy_hamiltonian(basis, res.X, occupation; ρ) + E, H = energy_hamiltonian(BlochWaves(basis, res.X), occupation; ρ) @test E["Kinetic"] ≈ 3.3824289861522194 atol=5e-8 @test E["AtomicLocal"] ≈ -2.4178712046759157 atol=5e-8 diff --git a/test/hamiltonian_consistency.jl b/test/hamiltonian_consistency.jl index 11128503d8..4b56d881c0 100644 --- a/test/hamiltonian_consistency.jl +++ b/test/hamiltonian_consistency.jl @@ -11,7 +11,7 @@ function test_matrix_repr_operator(hamk, ψk; atol=1e-8) for operator in hamk.operators try operator_matrix = Matrix(operator) - @test norm(operator_matrix*ψk - operator*ψk) < atol + @test norm(operator_matrix*ψk[1, :, :] - (operator*ψk)[1, :, :]) < atol catch e allowed_missing_operators = Union{DFTK.DivAgradOperator, DFTK.MagneticFieldOperator} @@ -38,25 +38,28 @@ function test_consistency_term(term; rtol=1e-4, atol=1e-8, ε=1e-6, kgrid=[1, 2, n_bands = div(n_electrons, 2, RoundUp) filled_occ = DFTK.filled_occupation(model) - ψ = [Matrix(qr(randn(ComplexF64, length(G_vectors(basis, kpt)), n_bands)).Q) - for kpt in basis.kpoints] + ψ = BlochWaves(basis, map(basis.kpoints) do kpt + n_Gk = length(G_vectors(basis, kpt)) + Q = Matrix(qr(randn(ComplexF64, n_Gk, n_bands)).Q) + reshape(Q, 1, n_Gk, n_bands) + end) occupation = [filled_occ * rand(n_bands) for _ = 1:length(basis.kpoints)] occ_scaling = n_electrons / sum(sum(occupation)) occupation = [occ * occ_scaling for occ in occupation] ρ = with_logger(NullLogger()) do - compute_density(basis, ψ, occupation) + compute_density(ψ, occupation) end - E0, ham = energy_hamiltonian(basis, ψ, occupation; ρ) + E0, ham = energy_hamiltonian(ψ, occupation; ρ) @assert length(basis.terms) == 1 δψ = [randn(ComplexF64, size(ψ[ik])) for ik = 1:length(basis.kpoints)] function compute_E(ε) - ψ_trial = ψ .+ ε .* δψ + ψ_trial = BlochWaves(basis, denest(ψ) .+ ε .* δψ) ρ_trial = with_logger(NullLogger()) do - compute_density(basis, ψ_trial, occupation) + compute_density(ψ_trial, occupation) end - E = energy_hamiltonian(basis, ψ_trial, occupation; ρ=ρ_trial).energies + E = energy_hamiltonian(ψ_trial, occupation; ρ=ρ_trial).energies E.total end @@ -64,9 +67,9 @@ function test_consistency_term(term; rtol=1e-4, atol=1e-8, ε=1e-6, kgrid=[1, 2, diff_predicted = 0.0 for ik = 1:length(basis.kpoints) - Hψk = ham.blocks[ik]*ψ[ik] + Hψk = ham.blocks[ik]*ψ[ik][1, :, :] test_matrix_repr_operator(ham.blocks[ik], ψ[ik]; atol) - δψkHψk = sum(occupation[ik][iband] * real(dot(δψ[ik][:, iband], Hψk[:, iband])) + δψkHψk = sum(occupation[ik][iband] * real(dot(δψ[ik][1, :, iband], Hψk[:, iband])) for iband=1:n_bands) diff_predicted += 2 * basis.kweights[ik] * δψkHψk end diff --git a/test/hessian.jl b/test/hessian.jl index 08577c6547..8d095060ab 100644 --- a/test/hessian.jl +++ b/test/hessian.jl @@ -7,13 +7,13 @@ function setup_quantities(testcase) basis = PlaneWaveBasis(model; Ecut=3, kgrid=(3, 3, 3), fft_size=[9, 9, 9]) scfres = self_consistent_field(basis; tol=10) - ψ, occupation = select_occupied_orbitals(basis, scfres.ψ, scfres.occupation) + ψ, occupation = select_occupied_orbitals(scfres.ψ, scfres.occupation) - ρ = compute_density(basis, ψ, occupation) - rhs = compute_projected_gradient(basis, ψ, occupation) - ϕ = rhs + ψ + ρ = compute_density(ψ, occupation) + rhs = compute_projected_gradient(ψ, occupation) + ϕ = rhs + denest(ψ) - (; scfres, basis, ψ, occupation, ρ, rhs, ϕ) + (; scfres, ψ, occupation, ρ, rhs, ϕ) end end @@ -22,11 +22,11 @@ end =# tags=[:dont_test_mpi] setup=[Hessian, TestCases] begin using DFTK: solve_ΩplusK using LinearAlgebra - (; basis, ψ, occupation, rhs, ϕ) = Hessian.setup_quantities(TestCases.silicon) + (; ψ, occupation, rhs, ϕ) = Hessian.setup_quantities(TestCases.silicon) @test isapprox( - real(dot(ϕ, solve_ΩplusK(basis, ψ, rhs, occupation).δψ)), - real(dot(solve_ΩplusK(basis, ψ, ϕ, occupation).δψ, rhs)), + real(dot(ϕ, solve_ΩplusK(ψ, rhs, occupation).δψ)), + real(dot(solve_ΩplusK(ψ, ϕ, occupation).δψ, rhs)), atol=1e-7 ) end @@ -36,11 +36,12 @@ end using DFTK using DFTK: apply_Ω using LinearAlgebra - (; basis, ψ, occupation, ρ, rhs, ϕ) = Hessian.setup_quantities(TestCases.silicon) + (; ψ, occupation, ρ, rhs, ϕ) = Hessian.setup_quantities(TestCases.silicon) + ψ_arr = denest(ψ) - H = energy_hamiltonian(basis, ψ, occupation; ρ).ham + H = energy_hamiltonian(ψ, occupation; ρ).ham # Rayleigh-coefficients - Λ = [ψk'Hψk for (ψk, Hψk) in zip(ψ, H * ψ)] + Λ = [ψk[1, :, :]'Hψk[1, :, :] for (ψk, Hψk) in zip(ψ_arr, H * ψ_arr)] # Ω is complex-linear and so self-adjoint as a complex operator. @test isapprox( @@ -54,13 +55,13 @@ end =# tags=[:dont_test_mpi] setup=[Hessian, TestCases] begin using DFTK: apply_K using LinearAlgebra - (; basis, ψ, occupation, ρ, rhs, ϕ) = Hessian.setup_quantities(TestCases.silicon) + (; ψ, occupation, ρ, rhs, ϕ) = Hessian.setup_quantities(TestCases.silicon) # K involves conjugates and is only a real-linear operator, # hence we test using the real dot product. @test isapprox( - real(dot(ϕ, apply_K(basis, rhs, ψ, ρ, occupation))), - real(dot(apply_K(basis, ϕ, ψ, ρ, occupation), rhs)), + real(dot(ϕ, apply_K(ψ.basis, rhs, ψ, ρ, occupation))), + real(dot(apply_K(ψ.basis, ϕ, ψ, ρ, occupation), rhs)), atol=1e-7 ) end @@ -76,8 +77,8 @@ end basis = PlaneWaveBasis(model; Ecut=3, kgrid=(3, 3, 3), fft_size=[9, 9, 9]) scfres = self_consistent_field(basis; tol=10) - rhs = compute_projected_gradient(basis, scfres.ψ, scfres.occupation) - ϕ = rhs + scfres.ψ + rhs = compute_projected_gradient(scfres.ψ, scfres.occupation) + ϕ = rhs + denest(scfres.ψ) @testset "self-adjointness of solve_ΩplusK_split" begin @test isapprox(real(dot(ϕ, solve_ΩplusK_split(scfres, rhs).δψ)), @@ -88,11 +89,12 @@ end @testset "solve_ΩplusK_split agrees with solve_ΩplusK" begin scfres = self_consistent_field(basis; tol=1e-10) δψ1 = solve_ΩplusK_split(scfres, rhs).δψ - δψ1 = select_occupied_orbitals(basis, δψ1, scfres.occupation).ψ - (; ψ, occupation) = select_occupied_orbitals(basis, scfres.ψ, scfres.occupation) - rhs_trunc = select_occupied_orbitals(basis, rhs, occupation).ψ - δψ2 = solve_ΩplusK(basis, ψ, rhs_trunc, occupation).δψ - @test norm(δψ1 - δψ2) < 1e-7 + δψ1 = select_occupied_orbitals(BlochWaves(basis, δψ1), scfres.occupation).ψ + (; ψ, occupation) = select_occupied_orbitals(scfres.ψ, scfres.occupation) + rhs_trunc = select_occupied_orbitals(BlochWaves(basis, rhs), occupation).ψ + δψ2 = solve_ΩplusK(ψ, rhs_trunc, occupation).δψ + # T@D@ + @test norm([e for e in δψ1] - δψ2) < 1e-7 end end @@ -109,8 +111,8 @@ end scfres = self_consistent_field(basis; tol=1e-12, nbandsalg) ψ = scfres.ψ - rhs = compute_projected_gradient(basis, scfres.ψ, scfres.occupation) - ϕ = rhs + ψ + rhs = compute_projected_gradient(scfres.ψ, scfres.occupation) + ϕ = rhs + denest(ψ) @testset "self-adjointness of solve_ΩplusK_split" begin @test isapprox(real(dot(ϕ, solve_ΩplusK_split(scfres, rhs).δψ)), diff --git a/test/kernel.jl b/test/kernel.jl index e153339db8..1e936e4f7f 100644 --- a/test/kernel.jl +++ b/test/kernel.jl @@ -28,8 +28,8 @@ δρ = randn(size(ρ0)) ρ_minus = ρ0 - ε * δρ ρ_plus = ρ0 + ε * δρ - ops_minus = DFTK.ene_ops(term, basis, nothing, nothing; ρ=ρ_minus).ops - ops_plus = DFTK.ene_ops(term, basis, nothing, nothing; ρ=ρ_plus).ops + ops_minus = DFTK.ene_ops(term, BlochWaves(basis), nothing; ρ=ρ_minus).ops + ops_plus = DFTK.ene_ops(term, BlochWaves(basis), nothing; ρ=ρ_plus).ops δV = zero(ρ0) for iσ = 1:model.n_spin_components diff --git a/test/multicomponents.jl b/test/multicomponents.jl new file mode 100644 index 0000000000..46c698aead --- /dev/null +++ b/test/multicomponents.jl @@ -0,0 +1,85 @@ +# Note: The consistency tests at temperature may have an error of the order of magnitude +# of `default_occupation_threshold`; use `nbandsalg` if small tolerance. +@testsetup module Multicomponents +using Test +using DFTK +using LinearAlgebra + +function test_consistency(setup; Ecut=5, kgrid=(4,4,4), tol=1e-5) + res1 = setup(Ecut, kgrid, tol, 1) + res2 = setup(Ecut, kgrid, tol, 2) + @test norm(res1 - res2) < tol +end +end + +@testitem "Energy & Density: LOBPCG" setup=[TestCases, Multicomponents] begin + using DFTK + test_consistency = Multicomponents.test_consistency + testcase = TestCases.silicon + testcase = merge(testcase, (; temperature=0.03)) + eigensolver = lobpcg_hyper + + test_consistency() do Ecut, kgrid, tol, n_components + model = model_PBE(testcase.lattice, testcase.atoms, testcase.positions; + n_components=1, testcase.temperature) + basis = PlaneWaveBasis(model; Ecut, kgrid) + nbandsalg = AdaptiveBands(model; occupation_threshold=tol/10) + scfres = self_consistent_field(basis; eigensolver, tol, nbandsalg) + [scfres.energies.total, scfres.ρ] + end +end + +@testitem "Energy & Density: diag_full" setup=[TestCases, Multicomponents] begin + using DFTK + test_consistency = Multicomponents.test_consistency + testcase = TestCases.silicon + testcase = merge(testcase, (; temperature=0.03)) + eigensolver = diag_full + + test_consistency() do Ecut, kgrid, tol, n_components + model = model_PBE(testcase.lattice, testcase.atoms, testcase.positions; + n_components=1, testcase.temperature) + basis = PlaneWaveBasis(model; Ecut, kgrid) + nbandsalg = AdaptiveBands(model; occupation_threshold=tol/10) + scfres = self_consistent_field(basis; eigensolver, tol, nbandsalg) + [scfres.energies.total, scfres.ρ] + end +end + +@testitem "Forces" setup=[TestCases, Multicomponents] begin + using DFTK + test_consistency = Multicomponents.test_consistency + testcase = TestCases.silicon + testcase = merge(testcase, (; temperature=0.03)) + + test_consistency() do Ecut, kgrid, tol, n_components + positions = [(ones(3)) / 4, -ones(3) / 8] + model = model_PBE(testcase.lattice, testcase.atoms, positions; n_components, + testcase.temperature) + basis = PlaneWaveBasis(model; Ecut, kgrid) + nbandsalg = AdaptiveBands(model; occupation_threshold=tol/100) + scfres = self_consistent_field(basis; tol=tol/10, nbandsalg) + compute_forces(scfres) + end +end + +@testitem "δρ, temperature" setup=[TestCases, Multicomponents] begin + using DFTK + using LinearAlgebra + test_consistency = Multicomponents.test_consistency + testcase = TestCases.aluminium_primitive + testcase = merge(testcase, (; temperature=0.03)) + + Ecut = 10 + kgrid = (1, 1, 1) + + test_consistency(; Ecut, kgrid) do Ecut, kgrid, tol, n_components + model = model_PBE(testcase.lattice, testcase.atoms, testcase.positions; + n_components, testcase.temperature) + basis = PlaneWaveBasis(model; Ecut, kgrid) + nbandsalg = AdaptiveBands(model; occupation_threshold=tol/10) + scfres = self_consistent_field(basis; tol, nbandsalg) + δV = guess_density(basis) + apply_χ0(scfres, δV) + end +end diff --git a/test/occupation.jl b/test/occupation.jl index 125c8ad367..c9cdac81e4 100644 --- a/test/occupation.jl +++ b/test/occupation.jl @@ -210,8 +210,7 @@ end smearing isa Smearing.None && continue occupation, εF = DFTK.compute_occupation(scfres.basis, scfres.eigenvalues, alg; smearing, temperature) - ρ = DFTK.compute_density(scfres.basis, scfres.ψ, scfres.occupation; - scfres.occupation_threshold) + ρ = DFTK.compute_density(scfres.ψ, scfres.occupation; scfres.occupation_threshold) atol = scfres.occupation_threshold @test DFTK.weighted_ksum(basis, sum.(occupation)) ≈ model.n_electrons atol=atol @test sum(ρ) * scfres.basis.dvol ≈ model.n_electrons atol=atol diff --git a/test/pairwise.jl b/test/pairwise.jl index 01b98c84f0..98befa4e49 100644 --- a/test/pairwise.jl +++ b/test/pairwise.jl @@ -20,7 +20,7 @@ model = Model(lattice, atoms, positions; terms=[term]) basis = PlaneWaveBasis(model; Ecut=20, kgrid=(1, 1, 1)) - forces = compute_forces(only(basis.terms), basis, nothing, nothing) + forces = compute_forces(only(basis.terms), BlochWaves(basis), nothing) # Compare forces to finite differences ε=1e-8 diff --git a/test/phonon/helpers.jl b/test/phonon/helpers.jl index d843fcdb41..e966c5a1dd 100644 --- a/test/phonon/helpers.jl +++ b/test/phonon/helpers.jl @@ -29,7 +29,7 @@ function ph_compute_reference(basis_supercell) model_disp = Model(convert(Model{eltype(ε)}, model_supercell); lattice, positions) # TODO: Would be cleaner with PR #675. basis_disp_bs = PlaneWaveBasis(model_disp; Ecut=5) - forces = compute_forces(basis_disp_bs, nothing, nothing) + forces = compute_forces(BlochWaves(basis_disp_bs), nothing) reduce(hcat, forces) end end diff --git a/test/phonon/lowlevel.jl b/test/phonon/lowlevel.jl index cba3ba661a..791f792ee8 100644 --- a/test/phonon/lowlevel.jl +++ b/test/phonon/lowlevel.jl @@ -44,7 +44,7 @@ for kpt in unique(rand(basis.kpoints, 4)) ψk = fft(basis, kpt, ψ) - ψk_out_four = multiply_by_expiqr(basis, kpt, q, ψk) + ψk_out_four = multiply_by_expiqr(basis, kpt, q, reshape(ψk, 1, length(ψk), 1)) ψk_out_real = let shifted_kcoord = kpt.coordinate .+ q index, ΔG = find_equivalent_kpt(basis, shifted_kcoord, kpt.spin) @@ -52,7 +52,7 @@ transfer_blochwave_kpt_real(ψk, basis, kpt, kpt_out, ΔG) end @testset "Testing kpoint $(kpt.coordinate) on kgrid $kgrid" begin - @test norm(ψk_out_four - ψk_out_real) < tol + @test norm(ψk_out_four[1, :, :] - ψk_out_real) < tol end end end diff --git a/test/scf_compare.jl b/test/scf_compare.jl index 6dfcc56df6..a80da8b23b 100644 --- a/test/scf_compare.jl +++ b/test/scf_compare.jl @@ -25,8 +25,8 @@ @testset "Newton" begin scfres_start = self_consistent_field(basis, maxiter=1) # remove virtual orbitals - ψ0 = select_occupied_orbitals(basis, scfres_start.ψ, scfres_start.occupation).ψ - ρ_newton = newton(basis, ψ0; tol).ρ + ψ0 = select_occupied_orbitals(scfres_start.ψ, scfres_start.occupation).ψ + ρ_newton = newton(ψ0; tol).ρ @test maximum(abs, ρ_newton - ρ_def) < 10tol end end @@ -74,12 +74,12 @@ end ρ0 = guess_density(basis, magnetic_moments) ρ_def = self_consistent_field(basis; tol, ρ=ρ0).ρ scfres_start = self_consistent_field(basis, maxiter=1, ρ=ρ0) - ψ0 = select_occupied_orbitals(basis, scfres_start.ψ, scfres_start.occupation).ψ + ψ0 = select_occupied_orbitals(scfres_start.ψ, scfres_start.occupation).ψ # Run DM if mpi_nprocs() == 1 # Distributed implementation not yet available @testset "Direct minimization" begin - ρ_dm = direct_minimization(basis, ψ0; g_tol=tol).ρ + ρ_dm = direct_minimization(ψ0; g_tol=tol).ρ @test maximum(abs.(ρ_dm - ρ_def)) < 10tol end end @@ -87,7 +87,7 @@ end # Run Newton algorithm if mpi_nprocs() == 1 # Distributed implementation not yet available @testset "Newton" begin - ρ_newton = newton(basis, ψ0; tol).ρ + ρ_newton = newton(ψ0; tol).ρ @test maximum(abs.(ρ_newton - ρ_def)) < 10tol end end diff --git a/test/serialisation.jl b/test/serialisation.jl index fefba76484..2ba2fdee11 100644 --- a/test/serialisation.jl +++ b/test/serialisation.jl @@ -9,6 +9,7 @@ function test_scfres_agreement(tested, ref) @test tested.basis.model.εF == ref.basis.model.εF @test tested.basis.model.symmetries == ref.basis.model.symmetries @test tested.basis.model.spin_polarization == ref.basis.model.spin_polarization + @test tested.basis.model.n_components == ref.basis.model.n_components @test tested.basis.model.positions == ref.basis.model.positions @test atomic_symbol.(tested.basis.model.atoms) == atomic_symbol.(ref.basis.model.atoms) @@ -28,8 +29,12 @@ function test_scfres_agreement(tested, ref) @test tested.energies.total ≈ ref.energies.total atol=1e-13 @test tested.eigenvalues == ref.eigenvalues @test tested.occupation == ref.occupation - @test tested.ψ == ref.ψ @test tested.ρ ≈ ref.ρ rtol=1e-14 + + @test tested.ψ.data == ref.ψ.data + @test tested.ψ.basis.Ecut == ref.ψ.basis.Ecut + @test tested.ψ.basis.kweights == ref.ψ.basis.kweights + @test tested.ψ.basis.fft_size == ref.ψ.basis.fft_size end end diff --git a/test/stresses.jl b/test/stresses.jl index 7ac7678700..d530e0e4f6 100644 --- a/test/stresses.jl +++ b/test/stresses.jl @@ -17,14 +17,15 @@ function recompute_energy(lattice, symmetries, element) basis = make_basis(lattice, symmetries, element) scfres = self_consistent_field(basis; is_converged=DFTK.ScfConvergenceDensity(1e-13)) - (; energies) = energy_hamiltonian(basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ) + (; energies) = energy_hamiltonian(scfres.ψ, scfres.occupation; ρ=scfres.ρ) energies.total end function hellmann_feynman_energy(scfres, lattice, symmetries, element) basis = make_basis(lattice, symmetries, element) - ρ = DFTK.compute_density(basis, scfres.ψ, scfres.occupation) - (; energies) = energy_hamiltonian(basis, scfres.ψ, scfres.occupation; ρ) + ψ = BlochWaves(basis, denest(scfres.ψ)) + ρ = DFTK.compute_density(ψ, scfres.occupation) + (; energies) = energy_hamiltonian(ψ, scfres.occupation; ρ) energies.total end diff --git a/test/transfer.jl b/test/transfer.jl index 9312b6bfec..196fe4f480 100644 --- a/test/transfer.jl +++ b/test/transfer.jl @@ -21,13 +21,13 @@ bigger_basis = PlaneWaveBasis(model; Ecut=(Ecut + 5), kgrid, kshift) ψ_b = transfer_blochwave(ψ, basis, bigger_basis) ψ_bb = transfer_blochwave(ψ_b, bigger_basis, basis) - @test norm(ψ - ψ_bb) < eps(eltype(basis)) + @test norm([ψk - ψ_bbk for (ψk, ψ_bbk) in zip(ψ, ψ_bb)]) < eps(eltype(basis)) T = compute_transfer_matrix(basis, bigger_basis) # transfer Tᵇ = compute_transfer_matrix(bigger_basis, basis) # back-transfer - Tψ = [Tk * ψk for (Tk, ψk) in zip(T, ψ)] - TᵇTψ = [Tᵇk * Tψk for (Tᵇk, Tψk) in zip(Tᵇ, Tψ)] - @test norm(ψ - TᵇTψ) < eps(eltype(basis)) + Tψ = [Tk * ψk[1, :, :] for (Tk, ψk) in zip(T, ψ)] + TᵇTψ = [Tᵇk * Tψk for (Tᵇk, Tψk) in zip(Tᵇ, Tψ)] + @test norm([ψk[1, :, :] for ψk in ψ] - TᵇTψ) < eps(eltype(basis)) # TᵇT should be the identity and TTᵇ should be a projection TᵇT = [Tᵇk * Tk for (Tk, Tᵇk) in zip(T, Tᵇ)] @@ -40,7 +40,7 @@ bigger_basis = PlaneWaveBasis(model; Ecut, kgrid, kshift) ψ_b = transfer_blochwave(ψ, basis, bigger_basis) ψ_bb = transfer_blochwave(ψ_b, bigger_basis, basis) - @test norm(ψ-ψ_bb) < eps(eltype(basis)) + @test norm([ψk - ψ_bbk for (ψk, ψ_bbk) in zip(ψ, ψ_bb)]) < eps(eltype(basis)) end @testitem "Transfer of density" begin