update DFTK integration for DFTK 0.5.0

Author: jrdegreeff

Project.toml

@@ -1,7 +1,7 @@
 name = "Atomistic"
 uuid = "3e762ff9-b8d8-45cf-a902-edcbc7124e71"
 authors = ["Jeremiah DeGreeff <[email protected]>"]
-version = "0.3.4"
+version = "0.3.5"
 
 [deps]
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
@@ -21,7 +21,7 @@ UnitfulRecipes = "42071c24-d89e-48dd-8a24-8a12d9b8861f"
 
 [compat]
 AtomsBase = "0.2"
-DFTK = "0.4"
+DFTK = "0.5"
 Distances = "0.10"
 InteratomicPotentials = "0.1"
 Molly = "0.9"

src/Atomistic.jl

@@ -18,7 +18,6 @@ using NBodySimulator
 using PeriodicTable
 using Plots, UnitfulRecipes
 
-import Base: @kwdef
 import DFTK: Mixing
 import NBodySimulator: Body, BoundaryConditions, NullThermostat, SimulationResult, Thermostat
 import Plots: Plot

src/integrations/dftk_integration.jl

@@ -1,84 +1,28 @@
-# Integrations with DFTK.jl
-
-# -----------------------------------------------------------------------------
-# Integration with AtomsBase 
-# -----------------------------------------------------------------------------
-
-# ! all functions in this section are copied directly from experimental DFTK code
-# ! this code will be deleted once it is live in DFTK
-
-function parse_system(system::AbstractSystem{3})
-    if !all(periodicity(system))
-        error("DFTK only supports calculations with periodic boundary conditions.")
-    end
-
-    # Parse abstract system and return data required to construct model
-    lattice = austrip.(hcat(bounding_box(system)...))
-    T = eltype(lattice)
-
-    # Cache for instantiated pseudopotentials
-    # (such that the respective objects are indistinguishable)
-    cached_pseudos = Dict{String,Any}()
-    atoms = map(system) do atom
-        if hasproperty(atom, :potential)
-            potential = atom.potential
-        elseif hasproperty(atom, :pseudopotential)
-            pspkey = atom.pseudopotential
-            potential = get!(cached_pseudos, pspkey) do
-                ElementPsp(AtomsBase.atomic_symbol(atom); psp = load_psp(pspkey))
-            end
-        else
-            potential = ElementCoulomb(AtomsBase.atomic_symbol(atom))
-        end
-
-        potential => SVector{3,T}(lattice \ T.(austrip.(position(atom))))
-    end
-
-    oldatoms = oldatoms_from_new(atoms)
-
-    (; lattice, atoms = oldatoms)
-end
-function oldatoms_from_new(atomic_potentials)
-    potentials = first.(atomic_potentials)
-    potential_groups = [findall(Ref(pot) .== potentials) for pot in Set(potentials)]
-    [first(atomic_potentials[first(group)]) => last.(atomic_potentials[group]) for group in potential_groups]
-end
-function DFTK.model_LDA(system::AbstractSystem; kwargs...)
-    parsed = parse_system(system)
-    model_LDA(parsed.lattice, parsed.atoms; kwargs...)
-end
-
 # -----------------------------------------------------------------------------
 # Integration with InteratomicPotentials
 # -----------------------------------------------------------------------------
 
 # ! This integration should ultimately move to the DFTK package itself
 
-@kwdef struct DFTKPotential <: ArbitraryPotential
+struct DFTKPotential <: ArbitraryPotential
     Ecut::Real
     kgrid::AbstractVector{<:Integer}
-    n_bands::Union{Integer,Nothing} = nothing
-    tol::Union{AbstractFloat,Nothing} = nothing
-    damping::Union{AbstractFloat,Nothing} = nothing
-    mixing::Union{Mixing,Nothing} = nothing
-    previous_scfres::Ref{Any} = Ref{Any}()
-    potential_energy_cache::Dict{Float64,Float64} = Dict{Float64,Float64}()
+    scf_args::Dict{Symbol,Any}
+    previous_scf::Ref{Any}
 end
 function DFTKPotential(Ecut::Real, kgrid::AbstractVector{<:Integer}; kwargs...)
-    DFTKPotential(; Ecut = Ecut, kgrid = kgrid, kwargs...)
+    DFTKPotential(Ecut, kgrid, Dict{Symbol,Any}(kwargs...), Ref{Any}())
 end
 function DFTKPotential(Ecut::Unitful.Energy, kgrid::AbstractVector{<:Integer}; kwargs...)
-    DFTKPotential(; Ecut = austrip(Ecut), kgrid = kgrid, kwargs...)
+    DFTKPotential(austrip(Ecut), kgrid, Dict{Symbol,Any}(kwargs...), Ref{Any}())
 end
 
 function InteratomicPotentials.energy_and_force(system::AbstractSystem, potential::DFTKPotential)
     model = model_LDA(system)
     basis = PlaneWaveBasis(model; Ecut = potential.Ecut, kgrid = potential.kgrid)
 
-    args = (f => getfield(potential, f) for f ∈ (:n_bands, :tol, :damping, :mixing) if !isnothing(getfield(potential, f)))
-    extra_args = isassigned(potential.previous_scfres) ? (ψ = potential.previous_scfres[].ψ, ρ = potential.previous_scfres[].ρ) : (;)
-    scfres = self_consistent_field(basis; args..., extra_args...)
-    potential.previous_scfres[] = scfres
-    # TODO: support multiple species
-    (; e = scfres.energies.total, f = compute_forces_cart(scfres)[1])
+    extra_args = isassigned(potential.previous_scf) ? (ψ = potential.previous_scf[].ψ, ρ = potential.previous_scf[].ρ) : (;)
+    scfres = self_consistent_field(basis; potential.scf_args..., extra_args...)
+    potential.previous_scf[] = scfres # cache previous scf as starting point for next calculation
+    (; e = scfres.energies.total, f = compute_forces_cart(scfres))
 end

test/integrations/dftk_integration.jl

@@ -16,15 +16,15 @@
     @test eandf1.e isa Float64
     @test eandf1.f isa AbstractVector{<:SVector{3,<:Float64}}
 
-    @test potential1.previous_scfres[].energies.total == eandf1.e
-    @test compute_forces_cart(potential1.previous_scfres[])[1] == eandf1.f
+    @test potential1.previous_scf[].energies.total == eandf1.e
+    @test compute_forces_cart(potential1.previous_scf[]) == eandf1.f
 
     system2 = System(system1)
     eandf2 = energy_and_force(system2, potential2)
 
     @test eandf2.e isa Float64
     @test eandf2.f isa AbstractVector{<:SVector{3,<:Float64}}
 
-    @test potential2.previous_scfres[].energies.total == eandf2.e
-    @test compute_forces_cart(potential2.previous_scfres[])[1] == eandf2.f
+    @test potential2.previous_scf[].energies.total == eandf2.e
+    @test compute_forces_cart(potential2.previous_scf[]) == eandf2.f
 end