Support isolated systems

src/Model.jl

@@ -14,6 +14,10 @@ struct Model{T <: Real, VT <: Real}
     recip_lattice::Mat3{T}
     # Dimension of the system; 3 unless `lattice` has zero columns
     n_dim::Int
+    # Whether the problem is periodic along each direction
+    # When not periodic, the system is assumed to be centered
+    # at the middle of the cell (1/2)
+    periodic::Vec3{Bool}
     # Useful for conversions between cartesian and reduced coordinates
     inv_lattice::Mat3{T}
     inv_recip_lattice::Mat3{T}
@@ -66,7 +70,7 @@ _is_well_conditioned(A; tol=1e5) = (cond(A) <= tol)
 
 """
     Model(lattice, atoms, positions; n_electrons, magnetic_moments, terms, temperature,
-          smearing, spin_polarization, symmetries)
+          smearing, spin_polarization, symmetries, periodic)
 
 Creates the physical specification of a model (without any discretization information).
 
@@ -109,6 +113,7 @@ function Model(lattice::AbstractMatrix{T},
                spin_polarization=default_spin_polarization(magnetic_moments),
                symmetries=default_symmetries(lattice, atoms, positions, magnetic_moments,
                                              spin_polarization, terms),
+               periodic=Vec3(true, true, true)
                ) where {T <: Real}
     # Validate εF and n_electrons
     if !isnothing(εF)  # fixed Fermi level
@@ -176,7 +181,8 @@ function Model(lattice::AbstractMatrix{T},
     @assert !isempty(symmetries)  # Identity has to be always present.
 
     Model{T,value_type(T)}(model_name,
-                           lattice, recip_lattice, n_dim, inv_lattice, inv_recip_lattice,
+                           lattice, recip_lattice, n_dim, periodic,
+                           inv_lattice, inv_recip_lattice,
                            unit_cell_volume, recip_cell_volume,
                            n_electrons, εF, spin_polarization, n_spin, temperature, smearing,
                            atoms, positions, atom_groups, terms, symmetries)
@@ -218,6 +224,7 @@ function Model{T}(model::Model;
                   kwargs...) where {T <: Real}
     Model(T.(lattice), atoms, positions;
           model.model_name,
+          model.periodic,
           model.n_electrons,
           magnetic_moments=[],  # not used because symmetries explicitly given
           terms=model.term_types,

src/elements.jl

@@ -205,3 +205,12 @@
 function local_potential_fourier(el::ElementGaussian, q::Real)
     -el.α * exp(- (q * el.L)^2 / 2)  # = ∫_ℝ³ V(x) exp(-ix⋅q) dx
 end
+# T@D@ {
+
+function charge_real(el::Element, r)
+    zero(r)
+ end
+function charge_fourier(el::Element, q)
+    zero(q)
+end
+# }

src/fft.jl

@@ -24,8 +24,8 @@ function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, f_fourier::Abstrac
     mul!(f_real, basis.opBFFT, f_fourier)
     f_real .*= basis.ifft_normalization
 end
-function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis,
-                 kpt::Kpoint, f_fourier::AbstractVector; normalize=true)
+function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, kpt::Kpoint,
+               f_fourier::AbstractVector; normalize=true)
     @assert length(f_fourier) == length(kpt.mapping)
     @assert size(f_real) == basis.fft_size
 
@@ -79,8 +79,8 @@ function fft!(f_fourier::AbstractArray3, basis::PlaneWaveBasis, f_real::Abstract
     mul!(f_fourier, basis.opFFT, f_real)
     f_fourier .*= basis.fft_normalization
 end
-function fft!(f_fourier::AbstractVector, basis::PlaneWaveBasis,
-                 kpt::Kpoint, f_real::AbstractArray3; normalize=true)
+function fft!(f_fourier::AbstractVector, basis::PlaneWaveBasis, kpt::Kpoint,
+              f_real::AbstractArray3; normalize=true)
     @assert size(f_real) == basis.fft_size
     @assert length(f_fourier) == length(kpt.mapping)
 

src/terms/ewald.jl

@@ -20,8 +20,14 @@
 @timing "precomp: Ewald" function TermEwald(basis::PlaneWaveBasis{T};
                                             η=default_η(basis.model.lattice)) where {T}
     model = basis.model
+    periodic = model.periodic
+    if all(.! periodic)
+        @assert all(periodic) || all(.! periodic)  # only periodic or isolated atm
+        η = zero(T)  # no Ewald splitting in this case: simple Coulomb for the unit cell only
+    end
     charges = charge_ionic.(model.atoms)
-    (; energy, forces) = energy_forces_ewald(model.lattice, charges, model.positions; η)
+    (; energy, forces) = energy_forces_ewald(model.lattice, charges, model.positions; η,
+                                             periodic)
     TermEwald(energy, forces, η)
 end
 
@@ -34,17 +40,17 @@
 Standard computation of energy and forces.
 """
 function energy_forces_ewald(lattice::AbstractArray{T}, charges::AbstractArray,
-                             positions; kwargs...) where {T}
-    energy_forces_ewald(T, lattice, charges, positions, zero(Vec3{T}), nothing)
+                             positions; periodic=Vec3(true, true, true), kwargs...) where {T}
+    energy_forces_ewald(T, lattice, charges, positions, zero(Vec3{T}), nothing; periodic)
 end
 
 """
 Computation for phonons; required to build the dynamical matrix.
 """
 function energy_forces_ewald(lattice::AbstractArray{T}, charges, positions, q,
-                             ph_disp; kwargs...) where{T}
+                             ph_disp; periodic=Vec3(true, true, true), kwargs...) where{T}
     S = promote_type(complex(T), eltype(ph_disp[1]))
-    energy_forces_ewald(S, lattice, charges, positions, q, ph_disp; kwargs...)
+    energy_forces_ewald(S, lattice, charges, positions, q, ph_disp; periodic, kwargs...)
 end
 
 # To compute the electrostatics of the system, we use the Ewald splitting method due to the
@@ -80,7 +86,7 @@
 potential term if you want to customise this treatment.
 """
 function energy_forces_ewald(S, lattice::AbstractArray{T}, charges, positions, q, ph_disp;
-                             η=default_η(lattice)) where {T}
+                             η=default_η(lattice), periodic=Vec3(true, true, true)) where {T}
     @assert length(charges) == length(positions)
     if isempty(charges)
         return (; energy=zero(T), forces=zero(positions))
@@ -115,27 +121,38 @@
     poslims = [maximum(rj[i] - rk[i] for rj in positions for rk in positions) for i in 1:3]
     Rlims = estimate_integer_lattice_bounds(lattice, max_erfc_arg / η, poslims)
 
+    # Do we want shells in all directions?
+    is_dim_trivial = .!periodic
+    max_shell(n, trivial) = trivial ? 0 : n
+    Rlims = max_shell.(Rlims, is_dim_trivial)
+
     #
     # Reciprocal space sum
     #
-    # Initialize reciprocal sum with correction term for charge neutrality
-    sum_recip::S = - (sum(charges)^2 / 4η^2)
+    if all(periodic)
+        # Initialize reciprocal sum with correction term for charge neutrality
+        sum_recip::S = - (sum(charges)^2 / 4η^2)
+    else
+        sum_recip = zero(T)
+    end
     forces_recip = zeros(Vec3{S}, length(positions))
 
-    for G1 in -Glims[1]:Glims[1], G2 in -Glims[2]:Glims[2], G3 in -Glims[3]:Glims[3]
-        G = Vec3(G1, G2, G3)
-        iszero(G) && continue
-        Gsq = norm2(recip_lattice * G)
-        cos_strucfac = sum(Z * cos2pi(dot(r, G)) for (r, Z) in zip(positions, charges))
-        sin_strucfac = sum(Z * sin2pi(dot(r, G)) for (r, Z) in zip(positions, charges))
-        sum_strucfac = cos_strucfac^2 + sin_strucfac^2
-        sum_recip += sum_strucfac * exp(-Gsq / 4η^2) / Gsq
-        for (ir, r) in enumerate(positions)
-            Z = charges[ir]
-            dc = -Z*2S(π)*G*sin2pi(dot(r, G))
-            ds = +Z*2S(π)*G*cos2pi(dot(r, G))
-            dsum = cos_strucfac*dc + sin_strucfac*ds
-            forces_recip[ir] -= dsum * exp(-Gsq / 4η^2)/Gsq
+    if all(periodic)  # if non-periodic, no reciprocal-space contribution
+        for G1 in -Glims[1]:Glims[1], G2 in -Glims[2]:Glims[2], G3 in -Glims[3]:Glims[3]
+            G = Vec3(G1, G2, G3)
+            iszero(G) && continue
+            Gsq = norm2(recip_lattice * G)
+            cos_strucfac = sum(Z * cos2pi(dot(r, G)) for (r, Z) in zip(positions, charges))
+            sin_strucfac = sum(Z * sin2pi(dot(r, G)) for (r, Z) in zip(positions, charges))
+            sum_strucfac = cos_strucfac^2 + sin_strucfac^2
+            sum_recip += sum_strucfac * exp(-Gsq / 4η^2) / Gsq
+            for (ir, r) in enumerate(positions)
+                Z = charges[ir]
+                dc = -Z*2S(π)*G*sin2pi(dot(r, G))
+                ds = +Z*2S(π)*G*cos2pi(dot(r, G))
+                dsum = cos_strucfac*dc + sin_strucfac*ds
+                forces_recip[ir] -= dsum * exp(-Gsq / 4η^2)/Gsq
+            end
         end
     end
 
@@ -146,8 +163,12 @@
     #
     # Real-space sum
     #
-    # Initialize real-space sum with correction term for uniform background
-    sum_real::S = -2η / sqrt(S(π)) * sum(Z -> Z^2, charges)
+    if all(periodic)
+        # Initialize real-space sum with correction term for uniform background
+        sum_real::S = -2η / sqrt(S(π)) * sum(Z -> Z^2, charges)
+    else
+        sum_real = zero(S)
+    end
     forces_real = zeros(Vec3{S}, length(positions))
 
     for R1 in -Rlims[1]:Rlims[1], R2 in -Rlims[2]:Rlims[2], R3 in -Rlims[3]:Rlims[3]

src/terms/hartree.jl

@@ -1,3 +1,4 @@
+using SpecialFunctions
 """
 Hartree term: for a decaying potential V the energy would be
 
@@ -39,11 +40,110 @@
     TermHartree(T(scaling_factor), T(scaling_factor) .* poisson_green_coeffs)
 end
 
+# a gaussian of exponent α and integral M
+ρref_real(r::T, M=1, α=1, d=3) where {T} = M * exp(-T(1)/2 * (α*r)^2) / ((2T(π))^(T(d)/2)) * α^d
+# solution of -ΔVref = 4π ρref
+function Vref_real(r::T, M=1, α=1, d=3) where {T}
+    if d == 3
+        r == 0 && return M * 2 / sqrt(T(pi)) * α / sqrt(T(2))
+        M * erf(α/sqrt(2)*r)/r
+    elseif d == 1
+        ## TODO find a way to compute this more stably when r >> 1
+        res = M * (-2π*r*erf(α*r/sqrt(2)) - 2sqrt(2π)*exp(-α^2 * r^2/2)/α)
+        res
+    else
+        error()
+    end
+end
+
+one_hot(i) = Vec3{Bool}(j == i for j=1:3)
+∂f∂α(f, α, r) = ForwardDiff.derivative(ε -> f(r + ε * one_hot(α)), zero(eltype(r)))
+
+function get_center(basis, ρ)
+    sumρ = sum(abs.(ρ))
+    center = map(1:3) do α
+        rα = [r[α] for r in r_vectors_cart(basis)]
+        sum(rα .* abs.(ρ)) / sumρ
+    end
+    Vec3(center...)
+end
+function get_dipole(α, center, basis, ρ)
+    rα = [r[α]-center[α] for r in r_vectors_cart(basis)]
+    dot(rα, ρ) * basis.dvol
+end
+function get_variance(center, basis, ρ)
+    rr = [sum(abs2, r-center) for r in r_vectors_cart(basis)]
+    dot(rr, abs.(ρ)) * basis.dvol
+end
+
 @timing "ene_ops: hartree" function ene_ops(term::TermHartree, basis::PlaneWaveBasis{T},
                                             ψ, occupation; ρ, kwargs...) where {T}
-    ρtot_fourier = fft(basis, total_density(ρ))
+    model = basis.model
+    ρtot_real = total_density(ρ)
+    # T@D@ {
+    ρtot_real .-= [sum(DFTK.charge_real(el, norm(r - vector_red_to_cart(model, pos)))
+                       for (el, pos) in zip(model.atoms, model.positions))
+                   for r in r_vectors_cart(basis)]
+    # println(sum(ρtot_real) * basis.dvol)
+    # }
+    ρtot_fourier = fft(basis, ρtot_real)
     pot_fourier = term.poisson_green_coeffs .* ρtot_fourier
     pot_real = irfft(basis, pot_fourier)
+
+    # For isolated systems, the above does not compute a good potential (e.g., it assumes
+    # zero DC component).
+    # We correct it by solving -Δ V = 4πρ in two steps: we split ρ into ρref and ρ-ρref,
+    # where ρref is a Gaussian has the same total charge as ρ.
+    # We compute the first potential in real space (explicitly since ρref is known), and the
+    # second (short-range) potential in Fourier space.
+    # Compared to the usual scheme (ρref = 0), this results in a correction potential equal
+    # to Vref computed in real space minus Vref computed in Fourier space
+    if any(.!model.periodic)
+        @assert all(.!model.periodic)
+        # Determine center and width from density.
+        # Strictly speaking, these computations should result in extra terms to guarantee
+        # energy/ham consistency.
+        center = get_center(basis, ρtot_real)
+        ρref_11 = [ρref_real(norm(r - center), 1, 1, basis.model.n_dim) for r in r_vectors_cart(basis)]
+        spread_11 = get_variance(center, basis, ρref_11)
+        spread_ρ = get_variance(center, basis, ρtot_real)
+        α = 1/sqrt(spread_ρ/spread_11)
+
+        total_charge = sum(ρtot_real) .* basis.dvol
+        ρrad_fun(r) = ρref_real(norm(r - center), 1, α, basis.model.n_dim)
+        ρrad = ρrad_fun.(r_vectors_cart(basis))
+
+
+        # At this point we've determined a gaussian of same width and center as the original.
+        # Now we get the dipole moments of ρ and match them.
+        ρders = [∂f∂α.(ρrad_fun, α, r_vectors_cart(basis)) for α=1:3]
+
+        coeffs_ders = map(1:3) do α
+            get_dipole(α, center, basis, ρtot_real) / get_dipole(α, center, basis, ρders[α])
+        end
+        # T@D@ {
+        # println(coeffs_ders)
+        if basis.model.n_dim == 1
+            coeffs_ders[2:3] .= 0
+        end
+        # }
+        ρref = total_charge * ρrad + sum([coeffs_ders[α]*ρders[α] for α=1:3])
+
+        # Compute corresponding solution of -ΔVref = 4π ρref.
+        Vref_rad_fun(r) = Vref_real(norm(r - center), 1, α, basis.model.n_dim)
+        Vref_rad = Vref_rad_fun.(r_vectors_cart(basis))
+        Vref_ders = [∂f∂α.(Vref_rad_fun, α, r_vectors_cart(basis)) for α=1:3]
+        Vref = total_charge * Vref_rad + sum([coeffs_ders[α]*Vref_ders[α] for α=1:3])
+
+        # TODO possibly optimize FFTs here
+        Vcorr_real = Vref - irfft(basis, term.poisson_green_coeffs .* fft(basis, ρref))
+        Vcorr_fourier = fft(basis, Vcorr_real)
+        pot_real .+= Vcorr_real
+        pot_fourier .+= Vcorr_fourier
+    end
+
+    pot_fourier *= term.scaling_factor
+    pot_real *= term.scaling_factor
     E = real(dot(pot_fourier, ρtot_fourier) / 2)
 
     ops = [RealSpaceMultiplication(basis, kpt, pot_real) for kpt in basis.kpoints]

src/terms/local.jl

@@ -83,43 +83,53 @@
 """
 struct AtomicLocal end
 function (::AtomicLocal)(basis::PlaneWaveBasis{T}) where {T}
-    # pot_fourier is <e_G|V|e_G'> expanded in a basis of e_{G-G'}
-    # Since V is a sum of radial functions located at atomic
-    # positions, this involves a form factor (`local_potential_fourier`)
-    # and a structure factor e^{-i G·r}
     model = basis.model
-    G_cart = to_cpu(G_vectors_cart(basis))
-    # TODO Bring G_cart on the CPU for compatibility with the pseudopotentials which
-    #      are not isbits ... might be able to solve this by restructuring the loop
-
-
-    # Pre-compute the form factors at unique values of |G| to speed up
-    # the potential Fourier transform (by a lot). Using a hash map gives O(1)
-    # lookup.
-    form_factors = IdDict{Tuple{Int,T},T}()  # IdDict for Dual compatability
-    for G in G_cart
-        q = norm(G)
-        for (igroup, group) in enumerate(model.atom_groups)
-            if !haskey(form_factors, (igroup, q))
-                element = model.atoms[first(group)]
-                form_factors[(igroup, q)] = local_potential_fourier(element, q)
+
+    if all(model.periodic)
+        # pot_fourier is <e_G|V|e_G'> expanded in a basis of e_{G-G'}
+        # Since V is a sum of radial functions located at atomic
+        # positions, this involves a form factor (`local_potential_fourier`)
+        # and a structure factor e^{-i G·r}
+        G_cart = to_cpu(G_vectors_cart(basis))
+        # TODO Bring G_cart on the CPU for compatibility with the pseudopotentials which
+        #      are not isbits ... might be able to solve this by restructuring the loop
+
+
+        # Pre-compute the form factors at unique values of |G| to speed up
+        # the potential Fourier transform (by a lot). Using a hash map gives O(1)
+        # lookup.
+        form_factors = IdDict{Tuple{Int,T},T}()  # IdDict for Dual compatability
+        for G in G_cart
+            q = norm(G)
+            for (igroup, group) in enumerate(model.atom_groups)
+                if !haskey(form_factors, (igroup, q))
+                    element = model.atoms[first(group)]
+                    form_factors[(igroup, q)] = local_potential_fourier(element, q)
+                end
             end
         end
-    end
 
-    Gs = to_cpu(G_vectors(basis))  # TODO Again for GPU compatibility
-    pot_fourier = map(enumerate(Gs)) do (iG, G)
-        q = norm(G_cart[iG])
-        pot = sum(enumerate(model.atom_groups)) do (igroup, group)
-            structure_factor = sum(r -> cis2pi(-dot(G, r)), @view model.positions[group])
-            form_factors[(igroup, q)] * structure_factor
+        Gs = to_cpu(G_vectors(basis))  # TODO Again for GPU compatibility
+        pot_fourier = map(enumerate(Gs)) do (iG, G)
+            q = norm(G_cart[iG])
+            pot = sum(enumerate(model.atom_groups)) do (igroup, group)
+                structure_factor = sum(r -> cis2pi(-dot(G, r)), @view model.positions[group])
+                form_factors[(igroup, q)] * structure_factor
+            end
+            pot / sqrt(model.unit_cell_volume)
         end
-        pot / sqrt(model.unit_cell_volume)
-    end
-
-    enforce_real!(basis, pot_fourier)  # Symmetrize Fourier coeffs to have real iFFT
-    pot_real = irfft(basis, to_device(basis.architecture, pot_fourier))
 
+        enforce_real!(basis, pot_fourier)  # Symmetrize Fourier coeffs to have real iFFT
+        pot_real = irfft(basis, to_device(basis.architecture, pot_fourier))
+    else
+        @assert all(.!model.periodic)
+        # simple real-space sum
+        pot_real = map(r_vectors(basis)) do r
+            sum(zip(model.atoms, model.positions)) do (element, r_ion)
+                local_potential_real(element, norm(model.lattice * (r-r_ion)))
+            end
+        end
+    end
     TermAtomicLocal(pot_real)
 end