BB: Collect strong branching data

src/MIPLearn.jl

@@ -21,6 +21,9 @@ global UserCutsComponent = PyNULL()
 global MemorySample = PyNULL()
 global Hdf5Sample = PyNULL()
 
+to_str_array(values) = py"to_str_array"(values)
+from_str_array(values) = py"from_str_array"(values)
+
 include("solvers/structs.jl")
 
 include("utils/log.jl")
@@ -65,9 +68,6 @@ function __init__()
     """
 end
 
-to_str_array(values) = py"to_str_array"(values)
-from_str_array(values) = py"from_str_array"(values)
-
 function convert(::Type{SparseMatrixCSC}, o::PyObject)
     I, J, V = pyimport("scipy.sparse").find(o)
     return sparse(I .+ 1, J .+ 1, V, o.shape...)

src/bb/BB.jl

@@ -10,6 +10,7 @@ frac(x) = x - floor(x)
 
 include("structs.jl")
 
+include("collect.jl")
 include("nodepool.jl")
 include("optimize.jl")
 include("log.jl")

src/bb/collect.jl

@@ -0,0 +1,61 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+using Printf
+using Base.Threads
+import Base.Threads: @threads, nthreads, threadid
+
+import ..load_data, ..Hdf5Sample
+
+function collect!(
+    optimizer,
+    filename::String;
+    time_limit::Float64 = Inf,
+    node_limit::Int = typemax(Int),
+    gap_limit::Float64 = 1e-4,
+    print_interval::Int = 5,
+    branch_rule::VariableBranchingRule = ReliabilityBranching(collect = true),
+)::NodePool
+    model = read_from_file(filename)
+    mip = init(optimizer)
+    load!(mip, model)
+
+    h5 = Hdf5Sample(replace(filename, ".mps.gz" => ".h5"), "r")
+    primal_bound = h5.get_scalar("mip_upper_bound")
+    if primal_bound === nothing
+        primal_bound = h5.get_scalar("mip_obj_value")
+    end
+    h5.file.close()
+
+    pool = solve!(
+        mip;
+        initial_primal_bound = primal_bound,
+        time_limit,
+        node_limit,
+        gap_limit,
+        print_interval,
+        branch_rule,
+    )
+
+    h5 = Hdf5Sample(replace(filename, ".mps.gz" => ".h5"))
+    pseudocost_up = [NaN for _ = 1:mip.nvars]
+    pseudocost_down = [NaN for _ = 1:mip.nvars]
+    priorities = [0.0 for _ = 1:mip.nvars]
+    for (var, var_hist) in pool.var_history
+        pseudocost_up[var.index] = var_hist.pseudocost_up
+        pseudocost_down[var.index] = var_hist.pseudocost_down
+        x = mean(var_hist.fractional_values)
+        f_up = x - floor(x)
+        f_down = ceil(x) - x
+        priorities[var.index] =
+            var_hist.pseudocost_up * f_up * var_hist.pseudocost_down * f_down
+    end
+    h5.put_array("bb_var_pseudocost_up", pseudocost_up)
+    h5.put_array("bb_var_pseudocost_down", pseudocost_down)
+    h5.put_array("bb_var_priority", priorities)
+    collect!(branch_rule, h5)
+    h5.file.close()
+
+    return pool
+end

src/bb/cplex.jl

@@ -19,15 +19,7 @@ function _probe(
     status = CPXlpopt(cpx.env, cpx.lp)
     status == 0 || error("CPXlpopt failed ($status)")
 
-    status = CPXstrongbranch(
-        cpx.env,
-        cpx.lp,
-        indices,
-        cnt,
-        downobj,
-        upobj,
-        itlim,
-    )
+    status = CPXstrongbranch(cpx.env, cpx.lp, indices, cnt, downobj, upobj, itlim)
     status == 0 || error("CPXstrongbranch failed ($status)")
 
     return upobj[1] * mip.sense, downobj[1] * mip.sense
@@ -38,4 +30,4 @@ function _relax_integrality!(cpx::CPLEX.Optimizer)::Nothing
     status = CPXchgprobtype(cpx.env, cpx.lp, CPLEX.CPXPROB_LP)
     status == 0 || error("CPXchgprobtype failed ($status)")
     return
-end
+end

src/bb/lp.jl

@@ -11,13 +11,14 @@ const MOI = MathOptInterface
 
 function init(constructor)::MIP
     return MIP(
-        constructor,
-        Any[nothing for t = 1:nthreads()],
-        Variable[],
-        Float64[],
-        Float64[],
-        1.0,
-        0,
+        constructor = constructor,
+        optimizers = Any[nothing for t = 1:nthreads()],
+        int_vars = Variable[],
+        int_vars_lb = Float64[],
+        int_vars_ub = Float64[],
+        sense = 1.0,
+        lp_iterations = 0,
+        nvars = 0,
     )
 end
 
@@ -27,10 +28,10 @@ function read!(mip::MIP, filename::AbstractString)::Nothing
 end
 
 function load!(mip::MIP, prototype::JuMP.Model)
+    mip.nvars = num_variables(prototype)
     _replace_zero_one!(backend(prototype))
     _assert_supported(backend(prototype))
-    mip.int_vars, mip.int_vars_lb, mip.int_vars_ub =
-        _get_int_variables(backend(prototype))
+    mip.int_vars, mip.int_vars_lb, mip.int_vars_ub = _get_int_variables(backend(prototype))
     mip.sense = _get_objective_sense(backend(prototype))
     _relax_integrality!(backend(prototype))
     @threads for t = 1:nthreads()
@@ -133,11 +134,7 @@ function _get_int_variables(
                 var_ub = constr.upper
                 MOI.delete(optimizer, _upper_bound_index(var))
             end
-            MOI.add_constraint(
-                optimizer,
-                var,
-                MOI.Interval(var_lb, var_ub),
-            )
+            MOI.add_constraint(optimizer, var, MOI.Interval(var_lb, var_ub))
         end
         push!(vars, var)
         push!(lb, var_lb)

src/bb/optimize.jl

@@ -19,7 +19,7 @@ function solve!(
     enable_plunging = true,
 )::NodePool
     time_initial = time()
-    pool = NodePool(mip=mip)
+    pool = NodePool(mip = mip)
     pool.primal_bound = initial_primal_bound
 
     root_node = _create_node(mip)
@@ -34,9 +34,9 @@ function solve!(
 
     offer(
         pool,
-        parent_node=nothing,
-        child_nodes=[root_node],
-        print_interval=print_interval,
+        parent_node = nothing,
+        child_nodes = [root_node],
+        print_interval = print_interval,
     )
     @threads for t = 1:nthreads()
         child_one, child_zero, suggestions = nothing, nothing, Node[]
@@ -47,10 +47,10 @@ function solve!(
             end
             node = take(
                 pool,
-                suggestions=suggestions,
-                time_remaining=time_limit - time_elapsed,
-                node_limit=node_limit,
-                gap_limit=gap_limit,
+                suggestions = suggestions,
+                time_remaining = time_limit - time_elapsed,
+                node_limit = node_limit,
+                gap_limit = gap_limit,
             )
             if node == :END
                 break
@@ -85,26 +85,26 @@ function solve!(
 
                 child_zero = _create_node(
                     mip,
-                    index=ids[2],
-                    parent=node,
-                    branch_var=branch_var,
-                    branch_var_lb=var_lb,
-                    branch_var_ub=floor(var_value),
+                    index = ids[2],
+                    parent = node,
+                    branch_var = branch_var,
+                    branch_var_lb = var_lb,
+                    branch_var_ub = floor(var_value),
                 )
                 child_one = _create_node(
                     mip,
-                    index=ids[1],
-                    parent=node,
-                    branch_var=branch_var,
-                    branch_var_lb=ceil(var_value),
-                    branch_var_ub=var_ub,
+                    index = ids[1],
+                    parent = node,
+                    branch_var = branch_var,
+                    branch_var_lb = ceil(var_value),
+                    branch_var_ub = var_ub,
                 )
                 offer(
                     pool,
-                    parent_node=node,
-                    child_nodes=[child_one, child_zero],
-                    time_elapsed=time_elapsed,
-                    print_interval=print_interval,
+                    parent_node = node,
+                    child_nodes = [child_one, child_zero],
+                    time_elapsed = time_elapsed,
+                    print_interval = print_interval,
                 )
             end
         end
@@ -114,11 +114,11 @@ end
 
 function _create_node(
     mip;
-    index::Int=0,
-    parent::Union{Nothing,Node}=nothing,
-    branch_var::Union{Nothing,Variable}=nothing,
-    branch_var_lb::Union{Nothing,Float64}=nothing,
-    branch_var_ub::Union{Nothing,Float64}=nothing
+    index::Int = 0,
+    parent::Union{Nothing,Node} = nothing,
+    branch_var::Union{Nothing,Variable} = nothing,
+    branch_var_lb::Union{Nothing,Float64} = nothing,
+    branch_var_ub::Union{Nothing,Float64} = nothing,
 )::Node
     if parent === nothing
         branch_vars = Variable[]
@@ -135,8 +135,9 @@ function _create_node(
     status, obj = solve_relaxation!(mip)
     if status == :Optimal
         vals = values(mip, mip.int_vars)
-        fractional_indices =
-            [j for j in 1:length(mip.int_vars) if 1e-6 < vals[j] - floor(vals[j]) < 1 - 1e-6]
+        fractional_indices = [
+            j for j in 1:length(mip.int_vars) if 1e-6 < vals[j] - floor(vals[j]) < 1 - 1e-6
+        ]
         fractional_values = vals[fractional_indices]
         fractional_variables = mip.int_vars[fractional_indices]
     else
@@ -159,51 +160,6 @@ function _create_node(
     )
 end
 
-function solve!(
-    optimizer,
-    filename::String;
-    time_limit::Float64=Inf,
-    node_limit::Int=typemax(Int),
-    gap_limit::Float64=1e-4,
-    print_interval::Int=5,
-    branch_rule::VariableBranchingRule=ReliabilityBranching()
-)::NodePool
-    model = read_from_file("$filename.mps.gz")
-    mip = init(optimizer)
-    load!(mip, model)
-
-    h5 = Hdf5Sample("$filename.h5")
-    primal_bound = h5.get_scalar("mip_obj_value")
-    nvars = length(h5.get_array("static_var_names"))
-
-    pool = solve!(
-        mip;
-        initial_primal_bound=primal_bound,
-        time_limit,
-        node_limit,
-        gap_limit,
-        print_interval,
-        branch_rule
-    )
-
-    pseudocost_up = [NaN for _ = 1:nvars]
-    pseudocost_down = [NaN for _ = 1:nvars]
-    priorities = [0.0 for _ in 1:nvars]
-    for (var, var_hist) in pool.var_history
-        pseudocost_up[var.index] = var_hist.pseudocost_up
-        pseudocost_down[var.index] = var_hist.pseudocost_down
-        x = mean(var_hist.fractional_values)
-        f_up = x - floor(x)
-        f_down = ceil(x) - x
-        priorities[var.index] = var_hist.pseudocost_up * f_up * var_hist.pseudocost_down * f_down
-    end
-    h5.put_array("bb_var_pseudocost_up", pseudocost_up)
-    h5.put_array("bb_var_pseudocost_down", pseudocost_down)
-    h5.put_array("bb_var_priority", priorities)
-
-    return pool
-end
-
 function _set_node_bounds(node::Node)
     set_bounds!(node.mip, node.branch_vars, node.branch_lb, node.branch_ub)
 end

src/bb/structs.jl

@@ -9,7 +9,7 @@ struct Variable
     index::Any
 end
 
-mutable struct MIP
+Base.@kwdef mutable struct MIP
     constructor::Any
     optimizers::Vector
     int_vars::Vector{Variable}
@@ -17,6 +17,7 @@ mutable struct MIP
     int_vars_ub::Vector{Float64}
     sense::Float64
     lp_iterations::Int64
+    nvars::Int
 end
 
 struct Node

src/bb/varbranch/reliability.jl

@@ -2,6 +2,16 @@
 #  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
+import ..to_str_array
+
+Base.@kwdef mutable struct ReliabilityBranchingStats
+    branched_count::Vector{Int} = []
+    num_strong_branch_calls = 0
+    score_var_names::Vector{String} = []
+    score_features::Vector{Vector{Float32}} = []
+    score_targets::Vector{Float32} = []
+end
+
 """
     ReliabilityBranching
 
@@ -13,12 +23,14 @@ Base.@kwdef mutable struct ReliabilityBranching <: VariableBranchingRule
     min_samples::Int = 8
     max_sb_calls::Int = 100
     look_ahead::Int = 10
-    n_sb_calls::Int = 0
     side_effect::Bool = true
     max_iterations::Int = 1_000_000
     aggregation::Symbol = :prod
+    stats::ReliabilityBranchingStats = ReliabilityBranchingStats()
+    collect::Bool = false
 end
 
+
 function _strong_branch_score(;
     node::Node,
     pool::NodePool,
@@ -28,7 +40,6 @@ function _strong_branch_score(;
     max_iterations::Int,
     aggregation::Symbol,
 )::Tuple{Float64,Int}
-
     # Find current variable lower and upper bounds
     offset = findfirst(isequal(var), node.mip.int_vars)
     var_lb = node.mip.int_vars_lb[offset]
@@ -68,6 +79,14 @@ function find_branching_var(
     node::Node,
     pool::NodePool,
 )::Variable
+    stats = rule.stats
+
+    # Initialize statistics
+    if isempty(stats.branched_count)
+        stats.branched_count = zeros(node.mip.nvars)
+    end
+
+    # Sort variables by pseudocost score
     nfrac = length(node.fractional_variables)
     pseudocost_scores = [
         _pseudocost_score(
@@ -79,10 +98,37 @@ function find_branching_var(
     ]
     σ = sortperm(pseudocost_scores, rev = true)
     sorted_vars = node.fractional_variables[σ]
+
+    if rule.collect
+        # Compute dynamic features for all fractional variables
+        features = []
+        for (i, var) in enumerate(sorted_vars)
+            branched_count = stats.branched_count[var.index]
+            branched_count_rel = 0.0
+            branched_count_sum = sum(stats.branched_count[var.index])
+            if branched_count_sum > 0
+                branched_count_rel = branched_count / branched_count_sum
+            end
+            push!(
+                features,
+                Float32[
+                    nfrac,
+                    node.fractional_values[σ[i]],
+                    node.depth,
+                    pseudocost_scores[σ[i]][1],
+                    branched_count,
+                    branched_count_rel,
+                ],
+            )
+        end
+    end
+
     _set_node_bounds(node)
     no_improv_count, n_sb_calls = 0, 0
-    max_score, max_var = pseudocost_scores[σ[1]], sorted_vars[1]
+    max_score, max_var = (-Inf, -Inf), sorted_vars[1]
     for (i, var) in enumerate(sorted_vars)
+
+        # Decide whether to use strong branching
         use_strong_branch = true
         if n_sb_calls >= rule.max_sb_calls
             use_strong_branch = false
@@ -95,9 +141,10 @@ function find_branching_var(
                 end
             end
         end
+
         if use_strong_branch
+            # Compute strong branching score
             n_sb_calls += 1
-            rule.n_sb_calls += 1
             score = _strong_branch_score(
                 node = node,
                 pool = pool,
@@ -107,6 +154,13 @@ function find_branching_var(
                 max_iterations = rule.max_iterations,
                 aggregation = rule.aggregation,
             )
+
+            if rule.collect
+                # Store training data
+                push!(stats.score_var_names, name(node.mip, var))
+                push!(stats.score_features, features[i])
+                push!(stats.score_targets, score[1])
+            end
         else
             score = pseudocost_scores[σ[i]]
         end
@@ -119,5 +173,16 @@ function find_branching_var(
         no_improv_count <= rule.look_ahead || break
     end
     _unset_node_bounds(node)
+
+    # Update statistics
+    stats.branched_count[max_var.index] += 1
+    stats.num_strong_branch_calls += n_sb_calls
+
     return max_var
 end
+
+function collect!(rule::ReliabilityBranching, h5)
+    h5.put_array("bb_score_var_names", to_str_array(rule.stats.score_var_names))
+    h5.put_array("bb_score_features", vcat(rule.stats.score_features'...))
+    h5.put_array("bb_score_targets", rule.stats.score_targets)
+end

src/bb/varbranch/strong.jl

@@ -21,12 +21,12 @@ end
 
 function find_branching_var(rule::StrongBranching, node::Node, pool::NodePool)::Variable
     rb_rule = ReliabilityBranching(
-        min_samples=typemax(Int),
-        max_sb_calls=rule.max_calls,
-        look_ahead=rule.look_ahead,
-        side_effect=rule.side_effect,
-        max_iterations=rule.max_iterations,
-        aggregation=rule.aggregation,
+        min_samples = typemax(Int),
+        max_sb_calls = rule.max_calls,
+        look_ahead = rule.look_ahead,
+        side_effect = rule.side_effect,
+        max_iterations = rule.max_iterations,
+        aggregation = rule.aggregation,
     )
     return find_branching_var(rb_rule, node, pool)
 end

src/solvers/jump_solver.jl

@@ -78,13 +78,9 @@ function _update_solution!(data::JuMPSolverData)
                 rc += shadow_price(FixRef(var))
             end
             push!(data.reduced_costs, rc)
-            
+
             # Basis status
-            data.basis_status[var] = MOI.get(
-                data.model,
-                MOI.VariableBasisStatus(),
-                var,
-            )
+            data.basis_status[var] = MOI.get(data.model, MOI.VariableBasisStatus(), var)
         end
 
         try

src/utils/benchmark.jl

@@ -17,7 +17,7 @@ function run_benchmarks(;
     solvers = OrderedDict(
         "baseline" => LearningSolver(optimizer),
         "ml-exact" => LearningSolver(optimizer),
-        "ml-heuristic" => LearningSolver(optimizer, mode="heuristic"),
+        "ml-heuristic" => LearningSolver(optimizer, mode = "heuristic"),
     )
 
     #solve!(solvers["baseline"], train_instances, build_model; progress)
@@ -43,7 +43,7 @@ function run_benchmarks(;
         end
     end
     CSV.write(output_filename, results)
-    
+
     # fig_filename = "$(tempname()).svg"
     # df = pyimport("pandas").read_csv(csv_filename)
     # miplearn.benchmark.plot(df, output=fig_filename)

test/bb/lp_test.jl

@@ -28,7 +28,7 @@ function runtests(optimizer_name, optimizer; large = true)
             @test BB.name(mip, mip.int_vars[1]) == "xab"
             @test BB.name(mip, mip.int_vars[2]) == "xac"
             @test BB.name(mip, mip.int_vars[3]) == "xad"
-            
+
             @test mip.int_vars_lb[1] == 0.0
             @test mip.int_vars_ub[1] == 1.0
 
@@ -70,25 +70,25 @@ function runtests(optimizer_name, optimizer; large = true)
         end
 
         @testset "varbranch" begin
-            branch_rules = [
-                BB.RandomBranching(),
-                BB.FirstInfeasibleBranching(),
-                BB.LeastInfeasibleBranching(),
-                BB.MostInfeasibleBranching(),
-                BB.PseudocostBranching(),
-                BB.StrongBranching(),
-                BB.ReliabilityBranching(),
-                BB.HybridBranching(),
-                BB.StrongBranching(aggregation=:min),
-                BB.ReliabilityBranching(aggregation=:min),
-            ]
-            for branch_rule in branch_rules
-                for instance in ["bell5", "vpm2"]
+            for instance in ["bell5", "vpm2"]
+                for branch_rule in [
+                    BB.RandomBranching(),
+                    BB.FirstInfeasibleBranching(),
+                    BB.LeastInfeasibleBranching(),
+                    BB.MostInfeasibleBranching(),
+                    BB.PseudocostBranching(),
+                    BB.StrongBranching(),
+                    BB.ReliabilityBranching(),
+                    BB.HybridBranching(),
+                    BB.StrongBranching(aggregation = :min),
+                    BB.ReliabilityBranching(aggregation = :min, collect = true),
+                ]
                     h5 = Hdf5Sample("$basepath/../fixtures/$instance.h5")
                     mip_lower_bound = h5.get_scalar("mip_lower_bound")
                     mip_upper_bound = h5.get_scalar("mip_upper_bound")
                     mip_sense = h5.get_scalar("mip_sense")
-                    mip_primal_bound = mip_sense == "min" ? mip_upper_bound : mip_lower_bound
+                    mip_primal_bound =
+                        mip_sense == "min" ? mip_upper_bound : mip_lower_bound
                     h5.file.close()
 
                     mip = BB.init(optimizer)
@@ -104,25 +104,35 @@ function runtests(optimizer_name, optimizer; large = true)
                 end
             end
         end
+
+        @testset "collect" begin
+            rule = BB.ReliabilityBranching(collect = true)
+            BB.collect!(
+                optimizer,
+                "$basepath/../fixtures/bell5.mps.gz",
+                node_limit = 100,
+                print_interval = 10,
+                branch_rule = rule,
+            )
+            n_sb = rule.stats.num_strong_branch_calls
+            h5 = Hdf5Sample("$basepath/../fixtures/bell5.h5")
+            @test size(h5.get_array("bb_var_pseudocost_up")) == (104,)
+            @test size(h5.get_array("bb_score_var_names")) == (n_sb,)
+            @test size(h5.get_array("bb_score_features")) == (n_sb, 6)
+            @test size(h5.get_array("bb_score_targets")) == (n_sb,)
+            h5.file.close()
+        end
     end
 end
 
 @testset "BB" begin
-    @time runtests(
-        "Clp",
-        optimizer_with_attributes(
-            Clp.Optimizer,
-        ),
-    )
+    @time runtests("Clp", optimizer_with_attributes(Clp.Optimizer))
 
     if is_gurobi_available
         using Gurobi
         @time runtests(
             "Gurobi",
-            optimizer_with_attributes(
-                Gurobi.Optimizer,
-                "Threads" => 1,
-            )
+            optimizer_with_attributes(Gurobi.Optimizer, "Threads" => 1),
         )
     end
 
@@ -130,10 +140,7 @@ end
         using CPLEX
         @time runtests(
             "CPLEX",
-            optimizer_with_attributes(
-                CPLEX.Optimizer,
-                "CPXPARAM_Threads" => 1,
-            ),
+            optimizer_with_attributes(CPLEX.Optimizer, "CPXPARAM_Threads" => 1),
         )
     end
 end

test/instance/file_instance_test.jl

@@ -33,9 +33,9 @@ using Cbc
     @testset "Save and load data" begin
         filename = tempname()
         data = KnapsackData(
-            weights=[5.0, 5.0, 5.0],
-            prices=[1.0, 1.0, 1.0],
-            capacity=3.0,
+            weights = [5.0, 5.0, 5.0],
+            prices = [1.0, 1.0, 1.0],
+            capacity = 3.0,
         )
         MIPLearn.save_data(filename, data)
         loaded = MIPLearn.load_data(filename)