Removes OptimizationFunction and add GalacticOptim as a dependency (#1563)

Author: Vaibhavdixit02

Project.toml

@@ -83,7 +83,6 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GalacticOptim = "a75be94c-b780-496d-a8a9-0878b188d577"
 GalacticOptimJL = "9d3c5eb1-403b-401b-8c0f-c11105342e6b"
-Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
@@ -94,4 +93,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BenchmarkTools", "ForwardDiff", "GalacticOptim", "GalacticOptimJL", "OrdinaryDiffEq", "Optim", "Random", "ReferenceTests", "SafeTestsets", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]
+test = ["BenchmarkTools", "ForwardDiff", "GalacticOptim", "GalacticOptimJL", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]

src/ModelingToolkit.jl

@@ -25,7 +25,6 @@ using DocStringExtensions
 using Base: RefValue
 using Combinatorics
 import IfElse
-
 import Distributions
 
 RuntimeGeneratedFunctions.init(@__MODULE__)
@@ -133,6 +132,7 @@ include("systems/diffeqs/basic_transformations.jl")
 include("systems/jumps/jumpsystem.jl")
 
 include("systems/nonlinear/nonlinearsystem.jl")
+include("systems/nonlinear/modelingtoolkitize.jl")
 
 include("systems/optimization/optimizationsystem.jl")
 

src/systems/nonlinear/modelingtoolkitize.jl

@@ -0,0 +1,47 @@
+"""
+$(TYPEDSIGNATURES)
+
+Generate `NonlinearSystem`, dependent variables, and parameters from an `NonlinearProblem`.
+"""
+function modelingtoolkitize(prob::NonlinearProblem; kwargs...)
+    p = prob.p
+    has_p = !(p isa Union{DiffEqBase.NullParameters,Nothing})
+
+    _vars = reshape([variable(:x, i) for i in eachindex(prob.u0)], size(prob.u0))
+
+    vars = prob.u0 isa Number ? _vars : ArrayInterface.restructure(prob.u0, _vars)
+    params = if has_p
+        _params = define_params(p)
+        p isa Number ? _params[1] : (p isa Tuple || p isa NamedTuple ? _params : ArrayInterface.restructure(p, _params))
+    else
+        []
+    end
+
+    if DiffEqBase.isinplace(prob)
+        rhs = ArrayInterface.restructure(prob.u0, similar(vars, Num))
+        prob.f(rhs, vars, params)
+    else
+        rhs = prob.f(vars, params)
+    end
+    out_def = prob.f(prob.u0, prob.p)
+    eqs = vcat([0.0 ~ rhs[i] for i in 1:length(out_def)]...)
+
+    sts = vec(collect(vars))
+
+    params = if params isa Number || (params isa Array && ndims(params) == 0)
+        [params[1]]
+    else
+        vec(collect(params))
+    end
+    default_u0 = Dict(sts .=> vec(collect(prob.u0)))
+    default_p = has_p ? Dict(params .=> vec(collect(prob.p))) : Dict()
+
+    de = NonlinearSystem(
+        eqs, sts, params,
+        defaults=merge(default_u0, default_p);
+        name=gensym(:MTKizedNonlinProb),
+        kwargs...
+    )
+
+    de
+end

src/systems/nonlinear/nonlinearsystem.jl

@@ -137,6 +137,23 @@ function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = paramete
     return build_function(jac, vs, ps; kwargs...)
 end
 
+function calculate_hessian(sys::NonlinearSystem; sparse=false, simplify=false)
+    rhs = [eq.rhs for eq ∈ equations(sys)]
+    vals = [dv for dv in states(sys)]
+    if sparse
+        hess = [sparsehessian(rhs[i], vals, simplify=simplify) for i in 1:length(rhs)]
+    else
+        hess = [hessian(rhs[i], vals, simplify=simplify) for i in 1:length(rhs)]
+    end
+    return hess
+end
+
+function generate_hessian(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys);
+                           sparse = false, simplify=false, kwargs...)
+    hess = calculate_hessian(sys,sparse=sparse, simplify=simplify)
+    return build_function(hess, vs, ps; kwargs...)
+end
+
 function generate_function(sys::NonlinearSystem, dvs = states(sys), ps = parameters(sys); kwargs...)
     rhss = [deq.rhs for deq ∈ equations(sys)]
     pre, sol_states = get_substitutions_and_solved_states(sys)

src/systems/optimization/optimizationsystem.jl

@@ -78,12 +78,11 @@ function OptimizationSystem(op, states, ps;
     process_variables!(var_to_name, defaults, states)
     process_variables!(var_to_name, defaults, ps)
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
-
     OptimizationSystem(
                        value(op), states, ps, var_to_name,
                        observed,
                        equality_constraints, inequality_constraints,
-                       name, systems, defaults, checks = checks
+                       name, systems, defaults; checks = checks
                       )
 end
 
@@ -122,8 +121,6 @@ namespace_expr(sys::OptimizationSystem) = namespace_expr(get_op(sys), sys)
 
 hessian_sparsity(sys::OptimizationSystem) = hessian_sparsity(get_op(sys), states(sys))
 
-struct AutoModelingToolkit <: DiffEqBase.AbstractADType end
-
 DiffEqBase.OptimizationProblem(sys::OptimizationSystem,args...;kwargs...) =
     DiffEqBase.OptimizationProblem{true}(sys::OptimizationSystem,args...;kwargs...)
 
@@ -175,7 +172,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         _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)
+    _f = DiffEqBase.OptimizationFunction{iip,SciMLBase.NoAD,typeof(f),typeof(_grad),typeof(_hess),Nothing,Nothing,Nothing,Nothing}(f, SciMLBase.NoAD(), _grad, _hess, nothing, nothing, nothing, nothing)
 
     defs = defaults(sys)
     defs = mergedefaults(defs,parammap,ps)
@@ -253,21 +250,7 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
         hess = $_hess
         lb = $lb
         ub = $ub
-        _f = OptimizationFunction{$iip,typeof(f),typeof(grad),typeof(hess),Nothing,Nothing,Nothing,Nothing}(f,grad,hess,nothing,AutoModelingToolkit(),nothing,nothing,nothing,0)
+        _f = OptimizationFunction{$iip,typeof(f),typeof(grad),typeof(hess),SciMLBase.NoAD,Nothing,Nothing,Nothing}(f,grad,hess,nothing,SciMLBase.NoAD(),nothing,nothing,nothing,0)
         OptimizationProblem{$iip}(_f,u0,p;lb=lb,ub=ub,kwargs...)
     end
 end
-
-function DiffEqBase.OptimizationFunction{iip}(f, ::AutoModelingToolkit, x, p = DiffEqBase.NullParameters();
-                              grad=false, hess=false, cons = nothing, cons_j = nothing, cons_h = nothing,
-                              num_cons = 0, chunksize = 1, hv = nothing) where iip
-
-    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

@@ -1,4 +1,4 @@
-using ModelingToolkit, GalacticOptim, Optim, GalacticOptimJL
+using ModelingToolkit, GalacticOptim, GalacticOptimJL
 
 @variables t x(t) v(t) u(t)
 @parameters p[1:2]

test/modelingtoolkitize.jl

@@ -1,5 +1,5 @@
 using OrdinaryDiffEq, ModelingToolkit, Test
-using GalacticOptim, Optim, RecursiveArrayTools, GalacticOptimJL
+using GalacticOptim, RecursiveArrayTools, GalacticOptimJL
 
 N = 32
 const xyd_brusselator = range(0,stop=1,length=N)

test/optimizationsystem.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, SparseArrays, Test, GalacticOptim, Optim, GalacticOptimJL
+using ModelingToolkit, SparseArrays, Test, GalacticOptim, GalacticOptimJL
 
 @variables x y
 @parameters a b
@@ -51,9 +51,9 @@ rosenbrock(x, p) =  (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
 x0 = zeros(2)
 _p  = [1.0, 100.0]
 
-f = OptimizationFunction(rosenbrock,ModelingToolkit.AutoModelingToolkit(),x0,_p,grad=true,hess=true)
+f = OptimizationFunction(rosenbrock,GalacticOptim.AutoModelingToolkit())
 prob = OptimizationProblem(f,x0,_p)
-sol = solve(prob,Optim.Newton())
+sol = solve(prob,Newton())
 
 # issue #819
 @testset "Combined system name collisions" begin