linting

jalving · Jul 17, 2024 · cf96915 · cf96915
1 parent 8c75a91
commit cf96915
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 74 deletions.
diff --git a/src/aggregate.jl b/src/aggregate.jl
@@ -251,7 +251,7 @@ function _copy_constraints!(
             ref_map[src_cref] = new_cref
         end
         # pass constraint attributes
-        
+
         cis_src = MOI.get(source_element, MOI.ListOfConstraintIndices{F_moi,S}())
         MOIU.pass_attributes(dest, src, index_map_FS, cis_src)
     end

diff --git a/src/node_variables.jl b/src/node_variables.jl
@@ -511,65 +511,31 @@ function JuMP.unset_binary(nvref::NodeVariableRef)
     return nothing
 end
 
-### Utilities for querying variables used in constraints
-
-function _extract_variables(func::NodeVariableRef)
-    return [func]
-end
-
-function _extract_variables(ref::ConstraintRef)
-    func = JuMP.jump_function(JuMP.constraint_object(ref))
-    return _extract_variables(func)
-end
-
-function _extract_variables(func::JuMP.GenericAffExpr)
-    return collect(keys(func.terms))
-end
-
-function _extract_variables(func::JuMP.GenericQuadExpr)
-    quad_vars = vcat([[term[2]; term[3]] for term in JuMP.quad_terms(func)]...)
-    aff_vars = _extract_variables(func.aff)
-    return union(quad_vars, aff_vars)
-end
-
-function _extract_variables(func::JuMP.GenericNonlinearExpr)
-    vars = NodeVariableRef[]
-    for i in 1:length(func.args)
-        func_arg = func.args[i]
-        if func_arg isa Number
-            continue
-        elseif typeof(func_arg) == NodeVariableRef
-            push!(vars, func_arg)
-        else
-            append!(vars, _extract_variables(func_arg))
-        end
-    end
-    return vars
-end
-
 # Extended from https://github.com/jump-dev/JuMP.jl/blob/301d46e81cb66c74c6e22cd89fb89ced740f157b/src/variables.jl#L2721
 function JuMP.set_normalized_coefficient(
-    con_ref::S,
-    variable::NodeVariableRef,
-    value::Number,
-) where {S<:Union{NodeConstraintRef, EdgeConstraintRef}}
+    con_ref::S, variable::NodeVariableRef, value::Number
+) where {S<:Union{NodeConstraintRef,EdgeConstraintRef}}
     graph = owner_model(con_ref).source_graph.x
     _backend = backend(graph)
 
     MOI.modify(
         _backend.moi_backend,
         graph_index(con_ref),
-        MOI.ScalarCoefficientChange(graph_index(_backend, variable), convert(Float64,value)),
+        MOI.ScalarCoefficientChange(
+            graph_index(_backend, variable), convert(Float64, value)
+        ),
     )
     graph.is_model_dirty = true
-    return
+    return nothing
 end
 
 function JuMP.set_normalized_coefficient(
     constraints::AbstractVector{<:S},
     variables::AbstractVector{<:NodeVariableRef},
     coeffs::AbstractVector{<:Number},
-) where {S<:Union{NodeConstraintRef, EdgeConstraintRef, Union{NodeConstraintRef, EdgeConstraintRef}}}
+) where {
+    S<:Union{NodeConstraintRef,EdgeConstraintRef,Union{NodeConstraintRef,EdgeConstraintRef}}
+}
     c, n, m = length(constraints), length(variables), length(coeffs)
     if !(c == n == m)
         msg = "The number of constraints ($c), variables ($n) and coefficients ($m) must match"
@@ -578,25 +544,27 @@ function JuMP.set_normalized_coefficient(
     graph = [owner_model(con).source_graph.x for con in constraints]
     _backends = unique(backend.(graph))
     if length(_backends) != 1
-        error("Constraints belong to different graph backends; make sure constraints all come from the same graph")
+        error(
+            "Constraints belong to different graph backends; make sure constraints all come from the same graph",
+        )
     end
     _backend = _backends[1]
     graph = graph[1]
 
     MOI.modify(
         _backend.moi_backend,
         graph_index.(Ref(_backend), constraints),
-        MOI.ScalarCoefficientChange.(graph_index.(Ref(_backend), variables), convert.(Float64, coeffs)),
+        MOI.ScalarCoefficientChange.(
+            graph_index.(Ref(_backend), variables), convert.(Float64, coeffs)
+        ),
     )
     graph.is_model_dirty = true
-    return
+    return nothing
 end
 
 function JuMP.set_normalized_coefficient(
-    constraint::S,
-    variable::NodeVariableRef,
-    new_coefficients::Vector{Tuple{Int64,T}},
-) where {T,S<:Union{NodeConstraintRef, EdgeConstraintRef}}
+    constraint::S, variable::NodeVariableRef, new_coefficients::Vector{Tuple{Int64,T}}
+) where {T,S<:Union{NodeConstraintRef,EdgeConstraintRef}}
     graph = owner_model(con_ref).source_graph.x
     _backend = backend(graph)
 
@@ -606,23 +574,18 @@ function JuMP.set_normalized_coefficient(
         MOI.MultirowChange(graph_index(_backend, variable), new_coefficients),
     )
     graph.is_model_dirty = true
-    return
+    return nothing
 end
 
 function JuMP.set_normalized_coefficients(
-    constraint::S,
-    variable::NodeVariableRef,
-    new_coefficients::Vector{Tuple{Int64,T}},
-) where {T,S<:Union{NodeConstraintRef, EdgeConstraintRef}}
+    constraint::S, variable::NodeVariableRef, new_coefficients::Vector{Tuple{Int64,T}}
+) where {T,S<:Union{NodeConstraintRef,EdgeConstraintRef}}
     return JuMP.set_normalized_coefficient(constraint, variable, new_coefficients)
 end
 
 function JuMP.set_normalized_coefficient(
-    constraint::S,
-    variable_1::NodeVariableRef,
-    variable_2::NodeVariableRef,
-    value::Number,
-) where {S<:Union{NodeConstraintRef, EdgeConstraintRef}}
+    constraint::S, variable_1::NodeVariableRef, variable_2::NodeVariableRef, value::Number
+) where {S<:Union{NodeConstraintRef,EdgeConstraintRef}}
     new_value = convert(Float64, value)
     if variable_1 == variable_2
         new_value *= Float64(2)
@@ -633,21 +596,19 @@ function JuMP.set_normalized_coefficient(
         _backend.moi_backend,
         graph_index(_backend, constraint),
         MOI.ScalarQuadraticCoefficientChange(
-            graph_index(_backend, variable_1),
-            graph_index(_backend, variable_2),
-            new_value,
+            graph_index(_backend, variable_1), graph_index(_backend, variable_2), new_value
         ),
     )
     graph.is_model_dirty = true
-    return
+    return nothing
 end
 
 function JuMP.set_normalized_coefficient(
     constraints::AbstractVector{S},
     variables_1::AbstractVector{<:NodeVariableRef},
     variables_2::AbstractVector{<:NodeVariableRef},
     coeffs::AbstractVector{<:Number},
-) where {S<:Union{NodeConstraintRef, EdgeConstraintRef}}
+) where {S<:Union{NodeConstraintRef,EdgeConstraintRef}}
     c, m = length(constraints), length(coeffs)
     n1, n2 = length(variables_1), length(variables_1)
     if !(c == n1 == n2 == m)
@@ -663,7 +624,9 @@ function JuMP.set_normalized_coefficient(
     graph = [owner_model(con).source_graph.x for con in constraints]
     _backends = unique(backend.(graph))
     if length(_backends) != 1
-        error("Constraints belong to different graph backends; make sure constraints all come from the same graph")
+        error(
+            "Constraints belong to different graph backends; make sure constraints all come from the same graph",
+        )
     end
     _backend = _backends[1]
     graph = graph[1]
@@ -677,5 +640,41 @@ function JuMP.set_normalized_coefficient(
         ),
     )
     graph.is_model_dirty = true
-    return
+    return nothing
+end
+
+### Utilities for querying variables used in constraints
+
+function _extract_variables(func::NodeVariableRef)
+    return [func]
+end
+
+function _extract_variables(ref::ConstraintRef)
+    func = JuMP.jump_function(JuMP.constraint_object(ref))
+    return _extract_variables(func)
+end
+
+function _extract_variables(func::JuMP.GenericAffExpr)
+    return collect(keys(func.terms))
+end
+
+function _extract_variables(func::JuMP.GenericQuadExpr)
+    quad_vars = vcat([[term[2]; term[3]] for term in JuMP.quad_terms(func)]...)
+    aff_vars = _extract_variables(func.aff)
+    return union(quad_vars, aff_vars)
+end
+
+function _extract_variables(func::JuMP.GenericNonlinearExpr)
+    vars = NodeVariableRef[]
+    for i in 1:length(func.args)
+        func_arg = func.args[i]
+        if func_arg isa Number
+            continue
+        elseif typeof(func_arg) == NodeVariableRef
+            push!(vars, func_arg)
+        else
+            append!(vars, _extract_variables(func_arg))
+        end
+    end
+    return vars
 end
diff --git a/test/test_optigraph.jl b/test/test_optigraph.jl
@@ -358,11 +358,7 @@ function test_variable_constraints()
 
     set_normalized_coefficient(n1[:con1], n1[:x], 2)
     @test normalized_coefficient(n1[:con1], n1[:x]) == 2
-    set_normalized_coefficient(
-        [n1[:con1], n1[:con2]],
-        [n1[:x], n1[:x]],
-        [3.0, 2.0]
-    )
+    set_normalized_coefficient([n1[:con1], n1[:con2]], [n1[:x], n1[:x]], [3.0, 2.0])
     @test normalized_coefficient(n1[:con1], n1[:x]) == 3
     @test normalized_coefficient(n1[:con2], n1[:x]) == 2
     set_normalized_coefficient(graph[:link_con2], n1[:y], n2[:x], 2.0)
@@ -371,7 +367,7 @@ function test_variable_constraints()
         [graph[:link_con2], graph[:link_con3]],
         [n1[:y], n1[:y]],
         [n2[:x], n2[:x]],
-        [3.0, 3.0]
+        [3.0, 3.0],
     )
     @test normalized_coefficient(graph[:link_con2], n1[:y], n2[:x]) == 3
     @test normalized_coefficient(graph[:link_con3], n1[:y], n2[:x]) == 3