Objective Sensitivity 2.0 (#303)

* implement dual of parameter anywhere

* fix api and more tests

* format

* update docs

* update for appropriate API

* update documemtation

* format

* rm legacy function

* rm unecessary functions

* read necessary functions

* Update DiffOpt.jl

* fix bug indexing

* format

* initial sketch

* rm method

* add JuMP API and test it

* add diff_model test

* Fix conic error (#284)

* fix conic error

* use jump psd problem

* [docs] update to Documenter@1 (#286)

* [docs] update to Documenter@1

* Update

* update

* format

* Pass attributes through Objective.FunctionConversionBridge (#287)

* Pass attributes through Objective.FunctionConversionBridge

* Fix

* add test

* fix tol

* fix test

* add reverse test

---------

Co-authored-by: joaquimg <[email protected]>

* bump POI

* cleanup

---------

Co-authored-by: Oscar Dowson <[email protected]>
Co-authored-by: Benoît Legat <[email protected]>
Co-authored-by: joaquimg <[email protected]>

* fix parameters support

* Update jump-api pr (#294)

* prepare more tests

* sketch

* temp examples

* remove printing

* remove printing

* remove duplicatie def

* move definition

* format

* format

* simplify conic model

* force nlp on nlp tests

* fix param tests

* cleanup parameter usage

* remove code

* cleanup usage of parameters

* format

* add temp dep

* format

* add temp dep

* fix PSDSquare

* temp fix for tests

* format

* Improve name handling, cleanup parameter error messages, remove slack

* re-enable constraint names

* integrate new POI

* remove test

* rm adhoc pkg add

* fix docs

* add POI test and docs

* rm test_solve_conflict

* format

* add parameter in cone test

* add back name constraint vect

* Update Project.toml

* implement dual of parameter anywhere

* update for appropriate API

* format

* fix

* fix docs

* add error for non implemented backends

* fix error handeling

* format

* fix tests

* fix tests

* format

* improve coverage

* format

* bump version

* update docs

* udpate docs

* fix error dual objective sensitivity mention

* format

* fix typo

* fix typo

* Update docs/src/usage.md

Co-authored-by: Benoît Legat <[email protected]>

* Apply suggestion from @andrewrosemberg

* pin MLDatasets

* Apply suggestion from @joaquimg

Co-authored-by: Joaquim <[email protected]>

* Apply suggestion from @joaquimg

Co-authored-by: Joaquim <[email protected]>

* Apply suggestion from @joaquimg

Co-authored-by: Joaquim <[email protected]>

* Apply suggestion from @andrewrosemberg

* Apply suggestion from @andrewrosemberg

* Apply suggestion from @andrewrosemberg

* Apply suggestion from @andrewrosemberg

* rename objective_sensitivity_p

* fix tests

* update tests

* format

* rm dual and addition overloads

---------

Co-authored-by: joaquimg <[email protected]>
Co-authored-by: Oscar Dowson <[email protected]>
Co-authored-by: Benoît Legat <[email protected]>

Author: 4 people

Project.toml

@@ -1,7 +1,7 @@
 name = "DiffOpt"
 uuid = "930fe3bc-9c6b-11ea-2d94-6184641e85e7"
 authors = ["Akshay Sharma", "Mathieu Besançon", "Joaquim Dias Garcia", "Benoît Legat", "Oscar Dowson", "Andrew Rosemberg"]
-version = "0.5.1"
+version = "0.5.2"
 
 [deps]
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"

docs/src/usage.md

@@ -117,4 +117,70 @@ DiffOpt.reverse_differentiate!(model)
 @show MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p)) == MOI.Parameter(direction_x * 3 / pc_val)
 @show abs(MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(pc)).value -
     -direction_x * 3 * p_val / pc_val^2) < 1e-5
-```
+```
+
+## Calculating objective sensitivity with respect to parameters (currently only supported for Nonlinear Programs)
+
+Consider a differentiable model with parameters `p` and `pc` as in the previous example:
+
+```julia
+using JuMP, DiffOpt, HiGHS
+
+model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
+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)
+@objective(model, Min, x^4)
+optimize!(model)
+
+direction_p = 3.0
+MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(p), Parameter(direction_p))
+DiffOpt.forward_differentiate!(model)
+```
+
+Using Lagrangian duality we can easily calculate the objective sensitivity with respect to parameters that appear as constants of the constraints (e.g, `cons` in this case for parameter `p`) - i.e. The objective sensitivity w.r.t. a constant parameter change is given by the optimal dual multiplier, under strong duality.
+
+On the other hand, if the parameter appears as a coefficient of the constraints, one can calculate the objective sensitivity with respect to the parameter using the sensitivities of the variables with respect to the parameter, \( \frac{\partial x}{\partial p} \), and the gradient of the objective with respect to the variables \( \frac{\partial f}{\partial x} \):
+
+```math
+\frac{\partial f}{\partial p} = \frac{\partial f}{\partial x} \frac{\partial x}{\partial p}
+```
+ - A consequence of the chain-rule.
+
+
+In order to calculate the objective perturbation with respect to the parameter perturbation vector, we can use the following code:
+
+```julia
+# Always a good practice to clear previously set sensitivities
+DiffOpt.empty_input_sensitivities!(model)
+
+MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(p), Parameter(3.0))
+MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(p_c), Parameter(3.0))
+DiffOpt.forward_differentiate!(model)
+
+MOI.get(model, DiffOpt.ForwardObjectiveSensitivity())
+```
+
+In reverse mode, we can calculate the parameter perturbation with respect to the objective perturbation:
+
+```julia
+# Always a good practice to clear previously set sensitivities
+DiffOpt.empty_input_sensitivities!(model)
+
+MOI.set(
+    model,
+    DiffOpt.ReverseObjectiveSensitivity(),
+    0.1,
+)
+
+DiffOpt.reverse_differentiate!(model)
+
+MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p))
+```
+
+It is important to note that the (reverse) parameter perturbation given an objective perturbation is somewhat equivalent to the perturbation with respect to solution (since one can be calculated from the other). Therefore, one cannot set both the objective sensitivity (`DiffOpt.ReverseObjectiveSensitivity`) and the solution sensitivity (e.g. `DiffOpt.ReverseVariablePrimal`) at the same time - the code will throw an error if you try to do so.

src/ConicProgram/ConicProgram.jl

@@ -450,4 +450,16 @@ function MOI.get(
     return MOI.get(model.model, attr, ci)
 end
 
+function MOI.get(::Model, ::DiffOpt.ForwardObjectiveSensitivity)
+    return error(
+        "ForwardObjectiveSensitivity is not implemented for the Conic Optimization backend",
+    )
+end
+
+function MOI.set(::Model, ::DiffOpt.ReverseObjectiveSensitivity, val)
+    return error(
+        "ReverseObjectiveSensitivity is not implemented for the Conic Optimization backend",
+    )
+end
+
 end

src/NonLinearProgram/NonLinearProgram.jl

@@ -27,10 +27,11 @@ end
 Base.@kwdef struct ForwCache
     primal_Δs::Dict{MOI.VariableIndex,Float64}  # Sensitivity for primal variables (indexed by VariableIndex)
     dual_Δs::Vector{Float64}           # Sensitivity for constraints and bounds (indexed by ConstraintIndex)
+    objective_sensitivity_p::Float64  # Objective Sensitivity wrt parameters
 end
 
 Base.@kwdef struct ReverseCache
-    Δp::Vector{Float64}  # Sensitivity for parameters
+    Δp::Dict{MOI.ConstraintIndex,Float64}  # Sensitivity for parameters
 end
 
 # Define the form of the NLP
@@ -525,15 +526,19 @@ function DiffOpt.forward_differentiate!(model::Model; tol = 1e-6)
         end
 
         # Compute Jacobian
-        Δs = _compute_sensitivity(model; tol = tol)
+        Δs, df_dp = _compute_sensitivity(model; tol = tol)
 
         # Extract primal and dual sensitivities
         primal_Δs = Δs[1:length(model.cache.primal_vars), :] * Δp # Exclude slacks
         dual_Δs = Δs[cache.index_duals, :] * Δp # Includes constraints and bounds
 
+        # obj sensitivity wrt parameters
+        objective_sensitivity_p = df_dp * Δp
+
         model.forw_grad_cache = ForwCache(;
             primal_Δs = Dict(model.cache.primal_vars .=> primal_Δs),
             dual_Δs = dual_Δs,
+            objective_sensitivity_p = objective_sensitivity_p,
         )
     end
     return nothing
@@ -545,50 +550,53 @@ function DiffOpt.reverse_differentiate!(model::Model; tol = 1e-6)
         form = model.model
 
         # Compute Jacobian
-        Δs = _compute_sensitivity(model; tol = tol)
-        num_primal = length(cache.primal_vars)
-        # Fetch primal sensitivities
-        Δx = zeros(num_primal)
-        for (i, var_idx) in enumerate(cache.primal_vars)
-            if haskey(model.input_cache.dx, var_idx)
-                Δx[i] = model.input_cache.dx[var_idx]
+        Δs, df_dp = _compute_sensitivity(model; tol = tol)
+        Δp = if !iszero(model.input_cache.dobj)
+            model.input_cache.dobj * df_dp
+        else
+            num_primal = length(cache.primal_vars)
+            # Fetch primal sensitivities
+            Δx = zeros(num_primal)
+            for (i, var_idx) in enumerate(cache.primal_vars)
+                if haskey(model.input_cache.dx, var_idx)
+                    Δx[i] = model.input_cache.dx[var_idx]
+                end
             end
-        end
-        # Fetch dual sensitivities
-        num_constraints = length(cache.cons)
-        num_up = length(cache.has_up)
-        num_low = length(cache.has_low)
-        Δdual = zeros(num_constraints + num_up + num_low)
-        for (i, ci) in enumerate(cache.cons)
-            idx = form.nlp_index_2_constraint[ci]
-            if haskey(model.input_cache.dy, idx)
-                Δdual[i] = model.input_cache.dy[idx]
+            # Fetch dual sensitivities
+            num_constraints = length(cache.cons)
+            num_up = length(cache.has_up)
+            num_low = length(cache.has_low)
+            Δdual = zeros(num_constraints + num_up + num_low)
+            for (i, ci) in enumerate(cache.cons)
+                idx = form.nlp_index_2_constraint[ci]
+                if haskey(model.input_cache.dy, idx)
+                    Δdual[i] = model.input_cache.dy[idx]
+                end
             end
-        end
-        for (i, var_idx) in enumerate(cache.primal_vars[cache.has_low])
-            idx = form.constraint_lower_bounds[var_idx.value].value
-            if haskey(model.input_cache.dy, idx)
-                Δdual[num_constraints+i] = model.input_cache.dy[idx]
+            for (i, var_idx) in enumerate(cache.primal_vars[cache.has_low])
+                idx = form.constraint_lower_bounds[var_idx.value].value
+                if haskey(model.input_cache.dy, idx)
+                    Δdual[num_constraints+i] = model.input_cache.dy[idx]
+                end
             end
-        end
-        for (i, var_idx) in enumerate(cache.primal_vars[cache.has_up])
-            idx = form.constraint_upper_bounds[var_idx.value].value
-            if haskey(model.input_cache.dy, idx)
-                Δdual[num_constraints+num_low+i] = model.input_cache.dy[idx]
+            for (i, var_idx) in enumerate(cache.primal_vars[cache.has_up])
+                idx = form.constraint_upper_bounds[var_idx.value].value
+                if haskey(model.input_cache.dy, idx)
+                    Δdual[num_constraints+num_low+i] = model.input_cache.dy[idx]
+                end
             end
+            # Extract Parameter sensitivities
+            Δw = zeros(size(Δs, 1))
+            Δw[1:num_primal] = Δx
+            Δw[cache.index_duals] = Δdual
+            Δp = Δs' * Δw
         end
-        # Extract Parameter sensitivities
-        Δw = zeros(size(Δs, 1))
-        Δw[1:num_primal] = Δx
-        Δw[cache.index_duals] = Δdual
-        Δp = Δs' * Δw
-
-        # Order by ConstraintIndex
-        varorder =
-            sort(collect(keys(form.var2ci)); by = x -> form.var2ci[x].value)
-        Δp = [Δp[form.var2param[var_idx].value] for var_idx in varorder]
-
-        model.back_grad_cache = ReverseCache(; Δp = Δp)
+
+        Δp_dict = Dict{MOI.ConstraintIndex,Float64}(
+            form.var2ci[var_idx] => Δp[form.var2param[var_idx].value]
+            for var_idx in keys(form.var2ci)
+        )
+        model.back_grad_cache = ReverseCache(; Δp = Δp_dict)
     end
     return nothing
 end
@@ -620,10 +628,16 @@ function MOI.get(
     ::DiffOpt.ReverseConstraintSet,
     ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}},
 ) where {T}
-    form = model.model
-    var_idx = MOI.VariableIndex(ci.value)
-    p_idx = form.var2param[var_idx].value
-    return MOI.Parameter{T}(model.back_grad_cache.Δp[p_idx])
+    return MOI.Parameter{T}(model.back_grad_cache.Δp[ci])
+end
+
+function MOI.get(model::Model, ::DiffOpt.ForwardObjectiveSensitivity)
+    return model.forw_grad_cache.objective_sensitivity_p
+end
+
+function MOI.set(model::Model, ::DiffOpt.ReverseObjectiveSensitivity, val)
+    model.input_cache.dobj = val
+    return
 end
 
 end # module NonLinearProgram

src/NonLinearProgram/nlp_utilities.jl

@@ -27,6 +27,13 @@ function _fill_off_diagonal(H::SparseMatrixCSC)
     return ret
 end
 
+function _compute_gradient(model::Model)
+    evaluator = model.cache.evaluator
+    grad = zeros(length(model.x))
+    MOI.eval_objective_gradient(evaluator, grad, model.x)
+    return grad
+end
+
 """
     _compute_optimal_hessian(evaluator::MOI.Nonlinear.Evaluator, rows::Vector{JuMP.ConstraintRef}, x::Vector{JuMP.VariableRef})
 
@@ -104,7 +111,7 @@ function _create_evaluator(form::Form)
         backend,
         MOI.VariableIndex.(1:form.num_variables),
     )
-    MOI.initialize(evaluator, [:Hess, :Jac])
+    MOI.initialize(evaluator, [:Hess, :Jac, :Grad])
     return evaluator
 end
 
@@ -480,6 +487,11 @@ function _compute_sensitivity(model::Model; tol = 1e-6)
     # Dual bounds lower
     ∂s[(num_w+num_cons+1):(num_w+num_cons+num_lower), :] *= _sense_multiplier
     # Dual bounds upper
-    ∂s[(num_w+num_cons+num_lower+1):end, :] *= -_sense_multiplier
-    return ∂s
+    ∂s[((num_w+num_cons+num_lower+1):end), :] *= -_sense_multiplier
+
+    # dual wrt parameter
+    primal_idx = [i.value for i in model.cache.primal_vars]
+    df_dx = _compute_gradient(model)[primal_idx]
+    df_dp = df_dx'∂s[1:num_vars, :]
+    return ∂s, df_dp
 end

src/QuadraticProgram/QuadraticProgram.jl

@@ -501,4 +501,16 @@ function MOI.set(model::Model, ::LinearAlgebraSolver, linear_solver)
     return model.linear_solver = linear_solver
 end
 
+function MOI.get(::Model, ::DiffOpt.ForwardObjectiveSensitivity)
+    return error(
+        "ForwardObjectiveSensitivity is not implemented for the Quadratic Optimization backend",
+    )
+end
+
+function MOI.set(::Model, ::DiffOpt.ReverseObjectiveSensitivity, val)
+    return error(
+        "ReverseObjectiveSensitivity is not implemented for the Quadratic Optimization backend",
+    )
+end
+
 end

src/diff_opt.jl

@@ -15,6 +15,7 @@ Base.@kwdef mutable struct InputCache
     dx::Dict{MOI.VariableIndex,Float64} = Dict{MOI.VariableIndex,Float64}()# dz for QP
     dy::Dict{MOI.ConstraintIndex,Float64} = Dict{MOI.ConstraintIndex,Float64}()
     # Dual sensitivity currently only works for NonLinearProgram
+    dobj::Float64 = 0.0 # Objective input sensitivity for reverse differentiation
     # ds
     # dy #= [d\lambda, d\nu] for QP
     # FIXME Would it be possible to have a DoubleDict where the value depends
@@ -35,6 +36,7 @@ end
 function Base.empty!(cache::InputCache)
     empty!(cache.dx)
     empty!(cache.dy)
+    cache.dobj = 0.0
     empty!(cache.parameter_constraints)
     empty!(cache.scalar_constraints)
     empty!(cache.vector_constraints)
@@ -191,6 +193,20 @@ MOI.set(model, DiffOpt.ReverseConstraintDual(), ci, value)
 """
 struct ReverseConstraintDual <: MOI.AbstractConstraintAttribute end
 
+"""
+    ReverseObjectiveSensitivity <: MOI.AbstractModelAttribute
+
+A `MOI.AbstractModelAttribute` to set input data for reverse differentiation.
+
+For instance, to set the sensitivity of the parameter perturbation with respect to the
+objective function perturbation, do the following:
+
+```julia
+MOI.set(model, DiffOpt.ReverseObjectiveSensitivity(), value)
+```
+"""
+struct ReverseObjectiveSensitivity <: MOI.AbstractModelAttribute end
+
 """
     ForwardConstraintDual <: MOI.AbstractConstraintAttribute
 
@@ -206,6 +222,21 @@ struct ForwardConstraintDual <: MOI.AbstractConstraintAttribute end
 
 MOI.is_set_by_optimize(::ForwardConstraintDual) = true
 
+"""
+    ForwardObjectiveSensitivity <: MOI.AbstractModelAttribute
+
+A `MOI.AbstractModelAttribute` to get output objective sensitivity data from forward differentiation.
+
+For instance, to get the sensitivity of the objective function with respect to the parameter perturbation, do the following:
+
+```julia
+MOI.get(model, DiffOpt.ForwardObjectiveSensitivity())
+```
+"""
+struct ForwardObjectiveSensitivity <: MOI.AbstractModelAttribute end
+
+MOI.is_set_by_optimize(::ForwardObjectiveSensitivity) = true
+
 """
     ReverseObjectiveFunction <: MOI.AbstractModelAttribute