BlochWaves structure

Author: epolack

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,16 +69,16 @@ 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
@@ -149,7 +149,7 @@ Mres = apply_metric(ψr.data, 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);
@@ -163,15 +163,15 @@ e2 = apply_metric(ψr, P, resHF, apply_inv_M);
 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
 
@@ -197,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

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.ψ

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

src/DFTK.jl

@@ -72,8 +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 to_composite_σG
-export from_composite_σG
+export BlochWaves, view_component, nest, denest
+export blochwave_as_matrix
+export blochwave_as_tensor
+export blochwaves_as_matrices
 export ifft
 export irfft
 export ifft!

src/densities.jl

@@ -1,23 +1,24 @@
 # Densities (and potentials) are represented by arrays
 # ρ[ix,iy,iz,iσ] in real space, where iσ ∈ [1:n_spin_components]
 
-# TODO: We reduce all components for the density. Will need to be though again when we merge
-# the components and the spins.
 """
-    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
@@ -66,21 +67,22 @@ 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{TT}, ψ,
-                                                       occupation) where {TT}
+@views @timing function compute_kinetic_energy_density(ψ::BlochWaves{T, Tψ},
+                                                       occupation) where {T, Tψ}
+    basis = ψ.basis
     @assert basis.model.n_components == 1
-    T = promote_type(TT, real(eltype(ψ[1])))
-    τ = similar(ψ[1], T, (basis.fft_size..., basis.model.n_spin_components))
+    TT = promote_type(T, real(Tψ))
+    τ = similar(ψ[1], TT, (basis.fft_size..., basis.model.n_spin_components))
     τ .= 0
     dαψnk_real = zeros(complex(T), basis.fft_size)
     for (ik, kpt) in enumerate(basis.kpoints)

src/orbitals.jl

@@ -4,19 +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)]
 
+    ψ = 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 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.

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

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

src/response/hessian.jl

@@ -40,9 +40,10 @@ 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)
@@ -62,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)
@@ -79,8 +81,8 @@ 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_ψ(ψ))
@@ -152,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(δρ)
@@ -183,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
+    δ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...)
 

src/scf/direct_minimization.jl

@@ -63,14 +63,16 @@ 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
@@ -81,7 +83,7 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0;
     Nk = length(basis.kpoints)
 
     if isnothing(ψ0)
-        ψ0 = [random_orbitals(basis, kpt, n_bands) for kpt in basis.kpoints]
+        ψ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]
@@ -100,8 +102,8 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0;
     # computes energies and gradients
     function fg!(E, 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
@@ -145,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

src/scf/nbands_algorithm.jl

@@ -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