use Parameter set

jump-dev · Jan 14, 2025 · 9ec668b · 9ec668b
1 parent 91865cf
commit 9ec668b
Show file tree

Hide file tree

Showing 5 changed files with 129 additions and 70 deletions.
diff --git a/README.md b/README.md
@@ -64,7 +64,7 @@ optimize!(model)
 
 # differentiate w.r.t. p
 direction_p = 3.0
-MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(p), direction_p)
+MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(p), Parameter(direction_p))
 DiffOpt.forward_differentiate!(model)
 @show MOI.get(model, DiffOpt.ForwardVariablePrimal(), x) == direction_p * 3 / pc_val
 
@@ -82,7 +82,7 @@ optimize!(model)
 DiffOpt.empty_input_sensitivities!(model)
 # differentiate w.r.t. pc
 direction_pc = 10.0
-MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(pc), direction_pc)
+MOI.set(model, DiffOpt.ForwardConstraintSet(), ParameterRef(pc), Parameter(direction_pc))
 DiffOpt.forward_differentiate!(model)
 @show abs(MOI.get(model, DiffOpt.ForwardVariablePrimal(), x) -
     -direction_pc * 3 * p_val / pc_val^2) < 1e-5
@@ -93,8 +93,8 @@ DiffOpt.empty_input_sensitivities!(model)
 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)) -
+@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
 ```
 

diff --git a/src/DiffOpt.jl b/src/DiffOpt.jl
@@ -19,8 +19,8 @@ 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("jump_moi_overloads.jl")
 
 include("copy_dual.jl")
 include("bridges.jl")

diff --git a/src/jump_moi_overloads.jl b/src/jump_moi_overloads.jl
@@ -80,6 +80,37 @@ function MOI.get(
     return _moi_get_result(JuMP.backend(model), attr, JuMP.index(var_ref))
 end
 
+# extras to handle model_dirty
+
+function MOI.get(
+    model::JuMP.Model,
+    attr::ReverseConstraintSet,
+    var_ref::JuMP.ConstraintRef,
+)
+    JuMP.check_belongs_to_model(var_ref, model)
+    return _moi_get_result(JuMP.backend(model), attr, JuMP.index(var_ref))
+end
+
+function MOI.set(
+    model::JuMP.Model,
+    attr::ForwardConstraintSet,
+    con_ref::JuMP.ConstraintRef,
+    set::MOI.AbstractScalarSet,
+)
+    JuMP.check_belongs_to_model(con_ref, model)
+    return MOI.set(JuMP.backend(model), attr, JuMP.index(con_ref), set)
+end
+
+function MOI.set(
+    model::JuMP.Model,
+    attr::ForwardConstraintSet,
+    con_ref::JuMP.ConstraintRef,
+    set::JuMP.AbstractScalarSet,
+)
+    JuMP.check_belongs_to_model(con_ref, model)
+    return MOI.set(JuMP.backend(model), attr, JuMP.index(con_ref), JuMP.moi_set(set))
+end
+
 """
     abstract type AbstractLazyScalarFunction <: MOI.AbstractScalarFunction end
 

diff --git a/src/parameters.jl b/src/parameters.jl
@@ -308,14 +308,14 @@ function MOI.set(
     model::POI.Optimizer,
     ::ForwardConstraintSet,
     ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}},
-    value::Number,
+    set::MOI.Parameter,
 ) where {T}
     variable = MOI.VariableIndex(ci.value)
     if _is_variable(model, variable)
         error("Trying to set a forward parameter sensitivity for a variable")
     end
     sensitivity_data = _get_sensitivity_data(model)
-    sensitivity_data.parameter_input_forward[variable] = value
+    sensitivity_data.parameter_input_forward[variable] = set.value
     return
 end
 
@@ -573,16 +573,7 @@ function MOI.get(
         error("Trying to get a backward parameter sensitivity for a variable")
     end
     sensitivity_data = _get_sensitivity_data(model)
-    return get(sensitivity_data.parameter_output_backward, variable, 0.0)
-end
-
-# extras to handle model_dirty
-
-function MOI.get(
-    model::JuMP.Model,
-    attr::ReverseConstraintSet,
-    var_ref::JuMP.ConstraintRef,
-)
-    JuMP.check_belongs_to_model(var_ref, model)
-    return _moi_get_result(JuMP.backend(model), attr, JuMP.index(var_ref))
+    return MOI.Parameter{T}(
+        get(sensitivity_data.parameter_output_backward, variable, 0.0),
+    )
 end