[WIP] Scalar affine function reloaded

src/Utilities/Utilities.jl

@@ -14,6 +14,7 @@ const MOIU = MOI.Utilities # used in macro
 const SVF = MOI.SingleVariable
 const VVF = MOI.VectorOfVariables
 const SAF{T} = MOI.ScalarAffineFunction{T}
+const GAF{T} = MOI.GenericScalarAffineFunction{T}
 const VAF{T} = MOI.VectorAffineFunction{T}
 const SQF{T} = MOI.ScalarQuadraticFunction{T}
 const VQF{T} = MOI.VectorQuadraticFunction{T}

src/Utilities/constraints.jl

@@ -15,7 +15,7 @@ function normalize_and_add_constraint(
     func::MOI.AbstractScalarFunction,
     set::MOI.AbstractScalarSet;
     allow_modify_function::Bool = false,
-) where {T}
+)
     return MOI.add_constraint(
         model,
         normalize_constant(
@@ -45,11 +45,11 @@ function normalize_constant(
     return func, set
 end
 function normalize_constant(
-    func::Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}},
+    func::Union{MOI.GenericScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}},
     set::MOI.AbstractScalarSet;
     allow_modify_function::Bool = false,
 ) where {T}
-    set = shift_constant(set, -func.constant)
+    set = shift_constant(set, -MOI.constant(func))
     if !allow_modify_function
         func = copy(func)
     end

src/Utilities/functions.jl

@@ -164,6 +164,13 @@ end
 function map_indices(index_map::F, f::MOI.VectorOfVariables) where {F<:Function}
     return MOI.VectorOfVariables(index_map.(f.variables))
 end
+function map_indices(index_map::Function, f::F) where {F <: VAF}
+    return F(map_indices.(index_map, f.terms), MOI.constant(f))
+end
+
+function map_indices(index_map::FI, f::F) where {FI <: Function, F <: MOI.GenericScalarAffineFunction}
+    return F(map_indices.(index_map, MOI.scalar_terms(f)), MOI.constant(f))
+end
 
 function map_indices(index_map::F, f::Union{SAF,VAF}) where {F<:Function}
     return typeof(f)(map_indices.(index_map, f.terms), MOI.constant(f))
@@ -286,11 +293,11 @@ end
 
 function substitute_variables(
     variable_map::F,
-    func::MOI.ScalarAffineFunction{T},
-) where {T,F<:Function}
-    g = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{T}[], MOI.constant(func))
+    func::FT,
+) where {T, F <: Function, FT <: MOI.GenericScalarAffineFunction{T}}
+    g = FT(MOI.ScalarAffineTerm{T}[], MOI.constant(func))
     for term in func.terms
-        operate!(
+    operate!(
             +,
             T,
             g,
@@ -714,7 +721,7 @@ canonicalize!(f::Union{MOI.VectorOfVariables,MOI.SingleVariable}) = f
 Convert a function to canonical form in-place, without allocating a copy to hold
 the result. See [`canonical`](@ref).
 """
-function canonicalize!(f::Union{SAF,VAF})
+function canonicalize!(f::Union{MOI.ScalarAffineFunction,VAF})
     sort_and_compress!(
         f.terms,
         MOI.term_indices,
@@ -724,6 +731,29 @@ function canonicalize!(f::Union{SAF,VAF})
     return f
 end
 
+function canonicalize!(f::MOI.ZippedAffineFunction)
+    delete_indices = BitSet()
+    for i in 1:length(f.terms.variable_indices)-1
+        if iszero(f.terms.coefficients[i])
+            push!(delete_indices, i)
+            continue
+        end
+        for j in i+1:length(f.terms.variable_indices)
+            if j in delete_indices
+                continue
+            end
+            if f.terms.variable_indices[i] == f.terms.variable_indices[j]
+                f.terms.coefficients[i] += f.terms.coefficients[j]
+                f.terms.coefficients[j] = 0
+                push!(delete_indices, j)
+            end
+        end
+    end
+    deleteat!(f.terms.coefficients, delete_indices)
+    deleteat!(f.terms.variable_indices, delete_indices)
+    return f
+end
+
 """
     canonicalize!(f::Union{ScalarQuadraticFunction, VectorQuadraticFunction})
 
@@ -873,21 +903,22 @@ function isapprox_zero(f::MOI.AbstractFunction, tol)
     return all_coefficients(α -> isapprox_zero(α, tol), f)
 end
 
-_is_constant(f::MOI.ScalarAffineFunction) = isempty(f.terms)
+_is_constant(f::MOI.GenericScalarAffineFunction) = isempty(f.terms)
+
 function _is_constant(f::MOI.ScalarQuadraticFunction)
     return isempty(f.affine_terms) && isempty(f.quadratic_terms)
 end
 
 Base.iszero(::MOI.SingleVariable) = false
 function Base.iszero(
-    f::Union{MOI.ScalarAffineFunction,MOI.ScalarQuadraticFunction},
+    f::Union{MOI.GenericScalarAffineFunction,MOI.ScalarQuadraticFunction},
 )
     return iszero(MOI.constant(f)) && _is_constant(canonical(f))
 end
 
 Base.isone(::MOI.SingleVariable) = false
 function Base.isone(
-    f::Union{MOI.ScalarAffineFunction,MOI.ScalarQuadraticFunction},
+    f::Union{MOI.GenericScalarAffineFunction,MOI.ScalarQuadraticFunction},
 )
     return isone(MOI.constant(f)) && _is_constant(canonical(f))
 end
@@ -1029,8 +1060,8 @@ end
 
 Return a new function `f` modified according to `change`.
 """
-function modify_function(f::SAF, change::MOI.ScalarConstantChange)
-    return SAF(f.terms, change.new_constant)
+function modify_function(f::F, change::MOI.ScalarConstantChange) where {F <: MOI.GenericScalarAffineFunction}
+    return F(MOI.scalar_terms(f), change.new_constant)
 end
 function modify_function(f::VAF, change::MOI.VectorConstantChange)
     return VAF(f.terms, change.new_constant)
@@ -1071,11 +1102,11 @@ function _modifycoefficient(
 end
 
 function modify_function(
-    f::MOI.ScalarAffineFunction{T},
+    f::F,
     change::MOI.ScalarCoefficientChange{T},
-) where {T}
-    terms = _modifycoefficient(f.terms, change.variable, change.new_coefficient)
-    return MOI.ScalarAffineFunction(terms, f.constant)
+) where {T, F <: MOI.GenericScalarAffineFunction{T}}
+    terms = _modifycoefficient(MOI.scalar_terms(f), change.variable, change.new_coefficient)
+    return F(terms, MOI.constant(f))
 end
 
 function modify_function(
@@ -1456,7 +1487,7 @@ end
 
 # Functions convertible to a ScalarAffineFunction
 const ScalarAffineLike{T} =
-    Union{T,MOI.SingleVariable,MOI.ScalarAffineFunction{T}}
+    Union{T,MOI.SingleVariable,MOI.GenericScalarAffineFunction{T}}
 # Functions convertible to a ScalarQuadraticFunction
 const ScalarQuadraticLike{T} =
     Union{ScalarAffineLike{T},MOI.ScalarQuadraticFunction{T}}
@@ -1465,7 +1496,7 @@ const ScalarQuadraticLike{T} =
 # `+(::SingleVariable, ::Any)` which should rather be
 # `+(::SingleVariable, ::Number)`.
 const TypedScalarLike{T} =
-    Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}
+    Union{MOI.GenericScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}
 # Used for overloading Base operator functions so `T` is not in the union to
 # avoid overloading e.g. `+(::Float64, ::Float64)`
 const ScalarLike{T} = Union{MOI.SingleVariable,TypedScalarLike{T}}
@@ -2334,7 +2365,7 @@ function Base.:*(f::TypedLike, g::Bool)
 end
 Base.:*(f::Bool, g::TypedLike) = g * f
 
-function Base.:^(func::MOI.ScalarAffineFunction{T}, p::Integer) where {T}
+function Base.:^(func::MOI.GenericScalarAffineFunction{T}, p::Integer) where {T}
     if iszero(p)
         return one(MOI.ScalarQuadraticFunction{T})
     elseif isone(p)
@@ -2374,7 +2405,7 @@ LinearAlgebra.symmetric(f::ScalarLike, ::Symbol) = f
 function promote_operation(
     ::typeof(/),
     ::Type{T},
-    ::Type{<:Union{MOI.SingleVariable,MOI.ScalarAffineFunction{T}}},
+    ::Type{<:Union{MOI.SingleVariable,MOI.GenericScalarAffineFunction{T}}},
     ::Type{T},
 ) where {T}
     return MOI.ScalarAffineFunction{T}
@@ -2567,7 +2598,7 @@ end
 number_of_affine_terms(::Type{T}, ::T) where {T} = 0
 number_of_affine_terms(::Type, ::SVF) = 1
 number_of_affine_terms(::Type, f::VVF) = length(f.variables)
-function number_of_affine_terms(::Type{T}, f::Union{SAF{T},VAF{T}}) where {T}
+function number_of_affine_terms(::Type{T}, f::Union{GAF{T},VAF{T}}) where {T}
     return length(f.terms)
 end
 function number_of_affine_terms(::Type{T}, f::Union{SQF{T},VQF{T}}) where {T}
@@ -2576,7 +2607,7 @@ end
 
 function number_of_quadratic_terms(
     ::Type{T},
-    ::Union{T,SVF,VVF,SAF{T},VAF{T}},
+    ::Union{T,SVF,VVF,GAF{T},VAF{T}},
 ) where {T}
     return 0
 end
@@ -2629,11 +2660,12 @@ function fill_terms(
     terms::Vector{MOI.VectorAffineTerm{T}},
     offset::Int,
     output_offset::Int,
-    func::Union{SAF{T},VAF{T}},
+    func::Union{GAF{T},VAF{T}},
 ) where {T}
     n = number_of_affine_terms(T, func)
     return terms[offset.+(1:n)] .= offset_term.(func.terms, output_offset)
 end
+
 function fill_terms(
     terms::Vector{MOI.VectorAffineTerm{T}},
     offset::Int,
@@ -2682,9 +2714,9 @@ function fill_constant(
     constant::Vector{T},
     offset::Int,
     output_offset::Int,
-    func::Union{SAF{T},SQF{T}},
+    func::Union{GAF{T},SQF{T}},
 ) where {T}
-    return constant[offset+1] = func.constant
+    return constant[offset+1] = MOI.constant(func)
 end
 function fill_constant(
     constant::Vector{T},
@@ -2708,12 +2740,12 @@ function vectorize(funcs::AbstractVector{MOI.SingleVariable})
 end
 
 """
-    vectorize(funcs::AbstractVector{MOI.ScalarAffineFunction{T}}) where T
+    vectorize(funcs::AbstractVector{<:MOI.GenericScalarAffineFunction{T}}) where {T}
 
 Returns the vector of scalar affine functions in the form of a
 `MOI.VectorAffineFunction{T}`.
 """
-function vectorize(funcs::AbstractVector{MOI.ScalarAffineFunction{T}}) where {T}
+function vectorize(funcs::AbstractVector{<:MOI.GenericScalarAffineFunction{T}}) where {T}
     nterms =
         mapreduce(func -> number_of_affine_terms(T, func), +, funcs, init = 0)
     out_dim = mapreduce(func -> output_dim(T, func), +, funcs, init = 0)

src/Utilities/mutable_arithmetics.jl

@@ -4,15 +4,16 @@
 
 MA.mutability(::Type{<:TypedLike}) = MA.IsMutable()
 
-function MA.mutable_copy(func::MOI.ScalarAffineFunction)
+function MA.mutable_copy(func::F) where {F <: MOI.GenericScalarAffineFunction}
     terms = [
         MOI.ScalarAffineTerm(
             MA.copy_if_mutable(t.coefficient),
             t.variable_index,
         ) for t in func.terms
     ]
-    return MOI.ScalarAffineFunction(terms, MA.copy_if_mutable(func.constant))
+    return F(terms, MA.copy_if_mutable(func.constant))
 end
+
 function MA.mutable_copy(func::MOI.ScalarQuadraticFunction)
     affine_terms = [
         MOI.ScalarAffineTerm(
@@ -75,14 +76,14 @@ function MA.promote_operation(
     op::PROMOTE_IMPLEMENTED_OP,
     F::Type{<:ScalarLike{T}},
     G::Type{<:ScalarLike{T}},
-) where {T,N}
+) where {T}
     return promote_operation(op, T, F, G)
 end
 function MA.promote_operation(
     op::PROMOTE_IMPLEMENTED_OP,
     F::Type{T},
     G::Type{<:TypedLike{T}},
-) where {T,N}
+) where {T}
     return promote_operation(op, T, F, G)
 end
 function MA.promote_operation(