public get_y, get_y_dual, y_shape

Author: klamike

src/L2ODLL.jl

@@ -132,18 +132,56 @@ function make_vector_data(cache::DLLCache; M=SparseArrays.SparseMatrixCSC{Float6
     completion_data = convert(VectorStandardFormData{M,V,T}, model_to_data(completion_model))
     return completion_data, y_sets, (p_ref, y_ref, ref_map)
 end
-function get_y(cache::DLLCache)
-    return get_y(cache.dual_model, cache.decomposition)
+
+"""
+    get_y(model::JuMP.Model)
+
+Get the primal constraints corresponding to the `y` variables in the decomposition.
+"""
+function get_y(model::JuMP.Model)
+    return get_cache(model).decomposition.y_ref
+end
+
+"""
+    get_y_dual(model::JuMP.Model)
+
+Get the dual variables corresponding to the `y` variables in the decomposition.
+    These are VariableRefs belonging to the dual model, not the passed-in `model`.
+"""
+function get_y_dual(model::JuMP.Model)
+    return get_y_dual(get_cache(model))
+end
+function get_y_dual(cache::DLLCache)
+    return get_y_dual(cache.dual_model, cache.decomposition)
 end
 
+"""
+    y_shape(model::JuMP.Model)
+
+Get the shape of the `y` variables in the decomposition.
+    This is a Vector{Int} where each entry is the number of dual variables for that constraint.
+"""
+function y_shape(model::JuMP.Model)
+    return y_shape(get_cache(model))
+end
 function y_shape(cache::DLLCache)
-    return length.(get_y(cache.dual_model, cache.decomposition))
+    return length.(get_y_dual(cache.dual_model, cache.decomposition))
 end
 
+"""
+    flatten_y(y::AbstractVector)
+
+Flatten a vector of `y` variables into a single vector, i.e. Vector{Vector{Float64}} -> Vector{Float64}.
+"""
 function flatten_y(y::AbstractVector)
     return reduce(vcat, y)
 end
 
+"""
+    unflatten_y(y::AbstractVector, y_shape::AbstractVector{Int})
+
+Unflatten a vector of flattened `y` variables into a vector of vectors, i.e. Vector{Float64} -> Vector{Vector{Float64}}.
+"""
 function unflatten_y(y::AbstractVector, y_shape::AbstractVector{Int})
     return [y[start_idx:start_idx + shape - 1] for (start_idx, shape) in enumerate(y_shape)]
 end

src/layers/bounded.jl

@@ -28,7 +28,7 @@ end
 
 function bounded_builder(decomposition::BoundDecomposition, proj_fn, dual_model::JuMP.Model; completion=:exact, μ=1.0)
     p_vars = get_p(dual_model, decomposition)
-    y_vars = get_y(dual_model, decomposition)
+    y_vars = get_y_dual(dual_model, decomposition)
     zl_vars = only.(get_zl(dual_model, decomposition))
     zu_vars = only.(get_zu(dual_model, decomposition))
     types = filter(

src/layers/convex_qp.jl

@@ -33,7 +33,7 @@ end
 
 function convex_qp_builder(decomposition::ConvexQP, proj_fn, dual_model::JuMP.Model)
     p_vars = get_p(dual_model, decomposition)
-    y_vars = get_y(dual_model, decomposition)
+    y_vars = get_y_dual(dual_model, decomposition)
     x_vars = get_x(decomposition)
     if !all(x -> has_quadslack(dual_model, x), x_vars)
         @warn "Some primal variables do not have a quadratic objective term, " *

src/layers/generic.jl

@@ -78,7 +78,7 @@ function _make_completion_model(decomposition::AbstractDecomposition, dual_model
 
     # mark y and p as parameters (optimizing over z only)
     p_ref = getindex.(ref_map, get_p(dual_model, decomposition))
-    y_ref = getindex.(ref_map, get_y(dual_model, decomposition))
+    y_ref = getindex.(ref_map, get_y_dual(dual_model, decomposition))
     y_ref_flat = reduce(vcat, y_ref)
     JuMP.@constraint(completion_model, y_ref_flat .∈ MOI.Parameter.(zeros(length(y_ref_flat))))
     JuMP.@constraint(completion_model, p_ref .∈ MOI.Parameter.(zeros(length(p_ref))))
@@ -93,7 +93,7 @@ function make_vector_data(decomposition::AbstractDecomposition, dual_model::JuMP
     return data, y_sets, (p_ref, y_ref, ref_map)
 end
 
-function get_y(dual_model, decomposition::AbstractDecomposition)
+function get_y_dual(dual_model, decomposition::AbstractDecomposition)
     return Dualization._get_dual_variables.(dual_model, decomposition.y_ref)
 end
 

test/runtests.jl

@@ -48,7 +48,7 @@ SOLVER = () -> ParametricOptInterface.Optimizer(HiGHS.Optimizer());
 
         L2ODLL.decompose!(m);
 
-        cqp_y_pred = randn_like(L2ODLL.get_y(L2ODLL.get_cache(m)));
+        cqp_y_pred = randn_like(L2ODLL.get_y_dual(m));
 
         dobj = L2ODLL.dual_objective(m, cqp_y_pred, param_value)
         dobj_wrt_y = L2ODLL.dual_objective_gradient(m, cqp_y_pred, param_value)
@@ -65,7 +65,7 @@ SOLVER = () -> ParametricOptInterface.Optimizer(HiGHS.Optimizer());
         JuMP.set_optimizer(m, Clarabel.Optimizer);
         JuMP.set_silent(m);
         JuMP.optimize!(m)
-        cqp_y_true = JuMP.dual.(L2ODLL.get_cache(m).decomposition.y_ref)
+        cqp_y_true = JuMP.dual.(L2ODLL.get_y(m))
         dobj1 = L2ODLL.dual_objective(m, cqp_y_true, param_value)
         @test isapprox(dobj1, JuMP.objective_value(m), atol=1e-6)
     end
@@ -91,7 +91,7 @@ SOLVER = () -> ParametricOptInterface.Optimizer(HiGHS.Optimizer());
 
         L2ODLL.decompose!(m);
 
-        blp_y_pred = randn_like(L2ODLL.get_y(L2ODLL.get_cache(m)));
+        blp_y_pred = randn_like(L2ODLL.get_y_dual(m));
 
         dobj = L2ODLL.dual_objective(m, blp_y_pred, param_value)
         dobj_wrt_y = L2ODLL.dual_objective_gradient(m, blp_y_pred, param_value)
@@ -113,7 +113,7 @@ SOLVER = () -> ParametricOptInterface.Optimizer(HiGHS.Optimizer());
         JuMP.set_optimizer(m, Clarabel.Optimizer);
         JuMP.set_silent(m);
         JuMP.optimize!(m)
-        blp_y_true = JuMP.dual.(L2ODLL.get_cache(m).decomposition.y_ref)
+        blp_y_true = JuMP.dual.(L2ODLL.get_y(m))
 
         dobj1 = L2ODLL.dual_objective(m, blp_y_true, param_value)
         @test isapprox(dobj1, JuMP.objective_value(m), atol=1e-6)