Reintroduce structure optimization

This reverts and updates parts of commit 2ff4988.

Signed-off-by: Thomas Kemmer <[email protected]>

Project.toml

@@ -19,6 +19,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mendeleev = "c116f080-063d-490a-9873-2b5b2cce4c34"
 MetaGraphs = "626554b9-1ddb-594c-aa3c-2596fe9399a5"
 MolecularGraph = "6c89ec66-9cd8-5372-9f91-fabc50dd27fd"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
@@ -49,6 +50,7 @@ LinearAlgebra = "1.9"
 Mendeleev = "1"
 MetaGraphs = "0.7.2, 0.8"
 MolecularGraph = "0.16, 0.17"
+Optimization = "4"
 PhysicalConstants = "0.2"
 PrettyTables = "2.3"
 Printf = "1.9"

src/BiochemicalAlgorithms.jl

@@ -17,6 +17,7 @@ using Quaternions: quat
 
 import JSON3
 import MolecularGraph
+import Optimization
 import PrettyTables
 import StaticArrays
 import Tables, TableOperations
@@ -90,6 +91,9 @@ include("preprocessing/build_bonds.jl")
 include("preprocessing/add_hydrogens.jl")
 include("preprocessing/reconstruct_fragments.jl")
 
+# optimization
+include("optimization/optimize_structure.jl")
+
 export
     ball_data_path
 

src/optimization/optimize_structure.jl

@@ -0,0 +1,26 @@
+export
+    optimize_structure!
+
+"""
+    TODO
+"""
+function optimize_structure!(ff::ForceField; alg = Optimization.LBFGS(), kwargs...)
+    r0 = collect(Float64, Iterators.flatten(atoms(ff.system).r))
+
+    optf = Optimization.OptimizationFunction(
+        (r, _ = nothing) -> begin
+            atoms(ff.system).r .= eachcol(reshape(r, 3, :))
+            update!(ff)
+            compute_energy!(ff)
+        end,
+        grad = (F, r, _) -> begin
+            update!(ff)
+            compute_forces!(ff)
+            F .= -collect(Float64, Iterators.flatten(atoms(ff.system).F)) ./ BiochemicalAlgorithms.force_prefactor
+            nothing
+        end
+    )
+    prob = Optimization.OptimizationProblem(optf, r0)
+
+    Optimization.solve(prob, alg; kwargs...)
+end