Implement filter_variables for VectorNonlinearFunction (#2556)

jump-dev · Oct 9, 2024 · a08c93c · a08c93c
1 parent 21da4b9
commit a08c93c
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 0 deletions.
diff --git a/src/Utilities/functions.jl b/src/Utilities/functions.jl
@@ -1281,6 +1281,14 @@ function filter_variables(keep::Function, f::MOI.ScalarNonlinearFunction)
     return MOI.ScalarNonlinearFunction(f.head, args)
 end
 
+function filter_variables(keep::Function, f::MOI.VectorNonlinearFunction)
+    return MOI.VectorNonlinearFunction(
+        MOI.ScalarNonlinearFunction[
+            filter_variables(keep, row) for row in f.rows
+        ],
+    )
+end
+
 """
     remove_variable(f::AbstractFunction, vi::VariableIndex)
 

diff --git a/test/Bridges/Constraint/slack.jl b/test/Bridges/Constraint/slack.jl
@@ -299,6 +299,18 @@ function test_vector_slack()
     return
 end
 
+function test_basic_VectorNonlinearFunction()
+    mock = MOI.Utilities.MockOptimizer(
+        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
+    )
+    bridged_mock = MOI.Bridges.Constraint.VectorSlack{Float64}(mock)
+    MOI.Test.test_basic_VectorNonlinearFunction_Nonnegatives(
+        bridged_mock,
+        MOI.Test.Config(),
+    )
+    return
+end
+
 function test_runtests()
     MOI.Bridges.runtests(
         MOI.Bridges.Constraint.ScalarSlackBridge,

diff --git a/test/Utilities/functions.jl b/test/Utilities/functions.jl
@@ -1933,6 +1933,19 @@ function test_filter_variables_scalarnonlinearfunction()
     return
 end
 
+function test_filter_variables_vectornonlinearfunction()
+    x = MOI.VariableIndex(1)
+    y = MOI.VariableIndex(2)
+    f = MOI.ScalarNonlinearFunction(:+, Any[x, 2, y])
+    g = MOI.ScalarNonlinearFunction(:-, Any[x, 2, y])
+    fg = MOI.VectorNonlinearFunction([f, g])
+    new_f = MOI.ScalarNonlinearFunction(:+, Any[2, y])
+    new_g = MOI.ScalarNonlinearFunction(:-, Any[0, 2, y])
+    new_fg = MOI.VectorNonlinearFunction([new_f, new_g])
+    @test new_fg ≈ MOI.Utilities.filter_variables(xi -> xi != x, fg)
+    return
+end
+
 function test_ScalarNonlinearFunction_count_map_indices_and_print()
     model = MOI.Utilities.CachingOptimizer(
         MOI.Utilities.Model{Bool}(),