Ultrasoft pseudopotentials

Project.toml

@@ -13,6 +13,7 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DftFunctionals = "6bd331d2-b28d-4fd3-880e-1a1c7f37947f"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+EzXML = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
@@ -44,6 +45,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
+WignerSymbols = "9f57e263-0b3d-5e2e-b1be-24f2bb48858b"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
@@ -90,6 +92,7 @@ DftFunctionals = "0.3"
 DifferentiationInterface = "0.6.39, 0.7"
 DocStringExtensions = "0.9"
 DoubleFloats = "1"
+EzXML = "1.2.1"
 FFTW = "1.5"
 FiniteDiff = "2"
 FiniteDifferences = "0.12"
@@ -135,6 +138,7 @@ TimerOutputs = "0.5.29"
 Unitful = "1"
 UnitfulAtomic = "1"
 Wannier = "0.3.2"
+WignerSymbols = "2.0.0"
 WriteVTK = "1"
 julia = "1.10"
 wannier90_jll = "3.1"

src/DFTK.jl

@@ -88,6 +88,7 @@ include("Kpoint.jl")
 include("PlaneWaveBasis.jl")
 include("orbitals.jl")
 include("input_output.jl")
+include("ultrasoft.jl")
 
 export create_supercell
 export cell_to_supercell

src/PlaneWaveBasis.jl

@@ -94,6 +94,8 @@ struct PlaneWaveBasis{T,
 
     ## Instantiated terms (<: Term). See Hamiltonian for high-level usage
     terms::Vector{Any}
+
+    augmentation_regions
 end
 
 
@@ -241,6 +243,8 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Real, fft_size::Tuple{Int, Int, I
     dvol  = model.unit_cell_volume ./ prod(fft_size)
     terms = Vector{Any}(undef, length(model.term_types))  # Dummy terms array, filled below
 
+    augmentation_regions = precompute_augmentation_regions(model, fft_grid)
+
     basis = PlaneWaveBasis{T, value_type(T), Arch, typeof(fft_grid),
                            typeof(kpoints[1].G_vectors)}(
         model, fft_size, dvol,
@@ -250,7 +254,7 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Real, fft_size::Tuple{Int, Int, I
         kcoords_global, kweights_global, n_irreducible_kpoints,
         comm_kpts, krange_thisproc, krange_allprocs, krange_thisproc_allspin,
         architecture, symmetries, symmetries_respect_rgrid,
-        use_symmetries_for_kpoint_reduction, terms)
+        use_symmetries_for_kpoint_reduction, terms, augmentation_regions)
 
     # Instantiate the terms with the basis
     for (it, t) in enumerate(model.term_types)

src/common/spherical_harmonics.jl

@@ -45,5 +45,57 @@ function ylm_real(l::Integer, m::Integer, rvec::AbstractVector{T}) where {T}
         (m ==  3) && return sqrt( 35 / 32T(π)) * (x^2 - 3y^2) * x / r^3
     end
 
+    if l == 4
+        (m == -4) && return sqrt( 315 / 16T(π)) * x * y * (x^2 - y^2) / r^4
+        (m == -3) && return sqrt( 315 / 32T(π)) * y * (3x^2 - y^2) * z / r^4
+        (m == -2) && return sqrt(  45 / 16T(π)) * x * y * (7z^2 - r^2) / r^4
+        (m == -1) && return sqrt(  45 / 32T(π)) * y * (7z^3 - 3*z*r^2) / r^4
+        (m ==  0) && return sqrt(   9 /256T(π)) * (35z^4 - 30*z^2*r^2 + 3*r^4) / r^4
+        (m ==  1) && return sqrt(  45 / 32T(π)) * x * (7z^3 - 3*z*r^2) / r^4
+        (m ==  2) && return sqrt(  45 / 64T(π)) * (x^2 - y^2) * (7z^2 - r^2) / r^4
+        (m ==  3) && return sqrt( 315 / 32T(π)) * x * (x^2 - 3y^2) * z / r^4
+        (m ==  4) && return sqrt( 315 /256T(π)) * (x^2 * (x^2 - 3y^2) - y^2 * (3x^2 - y^2)) / r^4
+    end
+
     error("The case l = $l and m = $m is not implemented")
 end
+
+function pack_lm(l::Integer, m::Integer)
+    @assert 0 ≤ l
+    @assert -l ≤ m ≤ l
+    return l^2 + l + m + 1
+end
+
+# TODO: check this
+function unpack_lm(lm::Integer)
+    @assert lm ≥ 1
+    l = floor(Int, sqrt(lm - 1))
+    m = lm - (l^2 + l + 1)
+    return l, m
+end
+
+function gaunt_coefficients(T::Type, l::Integer)
+    nharm = (2l+1)^2
+    # Sample points randomly on a sphere
+    rng = Xoshiro(42)
+    points = rand(rng, T, 3, nharm)
+    points ./= sqrt.(sum(points.^2, dims=1))
+    # Evaluate spherical harmonics at these points
+    Y = zeros(T, nharm, nharm)
+    for j in 1:nharm
+        for i in 1:nharm
+            Y[i, j] = ylm_real(unpack_lm(j)..., points[:, i])
+        end
+    end
+    fact = factorize(Y)
+
+    coefs = zeros(T, nharm, (l+1)^2, (l+1)^2)
+
+    for j in 1:(l+1)^2
+        for i in 1:(l+1)^2
+            coefs[:, i, j] = fact \ (Y[:, i] .* Y[:, j])
+        end
+    end
+
+    coefs
+end

src/densities.jl

@@ -42,7 +42,51 @@ using an optional `occupation_threshold`. By default all occupation numbers are
     end
     ρ = sum(getfield.(storages, :ρ))
 
+    # Hacky augmentation charges
+    ρ_augmentation = zero(ρ)
+    nonloc_ops = AtomicNonlocal()(basis).ops
+    aug_coeffs = zeros(Complex{T}, 36, 36)
+    for (ik, n) in range
+        P = nonloc_ops[ik].P
+
+        projections = P' * ψ[ik][:, n]
+        for iatom in 1:length(basis.model.positions)
+            # TODO: only loop over iproj:18?
+            for iproj=1:18, jproj=1:18
+                coeff = projections[(iatom-1) * 18 + iproj]' *
+                        projections[(iatom-1) * 18 + jproj]
+                aug_coeffs[(iatom-1) * 18 + iproj, (iatom-1) * 18 + jproj] +=
+                                (real(coeff) * occupation[ik][n] # TODO: is the real call correct?
+                                * basis.kweights[ik])
+            end
+        end
+    end
+
+    # TODO: hardcoded for the Si psp I have
+    psp = basis.model.atoms[1].psp
+    # proj_indices = [1, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6]
+    # proj_to_m = [0, 0, -1, 0, 1, -1, 0, 1, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2]
+    proj_indices = [1, 2, 3, 4, 3, 4, 3, 4, 5, 6, 5, 6, 5, 6, 5, 6, 5, 6]
+    proj_to_m = [0, 0, -1, -1, 0, 0, 1, 1, -2, -2, -1, -1, 0, 0, 1, 1, 2, 2]
+
+    for iatom in 1:length(basis.model.positions)
+        box = basis.augmentation_regions[iatom]
+        @info "Looping over grid for atom $iatom: box size $(length(box.grid_indices))"
+        for (xyz, Q) in zip(box.grid_indices, eachslice(box.augmentations; dims=1))
+            tot = zero(T)
+            # TODO: only loop over iproj:18?
+            for iproj=1:18, jproj=1:18
+                tot += aug_coeffs[(iatom-1) * 18 + iproj, (iatom-1) * 18 + jproj] * Q[iproj, jproj]
+            end
+            # TODO: assumes 1 spin channel
+            ρ_augmentation[xyz..., 1] += tot
+        end
+    end
+
+    ρ += ρ_augmentation
+
     mpi_sum!(ρ, basis.comm_kpts)
+    # TODO: is this fine to do after the charge augmentation?
     ρ = symmetrize_ρ(basis, ρ; do_lowpass=false)
 
     # There can always be small negative densities, e.g. due to numerical fluctuations

src/eigen/diag.jl

@@ -13,6 +13,28 @@ function diagonalize_all_kblocks(eigensolver, ham::Hamiltonian, nev_per_kpoint::
     kpoints = ham.basis.kpoints
     results = Vector{Any}(undef, length(kpoints))
 
+    # TODO: temporary plumbing :P
+    B = [I for _ in kpoints]
+    if any(ham.basis.model.atoms) do el el isa ElementPsp && el.psp isa PspUpf && el.psp.type == "US" end
+        model = ham.basis.model
+
+        # keep only pseudopotential atoms and positions
+        psp_groups = [group for group in model.atom_groups
+                        if model.atoms[first(group)] isa ElementPsp]
+        psps          = [model.atoms[first(group)].psp for group in psp_groups]
+        psp_positions = [model.positions[group] for group in psp_groups]
+
+        # We have an overlap term
+        B = map(ham.blocks) do block
+            op = block.nonlocal_op
+            Q = build_projection_coefficients(eltype(ham.basis), psps, psp_positions, "Q")
+            LinearMap(size(op.P, 1)) do Sψ, ψ
+                mul!(Sψ, op.P, (Q * (op.P' * ψ)), 1, 0)
+                Sψ .+= ψ # + identity
+            end
+        end
+    end
+
     for (ik, kpt) in enumerate(kpoints)
         n_Gk = length(G_vectors(ham.basis, kpt))
         if n_Gk < nev_per_kpoint
@@ -48,7 +70,7 @@ function diagonalize_all_kblocks(eigensolver, ham::Hamiltonian, nev_per_kpoint::
         prec = nothing
         !isnothing(prec_type) && (prec = prec_type(ham[ik]))
         results[ik] = eigensolver(ham[ik], ψguessk;
-                                  prec, tol, miniter, maxiter, n_conv_check)
+                                  B=B[ik], prec, tol, miniter, maxiter, n_conv_check)
     end
 
     # Transform results into a nicer datastructure

src/eigen/diag_lobpcg_hyper.jl

@@ -4,11 +4,11 @@ include("lobpcg_hyper_impl.jl")
 # but X and the history on the device (for GPU runs)
 function lobpcg_hyper(A, X0; maxiter=100, prec=nothing,
                       tol=20size(A, 2)*eps(real(eltype(A))),
-                      largest=false, n_conv_check=nothing, kwargs...)
+                      largest=false, n_conv_check=nothing, B=I, kwargs...)
     prec === nothing && (prec = I)
 
     @assert !largest "Only seeking the smallest eigenpairs is implemented."
-    result = LOBPCG(A, X0, I, prec, tol, maxiter; n_conv_check, kwargs...)
+    result = LOBPCG(A, X0, B, prec, tol, maxiter; n_conv_check, kwargs...)
 
     n_conv_check === nothing && (n_conv_check = size(X0, 2))
     converged = maximum(result.residual_norms[1:n_conv_check]) < tol