Integrate with POI to improve UX

Project.toml

@@ -12,7 +12,9 @@ LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 MathOptSetDistances = "3b969827-a86c-476c-9527-bb6f1a8fbad5"
+ParametricOptInterface = "0ce4ce61-57bf-432b-a095-efac525d185e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [compat]
 BlockDiagonals = "0.1"
@@ -22,4 +24,5 @@ JuMP = "1"
 LazyArrays = "0.21, 0.22, 1"
 MathOptInterface = "1.18"
 MathOptSetDistances = "0.2.7"
+ParametricOptInterface = "0.9.0"
 julia = "1.6"

README.md

@@ -31,7 +31,76 @@ examples, tutorials, and an API reference.
 
 ## Use with JuMP
 
-Use DiffOpt with JuMP by following this brief example:
+### DiffOpt-JuMP API with `Parameters`
+
+```julia
+using JuMP, DiffOpt, HiGHS
+
+model = Model(
+    () -> DiffOpt.diff_optimizer(
+        HiGHS.Optimizer;
+        with_parametric_opt_interface = true,
+    ),
+)
+set_silent(model)
+
+p_val = 4.0
+pc_val = 2.0
+@variable(model, x)
+@variable(model, p in Parameter(p_val))
+@variable(model, pc in Parameter(pc_val))
+@constraint(model, cons, pc * x >= 3 * p) #??? InvalidConstraintRef TODO
+@objective(model, Min, 2x)
+optimize!(model)
+@show value(x) == 3 * p_val / pc_val
+
+# the function is
+# x(p, pc) = 3p / pc
+# hence,
+# dx/dp = 3 / pc
+# dx/dpc = -3p / pc^2
+
+# First, try forward mode AD
+
+# differentiate w.r.t. p
+direction_p = 3.0
+MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(p), direction_p)
+DiffOpt.forward_differentiate!(model)
+@show MOI.get(model, DiffOpt.ForwardVariablePrimal(), x) == direction_p * 3 / pc_val
+
+# update p and pc
+p_val = 2.0
+pc_val = 6.0
+set_parameter_value(p, p_val)
+set_parameter_value(pc, pc_val)
+# re-optimize
+optimize!(model)
+# check solution
+@show value(x) ≈ 3 * p_val / pc_val
+
+# stop differentiating with respect to p
+DiffOpt.empty_input_sensitivities!(model)
+# differentiate w.r.t. pc
+direction_pc = 10.0
+MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(pc), direction_pc)
+DiffOpt.forward_differentiate!(model)
+@show abs(MOI.get(model, DiffOpt.ForwardVariablePrimal(), x) -
+    -direction_pc * 3 * p_val / pc_val^2) < 1e-5
+
+# always a good practice to clear previously set sensitivities
+DiffOpt.empty_input_sensitivities!(model)
+# Now, reverse model AD
+direction_x = 10.0
+MOI.set(model, DiffOpt.ReverseVariablePrimal(), x, direction_x)
+DiffOpt.reverse_differentiate!(model)
+@show MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p)) == direction_x * 3 / pc_val
+@show abs(MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(pc)) -
+    -direction_x * 3 * p_val / pc_val^2) < 1e-5
+```
+
+### Low level DiffOpt-JuMP API:
+
+A brief example:
 
 ```julia
 using JuMP, DiffOpt, HiGHS

docs/src/manual.md

@@ -4,9 +4,9 @@
     As of now, this package only works for optimization models that can be written either in convex conic form or convex quadratic form.
 
 
-## Supported objectives & constraints - scheme 1
+## Supported objectives & constraints - `QuadraticProgram` backend
 
-For `QPTH`/`OPTNET` style backend, the package supports following `Function-in-Set` constraints: 
+For `QuadraticProgram` backend, the package supports following `Function-in-Set` constraints: 
 
 |  MOI Function | MOI Set |
 |:-------|:---------------|
@@ -26,9 +26,9 @@ and the following objective types:
 | `ScalarQuadraticFunction`  | 
 
 
-## Supported objectives & constraints - scheme 2
+## Supported objectives & constraints - `ConicProgram` backend
 
-For `DiffCP`/`CVXPY` style backend, the package supports following `Function-in-Set` constraints: 
+For the `ConicProgram` backend backend, the package supports following `Function-in-Set` constraints: 
 
 |  MOI Function | MOI Set |
 |:-------|:---------------|
@@ -50,18 +50,22 @@ and the following objective types:
 |   `VariableIndex`   |
 |   `ScalarAffineFunction`   |
 
+Other conic sets such as `RotatedSecondOrderCone` and `PositiveSemidefiniteConeSquare` are supported through bridges.
 
-## Creating a differentiable optimizer
+
+## Creating a differentiable MOI optimizer
 
 You can create a differentiable optimizer over an existing MOI solver by using the `diff_optimizer` utility. 
 ```@docs
 diff_optimizer
 ```
 
-## Adding new sets and constraints
+## Creating a differentiable JuMP model
 
-The DiffOpt `Optimizer` behaves similarly to other MOI Optimizers
-and implements the `MOI.AbstractOptimizer` API.
+You initialize a differentiable JuMP model by using the `diff_model` utility. 
+```@docs
+diff_model
+```
 
 ## Projections on cone sets
 
@@ -104,6 +108,4 @@ In the light of above, DiffOpt differentiates program variables ``x``, ``s``, ``
 - OptNet: Differentiable Optimization as a Layer in Neural Networks
 
 ### Backward Pass vector
-One possible point of confusion in finding Jacobians is the role of the backward pass vector - above eqn (7), *OptNet: Differentiable Optimization as a Layer in Neural Networks*. While differentiating convex programs, it is often the case that we don't want to find the actual derivatives, rather we might be interested in computing the product of Jacobians with a *backward pass vector*, often used in backprop in machine learning/automatic differentiation. This is what happens in scheme 1 of `DiffOpt` backend.
-
-But, for the conic system (scheme 2), we provide perturbations in conic data (`dA`, `db`, `dc`) to compute pertubations (`dx`, `dy`, `dz`) in input variables. Unlike the quadratic case, these perturbations are actual derivatives, not the product with a backward pass vector. This is an important distinction between the two schemes of differential optimization.
+One possible point of confusion in finding Jacobians is the role of the backward pass vector - above eqn (7), *OptNet: Differentiable Optimization as a Layer in Neural Networks*. While differentiating convex programs, it is often the case that we don't want to find the actual derivatives, rather we might be interested in computing the product of Jacobians with a *backward pass vector*, often used in backpropagation in machine learning/automatic differentiation. This is what happens in `DiffOpt` backends.

src/DiffOpt.jl

@@ -12,13 +12,15 @@ import LazyArrays
 import LinearAlgebra
 import MathOptInterface as MOI
 import MathOptSetDistances as MOSD
+import ParametricOptInterface as POI
 import SparseArrays
 
 include("utils.jl")
 include("product_of_sets.jl")
 include("diff_opt.jl")
 include("moi_wrapper.jl")
 include("jump_moi_overloads.jl")
+include("parameters.jl")
 
 include("copy_dual.jl")
 include("bridges.jl")
@@ -40,4 +42,7 @@ end
 
 export diff_optimizer
 
+# TODO
+# add precompilation statements
+
 end # module

src/bridges.jl

@@ -43,6 +43,25 @@
         MOI.get(model, attr, bridge.vector_constraint),
     )[1]
 end
+
+function MOI.set(
+    model::MOI.ModelLike,
+    attr::ForwardConstraintFunction,
+    bridge::MOI.Bridges.Constraint.ScalarizeBridge{T},
+    value,
+) where {T}
+    MOI.set.(model, attr, bridge.scalar_constraints, value)
+    return
+end
+
+function MOI.get(
+    model::MOI.ModelLike,
+    attr::ReverseConstraintFunction,
+    bridge::MOI.Bridges.Constraint.ScalarizeBridge,
+)
+    return _vectorize(MOI.get.(model, attr, bridge.scalar_constraints))
+end
+
 function MOI.get(
     model::MOI.ModelLike,
     attr::DiffOpt.ReverseConstraintFunction,

src/diff_opt.jl

@@ -58,6 +58,17 @@ The output solution differentials can be queried with the attribute
 """
 function forward_differentiate! end
 
+"""
+    empty_input_sensitivities!(model::MOI.ModelLike)
+
+Empty the input sensitivities of the model.
+Sets to zero all the sensitivities set by the user with method such as:
+- `MOI.set(model, DiffOpt.ReverseVariablePrimal(), variable_index, value)`
+- `MOI.set(model, DiffOpt.ForwardObjectiveFunction(), expression)`
+- `MOI.set(model, DiffOpt.ForwardConstraintFunction(), index, expression)`
+"""
+function empty_input_sensitivities! end
+
 """
     ForwardObjectiveFunction <: MOI.AbstractModelAttribute
 

src/jump_moi_overloads.jl

@@ -307,6 +307,11 @@
     return forward_differentiate!(JuMP.backend(model))
 end
 
+function empty_input_sensitivities!(model::JuMP.Model)
+    empty_input_sensitivities!(JuMP.backend(model))
+    return
+end
+
 # MOI.Utilities
 
 function reverse_differentiate!(model::MOI.Utilities.CachingOptimizer)
@@ -317,6 +322,11 @@
     return forward_differentiate!(model.optimizer)
 end
 
+function empty_input_sensitivities!(model::MOI.Utilities.CachingOptimizer)
+    empty_input_sensitivities!(model.optimizer)
+    return
+end
+
 # MOIB
 
 function reverse_differentiate!(model::MOI.Bridges.AbstractBridgeOptimizer)
@@ -326,3 +336,8 @@
 function forward_differentiate!(model::MOI.Bridges.AbstractBridgeOptimizer)
     return forward_differentiate!(model.model)
 end
+
+function empty_input_sensitivities!(model::MOI.Bridges.AbstractBridgeOptimizer)
+    empty_input_sensitivities!(model.model)
+    return
+end

src/moi_wrapper.jl

@@ -4,7 +4,7 @@
 # in the LICENSE.md file or at https://opensource.org/licenses/MIT.
 
 """
-    diff_optimizer(optimizer_constructor)::Optimizer
+    diff_optimizer(optimizer_constructor)
 
 Creates a `DiffOpt.Optimizer`, which is an MOI layer with an internal optimizer
 and other utility methods. Results (primal, dual and slack values) are obtained
@@ -21,19 +21,35 @@
 julia> model.add_constraint(model, ...)
 ```
 """
-function diff_optimizer(optimizer_constructor)::Optimizer
-    optimizer =
-        MOI.instantiate(optimizer_constructor; with_bridge_type = Float64)
+function diff_optimizer(
+    optimizer_constructor;
+    method = nothing,
+    with_parametric_opt_interface::Bool = false,
+    with_bridge_type = Float64,
+    with_cache::Bool = true,
+)
+    optimizer = MOI.instantiate(
+        optimizer_constructor;
+        with_bridge_type = with_bridge_type,
+    )
     # When we do `MOI.copy_to(diff, optimizer)` we need to efficiently `MOI.get`
     # the model information from `optimizer`. However, 1) `optimizer` may not
     # implement some getters or it may be inefficient and 2) the getters may be
     # unimplemented or inefficient through some bridges.
     # For this reason we add a cache layer, the same cache JuMP adds.
-    caching_opt = MOI.Utilities.CachingOptimizer(
-        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
-        optimizer,
-    )
-    return Optimizer(caching_opt)
+    caching_opt = if with_cache
+        MOI.Utilities.CachingOptimizer(
+            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
+            optimizer,
+        )
+    else
+        optimizer
+    end
+    if with_parametric_opt_interface
+        return POI.Optimizer(Optimizer(caching_opt; method = method))
+    else
+        return Optimizer(caching_opt; method = method)
+    end
 end
 
 mutable struct Optimizer{OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
@@ -49,10 +65,16 @@
     # sensitivity input cache using MOI like sparse format
     input_cache::InputCache
 
-    function Optimizer(optimizer::OT) where {OT<:MOI.ModelLike}
+    function Optimizer(
+        optimizer::OT;
+        method = nothing,
+    ) where {OT<:MOI.ModelLike}
         output =
             new{OT}(optimizer, Any[], nothing, nothing, nothing, InputCache())
         add_all_model_constructors(output)
+        if method !== nothing
+            output.model_constructor = method
+        end
         return output
     end
 end
@@ -552,6 +574,11 @@
     return forward_differentiate!(diff)
 end
 
+function empty_input_sensitivities!(model::Optimizer)
+    empty!(model.input_cache)
+    return
+end
+
 function _instantiate_with_bridges(model_constructor)
     model = MOI.Bridges.LazyBridgeOptimizer(MOI.instantiate(model_constructor))
     # We don't add any variable bridge here because: