[WIP-WIP] BlochWaves structure

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))
 ```
 

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
 

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]

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

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/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,15 +86,15 @@ 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
 # have one ``k``-point:
 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

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

ext/DFTKGenericLinearAlgebraExt/DFTKGenericLinearAlgebraExt.jl

@@ -29,29 +29,29 @@ 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}
     subplans
     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

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)

ext/DFTKJLD2Ext.jl

@@ -34,17 +34,20 @@ 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
         sym in kpt_properties && continue
         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

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

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")