MOAI decision trees as GDP disjunctions

Project.toml

@@ -11,10 +11,12 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 [weakdeps]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+MathOptAI = "e52c2cb8-508e-4e12-9dd2-9c4755b60e73"
 
 [extensions]
 AbstractGPsDisjunctiveProgramming = "AbstractGPs"
 InfiniteDisjunctiveProgramming = "InfiniteOpt"
+MathOptAIDisjunctiveProgramming = "MathOptAI"
 
 [compat]
 AbstractGPs = "0.5"
@@ -25,6 +27,7 @@ Ipopt = "1.9.0"
 JuMP = "1.18"
 Juniper = "0.9.3"
 LinearAlgebra = "1"
+MathOptAI = "1"
 Pajarito = "0.8"
 Reexport = "1"
 julia = "1.10"
@@ -39,4 +42,4 @@ Pajarito = "2f354839-79df-5901-9f0a-cdb2aac6fe30"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "HiGHS", "Test", "Juniper", "Ipopt", "InfiniteOpt", "AbstractGPs", "Hypatia", "Pajarito"]
+test = ["Aqua", "HiGHS", "Test", "Juniper", "Ipopt", "InfiniteOpt", "AbstractGPs", "Hypatia", "Pajarito", "MathOptAI"]

ext/MathOptAIDisjunctiveProgramming.jl

@@ -0,0 +1,97 @@
+module MathOptAIDisjunctiveProgramming
+
+import MathOptAI
+import DisjunctiveProgramming as DP
+import JuMP
+const _MOI = JuMP.MOI
+
+################################################################################
+#                            DISJUNCTIVE PREDICTOR
+################################################################################
+# A MathOptAI predictor to embed as a GDP disjunction
+struct _Disjunctive{P <: MathOptAI.AbstractPredictor} <: MathOptAI.AbstractPredictor
+    predictor::P
+end
+DP.Disjunctive(predictor) = _Disjunctive(predictor)
+
+################################################################################
+#                              DECISION TREES
+################################################################################
+# Leaf => root to leaf split rows as (feature, set) pairs; the left bound is
+# threshold - atol, MathOptAI's disjoint-branch convention
+_leaf_paths(leaf, path) = [leaf => path]
+function _leaf_paths(
+    node::MathOptAI.BinaryDecisionTree,
+    path = Pair{Int, _MOI.AbstractScalarSet}[]
+    )
+    left = node.feat_id => _MOI.LessThan(node.feat_value - node.atol)
+    right = node.feat_id => _MOI.GreaterThan(node.feat_value)
+    return vcat(_leaf_paths(node.lhs, [path; left]),
+        _leaf_paths(node.rhs, [path; right]))
+end
+
+# Leaf rule as an `Affine`: a constant leaf has zero coefficients
+function _affine_leaf(leaf::Real, num_inputs)
+    value = float(leaf)
+    return MathOptAI.Affine(zeros(typeof(value), 1, num_inputs), [value])
+end
+_affine_leaf(leaf::MathOptAI.Affine, num_inputs) = leaf
+
+# Bounds of an input read through JuMP, since InfiniteOpt's MathOptAI hook
+# reports none for its variables; an infinite parameter gives its domain
+function _input_bounds(x::JuMP.AbstractVariableRef)
+    JuMP.is_fixed(x) && return JuMP.fix_value(x), JuMP.fix_value(x)
+    value_type = JuMP.value_type(typeof(x))
+    JuMP.is_binary(x) && return zero(value_type), one(value_type)
+    lower = JuMP.has_lower_bound(x) ? JuMP.lower_bound(x) : missing
+    upper = JuMP.has_upper_bound(x) ? JuMP.upper_bound(x) : missing
+    return lower, upper
+end
+_input_bounds(x) = MathOptAI.get_variable_bounds(x)
+
+# Envelope of the leaf rules over the input box; `missing` when an input
+# with a nonzero coefficient is unbounded, which Hull then rejects itself
+function _output_bounds(leaves, x)
+    bounds = _input_bounds.(x)
+    lower, upper = Inf, -Inf
+    for leaf in leaves
+        leaf_lower = leaf_upper = only(leaf.b)
+        for (coeff, (input_lower, input_upper)) in zip(leaf.A, bounds)
+            iszero(coeff) && continue
+            leaf_lower += coeff * (coeff > 0 ? input_lower : input_upper)
+            leaf_upper += coeff * (coeff > 0 ? input_upper : input_lower)
+        end
+        lower, upper = min(lower, leaf_lower), max(upper, leaf_upper)
+    end
+    return lower, upper
+end
+
+function MathOptAI.add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::_Disjunctive{<:MathOptAI.BinaryDecisionTree},
+    x::Vector
+    )
+    paths = _leaf_paths(predictor.predictor)
+    leaves = [_affine_leaf(leaf, length(x)) for (leaf, _) in paths]
+    y = only(MathOptAI.add_variables(model, x, 1, "moai_Disjunctive"))
+    cons = Any[]
+    MathOptAI.set_variable_bounds(cons, y, _output_bounds(leaves, x)...)
+    # indicators vary over the parameters of y (if any), so each point
+    # picks its own leaf
+    tag = DP.VariableProperties(y).variable_type
+    logical = isnothing(tag) ? DP.Logical : DP.Logical(tag)
+    indicators = JuMP.@variable(model, [1:length(paths)], logical,
+        base_name = "moai_Disjunctive_Z")
+    for (indicator, (_, path), leaf) in zip(indicators, paths, leaves)
+        for (feature, set) in path # affine row, the form DP.jl reformulates
+            push!(cons, JuMP.@constraint(model, 1 * x[feature] in set,
+                DP.Disjunct(indicator)))
+        end
+        push!(cons, JuMP.@constraint(model,
+            y == only(leaf.A * x) + only(leaf.b), DP.Disjunct(indicator)))
+    end
+    push!(cons, DP.disjunction(model, indicators))
+    return [y], MathOptAI.Formulation(predictor, Any[y; indicators], cons)
+end
+
+end # module

src/extension_api.jl

@@ -117,3 +117,37 @@ function sample_M_values(sampler, objectives, sub, method, support_grids)
           "AbstractGPs loaded) to estimate M values with a " *
           "Gaussian process.")
 end
+
+################################################################################
+#                           MATHOPTAI PREDICTORS
+################################################################################
+"""
+    Disjunctive(predictor::MathOptAI.AbstractPredictor)
+
+Wrap a MathOptAI predictor so that `MathOptAI.add_predictor` embeds it in
+a [`GDPModel`](@ref) as a disjunction instead of MathOptAI's own
+mixed-integer formulation. The disjunction is reformulated like any other
+when the model is solved with `optimize!(model, gdp_method = ...)` or by
+[`reformulate_model`](@ref). Requires `MathOptAI` to be loaded. Supports
+`BinaryDecisionTree` with constant or `Affine` leaves: one disjunct per
+leaf, carrying the root-to-leaf split rows and the leaf model.
+
+The returned `Formulation` holds the output variable, the leaf indicators,
+the output bounds, the disjunct constraints and the disjunction.
+
+**Example**
+```julia
+julia> using DisjunctiveProgramming, MathOptAI, HiGHS
+
+julia> tree = MathOptAI.BinaryDecisionTree(1, 0.0, -1, 1);
+
+julia> model = GDPModel(HiGHS.Optimizer);
+
+julia> @variable(model, -1 <= x[1:1] <= 1);
+
+julia> y, f = MathOptAI.add_predictor(model, Disjunctive(tree), x);
+
+julia> optimize!(model, gdp_method = Hull())
+```
+"""
+function Disjunctive end

test/extensions/MathOptAIDisjunctiveProgramming.jl

@@ -0,0 +1,191 @@
+using HiGHS, MathOptAI, InfiniteOpt
+import DisjunctiveProgramming as DP
+
+# x[1] <= 0 ? -1 : (x[1] <= 1 ? 0 : 1), MathOptAI's own docstring example.
+# The atol gap sits well above HiGHS's MIP feasibility tolerance (1e-6).
+_step_tree(atol = 1e-3) = MathOptAI.BinaryDecisionTree{Float64, Int}(1, 0.0,
+    -1, MathOptAI.BinaryDecisionTree{Float64, Int}(1, 1.0, 0, 1, atol), atol)
+
+# linear model tree: x[1] <= 1 ? 2 x[1] - x[2] : x[1] + 3 x[2] + 1
+const _AffineTree = MathOptAI.BinaryDecisionTree{Float64,
+    MathOptAI.Affine{Float64}}
+_affine_tree(atol = 1e-3) = _AffineTree(1, 1.0,
+    MathOptAI.Affine([2.0 -1.0], [0.0]),
+    MathOptAI.Affine([1.0 3.0], [1.0]), atol)
+_affine_truth(x) = x[1] < 1 ? 2x[1] - x[2] : x[1] + 3x[2] + 1
+
+# surrogate value at pinned inputs, from the GDP embedding or MathOptAI's own
+# MIP; the pin is a global equality so the box the M values are read from
+# stays [lb, ub] and every leaf keeps a feasible region
+function _surrogate_value(model, predictor, xval, lb, ub; kwargs...)
+    set_silent(model)
+    @variable(model, lb[i] <= x[i in eachindex(xval)] <= ub[i])
+    y, _ = MathOptAI.add_predictor(model, predictor, x)
+    @constraint(model, x .== xval)
+    optimize!(model; kwargs...)
+    return is_solved_and_feasible(model) ? value(y[1]) : nothing
+end
+_gdp_value(predictor, method, xval, lb, ub) = _surrogate_value(
+    GDPModel(HiGHS.Optimizer), Disjunctive(predictor), xval, lb, ub,
+    gdp_method = method)
+_moai_value(predictor, xval, lb, ub) =
+    _surrogate_value(Model(HiGHS.Optimizer), predictor, xval, lb, ub)
+
+function test_disjunctive_wrapper()
+    predictor = Disjunctive(_step_tree())
+    @test predictor isa MathOptAI.AbstractPredictor
+    @test predictor.predictor isa MathOptAI.BinaryDecisionTree
+end
+
+function test_tree_formulation_contents()
+    model = GDPModel(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, -5 <= x[1:1] <= 5)
+    y, formulation = MathOptAI.add_predictor(model, Disjunctive(_step_tree()),
+        x)
+    @test length(y) == 1
+    @test lower_bound(y[1]) == -1 && upper_bound(y[1]) == 1
+    @test formulation isa MathOptAI.Formulation
+    @test formulation.predictor isa typeof(Disjunctive(_step_tree()))
+    # the output and the three leaf indicators
+    @test length(formulation.variables) == 1 + 3
+    @test formulation.variables[1] === y[1]
+    @test all(v -> v isa LogicalVariableRef, formulation.variables[2:end])
+    # two output bounds, 5 path rows, 3 leaf rows and the disjunction
+    @test length(formulation.constraints) == 2 + 5 + 3 + 1
+    @test all(c -> is_valid(model, c), formulation.constraints)
+    dref = formulation.constraints[end]
+    @test dref isa DisjunctionRef
+    @test length(DP._disjunctions(model)) == 1
+    @test length(DP._disjunct_constraints(model)) == 5 + 3
+    @test haskey(DP._exactly1_constraints(model), dref)
+    # the GDP method is chosen at solve time like for any disjunction
+    @objective(model, Max, y[1] - 0.1 * x[1])
+    optimize!(model, gdp_method = BigM())
+    @test is_solved_and_feasible(model)
+    @test value(y[1]) ≈ 1 && value(x[1]) ≈ 1
+    @objective(model, Min, y[1] + 0.1 * x[1])
+    optimize!(model, gdp_method = Hull())
+    @test value(y[1]) ≈ -1 && value(x[1]) ≈ -5
+end
+
+# every method reproduces MathOptAI's leaf regions, including the atol gap
+function test_tree_matches_mathoptai()
+    methods = [BigM(), Hull(), MBM(HiGHS.Optimizer), Indicator()]
+    for xval in [-2.0, -0.01, 0.0, 0.5, 0.99, 1.0, 3.0]
+        expected = _moai_value(_step_tree(), [xval], [-5], [5])
+        @test expected !== nothing
+        for method in methods
+            @test _gdp_value(_step_tree(), method, [xval], [-5], [5]) ≈
+                expected atol = 1e-8
+        end
+    end
+    @test _moai_value(_step_tree(), [1 - 5e-4], [-5], [5]) === nothing
+    for method in methods
+        @test _gdp_value(_step_tree(), method, [1 - 5e-4], [-5], [5]) ===
+            nothing
+    end
+end
+
+# whole-model methods apply too, since nothing is reformulated early
+function test_cutting_planes()
+    model = GDPModel(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, -5 <= x[1:1] <= 5)
+    y, _ = MathOptAI.add_predictor(model, Disjunctive(_step_tree()), x)
+    @objective(model, Max, y[1] - 0.1 * x[1])
+    optimize!(model, gdp_method = CuttingPlanes(HiGHS.Optimizer))
+    @test is_solved_and_feasible(model)
+    @test value(y[1]) ≈ 1 atol = 1e-4
+    @test value(x[1]) ≈ 1 atol = 1e-4
+end
+
+function test_affine_leaves()
+    for method in [BigM(), Hull()], xval in
+        [[-2.0, 0.5], [0.0, 0.0], [0.999, -1.0], [1.0, 2.0], [3.0, 1.5]]
+        @test _gdp_value(_affine_tree(), method, xval, [-3, -3], [3, 3]) ≈
+            _affine_truth(xval) atol = 1e-6
+    end
+    # envelope bounds and a free optimum
+    model = GDPModel(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, -3 <= x[1:2] <= 3)
+    y, formulation = MathOptAI.add_predictor(model,
+        Disjunctive(_affine_tree()), x)
+    @test lower_bound(y[1]) == -11 && upper_bound(y[1]) == 13
+    @objective(model, Max, y[1])
+    optimize!(model, gdp_method = Hull())
+    @test is_solved_and_feasible(model)
+    @test value(y[1]) ≈ 13 && value(x[1]) ≈ 3 && value(x[2]) ≈ 3
+    @objective(model, Min, y[1])
+    optimize!(model, gdp_method = Hull())
+    @test value(y[1]) ≈ -9 && value(x[1]) ≈ -3 && value(x[2]) ≈ 3
+end
+
+# an expression input takes MathOptAI's bounds: 2 x[1] - 1 in [-7, 5]
+function test_expression_input()
+    model = GDPModel(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, -3 <= x[1:2] <= 3)
+    inputs = [2 * x[1] - 1, x[2]]
+    y, _ = MathOptAI.add_predictor(model, Disjunctive(_step_tree()), inputs)
+    @test lower_bound(y[1]) == -1 && upper_bound(y[1]) == 1
+    z, _ = MathOptAI.add_predictor(model, Disjunctive(_affine_tree()), inputs)
+    @test lower_bound(z[1]) == -17 && upper_bound(z[1]) == 15
+    @constraint(model, x .== [1.5, 2.0])
+    optimize!(model, gdp_method = BigM())
+    @test is_solved_and_feasible(model)
+    @test value(y[1]) ≈ 1 && value(z[1]) ≈ _affine_truth([2.0, 2.0])
+end
+
+# tree evaluated pointwise: y(t) and the leaf indicators follow x(t), so the
+# active leaf changes across the supports
+function test_infinite_tree()
+    for method in [BigM(), Hull(), MBM(HiGHS.Optimizer)]
+        model = InfiniteGDPModel(HiGHS.Optimizer)
+        set_silent(model)
+        @infinite_parameter(model, t in [0, 1], num_supports = 11)
+        @variable(model, -5 <= x[1:1] <= 5, Infinite(t))
+        y, formulation = MathOptAI.add_predictor(model,
+            Disjunctive(_step_tree()), x)
+        @test parameter_refs(y[1]) == (t,)
+        @test lower_bound(y[1]) == -1 && upper_bound(y[1]) == 1
+        @test all(v -> v isa LogicalVariableRef, formulation.variables[2:end])
+        @constraint(model, x[1] == 4t - 2)
+        optimize!(model, gdp_method = method)
+        @test is_solved_and_feasible(model)
+        xvals = value(x[1])
+        @test value(y[1]) ≈ _step_value.(xvals) atol = 1e-6
+        @test length(value(formulation.variables[2])) == 11
+    end
+end
+_step_value(x) = x <= -1e-3 ? -1.0 : (x <= 1 - 1e-3 ? 0.0 : 1.0)
+
+# an infinite parameter as an input takes its domain as bounds
+function test_infinite_tree_parameter_input()
+    for method in [BigM(), Hull()]
+        model = InfiniteGDPModel(HiGHS.Optimizer)
+        set_silent(model)
+        @infinite_parameter(model, t in [0, 1], num_supports = 11)
+        @variable(model, -3 <= x <= 3, Infinite(t))
+        y, _ = MathOptAI.add_predictor(model, Disjunctive(_affine_tree()),
+            [x, t])
+        @test lower_bound(y[1]) == -7 && upper_bound(y[1]) == 7
+        @constraint(model, x == 3t - 1)
+        optimize!(model, gdp_method = method)
+        @test is_solved_and_feasible(model)
+        expected = _affine_truth.(zip(value(x), supports(t)))
+        @test value(y[1]) ≈ expected atol = 1e-6
+    end
+end
+
+@testset "MathOptAIDisjunctiveProgramming" begin
+    test_disjunctive_wrapper()
+    test_tree_formulation_contents()
+    test_tree_matches_mathoptai()
+    test_cutting_planes()
+    test_affine_leaves()
+    test_expression_input()
+    test_infinite_tree()
+    test_infinite_tree_parameter_input()
+end

test/runtests.jl

@@ -27,4 +27,5 @@ include("print.jl")
 include("solve.jl")
 include("extensions/InfiniteDisjunctiveProgramming.jl")
 include("extensions/AbstractGPsDisjunctiveProgramming.jl")
+include("extensions/MathOptAIDisjunctiveProgramming.jl")