big squash

docs/src/developer/data_structures.md

@@ -140,10 +140,13 @@ 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))
 ```
 
+The first index of `ψtest` is the component (e.g., the spin) and the second is the
+``G``-component.
+
 We now perform an IFFT to get ψ in real space. The ``k``-point has to be
 passed because ψ is expressed on the ``k``-dependent basis.
 Again the function is normalised:

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

@@ -85,50 +85,60 @@ res, occ = DFTK.select_occupied_orbitals(basis_ref, res, scfres.occupation)
 #   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, ϕ, ψ)
-    map(zip(ϕ, ψ)) do (ϕk, ψk)
+    ϕ = to_composite_σG(basis, ϕ)
+    ψ = to_composite_σG(basis, ψ)
+    map(zip(basis.kpoints, ϕ, ψ)) do (kpt, ϕk, ψk)
         S = ψk'ϕk
         U = S*(S'S)^(-1/2)
-        ϕk - ψk*U
+        from_composite_σG(basis, kpt, ϕk - ψk*U)
     end
 end
 err = compute_error(basis_ref, ψ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)
+map(zip(P, ψr, basis_ref.kpoints)) do (Pk, ψk, kpt)
+    DFTK.precondprep!(Pk, to_composite_σG(basis_ref, kpt, ψk))
 end
-function apply_M(φk, Pk, δφnk, n)
-    DFTK.proj_tangent_kpt!(δφnk, φk)
-    δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk
-    DFTK.proj_tangent_kpt!(δφnk, φk)
-    δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk
-    DFTK.proj_tangent_kpt!(δφnk, φk)
+function apply_M(basis, kpt, φk, Pk, δφnk, n)
+    fact = reshape(sqrt.(Pk.mean_kin[n] .+ Pk.kin), size(δφnk)...)
+    DFTK.proj_tangent_kpt!(δφnk, basis, kpt, φk)
+    δφnk = fact .* δφnk
+    DFTK.proj_tangent_kpt!(δφnk, basis, kpt, φk)
+    δφnk = fact .* δφnk
+    DFTK.proj_tangent_kpt!(δφnk, basis, kpt, φk)
 end
-function apply_inv_M(φk, Pk, δφnk, n)
-    DFTK.proj_tangent_kpt!(δφnk, φk)
-    op(x) = apply_M(φk, Pk, x, n)
+function apply_inv_M(basis, kpt, φk, Pk, δφnk, n)
+    DFTK.proj_tangent_kpt!(δφnk, basis, kpt, φk)
+    op(x) = let
+        x = from_composite_σG(basis, kpt, x)
+        to_composite_σG(basis, kpt, apply_M(basis, kpt, φk, Pk, x, n))
+    end
     function f_ldiv!(x, y)
-        x .= DFTK.proj_tangent_kpt(y, φk)
+        x_σG = from_composite_σG(basis, kpt, x)
+        y_σG = from_composite_σG(basis, kpt, y)
+        x_σG .= DFTK.proj_tangent_kpt(basis, kpt, y_σG, φk)
         x ./= (Pk.mean_kin[n] .+ Pk.kin)
-        DFTK.proj_tangent_kpt!(x, φk)
+        DFTK.proj_tangent_kpt!(x_σG, basis, kpt, φk)
+        x
     end
-    J = LinearMap{eltype(φk)}(op, size(δφnk, 1))
-    δφnk = cg(J, δφnk; Pl=DFTK.FunctionPreconditioner(f_ldiv!),
-              verbose=false, reltol=0, abstol=1e-15)
-    DFTK.proj_tangent_kpt!(δφnk, φk)
+    J = LinearMap{eltype(φk)}(op, prod(size(δφnk)[1:2]))
+    δφnk_vec = to_composite_σG(basis, kpt, δφnk)
+    δφnk_vec = cg(J, δφnk_vec; Pl=DFTK.FunctionPreconditioner(f_ldiv!),
+                  verbose=false, reltol=0, abstol=1e-15)
+    DFTK.proj_tangent_kpt!(δφnk, basis, kpt, φk)
 end
-function apply_metric(φ, P, δφ, A::Function)
+function apply_metric(basis, φ, 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(basis, basis.kpoints[ik], φk, P[ik], δφk[:, :, n], n)
         end
         Aδφk
     end
 end
-Mres = apply_metric(ψr, P, res, apply_inv_M);
+Mres = apply_metric(basis, ψr, 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]:
@@ -151,14 +161,14 @@ resLF = DFTK.transfer_blochwave(res, basis_ref, basis)
 resHF = res - DFTK.transfer_blochwave(resLF, basis, basis_ref);
 
 # - Compute ``{\boldsymbol M}^{-1}_{22}R_2(P)``:
-e2 = apply_metric(ψr, P, resHF, apply_inv_M);
+e2 = apply_metric(basis, ψr, P, resHF, apply_inv_M);
 
 # - Compute the right hand side of the Schur system:
 ## Rayleigh coefficients needed for `apply_Ω`
-Λ = map(enumerate(ψr)) do (ik, ψk)
-    Hk = hamr.blocks[ik]
-    Hψk = Hk * ψk
-    ψk'Hψk
+Λ = map(enumerate(zip(ψr, hamr * ψr))) do (ik, (ψk, Hψk))
+    ψk = to_composite_σG(basis, basis.kpoints[ik], ψk)
+    Hψk = to_composite_σG(basis, basis.kpoints[ik], Hψk)
+    from_composite_σG(basis, basis.kpoints[ik], ψk'Hψk)
 end
 ΩpKe2 = DFTK.apply_Ω(e2, ψr, hamr, Λ) .+ DFTK.apply_K(basis_ref, e2, ψr, ρr, occ)
 ΩpKe2 = DFTK.transfer_blochwave(ΩpKe2, basis_ref, basis)
@@ -204,7 +214,7 @@ end;
 #   the actual error ``P-P_*``. Usually this is of course not the case, but this
 #   is the "best" improvement we can hope for with a linearisation, so we are
 #   aiming for this precision.
-df_err = df(basis_ref, occ, ψr, DFTK.proj_tangent(err, ψr), ρr)
+df_err = df(basis_ref, occ, ψr, DFTK.proj_tangent(basis, err, ψr), ρr)
 forces["F(P) - df(P)⋅(P-P_*)"]   = f - df_err
 relerror["F(P) - df(P)⋅(P-P_*)"] = compute_relerror(f - df_err);
 

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]
@@ -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
@@ -107,4 +108,4 @@ A[1, end] = A[end, 1] = -1
 K = A / dx^2 / 2
 V = Diagonal(pot.(x) + C .* α .* (ρ.^(α-1)))
 H_findiff = K + V;
-maximum(abs.(H_findiff*ψ - (dot(ψ, H_findiff*ψ) / dot(ψ, ψ)) * ψ))
+maximum(abs, H_findiff*ψ - (dot(ψ, H_findiff*ψ) / dot(ψ, ψ)) * ψ)

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

@@ -89,6 +89,12 @@ function load_scfres_jld2(jld, basis; skip_hamiltonian, strict)
                             "($(jld["n_spin_components"])) and supplied basis " *
                             "($(basis.model.n_spin_components))")
             end
+            if jld["n_components"] != basis.model.n_components
+                consistent_kpts = false
+                errormsg = ("Mismatch in number of spin components between file " *
+                            "($(jld["n_components"])) and supplied basis " *
+                            "($(basis.model.n_components))")
+            end
         end
         errormsg = MPI.bcast(errormsg, 0, MPI.COMM_WORLD)
         isempty(errormsg) || error(errormsg)

ext/DFTKWriteVTKExt.jl

@@ -24,12 +24,14 @@ function DFTK.save_scfres(::Val{:vts}, filename::AbstractString, scfres::NamedTu
         ψblock = gather_kpts_block(scfres.basis, ψblock_dist)
 
         if mpi_master()
-            for ik = 1:length(basis.kpoints), n = 1:size(ψblock, 2)
-                kpt_n_G = length(G_vectors(basis, basis.kpoints[ik]))
-                ψnk_real = ifft(basis, basis.kpoints[ik], ψblock[1:kpt_n_G, n, ik])
-                prefix = @sprintf "ψ_k%03i_n%03i" ik n
-                vtkfile["$(prefix)_real"] = real.(ψnk_real)
-                vtkfile["$(prefix)_imag"] = imag.(ψnk_real)
+            for ik = 1:length(basis.kpoints)
+                for σ = 1:basis.model.n_components, n = 1:size(ψblock, 3)
+                    kpt_n_G = length(G_vectors(basis, basis.kpoints[ik]))
+                    ψσnk_real = ifft(basis, basis.kpoints[ik], ψblock[:, 1:kpt_n_G, n, ik])
+                    prefix = @sprintf "ψ_k%03i_n%03i_σ%03i" ik n σ
+                    vtkfile["$(prefix)_real"] = real.(ψσnk_real)
+                    vtkfile["$(prefix)_imag"] = imag.(ψσnk_real)
+                end
             end
         end
     end

src/DFTK.jl

@@ -75,6 +75,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!
@@ -86,6 +88,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

@@ -73,6 +73,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
@@ -250,7 +255,10 @@ 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
@@ -338,6 +346,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,