New AtomicPotentials.jl + PseudoPotentialIO.jl backend

[azadoks]

This draft replaces the existing pseudopotential framework in pseudos with the packages

• AtomicPotentials.jl which implements numerical pseudopotentials and the existing analytical Element types under a common interface.
• PseudoPotentialIOExperimental.jl which removes all pseudopotential evaluation code from PseudoPotentialIO.jl and instead focuses only on parsing pseudopotential file formats.

Ideally, this refactor aims to

• Homogenize use of atomic potentials to subsequently
  • increase code sharing among Hamiltonian terms which are derived from the atomic potentials (AtomicLocal, AtomicNonlocal, PspCorrection, and Xc) which will
  • facilitate code readability and performance optimization.
• Facilitate support for evaluating atomic potentials on accelerators (GPU)

There's a lot to this PR, so it will take a while to write up, but here are the key

Progress points

• All Element subtypes are gone, replaced by Element(symbol, potential)

• All pseudos code is gone, replaced by AtomicPotentials.jl

• Loading pseudopotentials is done by load_psp, which supports the DFTK "hgh/.../..." syntax as well as "/path/to/file" and the PseudoLibrary ("family", "filename")

• Form factor calculation is centralized in form_factors.jl: angular + radial parts

  • Radial form factors in Fourier space are now interpolated using either Linear or CubicSplines interpolation from Interpolations.jl (GPU-supported)

• Structure factor and structure factor gradients are centralized in structure_factors.jl

• Projection vector construction is centralized in projection_vectors.jl

  • Projection coefficients are now stored in sparse block-diagonal arrays

• Local atomic superposition (i.e. valence density, core density, and local potential) construction and force calculations (for AtomicLocal and Xc terms) are centralized in atomic_superposition.jl

• Density guess code is refactored and much much shorter (using atomic_superposition.jl)

• New orbital guess code is in orbitals.jl, using UPF's PSWFC

  • TODO: need to parse the pseudo-energy of the orbitals so that they can be sorted by lowest energy if fewer orbitals than available are requested for the guess

• Construction of all the affected Hamiltonian terms with HGH, UPF, and PSP8 pseudopotentials is working

• Total energies, forces, and stresses (with AD) are working and agree for at least silicon.jl

• PlaneWaveBasis construction time is 5~20x faster depending on the discretization parameters, system size, pseudopotential type, interpolation method, quadrature method, ....

• compute_forces and compute_stresses also see a good speedup

• guess_density is faster, but it was never slow

• A non-zero amount of effort has been put in to get pseudos-on-GPU working: radial Fourier transforms, interpolation, etc. can all be done on GPU. Some deep sleuthing will need to be done to make the code compile with CUDA though (I get all sorts of invalid IR problems, mostly to do with compute_angular_form_factors and compute_radial_form_factors for the moment).

Pain points

• AtomsBase support is broken; need to think about how to serialize the Elements with their potentials. The potentials have the file content hash and DFTK identifier / filepath / PseudoLibrary identifier from loading, which should be enough.
• GPU is currently broken; I can revert to CPU-computed pseudo quantities, but that's not the goal, is it?

In-depth discussion

Summary

This draft replaces the existing framework for (pseudo-)atomic potentials with the packages AtomicPotentials.jl and PseudoPotentialIOExperimental.jl.

Previous implementation

pseudo/
	NormConservingPsp.jl  # Pseudopotential interface
		eval_psp_projector_real           # beta_li(r)
		eval_psp_projector_fourier        # beta_li(q)
		eval_psp_local_real               # Vloc(r)
		eval_psp_local_fourier            # Vloc(q)
		eval_psp_energy_correction        # lim(q->0)[Vloc(q)] 
		eval_psp_density_valence_real     # rhov(r)
		eval_psp_density_valence_fourier  # rhov(q)
		eval_psp_density_core_real        # rhoc(r)
		eval_psp_density_core_fourier     # rhoc(q)
		projector_indices
		count_n_proj_radial               # Nbeta(l)
		count_n_proj                      # Nbeta(l) * (2l + 1)
	PspHgh.jl                             # Implementation for HGH
		*: real, fourier
		eval_psp_projector_*              # Analytical
		eval_psp_local_*                  # Analytical
	PspUpf.jl                             # Implementation for UPF
		*: Vloc, beta, chi, rhoc, rhov
		eval_psp_*_real                   # Interpolate
		eval_psp_*_fourier                # Radial FT w/ Rect. quadrature
	attach_psp.jl                         # AtomsBase interop.
	list_psp.jl                           # List/filter included HGH pseudos
	load_psp.jl                           # Load incl. HGH, arb. HGH, UPF
terms/
	local.jl              # Local atomic potential
		Builds a superposition of atomic Vloc(Q) -IFFT-> Vloc(R)
	nonlocal.jl           # Non-local part of separable KB pseudo
		Builds projection vectors beta_k(Q) and dense coupling/coefficient
		matrices D_k[i,j]
	psp_correction.jl     # DC correction from local potential
		Calls eval_psp_energy_correction and does the appropriate sum
		over sites
	xc.jl                 # Non-linear core correction
		Calls density_methods.jl for rhoc(R)
density_methods.jl        # Superposition of atomic densities
	Builds a superposition of atomic F(Q) -IFFT-> f(R)
		Pseudo rhov, Gaussian rhov, pseudo rhoc
elements.jl               # Type of the atoms stored by Model
	charge_nuclear
	atomic_symbol
	charge_ionic
	n_elec_valence
	n_elec_core
	has_core_density
	valence_charge_density_fourier
	gaussian_valence_charge_density_fourier
	core_charge_density_fourier
	ElementCoulomb(Z, symbol)
	ElementPsp(Z, symbol, psp)
	ElementCohenBergstresser(Z, symbol, V(q), a)
	ElementGaussian(alpha, L, symbol)

Proposed impelmentation

PseudoPotentialIO.jl
	file/                          # Loaders and writers
		file.jl                    # Interface
		hgh.jl                     # HGH format
		psp8.jl                    # ABINIT PSP8 format
		upf.jl                     # Universal Pseudopotential Format
		upf1.jl                    # UPF vOld (pseudo-XML)
		upf2.jl                    # UPF v2.0.1 (XML)
	io/
		list.jl
			list_families          # List arifacts from PseudoLibrary
			list_family_psps       # List files in an artifact
			resolve_family         # Resolve artifact filepaths
		load.jl
			load_psp_file          # Load a psp file (artifact, path)
			load_family_psp_files  # Load all fiels in an artifact
		show.jl                    # Various summaries of artifacts
			show_family_summary
			show_family_table
			show_family_periodic_table

AtomicPotentials.jl
	atomic.jl
		AtomicPotential                # Generic container
			local_potential            # Vloc (REQUIRED)
				AbstractLocalPotential
			non_local_potential        # Vnl
			  projectors               # beta[l][i]
				  AbstractQuantity
			  coupling                 # D[l][i,i']
				  AbstractMatrix
			  augmentation
				  functions            # Q[l][i,i']
					  AbstractQuantity
				  coupling             # q[l][i,i']
					  AbstractMatrix
			core_charge_density        # rhoc
				AbstractQuantity
			valence_charge_density     # rhov
				AbstractQuantity
			orbitals                   # chi[l][i]
				AbstractQuantity
	quantity.jl
		EvaluationSpace
		RealSpace
		FourierSpace
		--
		Analyticity
		Analtyical
		Numerical
		--
		AbstractQuantity
		NumericalQuantity
		HghProjector
		HydrogenicProjector
		GaussianChargeDensity
	local_potential.jl
		AbstractLocalPotential
		NumericalLocalPotential
		HghLocalPotential
		CoulombLocalPotential
		CoulombErfLocalPotential
		CohenBergstresserLocalPotential
		--
		AbstractLocalPotentialCorrection
			# Coulomb-like tail used to perform radial Fourier transform
			# of numerical local potentials
		CoulombLocalPotentialCorrection     # -Z / r
		CoulombErfLocalPotentialCorrection  # -Z erf(r) / r 
	augmentation.jl
		Augmentation                   # Ultrasoft augmentation
	non_local.jl
		NonLocalPotential  # Container for beta_li(r), D_l[i,i'], augment.
	pseudopotentialio.jl   # Load AtomicPotentials from pseudopot. files
		AtomicPotential(::PseudoPotentialIO.HghFile)
		AtomicPotential(::PseudoPotentialIO.UpfFile)
		AtomicPotential(::PseudoPotentialIO.Psp8File)
	radial_fourier_transform.jl
		rft   # Radial Fourier transform f(r) -> F(q)
		irft  # Inverse radial Fourier transform F(q) -> f(r)
	fast_sphericalbesselj.jl           # j_l(x) for l = 0:1:5
		# Used by rft / irft
	MeshClassification/
		# Used by Interpolation and NumericalQuadrature to determine
		# whether a mesh is uniform (optimized methods available)
		# or non-uniform (generalized methods necessary)
		MeshClassification.jl
			gradient    # dx(i)/di
				# Used by NumericalQuadrature.QESimpson, which uses the
				# mesh gradient rather than finite differences
			Mesh        # Abstract type
			Uniform     # x(i) = b +    a (i - 1)
			Log         # x(i) = b  exp(a (i - 1))
			ShiftedLog  # x(i) = b (exp(a (i - 1)) - 1)
	Interpolation/                    # Interpolation interface
		# Allows for passing around a configuration used to build
		# interpolators, i.e. when instantiating `Term`s in DFTK.
		# GPU support is needed for interpolation Fourier-space
		# quantities; while non-uniform mesh support is needed
		# for interpolating the real-space quantities from UPFs,
		# which sometimes use log and shifted-log grids. Interpolating
		# real-space quantities before radial Fourier transforms can
		# often greatly improve the quality of the transform.
		Interpolation.jl
			# Interpolations.jl: supports GPU -- useful for Fourier-space
			# interpolation
			InterpolationsLinear      # Interpolations.Linear
				# Uniform + non-uniform meshes
			InterpolationsCubicSpline # Interpolations.CubicSpline
				# Uniform mesh only
			# BSplineKit.jl: no GPU support
			BSplineKitSpline          # BSplineKit.Spline
				# Uniform + non-uniform meshes -- useful for real-space
				# interpolation
	NumericalQuadrature/
		# Allows for selecting a quadrature method in the (inverse) radial
		# Fourier transform. Some methods reproduce the behavior of other
		# codes but have restrictions/drawbacks. `Trapezoid` and `Simpson`
		# methods are bespoke and support non-uniform meshes and both even
		# and odd numbers of intervals.
		NumericalQuadrature.jl
			integration_weights(x, method)
				# Build a vector of weights such that integration for a
				# function f(x) = y is performed by dot(y, weights)
			integrate(x, y, method)
		abinit_corrected_trapezoid.jl  # ABINIT's integrator (uniform only)
		qe_simpson.jl                  # QE's integrator (uniform + non)
			# Performs poorly on non-uniform meshes and is fundamentally
			# incorrect in its use of the mesh derivative
		simpson.jl                     # Composite Simpson (uniform + non)
		trapezoid.jl                   # Composite Trapezoid (unif. + non)

DFTK.jl/
	elements.jl
		Element          # Generic element type
			Z
			symbol
			potential
				AtomicPotential
		charge_nuclear
		atomic_symbol
		charge_ionic
		n_elec_valence
		n_elec_core
	structure_factors.jl
		compute_structure_factor           # cis2pi(Q * R)
		compute_structure_factors
			# Compute structure factors for vectors and singly-nested
			# vectors of positions
		compute_structure_factor_gradient   # 2pi i Q cis2pi(Q * R)
		compute_structure_factor_gradients
			# Vectors and singly-nested vectors of positions
	form_factors.jl
		compute_form_factor_angular
		compute_form_factors_angular
		compute_form_factor_radial
		compute_form_factors_radial
	atomic_superposition.jl
		build_atomic_superposition
		
	projection_vectors.jl
	density_methods.jl
		guess_density
		random_density
	orbitals.jl
		guess_orbitals
	terms/

Loading pseudopotentials / construction atomic potentials

For pseudopotentials, the target interface is:

using DFTK
# DFTK style
Si_1 = Element(:Si, load_psp("hgh/lda/si-q4"))
Si_1.potential::AtomicPotential{RealSpace}
Si_1.potential.local_potential::HghLocalPotential{RealSpace}
# PseudoLibrary family, file
Si_2 = Element(:Si, load_psp("hgh_lda_upf", "Si.pz-hgh.UPF"))
Si_2.potential::AtomicPotential{RealSpace}
Si_2.potential.local_potential::NumericalLocalPotential{RealSpace}
# absolute file-path
Si_3 = Element(:Si, load_psp("/path/to/Si.psp8"))
Si_3.potential::AtomicPotential{RealSpace}
Si_3.potential.local_potential::NumericalLocalPotential{RealSpace}
# relative file-path
Si_4 = Element(:Si, load_psp("../path/to/Si.hgh"))
Si_4.potential::AtomicPotential{RealSpace}
Si_4.potential.local_potential::HghLocalPotential{RealSpace}

Internally, load_psp determines whether the string is a DFTK-style identifier or not; it then resolves the identifier to a file path or (family, file) pair and calls load_psp_file from PseudoPotentialIO.jl. The resulting PsPFile is then loaded into an AtomicPotential by AtomicPotentials.jl.

For arbitrary potentials, the interface is:

using DFTK
using AtomicPotentials
Si_5::Element = Element(
	:Si,
	AtomicPotential(
		CoulombLocalPotential{RealSpace}(:Si),
		ionic_charge_density=GaussianChargeDensity{RealSpace}(14, 14),
		# orbitals, etc...
	)
)

Make-up and functionality of atomic potentials

Currently AtomicPotentials can hold all relevant quantities for separable norm-conserving and ultrasoft pseudopotentials:

struct AtomicPotential{S<:EvaluationSpace}
	local_potential::AbstractLocalPotential{S}
	non_local_potential::NonLocalPotential{S}
		projectors::Vector{Vector{AbstractQuantitiy{S}}}
		coupiling::Vector{Matrix}
		augmentation::Augmentation{S}
			functions::Vector{Matrix{AbstractQuantity{S}}}
			coupling::Vector{Matrix{Matrix}}
	core_charge_density::AbstractQuantity{S}
	ionic_charge_density::AbstractQuantity{S}
	orbitals::Vector{Vector{AbstractQuantity{S}}}
end

It keeps track of the evaluation space of all its component quantities and provides some convenience functions for accessing quantities and their metadata and for transforming the potential as a whole.

Quantities

Quantities can exist in real or Fourier space, and be flagged as analytical or numerical. Both these flags affect how the quantity is evaluated and Fourier transformed.

Real-space and Fourier-space

These flags are used to dispatch on evaluation functions for analytical potentials.

Numerical and Analytical

These flags are used to dispatch on radial Fourier transform functions; either using numerical integration or for simply swapping the RealSpace/FourierSpace type parameter so that evaluation dispatches to the appropriate method.

Efficient evaluation

One of, if not the, most expensive part of using numerical pseudo-potentials is computing the radial Fourier transform

$$f(q) = 4\pi \int_0^{\infty} r f(r) j_l(qr) r dr.$$

AtomicPotentials.jl provides the function rft to evaluate this transform using a variety of quadrature methods. It also provides a variety of interpolation methods so that the transform can be evaluated on a selected set of points and then interpolated onto, e.g., the full G-vector grid of a plane-wave basis.

Efficient and reusable code in DFTK

Fourier transforms of atomic quantities

Local quantities, e.g. core/valence charge density or local potential, all behave in similar ways and are used in similar ways in DFTK.
Building density guesses, and instantiating the AtomicLocalPotential and Xc terms share a common motif:

$$F(Q) = \sum_{I} \underbrace{e^{-2\pi i (Q \cdot R_{I})}}_{structure-factor} %\underbrace{(-i)^l Y_{lm}(Q)}_{form-factor~angular} \underbrace{4\pi \int_0^{\infty} r f_{a(I)}(r) j_0(|Q|r) r dr}_{form-factor~radial}$$

Non-local quantities, e.g. Kleinman-Bylander projectors or (pseudo-)atomic orbitals, behave in a slightly different way in that their form-factor has an additional angular part:

$$F_{Ilmi}(Q) = \underbrace{e^{-2\pi i (Q \cdot R_{I})}}_{structure-factor} \underbrace{(-i)^l Y_{lm}(Q)}_{form-factor~angular} \underbrace{4\pi \int_0^{\infty} r f_{a(I)li}(r) j_l(|Q|r) r dr}_{form-factor~radial}$$

An obvious method for reducing code complexity here is to have a minimal set of functions which are shared:

function structure_factor(R, Q) end
function form_factor_angular(l, m, Q) end
function form_factor_radial(f, Q) end

form_factor_radial is particularly expensive due to it being a radial Fourier transform of a real-space function. Therefore, in practice, form_factor_radial is used as follows.

Numerical quantities are radial Fourier transformed before the inner loops described below. They are evaluated on a uniform mesh of points taken as input to the PlaneWaveBasis using a quadrature method also taken as input. The resulting Fourier-space quantity is interpolated using a method taken as input by the PlaneWaveBasis. The resulting interpolator is then passed to form_factor_radial which evaluates it.

Analytical quantities are passed around with no changes. If they are in Fourier space, form_factor_radial simply evaluates them. If they are in Real space, it first calls rft to retrieve a copy of the quantity with the EvaluationSpace type parameter changed. Then it evaluates the quantity as above.

Computing each of these terms for all G-vectors $Q$ and all of the various indices becomes increasingly expensive, so minimizing or removing repeated calculations is important. Therefore, each term is calculated separately only for indices and arguments on which it actually depends. A mostly true-to-life example for building projection vectors for the non-local potential is shown below:

# given: kpt::Kpoint

# This grouping by species is done often but can't be completely
# generalized at the moment; see `group_local_quantities`,
# `_group_non_local`, `guess_orbitals`
atom_groups = model.atom_groups
atoms = [model.atoms[first(group)] for group in atom_groups]
positions = [model.positions[group] for group in atom_groups]
projectors = [atom.potential.non_local_potential.projectors
			  for group in atom_groups]
couplings = [atom.potential.non_local_potential.coupling
			 for group in atom_groups]

# Compute structure factors once for all G-vectors and positions
# Implemented in `compute_structure_factors`
structure_factors = map(positions) do R_a
	map(R_a) do R_aI
		map(Gplusk_vectors(basis, kpt)) do Q
			structure_factor(R_aI, Q)
		end::Array{Complex{T}}  # G-vector
	end::Vector{Array{Complex{T}}}  # site
end::Vector{Vector{Array{Complex{T}}}}  # species

# Compute angular part of the form factor for each l, m
# Implemented in `compute_form_factors_angular`
l_max = map(atoms) do atom
	maximum_angular_momentum(atom.potential.non_local_potential)
end
form_factors_angular = map(0:l_max) do l
	map(-l:+l) do m
		map(Gplusk_vectors_cart(basis, kpt)) do Q
			form_factor_angular(l, m, Q)
		end::Array{Complex{T}}  # G-vector
	end::Vector{Array{Complex{T}}}  # magnetic quantum number
end::Vector{Vector{Array{Complex{T}}}}  # angular momentum

# Compute radial part of the form factor for each species
# Implemented in `compute_form_factors_radial`
form_factors_radial = map(projectors) do f_a
    map(f_a) do f_al
        map(f_al) do f_ali
	        map(Gplusk_vectors_cart(basis, kpt)) do Q
		        form_factor_radial(f_ali, Q)
		    end::Array{T}  # G-vector
		end::Vector{Array{T}}  # projector
	end::Vector{Vector{Array{T}}}  # angular momentum
end::Vector{Vector{Vector{Array{T}}}}  # species

# Combine the quantities to compute the projection vectors
# Implemented in `build_projection_vectors`
N = 1 / sqrt(model.unit_cell_volume)
P = map(zip(structure_factors, form_factors_radial)) do (sf_a, ffr_a)
    map(sf_a) do sf_aI
        map(zip(ffr_a, form_factors_angular)) do (ffr_al, ffa_l)
            map(ffa_l) do ffa_lm
	            map(ffr_al) do ffr_ali
	                N .* sf_aI .* ffa_lm .* ffr_ali
	            end::Array{Complex{T}}  # projector
	        end::Vector{Array{Complex{T}}}  # mangetic quantum number
	    end::Vector{::Vector{Array{Complex{T}}}}  # angular momentum
	end::Vector{::Vector{::Vector{Array{Complex{T}}}}}  # site
end::Vector{::Vector{::Vector{::Vector{Array{Complex{T}}}}}}  # species

# Flatten to a multi-indexed matrix for efficient calculation of P D P'
# Implemented in `projection_vectors_to_matrix`
reduce_hcat = Base.Fix1(reduce, hcat)
P::Matrix{Complex{T}} = (  # P[G-vector, aIlmi]
#        species    site       l          m          i
	P |> flatten |> flatten |> flatten |> flatten |> flatten |> reduce_hcat
)

# Build the coefficient/coupling matrix D
# Implemented in `build_projection_coupling`
n_sites_per_atom = map(length, positions)
D = map(zip(couplings, n_sites_per_atom)) do (D_a, n_a)
	map(1:n_a) do _
		map(enumerate(D_a)) do (l_index, D_al)
			l = l_index - 1
			map(-l:+l) do _
				Symmetric(D_al)  # projector,projector
			end  # magnetic quantum number
		end  # angular momentum
	end  # site
end  # species

# Flatten to a sparse block-diagonal matrix for efficient calculation of
# P D P'; Implemented in `projection_coupling_to_matrix`
spl_bd = splat(blockdiag)
D::SparseArrayCSC{T} = (  # D[aIlmi,a'I'l'm'i']
#        species    site       l          m
	D |> flatten |> flatten |> flatten |> flatten .|> sparse |> spl_bd
)

Vnl_k = P * D * P'