[WIP] Integrate with POI to improve UX

Author: joaquimg

Project.toml

@@ -12,6 +12,7 @@ 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"
 
 [compat]

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")
@@ -38,6 +40,16 @@ function add_all_model_constructors(model)
     return
 end
 
+function add_conic_model_constructor(model)
+    add_model_constructor(model, ConicProgram.Model)
+    return
+end
+
+function add_quadratic_model_constructor(model)
+    add_model_constructor(model, QuadraticProgram.Model)
+    return
+end
+
 export diff_optimizer
 
 end # module

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 @@ function forward_differentiate!(model::JuMP.Model)
     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 @@ function forward_differentiate!(model::MOI.Utilities.CachingOptimizer)
     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 @@ end
 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,37 @@
 # in the LICENSE.md file or at https://opensource.org/licenses/MIT.
 
 """
-    diff_optimizer(optimizer_constructor)::Optimizer
+    DiffMethod
+
+An enum to define the differentiation method.
+
+## Values
+
+Possible values are:
+
+ * [`DIFF_AUTOMATIC`]: Automatic differentiation: tries all differentiation methods and chooses the first that works. This can be slower than manually choosing a method.
+ * [`DIFF_CONIC`]: Conic optimization based differentiation: works for conic programs or programs that can be transformed into conic programs by JuMP.
+ * [`DIFF_QUADRATIC`]: Quadratic optimization based differentiation: works for quadratic programs or programs that can be transformed into conic programs by JuMP.
+"""
+@enum(DiffMethod, DIFF_AUTOMATIC, DIFF_CONIC, DIFF_QUADRATIC)
+
+@doc(
+    "Automatic differentiation: tries all differentiation methods and chooses the first that works. This can be slower than manually choosing a method.",
+    DIFF_AUTOMATIC
+)
+
+@doc(
+    "Conic optimization based differentiation: works for conic programs or programs that can be transformed into conic programs by JuMP.",
+    DIFF_CONIC
+)
+
+@doc(
+    "Quadratic optimization based differentiation: works for quadratic programs or programs that can be transformed into conic programs by JuMP.",
+    DIFF_QUADRATIC
+)
+
+"""
+    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 +51,33 @@ julia> x = model.add_variable(model)
 julia> model.add_constraint(model, ...)
 ```
 """
-function diff_optimizer(optimizer_constructor)::Optimizer
+function diff_optimizer(
+    optimizer_constructor;
+    method = DIFF_AUTOMATIC,
+    with_parametric_opt_interface::Bool = false,
+    with_bridge_type = Float64,
+    with_cache::Bool = true,
+)
     optimizer =
-        MOI.instantiate(optimizer_constructor; with_bridge_type = Float64)
+        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 +93,18 @@ mutable struct Optimizer{OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
     # sensitivity input cache using MOI like sparse format
     input_cache::InputCache
 
-    function Optimizer(optimizer::OT) where {OT<:MOI.ModelLike}
+    function Optimizer(optimizer::OT; method = DIFF_AUTOMATIC) where {OT<:MOI.ModelLike}
         output =
             new{OT}(optimizer, Any[], nothing, nothing, nothing, InputCache())
-        add_all_model_constructors(output)
+        if method == DIFF_CONIC
+            add_conic_model_constructor(output)
+        elseif method == DIFF_QUADRATIC
+            add_quadratic_model_constructor(output)
+        elseif method == DIFF_AUTOMATIC
+            add_all_model_constructors(output)
+        else
+            add_model_constructor(output, method)
+        end
         return output
     end
 end
@@ -552,6 +604,11 @@ function forward_differentiate!(model::Optimizer)
     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: