From 4357b5d42778b83a372563fc6b7124ba9572c279 Mon Sep 17 00:00:00 2001
From: Vaibhav Dixit <vaibhavyashdixit@gmail.com>
Date: Thu, 19 Oct 2023 17:46:02 -0400
Subject: [PATCH] Use linearity detection, wip

---
 .gitignore                                    |  2 +-
 lib/OptimizationPRIMA/Project.toml            |  2 ++
 .../src/OptimizationPRIMA.jl                  | 33 +++++++++++++++++--
 lib/OptimizationPRIMA/test/runtests.jl        |  5 +++
 4 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index bd2b72711..1f9682187 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,5 @@
 .DS_Store
-/Manifest.toml
+Manifest.toml
 /dev/
 /docs/build/
 .vscode
\ No newline at end of file
diff --git a/lib/OptimizationPRIMA/Project.toml b/lib/OptimizationPRIMA/Project.toml
index 67fd4e60e..1a6980637 100644
--- a/lib/OptimizationPRIMA/Project.toml
+++ b/lib/OptimizationPRIMA/Project.toml
@@ -4,8 +4,10 @@ authors = ["Vaibhav Dixit <vaibhavyashdixit@gmail.com> and contributors"]
 version = "1.0.0-DEV"
 
 [deps]
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 PRIMA = "0a7d04aa-8ac2-47b3-b7a7-9dbd6ad661ed"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 julia = "1"
diff --git a/lib/OptimizationPRIMA/src/OptimizationPRIMA.jl b/lib/OptimizationPRIMA/src/OptimizationPRIMA.jl
index 8ff76ea4f..884136625 100644
--- a/lib/OptimizationPRIMA/src/OptimizationPRIMA.jl
+++ b/lib/OptimizationPRIMA/src/OptimizationPRIMA.jl
@@ -1,6 +1,7 @@
 module OptimizationPRIMA
 
-using PRIMA, Optimization, Optimization.SciMLBase
+using PRIMA, Optimization, Optimization.SciMLBase, SparseArrays
+import ModelingToolkit, ModelingToolkit.Symbolics
 
 abstract type PRIMASolvers end
 
@@ -19,6 +20,10 @@ function SciMLBase.__init(prob::SciMLBase.OptimizationProblem, opt::PRIMASolvers
     data = Optimization.DEFAULT_DATA;
     callback = (args...) -> (false),
     progress = false, kwargs...)
+    if opt isa COBYLA
+        sys = ModelingToolkit.modelingtoolkitize(prob)
+        prob = OptimizationProblem(sys, prob.u0, prob.p; lb = prob.lb, ub = prob.ub, cons_sparse = true)
+    end
     return OptimizationCache(prob, opt, data; callback, progress,
         kwargs...)
 end
@@ -127,11 +132,33 @@ function SciMLBase.__solve(cache::OptimizationCache{
 
     t0 = time()
     if cache.opt isa COBYLA
+        lininds = Int[]
+        nonlininds = Int[]
+        symboliclinexpr = []
+        for i in eachindex(cache.f.cons_expr)
+            println(cache.f.cons_expr[i])
+            println(isempty(cache.f.cons_hess_prototype[i]))
+            if isempty(cache.f.cons_hess_prototype[i])
+                push!(lininds, i)
+                symbolicexpr = Symbolics.parse_expr_to_symbolic(cache.f.cons_expr[i], (@__MODULE__,))
+                println(symbolicexpr)
+                push!(symboliclinexpr, symbolicexpr)
+            else
+                push!(nonlininds, i)
+            end
+        end
+        res1 = zeros(length(cache.f.cons_expr))
+        nonlincons = (res, θ) -> (cache.f.cons(res1, θ); res .= res1[nonlininds])
+        println(symboliclinexpr)
+        A = hcat(res1[lininds]...)
+        println(A)
+        b = cache.lcons[lininds]
+        println(b)
         function fwcons(θ, res)
-            cache.f.cons(res, θ, cache.p)
+            nonlincons(res, θ)
             return _loss(θ)
         end
-        (minx, minf, nf, rc, cstrv) = optfunc(fwcons, cache.u0; kws...)
+        (minx, minf, nf, rc, cstrv) = optfunc(fwcons, cache.u0; linear_eq = (A, b), nonlinear_ineq = length(nonlininds), kws...)
     elseif cache.opt isa LINCOA
         (minx, minf, nf, rc, cstrv) = optfunc(_loss, cache.u0; kws...)
     else
diff --git a/lib/OptimizationPRIMA/test/runtests.jl b/lib/OptimizationPRIMA/test/runtests.jl
index 2cec2c9e1..fc7c940da 100644
--- a/lib/OptimizationPRIMA/test/runtests.jl
+++ b/lib/OptimizationPRIMA/test/runtests.jl
@@ -18,4 +18,9 @@ using Test
     @test 10 * sol.objective < l1
     @test_throws SciMLBase.IncompatibleOptimizerError Optimization.solve(prob, COBYLA(), maxiters = 1000)
 
+    function con2_c(res, x, p)
+        res .= [x[1] + x[2], x[2] * sin(x[1]) - x[1]]
+    end
+    optprob = OptimizationFunction(rosenbrock, cons = con2_c)
+    prob = OptimizationProblem(, x0, _p, lcons = [-Inf, -Inf], ucons = [Inf, Inf])
 end