Objective Sensitivity

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 Lagrandian duality we could already calculate the objective sensitivity with respect to parameters that appear in the RHS of the constraints (e.g, `cons` in this case for parameter `p`).
+
+On the other hand, if the parameter appears in the LHS of the constraints, we 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}
+```
+
+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 the backward 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.
 ```

src/DiffOpt.jl

@@ -14,6 +14,7 @@ import MathOptInterface as MOI
 import MathOptSetDistances as MOSD
 import ParametricOptInterface as POI
 import SparseArrays
+import JuMP: dual
 
 include("utils.jl")
 include("product_of_sets.jl")

src/NonLinearProgram/NonLinearProgram.jl

@@ -27,6 +27,7 @@
 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)
+    dual_p::Float64  # Objective Sensitivity wrt parameters
 end
 
 Base.@kwdef struct ReverseCache
@@ -513,15 +514,19 @@
         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
+        dual_p = df_dp * Δp
+
         model.forw_grad_cache = ForwCache(;
             primal_Δs = Dict(model.cache.primal_vars .=> primal_Δs),
             dual_Δs = dual_Δs,
+            dual_p = dual_p,
         )
     end
     return nothing
@@ -533,43 +538,47 @@
         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 =
@@ -610,4 +619,8 @@
     return MOI.Parameter{T}(model.back_grad_cache.Δp[ci.value])
 end
 
+function MOI.get(model::Model, ::DiffOpt.ForwardObjectiveSensitivity)
+    return model.forw_grad_cache.dual_p
+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
 
@@ -496,5 +503,10 @@ function _compute_sensitivity(model::Model; tol = 1e-6)
     ∂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
+
+    # 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/diff_opt.jl

@@ -16,6 +16,7 @@
     dp::Dict{MOI.ConstraintIndex,Float64} = Dict{MOI.ConstraintIndex,Float64}() # Specifically for NonLinearProgram
     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 @@
     empty!(cache.dx)
     empty!(cache.dp)
     empty!(cache.dy)
+    cache.dobj = 0.0
     empty!(cache.scalar_constraints)
     empty!(cache.vector_constraints)
     cache.objective = nothing
@@ -184,6 +186,20 @@
 """
 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
 
@@ -199,6 +215,21 @@
 
 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
 
@@ -403,6 +434,11 @@
     return
 end
 
+function MOI.set(model::AbstractModel, ::ReverseObjectiveSensitivity, val)
+    model.input_cache.dobj = val
+    return
+end
+
 function MOI.set(
     model::AbstractModel,
     ::ForwardConstraintFunction,

src/jump_moi_overloads.jl

@@ -66,6 +66,10 @@ function MOI.get(
     return JuMP.jump_function(model, moi_func)
 end
 
+function MOI.get(model::JuMP.Model, attr::ForwardObjectiveSensitivity)
+    return MOI.get(JuMP.backend(model), attr)
+end
+
 function MOI.get(model::JuMP.Model, attr::ReverseObjectiveFunction)
     func = MOI.get(JuMP.backend(model), attr)
     return JuMP.jump_function(model, func)
@@ -107,6 +111,34 @@ function MOI.get(
     return _moi_get_result(JuMP.backend(model), attr, JuMP.index(var_ref))
 end
 
+function MOI.set(
+    model::JuMP.Model,
+    attr::ReverseVariablePrimal,
+    var_ref::JuMP.VariableRef,
+    val::Number,
+)
+    JuMP.check_belongs_to_model(var_ref, model)
+    return MOI.set(JuMP.backend(model), attr, JuMP.index(var_ref), val)
+end
+
+function MOI.set(
+    model::JuMP.Model,
+    attr::ReverseConstraintDual,
+    con_ref::JuMP.ConstraintRef,
+    val::Number,
+)
+    JuMP.check_belongs_to_model(con_ref, model)
+    return MOI.set(JuMP.backend(model), attr, JuMP.index(con_ref), val)
+end
+
+function MOI.set(
+    model::JuMP.Model,
+    attr::ReverseObjectiveSensitivity,
+    val::Number,
+)
+    return MOI.set(JuMP.backend(model), attr, val)
+end
+
 function MOI.get(
     model::JuMP.Model,
     attr::ReverseConstraintSet,