Add fields for manifold and curvature results in cache

Author: Vaibhavdixit02

Project.toml

@@ -1,7 +1,7 @@
 name = "OptimizationBase"
 uuid = "bca83a33-5cc9-4baa-983d-23429ab6bcbb"
 authors = ["Vaibhav Dixit <[email protected]> and contributors"]
-version = "1.3.1"
+version = "1.3.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/OptimizationBase.jl

@@ -10,8 +10,8 @@ end
 
 using ArrayInterface, Base.Iterators, SparseArrays, LinearAlgebra
 using SymbolicIndexingInterface
-using SymbolicAnalysis: propagate_sign, propagate_curvature, propagate_gcurvature,
-                        getcurvature, getgcurvature, getsign
+using SymbolicAnalysis
+using SymbolicAnalysis: AnalysisResult
 import Symbolics
 import Manifolds
 import Symbolics: variable, Equation, Inequality, unwrap, @variables

src/cache.jl

@@ -1,5 +1,11 @@
 import Symbolics: ≲, ~
-struct OptimizationCache{F, RC, LB, UB, LC, UC, S, O, D, P, C} <:
+
+struct AnalysisResults
+    objective::Union{Nothing, AnalysisResult}
+    constraints::Union{Nothing, Vector{AnalysisResult}}
+end
+
+struct OptimizationCache{F, RC, LB, UB, LC, UC, S, O, D, P, C, M} <:
        SciMLBase.AbstractOptimizationCache
     f::F
     reinit_cache::RC
@@ -12,6 +18,8 @@ struct OptimizationCache{F, RC, LB, UB, LC, UC, S, O, D, P, C} <:
     data::D
     progress::P
     callback::C
+    manifold::M
+    analysis_results::AnalysisResults
     solver_args::NamedTuple
 end
 
@@ -101,32 +109,44 @@ function OptimizationCache(prob::SciMLBase.OptimizationProblem, opt, data = DEFA
     end
 
     if obj_expr !== nothing
-        try
-            obj_expr = obj_expr |> Symbolics.unwrap
-            obj_expr = propagate_curvature(propagate_sign(obj_expr))
-            @info "Objective Euclidean curvature: $(getcurvature(obj_expr))"
-        catch
-            @info "No euclidean atom available"
+        obj_expr = obj_expr |> Symbolics.unwrap
+        if manifold === nothing
+            obj_res = analyze(obj_expr)
+        else
+            obj_res = analyze(obj_expr, manifold)
         end
 
-        try
-            obj_expr = propagate_gcurvature(propagate_sign(obj_expr), manifold)
-            @info "Objective Geodesic curvature: $(getgcurvature(obj_expr))"
-        catch
-            @info "No geodesic atom available"
+        @info "Objective Euclidean curvature: $(obj_res.curvature)"
+
+        if obj_res.gcurvature !== nothing
+            @info "Objective Geodesic curvature: $(obj_res.gcurvature)"
         end
+    else
+        obj_res = nothing
     end
 
     if cons_expr !== nothing
         cons_expr = cons_expr .|> Symbolics.unwrap
-        cons_expr = propagate_curvature.(propagate_sign.(cons_expr))
-        @info "Constraints Euclidean curvature: $(getcurvature.(cons_expr))"
+        if manifold === nothing
+            cons_res = analyze.(cons_expr)
+        else
+            cons_res = analyze.(cons_expr, Ref(manifold))
+        end
+        for i in 1:num_cons
+            @info "Constraints Euclidean curvature: $(cons_res[i].curvature)"
+
+            if cons_res[i].gcurvature !== nothing
+                @info "Constraints Geodesic curvature: $(cons_res[i].gcurvature)"
+            end
+        end
+    else
+        cons_res = nothing
     end
 
     return OptimizationCache(f, reinit_cache, prob.lb, prob.ub, prob.lcons,
         prob.ucons, prob.sense,
-        opt, data, progress, callback,
-        merge((; maxiters, maxtime, abstol, reltol, manifold),
+        opt, data, progress, callback, manifold, AnalysisResults(obj_res, cons_res),
+        merge((; maxiters, maxtime, abstol, reltol),
             NamedTuple(kwargs)))
 end
 

test/cvxtest.jl

@@ -9,6 +9,8 @@ optf = OptimizationFunction(f, Optimization.AutoForwardDiff())
 prob = OptimizationProblem(optf, [0.4], structural_analysis = true)
 
 @time sol = solve(prob, Optimization.LBFGS(), maxiters = 1000)
+@test sol.cache.analysis_results.objective.curvature == SymbolicAnalysis.Convex
+@test sol.cache.analysis_results.constraints === nothing
 
 x0 = zeros(2)
 rosenbrock(x, p = nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
@@ -17,6 +19,7 @@ l1 = rosenbrock(x0)
 optf = OptimizationFunction(rosenbrock, AutoEnzyme())
 prob = OptimizationProblem(optf, x0, structural_analysis = true)
 @time res = solve(prob, Optimization.LBFGS(), maxiters = 100)
+@test res.cache.analysis_results.objective.curvature == SymbolicAnalysis.UnknownCurvature
 
 function con2_c(res, x, p)
     res .= [x[1]^2 + x[2]^2, (x[2] * sin(x[1]) + x[1]) - 5]
@@ -26,6 +29,10 @@ optf = OptimizationFunction(rosenbrock, AutoZygote(), cons = con2_c)
 prob = OptimizationProblem(optf, x0, lcons = [1.0, -Inf], ucons = [1.0, 0.0],
     lb = [-1.0, -1.0], ub = [1.0, 1.0], structural_analysis = true)
 @time res = solve(prob, Optimization.LBFGS(), maxiters = 100)
+@test res.cache.analysis_results.objective.curvature == SymbolicAnalysis.Convex
+@test res.cache.analysis_results.constraints[1].curvature == SymbolicAnalysis.Convex
+@test res.cache.analysis_results.constraints[2].curvature ==
+      SymbolicAnalysis.UnknownCurvature
 
 m = 100
 σ = 0.005
@@ -41,3 +48,5 @@ prob = OptimizationProblem(optf, data2[1]; manifold = M, structural_analysis = t
 opt = OptimizationManopt.GradientDescentOptimizer()
 @time sol = solve(prob, Optimization.LBFGS(), maxiters = 100)
 @test sol.minimizer < 1e-1
+@test sol.cache.analysis_results.objective.curvature == SymbolicAnalysis.UnknownCurvature
+@test sol.cache.analysis_results.objective.gcurvature == SymbolicAnalysis.GConvex