Integrate with AtomsBase (#558)

Project.toml

@@ -5,6 +5,7 @@ version = "0.4.7"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0"
 Conda = "8f4d0f93-b110-5947-807f-2305c1781a2d"
@@ -43,6 +44,7 @@ spglib_jll = "ac4a9f1e-bdb2-5204-990c-47c8b2f70d4e"
 
 [compat]
 AbstractFFTs = "1"
+AtomsBase = "0.2.1"
 BlockArrays = "0.16.2"
 Brillouin = "0.5"
 Conda = "1"

docs/src/guide/input_output.jl

@@ -40,7 +40,10 @@ magnetic_moments = load_magnetic_moments(qe_input)
 # not exposed inside the ASE datastructures.
 # See [Creating slabs with ASE](@ref) for more details.
 
-atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="pbe")) for el in atoms]
+atoms = map(atoms) do el
+    @assert el.symbol == :Fe
+    ElementPsp(:Fe, psp=load_psp("hgh/pbe/fe-q16.hgh"))
+end;
 
 # Finally we run the calculation.
 

examples/arbitrary_floattype.jl

@@ -29,7 +29,7 @@ using DFTK
 ## Setup silicon lattice
 a = 10.263141334305942  # lattice constant in Bohr
 lattice = a / 2 .* [[0 1 1.]; [1 0 1.]; [1 1 0.]]
-Si = ElementPsp(:Si, psp=load_psp(:Si, functional="lda"))
+Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
 atoms = [Si, Si]
 positions = [ones(3)/8, -ones(3)/8]
 

examples/ase.jl

@@ -68,10 +68,17 @@ pyimport("ase.io").write("surface.png", surface * (3, 3, 1),
 
 using DFTK
 
-atoms = load_atoms(surface)
-atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="pbe")) for el in atoms]
 positions = load_positions(surface)
-lattice = load_lattice(surface);
+lattice   = load_lattice(surface)
+atoms = map(load_atoms(surface)) do el
+    if el.symbol == :Ga
+        ElementPsp(:Ga, psp=load_psp("hgh/pbe/ga-q3.hgh"))
+    elseif el.symbol == :As
+        ElementPsp(:As, psp=load_psp("hgh/pbe/as-q5.hgh"))
+    else
+        error("Unsupported element: $el")
+    end
+end;
 
 # We model this surface with (quite large a) temperature of 0.01 Hartree
 # to ease convergence. Try lowering the SCF convergence tolerance (`tol`)

examples/atomsbase.jl

@@ -0,0 +1,26 @@
+using DFTK
+using AtomsBase
+using Unitful
+using UnitfulAtomic
+using PyCall
+
+# Construct system in the AtomsBase world
+a = 10.26u"bohr"  # Silicon lattice constant
+lattice = a / 2 * [[0, 1, 1.],  # Lattice as vector of vectors
+                   [1, 0, 1.],
+                   [1, 1, 0.]]
+atoms  = [:Si => ones(3)/8, :Si => -ones(3)/8]
+system = periodic_system(atoms, lattice; fractional=true)
+
+# Use it inside DFTK:
+# Attach pseudopotentials, construct appropriate model,
+# discretise and solve.
+system = attach_psp(system; family="hgh", functional="lda")
+
+model  = model_LDA(system; temperature=1e-3)
+basis  = PlaneWaveBasis(model; Ecut=15, kgrid=[4, 4, 4])
+scfres = self_consistent_field(basis, tol=1e-8)
+
+# Get the DFTK model back into the AtomsBase world:
+newsystem = atomic_system(model)
+@show atomic_number(newsystem)

examples/scf_callbacks.jl

@@ -16,7 +16,7 @@ using DFTK
 using PyCall
 
 silicon = pyimport("ase.build").bulk("Si")
-atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="lda"))
+atoms = [ElementPsp(el.symbol, psp=load_psp("hgh/lda/si-q4.hgh"))
          for el in load_atoms(silicon)]
 positions = load_positions(silicon)
 lattice   = load_lattice(silicon);

src/DFTK.jl

@@ -4,9 +4,9 @@ with plane-wave density-functional theory algorithms.
 """
 module DFTK
 
-using Printf
-using Markdown
 using LinearAlgebra
+using Markdown
+using Printf
 using Requires
 using TimerOutputs
 using spglib_jll
@@ -143,9 +143,11 @@ export guess_density
 export random_density
 export load_psp
 export list_psp
+export attach_psp
 include("guess_density.jl")
 include("pseudo/load_psp.jl")
 include("pseudo/list_psp.jl")
+include("pseudo/attach_psp.jl")
 
 export pymatgen_structure
 export ase_atoms
@@ -154,6 +156,7 @@ export load_atoms
 export load_positions
 export load_magnetic_moments
 export run_wannier90
+include("external/atomsbase.jl")
 include("external/load_from_file.jl")
 include("external/ase.jl")
 include("external/pymatgen.jl")

src/Model.jl

@@ -168,9 +168,11 @@ function Model(lattice::AbstractMatrix{<:Quantity}, args...; kwargs...)
     Model(austrip.(lattice), args...; kwargs...)
 end
 
-normalize_magnetic_moment(::Nothing)  = Vec3{Float64}(zeros(3))
-normalize_magnetic_moment(mm::Number) = Vec3{Float64}(0, 0, mm)
-normalize_magnetic_moment(mm::AbstractVector) = Vec3{Float64}(mm)
+
+normalize_magnetic_moment(::Nothing)::Vec3{Float64}          = (0, 0, 0)
+normalize_magnetic_moment(mm::Number)::Vec3{Float64}         = (0, 0, mm)
+normalize_magnetic_moment(mm::AbstractVector)::Vec3{Float64} = mm
+
 
 """
 :none if no element has a magnetic moment, else :collinear or :full
@@ -211,7 +213,6 @@ function default_symmetries(lattice, atoms, positions, magnetic_moments, spin_po
 end
 
 
-
 """
 Maximal occupation of a state (2 for non-spin-polarized electrons, 1 otherwise).
 """

src/elements.jl

@@ -1,4 +1,5 @@
 import PeriodicTable
+using AtomsBase
 
 # Alias to avoid similarity of elements and Element in DFTK module namespace
 periodic_table = PeriodicTable.elements
@@ -14,7 +15,7 @@ abstract type Element end
 charge_nuclear(::Element) = 0
 
 """Chemical symbol corresponding to an element"""
-atomic_symbol(::Element) = :X
+AtomsBase.atomic_symbol(::Element) = :X
 # The preceeding functions are fallback implementations that should be altered as needed.
 
 """Return the total ionic charge of an atom type (nuclear charge - core electrons)"""
@@ -46,7 +47,7 @@ struct ElementCoulomb <: Element
 end
 charge_ionic(el::ElementCoulomb)   = el.Z
 charge_nuclear(el::ElementCoulomb) = el.Z
-atomic_symbol(el::ElementCoulomb)  = el.symbol
+AtomsBase.atomic_symbol(el::ElementCoulomb) = el.symbol
 
 """
 Element interacting with electrons via a bare Coulomb potential
@@ -87,7 +88,7 @@ function ElementPsp(key; psp)
 end
 charge_ionic(el::ElementPsp)   = el.psp.Zion
 charge_nuclear(el::ElementPsp) = el.Z
-atomic_symbol(el::ElementPsp)  = el.symbol
+AtomsBase.atomic_symbol(el::ElementPsp) = el.symbol
 
 function local_potential_fourier(el::ElementPsp, q::T) where {T <: Real}
     q == 0 && return zero(T)  # Compensating charge background
@@ -108,7 +109,7 @@ struct ElementCohenBergstresser <: Element
 end
 charge_ionic(el::ElementCohenBergstresser)   = 2
 charge_nuclear(el::ElementCohenBergstresser) = el.Z
-atomic_symbol(el::ElementCohenBergstresser)  = el.symbol
+AtomsBase.atomic_symbol(el::ElementCohenBergstresser) = el.symbol
 
 """
 Element where the interaction with electrons is modelled

src/external/ase.jl

@@ -3,8 +3,11 @@ ase_cell(lattice) = pyimport("ase").cell.Cell(Array(lattice)' / austrip(1u"Å"))
 ase_cell(model::Model) = ase_cell(model.lattice)
 
 function ase_atoms(lattice_or_model, atoms, positions, magnetic_moments=[])
-    # Collect Z magnetic moments
-    magmoms = [normalize_magnetic_moment(mom)[3] for mom in magnetic_moments]
+    if isempty(magnetic_moments)  # Collect Z magnetic moments
+        magmoms = nothing
+    else
+        magmoms = [normalize_magnetic_moment(mom)[3] for mom in magnetic_moments]
+    end
     cell    = ase_cell(lattice_or_model)
     symbols = string.(atomic_symbol.(atoms))
     scaled_positions = reduce(hcat, positions)'

src/external/atomsbase.jl

@@ -0,0 +1,116 @@
+using AtomsBase
+# Key functionality to integrate DFTK and AtomsBase
+
+function construct_system(lattice::AbstractMatrix, atoms::Vector, positions::Vector, magnetic_moments::Vector)
+    @assert length(atoms) == length(positions)
+    @assert isempty(magnetic_moments) || length(magnetic_moments) == length(atoms)
+    atomsbase_atoms = map(enumerate(atoms)) do (i, element)
+        kwargs = Dict{Symbol, Any}(:potential => element, )
+        if !isempty(magnetic_moments)
+            kwargs[:magnetic_moment] = normalize_magnetic_moment(magnetic_moments[i])
+        end
+        if element isa ElementPsp
+            kwargs[:pseudopotential] = element.psp.identifier
+        end
+
+        position = lattice * positions[i] * u"bohr"
+        if atomic_symbol(element) == :X  # dummy element ... should solve this upstream
+            Atom(:X, position; atomic_symbol=:X, atomic_number=0, atomic_mass=0u"u", kwargs...)
+        else
+            Atom(atomic_symbol(element), position; kwargs...)
+        end
+    end
+    periodic_system(atomsbase_atoms, collect(eachcol(lattice)) * u"bohr")
+end
+
+"""
+    atomic_system(model::DFTK.Model, magnetic_moments=[])
+
+Construct an AtomsBase atomic system from a DFTK model and associated magnetic moments.
+"""
+function AtomsBase.atomic_system(model::Model, magnetic_moments=[])
+    construct_system(model.lattice, model.atoms, model.positions, magnetic_moments)
+end
+
+"""
+    periodic_system(model::DFTK.Model, magnetic_moments=[])
+
+Construct an AtomsBase atomic system from a DFTK model and associated magnetic moments.
+"""
+function AtomsBase.periodic_system(model::Model, magnetic_moments=[])
+    atomic_system(model, magnetic_moments)
+end
+
+
+function parse_system(system::AbstractSystem{D}) where {D}
+    if !all(periodicity(system))
+        error("DFTK only supports calculations with periodic boundary conditions.")
+    end
+
+    # Parse abstract system and return data required to construct model
+    mtx = austrip.(reduce(hcat, bounding_box(system)))
+    T = eltype(mtx)
+    lattice = zeros(T, 3, 3)
+    lattice[1:D, 1:D] .= mtx
+
+    # Cache for instantiated pseudopotentials. This is done to ensure that identical
+    # atoms are indistinguishable in memory, which is used in the Model constructor
+    # to deduce the atom_groups.
+    cached_pspelements = Dict{String, ElementPsp}()
+    atoms = map(system) do atom
+        if hasproperty(atom, :potential)
+            atom.potential
+        elseif hasproperty(atom, :pseudopotential)
+            get!(cached_pspelements, atom.pseudopotential) do
+                ElementPsp(atomic_symbol(atom); psp=load_psp(atom.pseudopotential))
+            end
+        else
+            ElementCoulomb(atomic_symbol(atom))
+        end
+    end
+
+    positions = map(system) do atom
+        coordinate = zeros(T, 3)
+        coordinate[1:D] = lattice[1:D, 1:D] \ T.(austrip.(position(atom)))
+        Vec3{T}(coordinate)
+    end
+
+    magnetic_moments = map(system) do atom
+        hasproperty(atom, :magnetic_moment) || return nothing
+        getproperty(atom, :magnetic_moment)
+    end
+    if all(m -> isnothing(m) || iszero(m) || isempty(m), magnetic_moments)
+        empty!(magnetic_moments)
+    else
+        magnetic_moments = normalize_magnetic_moment.(magnetic_moments)
+    end
+
+    # TODO Use system to determine n_electrons
+
+    (; lattice, atoms, positions, magnetic_moments)
+end
+
+
+function _call_with_system(f, system::AbstractSystem, args...; kwargs...)
+    @assert !(:magnetic_moments in keys(kwargs))
+    parsed = parse_system(system)
+    f(parsed.lattice, parsed.atoms, parsed.positions, args...;
+      parsed.magnetic_moments, kwargs...)
+end
+
+
+"""
+    Model(system::AbstractSystem; kwargs...)
+
+AtomsBase-compatible Model constructor. Sets structural information (`atoms`, `positions`,
+`lattice`, `n_electrons` etc.) from the passed `system`.
+"""
+Model(system::AbstractSystem; kwargs...) = _call_with_system(Model, system; kwargs...)
+
+
+# Generate equivalent functions for AtomsBase
+for fun in (:model_atomic, :model_DFT, :model_LDA, :model_PBE, :model_SCAN)
+    @eval function $fun(system::AbstractSystem, args...; kwargs...)
+        _call_with_system($fun, system, args...; kwargs...)
+    end
+end

src/guess_density.jl

@@ -16,11 +16,12 @@ end
 
 
 @doc raw"""
-    guess_density(basis, magnetic_moments)
+    guess_density(basis, magnetic_moments=[])
+    guess_density(basis, system)
 
 Build a superposition of atomic densities (SAD) guess density.
 
-We take for the guess density a gaussian centered around the atom, of
+We take for the guess density a Gaussian centered around the atom, of
 length specified by `atom_decay_length`, normalized to get the right number of electrons
 ```math
 \hat{ρ}(G) = Z \exp\left(-(2π \text{length} |G|)^2\right)
@@ -31,7 +32,11 @@ The magnetic moments should be specified in units of ``μ_B``.
 function guess_density(basis::PlaneWaveBasis, magnetic_moments=[])
     guess_density(basis, basis.model.atoms, basis.model.positions, magnetic_moments)
 end
-@timing function guess_density(basis::PlaneWaveBasis, atoms, positions, magnetic_moments=[])
+function guess_density(basis::PlaneWaveBasis, system::AbstractSystem)
+    parsed = parse_system(system)
+    guess_density(basis, parsed.atoms, parsed.positions, parsed.magnetic_moments)
+end
+@timing function guess_density(basis::PlaneWaveBasis, atoms, positions, magnetic_moments)
     ρtot = _guess_total_density(basis, atoms, positions)
     if basis.model.n_spin_components == 1
         ρspin = nothing

src/interpolation_transfer.jl

@@ -1,4 +1,5 @@
-using Interpolations
+import Interpolations
+import Interpolations: interpolate, extrapolate, scale, BSpline, Quadratic, OnCell
 using SparseArrays
 
 """
@@ -46,7 +47,7 @@ function interpolate_density(ρ_in::AbstractArray, grid_in, grid_out, lattice_in
 
     # interpolate ρ_in_supercell from grid grid_supercell to grid_out
     axes_in = (range(0, 1, length=grid_supercell[i]+1)[1:end-1] for i=1:3)
-    itp = interpolate(ρ_in_supercell, BSpline(Quadratic(Periodic(OnCell()))))
+    itp = interpolate(ρ_in_supercell, BSpline(Quadratic(Interpolations.Periodic(OnCell()))))
     sitp = scale(itp, axes_in...)
     ρ_interp = extrapolate(sitp, Periodic())
     ρ_out = similar(ρ_in, grid_out)