Add verbosity flag and callback for structured printing (#16)

Author: mfherbst

Project.toml

@@ -1,7 +1,7 @@
 name = "GeometryOptimization"
 uuid = "673bf261-a53d-43b9-876f-d3c1fc8329c2"
 authors = ["JuliaMolSim community"]
-version = "0.0.2"
+version = "0.1.0"
 
 [deps]
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
@@ -11,7 +11,10 @@ LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
+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"
 
@@ -23,16 +26,20 @@ DocStringExtensions = "0.9"
 LineSearches = "7"
 Optimization = "3"
 OptimizationOptimJL = "0.3"
+PrettyTables = "2"
 StaticArrays = "1"
 TestItemRunner = "0.2"
+TimerOutputs = "0.5.12"
 Unitful = "1"
 UnitfulAtomic = "1"
 julia = "1.10"
 
 [extras]
+ASEconvert = "3da9722f-58c2-4165-81be-b4d7253e8fd2"
 AtomsBuilder = "f5cc8831-eeb7-4288-8d9f-d6c1ddb77004"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 
 [targets]
-test = ["AtomsBuilder", "Test", "TestItemRunner"]
+test = ["ASEconvert", "AtomsBuilder", "PythonCall", "Test", "TestItemRunner"]

docs/src/examples/aluminium_dftk.md

@@ -20,8 +20,8 @@ using DFTK
 
 model_kwargs = (; functionals=[:lda_x, :lda_c_pw], temperature=1e-3)
 basis_kwargs = (; kgrid=(3, 3, 3), Ecut=10.0)
-scf_kwargs   = (; tol=1e-6, mixing=KerkerMixing())
-calc = DFTKCalculator(; model_kwargs, basis_kwargs, scf_kwargs, verbose=true)
+scf_kwargs   = (; mixing=KerkerMixing())
+calc = DFTKCalculator(; model_kwargs, basis_kwargs, scf_kwargs)
 nothing
 ```
 
@@ -42,18 +42,24 @@ nothing
 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)
-solver as the third argument.
+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();
-                           tol_force=1e-4u"eV/Å",
-                           show_trace=true)
+                           tol_forces=1e-4u"eV/Å", verbosity=2)
 nothing
 ```
 
+!!! tip "Automatically adapted calculator parameters"
+    Some calculators (such as DFTK) are able to adapt to the keyword arguments
+    and parameters passed to `minimize_energy!`. In this case the SCF tolerance
+    is automatically adapted according to the convergence parameters
+    (here `tol_forces`) passed to `minimize_energy!`.
+
 The final energy is
 ```@example dftk-aluminium
 results.energy

docs/src/examples/other_solvers.md

@@ -46,7 +46,7 @@ using OptimizationNLopt
 solver = NLopt.LD_TNEWTON()
 
 results = minimize_energy!(system, calc, solver;
-                           tol_force=1e-4u"eV/Å", maxeval=100)
+                           tol_forces=1e-4u"eV/Å", maxeval=100)
 nothing
 ```
 

docs/src/examples/tial_lj.md

@@ -26,7 +26,7 @@ Minimise energy:
 ```julia
 ## TODO: Should run as @example once EmpiricalPotentials is compatible
 
-results = minimize_energy!(system, calc, GO.OptimCG(); maxiters=10, show_trace=true)
+results = minimize_energy!(system, calc, GO.OptimCG(); maxiters=10, verbosity=1)
 results.energy
 ```
 

docs/src/index.md

@@ -27,8 +27,8 @@ system = isolated_system([:H => [0, 0, 0.0]u"bohr",
                           :H => [0, 0, 1.9]u"bohr"])
 calc = LennardJones(-1.17u"hartree", 0.743u"angstrom", 1, 1, 0.6u"nm")
 
-# Run the geometry optimisation
-results = minimize_energy!(system, calc)
+# Run the geometry optimisation (using verbosity=1 to print the progress)
+results = minimize_energy!(system, calc; verbosity=1)
 
 # Inspect the results
 optimised_system = results.system

src/GeometryOptimization.jl

@@ -24,5 +24,9 @@ $(DOCSTRING)
 
 include("clamping_updating_positions.jl")
 include("optimization.jl")
+include("callbacks.jl")
+
+export minimize_energy!
+export GeoOptDefaultCallback
 
 end

src/callbacks.jl

@@ -0,0 +1,69 @@
+import TimerOutputs
+using PrettyTables
+using Printf
+
+"""
+Callback producing a convergence table summarising the geometry
+optimisation convergence. If `always_show_header=true` the
+header is shown in each iteration. This is helpful if the calculator
+produces output as well.
+"""
+struct GeoOptDefaultCallback
+    verbosity::Int
+    always_show_header::Bool
+    prev_time::Ref{UInt64}
+end
+function GeoOptDefaultCallback(verbosity=1; always_show_header=verbosity > 1)
+    GeoOptDefaultCallback(verbosity, always_show_header, Ref{UInt64}(0))
+end
+
+format_log8(value) = (value < 0 ? " " : "+") * (@sprintf "%8.2f" log10(abs(value)))
+
+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)
+    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]
+    logΔE = optim_state.iter < 1 ? "" : format_log8(austrip(geoopt_state.energy_change))
+
+    maxforce = austrip(maximum(norm, geoopt_state.forces))
+    fstr = iszero(maxforce) ? "" : round(maxforce, sigdigits=8)
+
+    fields = [  # Header, width, value
+        ("n",           3, optim_state.iter),
+        ("Energy",     15, Estr),
+        ("log10(ΔE)",   9, logΔE),
+        ("max(Force)", 10, fstr),
+        # TODO Maximal atomic displacement
+        # TODO Current virial, trace of lattice deformation matrix
+        ("Δtime",       6, tstr),
+        # TODO Would be nice to have some simple way to add in
+        #      a few calculator-specific things (e.g. total number of SCF iterations)
+    ]
+
+    if cb.always_show_header
+        hlines = :all
+        show_header = true
+    elseif optim_state.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)" : ""
+
+    cb.always_show_header && println(stdout)
+    pretty_table(stdout, reshape(getindex.(fields, 3), 1, length(fields));
+                 header=Symbol.(getindex.(fields, 1)), columns_width=getindex.(fields, 2),
+                 title, show_header, hlines)
+    cb.always_show_header && println(stdout)
+
+    flush(stdout)
+    return false
+end

src/clamping_updating_positions.jl

@@ -2,7 +2,6 @@
 # Interface between AtomsBase.jl and GeometryOptimization.jl that provides 
 # utility functions for manipulating systems.
 #
-export clamp_atoms
 
 """
 Convert the input system to a kind of system we can work with, i.e. one where
@@ -48,6 +47,7 @@ indices corresponding to their positions in the system.
     in the future.
 """
 function clamp_atoms(system, clamped_indexes::Union{AbstractVector{<:Integer},Nothing})
+    system = convert_to_updatable(system)
     clamped = falses(length(system))
     clamped[clamped_indexes] .= true
     particles = [Atom(atom; clamped=msk) for (atom, msk) in zip(system, clamped)]

src/optimization.jl

@@ -3,21 +3,27 @@
 # We always work in cartesian coordinates.
 #
 
-export minimize_energy!
-
 using AtomsCalculators: Energy, Forces, calculate
 
 mutable struct GeometryOptimizationState
-    calc_state  # Reference to the current calculator state
-    energy      # Most recent energy value
-    forces      # Most recent force value
-    virial      # Most recent virial value
+    calculator
+    calc_state          # Reference to the most recent calculator state
+    energy              # Current energy value
+    forces              # Current force value
+    virial              # Current virial value
+    energy_change       # Change in energy since previous step
+    start_time::UInt64  # Time when the calculation started from time_ns()
 end
-function GeometryOptimizationState(system, calculator)
-    GeometryOptimizationState(AC.get_state(calculator),
-                              AC.zero_energy(system, calculator),
-                              AC.zero_forces(system, calculator),
-                              AC.zero_virial(system, calculator))
+function GeometryOptimizationState(system, calculator; start_time=time_ns())
+    GeometryOptimizationState(
+        calculator,
+        AC.get_state(calculator),
+        AC.zero_energy(system, calculator),
+        AC.zero_forces(system, calculator),
+        AC.zero_virial(system, calculator),
+        AC.zero_energy(system, calculator),
+        start_time,
+    )
 end
 
 """
@@ -48,7 +54,6 @@ function Optimization.OptimizationProblem(system, calculator, geoopt_state; kwar
         new_system = update_not_clamped_positions(system, x * u"bohr")
         res = calculate(Energy(), new_system, calculator, ps, geoopt_state.calc_state)
         geoopt_state.calc_state = res.state
-        geoopt_state.energy = res.energy
         austrip(res.energy), geoopt_state
     end
     g! = function(G::AbstractVector{<:Real}, x::AbstractVector{<:Real}, ps)
@@ -89,13 +94,18 @@ can also be employed here.
 
 ## Keyword arguments:
 - `maxiters`: Maximal number of iterations
+- `maxtime`:  Maximal allowed runtime (in seconds)
 - `tol_energy`: Tolerance in the energy to stop the minimisation (all `tol_*` need to be satisfied)
-- `tol_force`:  Tolerance in the force  to stop the minimisation (all `tol_*` need to be satisfied)
+- `tol_forces`:  Tolerance in the force  to stop the minimisation (all `tol_*` need to be satisfied)
 - `tol_virial`: Tolerance in the virial to stop the minimisation (all `tol_*` need to be satisfied)
-- `maxstep`: Maximal step size (in AU) to be taken in a single optimisation step
+- `maxstep`: Maximal step size (in AU or length units) to be taken in a single optimisation step
   (not supported for all `solver`s)
-- `callback`: A custom callback, which obtains the pair `(optimization_state, geoopt_state)` and is
-  expected to return `false` (continue iterating) or `true` (halt iterations).
+- `verbosity`: Printing level. The idea is that `0` is silent, `1` displays the optimisation
+  progress and `≥ 2` starts displaying things from the calculator as well (e.g SCF iterations).
+- `callback`: A custom callback, which obtains the pair `(optimization_state, geoopt_state)`
+  and is expected to return `false` (continue iterating) or `true` (halt iterations). Note
+  that specifying this overwrites the default printing callback. The calculation thus becomes
+  silent unless a [`GeoOptDefaultCallback`](@ref) is included in the callback.
 - `kwargs`: All other keyword arguments are passed to the call to `solve`. Note, that
   if special `kwargs` should be passed to the `Optimization.OptimizationProblem` the user
   needs to setup the problem manually (e.g. `OptimizationProblem(system, calculator)`)
@@ -110,12 +120,14 @@ end
 # do some additional calculator-specific setup (e.g. callbacks) and so on. Then
 # by calling this function the actual minimisation is started off.
 function _minimize_energy!(system, calculator, solver;
-                           maxiters=100,
+                           maxiters::Integer=100,
+                           maxtime::Integer=60*60*24*365,  # 1 year
                            tol_energy=Inf*u"eV",
-                           tol_force=1e-4u"eV/Å",  # VASP default
-                           tol_virial=1e-6u"eV",   # TODO How reasonable ?
+                           tol_forces=1e-4u"eV/Å",  # VASP default
+                           tol_virial=1e-6u"eV",    # TODO How reasonable ?
                            maxstep=0.8u"bohr",
-                           callback=(x,y) -> false,
+                           verbosity::Integer=0,
+                           callback=GeoOptDefaultCallback(verbosity),
                            kwargs...)
     solver = setup_solver(system, calculator, solver; maxstep)
     system = convert_to_updatable(system)
@@ -124,23 +136,27 @@ function _minimize_energy!(system, calculator, solver;
     problem = OptimizationProblem(system, calculator, geoopt_state; sense=Optimization.MinSense)
     converged = false
 
-    Eold = AC.zero_energy(system, calculator)
+    Eold = 0  # Atomic units
     function inner_callback(optim_state, ::Any, geoopt_state)
+        geoopt_state.energy        = optim_state.objective * u"hartree"
+        geoopt_state.energy_change = (optim_state.objective - Eold) * u"hartree"
+
         halt = callback(optim_state, geoopt_state)
         halt && return true
 
-        energy_converged = austrip(abs(geoopt_state.energy - Eold))    < austrip(tol_energy)
-        force_converged  = austrip(maximum(norm, geoopt_state.forces)) < austrip(tol_force)
+        energy_converged = optim_state.objective - Eold < austrip(tol_energy)
+        force_converged  = austrip(maximum(norm, geoopt_state.forces)) < austrip(tol_forces)
         virial_converged = austrip(maximum(abs,  geoopt_state.virial)) < austrip(tol_virial)
 
-        Eold = geoopt_state.energy
+        Eold = optim_state.objective
         converged = energy_converged && force_converged && virial_converged
         return converged
     end
 
-    optimres = solve(problem, solver; maxiters, callback=inner_callback, kwargs...)
+    optimres = solve(problem, solver; maxiters, maxtime, callback=inner_callback, kwargs...)
+    converged || @warn "Geometry optimisation not converged."
     (; system=update_not_clamped_positions(system, optimres.u * u"bohr"), converged,
-       energy=optimres.objective, geoopt_state.forces, geoopt_state.virial,
+       energy=optimres.objective * u"hartree", geoopt_state.forces, geoopt_state.virial,
        state=geoopt_state.calc_state, optimres.stats, optimres.alg, optimres)
 end
 

test/interface_test.jl

@@ -1,21 +1,29 @@
-@testitem "Test GeometryOptimization with AtomsCalculators interface" #=
-        =# setup=[TestCalculators] begin
+@testitem "Test GeometryOptimization with AtomsCalculators interface" begin
     # This is a fake test. It does not really test anything,
     # just that the interface code compiles and works properly.
 
     using AtomsBuilder
+    using AtomsCalculators
     using GeometryOptimization
     using LinearAlgebra
     using Unitful
     using UnitfulAtomic
+    AC = AtomsCalculators
     GO = GeometryOptimization
 
-    system = rattle!(bulk(:Al; cubic=true), 0.1u"Å")
-    system = clamp_atoms(system, [1])
+    struct DummyCalc end
+    AC.energy_unit(::DummyCalc) = u"hartree"
+    AC.length_unit(::DummyCalc) = u"bohr"
+    AC.@generate_interface function AC.potential_energy(system, calc::DummyCalc; kwargs...)
+        AC.zero_energy(system, calc)
+    end
+    AC.@generate_interface function AC.forces(system, calc::DummyCalc; kwargs...)
+        AC.zero_forces(system, calc)
+    end
 
-    calculator = TestCalculators.DummyCalc()
+    system = GO.clamp_atoms(bulk(:Al; cubic=true), [1])
     for solver in (GO.Autoselect(), GO.OptimCG(), GO.OptimLBFGS(), GO.OptimSD())
-        results = minimize_energy!(system, calculator, solver)
+        results = minimize_energy!(system, DummyCalc(), solver; verbosity=1)
         @test norm(position(results.system[1]) - position(system[1])) < 1e-5u"bohr"
     end
 end