big squash

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

@@ -84,51 +84,52 @@ res, occ = DFTK.select_occupied_orbitals(basis_ref, res, scfres.occupation)
 #   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]:
@@ -158,7 +159,7 @@ 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)

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

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,8 @@ 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 ifft
 export irfft
 export ifft!
@@ -83,6 +85,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")

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

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,