ExaModels support

Author: klamike

README.md

@@ -12,7 +12,7 @@ MadDiff implements forward and reverse mode implicit differentiation for MadSuit
 ## NLPModels interface
 
 > [!NOTE]
-> The [NLPModels](https://github.com/JuliaSmoothOptimizers/NLPModels.jl) interface requires that your `AbstractNLPModel` implementation includes the [`ParametricNLPModels`](https://github.com/klamike/ParametricNLPModels.jl/tree/mk/pnlpm) API. Currently, this is automated only for the case when using MadNLP through JuMP, but support for ExaModels, ADNLPModels, and NLPModelsJuMP is planned.
+> The [NLPModels](https://github.com/JuliaSmoothOptimizers/NLPModels.jl) interface requires that your `AbstractNLPModel` implementation includes the [`ParametricNLPModels`](https://github.com/klamike/ParametricNLPModels.jl/tree/mk/pnlpm) API. Currently, this is automated only for the case when using MadNLP through JuMP or when using ExaModels; support for other solvers and frameworks is planned.
 
 
 ```julia

docs/src/index.md

@@ -10,7 +10,7 @@ MadDiff implements forward and reverse mode implicit differentiation for MadSuit
 
 ## NLPModels interface
 
-> The [NLPModels](https://github.com/JuliaSmoothOptimizers/NLPModels.jl) interface requires that your `AbstractNLPModel` implementation includes the [`ParametricNLPModels`](https://github.com/klamike/ParametricNLPModels.jl/tree/mk/pnlpm) API. Currently, this is automated only for the case when using MadNLP through JuMP, but support for ExaModels, ADNLPModels, and NLPModelsJuMP is planned.
+> The [NLPModels](https://github.com/JuliaSmoothOptimizers/NLPModels.jl) interface requires that your `AbstractNLPModel` implementation includes the [`ParametricNLPModels`](https://github.com/klamike/ParametricNLPModels.jl/tree/mk/pnlpm) API. Currently, this is automated only for the case when using MadNLP through JuMP or when using ExaModels; support for other solvers and frameworks is planned.
 
 
 ```julia

test/Project.toml

@@ -22,6 +22,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 DiffOpt = {rev = "mk/allow_obj_and_sol", url = "https://github.com/klamike/DiffOpt.jl.git"}
+ExaModels = {rev = "mk/param_ad", url = "https://github.com/klamike/ExaModels.jl"}
 HybridKKT = {rev = "mk/latest", url = "https://github.com/klamike/HybridKKT.jl.git"}
 MadDiff = {path = ".."}
 MadIPM = {rev = "mk/flip", url = "https://github.com/klamike/MadIPM.jl.git"}

test/runtests.jl

@@ -1,7 +1,7 @@
 using Test, Random, LinearAlgebra
 using MadDiff
 using MadNLP, MadIPM, MadNCL, HybridKKT
-using NLPModels, CUDA, MadNLPGPU, MadNLPTests, QuadraticModels
+using NLPModels, CUDA, MadNLPGPU, MadNLPTests, QuadraticModels, ExaModels
 using JuMP, DiffOpt, MathOptInterface
 const MOI = MathOptInterface
 

test/test_jacobian.jl

@@ -1,4 +1,4 @@
-using NLPModels
+using NLPModels, LinearAlgebra, ExaModels
 
 @testset "jac/jact/jvp/vjp" begin
     function _check_consistency(sens; atol = 1e-8)
@@ -91,3 +91,51 @@ using NLPModels
     optimize!(model_rect)
     _check_consistency(MadDiff.MadDiffSolver(unsafe_backend(model_rect).inner.solver); atol = 1e-8)
 end
+
+@testset "ExaModels JVP/VJP vs FiniteDiff" begin
+    p0 = [1.0, 3.0]
+    h  = 1e-5
+    atol = sqrt(h)
+
+    function _make_exa(p_vals)
+        c = ExaCore()
+        p = ExaModels.parameter(c, p_vals)
+        x = ExaModels.variable(c, 2)
+        ExaModels.objective(c, x[1]^2 + x[2]^2 + p[1] * x[1])
+        ExaModels.constraint(c, x[1] + x[2] - p[2])
+        return ExaModel(c)
+    end
+
+    function _solve_exa(p_vals)
+        m = _make_exa(p_vals)
+        solver = MadNLP.MadNLPSolver(m; print_level = MadNLP.ERROR)
+        MadNLP.solve!(solver)
+        return solver
+    end
+
+    solver0 = _solve_exa(p0)
+    sens = MadDiffSolver(solver0)
+    x0   = Vector(MadNLP.variable(solver0.x))
+    y0   = Vector(solver0.y)
+    n_p  = sens.n_p
+
+    for j in 1:n_p
+        Δp    = zeros(n_p); Δp[j] = 1.0
+        jvp   = MadDiff.jacobian_vector_product!(sens, Δp)
+        s_plus = _solve_exa(p0 .+ h .* Δp)
+        @test isapprox(jvp.dx, (Vector(MadNLP.variable(s_plus.x)) .- x0) ./ h; atol)
+        @test isapprox(jvp.dy, (Vector(s_plus.y)                  .- y0) ./ h; atol)
+    end
+
+    rng   = MersenneTwister(42)
+    dL_dx = randn(rng, length(x0))
+    dL_dy = randn(rng, length(y0))
+    vjp   = MadDiff.vector_jacobian_product!(sens; dL_dx, dL_dy)
+    for j in 1:n_p
+        Δp     = zeros(n_p); Δp[j] = 1.0
+        s_plus = _solve_exa(p0 .+ h .* Δp)
+        dx_fd  = (Vector(MadNLP.variable(s_plus.x)) .- x0) ./ h
+        dy_fd  = (Vector(s_plus.y)                  .- y0) ./ h
+        @test isapprox(vjp.grad_p[j], dot(dL_dx, dx_fd) + dot(dL_dy, dy_fd); atol)
+    end
+end