Implement two-stage spanning tree benchmark

Project.toml

@@ -10,10 +10,12 @@ projects = ["docs", "test"]
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 ConstrainedShortestPaths = "b3798467-87dc-4d99-943d-35a1bd39e395"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+GLPK = "60bf3e95-4087-53dc-ae20-288a0d20c6a6"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
@@ -45,10 +47,12 @@ DFLBenchmarksPlotsExt = "Plots"
 Combinatorics = "1.0.3"
 ConstrainedShortestPaths = "0.6.0"
 DataDeps = "0.7"
+DataStructures = "0.18.22"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 FileIO = "1.17.0"
 Flux = "0.16"
+GLPK = "1.2.1"
 Graphs = "1.11"
 HiGHS = "1.9"
 Images = "0.26.1"
@@ -65,7 +69,7 @@ Plots = "1"
 Printf = "1"
 Random = "1"
 Requires = "1.3.0"
-SCIP = "0.12"
+SCIP = "0.12.8"
 SimpleWeightedGraphs = "1.4"
 SparseArrays = "1"
 Statistics = "1"

docs/src/benchmarks/two_stage_spanning_tree.md

@@ -0,0 +1,28 @@
+# Problem statement
+
+We consider a two-stage stochastic variant of the classic [minimum spanning tree problem](https://en.wikipedia.org/wiki/Minimum_spanning_tree).
+
+Rather than immediately constructing a spanning tree and incurring a cost ``c_e`` for each selected edge in the tree, we instead can build only a partial tree (forest) during the first stage and paying first stage costs ``c_e`` for the selected edges. Then, second stage costs ``d_e`` are revealed and replace first stage costs. The task then involves completing the first stage forest into a spanning tree.
+
+The objective is to minimize the total incurred cost in expectation.
+
+## Instance
+Let ``G = (V,E)`` be an undirected **graph**, and ``S`` be a finite set of **scenarios**.
+
+For each edge ``e`` in ``E``, we have a **first stage cost** ``c_e\in\mathbb{R}``.
+
+For each edge ``e`` in ``E`` and scenario ``s`` in ``S``, we have a **second stage cost** ``d_{es}\in\mathbb{R}``.
+
+# MIP formulation
+Unlike the regular minimum spanning tree problem, this two-stage variant is NP-hard.
+However, it can still be formulated as linear program with binary variables, and exponential number of constraints.
+```math
+\begin{array}{lll}
+\min\limits_{y, z}\, & \sum\limits_{e\in E}c_e y_e + \frac{1}{|S|}\sum\limits_{s \in S}d_{es}z_{es} & \\
+\mathrm{s.t.}\, & \sum\limits_{e\in E}y_e + z_{es} = |V| - 1, & \forall s \in S\\
+& \sum\limits_{e\in E(Y)} y_e + z_{es} \leq |Y| - 1,\quad & \forall \emptyset \subsetneq Y \subsetneq V,\, \forall s\in S\\
+& y_e\in \{0, 1\}, & \forall e\in E\\
+& z_{es}\in \{0, 1\}, & \forall e\in E, \forall s\in S
+\end{array}
+```
+where ``y_e`` is a binary variable indicating if ``e`` is in the first stage solution, and ``z_{es}`` is a binary variable indicating if ``e`` is in the second stage solution for scenario ``s``.

ext/DFLBenchmarksPlotsExt.jl

@@ -11,6 +11,7 @@ include("plots/argmax2d_plots.jl")
 include("plots/warcraft_plots.jl")
 include("plots/svs_plots.jl")
 include("plots/dvs_plots.jl")
+include("plots/tst_plots.jl")
 
 """
     plot_solution(bench::AbstractBenchmark, sample::DataSample, y; kwargs...)

ext/plots/dvs_plots.jl

@@ -3,8 +3,6 @@ using Printf: @sprintf
 
 has_visualization(::DynamicVehicleSchedulingBenchmark) = true
 
-# ── helpers (moved from static_vsp/plot.jl) ─────────────────────────────────
-
 function _plot_static_instance(
     x_depot,
     y_depot,

ext/plots/tst_plots.jl

@@ -0,0 +1,110 @@
+import DecisionFocusedLearningBenchmarks.TwoStageSpanningTree:
+    TwoStageSpanningTreeInstance,
+    TwoStageSpanningTreeSolution,
+    nb_scenarios,
+    kruskal,
+    solution_from_first_stage_forest
+using Graphs: ne, nv, edges, src, dst
+
+has_visualization(::TwoStageSpanningTreeBenchmark) = true
+
+function _plot_grid_graph(
+    graph,
+    n,
+    m,
+    weights=nothing;
+    edge_colors=fill(:black, ne(graph)),
+    edge_widths=fill(1, ne(graph)),
+    edge_labels=fill(nothing, ne(graph)),
+    δ=0.25,
+    δ₂=0.13,
+    space_for_legend=0,
+)
+    node_pos = [((i - 1) % n, floor((i - 1) / n)) for i in 1:nv(graph)]
+    function segment(i, j)
+        return [node_pos[i][1], node_pos[j][1]], [node_pos[i][2], node_pos[j][2]]
+    end
+    fig = Plots.plot(;
+        axis=([], false),
+        ylimits=(-δ, m - 1 + δ + space_for_legend),
+        xlimits=(-δ, n - 1 + δ),
+        aspect_ratio=:equal,
+        leg=:top,
+    )
+    for (color, width, label, e) in zip(edge_colors, edge_widths, edge_labels, edges(graph))
+        Plots.plot!(fig, segment(src(e), dst(e))...; color, width, label)
+    end
+    Plots.scatter!(fig, node_pos; label=nothing, markersize=15, color=:lightgrey)
+    if !isnothing(weights)
+        for (w, e) in zip(weights, edges(graph))
+            i, j = src(e), dst(e)
+            x = (node_pos[j][1] + node_pos[i][1]) / 2
+            y = (node_pos[j][2] + node_pos[i][2]) / 2
+            j == i + 1 ? (y += δ₂) : (x -= δ₂)
+            Plots.annotate!(fig, x, y, Int(w))
+        end
+    end
+    return fig
+end
+
+function plot_instance(bench::TwoStageSpanningTreeBenchmark, sample::DataSample; kwargs...)
+    (; n, m) = bench
+    return _plot_grid_graph(
+        sample.instance.graph, n, m, sample.instance.first_stage_costs; kwargs...
+    )
+end
+
+function plot_solution(bench::TwoStageSpanningTreeBenchmark, sample::DataSample; kwargs...)
+    (; n, m) = bench
+    (; instance) = sample.context
+    y = sample.y
+    isnothing(y) && error("sample.y is nothing — provide a labeled sample")
+
+    # Use the evaluation scenario if present, otherwise fall back to first feature scenario
+    d_plot = if hasproperty(sample.extra, :scenario)
+        sample.extra.scenario
+    else
+        instance.second_stage_costs[:, 1]
+    end
+
+    # Complete first-stage forest to a spanning tree for display
+    inst_s = TwoStageSpanningTreeInstance(
+        instance.graph, instance.first_stage_costs, reshape(d_plot, :, 1)
+    )
+    full_sol = solution_from_first_stage_forest(BitVector(y .> 0), inst_s)
+
+    yv, zv = full_sol.y, full_sol.z[:, 1]
+    is_labeled_1 = is_labeled_2 = false
+    edge_labels = fill(nothing, ne(instance.graph))
+    for i in 1:ne(instance.graph)
+        if !is_labeled_1 && yv[i]
+            edge_labels[i] = "First stage"
+            is_labeled_1 = true
+        elseif !is_labeled_2 && zv[i]
+            edge_labels[i] = "Second stage"
+            is_labeled_2 = true
+        end
+    end
+    edge_colors = [
+        if yv[i]
+            :red
+        elseif zv[i]
+            :green
+        else
+            :black
+        end for i in 1:ne(instance.graph)
+    ]
+    edge_widths = [(yv[i] || zv[i]) ? 3 : 1 for i in 1:ne(instance.graph)]
+    weights = yv .* instance.first_stage_costs .+ .!yv .* d_plot
+    return _plot_grid_graph(
+        instance.graph,
+        n,
+        m,
+        weights;
+        edge_colors,
+        edge_widths,
+        edge_labels,
+        space_for_legend=0.75,
+        kwargs...,
+    )
+end

src/DecisionFocusedLearningBenchmarks.jl

@@ -55,6 +55,7 @@ include("Warcraft/Warcraft.jl")
 include("FixedSizeShortestPath/FixedSizeShortestPath.jl")
 include("PortfolioOptimization/PortfolioOptimization.jl")
 include("StochasticVehicleScheduling/StochasticVehicleScheduling.jl")
+include("TwoStageSpanningTree/TwoStageSpanningTree.jl")
 include("DynamicVehicleScheduling/DynamicVehicleScheduling.jl")
 include("DynamicAssortment/DynamicAssortment.jl")
 include("Maintenance/Maintenance.jl")
@@ -91,6 +92,7 @@ using .Warcraft
 using .FixedSizeShortestPath
 using .PortfolioOptimization
 using .StochasticVehicleScheduling
+using .TwoStageSpanningTree
 using .DynamicVehicleScheduling
 using .DynamicAssortment
 using .Maintenance
@@ -104,6 +106,7 @@ export PortfolioOptimizationBenchmark
 export RankingBenchmark
 export StochasticVehicleSchedulingBenchmark
 export SubsetSelectionBenchmark
+export TwoStageSpanningTreeBenchmark
 export WarcraftBenchmark
 export MaintenanceBenchmark
 

src/TwoStageSpanningTree/TwoStageSpanningTree.jl

@@ -0,0 +1,39 @@
+module TwoStageSpanningTree
+
+using DataStructures: IntDisjointSets, in_same_set, union!
+using DocStringExtensions
+using Flux
+using Graphs
+using GLPK
+using HiGHS
+using JuMP
+using JuMP: MOI
+using LinearAlgebra: dot
+using Random
+using Statistics: mean, quantile
+
+using ..Utils
+
+include("utils.jl")
+include("instance.jl")
+include("solution.jl")
+
+include("algorithms/anticipative.jl")
+include("algorithms/cut_generation.jl")
+include("algorithms/benders_decomposition.jl")
+include("algorithms/column_generation.jl")
+include("algorithms/lagrangian_relaxation.jl")
+
+include("learning/features.jl")
+
+include("benchmark.jl")
+
+export TwoStageSpanningTreeBenchmark
+export TwoStageSpanningTreeInstance, nb_scenarios
+export TwoStageSpanningTreeSolution,
+    solution_value, is_feasible, solution_from_first_stage_forest
+export kruskal, anticipative_solution
+export cut_generation,
+    benders_decomposition, column_generation, column_heuristic, lagrangian_relaxation
+
+end

src/TwoStageSpanningTree/algorithms/anticipative.jl

@@ -0,0 +1,16 @@
+"""
+$TYPEDSIGNATURES
+
+Compute an anticipative solution for given scenario.
+"""
+function anticipative_solution(instance::TwoStageSpanningTreeInstance, scenario::Int=1)
+    (; graph, first_stage_costs, second_stage_costs) = instance
+    scenario_second_stage_costs = @view second_stage_costs[:, scenario]
+    weights = min.(first_stage_costs, scenario_second_stage_costs)
+    (; value, tree) = kruskal(graph, weights)
+
+    slice = first_stage_costs .<= scenario_second_stage_costs
+    y = tree[slice]
+    z = tree[.!slice]
+    return (; value, y, z)
+end