Merge pull request #1 from CedricTravelletti/main

Basic implementation of geometry optimization interface.

Author: cortner

.github/workflows/CI.yml

@@ -19,7 +19,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
           - '1.9'
           - 'nightly'
         os:

Project.toml

@@ -3,11 +3,36 @@ uuid = "673bf261-a53d-43b9-876f-d3c1fc8329c2"
 authors = ["JuliaMolSim community"]
 version = "0.0.1"
 
+[deps]
+AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
+AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
+UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
+
 [compat]
-julia = "1"
+AtomsBase = "0.3"
+AtomsCalculators = "0.1"
+Optimization = "3.20"
+OptimizationOptimJL = "0.1"
+StaticArrays = "1.8"
+TestItemRunner = "0.2"
+Unitful = "1.19"
+UnitfulAtomic = "1.0"
+julia = "1.9"
+DFTK = "0.6"
+ASEconvert = "0.1"
+EmpiricalPotentials = "0.0.1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
+DFTK = "acf6eb54-70d9-11e9-0013-234b7a5f5337"
+ASEconvert = "3da9722f-58c2-4165-81be-b4d7253e8fd2"
+EmpiricalPotentials = "38527215-9240-4c91-a638-d4250620c9e2"
 
 [targets]
-test = ["Test"]
+examples = ["DFTK", "ASEconvert", "EmpiricalPotentials"]
+test = ["Test", "TestItemRunner"]

examples/Al_supercell.jl

@@ -0,0 +1,60 @@
+#= 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(al_supercell; 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)

examples/H2.jl

@@ -0,0 +1,32 @@
+using Printf
+using LinearAlgebra
+using DFTK
+using Unitful
+using UnitfulAtomic
+using OptimizationOptimJL
+
+using GeometryOptimization
+
+
+a = 10.                  # Big box around the atoms.
+lattice = a * I(3)
+H = ElementPsp(:H; psp=load_psp("hgh/lda/h-q1"));
+atoms = [H, H];
+positions = [[0, 0, 0], [0, 0, .16]]
+system = periodic_system(lattice, atoms, positions)
+
+# Set everything to optimizable.
+system = clamp_atoms(system, [1])
+
+# Create a simple calculator for the model.
+model_kwargs = (; functionals = [:lda_x, :lda_c_pw])
+basis_kwargs = (; kgrid = [1, 1, 1], Ecut = 10.0)
+scf_kwargs = (; tol = 1e-7)
+calculator = DFTKCalculator(system; model_kwargs, basis_kwargs, scf_kwargs)
+
+solver = OptimizationOptimJL.LBFGS()
+optim_options = (f_tol=1e-32, iterations=20, show_trace=true)
+
+results = minimize_energy!(system, calculator; solver=solver, optim_options...)
+println(results)
+@printf "Bond length: %3f bohrs.\n" norm(results.minimizer[1:end])

examples/H2_LJ.jl

@@ -0,0 +1,23 @@
+#= Test Geometry Optimization on an aluminium supercell.
+=#
+using LinearAlgebra
+using EmpiricalPotentials
+using Unitful
+using UnitfulAtomic
+using OptimizationOptimJL
+using AtomsBase
+
+using GeometryOptimization
+
+
+bounding_box = 10.0u"angstrom" .* [[1, 0, 0.], [0., 1, 0], [0., 0, 1]]
+atoms = [:H => [0, 0, 0.0]u"bohr", :H => [0, 0, 1.9]u"bohr"]
+system = periodic_system(atoms, bounding_box)
+
+lj = LennardJones(-1.17u"hartree", 0.743u"angstrom", 1, 1, 0.6u"nm")
+
+solver = OptimizationOptimJL.LBFGS()
+optim_options = (f_tol=1e-6, iterations=100, show_trace=false)
+
+results = minimize_energy!(system, lj; solver=solver, optim_options...)
+println("Bond length: $(norm(results.minimizer[1:3] - results.minimizer[4:end])).")

examples/TiAl_EmpiricalPotentials.jl

@@ -0,0 +1,26 @@
+using AtomsBase
+using AtomsCalculators
+using EmpiricalPotentials
+using ExtXYZ
+using Unitful
+using UnitfulAtomic
+using OptimizationOptimJL
+
+using GeometryOptimization
+
+fname = joinpath(pkgdir(EmpiricalPotentials), "data/", "TiAl-1024.xyz")
+data = ExtXYZ.load(fname) |> FastSystem
+
+lj = LennardJones(-1.0u"meV", 3.1u"Å",  13, 13, 6.0u"Å")
+
+# Convert to AbstractSystem, so we have the `particles` attribute.
+particles = map(data) do atom
+    Atom(; pairs(atom)...)
+end
+system = AbstractSystem(data; particles)
+
+solver = OptimizationOptimJL.LBFGS()
+optim_options = (f_tol=1e-8, g_tol=1e-8, iterations=10, show_trace=true)
+
+results = minimize_energy!(system, lj; solver, optim_options...)
+println(results)

src/GeometryOptimization.jl

@@ -1,5 +1,14 @@
 module GeometryOptimization
 
-# Write your package code here.
+using StaticArrays
+using Optimization
+using OptimizationOptimJL
+using AtomsBase
+using AtomsCalculators
+using Unitful
+using UnitfulAtomic
+
+include("atomsbase_interface.jl")
+include("optimization.jl")
 
 end

src/atomsbase_interface.jl

@@ -0,0 +1,57 @@
+#
+# Interface between AtomsBase.jl and GeometryOptimization.jl that provides 
+# utility functions for manipulating systems.
+#
+# IMPORTANT: Note that we always work in cartesian coordinates.
+#
+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.
+
+"""
+function update_positions(system, positions::AbstractVector{<:AbstractVector{<:Unitful.Length}})
+    particles = [Atom(atom; position) for (atom, position) in zip(system, positions)]
+    AbstractSystem(system; particles)
+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).
+"""
+function update_not_clamped_positions(system, positions::AbstractVector{<:Unitful.Length})
+    mask = not_clamped_mask(system)
+    new_positions = deepcopy(position(system))
+    new_positions[mask] = reinterpret(reshape, SVector{3, eltype(positions)},
+                                      reshape(positions, 3, :))
+    update_positions(system, new_positions)
+end
+
+@doc raw"""
+Returns a mask for selecting the not clamped atoms in the system.
+
+"""
+function not_clamped_mask(system)
+    # If flag not set, the atom is considered not clamped.
+    [haskey(a, :clamped) ? !a[:clamped] : true for a in system]
+end
+
+function not_clamped_positions(system)
+    mask = not_clamped_mask(system)
+    Iterators.flatten(system[mask, :position])
+end
+
+@doc raw"""
+    Clamp given atoms in the system. Clamped atoms are fixed and their positions 
+    will not be optimized. The atoms to be clamped should be given as a list of 
+    indices corresponding to their positions in the system.
+
+    """
+function clamp_atoms(system, clamped_indexes::Union{AbstractVector{<:Integer},Nothing})
+    clamped = falses(length(system))
+    clamped[clamped_indexes] .= true
+    particles = [Atom(atom; clamped=m) for (atom, m) in zip(system, clamped)]
+    AbstractSystem(system, particles=particles)
+end

src/optimization.jl

@@ -0,0 +1,43 @@
+#=
+# Note that by default all particles in the system are assumed optimizable.
+# IMPORTANT: Note that we always work in cartesian coordinates.
+=#
+
+export minimize_energy!
+
+
+""" 
+By default we work in cartesian coordinaes.
+Note that internally, when optimizing the cartesian positions, atomic units 
+are used.
+"""
+function Optimization.OptimizationFunction(system, calculator; kwargs...)
+    mask = not_clamped_mask(system)  # mask is assumed not to change during optim.
+
+    f = function(x::AbstractVector{<:Real}, p)
+        new_system = update_not_clamped_positions(system, x * u"bohr")
+        energy = AtomsCalculators.potential_energy(new_system, calculator; kwargs...)
+        austrip(energy)
+    end
+
+    g! = function(G::AbstractVector{<:Real}, x::AbstractVector{<:Real}, p)
+        new_system = update_not_clamped_positions(system, x * u"bohr")
+        energy = AtomsCalculators.potential_energy(new_system, calculator; kwargs...)
+
+        forces = AtomsCalculators.forces(new_system, calculator; kwargs...)
+        # Translate the forces vectors on each particle to a single gradient for the optimization parameter.
+        forces_concat = collect(Iterators.flatten(forces[mask]))
+
+        # NOTE: minus sign since forces are opposite to gradient.
+        G .= - austrip.(forces_concat)
+    end
+    OptimizationFunction(f; grad=g!)
+end
+
+function minimize_energy!(system, calculator; 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)
+    problem = OptimizationProblem(f_opt, x0, nothing)  # Last argument needed in Optimization.jl.
+    solve(problem, solver; kwargs...)
+end

test/dummy_calculator.jl

@@ -0,0 +1,18 @@
+# Create a dummy AtomsCalculator to test the geometry optimization interfcae.
+@testsetup module TestCalculators
+    using AtomsCalculators
+    using Unitful
+    using UnitfulAtomic
+    
+    struct DummyCalculator end
+    
+    AtomsCalculators.@generate_interface function AtomsCalculators.potential_energy(
+            system, calculator::DummyCalculator; kwargs...)
+        0.0u"eV"
+    end
+        
+    AtomsCalculators.@generate_interface function AtomsCalculators.forces(
+            system, calculator::DummyCalculator; kwargs...)
+        AtomsCalculators.zero_forces(system, calculator)
+    end
+end