FFTs only work one band at a time

Project.toml

@@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
+NLSolversBase = "d41bc354-129a-5804-8e4c-c37616107c6c"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
@@ -27,6 +28,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
+NLSolversBase = "= 7.5.0"
 julia = "1.2"
 
 [extras]

examples/graphene.jl

@@ -0,0 +1,33 @@
+using DFTK
+using Printf
+
+kgrid = [12, 12, 4]
+Tsmear = 0.0009500431544769484
+Ecut = 35
+
+lattice = [4.659533614391621 -2.3297668071958104 0.0;
+           0.0 4.035274479829987 0.0;
+           0.0 0.0 15.117809010356462]
+C = Species(6, load_psp("hgh/pbe/c-q4"))
+composition = [C => [[0.0, 0.0, 0.0], [0.33333333333, 0.66666666667, 0.0]]]
+
+model = model_dft(lattice, [:gga_x_pbe, :gga_c_pbe], composition...;
+                  temperature=Tsmear, smearing=DFTK.Smearing.Gaussian())
+kcoords, ksymops = bzmesh_ir_wedge(kgrid, lattice, composition...)
+basis = PlaneWaveBasis(model, Ecut, kcoords, ksymops)
+
+# Run SCF
+n_bands = 6
+ham = Hamiltonian(basis, guess_density(basis, composition...))
+scfres = self_consistent_field(ham, n_bands)
+ham = scfres.ham
+
+# Print obtained energies
+energies = scfres.energies
+energies[:Ewald] = energy_nuclear_ewald(model.lattice, composition...)
+energies[:PspCorrection] = energy_nuclear_psp_correction(model.lattice, composition...)
+println("\nEnergy breakdown:")
+for key in sort([keys(energies)...]; by=S -> string(S))
+    @printf "    %-20s%-10.7f\n" string(key) energies[key]
+end
+@printf "\n    %-20s%-15.12f\n\n" "total" sum(values(energies))

src/HamiltonianBlock.jl

@@ -23,7 +23,7 @@ Base.eltype(block::HamiltonianBlock) = complex(eltype(block.values_kinetic))
 
 
 function Matrix(block::HamiltonianBlock)
-    n_bas = length(block.kpt.basis)
+    n_bas = length(G_vectors(block.kpt))
     T = eltype(block)
     mat = Matrix{T}(undef, (n_bas, n_bas))
     v = fill(zero(T), n_bas)
@@ -47,12 +47,16 @@ function LinearAlgebra.mul!(Y, block::HamiltonianBlock, X)
     if Vloc === nothing
         Y .= kin.diag .* X
     else
-        Xreal = G_to_r(block.basis, block.kpt, X)
-        Xreal .*= Vloc
-        r_to_G!(Y, block.basis, block.kpt, Xreal)
-        Y .+= kin.diag .* X
+        @views Threads.@threads for n = 1:size(X, 2)
+            Xreal = G_to_r(block.basis, block.kpt, X[:, n])
+            Xreal .*= Vloc
+            r_to_G!(Y[:, n], block.basis, block.kpt, Xreal)
+            Y[:,n] .+= kin.diag .* X[:, n]
+        end
     end
+
     Vnloc !== nothing && (Y .+= Vnloc * X)
+
     Y
 end
 

src/Kinetic.jl

@@ -21,5 +21,5 @@ function kblock(kin::Kinetic, kpt::Kpoint)
 
     T = eltype(basis.kpoints[1].coordinate)
     Diagonal(Vector{T}([sum(abs2, model.recip_lattice * (G + kpt.coordinate))
-                        for G in kpt.basis] ./ 2))
+                        for G in G_vectors(kpt)] ./ 2))
 end

src/PlaneWaveBasis.jl

@@ -11,15 +11,17 @@ include("fft.jl")
 #
 # G_to_r and r_to_G convert between these.
 
-# Each Kpoint has its own `basis`, consisting of all G vectors such that |k+G|^2 ≤ 1/2 Ecut
+# Each Kpoint has its own `basis`, consisting of all G vectors such that 1/2 |k+G|^2 ≤ Ecut
 struct Kpoint{T <: Real}
-    spin::Symbol              # :up, :down, :both or :spinless
-    coordinate::Vec3{T}       # Fractional coordinate of k-Point
-    mapping::Vector{Int}      # Index of basis[i] on FFT grid
-    basis::Vector{Vec3{Int}}  # Wave vectors in integer coordinates
+    spin::Symbol                  # :up, :down, :both or :spinless
+    coordinate::Vec3{T}           # Fractional coordinate of k-Point
+    mapping::Vector{Int}          # Index of G_vectors[i] on the FFT grid:
+                                  # G_vectors(pwbasis)[kpt.mapping[i]] == G_vectors(kpt)[i]
+    G_vectors::Vector{Vec3{Int}}  # Wave vectors in integer coordinates
 end
+G_vectors(kpt::Kpoint) = kpt.G_vectors
 
-struct PlaneWaveBasis{T <: Real, Tgrid, TopFFT, TipFFT}
+struct PlaneWaveBasis{T <: Real, Tgrid, TopFFT, TipFFT, TopIFFT, TipIFFT}
     model::Model{T}
     Ecut::T
 
@@ -41,6 +43,8 @@ struct PlaneWaveBasis{T <: Real, Tgrid, TopFFT, TipFFT}
     # Plans for forward and backward FFT on C_ρ
     opFFT::TopFFT  # out-of-place FFT plan
     ipFFT::TipFFT  # in-place FFT plan
+    opIFFT::TopIFFT
+    ipIFFT::TipIFFT
 end
 
 """
@@ -108,6 +112,8 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Number,
     # fft must be normalized by sqrt(Ω)/length
     ipFFT *= sqrt(model.unit_cell_volume) / length(ipFFT)
     opFFT *= sqrt(model.unit_cell_volume) / length(opFFT)
+    ipIFFT = inv(ipFFT)
+    opIFFT = inv(opFFT)
 
     # Default to no symmetry
     ksymops === nothing && (ksymops = [[(Mat3{Int}(I), Vec3(zeros(3)))]
@@ -127,9 +133,9 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Number,
             "the grid supports (== $max_E)")
 
     grids = tuple((range(T(0), T(1), length=fft_size[i] + 1)[1:end-1] for i=1:3)...)
-    PlaneWaveBasis{T, typeof(grids), typeof(opFFT), typeof(ipFFT)}(
+    PlaneWaveBasis{T, typeof(grids), typeof(opFFT), typeof(ipFFT), typeof(opIFFT), typeof(ipIFFT)}(
         model, Ecut, build_kpoints(model, fft_size, kcoords, Ecut),
-        kweights, ksymops, fft_size, grids, opFFT, ipFFT
+        kweights, ksymops, fft_size, grids, opFFT, ipFFT, opIFFT, ipIFFT
     )
 end
 
@@ -164,23 +170,15 @@ function index_G_vectors(pw::PlaneWaveBasis, G::AbstractVector{T}) where {T <: I
     end
 end
 
+AbstractArray3{T} = AbstractArray{T, 3}
 
 #
 # Perform FFT.
 #
 # There are two kinds of FFT/IFFT functions: the ones on `C_ρ` (that
 # don't take a kpoint as input) and the ones on `B_k` (that do)
-# 
 # Quantities in real-space are 3D arrays, quantities in reciprocal
-# space are 3D (`C_ρ`) or 1D (`B_k`). We take as input either a single
-# array (3D or 1D), or a bunch of them (4D or 2D)
-
-AbstractArray3or4D = Union{AbstractArray{T, 3}, AbstractArray{T, 4}} where T
-AbstractArray1or2D = Union{AbstractArray{T, 1}, AbstractArray{T, 2}} where T
-
-
-# The methods below use the fact that in julia extra dimensions are ignored.
-# Eg with x = randn(3), x[:,1] == x and size(x,2) == 1
+# space are 3D (`C_ρ`) or 1D (`B_k`).
 
 @doc raw"""
     G_to_r!(f_real, pw::PlaneWaveBasis, [kpt::Kpoint, ], f_fourier)
@@ -189,32 +187,21 @@ Perform an iFFT to translate between `f_fourier`, a fourier representation
 of a function either on ``B_k`` (if `kpt` is given) or on ``C_ρ`` (if not),
 and `f_real`. The function will destroy all data in `f_real`.
 """
-function G_to_r!(f_real::AbstractArray3or4D, pw::PlaneWaveBasis,
-                 f_fourier::AbstractArray3or4D)
-    n_bands = size(f_fourier, 4)
-    @views Threads.@threads for iband in 1:n_bands
-        ldiv!(f_real[:, :, :, iband], pw.opFFT, f_fourier[:, :, :, iband])
-    end
-    f_real
+function G_to_r!(f_real::AbstractArray3, pw::PlaneWaveBasis,
+                 f_fourier::AbstractArray3) where {Tr, Tf}
+    mul!(f_real, pw.opIFFT, f_fourier)
 end
-function G_to_r!(f_real::AbstractArray3or4D, pw::PlaneWaveBasis, kpt::Kpoint,
-                 f_fourier::AbstractArray1or2D)
-    n_bands = size(f_fourier, 2)
-    @assert size(f_fourier, 1) == length(kpt.mapping)
-    @assert size(f_real)[1:3] == pw.fft_size
-    @assert size(f_real, 4) == n_bands
+function G_to_r!(f_real::AbstractArray3, pw::PlaneWaveBasis, kpt::Kpoint,
+                 f_fourier::AbstractVector)
+    @assert length(f_fourier) == length(kpt.mapping)
+    @assert size(f_real) == pw.fft_size
 
+    # Pad the input data from B_k to C_ρ
+    fill!(f_real, 0)
+    f_real[kpt.mapping] = f_fourier
 
-    @views Threads.@threads for iband in 1:n_bands
-        fill!(f_real[:, :, :, iband], 0)
-        # Pad the input data from B_k to C_ρ
-        reshape(f_real, :, n_bands)[kpt.mapping, iband] = f_fourier[:, iband]
-
-        # Perform an FFT on C_ρ -> C_ρ^*
-        ldiv!(f_real[:, :, :, iband], pw.ipFFT, f_real[:, :, :, iband])
-    end
-
-    f_real
+    # Perform an FFT on C_ρ -> C_ρ^*
+    mul!(f_real, pw.ipIFFT, f_real)
 end
 
 @doc raw"""
@@ -224,15 +211,13 @@ Perform an iFFT to translate between `f_fourier`, a fourier representation
 of a function either on ``B_k`` (if `kpt` is given) or on ``C_ρ`` (if not)
 and return the values on the real-space grid `C_ρ^*`.
 """
-function G_to_r(pw::PlaneWaveBasis, f_fourier::AbstractArray3or4D)
+function G_to_r(pw::PlaneWaveBasis, f_fourier::AbstractArray3)
     G_to_r!(similar(f_fourier), pw, f_fourier)
 end
 function G_to_r(pw::PlaneWaveBasis, kpt::Kpoint, f_fourier::AbstractVector)
     G_to_r!(similar(f_fourier, pw.fft_size...), pw, kpt, f_fourier)
 end
-function G_to_r(pw::PlaneWaveBasis, kpt::Kpoint, f_fourier::AbstractMatrix)
-    G_to_r!(similar(f_fourier, pw.fft_size..., size(f_fourier, 2)), pw, kpt, f_fourier)
-end
+
 
 
 
@@ -244,30 +229,21 @@ Perform an FFT to translate between `f_real`, a function represented on
 coefficients to ``B_k`` (if `kpt` is given). Note: If `kpt` is given, all data
 in ``f_real`` will be distroyed as well.
 """
-function r_to_G!(f_fourier::AbstractArray3or4D, pw::PlaneWaveBasis,
-                 f_real::AbstractArray3or4D)
-    n_bands = size(f_fourier, 4)
-    @views Threads.@threads for iband in 1:n_bands
-        mul!(f_fourier[:, :, :, iband], pw.opFFT, f_real[:, :, :, iband])
-    end
-    f_fourier
+function r_to_G!(f_fourier::AbstractArray3, pw::PlaneWaveBasis,
+                 f_real::AbstractArray3)
+    mul!(f_fourier, pw.opFFT, f_real)
 end
-function r_to_G!(f_fourier::AbstractArray1or2D, pw::PlaneWaveBasis, kpt::Kpoint,
-                 f_real::AbstractArray3or4D)
-    n_bands = size(f_real, 4)
-    @assert size(f_real)[1:3] == pw.fft_size
-    @assert size(f_fourier, 1) == length(kpt.mapping)
-    @assert size(f_fourier, 2) == n_bands
-
-    @views Threads.@threads for iband in 1:n_bands
-        # FFT on C_ρ^∗ -> C_ρ
-        mul!(f_real[:, :, :, iband], pw.ipFFT, f_real[:, :, :, iband])
-
-        # Truncate the resulting frequency range to B_k
-        fill!(f_fourier[:, iband], 0)
-        f_fourier[:, iband] = reshape(f_real, :, n_bands)[kpt.mapping, iband]
-    end
-    f_fourier
+function r_to_G!(f_fourier::AbstractVector, pw::PlaneWaveBasis, kpt::Kpoint,
+                 f_real::AbstractArray3)
+    @assert size(f_real) == pw.fft_size
+    @assert length(f_fourier) == length(kpt.mapping)
+
+    # FFT on C_ρ^∗ -> C_ρ
+    mul!(f_real, pw.ipFFT, f_real)
+
+    # Truncate the resulting frequency range to B_k
+    fill!(f_fourier, 0)
+    f_fourier[:] = f_real[kpt.mapping]
 end
 
 @doc raw"""
@@ -276,13 +252,10 @@ end
 Perform an FFT to translate between `f_fourier`, a fourier representation
 on ``C_ρ^\ast`` and its fourier representation on ``C_ρ``.
 """
-function r_to_G(pw::PlaneWaveBasis, f_real::AbstractArray3or4D)
+function r_to_G(pw::PlaneWaveBasis, f_real::AbstractArray3)
     r_to_G!(similar(f_real), pw, f_real)
 end
 # TODO optimize this
-function r_to_G(pw::PlaneWaveBasis, kpt::Kpoint, f_real::AbstractArray{T, 3}) where {T}
+function r_to_G(pw::PlaneWaveBasis, kpt::Kpoint, f_real::AbstractArray3) where {T}
     r_to_G!(similar(f_real, length(kpt.mapping)), pw, kpt, copy(f_real))
 end
-function r_to_G(pw::PlaneWaveBasis, kpt::Kpoint, f_real::AbstractArray{T, 4}) where {T}
-    r_to_G!(similar(f_real, length(kpt.mapping), size(f_real, 4)), pw, kpt, copy(f_real))
-end

src/PotNonLocal.jl

@@ -21,7 +21,7 @@ struct PotNonLocalBlock
     kpt::Kpoint
 
     # Projection vectors and coefficients for this basis and k-Point
-    # n_Gk = length(kpoint.basis)
+    # n_Gk = length(G_vectors(kpt))
     # n_proj = ∑_atom ∑_l n_proj_per_l_for_atom * (2l + 1)
     proj_vectors  # n_proj × n_proj
     proj_coeffs   # n_Gk × n_proj

src/bzmesh.jl

@@ -44,9 +44,14 @@ function bzmesh_ir_wedge(kgrid_size, lattice, composition...; tol_symmetry=1e-5)
         "spglib failed to get the symmetries. Check your lattice, or use a uniform BZ mesh."
     )
 
-    # TODO checks
-    # - S is unitary
-    # - check that S * lattice + τ ∈ lattice
+    # Checks: S is unitary (noting that S is given in lattice coordinates)
+    for S in eachslice(spg_symops["rotations"]; dims=1)
+        Scart = lattice * S * inv(lattice)  # Form S in cartesian coords
+        if maximum(abs, Scart'Scart - I) > sqrt(eps(Float64))
+            error("spglib returned non-unitary rotation matrix")
+        end
+    end
+    # TODO check that S * lattice + τ ∈ lattice
 
     mapping, grid = spglib.get_stabilized_reciprocal_mesh(
         kgrid_size, spg_symops["rotations"], is_shift=[0, 0, 0], is_time_reversal=true
@@ -88,6 +93,10 @@ function bzmesh_ir_wedge(kgrid_size, lattice, composition...; tol_symmetry=1e-5)
         end
     end
 
+    # The symmetry operation (S == I and τ == 0) should be present for each k-Point
+    @assert all(nothing !== findfirst(Sτ -> iszero(Sτ[1] - I) && iszero(Sτ[2]), ops)
+                for ops in ksymops)
+
     kcoords, ksymops
 end