Support P-Split Reformulation

README.md

@@ -172,6 +172,10 @@ The following reformulation methods are currently supported:
 
 3. [Indicator](https://jump.dev/JuMP.jl/stable/manual/constraints/#Indicator-constraints): This method reformulates each disjunct constraint into an indicator constraint with the Boolean reformulation counterpart of the Logical variable used to define the disjunct constraint.
 
+4. [P-Split](https://arxiv.org/abs/2202.05198): This method reformulates each disjunct constraint into P constraints, each with a partitioned group defined by the user. This method requires that terms in the constraint be convex additively seperable with respect to each variable. The `PSplit` struct is created with the following required arguments:
+
+    - `partition`: Partition of the variables to be split. All variables must be in exactly one partition. (e.g., The variables `x[1:4]` can be partitioned into two groups ` partition = [[x[1], x[2]], [x[3], x[4]]]`)
+
 ## Release Notes
 
 Prior to `v0.4.0`, the package did not leverage the JuMP extension capabilities and was not as robust. For these earlier releases, refer to [Perez, Joshi, and Grossmann, 2023](https://arxiv.org/abs/2304.10492v1) and the following [JuliaCon 2022 Talk](https://www.youtube.com/watch?v=AMIrgTTfUkI).

src/DisjunctiveProgramming.jl

@@ -24,6 +24,7 @@ include("bigm.jl")
 include("hull.jl")
 include("indicator.jl")
 include("print.jl")
+include("psplit.jl")
 
 # Define additional stuff that should not be exported
 const _EXCLUDE_SYMBOLS = [Symbol(@__MODULE__), :eval, :include]

src/datatypes.jl

@@ -383,6 +383,23 @@ struct Hull{T} <: AbstractReformulationMethod
     end
 end
 
+"""
+    PSplit <: AbstractReformulationMethod
+
+A type for using the p-split reformulation approach for disjunctive 
+constraints.
+
+**Fields**
+- `partition::Vector{Vector{V}}`: The partition of variables
+"""
+struct PSplit{V <: JuMP.AbstractVariableRef} <: AbstractReformulationMethod
+    partition::Vector{Vector{V}}
+
+    function PSplit(partition::Vector{Vector{V}}) where {V <: JuMP.AbstractVariableRef}
+        new{V}(partition)
+    end
+end
+
 # temp struct to store variable disaggregations (reset for each disjunction)
 mutable struct _Hull{V <: JuMP.AbstractVariableRef, T} <: AbstractReformulationMethod
     value::T

src/psplit.jl

@@ -0,0 +1,320 @@
+function _build_partitioned_expression(
+    expr::JuMP.AffExpr,
+    partition_variables::Vector{<:JuMP.AbstractVariableRef}
+)
+    constant = JuMP.constant(expr)
+    new_affexpr = AffExpr(0.0, Dict{JuMP.AbstractVariableRef,Float64}())
+    for var in partition_variables
+        add_to_expression!(new_affexpr, coefficient(expr, var), var) 
+    end
+    return new_affexpr, constant
+end
+
+function _build_partitioned_expression(
+    expr::JuMP.QuadExpr,
+    partition_variables::Vector{<:JuMP.AbstractVariableRef}
+)
+    new_quadexpr = QuadExpr(0.0, Dict{JuMP.AbstractVariableRef,Float64}())
+    constant = JuMP.constant(expr)
+    for var in partition_variables
+        add_to_expression!(new_quadexpr, get(expr.terms, JuMP.UnorderedPair(var, var), 0.0), var,var) 
+        add_to_expression!(new_quadexpr, coefficient(expr, var), var) 
+    end
+
+    return new_quadexpr, constant
+end
+
+function _build_partitioned_expression(
+    expr::JuMP.AbstractVariableRef,
+    partition_variables::Vector{<:JuMP.AbstractVariableRef}
+)
+    if expr in partition_variables
+        return expr, 0
+    else
+        return 0, 0
+    end
+end
+
+function _build_partitioned_expression(
+    expr::Number,
+    partition_variables::Vector{<:JuMP.AbstractVariableRef}
+)
+    return expr, 0
+end
+
+function _build_partitioned_expression(
+    expr::JuMP.NonlinearExpr,
+    partition_variables::Vector{<:JuMP.AbstractVariableRef}
+)
+    @warn "$expr is nonlinear and reformulation not be equivalent. 
+    Partitioned functions must be convex additively seperarable with one constant" maxlog = 1
+
+    new_args = Vector{Any}(undef, length(expr.args))
+    for i in 1:length(expr.args)
+        new_args[i] = _build_partitioned_expression(expr.args[i], partition_variables)[1]
+        if expr.head in (:exp, :cos, :cosh, :log, :log10, :log2) && new_args[i] == 0
+                return 0, 0
+        end
+    end
+    if expr.head == :/ && new_args[2] == 0
+        return 0, 0
+    end
+    constant = 0
+    if expr.head in (:+, :-)
+        constant = get(filter(x -> isa(x, Number), new_args), 1, 0.0)
+        if expr.head == :-
+            constant = -constant
+        end
+    end
+    return JuMP.NonlinearExpr(expr.head, new_args...) - constant, constant
+end
+
+
+function _bound_auxiliary(
+    model::JuMP.AbstractModel,
+    v::JuMP.AbstractVariableRef,
+    func::JuMP.AffExpr,
+    method::PSplit
+)   
+    lower_bound = 0
+    upper_bound = 0
+    for (var, coeff) in func.terms
+        if var != v
+            JuMP.is_binary(var) && continue
+            if coeff > 0
+                lower_bound += coeff * variable_bound_info(var)[1]
+                upper_bound += coeff * variable_bound_info(var)[2]
+            else
+                lower_bound += coeff * variable_bound_info(var)[2]
+                upper_bound += coeff * variable_bound_info(var)[1]
+            end
+        end
+    end
+    JuMP.set_lower_bound(v, lower_bound)
+    JuMP.set_upper_bound(v, upper_bound)
+end    
+
+function _bound_auxiliary(
+    model::JuMP.AbstractModel,
+    v::JuMP.AbstractVariableRef,
+    func::JuMP.AbstractVariableRef,
+    method::PSplit
+)   
+    if func != v
+        lower_bound = variable_bound_info(func)[1]
+        upper_bound = variable_bound_info(func)[2]
+        JuMP.set_lower_bound(v, lower_bound)
+        JuMP.set_upper_bound(v, upper_bound)
+    else
+        JuMP.set_lower_bound(v,0)
+        JuMP.set_upper_bound(v,0)
+    end
+end
+
+function _bound_auxiliary(
+    model::JuMP.AbstractModel,
+    v::JuMP.AbstractVariableRef,
+    func::Union{JuMP.NonlinearExpr, JuMP.QuadExpr, Number},
+    method::PSplit
+)   
+    @warn "Unable to calculate explicit bounds for auxiliary variables inside of nonlinear or quadratic expressions." maxlog = 1
+end
+
+requires_variable_bound_info(method::PSplit) = true
+
+function set_variable_bound_info(vref::JuMP.AbstractVariableRef, ::PSplit)
+    if !has_lower_bound(vref) || !has_upper_bound(vref)
+        error("Variable $vref must have both lower and upper bounds defined when using the PSplit reformulation.")
+    else
+        lb = lower_bound(vref)
+        ub = upper_bound(vref)
+    end
+    return lb, ub
+end
+
+#DONE WITH NL
+function reformulate_disjunct_constraint(
+    model::JuMP.AbstractModel,
+    con::JuMP.ScalarConstraint{T, S},
+    bvref::Union{JuMP.AbstractVariableRef, JuMP.AffExpr},
+    method::PSplit
+) where {T, S <: _MOI.LessThan}
+    p = length(method.partition)
+    v = [@variable(model, base_name = "v_$(hash(con))_$(i)") for i in 1:p]
+    _, constant = _build_partitioned_expression(con.func, method.partition[p])
+    reform_con = Vector{JuMP.AbstractConstraint}(undef, p + 1)
+    for i in 1:p
+        func, _ = _build_partitioned_expression(con.func, method.partition[i])
+        reform_con[i] = JuMP.build_constraint(error, func - v[i], MOI.LessThan(0.0))
+        _bound_auxiliary(model, v[i], func, method)
+    end
+    reform_con[end] = JuMP.build_constraint(error, sum(v[i] * bvref for i in 1:p) - (con.set.upper - constant) * bvref, MOI.LessThan(0.0))
+    return reform_con
+end
+#DONE WITH NL
+function reformulate_disjunct_constraint(
+    model::JuMP.AbstractModel,
+    con::JuMP.ScalarConstraint{T, S},
+    bvref::Union{JuMP.AbstractVariableRef, JuMP.AffExpr},
+    method::PSplit
+) where {T, S <: _MOI.GreaterThan}
+    p = length(method.partition)
+    reform_con = Vector{JuMP.AbstractConstraint}(undef, p + 1)
+    v = [@variable(model, base_name = "v_$(hash(con))_$(i)") for i in 1:p]
+    _, constant = _build_partitioned_expression(con.func, method.partition[p])
+
+    for i in 1:p
+        func, _ = _build_partitioned_expression(con.func, method.partition[i])
+        reform_con[i] = JuMP.build_constraint(error, -func - v[i], MOI.LessThan(0.0))
+        _bound_auxiliary(model, v[i], -func, method)
+    end
+    reform_con[end] = JuMP.build_constraint(error, sum(v[i] * bvref for i in 1:p) - (-con.set.lower + constant) * bvref, MOI.LessThan(0.0))
+    return reform_con
+end
+
+#Works with NL
+function reformulate_disjunct_constraint(
+    model::JuMP.AbstractModel,
+    con::JuMP.ScalarConstraint{T, S},
+    bvref::Union{JuMP.AbstractVariableRef, JuMP.AffExpr},
+    method::PSplit
+) where {T, S <: Union{_MOI.Interval, _MOI.EqualTo}}
+    p = length(method.partition)
+    reform_con_lt = Vector{JuMP.AbstractConstraint}(undef, p + 1)
+    reform_con_gt = Vector{JuMP.AbstractConstraint}(undef, p + 1)
+    #let [_, 1] be the upper bound and [_, 2] be the lower bound
+    _, constant = _build_partitioned_expression(con.func, method.partition[p]) 
+    v = [@variable(model, base_name = "v_$(hash(con))_$(i)_$(j)") for i in 1:p, j in 1:2]
+    for i in 1:p
+        func, _= _build_partitioned_expression(con.func, method.partition[i])
+        reform_con_lt[i] = JuMP.build_constraint(error, func - v[i,1], MOI.LessThan(0.0))
+        reform_con_gt[i] = JuMP.build_constraint(error, -func - v[i,2], MOI.LessThan(0.0))
+        _bound_auxiliary(model, v[i,1], func, method)
+        _bound_auxiliary(model, v[i,2], -func, method)
+    end
+    set_values = _set_values(con.set)
+    reform_con_lt[end] = JuMP.build_constraint(error, sum(v[i,1] * bvref for i in 1:p) - (set_values[2] - constant) * bvref, MOI.LessThan(0.0))
+    reform_con_gt[end] = JuMP.build_constraint(error, sum(v[i,2] * bvref for i in 1:p) - (-set_values[1] + constant) * bvref, MOI.LessThan(0.0))
+    #TODO: how do i avoid the vcat?
+    return vcat(reform_con_lt, reform_con_gt)
+end
+#Functions
+function reformulate_disjunct_constraint(
+    model::JuMP.AbstractModel,
+    con::JuMP.VectorConstraint{T, S, R},
+    bvref::Union{JuMP.AbstractVariableRef, JuMP.AffExpr},
+    method::PSplit
+) where {T, S <: _MOI.Nonpositives, R}
+    p = length(method.partition)
+    d = con.set.dimension
+    v = [@variable(model, base_name = "v_$(hash(con))_$(i)_$(j)") for i in 1:p, j in 1:d]
+    reform_con = Vector{JuMP.AbstractConstraint}(undef, p + 1)
+    constants = Vector{Number}(undef, d)
+    for i in 1:p
+        partitioned_expressions = [_build_partitioned_expression(con.func[j], method.partition[i]) for j in 1:d]
+        func = JuMP.@expression(model, [j = 1:d], partitioned_expressions[j][1])
+        constants .= [partitioned_expressions[j][2] for j in 1:d]
+        reform_con[i] = JuMP.build_constraint(error, func - v[i,:], _MOI.Nonpositives(d))
+        for j in 1:d
+            _bound_auxiliary(model, v[i,j], func[j], method)
+        end
+    end
+    new_func = JuMP.@expression(model,[j = 1:d], sum(v[i,j] * bvref for i in 1:p) + constants[j]*bvref)
+    reform_con[end] = JuMP.build_constraint(error, new_func, _MOI.Nonpositives(d))
+    return vcat(reform_con)
+end
+
+function reformulate_disjunct_constraint(
+    model::JuMP.AbstractModel,
+    con::JuMP.VectorConstraint{T, S, R},
+    bvref::Union{JuMP.AbstractVariableRef, JuMP.AffExpr},
+    method::PSplit
+) where {T, S <: _MOI.Nonnegatives, R}
+    p = length(method.partition)
+    d = con.set.dimension
+    v = [@variable(model, base_name = "v_$(hash(con))_$(i)_$(j)") for i in 1:p, j in 1:d]
+    reform_con = Vector{JuMP.AbstractConstraint}(undef, p + 1)
+    constants = Vector{Number}(undef, d)
+    for i in 1:p
+        #I should be subtracting the constant here.
+        partitioned_expressions = [_build_partitioned_expression(con.func[j], method.partition[i]) for j in 1:d]
+        func = JuMP.@expression(model, [j = 1:d], -partitioned_expressions[j][1])
+        constants .= [-partitioned_expressions[j][2] for j in 1:d]
+        reform_con[i] = JuMP.build_constraint(error, func - v[i,:], _MOI.Nonpositives(d))
+        for j in 1:d
+            _bound_auxiliary(model, v[i,j], func[j], method)
+        end
+    end
+    new_func = JuMP.@expression(model,[j = 1:d], sum(v[i,j] * bvref for i in 1:p) + constants[j]*bvref)
+    reform_con[end] = JuMP.build_constraint(error, new_func, _MOI.Nonpositives(d))
+    #TODO: how do i avoid the vcat?
+    return vcat(reform_con)
+end
+#TODO:
+function reformulate_disjunct_constraint(
+    model::JuMP.AbstractModel,
+    con::JuMP.VectorConstraint{T, S, R},
+    bvref::Union{JuMP.AbstractVariableRef, JuMP.AffExpr},
+    method::PSplit
+) where {T, S <: _MOI.Zeros, R}
+    p = length(method.partition)
+    d = con.set.dimension
+    reform_con_np = Vector{JuMP.AbstractConstraint}(undef, p + 1)  # nonpositive (≤ 0)
+    reform_con_nn = Vector{JuMP.AbstractConstraint}(undef, p + 1)  # nonnegative (≥ 0)
+    v = [@variable(model, base_name = "v_$(hash(con))_$(i)_$(j)_$(k)") for i in 1:p, j in 1:d, k in 1:2]
+    constants = Vector{Number}(undef, d)
+    for i in 1:p
+        partitioned_expressions = [_build_partitioned_expression(con.func[j], method.partition[i]) for j in 1:d]
+
+        # Nonpositive part: func ≤ 0 → func - v ≤ 0
+        func = JuMP.@expression(model, [j = 1:d], partitioned_expressions[j][1])        
+        constants .= [partitioned_expressions[j][2] for j in 1:d]
+
+        reform_con_np[i] = JuMP.build_constraint(error, func - v[i,:,1], _MOI.Nonpositives(d))
+        reform_con_nn[i] = JuMP.build_constraint(error, -func - v[i,:,2], _MOI.Nonpositives(d))
+
+        for j in 1:d
+            _bound_auxiliary(model, v[i,j,1], func[j], method)
+            _bound_auxiliary(model, v[i,j,2], -func[j], method)
+        end
+    end
+
+    # Final constraints: combine auxiliary variables with constants
+    new_func_np = JuMP.@expression(model,[j = 1:d], sum(v[i,j,1] * bvref for i in 1:p) + constants[j]*bvref)
+    new_func_nn = JuMP.@expression(model,[j = 1:d], -sum(v[i,j,2] * bvref for i in 1:p) - constants[j]*bvref)
+    reform_con_np[end] = JuMP.build_constraint(error, new_func_np, _MOI.Nonpositives(d))
+    reform_con_nn[end] = JuMP.build_constraint(error, new_func_nn, _MOI.Nonpositives(d))
+    return vcat(reform_con_np, reform_con_nn)
+end
+
+################################################################################
+#                          FALLBACK WARNING DISPATCHES
+################################################################################
+
+# Generic fallback for _build_partitioned_expression
+function _build_partitioned_expression(
+    expr::Any,
+    ::Vector{<:JuMP.AbstractVariableRef}
+)
+    error("PSplit: _build_partitioned_expression not implemented for expression type $(typeof(expr)). Supported types: GenericAffExpr, GenericQuadExpr, VariableRef, Number, NonlinearExpr.")
+end
+
+# Generic fallback for _bound_auxiliary
+function _bound_auxiliary(
+    ::JuMP.AbstractModel,
+    v::JuMP.AbstractVariableRef,
+    func::Any,
+    ::PSplit
+)
+    error("PSplit: _bound_auxiliary not implemented for function type $(typeof(func)). Auxiliary variable bounds may be suboptimal. Supported types: GenericAffExpr, VariableRef.")
+end
+
+# Generic fallback for reformulate_disjunct_constraint (vector)
+function reformulate_disjunct_constraint(
+    ::JuMP.AbstractModel,
+    con::Any,
+    ::Union{JuMP.AbstractVariableRef, JuMP.AffExpr},
+    ::PSplit
+)
+    error("PSplit: reformulate_disjunct_constraint not implemented for vector constraint set type $(typeof(con)). Supported types: VectorConstraint of _MOI.Nonnegatives, _MOI.Nonpositives, _MOI.Zeros.")
+end