Make Optimization.jl optional (#27)

Author: mfherbst

Project.toml

@@ -1,31 +1,36 @@
 name = "GeometryOptimization"
 uuid = "673bf261-a53d-43b9-876f-d3c1fc8329c2"
 authors = ["JuliaMolSim community"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
-OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
+[weakdeps]
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+
+[extensions]
+GeometryOptimizationOptimizationExt = "Optimization"
+
 [compat]
 AtomsBase = "0.5"
 AtomsBuilder = "0.2.2"
 AtomsCalculators = "0.2.3"
 DocStringExtensions = "0.9"
 LineSearches = "7"
+Optim = "1.11.0"
 Optimization = "3"
-OptimizationOptimJL = "0.3"
 PrettyTables = "2"
 StaticArrays = "1"
 TestItemRunner = "0.2"
@@ -37,8 +42,9 @@ julia = "1.10"
 [extras]
 AtomsBuilder = "f5cc8831-eeb7-4288-8d9f-d6c1ddb77004"
 EmpiricalPotentials = "38527215-9240-4c91-a638-d4250620c9e2"
+OptimizationNLopt = "4e6fcdb7-1186-4e1f-a706-475e75c168bb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 
 [targets]
-test = ["AtomsBuilder", "EmpiricalPotentials", "Test", "TestItemRunner"]
+test = ["AtomsBuilder", "EmpiricalPotentials", "Optimization", "OptimizationNLopt", "Test", "TestItemRunner"]

README.md

@@ -9,6 +9,7 @@ A package for optimising the structural parameters of an atomistic system,
 i.e. the step usually referred to as
 **geometry optimisation** or **structural relaxation**
 in electronic structure theory and atomistic modelling.
+Both relaxing **atomic positions** as well as the **unit cell** is supported.
 
 The package is generic in the datastructures used to represent the geometry,
 the calculator used to evaluate energies and forces as well as the solver algorithm.

docs/src/examples/aluminium_dftk.md

@@ -36,15 +36,14 @@ calc = DFTKCalculator(; model_kwargs, basis_kwargs, scf_kwargs)
 
 We perform the structure optimisation using the LBFGS solver
 from Optim with solver parameters adapted for our geometry optimisation setting.
-This is selected by passing the [GeometryOptimization.OptimLBFGS](@ref)
+This is selected by passing the [OptimLBFGS](@ref)
 solver as the third argument. The `verbosity=2` flag makes sure we get
 output from both the geometry optimisation as well as the inner SCF solver.
 
 ```@example dftk-aluminium
 using GeometryOptimization
-GO = GeometryOptimization
 
-results = minimize_energy!(system, calc, GO.OptimLBFGS();
+results = minimize_energy!(system, calc, OptimLBFGS();
                            tol_forces=1e-4u"eV/Å", verbosity=2)
 nothing
 ```
@@ -65,13 +64,10 @@ We can view the final structure
 results.system
 ```
 
-Some statistics about the optimisation
+The results object returned from Optim (containing some statistics
+about the optimisation):
 ```@example dftk-aluminium
-results.stats
-```
-or the details about the selected algorithm:
-```@example dftk-aluminium
-results.alg
+results.optimres
 ```
 
 The final state of the calculator object is also accessible

docs/src/examples/other_solvers.md

@@ -1,37 +1,23 @@
-# Using a different Optimization.jl compatible solver
+# Using a Optimization.jl compatible solver
 
-In this example we perform the simplistic optimisation
-the bond length of a Hydrogen molecule using a trust region
-quasi-Newton method from
+In this example we optimize a bulk silicon structure
+using a trust region quasi-Newton method from
 [NLopt](https://github.com/JuliaOpt/NLopt.jl).
 
-We create a calculator employing the
-[density-functional toolkit](https://dftk.org/)
-to compute energies and forces at using the LDA density functional.
+We create a Stillinger-Weber calculator
 
 ```@example other-solvers
-using DFTK
-using PseudoPotentialData
-
-pseudopotentials = PseudoFamily("dojo.nc.sr.lda.v0_4_1.oncvpsp3.standard.upf")
-model_kwargs = (; functionals=LDA(), pseudopotentials)
-basis_kwargs = (; kgrid=(1, 1, 1), Ecut=20.0)
-calc = DFTKCalculator(; model_kwargs, basis_kwargs)
+using EmpiricalPotentials
+sw = StillingerWeber()
+nothing
 ```
 
-and we build the hydrogen molecular system,
-where we attach pseudopotential information for DFTK:
+and we build a slightly rattled silicon structure
 
 ```@example other-solvers
 using AtomsBuilder
 using Unitful
-using UnitfulAtomic
-
-bounding_box = [[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]u"Å"
-system = periodic_system([:H => [0, 0, 1.]u"bohr",
-                          :H => [0, 0, 3.]u"bohr"],
-                         bounding_box)
-nothing
+system = rattle!(bulk(:Si, cubic=true) * (2, 2, 2), 0.1u"Å")
 ```
 
 We now run [`GeometryOptimization.minimize_energy!`](@ref), but notably
@@ -44,15 +30,15 @@ using GeometryOptimization
 using OptimizationNLopt
 solver = NLopt.LD_TNEWTON()
 
-results = minimize_energy!(system, calc, solver;
+results = minimize_energy!(system, sw, solver;
                            tol_forces=1e-4u"eV/Å", verbosity=1,
                            maxeval=100)
 nothing
 ```
 
-The final hydrogen bond length is:
-
-```@example other-solvers
-using LinearAlgebra
-norm(position(results.system[1]) - position(results.system[2]))
-```
+!!! info ""
+    While in principle all *first-order* solvers supported
+    by Optimization.jl can be employed, only few have been tested with this
+    package so far. Given the complexity of the Optimization.jl ecosystem
+    we expect that minor changes of our integration may be needed to make
+    specific solvers work well. We welcome any PRs.

docs/src/examples/tial_lj.md

@@ -34,9 +34,8 @@ nothing  # hide
 Minimise energy:
 ```@example tial
 using GeometryOptimization
-GO = GeometryOptimization
 
-results = minimize_energy!(system, calc, GO.OptimCG(); maxiters=10, verbosity=1)
+results = minimize_energy!(system, calc, OptimCG(); maxiters=10, verbosity=1)
 nothing  # hide
 ```
 

ext/GeometryOptimizationOptimizationExt.jl

@@ -0,0 +1,76 @@
+module GeometryOptimizationOptimizationExt
+using AtomsCalculators
+using Optimization
+using GeometryOptimization
+import GeometryOptimization: GeoOptProblem, GeoOptConvergence
+const GO = GeometryOptimization
+
+function GeometryOptimization.solve_problem(
+    prob::GeoOptProblem, solver, cvg::GeoOptConvergence;
+    callback, maxiters, maxtime, kwargs...)
+
+    function inner_callback(optim_state, ::Any, geoopt_state)
+        cache_evaluations = geoopt_state.cache_evaluations
+
+        # Find position in the cache matching the current state
+        i_match = findlast(cache_evaluations) do eval
+            isnothing(eval.gradnorm) && return false
+
+            tol = 10eps(typeof(optim_state.objective))
+            obj_matches = abs(eval.objective - optim_state.objective)  < tol
+
+            if isnothing(optim_state.grad)
+                # Nothing we can do, let's just hope it's ok
+                grad_matches = true
+            else
+                g_norm = maximum(abs, optim_state.grad)
+                grad_matches = abs(eval.gradnorm  - g_norm) < tol
+            end
+
+            obj_matches && grad_matches
+        end
+        i_match = @something i_match length(cache_evaluations)
+
+        # Commit data from state and discard the rest
+        geoopt_state.n_iter = optim_state.iter
+        push!(geoopt_state.history_energy, cache_evaluations[i_match].energy)
+        if !isnothing(cache_evaluations[i_match].forces)
+            geoopt_state.forces .= cache_evaluations[i_match].forces
+        end
+        if !isnothing(cache_evaluations[i_match].virial)
+            geoopt_state.virial .= cache_evaluations[i_match].virial
+        end
+        empty!(cache_evaluations)
+
+        # Check for convergence
+        geoopt_state.converged = GO.is_converged(cvg, geoopt_state)
+
+        # Callback and possible abortion
+        halt = callback(optim_state, geoopt_state)
+        halt && return true
+
+        geoopt_state.converged
+    end
+
+    optimres = solve(Optimization.OptimizationProblem(prob), solver;
+                     maxiters, maxtime, callback=inner_callback, kwargs...)
+    (; minimizer=optimres.u, minimum=optimres.objective, optimres)
+end
+
+
+function Optimization.OptimizationProblem(prob::GeoOptProblem; kwargs...)
+    f = function(x::AbstractVector{<:Real}, ps)
+        GO.eval_objective_gradient!(nothing, prob, ps, x), prob.geoopt_state
+    end
+    g! = function(G::AbstractVector{<:Real}, x::AbstractVector{<:Real}, ps)
+        GO.eval_objective_gradient!(G, prob, ps, x), prob.geoopt_state
+        G
+    end
+    f_opt = OptimizationFunction(f; grad=g!)
+
+    # Note: Some optimisers modify Dofs x0 in-place, so x0 needs to be mutable type.
+    x0 = GO.get_dofs(prob.system, prob.dofmgr)
+    OptimizationProblem(f_opt, x0, AtomsCalculators.get_parameters(prob.calculator);
+                        sense=Optimization.MinSense, kwargs...)
+end
+end

src/GeometryOptimization.jl

@@ -4,18 +4,14 @@ using AtomsBase
 using AtomsCalculators
 using DocStringExtensions
 using LinearAlgebra
-using Optimization
+using LineSearches
+using Optim
 using StaticArrays
 using Unitful
 using UnitfulAtomic
 
-# Make sure Optim is always available
-using OptimizationOptimJL
-using LineSearches
-
-# Useful shortcuts
 using AtomsCalculators: Energy, Forces, Virial
-AC = AtomsCalculators
+const AC = AtomsCalculators
 
 @template METHODS =
 """
@@ -24,11 +20,13 @@ $(TYPEDSIGNATURES)
 $(DOCSTRING)
 """
 
-include("dof_management.jl")
-include("optimization.jl")
-include("callbacks.jl")
-
 export minimize_energy!
 export GeoOptDefaultCallback
+export Autoselect, OptimLBFGS, OptimCG, OptimSD
+
+include("dof_management.jl")
+include("minimize_energy.jl")
+include("optim.jl")
+include("callbacks.jl")
 
 end

src/callbacks.jl

@@ -25,23 +25,25 @@ function (cb::GeoOptDefaultCallback)(optim_state, geoopt_state)
     cb.verbosity ≤ 0 && return false  # No printing, just continue iterations
 
     # If first iteration clear a potentially cached previous time
-    optim_state.iter ≤ 0 && (cb.prev_time[] = 0)
+    geoopt_state.n_iter ≤ 0 && (cb.prev_time[] = 0)
     runtime_ns = time_ns() - geoopt_state.start_time
     tstr = @sprintf "% 6s" TimerOutputs.prettytime(runtime_ns - cb.prev_time[])
     cb.prev_time[] = runtime_ns
 
-    Estr  = (@sprintf "%+15.12f" round(austrip(geoopt_state.energy), sigdigits=13))[1:15]
-    if iszero(geoopt_state.energy_change) && optim_state.iter < 1
+    energy = austrip(geoopt_state.history_energy[end])
+    Estr  = (@sprintf "%+15.12f" round(energy, sigdigits=13))[1:15]
+    if geoopt_state.n_iter < 2
         logΔE = ""
     else
-        logΔE = format_log8(austrip(geoopt_state.energy_change))
+        energy_change = geoopt_state.history_energy[end] - geoopt_state.history_energy[end-1]
+        logΔE = format_log8(austrip(energy_change))
     end
 
     maxforce = austrip(maximum(norm, geoopt_state.forces))
     fstr = iszero(maxforce) ? "" : round(maxforce, sigdigits=8)
 
     fields = [  # Header, width, value
-        ("n",           3, optim_state.iter),
+        ("n",           3, geoopt_state.n_iter),
         ("Energy",     15, Estr),
         ("log10(ΔE)",   9, logΔE),
         ("max(Force)", 11, fstr),
@@ -62,14 +64,14 @@ function (cb::GeoOptDefaultCallback)(optim_state, geoopt_state)
     if cb.always_show_header
         hlines = :all
         show_header = true
-    elseif optim_state.iter == 0
+    elseif geoopt_state.n_iter == 0
         hlines = [0, 1]
         show_header = true
     else
         hlines = :none
         show_header = false
     end
-    title = iszero(optim_state.iter) ? "Geometry optimisation convergence (in atomic units)" : ""
+    title = iszero(geoopt_state.n_iter) ? "Geometry optimisation convergence (in atomic units)" : ""
 
     cb.always_show_header && println(stdout)
     pretty_table(stdout, reshape(getindex.(fields, 3), 1, length(fields));

src/dof_management.jl

@@ -29,14 +29,15 @@ One aspect of this definition is that clamped atom positions still change via
 the deformation `F`. This is natural in the context of optimizing the
 cell shape.
 """
-mutable struct DofManager{D,T}
+struct DofManager{D,T}
     variablecell::Bool
     ifree::Vector{Int}   # extract the free position dofs
     r0::T
     X0::Vector{SVector{D,T}}       # reference positions
     C0::NTuple{D, SVector{D, T}}   # reference cell
 end
-# TODO Why is this  mutable ?
+
+# TODO: Change F to F-I  (such that the initial guess is exactly 0 everywhere)
 
 
 # NOTES:
@@ -181,12 +182,12 @@ end
 #   Compute the gradient with respect to dofs
 #   from forces and virials
 
-function energy_dofs(system, calculator, dofmgr, x::AbstractVector, ps, state)
+function eval_objective(system, calculator, dofmgr, x::AbstractVector, ps, state)
     res = AC.calculate(Energy(), set_dofs(system, dofmgr, x), calculator, ps, state)
     (; energy_unitless=austrip(res.energy), res...)
 end
 
-function gradient_dofs(system, calculator, dofmgr, x::AbstractVector{T}, ps, state) where {T}
+function eval_gradient(system, calculator, dofmgr, x::AbstractVector{T}, ps, state) where {T}
     # Compute and transform forces and virial into a gradient w.r.t. x
     if fixedcell(dofmgr)
         # fixed cell version
@@ -212,3 +213,18 @@ function gradient_dofs(system, calculator, dofmgr, x::AbstractVector{T}, ps, sta
     end
     (; grad, res...)
 end
+
+# =======================
+#   Idea for making this more flexible:
+#
+# struct FixedCell   end
+# struct UnitaryCell end
+#
+# struct DofManager{D,T,CellParam,Sys}
+#     system0::Sys  # Reference system
+#     r0::T         # Reference unit length
+#     X0::Vector{SVector{D,T}}      # Reference positions
+#     C0::NTuple{D, SVector{D, T}}  # Reference cell
+#     cellparam::CellParam  # Parametrisation used for the unit cell (fixed, voigt, unitary)
+#     ifree::Vector{Int}    # Degrees of freedom, which can be changed
+# end