diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 86e4d95..35c6603 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -19,7 +19,6 @@ jobs: fail-fast: false matrix: version: - - '1.0' - '1.9' - 'nightly' os: diff --git a/Project.toml b/Project.toml index 8953c9e..11547d8 100644 --- a/Project.toml +++ b/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"] diff --git a/examples/Al_supercell.jl b/examples/Al_supercell.jl new file mode 100644 index 0000000..aea2658 --- /dev/null +++ b/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) diff --git a/examples/H2.jl b/examples/H2.jl new file mode 100644 index 0000000..4e97f41 --- /dev/null +++ b/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]) diff --git a/examples/H2_LJ.jl b/examples/H2_LJ.jl new file mode 100644 index 0000000..594f05d --- /dev/null +++ b/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])).") diff --git a/examples/TiAl_EmpiricalPotentials.jl b/examples/TiAl_EmpiricalPotentials.jl new file mode 100644 index 0000000..0aa3a85 --- /dev/null +++ b/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) diff --git a/src/GeometryOptimization.jl b/src/GeometryOptimization.jl index cde27c9..c85ffce 100644 --- a/src/GeometryOptimization.jl +++ b/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 diff --git a/src/atomsbase_interface.jl b/src/atomsbase_interface.jl new file mode 100644 index 0000000..13c31d4 --- /dev/null +++ b/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 diff --git a/src/optimization.jl b/src/optimization.jl new file mode 100644 index 0000000..87f79af --- /dev/null +++ b/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 diff --git a/test/dummy_calculator.jl b/test/dummy_calculator.jl new file mode 100644 index 0000000..ba29c0e --- /dev/null +++ b/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 diff --git a/test/geometry_optimization.jl b/test/geometry_optimization.jl new file mode 100644 index 0000000..b7fd6e2 --- /dev/null +++ b/test/geometry_optimization.jl @@ -0,0 +1,23 @@ +#= Test Geometry Optimization on an aluminium supercell. +=# +@testitem "Test AtomsCalculators energy and forces interface" tags=[:dont_test_mpi, :dont_test_windows] setup=[TestCalculators] begin + 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) + system = clamp_atoms(system, [1]) + + calculator = TestCalculators.DummyCalculator() + + solver = OptimizationOptimJL.LBFGS() + optim_options = (f_tol=1e-6, iterations=4, show_trace=false) + + results = minimize_energy!(system, calculator; solver=solver, optim_options...) + @test isapprox(results.u[1], 0; atol=1e-5) +end diff --git a/test/runtests.jl b/test/runtests.jl index 0f0669e..deed9a7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,3 @@ -using GeometryOptimization -using Test +using TestItemRunner -@testset "GeometryOptimization.jl" begin - # Write your tests here. -end +@run_package_tests verbose=true