Merge pull request #599 from SciML/optprob

Add GalacticOptim-compatible OptimizationProblem generation

Author: ChrisRackauckas

Project.toml

@@ -59,11 +59,13 @@ julia = "1.2"
 [extras]
 Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+GalacticOptim = "a75be94c-b780-496d-a8a9-0878b188d577"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Dagger", "ForwardDiff", "OrdinaryDiffEq", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]
+test = ["Dagger", "ForwardDiff", "GalacticOptim", "OrdinaryDiffEq", "Optim", "Random", "SteadyStateDiffEq", "Test", "StochasticDiffEq"]

src/systems/diffeqs/modelingtoolkitize.jl

@@ -84,3 +84,27 @@ function modelingtoolkitize(prob::DiffEqBase.SDEProblem)
 
     de
 end
+
+
+"""
+$(TYPEDSIGNATURES)
+
+Generate `OptimizationSystem`, dependent variables, and parameters from an `OptimizationProblem`.
+"""
+function modelingtoolkitize(prob::DiffEqBase.OptimizationProblem)
+
+    if prob.p isa Tuple || prob.p isa NamedTuple
+        p = [x for x in prob.p]
+    else
+        p = prob.p
+    end
+
+    vars = reshape([Variable(:x, i)() for i in eachindex(prob.u0)],size(prob.u0))
+    params = p isa DiffEqBase.NullParameters ? [] :
+             reshape([Variable(:α,i)() for i in eachindex(p)],size(Array(p)))
+
+
+    eqs = prob.f(vars, params)
+    de = OptimizationSystem(eqs,vec(vars),vec(params))
+    de
+end

src/systems/optimization/optimizationsystem.jl

@@ -63,11 +63,12 @@ end
 
 function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys);
                           sparse = false, kwargs...)
-    hes = calculate_hessian(sys)
     if sparse
-        hes = sparse(hes)
+        hess = sparsehessian(equations(sys),[dv() for dv in states(sys)])
+    else
+        hess = calculate_hessian(sys)
     end
-    return build_function(hes, convert.(Variable,vs), convert.(Variable,ps);
+    return build_function(hess, convert.(Variable,vs), convert.(Variable,ps);
                           conv = AbstractSysToExpr(sys),kwargs...)
 end
 
@@ -84,12 +85,18 @@ namespace_operation(sys::OptimizationSystem) = namespace_operation(sys.op,sys.na
 hessian_sparsity(sys::OptimizationSystem) =
     hessian_sparsity(sys.op,[dv() for dv in states(sys)])
 
+struct AutoModelingToolkit <: DiffEqBase.AbstractADType end
+
+DiffEqBase.OptimizationProblem(sys::OptimizationSystem,args...;kwargs...) =
+    DiffEqBase.OptimizationProblem{true}(sys::OptimizationSystem,args...;kwargs...)
+
 """
 ```julia
 function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
                                           parammap=DiffEqBase.NullParameters();
                                           u0=nothing, lb=nothing, ub=nothing,
-                                          hes = false, sparse = false,
+                                          grad = false,
+                                          hess = false, sparse = false,
                                           checkbounds = false,
                                           linenumbers = true, parallel=SerialForm(),
                                           kwargs...) where iip
@@ -98,30 +105,53 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
 Generates an OptimizationProblem from an OptimizationSystem and allows for automatically
 symbolically calculating numerical enhancements.
 """
-function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
+function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0,
                                           parammap=DiffEqBase.NullParameters();
-                                          u0=nothing, lb=nothing, ub=nothing,
-                                          hes = false, sparse = false,
+                                          lb=nothing, ub=nothing,
+                                          grad = false,
+                                          hess = false, sparse = false,
                                           checkbounds = false,
                                           linenumbers = true, parallel=SerialForm(),
                                           kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
 
     f = generate_function(sys,checkbounds=checkbounds,linenumbers=linenumbers,
-                              parallel=parallel,expression=Val{false})
+                              expression=Val{false})
     u0 = varmap_to_vars(u0,dvs)
+
+    if grad
+        grad_oop,grad_iip = generate_gradient(sys,checkbounds=checkbounds,linenumbers=linenumbers,
+                                  parallel=parallel,expression=Val{false})
+        _grad(u,p) = grad_oop(u,p)
+        _grad(J,u,p) = (grad_iip(J,u,p); J)
+    else
+        _grad = nothing
+    end
+
+    if hess
+        hess_oop,hess_iip = generate_hessian(sys,checkbounds=checkbounds,linenumbers=linenumbers,
+                                 sparse=sparse,parallel=parallel,expression=Val{false})
+       _hess(u,p) = hess_oop(u,p)
+       _hess(J,u,p) = (hess_iip(J,u,p); J)
+    else
+        _hess = nothing
+    end
+
+    _f = DiffEqBase.OptimizationFunction{iip,AutoModelingToolkit,typeof(f),typeof(_grad),typeof(_hess),Nothing,Nothing,Nothing,Nothing}(f,AutoModelingToolkit(),_grad,_hess,nothing,nothing,nothing,nothing,0)
+
     p = varmap_to_vars(parammap,ps)
     lb = varmap_to_vars(lb,dvs)
     ub = varmap_to_vars(ub,dvs)
-    OptimizationProblem(f,p;u0=u0,lb=lb,ub=ub,kwargs...)
+    OptimizationProblem{iip}(_f,u0,p;lb=lb,ub=ub,kwargs...)
 end
 
 """
 ```julia
 function DiffEqBase.OptimizationProblemExpr{iip}(sys::OptimizationSystem,
                                           parammap=DiffEqBase.NullParameters();
                                           u0=nothing, lb=nothing, ub=nothing,
+                                          grad = false,
                                           hes = false, sparse = false,
                                           checkbounds = false,
                                           linenumbers = true, parallel=SerialForm(),
@@ -134,18 +164,36 @@ calculating numerical enhancements.
 """
 struct OptimizationProblemExpr{iip} end
 
-function OptimizationProblemExpr{iip}(sys::OptimizationSystem,
+OptimizationProblemExpr(sys::OptimizationSystem,args...;kwargs...) =
+    OptimizationProblemExpr{true}(sys::OptimizationSystem,args...;kwargs...)
+    
+function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
                                           parammap=DiffEqBase.NullParameters();
-                                          u0=nothing, lb=nothing, ub=nothing,
-                                          hes = false, sparse = false,
+                                          lb=nothing, ub=nothing,
+                                          grad = false,
+                                          hess = false, sparse = false,
                                           checkbounds = false,
                                           linenumbers = false, parallel=SerialForm(),
                                           kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
-
+    idx = iip ? 2 : 1
     f = generate_function(sys,checkbounds=checkbounds,linenumbers=linenumbers,
-                              parallel=parallel,expression=Val{true})
+                              expression=Val{true})
+    if grad
+        _grad = generate_gradient(sys,checkbounds=checkbounds,linenumbers=linenumbers,
+                             parallel=parallel,expression=Val{false})[idx]
+    else
+        _grad = :nothing
+    end
+
+    if hess
+        _hess = generate_hessian(sys,checkbounds=checkbounds,linenumbers=linenumbers,
+                                         sparse=sparse,parallel=parallel,expression=Val{false})[idx]
+    else
+        _hess = :nothing
+    end
+
     u0 = varmap_to_vars(u0,dvs)
     p = varmap_to_vars(parammap,ps)
     lb = varmap_to_vars(lb,dvs)
@@ -154,8 +202,25 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem,
         f = $f
         p = $p
         u0 = $u0
+        grad = $_grad
+        hess = $_hess
         lb = $lb
         ub = $ub
-        OptimizationProblem(f,p;u0=u0,lb=lb,ub=ub,kwargs...)
+        _f = OptimizationFunction{$iip,typeof(f),typeof(grad),typeof(hess),Nothing,Nothing,Nothing,Nothing}(f,grad,hess,nothing,AutoModelingToolkit(),nothing,nothing,nothing,0)
+        OptimizationProblem{$iip}(_f,u0,p;lb=lb,ub=ub,kwargs...)
+    end
+end
+
+function OptimizationFunction(f, x, ::AutoModelingToolkit,p = DiffEqBase.NullParameters();
+                              grad=false, hess=false, cons = nothing, cons_j = nothing, cons_h = nothing,
+                              num_cons = 0, chunksize = 1, hv = nothing)
+
+    sys = modelingtoolkitize(OptimizationProblem(f,x,p))
+    u0map = states(sys) .=> x
+    if p == DiffEqBase.NullParameters()
+        parammap = DiffEqBase.NullParameters()
+    else
+        parammap = parameters(sys) .=> p
     end
+    OptimizationProblem(sys,u0map,parammap,grad=grad,hess=hess).f
 end

test/controlsystem.jl

@@ -13,4 +13,7 @@ eqs = [
 sys = ControlSystem(loss,eqs,t,[x,v],[u],[p])
 dt = 0.1
 tspan = (0.0,1.0)
-runge_kutta_discretize(sys,dt,tspan)
+sys = runge_kutta_discretize(sys,dt,tspan)
+
+u0 = rand(112) # guess for the state values
+prob = OptimizationProblem(sys,u0,[0.1],grad=true)

test/modelingtoolkitize.jl

@@ -1,4 +1,6 @@
-using OrdinaryDiffEq, ModelingToolkit
+using OrdinaryDiffEq, ModelingToolkit, Test
+using GalacticOptim, Optim
+
 const N = 32
 const xyd_brusselator = range(0,stop=1,length=N)
 brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
@@ -38,3 +40,22 @@ prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
                                      u0,(0.,11.5),p)
 
 modelingtoolkitize(prob_ode_brusselator_2d)
+
+## Optimization
+
+rosenbrock(x,p) =  (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+x0 = zeros(2)
+p  = [1.0,100.0]
+
+prob = OptimizationProblem(rosenbrock,x0,p)
+sys = modelingtoolkitize(prob) # symbolicitize me captain!
+
+prob = OptimizationProblem(sys,x0,p,grad=true,hess=true)
+sol = solve(prob,NelderMead())
+@test sol.minimum < 1e-8
+
+sol = solve(prob,BFGS())
+@test sol.minimum < 1e-8
+
+sol = solve(prob,Newton())
+@test sol.minimum < 1e-8

test/optimizationsystem.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, SparseArrays
+using ModelingToolkit, SparseArrays, Test, GalacticOptim, Optim
 
 @variables x y
 @parameters a b
@@ -21,3 +21,28 @@ generate_function(combinedsys)
 generate_gradient(combinedsys)
 generate_hessian(combinedsys)
 ModelingToolkit.hessian_sparsity(combinedsys)
+
+u0 = [
+    sys1.x=>1.0
+    sys1.y=>2.0
+    sys2.x=>3.0
+    sys2.y=>4.0
+    z=>5.0
+]
+p = [
+    sys1.a => 6.0
+    sys1.b => 7.0
+    sys2.a => 8.0
+    sys2.b => 9.0
+    β => 10.0
+]
+
+prob = OptimizationProblem(combinedsys,u0,p,grad=true)
+sol = solve(prob,NelderMead())
+@test sol.minimum < -1e5
+
+prob2 = remake(prob,u0=sol.minimizer)
+sol = solve(prob,BFGS(initial_stepnorm=0.0001),allow_f_increases=true)
+@test sol.minimum < -1e8
+sol = solve(prob2,BFGS(initial_stepnorm=0.0001),allow_f_increases=true)
+@test sol.minimum < -1e9