Implement kgrid_from_minimal_kpoints (#523)

examples/metallic_systems.jl

@@ -37,7 +37,7 @@ temperature = 0.01              # Smearing temperature in Hartree
 model = model_DFT(lattice, atoms, [:gga_x_pbe, :gga_c_pbe];
                   temperature=temperature,
                   smearing=DFTK.Smearing.FermiDirac())
-kgrid = kgrid_size_from_minimal_spacing(lattice, kspacing)
+kgrid = kgrid_from_minimal_spacing(lattice, kspacing)
 basis = PlaneWaveBasis(model; Ecut, kgrid);
 
 # Finally we run the SCF. Two magnesium atoms in

src/DFTK.jl

@@ -129,7 +129,7 @@ export symmetry_operations
 export standardize_atoms
 export bzmesh_uniform
 export bzmesh_ir_wedge
-export kgrid_size_from_minimal_spacing
+export kgrid_from_minimal_spacing, kgrid_from_minimal_n_kpoints
 include("symmetry.jl")
 include("bzmesh.jl")
 

src/Model.jl

@@ -97,7 +97,7 @@ function Model(lattice::AbstractMatrix{T};
     end
 
     # Special handling of 1D and 2D systems, and sanity checks
-    n_dim = 3 - count(iszero, eachcol(lattice))
+    n_dim = count(!iszero, eachcol(lattice))
     n_dim > 0 || error("Check your lattice; we do not do 0D systems")
     for i = n_dim+1:3
         norm(lattice[:, i]) == norm(lattice[i, :]) == 0 || error(
@@ -106,15 +106,10 @@ function Model(lattice::AbstractMatrix{T};
     _is_well_conditioned(lattice[1:n_dim, 1:n_dim]) || @warn (
         "Your lattice is badly conditioned, the computation is likely to fail.")
 
-    # Compute reciprocal lattice and volumes.
-    # recall that the reciprocal lattice is the set of G vectors such
-    # that G.R ∈ 2π ℤ for all R in the lattice
-    recip_lattice = zeros(T, 3, 3)
-    recip_lattice[1:n_dim, 1:n_dim] = 2T(π)*inv(lattice[1:n_dim, 1:n_dim]')
-    recip_lattice = Mat3{T}(recip_lattice)
-    # in the 1D or 2D case, the volume is the length/surface
-    unit_cell_volume  = abs(det(lattice[1:n_dim, 1:n_dim]))
-    recip_cell_volume = abs(det(recip_lattice[1:n_dim, 1:n_dim]))
+    # Note: In the 1D or 2D case, the volume is the length/surface
+    recip_lattice = compute_recip_lattice(lattice)
+    unit_cell_volume  = compute_unit_cell_volume(lattice)
+    recip_cell_volume = compute_unit_cell_volume(recip_lattice)
 
     spin_polarization in (:none, :collinear, :full, :spinless) ||
         error("Only :none, :collinear, :full and :spinless allowed for spin_polarization")
@@ -171,7 +166,7 @@ Default logic to determine the symmetry operations to be used in the model.
 """
 function default_symmetries(lattice, atoms, magnetic_moments, terms, spin_polarization;
                         tol_symmetry=1e-5)
-    dimension = 3 - count(iszero, eachcol(lattice))
+    dimension = count(!iszero, eachcol(lattice))
     if spin_polarization == :full || dimension != 3
         return [identity_symop()]  # Symmetry not supported in spglib
     elseif spin_polarization == :collinear && isempty(magnetic_moments)
@@ -215,3 +210,33 @@ end
 spin_components(model::Model) = spin_components(model.spin_polarization)
 
 _is_well_conditioned(A; tol=1e5) = (cond(A) <= tol)
+
+
+"""
+Compute the reciprocal lattice. Takes special care of 1D or 2D cases.
+We use the convention that the reciprocal lattice is the set of G vectors such
+that G ⋅ R ∈ 2π ℤ for all R in the lattice.
+"""
+function compute_recip_lattice(lattice::AbstractMatrix{T}) where {T}
+    # Note: pinv pretty much does the same, but the implied SVD causes trouble
+    #       with interval arithmetic and dual numbers, so we go for this version.
+    n_dim = count(!iszero, eachcol(lattice))
+    @assert 1 ≤ n_dim ≤ 3
+    if n_dim == 3
+        2T(π) * inv(lattice')
+    else
+        2T(π) * Mat3{T}([
+            inv(lattice[1:n_dim, 1:n_dim]')   zeros(T, n_dim, 3 - n_dim);
+            zeros(T, 3 - n_dim, 3)
+        ])
+    end
+end
+
+
+"""
+Compute unit cell volume volume. In case of 1D or 2D case, the volume is the length/surface.
+"""
+function compute_unit_cell_volume(lattice)
+    n_dim = count(!iszero, eachcol(lattice))
+    abs(det(lattice[1:n_dim, 1:n_dim]))
+end

src/PlaneWaveBasis.jl

@@ -306,7 +306,7 @@ end
 Creates a `PlaneWaveBasis` using the kinetic energy cutoff `Ecut` and a Monkhorst-Pack
 kpoint grid. The MP grid can either be specified directly with `kgrid` providing the
 number of points in each dimension and `kshift` the shift (0 or 1/2 in each direction).
-If not specified a grid is generated using `kgrid_size_from_minimal_spacing` with
+If not specified a grid is generated using `kgrid_from_minimal_spacing` with
 a minimal spacing of `2π * 0.022` per Bohr.
 
 If `use_symmetry` is `true` (default) the symmetries of the
@@ -317,7 +317,7 @@ undefined.
 """
 function PlaneWaveBasis(model::Model;
                         Ecut,
-                        kgrid=kgrid_size_from_minimal_spacing(model.lattice, 2π * 0.022),
+                        kgrid=kgrid_from_minimal_spacing(model, 2π * 0.022),
                         kshift=[iseven(nk) ? 1/2 : 0 for nk in kgrid],
                         use_symmetry=true,
                         kwargs...)

src/bzmesh.jl

@@ -145,20 +145,64 @@ const standardize_atoms = spglib_standardize_cell
 # TODO Maybe maximal spacing is actually a better name as the kpoints are spaced
 #      at most that far apart
 @doc raw"""
-Selects a kgrid_size to ensure a minimal spacing (in inverse Bohrs) between kpoints.
-Default is ``2π * 0.04 \AA^{-1}``.
+Selects a kgrid size to ensure a minimal spacing (in inverse Bohrs) between kpoints.
+A reasonable spacing is `0.25` inverse Bohrs (around ``2π * 0.04 \AA^{-1}``).
 """
-function kgrid_size_from_minimal_spacing(lattice, spacing=2π * 0.022)
-    lattice = austrip.(lattice)
-    spacing = austrip(spacing)
+function kgrid_from_minimal_spacing(lattice, spacing)
+    lattice       = austrip.(lattice)
+    spacing       = austrip(spacing)
+    recip_lattice = compute_recip_lattice(lattice)
     @assert spacing > 0
     isinf(spacing) && return [1, 1, 1]
 
-    for d in 1:3
-        @assert norm(lattice[:, d]) != 0
-        # Otherwise the formula for the reciprocal lattice
-        # computation is not correct
+    [max(1, ceil(Int, norm(vec) / spacing)) for vec in eachcol(recip_lattice)]
+end
+function kgrid_from_minimal_spacing(model::Model, args...)
+    kgrid_from_minimal_spacing(model.lattice, args...)
+end
+
+@doc raw"""
+Selects a kgrid size which ensures that at least a `n_kpoints` total number of ``k``-Points
+are used. The distribution of ``k``-Points amongst coordinate directions is as uniformly
+as possible, trying to achieve an identical minimal spacing in all directions.
+"""
+function kgrid_from_minimal_n_kpoints(lattice, n_kpoints::Integer)
+    lattice = austrip.(lattice)
+    n_dim   = count(!iszero, eachcol(lattice))
+    @assert n_kpoints > 0
+    n_kpoints == 1 && return [1, 1, 1]
+    n_dim == 1 && return [n_kpoints, 1, 1]
+
+    # Compute truncated reciprocal lattice
+    recip_lattice_nD = 2π * inv(lattice[1:n_dim, 1:n_dim]')
+    n_kpt_per_dim = n_kpoints^(1/n_dim)
+
+    # Start from a cubic lattice. If it is one, we are done. Otherwise the resulting
+    # spacings in each dimension bracket the ideal kpoint spacing.
+    spacing_per_dim = [norm(vec) / n_kpt_per_dim for vec in eachcol(recip_lattice_nD)]
+    min_spacing, max_spacing = extrema(spacing_per_dim)
+    if min_spacing ≈ max_spacing
+        return kgrid_from_minimal_spacing(lattice, min_spacing)
+    else
+        number_of_kpoints(spacing) = prod(vec -> norm(vec) / spacing, eachcol(recip_lattice_nD))
+        @assert number_of_kpoints(min_spacing) + 0.05 ≥ n_kpoints
+        @assert number_of_kpoints(max_spacing) - 0.05 ≤ n_kpoints
+
+        # TODO This is not optimal and sometimes finds too large grids
+        spacing = Roots.find_zero(sp -> number_of_kpoints(sp) - n_kpoints,
+                                  (min_spacing, max_spacing), Roots.Bisection(),
+                                  xatol=1e-4, atol=0, rtol=0)
+
+        # Sanity check: Sometimes root finding is just across the edge towards
+        # a larger number of k-Points than needed. This attempts a slightly larger spacing.
+        kgrid_larger = kgrid_from_minimal_spacing(lattice, spacing + 1e-4)
+        if prod(kgrid_larger) ≥ n_kpoints
+            return kgrid_larger
+        else
+            return kgrid_from_minimal_spacing(lattice, spacing)
+        end
     end
-    recip_lattice = 2π * inv(lattice')
-    [ceil(Int, norm(recip_lattice[:, i]) ./ spacing) for i = 1:3]
+end
+function kgrid_from_minimal_kpoints(model::Model, args...)
+    kgrid_from_minimal_kpoints(model.lattice, args...)
 end

src/fft.jl

@@ -62,7 +62,7 @@ end
 #      It should be merged with build_kpoints somehow
 function compute_fft_size_precise(lattice::AbstractMatrix{T}, Ecut, kpoints;
                                   supersampling=2, ensure_smallprimes=true) where T
-    recip_lattice = 2T(π)*pinv(lattice')  # pinv in case one of the dimension is trivial
+    recip_lattice  = compute_recip_lattice(lattice)
     recip_diameter = diameter(recip_lattice)
     Glims = [0, 0, 0]
     # get the bounding rectangle that contains all G-G' vectors

src/terms/ewald.jl

@@ -75,7 +75,7 @@ function energy_ewald(lattice, charges, positions; η=nothing, forces=nothing)
             return T(0)
         end
     end
-    energy_ewald(lattice, T(2π) * inv(lattice'), charges, positions; η=η, forces=forces)
+    energy_ewald(lattice, compute_recip_lattice(lattice), charges, positions; η, forces)
 end
 
 function energy_ewald(lattice, recip_lattice, charges, positions; η=nothing, forces=nothing)
@@ -162,9 +162,9 @@ function energy_ewald(lattice, recip_lattice, charges, positions; η=nothing, fo
         gsh += 1
     end
     # Amend sum_recip by proper scaling factors:
-    sum_recip *= 4T(π) / abs(det(lattice))
+    sum_recip *= 4T(π) / compute_unit_cell_volume(lattice)
     if forces !== nothing
-        forces_recip .*= 4T(π) / abs(det(lattice))
+        forces_recip .*= 4T(π) / compute_unit_cell_volume(lattice)
     end
 
     #

src/terms/psp_correction.jl

@@ -44,5 +44,5 @@ function energy_psp_correction(lattice, atoms)
         if attype isa ElementPsp
     )
 
-    correction_per_cell / abs(det(lattice))
+    correction_per_cell / compute_unit_cell_volume(lattice)
 end

test/bzmesh.jl

@@ -108,8 +108,22 @@ end
     @test atoms[1][2][1] - atoms[1][2][2] == ones(3) ./ 4
 end
 
-@testset "kgrid_size_from_minimal_spacing" begin
+@testset "kgrid_from_minimal_spacing" begin
     # Test that units are stripped from both the lattice and the spacing
-    lattice = [[-1 1 1]; [1 -1  1]; [1 1 -1]]
-    @test kgrid_size_from_minimal_spacing(lattice * u"angstrom", 0.5 / u"angstrom") == [9; 9; 9]
+    lattice = [[-1.0 1 1]; [1 -1  1]; [1 1 -1]]
+    @test kgrid_from_minimal_spacing(lattice * u"angstrom", 0.5 / u"angstrom") == [9; 9; 9]
+end
+
+@testset "kgrid_from_minimal_n_kpoints" begin
+    lattice = [[-1.0 1 1]; [1 -1  1]; [1 1 -1]]
+    @test kgrid_from_minimal_n_kpoints(lattice * u"Å", 1000) == [10, 10, 10]
+
+    @test kgrid_from_minimal_n_kpoints(magnesium.lattice, 1) == [1, 1, 1]
+    for n_kpt in [10, 20, 100, 400, 900, 1200]
+        @test prod(kgrid_from_minimal_n_kpoints(magnesium.lattice, n_kpt)) ≥ n_kpt
+    end
+
+    lattice = diagm([4., 10, 0])
+    @test kgrid_from_minimal_n_kpoints(lattice, 1000) == [50, 20, 1]
+    @test kgrid_from_minimal_n_kpoints(diagm([10, 0, 0]), 913) == [913, 1, 1]
 end