[WIP] Additional exx kernels: Voxel averaged and Wigner-Seitz truncation

Project.toml

@@ -47,6 +47,7 @@ UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a"
 GeometryOptimization = "673bf261-a53d-43b9-876f-d3c1fc8329c2"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -59,6 +60,7 @@ wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 [extensions]
 DFTKAMDGPUExt = ["MPI", "AMDGPU"]
 DFTKCUDAExt = ["Libxc", "MPI", "CUDA"]
+DFTKFastGaussQuadratureExt = "FastGaussQuadrature"
 DFTKGenericLinearAlgebraExt = "GenericLinearAlgebra"
 DFTKGeometryOptimizationExt = "GeometryOptimization"
 DFTKJLD2Ext = "JLD2"
@@ -88,6 +90,7 @@ DifferentiationInterface = "0.6.39, 0.7"
 DocStringExtensions = "0.9"
 DoubleFloats = "1"
 FFTW = "1.5"
+FastGaussQuadrature = "1.1.0"
 FiniteDiff = "2"
 FiniteDifferences = "0.12"
 ForwardDiff = "1"
@@ -163,4 +166,4 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 
 [targets]
-test = ["Test", "TestItemRunner", "AMDGPU", "ASEconvert", "AtomsBuilder", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "GeometryOptimization", "JLD2", "JSON3", "Logging", "Plots", "PythonCall", "QuadGK", "Random", "KrylovKit", "Wannier", "WriteVTK", "wannier90_jll"]
+test = ["Test", "TestItemRunner", "AMDGPU", "ASEconvert", "AtomsBuilder", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FastGaussQuadrature", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "GeometryOptimization", "JLD2", "JSON3", "Logging", "Plots", "PythonCall", "QuadGK", "Random", "KrylovKit", "Wannier", "WriteVTK", "wannier90_jll"]

ext/DFTKFastGaussQuadratureExt.jl

@@ -0,0 +1,116 @@
+module DFTKFastGaussQuadratureExt
+using FastGaussQuadrature
+using DFTK
+using DFTK: to_cpu, eval_kernel_fourier, VoxelAveraged
+using LinearAlgebra
+
+
+@views function DFTK._compute_kernel_fourier(kernel, regularization::VoxelAveraged,
+                                             basis::PlaneWaveBasis{T}, qpt, q) where {T}
+    model = basis.model
+    q = qpt.coordinate
+
+    # Get kgrid_size
+    if isnothing(basis.kgrid)
+        kgrid_size = Vec3{Int}(1, 1, 1)
+    elseif basis.kgrid isa AbstractVector
+        kgrid_size = Vec3{Int}(basis.kgrid)
+    elseif basis.kgrid isa MonkhorstPack
+        kgrid_size = Vec3{Int}(basis.kgrid.kgrid_size)
+    else
+        @error "Cannot determine kgrid_size for VoxelAveraged Coulomb model."
+    end
+
+    # Define Voxels as reciprocal cell deivided by k-mesh
+    voxel_basis = model.recip_lattice * Diagonal(1 ./ Vec3{T}(kgrid_size))
+    voxel_vol = abs(det(voxel_basis))
+
+    # get Gauss-Legendre nodes and weights
+    nodes_std, weights_std = gausslegendre(regularization.n_quadrature_points)
+    # Scale from [-1, 1] to [-0.5, 0.5]
+    nodes = T.(nodes_std ./ 2)
+    weights = T.(weights_std ./ 2)
+
+    kernel_fourier = zeros(T, length(qpt.G_vectors))
+
+    for (iG, G) in enumerate(to_cpu(qpt.G_vectors))
+        G_cart = model.recip_lattice * (G+q)
+
+        found_singularity = (iG==1 && iszero(q))
+
+        if found_singularity 
+            # === At Singularity (G+q=0) use surface reduction method ===
+
+            # We only do that for the 4π/(G+q)^2 kernel, hence the quadrature below
+            # covers the rest if kernel_fourier is different from Coulomb().
+
+            # Transforms volume integral ∫ 1/G^2 dV to surface integral Σ h * ∫ 1/G^2 dS
+            integral = zero(T)
+            for i in 1:3
+                u_i = voxel_basis[:, i]
+                u_j = voxel_basis[:, mod1(i+1, 3)]
+                u_k = voxel_basis[:, mod1(i+2, 3)]
+
+                # Height of the face from origin
+                normal = cross(u_j, u_k)
+                area_norm = norm(normal)
+
+                # Calculate distance h from center to face.
+                # Since face is at u_i/2, h = |(u_i/2) . n|
+                h = abs(dot(u_i, normal)) / (2 * area_norm)
+
+                # Integrate 1/G^2 over the face parallelogram
+                face_integral = zero(T)
+                for (wa, a) in zip(weights, nodes)
+                    for (wb, b) in zip(weights, nodes)
+                        # Parametrization of the face: r = u_i/2 + a*u_j + b*u_k
+                        r_vec = 0.5f0 * u_i + a * u_j + b * u_k 
+                        r_sq  = dot(r_vec, r_vec)
+                        face_integral += wa * wb / r_sq
+                    end
+                end
+
+                # Add contribution: 2 faces * h * Area * Mean(1/r^2)
+                # Area factor (area_norm) comes from the Jacobian.
+                integral += 2 * h * area_norm * face_integral
+            end
+
+            kernel_fourier[iG] = 4T(π) * (integral / voxel_vol)
+        end
+
+        # === Use smooth 3D Gaussian Quadrature ===
+        integral = zero(T)
+        for (wx, x) in zip(weights, nodes)
+            for (wy, y) in zip(weights, nodes)
+                for (wz, z) in zip(weights, nodes)
+                    # q vector inside voxel
+                    q_local = x * voxel_basis[:, 1] + 
+                              y * voxel_basis[:, 2] + 
+                              z * voxel_basis[:, 3]
+
+                    G_total = G_cart + q_local
+                    Gsq = dot(G_total, G_total)
+
+                    # switch temporarily to BigFloat
+                    Gsq_big = BigFloat(Gsq)
+                    val_big = eval_kernel_fourier(kernel, Gsq_big)
+
+                    # At singularity, already captured the 4π/(G+q)^2 contribution above: subtract
+                    if found_singularity
+                        val_big -= 4π/Gsq_big
+                    end
+
+                    # back to type T
+                    val = T(val_big)
+
+                    integral += wx * wy * wz * val
+                end
+            end
+        end
+        kernel_fourier[iG] += integral
+    end
+
+    kernel_fourier
+end
+
+end

src/DFTK.jl

@@ -101,9 +101,11 @@ include("supercell.jl")
 export Energies
 include("Energies.jl")
 
-export Coulomb, SphericallyTruncatedCoulomb
+export Coulomb
+export SphericallyTruncatedCoulomb
+export WignerSeitzTruncatedCoulomb
 export ShortRangeCoulomb, LongRangeCoulomb
-export ProbeCharge, ReplaceSingularity
+export ProbeCharge, ReplaceSingularity, VoxelAveraged
 include("coulomb.jl")
 
 export Hamiltonian