Lattice Optimization

Project.toml

@@ -1,14 +1,21 @@
 name = "GeometryOptimization"
 uuid = "673bf261-a53d-43b9-876f-d3c1fc8329c2"
 authors = ["JuliaMolSim community"]
-version = "0.0.1"
+version = "0.0.2"
 
 [deps]
+ASEconvert = "3da9722f-58c2-4165-81be-b4d7253e8fd2"
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1"
-EmpiricalPotentials = "38527215-9240-4c91-a638-d4250620c9e2"
+AtomsIO = "1692102d-eeb4-4df9-807b-c9517f998d44"
+AtomsIOPython = "9e4c859b-2281-48ef-8059-f50fe53c37b0"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"

examples/Al_strain.jl

@@ -0,0 +1,84 @@
+using DFTK: occupation
+#= Test Geometry Optimization on an aluminium supercell.
+=#
+using LinearAlgebra
+using DFTK
+using ASEconvert
+using LazyArtifacts
+using AtomsCalculators
+using Unitful
+using UnitfulAtomic
+using Random
+using OptimizationOptimJL
+
+using GeometryOptimization
+
+
+function build_al_supercell(rep=1)
+    pseudodojo_psp = artifact"pd_nc_sr_lda_standard_0.4.1_upf/Al.upf"
+    a = 7.65339 # true lattice constant.
+    lattice = a * Matrix(I, 3, 3)
+    Al = ElementPsp(:Al; psp=load_psp(pseudodojo_psp))
+    atoms     = [Al, Al, Al, Al]
+    positions = [[0.0, 0.0, 0.0], [0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]]
+    unit_cell = periodic_system(lattice, atoms, positions)
+
+    # Make supercell in ASE:
+    # We convert our lattice to the conventions used in ASE, make the supercell
+    # and then convert back ...
+    supercell_ase = convert_ase(unit_cell) * pytuple((rep, 1, 1))
+    supercell     = pyconvert(AbstractSystem, supercell_ase)
+
+    # Unfortunately right now the conversion to ASE drops the pseudopotential information,
+    # so we need to reattach it:
+    supercell = attach_psp(supercell; Al=pseudodojo_psp)
+    return supercell
+end;
+
+al_supercell = build_al_supercell(1)
+
+# Create a simple calculator for the model.
+model_kwargs = (; functionals = [:lda_x, :lda_c_pw], temperature = 1e-4)
+basis_kwargs = (; kgrid = [6, 6, 6], Ecut = 30.0)
+scf_kwargs = (; tol = 1e-6)
+calculator = DFTKCalculator(; model_kwargs, basis_kwargs, scf_kwargs, verbose=true)
+
+energy_true = AtomsCalculators.potential_energy(al_supercell, calculator)
+
+# Starting position is a random perturbation of the equilibrium one.
+Random.seed!(1234)
+x0 = vcat(position(al_supercell)...)
+σ = 0.5u"angstrom"; x0_pert = x0 + σ * rand(Float64, size(x0))
+al_supercell = update_not_clamped_positions(al_supercell, x0_pert)
+energy_pert = AtomsCalculators.potential_energy(al_supercell, calculator)
+
+println("Initial guess distance (norm) from true parameters $(norm(x0 - x0_pert)).")
+println("Initial regret $(energy_pert - energy_true).")
+
+optim_options = (f_tol=1e-6, iterations=6, show_trace=true)
+
+results = minimize_energy!(al_supercell, calculator; optim_options...)
+println(results)
+
+""" Returns the index of the highest occupied band. 
+    `atol` specifies the (fractional) occupations below which 
+    a band is considered unoccupied.
+"""
+function valence_band_index(occupations; atol=1e-36)
+    filter = x -> isapprox(x, 0.; atol)
+    maximum(maximum.(findall.(!filter, occupations)))
+end
+
+function band_gaps(scfres)
+    vi = valence_band_index(occupations; atol)
+    # If the system is metallic, by convenction band gap is zero.
+    if DFTK.is_metal(scfres.eigenvalues, scfres.εF; tol=1e-4)
+        return (; εMax_valence=0.0u"hartree", εMin_conduction=0.0u"hartree",
+                direct_bandgap=0.0u"hartree", valence_band_index=vi)
+    else
+        εMax_valence = maximum([εk[vi] for εk in scfres.eigenvalues]) * u"hartree"
+        εMin_conduction = minimum([εk[vi + 1] for εk in scfres.eigenvalues]) * u"hartree"
+        direct_bandgap = minimum([εk[vi + 1] - εk[vi] for εk in scfres.eigenvalues]) * u"hartree"
+        return (; εMax_valence, εMin_conduction, direct_bandgap, valence_band_index=vi)
+    end
+end

examples/Al_supercell.jl

@@ -40,7 +40,7 @@ al_supercell = build_al_supercell(1)
 model_kwargs = (; functionals = [:lda_x, :lda_c_pw], temperature = 1e-4)
 basis_kwargs = (; kgrid = [6, 6, 6], Ecut = 30.0)
 scf_kwargs = (; tol = 1e-6)
-calculator = DFTKCalculator(al_supercell; model_kwargs, basis_kwargs, scf_kwargs, verbose=true)
+calculator = DFTKCalculator(; model_kwargs, basis_kwargs, scf_kwargs, verbose=true)
 
 energy_true = AtomsCalculators.potential_energy(al_supercell, calculator)
 

examples/Al_variable_cell.jl

@@ -0,0 +1,46 @@
+using LinearAlgebra
+using DFTK
+using ASEconvert
+using LazyArtifacts
+using AtomsBase
+using AtomsCalculators
+using Unitful
+using UnitfulAtomic
+using Random
+using OptimizationOptimJL
+using ComponentArrays
+
+using GeometryOptimization
+
+
+function build_al_supercell(rep=1)
+    pseudodojo_psp = artifact"pd_nc_sr_lda_standard_0.4.1_upf/Al.upf"
+    a = 7.65339 # true lattice constant.
+    lattice = a * Matrix(I, 3, 3)
+    Al = ElementPsp(:Al; psp=load_psp(pseudodojo_psp))
+    atoms     = [Al, Al, Al, Al]
+    positions = [[0.0, 0.0, 0.0], [0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]]
+    unit_cell = periodic_system(lattice, atoms, positions)
+
+    # Make supercell in ASE:
+    # We convert our lattice to the conventions used in ASE, make the supercell
+    # and then convert back ...
+    supercell_ase = convert_ase(unit_cell) * pytuple((rep, 1, 1))
+    supercell     = pyconvert(AbstractSystem, supercell_ase)
+
+    # Unfortunately right now the conversion to ASE drops the pseudopotential information,
+    # so we need to reattach it:
+    supercell = attach_psp(supercell; Al=pseudodojo_psp)
+    return supercell
+end;
+
+al_supercell = build_al_supercell(1)
+
+# Create a simple calculator for the model.
+model_kwargs = (; functionals = [:lda_x, :lda_c_pw], temperature = 1e-2)
+basis_kwargs = (; kgrid = [2, 2, 2], Ecut = 10.0)
+scf_kwargs = (; tol = 1e-3)
+calculator = DFTKCalculator(; model_kwargs, basis_kwargs, scf_kwargs, verbose=true)
+
+optim_options = (f_tol=1e-6, iterations=6, show_trace=true)
+results = minimize_energy!(al_supercell, calculator; procedure="vc_relax", optim_options...)

examples/Si_verification.jl

@@ -0,0 +1,43 @@
+# Verify correctness of lattice optimization on silicon.
+using DFTK
+using AtomsBase
+using AtomsCalculators
+using ComponentArrays
+using LazyArtifacts
+using Random
+using Unitful
+using UnitfulAtomic
+using OptimizationOptimJL
+using GeometryOptimization
+
+
+lattice = [0.0  5.131570667152971 5.131570667152971;
+           5.131570667152971 0.0 5.131570667152971;
+           5.131570667152971 5.131570667152971  0.0]
+hgh_lda_family = artifact"hgh_lda_hgh"
+psp_hgh = joinpath(hgh_lda_family, "si-q4.hgh")
+
+positions = [ones(3)/8, -ones(3)/8]
+atoms = fill(ElementPsp(:Si; psp=load_psp(psp_hgh)), 2)
+system = periodic_system(lattice, atoms, positions)
+
+# Create a simple calculator for the model.
+model_kwargs = (; functionals = [:lda_x, :lda_c_pw], temperature = 1e-5)
+basis_kwargs = (; kgrid = [5, 5, 5], Ecut = 30.0)
+scf_kwargs = (; tol = 1e-6)
+calculator = DFTKCalculator(; model_kwargs, basis_kwargs, scf_kwargs, verbose=true)
+
+# Perturb unit cell.
+Random.seed!(1234)
+σ = 0.2u"bohr"
+bounding_box_pert = [v + σ * rand(Float64, size(v)) for v in bounding_box(system)]
+system_pert = update_positions(system, position(system); bounding_box=bounding_box_pert)
+
+
+using LineSearches
+linesearch =  BackTracking(c_1= 1e-4, ρ_hi= 0.8, ρ_lo= 0.1, order=2, maxstep=Inf)
+solver = OptimizationOptimJL.LBFGS(; linesearch)
+optim_options = (; solver, f_tol=1e-10, g_tol=1e-5, iterations=30,
+                 show_trace=true, store_trace = true, allow_f_increases=true)
+
+results = minimize_energy!(system_pert, calculator; procedure="vc_relax", optim_options...)

src/GeometryOptimization.jl

@@ -1,5 +1,6 @@
 module GeometryOptimization
 
+using LinearAlgebra
 using StaticArrays
 using Optimization
 using OptimizationOptimJL
@@ -9,6 +10,7 @@ using Unitful
 using UnitfulAtomic
 
 include("atomsbase_interface.jl")
+include("strain.jl")
 include("optimization.jl")
 
 end

src/atomsbase_interface.jl

@@ -9,15 +9,35 @@ export update_positions, update_not_clamped_positions, clamp_atoms
 
 @doc raw"""
 Creates a new system based on ``system`` but with atoms positions updated 
-to the ones provided.
+to the ones provided. Can also update lattice vectors if `bounding_box` is provided.
 
 """
-function update_positions(system, positions::AbstractVector{<:AbstractVector{<:Unitful.Length}})
+function update_positions(system, positions::AbstractVector{<:AbstractVector{<:Unitful.Length}};
+                          bounding_box=bounding_box(system))
     particles = [Atom(atom; position) for (atom, position) in zip(system, positions)]
-    AbstractSystem(system; particles)
+    AbstractSystem(system; particles, bounding_box)
 end
 
 @doc raw"""
+Creates a new system based on ``system`` but with atoms positions 
+updated to the ones provided and lattice vectors deformed according 
+to the provided strain.
+New generalized coordinates should be provided in a ComponentArray with 
+component `atoms` and `strain`.
+
+"""
+function update_positions(system, positions::ComponentVector)
+    deformation_tensor = I + voigt_to_full(austrip.(positions.strain))
+    # TODO: Do we want to apply the strain to the atoms too?
+    particles = [Atom(atom; position = deformation_tensor * position) for (atom, position)
+                 in zip(system, collect.(positions.atoms))]
+
+    bbox = eachcol(deformation_tensor * bbox_to_matrix(bounding_box(system)))
+    AbstractSystem(system; particles, bounding_box=bbox)
+end
+
+@doc raw"""
+=======
 Creates a new system based on ``system`` where the non clamped positions are
 updated to the ones provided (in the order in which they appear in the system).
 """
@@ -30,6 +50,23 @@ function update_not_clamped_positions(system, positions::AbstractVector{<:Unitfu
 end
 
 @doc raw"""
+Creates a new system based on ``system`` where the non clamped positions and 
+lattice vectors are updated to the ones provided.
+Note that the `atoms`component of `positions` should be a vector of 
+coordinates and that the `strain` component should be a 6-vector.
+
+"""
+function update_not_clamped_positions(system, positions::ComponentVector)
+    mask = not_clamped_mask(system)
+    new_positions = deepcopy(position(system))
+    atoms_positions = collect(positions.atoms)
+    new_positions[mask] = reinterpret(reshape, SVector{3, eltype(atoms_positions)},
+                                      reshape(atoms_positions, 3, :))
+    update_positions(system, ComponentVector(; atoms=new_positions, positions.strain))
+end
+
+@doc raw"""
+=======
 Returns a mask for selecting the not clamped atoms in the system.
 
 """

src/optimization.jl

@@ -2,26 +2,38 @@
 # Note that by default all particles in the system are assumed optimizable.
 # IMPORTANT: Note that we always work in cartesian coordinates.
 =#
-
+using DFTK
 export minimize_energy!
 
 
+function update_state(original_system, new_system, state)
+    if bounding_box(original_system) != bounding_box(new_system)
+        return DFTK.DFTKState()
+    else
+        return state
+    end
+end
+
 """ 
-By default we work in cartesian coordinaes.
+By default we work in cartesian coordinates.
 Note that internally, when optimizing the cartesian positions, atomic units 
 are used.
 """
-function Optimization.OptimizationFunction(system, calculator; kwargs...)
+function Optimization.OptimizationFunction(system, calculator; pressure=0.0, kwargs...)
     mask = not_clamped_mask(system)  # mask is assumed not to change during optim.
 
-    f = function(x::AbstractVector{<:Real}, p)
+    # TODO: Note that this function will dispatch appropriately when called with 
+    # a Component vector.
+    f = function(x, p)
         new_system = update_not_clamped_positions(system, x * u"bohr")
-        energy = AtomsCalculators.potential_energy(new_system, calculator; kwargs...)
+        state = update_state(system, new_system, calculator.state)
+        energy = AtomsCalculators.potential_energy(new_system, calculator; state, kwargs...)
         austrip(energy)
     end
 
-    g! = function(G::AbstractVector{<:Real}, x::AbstractVector{<:Real}, p)
+    function g!(G, x, p)
         new_system = update_not_clamped_positions(system, x * u"bohr")
+        # TODO: Determine if here we need a call to update_state.
         energy = AtomsCalculators.potential_energy(new_system, calculator; kwargs...)
 
         forces = AtomsCalculators.forces(new_system, calculator; kwargs...)
@@ -31,13 +43,34 @@ function Optimization.OptimizationFunction(system, calculator; kwargs...)
         # NOTE: minus sign since forces are opposite to gradient.
         G .= - austrip.(forces_concat)
     end
+    function g!(G::ComponentVector, x::ComponentVector, p)
+        deformation_tensor = I + voigt_to_full(austrip.(x.strain))
+        new_system = update_not_clamped_positions(system, x * u"bohr")
+
+        state = update_state(system, new_system, calculator.state)
+        forces = AtomsCalculators.forces(new_system, calculator; state, kwargs...)
+        # Translate the forces vectors on each particle to a single gradient for the optimization parameter.
+        forces_concat = collect(Iterators.flatten([deformation_tensor * f for f in forces[mask]]))
+
+        # NOTE: minus sign since forces are opposite to gradient.
+        G.atoms .= - austrip.(forces_concat)
+        virial = AtomsCalculators.virial(new_system, calculator)
+        G.strain .= - full_to_voigt(virial / deformation_tensor)
+    end
     OptimizationFunction(f; grad=g!)
 end
 
-function minimize_energy!(system, calculator; solver=Optim.LBFGS(), kwargs...)
+function minimize_energy!(system, calculator; pressure=0.0, procedure="relax", 
+                          solver=Optim.LBFGS(), kwargs...)
     # Use current system parameters as starting positions.
-    x0 = austrip.(not_clamped_positions(system)) # Optim modifies x0 in-place, so need a mutable type.
-    f_opt = OptimizationFunction(system, calculator)
+    if procedure == "relax"
+        x0 = austrip.(not_clamped_positions(system))
+    elseif procedure == "vc_relax"
+        x0 = ComponentVector(atoms = austrip.(reduce(vcat, position(system))), strain = zeros(6))
+    else
+        print("Error: unknown optimization procedure. Please use one of [`relax`, `vc_relax`].")
+    end
+    f_opt = OptimizationFunction(system, calculator; pressure)
     problem = OptimizationProblem(f_opt, x0, nothing)  # Last argument needed in Optimization.jl.
     solve(problem, solver; kwargs...)
 end