Minor tweaks

jump-dev · Mar 24, 2024 · e121140 · e121140
1 parent 96bd014
commit e121140
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 60 deletions.
diff --git a/src/objective.jl b/src/objective.jl
@@ -498,7 +498,8 @@ end
         coefficients::Vector{<:Real},
     )
 
-Set multiple linear objective coefficients associated with `variables` to `coefficients`, in a single call.
+Set multiple linear objective coefficients associated with `variables` to
+`coefficients`, in a single call.
 
 Note: this function will throw an error if a nonlinear objective is set.
 
@@ -527,17 +528,12 @@ function set_objective_coefficient(
 ) where {T}
     if _nlp_objective_function(model) !== nothing
         error("A nonlinear objective is already set in the model")
+    elseif length(variables) != length(coeffs)
+        msg = "The number of variables and coefficients must match"
+        throw(DimensionMismatch(msg))
     end
-    if length(variables) != length(coeffs)
-        throw(
-            DimensionMismatch(
-                "The number of variables and coefficients must match",
-            ),
-        )
-    end
-    coeffs_t = convert.(T, coeffs)::AbstractVector{T}
     F = objective_function_type(model)
-    _set_objective_coefficient(model, variables, coeffs_t, F)
+    _set_objective_coefficient(model, variables, convert.(T, coeffs), F)
     model.is_model_dirty = true
     return
 end
@@ -548,22 +544,15 @@ function _set_objective_coefficient(
     coeffs::AbstractVector{<:T},
     ::Type{GenericVariableRef{T}},
 ) where {T}
-    current_obj = objective_function(model)
-    current_obj_index = index(current_obj)
-    if length(variables) > 0
-        position = findfirst(x -> index(x) == current_obj_index, variables)
-        if position === nothing
-            set_objective_function(
-                model,
-                add_to_expression!(
-                    LinearAlgebra.dot(coeffs, variables),
-                    current_obj,
-                ),
-            )
-        else
-            set_objective_function(model, LinearAlgebra.dot(coeffs, variables))
-        end
+    if length(variables) == 0
+        return
+    end
+    new_objective = coeffs' * variables
+    current_obj = objective_function(model)::GenericVariableRef{T}
+    if !(current_obj in variables)
+        add_to_expression!(new_objective, current_obj)
     end
+    set_objective_function(model, new_objective)
     return
 end
 
@@ -700,13 +689,9 @@ function set_objective_coefficient(
 ) where {T}
     if _nlp_objective_function(model) !== nothing
         error("A nonlinear objective is already set in the model")
-    end
-    if !(length(variables_1) == length(variables_2) == length(coeffs))
-        throw(
-            DimensionMismatch(
-                "The number of variables and coefficients must match",
-            ),
-        )
+    elseif !(length(variables_1) == length(variables_2) == length(coeffs))
+        msg = "The number of variables and coefficients must match"
+        throw(DimensionMismatch(msg))
     end
     coeffs_t = convert.(T, coeffs)::AbstractVector{<:T}
     F = moi_function_type(objective_function_type(model))
@@ -722,14 +707,8 @@ function _set_objective_coefficient(
     coeffs::AbstractVector{<:T},
     ::Type{F},
 ) where {T,F}
-    current_obj = objective_function(model)
-    new_obj = add_to_expression!(
-        sum(
-            coeffs[i] * variables_1[i] * variables_2[i] for
-            i in eachindex(coeffs)
-        ),
-        current_obj,
-    )
+    new_obj = @expression(model, sum(coeffs .* variables_1 .* variables_2))
+    add_to_expression!(new_obj, objective_function(model))
     set_objective_function(model, new_obj)
     return
 end
@@ -741,8 +720,8 @@ function _set_objective_coefficient(
     coeffs::AbstractVector{<:T},
     ::Type{MOI.ScalarQuadraticFunction{T}},
 ) where {T}
-    for i in eachindex(variables_1)
-        if variables_1[i] == variables_2[i]
+    for (i, x, y) in zip(eachindex(coeffs), variables_1, variables_2)
+        if x == y
             coeffs[i] *= T(2)
         end
     end

diff --git a/src/variables.jl b/src/variables.jl
@@ -2573,7 +2573,8 @@ end
         values::AbstractVector{<:Number},
     )
 
-Set multiple coefficient of `variables` in the constraints `constraints` to `values`.
+Set multiple coefficient of `variables` in the constraints `constraints` to
+`values`.
 
 Note that prior to this step, JuMP will aggregate multiple terms containing the
 same variable. For example, given a constraint `2x + 3x <= 2`,
@@ -2607,11 +2608,8 @@ function set_normalized_coefficient(
     coeffs::AbstractVector{<:Number},
 ) where {T,F<:Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}}
     if !(length(constraints) == length(variables) == length(coeffs))
-        throw(
-            DimensionMismatch(
-                "The number of constraints, variables and coefficients must match",
-            ),
-        )
+        msg = "The number of constraints, variables and coefficients must match"
+        throw(DimensionMismatch(msg))
     end
     model = owner_model(first(constraints))
     MOI.modify(
@@ -2734,12 +2732,13 @@ end
         values::AbstractVector{<:Number},
     )
 
-Set multiple quadratic coefficients associated with `variables_1` and `variables_2` in
-the constraints `constraints` to `values`.
+Set multiple quadratic coefficients associated with `variables_1` and
+`variables_2` in the constraints `constraints` to `values`.
 
 Note that prior to this step, JuMP will aggregate multiple terms containing the
 same variable. For example, given a constraint `2x^2 + 3x^2 <= 2`,
-`set_normalized_coefficient(con, [x], [x], [4])` will create the constraint `4x^2 <= 2`.
+`set_normalized_coefficient(con, [x], [x], [4])` will create the constraint
+`4x^2 <= 2`.
 
 ## Example
 
@@ -2765,21 +2764,18 @@ function set_normalized_coefficient(
     variables_2::AbstractVector{<:AbstractVariableRef},
     coeffs::AbstractVector{<:Number},
 ) where {T,F<:MOI.ScalarQuadraticFunction{T}}
-    if !(
+    dimension_match =
         length(constraints) ==
         length(variables_1) ==
         length(variables_2) ==
         length(coeffs)
-    )
-        throw(
-            DimensionMismatch(
-                "The number of constraints, variables and coefficients must match",
-            ),
-        )
+    if !dimension_match
+        msg = "The number of constraints, variables and coefficients must match"
+        throw(DimensionMismatch(msg))
     end
     new_coeffs = convert.(T, coeffs)
-    for i in eachindex(variables_1)
-        if variables_1[i] == variables_2[i]
+    for (i, x, y) in zip(eachindex(new_coeffs), variables_1, variables_2)
+        if x == y
             new_coeffs[i] *= T(2)
         end
     end