Make DFTK integration an extension module

Project.toml

@@ -12,9 +12,16 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
+[weakdeps]
+DFTK = "acf6eb54-70d9-11e9-0013-234b7a5f5337"
+
+[extensions]
+InteratomicPotentialsDFTKExt = "DFTK"
+
 [compat]
 AtomsBase = "0.2, 0.3"
 Distances = "0.10"
+DFTK = "0.6"
 NearestNeighbors = "0.4.9"
 StaticArrays = "1"
 Unitful = "1"

ext/InteratomicPotentialsDFTKExt.jl

@@ -0,0 +1,30 @@
+module InteratomicPotentialsDFTKExt
+
+# Extension module compatibility
+if isdefined(Base, :get_extension)
+    using AtomsBase
+    using DFTK
+    using InteratomicPotentials: energy_and_force
+    using Unitful
+    using UnitfulAtomic
+else
+    using ..AtomsBase
+    using ..DFTK
+    using ..InteratomicPotentials: energy_and_force
+    using ..Unitful
+    using ..UnitfulAtomic
+end
+
+function energy_and_force(system::AbstractSystem, potential::DFTKPotential)
+    model = model_DFT(system, potential.functionals; potential.model_kwargs...)
+    basis = PlaneWaveBasis(model; potential.basis_kwargs...)
+
+    scfres = self_consistent_field(basis; potential.scf_kwargs...)
+    # cache ψ and ρ as starting point for next calculation
+    potential.scf_kwargs[:ψ] = scfres.ψ
+    potential.scf_kwargs[:ρ] = scfres.ρ
+
+    (; e=scfres.energies.total * u"hartree", f=compute_forces_cart(scfres) * u"hartree/bohr")
+end
+
+end

src/InteratomicPotentials.jl

@@ -1,4 +1,5 @@
 module InteratomicPotentials
+using Requires  # Only needed in Julia < 1.9
 
 using AtomsBase
 using Base.Threads
@@ -16,4 +17,13 @@ include("nnlist.jl")
 include("api.jl")
 include("types.jl")
 
+# Requires-based dependency management for Julia < 1.9
+if !isdefined(Base, :get_extension)
+    function __init__()
+        @require DFTK="acf6eb54-70d9-11e9-0013-234b7a5f5337" begin
+            include("../ext/InteratomicPotentialsDFTKExt.jl")
+        end
+    end
+end
+
 end

src/types.jl

@@ -68,4 +68,6 @@ abstract type LinearBasisPotential{P<:NamedTuple, HP<:NamedTuple} <: BasisPotent
 
 Abstract type to define methods for producing a set of local and force descriptors for a given configuration. Examples include the Atomic Cluster Expansion, SOAP descriptors, and SNAP descriptors. See the package InteratomicBasisPotentials.jl for implementation.
 """
-abstract type BasisSystem end 
+abstract type BasisSystem end 
+
+include("types/dftk_potential.jl")

src/types/dftk_potential.jl

@@ -0,0 +1,31 @@
+export DFTKPotential
+
+"""
+    DFTKPotential
+
+Integration to use DFTK as an interatomic potential. Make sure to issue
+a `using DFTK` to ensure the `energy_and_force` method using this type
+is available.
+"""
+Base.@kwdef struct DFTKPotential <: NonTrainablePotential
+    functionals::Vector{Symbol}    = [:gga_x_pbe, :gga_c_pbe]  # default to model_PBE
+    model_kwargs::Dict{Symbol,Any} = Dict{Symbol,Any}()
+    basis_kwargs::Dict{Symbol,Any} = Dict{Symbol,Any}()
+    scf_kwargs::Dict{Symbol,Any}   = Dict{Symbol,Any}()
+end
+
+"""
+    DFTKPotential(Ecut, kgrid; kwargs...)
+
+Construct a DFTK potential. `Ecut` is the kinetic energy cutoff,
+`kgrid` the k-point grid and other `kwargs` are directly passed
+to the `DFTKPotential` constructor. See
+[https://docs.dftk.org](https://docs.dftk.org) for more details
+on using DFTK.
+"""
+function DFTKPotential(Ecut, kgrid; kwargs...)
+    p = DFTKPotential(; kwargs...)
+    p.basis_kwargs[:Ecut]  = Ecut
+    p.basis_kwargs[:kgrid] = kgrid
+    p
+end

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
+DFTK = "acf6eb54-70d9-11e9-0013-234b7a5f5337"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/integration/dftk.jl

@@ -0,0 +1,33 @@
+@testset "DFTK Tests" begin
+    using AtomsBase
+    using DFTK
+    using InteratomicPotentials
+    using StaticArrays
+    using Unitful
+    using UnitfulAtomic
+
+    functionals = [:lda_x, :lda_c_pw]
+    scf_kwargs  = Dict(:damping => 0.7, :tol => 1e-4)
+    potential   = DFTKPotential(5u"hartree", [1, 1, 1]; functionals, scf_kwargs)
+
+    particles = [
+        :Ar => [21.0, 21.0, 21.0]u"bohr",
+        :Ar => [7.0, 21.0, 21.0]u"bohr",
+        :Ar => [21.0, 7.0, 21.0]u"bohr",
+        :Ar => [7.0, 7.0, 21.0]u"bohr",
+        :Ar => [21.0, 21.0, 7.0]u"bohr",
+        :Ar => [7.0, 21.0, 7.0]u"bohr",
+        :Ar => [21.0, 7.0, 7.0]u"bohr",
+        :Ar => [7.0, 7.0, 7.0]u"bohr"
+    ]
+    box = [[28.0, 0.0, 0.0], [0.0, 28.0, 0.0], [0.0, 0.0, 28.0]]u"bohr"
+    system = attach_psp(periodic_system(particles, box); Ar="hgh/lda/ar-q8.hgh")
+
+    eandf = energy_and_force(system, potential)
+    @test eandf.e isa Unitful.Energy
+    @test eandf.f isa AbstractVector{<:SVector{3,<:Unitful.Force}}
+    @test austrip(abs(eandf.e - -164.17u"hartree")) < 0.1
+
+    @test haskey(potential.scf_kwargs, :ψ)
+    @test haskey(potential.scf_kwargs, :ρ)
+end

test/runtests.jl

@@ -41,5 +41,6 @@ include("mocks.jl")
     @time @testset "Integration Tests" begin
         include("integration/lj_clusters/lj_150.jl")
         include("integration/lj_clusters/lj_1000.jl")
+        include("integration/dftk.jl")
     end
 end