Merge branch 'daenew' of https://github.com/ParasPuneetSingh/Optimization.jl into daenew

Author: ParasPuneetSingh

lib/OptimizationODE/Project.toml

@@ -6,10 +6,14 @@ version = "0.1.0"
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
+Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 
 [compat]
 ForwardDiff = "1"

lib/OptimizationODE/src/OptimizationODE.jl

@@ -23,7 +23,7 @@ struct DAEOptimizer{T}
     solver::T
 end
 
-DAEMassMatrix() = DAEOptimizer(Rodas5())
+DAEMassMatrix() = DAEOptimizer(Rosenbrock23(autodiff = false))
 DAEIndexing() = DAEOptimizer(IDA())
 
 
@@ -63,32 +63,6 @@ function SciMLBase.__init(prob::OptimizationProblem, opt::DAEOptimizer;
 end
 
 
-function solve_constrained_root(cache, u0, p)
-    n = length(u0)
-    cons_vals = cache.f.cons(u0, p)
-    m = length(cons_vals)
-    function resid!(res, u)
-        temp = similar(u)
-        f_mass!(temp, u, p, 0.0)
-        res .= temp
-    end
-    u0_ext = vcat(u0, zeros(m))
-    prob_nl = NonlinearProblem(resid!, u0_ext, p)
-    sol_nl = solve(prob_nl, Newton(); tol = 1e-8, maxiters = 100000,
-        callback = cache.callback, progress = get(cache.solver_args, :progress, false))
-    u_ext = sol_nl.u
-    return u_ext[1:n], sol_nl.retcode
-end
-
-
-function get_solver_type(opt::DAEOptimizer)
-    if opt.solver isa Union{Rodas5, RadauIIA5, ImplicitEuler, Trapezoid}
-        return :mass_matrix
-    else
-        return :indexing
-    end
-end
-
 function handle_parameters(p)
     if p isa SciMLBase.NullParameters
         return Float64[]
@@ -110,45 +84,6 @@ function setup_progress_callback(cache, solve_kwargs)
     return solve_kwargs
 end
 
-function finite_difference_jacobian(f, x; ϵ = 1e-8)
-    n = length(x)
-    fx = f(x)
-    if fx === nothing
-        return zeros(eltype(x), 0, n)
-    elseif isa(fx, Number)
-        J = zeros(eltype(fx), 1, n)
-        for j in 1:n
-            xj = copy(x)
-            xj[j] += ϵ
-            diff = f(xj)
-            if diff === nothing
-                diffval = zero(eltype(fx))
-            else
-                diffval = diff - fx
-            end
-            J[1, j] = diffval / ϵ
-        end
-        return J
-    else
-        m = length(fx)
-        J = zeros(eltype(fx), m, n)
-        for j in 1:n
-            xj = copy(x)
-            xj[j] += ϵ
-            fxj = f(xj)
-            if fxj === nothing
-                @inbounds for i in 1:m
-                    J[i, j] = -fx[i] / ϵ
-                end
-            else
-                @inbounds for i in 1:m
-                    J[i, j] = (fxj[i] - fx[i]) / ϵ
-                end
-            end
-        end
-        return J
-    end
-end
 
 function SciMLBase.__solve(
     cache::OptimizationCache{F,RC,LB,UB,LC,UC,S,O,D,P,C}
@@ -163,8 +98,7 @@ function SciMLBase.__solve(
     if cache.opt isa ODEOptimizer
         return solve_ode(cache, dt, maxit, u0, p)
     else
-        solver_method = get_solver_type(cache.opt)
-        if solver_method == :mass_matrix
+        if cache.opt.solver == Rosenbrock23(autodiff = false)
             return solve_dae_mass_matrix(cache, dt, maxit, u0, p)
         else
             return solve_dae_indexing(cache, dt, maxit, u0, p, differential_vars)
@@ -240,8 +174,8 @@ function solve_dae_mass_matrix(cache, dt, maxit, u0, p)
     n = length(u0)
     m = length(cons_vals)
     u0_extended = vcat(u0, zeros(m))
-    M = zeros(n + m, n + m)
-    M[1:n, 1:n] = I(n)
+    M = Diagonal(ones(n + m))
+
 
     function f_mass!(du, u, p_, t)
         x = @view u[1:n]
@@ -253,31 +187,18 @@ function solve_dae_mass_matrix(cache, dt, maxit, u0, p)
             grad_f .= ForwardDiff.gradient(z -> cache.f.f(z, p_), x)
         end
         J = Matrix{eltype(x)}(undef, m, n)
-        if cache.f.cons_j !== nothing
-            cache.f.cons_j(J, x)
-        else
-            J .= finite_difference_jacobian(z -> cache.f.cons(z, p_), x)
-        end
+        cache.f.cons_j !== nothing && cache.f.cons_j(J, x)
+
         @. du[1:n] = -grad_f - (J' * λ)
         consv = cache.f.cons(x, p_)
-        if consv === nothing
-            fill!(du[n+1:end], zero(eltype(x)))
-        else
-            if isa(consv, Number)
-                @assert m == 1
-                du[n+1] = consv
-            else
-                @assert length(consv) == m
-                @. du[n+1:end] = consv
-            end
-        end
+        @. du[n+1:end] = consv
         return nothing
     end
 
     if m == 0
-        optf = ODEFunction(f_mass!, mass_matrix = I(n))
+        optf = ODEFunction(f_mass!)
         prob = ODEProblem(optf, u0, (0.0, 1.0), p)
-        return solve(prob, HighOrderDescent(); dt=dt, maxiters=maxit)
+        return solve(prob, cache.opt.solver; dt=dt, maxiters=maxit)
     end
 
     ss_prob = SteadyStateProblem(ODEFunction(f_mass!, mass_matrix = M), u0_extended, p)
@@ -327,11 +248,8 @@ function solve_dae_indexing(cache, dt, maxit, u0, p, differential_vars)
         grad_f = similar(x)
         cache.f.grad(grad_f, x, p_)
         J = zeros(m, n)
-        if cache.f.cons_j !== nothing
-            cache.f.cons_j(J, x)
-        else
-            J .= finite_difference_jacobian(z -> cache.f.cons(z,p_), x)
-        end
+        cache.f.cons_j !== nothing && cache.f.cons_j(J, x)
+
         @. res[1:n] = du_x + grad_f + J' * λ
         consv = cache.f.cons(x, p_)
         @. res[n+1:end] = consv
@@ -364,4 +282,4 @@ function solve_dae_indexing(cache, dt, maxit, u0, p, differential_vars)
 end
 
 
-end 
+end 

lib/OptimizationODE/test/runtests.jl

@@ -5,40 +5,40 @@ using LinearAlgebra, ForwardDiff
 using OrdinaryDiffEq, DifferentialEquations, SteadyStateDiffEq, Sundials
 
 # Test helper functions
-function rosenbrock(x, p,args...)
+function rosenbrock(x, p)
     return (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
 end
 
-function rosenbrock_grad!(grad, x, p,args...)
+function rosenbrock_grad!(grad, x, p)
     grad[1] = -2.0 * (p[1] - x[1]) - 4.0 * p[2] * (x[2] - x[1]^2) * x[1]
     grad[2] = 2.0 * p[2] * (x[2] - x[1]^2)
 end
 
-function quadratic(x, p,args...)
+function quadratic(x, p)
     return (x[1] - p[1])^2 + (x[2] - p[2])^2
 end
 
-function quadratic_grad!(grad, x, p,args...)
+function quadratic_grad!(grad, x, p)
     grad[1] = 2.0 * (x[1] - p[1])
     grad[2] = 2.0 * (x[2] - p[2])
 end
 
 # Constrained optimization problem
-function constrained_objective(x, p,args...)
+function constrained_objective(x, p)
     return x[1]^2 + x[2]^2
 end
 
-function constrained_objective_grad!(grad, x, p,args...)
+function constrained_objective_grad!(grad, x, p)
     grad[1] = 2.0 * x[1]
     grad[2] = 2.0 * x[2]
 end
 
-function constraint_func(res, x, p,args...)
+function constraint_func(res, x, p)
     res[1] = x[1] + x[2] - 1.0  # x[1] + x[2] = 1
     return x[1] + x[2] - 1.0
 end
 
-function constraint_jac!(jac, x, p,args...)
+function constraint_jac!(jac, x, p)
     jac[1, 1] = 1.0
     jac[1, 2] = -1.0
 end
@@ -102,21 +102,21 @@ end
     # Minimize f(x) = x₁² + x₂²
     # Subject to x₁ - x₂ = 1
 
-    function constrained_objective(x, p,args...)
+    function constrained_objective(x, p)
         return x[1]^2 + x[2]^2
     end
 
-    function constrained_objective_grad!(g, x, p, args...)
+    function constrained_objective_grad!(g, x, p)
         g .= 2 .* x .* p[1]
         return nothing
     end
 
     # Constraint: x₁ - x₂ - p[1] = 0  (p[1] = 1 → x₁ - x₂ = 1)
-    function constraint_func(x, p, args...)
+    function constraint_func(x, p)
         return x[1] - x[2] - p[1]
     end
 
-    function constraint_jac!(J, x,args...)
+    function constraint_jac!(J, x)
         J[1, 1] = 1.0
         J[1, 2] = -1.0
         return nothing
@@ -159,8 +159,8 @@ end
             x0 = [0.0, 0.0]
             p=Float64[]  # No parameters provided
             # Create a problem with NullParameters
-            optf = OptimizationFunction((x, p, args...) -> sum(x.^2), 
-                                      grad=(grad, x, p, args...) -> (grad .= 2.0 .* x))
+            optf = OptimizationFunction((x, p) -> sum(x.^2), 
+                                      grad=(grad, x, p) -> (grad .= 2.0 .* x))
             prob = OptimizationProblem(optf, x0,p)  # No parameters provided
 
             opt = ODEGradientDescent()
@@ -233,26 +233,14 @@ end
             x = [1.0, 2.0]
             f(x) = [x[1]^2 + x[2], x[1] * x[2]]
 
-            J = OptimizationODE.finite_difference_jacobian(f, x)
+            J = ForwardDiff.jacobian(f, x)
 
             expected_J = [2.0 1.0; 2.0 1.0]
 
             @test isapprox(J, expected_J, atol=1e-6)
         end
     end
-    
-    @testset "Solver Type Detection" begin
-        @testset "Mass Matrix Solvers" begin
-            opt = DAEMassMatrix()
-            @test OptimizationODE.get_solver_type(opt) == :mass_matrix
-        end
-        
-        @testset "Index Method Solvers" begin
-            opt = DAEIndexing()
-            @test OptimizationODE.get_solver_type(opt) == :indexing
-        end
-    end
-    
+     
     @testset "Error Handling and Edge Cases" begin
         @testset "Empty Constraints" begin
             x0 = [1.5, 0.5]
@@ -274,8 +262,8 @@ end
             x0 = [0.5]
             p = [1.0]
 
-            single_var_func(x, p,args...) = (x[1] - p[1])^2
-            single_var_grad!(grad, x, p,args...) = (grad[1] = 2.0 * (x[1] - p[1]))
+            single_var_func(x, p) = (x[1] - p[1])^2
+            single_var_grad!(grad, x, p) = (grad[1] = 2.0 * (x[1] - p[1]))
 
             optf = OptimizationFunction(single_var_func; grad=single_var_grad!)
             prob = OptimizationProblem(optf, x0, p)
@@ -287,4 +275,4 @@ end
             @test isapprox(sol.u[1], p[1], atol=1e-1)
         end
     end
-end
+end