New AtomicPotentials.jl + PseudoPotentialIO.jl backend

Project.toml

@@ -5,9 +5,13 @@ version = "0.6.9"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+AtomicPotentials = "0015ad10-6e07-47eb-baf7-dbb079bee869"
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DftFunctionals = "6bd331d2-b28d-4fd3-880e-1a1c7f37947f"
@@ -25,6 +29,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
@@ -33,13 +38,14 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
-PseudoPotentialIO = "cb339c56-07fa-4cb2-923a-142469552264"
+PseudoPotentialIOExperimental = "3c1e8497-ccc6-4bf7-b29e-53e04e06307e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Spglib = "f761d5c5-86db-4880-b97f-9680a7cccfb5"
+StatProfilerHTML = "a8a75453-ed82-57c9-9e16-4cd1196ecbf5"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
@@ -49,6 +55,7 @@ spglib_jll = "ac4a9f1e-bdb2-5204-990c-47c8b2f70d4e"
 
 [compat]
 AbstractFFTs = "1"
+Adapt = "3.6"
 AtomsBase = "0.3.1"
 Brillouin = "0.5.11"
 ChainRulesCore = "1.15"
@@ -72,7 +79,6 @@ Polynomials = "3"
 PrecompileTools = "1"
 Primes = "0.5"
 ProgressMeter = "1"
-PseudoPotentialIO = "0.1"
 Requires = "1"
 Roots = "2"
 SpecialFunctions = "2"

src/DFTK.jl

@@ -15,10 +15,15 @@ using UnitfulAtomic
 using ForwardDiff
 using AbstractFFTs
 using GPUArraysCore
+using Adapt
 using Random
 using ChainRulesCore
 using PrecompileTools
 
+import PseudoPotentialIOExperimental
+using AtomicPotentials
+using AtomsBase
+
 export Vec3
 export Mat3
 export mpi_nprocs
@@ -38,21 +43,11 @@ include("architecture.jl")
 include("common/zeros_like.jl")
 include("common/norm.jl")
 
-export PspHgh
-export PspUpf
-include("pseudo/NormConservingPsp.jl")
-include("pseudo/PspHgh.jl")
-include("pseudo/PspUpf.jl")
-
-export ElementPsp
-export ElementCohenBergstresser
-export ElementCoulomb
-export ElementGaussian
-export charge_nuclear
-export charge_ionic
-export atomic_symbol
+export Element
+export nuclear_charge
 export n_elec_valence
 export n_elec_core
+export atomic_symbol
 include("elements.jl")
 
 export SymOp
@@ -174,15 +169,30 @@ export guess_density
 export random_density
 include("density_methods.jl")
 
+export compute_structure_factors
+export compute_structure_factor_gradients
+include("structure_factors.jl")
+
+export compute_form_factors_angular
+export compute_form_factors_radial
+include("form_factors.jl")
+
+export build_atomic_superposition
+include("atomic_superposition.jl")
+
+export build_projection_vectors
+export build_projection_coupling
+include("projection_vectors.jl")
+
 export load_psp
-export list_psp
-export attach_psp
+# export list_psp
+# export attach_psp
 include("pseudo/load_psp.jl")
-include("pseudo/list_psp.jl")
-include("pseudo/attach_psp.jl")
+# include("pseudo/list_psp.jl")
+# include("pseudo/attach_psp.jl")
 
 export DFTKPotential
-export atomic_system, periodic_system  # Reexport from AtomsBase
+# export atomic_system, periodic_system  # Reexport from AtomsBase
 export run_wannier90
 include("external/atomsbase.jl")
 include("external/interatomicpotentials.jl")
@@ -242,23 +252,23 @@ function __init__()
 end
 
 # Precompilation block with a basic workflow
-@setup_workload begin
-    # very artificial silicon ground state example
-    a = 10.26
-    lattice = a / 2 * [[0 1 1.];
-                       [1 0 1.];
-                       [1 1 0.]]
-    Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
-    atoms     = [Si, Si]
-    positions = [ones(3)/8, -ones(3)/8]
-    magnetic_moments = [2, -2]
+# @setup_workload begin
+#     # very artificial silicon ground state example
+#     a = 10.26
+#     lattice = a / 2 * [[0 1 1.];
+#                        [1 0 1.];
+#                        [1 1 0.]]
+#     Si = AtomicPotential(load_psp_file("hgh_lda_hgh", "si-q4.hgh"))
+#     atoms     = [Si, Si]
+#     positions = [ones(3)/8, -ones(3)/8]
+#     magnetic_moments = [2, -2]
 
-    @compile_workload begin
-        model = model_LDA(lattice, atoms, positions;
-                          magnetic_moments, temperature=0.1, spin_polarization=:collinear)
-        basis = PlaneWaveBasis(model; Ecut=5, kgrid=[2, 2, 2])
-        ρ0 = guess_density(basis, magnetic_moments)
-        scfres = self_consistent_field(basis; ρ=ρ0, tol=1e-2, maxiter=3, callback=identity)
-    end
-end
+#     @compile_workload begin
+#         model = model_LDA(lattice, atoms, positions;
+#                           magnetic_moments, temperature=0.1, spin_polarization=:collinear)
+#         basis = PlaneWaveBasis(model; Ecut=5, kgrid=[2, 2, 2])
+#         ρ0 = guess_density(basis, magnetic_moments)
+#         scfres = self_consistent_field(basis; ρ=ρ0, tol=1e-2, maxiter=3, callback=identity)
+#     end
+# end
 end # module DFTK

src/DispatchFunctional.jl

@@ -140,7 +140,7 @@ struct DispatchFunctional{Family,Kind} <: Functional{Family,Kind}
     inner::LibxcFunctional{Family,Kind}
 end
 DispatchFunctional(identifier::Symbol) = DispatchFunctional(LibxcFunctional(identifier))
-DftFunctionals.identifier(fun::DispatchFunctional) = identifier(fun.inner)
+DftFunctionals.identifier(fun::DispatchFunctional) = DftFunctionals.identifier(fun.inner)
 DftFunctionals.has_energy(fun::DispatchFunctional) = has_energy(fun.inner)
 
 for fun in (:potential_terms, :kernel_terms)

src/Model.jl

@@ -49,7 +49,7 @@ struct Model{T <: Real, VT <: Real}
     # `atom_groups` contains the groups of indices into atoms and positions, which
     # point to identical atoms. It is computed automatically on Model construction and may
     # be used to optimise the term instantiation.
-    atoms::Vector{Element}
+    atoms::Vector{Element{RealSpace}}
     positions::Vector{Vec3{T}}  # positions[i] is the location of atoms[i] in fract. coords
     atom_groups::Vector{Vector{Int}}  # atoms[i] == atoms[j] for all i, j in atom_group[α]
 
@@ -93,15 +93,15 @@ external potential breaks some of the passed symmetries. Use `false` to turn off
 symmetries completely.
 """
 function Model(lattice::AbstractMatrix{T},
-               atoms::Vector{<:Element}=Element[],
+               atoms::Vector{<:Element{RealSpace}}=Element[],
                positions::Vector{<:AbstractVector}=Vec3{T}[];
                model_name="custom",
                εF=nothing,
                n_electrons::Union{Int,Nothing}=isnothing(εF) ?
                                                n_electrons_from_atoms(atoms) : nothing,
                # Force electrostatics with non-neutral cells; results not guaranteed.
                # Set to `true` by default for charged systems.
-               disable_electrostatics_check=all(iszero, charge_ionic.(atoms)),
+               disable_electrostatics_check=all(iszero, ionic_charge.(atoms)),
                magnetic_moments=T[],
                terms=[Kinetic()],
                temperature=zero(T),
@@ -181,11 +181,11 @@ function Model(lattice::AbstractMatrix{T},
                            n_electrons, εF, spin_polarization, n_spin, temperature, smearing,
                            atoms, positions, atom_groups, terms, symmetries)
 end
-function Model(lattice::AbstractMatrix{<:Integer}, atoms::Vector{<:Element},
+function Model(lattice::AbstractMatrix{<:Integer}, atoms::Vector{<:Element{RealSpace}},
                positions::Vector{<:AbstractVector}; kwargs...)
     Model(Float64.(lattice), atoms, positions; kwargs...)
 end
-function Model(lattice::AbstractMatrix{<:Quantity}, atoms::Vector{<:Element},
+function Model(lattice::AbstractMatrix{<:Quantity}, atoms::Vector{<:Element{RealSpace}},
                positions::Vector{<:AbstractVector}; kwargs...)
     Model(austrip.(lattice), atoms, positions; kwargs...)
 end
@@ -214,7 +214,7 @@ or for changing `lattice` or `positions` in geometry/lattice optimisations.
 function Model{T}(model::Model;
                   lattice::AbstractMatrix=model.lattice,
                   positions::Vector{<:AbstractVector}=model.positions,
-                  atoms::Vector{<:Element}=model.atoms,
+                  atoms::Vector{<:Element{RealSpace}}=model.atoms,
                   kwargs...) where {T <: Real}
     Model(T.(lattice), atoms, positions;
           model.model_name,

src/PlaneWaveBasis.jl

@@ -45,6 +45,7 @@ struct PlaneWaveBasis{T,
                       T_G_vectors  <: AbstractArray{Vec3{Int}, 3},
                       T_r_vectors  <: AbstractArray{Vec3{VT},  3},
                       T_kpt_G_vecs <: AbstractVector{Vec3{Int}},
+                      T_q_grid     <: AbstractVector{T}
                      } <: AbstractBasis{T}
 
     # T is the default type to express data, VT the corresponding bare value type (i.e. not dual)
@@ -111,6 +112,17 @@ struct PlaneWaveBasis{T,
     # Therefore, all quantities should be symmetric to machine precision
     symmetries_respect_rgrid::Bool
 
+    ## Atomic potential parameters
+    # Quadrature method for computing radial Fourier transforms of numerical atomic
+    # quantities.
+    atom_rft_quadrature_method::AtomicPotentials.NumericalQuadrature.QuadratureMethod
+    # Radial grid in Fourier-space on which to evaluate the radial Fourier transforms
+    # of numerical atomic quantities.
+    atom_qgrid::T_q_grid
+    # Method for interpolating Fourier-space numerical atomic quantities for evaluation
+    # on the full G-grid.
+    atom_q_interpolation_method::AtomicPotentials.Interpolation.InterpolationMethod
+
     ## Instantiated terms (<: Term). See Hamiltonian for high-level usage
     terms::Vector{Any}
 end
@@ -159,7 +171,9 @@ end
 # and are stored in PlaneWaveBasis for easy reconstruction.
 function PlaneWaveBasis(model::Model{T}, Ecut::Number, fft_size, variational,
                         kcoords, kweights, kgrid, kshift,
-                        symmetries_respect_rgrid, comm_kpts,
+                        symmetries_respect_rgrid, atom_rft_quadrature_method,
+                        atom_minimal_dq,
+                        atom_q_interpolation_method, comm_kpts,
                         architecture::Arch
                        ) where {T <: Real, Arch <: AbstractArchitecture}
     # Validate fft_size
@@ -216,6 +230,14 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Number, fft_size, variational,
     ifft_normalization = 1/sqrt(model.unit_cell_volume)
     fft_normalization  = sqrt(model.unit_cell_volume) / length(ipFFT)
 
+    # Atomic potential Fourier transform grid
+    # avoiding model (not isbits) in closure for GPU compat
+    recip_lattice = model.recip_lattice
+    Gs_cart = map(G -> recip_lattice * G, Gs)
+    (qmin, qmax) = extrema(norm.(Gs_cart))
+    nq = ceil(Int, (qmax - qmin) / atom_minimal_dq)
+    atom_qgrid = range(qmin, qmax, nq)
+
     # Compute k-point information and spread them across processors
     # Right now we split only the kcoords: both spin channels have to be handled
     # by the same process
@@ -278,18 +300,20 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Number, fft_size, variational,
     r_vectors = [Vec3{VT}(VT(i-1) / N1, VT(j-1) / N2, VT(k-1) / N3)
                  for i = 1:N1, j = 1:N2, k = 1:N3]
     r_vectors = to_device(architecture, r_vectors)
+    atom_qgrid = to_device(architecture, atom_qgrid)
     terms = Vector{Any}(undef, length(model.term_types))  # Dummy terms array, filled below
 
     basis = PlaneWaveBasis{T, value_type(T), Arch, typeof(Gs), typeof(r_vectors),
-                           typeof(kpoints[1].G_vectors)}(
-        model, fft_size, dvol,
-        Ecut, variational,
+                           typeof(kpoints[1].G_vectors), typeof(atom_qgrid)}(
+        model,
+        fft_size, dvol, Ecut, variational,
         opFFT, ipFFT, opBFFT, ipBFFT,
         fft_normalization, ifft_normalization,
         Gs, r_vectors,
         kpoints, kweights_thisproc, kgrid, kshift,
         kcoords_global, kweights_global, comm_kpts, krange_thisproc, krange_allprocs,
-        architecture, symmetries, symmetries_respect_rgrid, terms)
+        architecture, symmetries, symmetries_respect_rgrid,
+        atom_rft_quadrature_method, atom_qgrid, atom_q_interpolation_method, terms)
     # Instantiate the terms with the basis
     for (it, t) in enumerate(model.term_types)
         term_name = string(nameof(typeof(t)))
@@ -304,6 +328,9 @@ end
 @timing function PlaneWaveBasis(model::Model{T}, Ecut::Number,
                                 kcoords ::Union{Nothing, AbstractVector},
                                 kweights::Union{Nothing, AbstractVector};
+                                atom_rft_quadrature_method=AtomicPotentials.NumericalQuadrature.Simpson(),
+                                atom_minimal_dq=0.01,
+                                atom_q_interpolation_method=AtomicPotentials.Interpolation.InterpolationsCubicSpline(),
                                 variational=true, fft_size=nothing,
                                 kgrid=nothing, kshift=nothing,
                                 symmetries_respect_rgrid=isnothing(fft_size),
@@ -324,8 +351,10 @@ end
         end
         fft_size = compute_fft_size(model, Ecut, kcoords; factors)
     end
-    PlaneWaveBasis(model, Ecut, fft_size, variational, kcoords, kweights,
-                   kgrid, kshift, symmetries_respect_rgrid, comm_kpts, architecture)
+    PlaneWaveBasis(model, Ecut, fft_size, variational, kcoords,
+                   kweights, kgrid, kshift, symmetries_respect_rgrid,
+                   atom_rft_quadrature_method, atom_minimal_dq, atom_q_interpolation_method,
+                   comm_kpts, architecture)
 end
 
 @doc raw"""
@@ -353,7 +382,11 @@ Creates a new basis identical to `basis`, but with a custom set of kpoints
     PlaneWaveBasis(basis.model, basis.Ecut,
                    basis.fft_size, basis.variational,
                    kcoords, kweights, kgrid, kshift,
-                   basis.symmetries_respect_rgrid, basis.comm_kpts, basis.architecture)
+                   basis.symmetries_respect_rgrid,
+                   basis.atom_rft_quadrature_method,
+                   basis.atom_minimal_dq,
+                   basis.atom_q_interpolation_method,
+                   basis.comm_kpts, basis.architecture)
 end
 
 """
@@ -553,3 +586,18 @@ function gather_kpts(data::AbstractArray, basis::PlaneWaveBasis)
         nothing
     end
 end
+
+# TODO (AZ): This should probably go somewhere else
+# function AtomicPotentials.rft(basis::PlaneWaveBasis, qty)
+#     return rft(qty, basis.atom_qgrid; quadrature_method=basis.atom_rft_quadrature_method)
+# end
+
+# function AtomicPotentials.rft(basis::PlaneWaveBasis, qty;
+#                               quadrature_method=basis.atom_rft_quadrature_method)
+#     return rft(qty, norm.(G_vectors_cart); quadrature_method)
+# end
+
+# function AtomicPotentials.rft(basis::PlaneWaveBasis, kpt::Kpoint, qty;
+#                               quadrature_method=basis.atom_rft_quadrature_method)
+#     return rft(qty, norm.(Gplusk_vectors_cart(basis, kpt)); quadrature_method)
+# end

src/architecture.jl

@@ -16,14 +16,16 @@ GPU(::Type{T}) where {T <: AbstractArray} = GPU{T}()
 Transfer an array from a device (typically a GPU) to the CPU.
 """
 to_cpu(x::AbstractArray) = Array(x)
-to_cpu(x::Array) = x
+to_cpu(x::Array)         = x
+to_cpu(x       )         = adapt(Array, x)
 
 """
 Transfer an array to a particular device (typically a GPU)
 """
 to_device(::CPU, x) = to_cpu(x)
 to_device(::GPU{ArrayType}, x::AbstractArray) where {ArrayType} = ArrayType(x)
 to_device(::GPU{ArrayType}, x::ArrayType)     where {ArrayType} = x
+to_device(::GPU{ArrayType}, x)                where {ArrayType} = adapt(ArrayType, x)
 
 """
 Synchronize data and finish all operations on the execution stream of the device.