replace InferOptDataset by Vector{DataSample}

Author: BatyLeo

docs/src/tutorials/warcraft.jl

@@ -45,7 +45,7 @@ model = generate_statistical_model(b)
 θ = model(x)
 # Note that the model is not trained yet, and its parameters are randomly initialized.
 
-# Finally, the [`generate_maximizer`](@ref) method can be used generates a combinatorial optimization algorithm that takes the predicted cell weights as input and returns the corresponding shortest path:
+# Finally, the [`generate_maximizer`](@ref) method can be used to generate a combinatorial optimization algorithm that takes the predicted cell weights as input and returns the corresponding shortest path:
 maximizer = generate_maximizer(b; dijkstra=true)
 # In the case o fthe Warcraft benchmark, the method has an additioonal keyword argument to chose the algorithm to use: Dijkstra's algorithm or Bellman-Ford algorithm.
 y = maximizer(θ)

src/FixedSizeShortestPath/shortest_paths.jl

@@ -121,7 +121,7 @@ function Utils.generate_dataset(
 
     # Label solutions
     solutions = shortest_path_maximizer.(.-costs)
-    return InferOptDataset(; features, costs, solutions)
+    return [DataSample(; x=x, θ=θ, y=y) for (x, θ, y) in zip(features, costs, solutions)]
 end
 
 """

src/InferOptBenchmarks.jl

@@ -31,7 +31,7 @@ using .PortfolioOptimization
 using .SubsetSelection
 
 # Interface
-export AbstractBenchmark, InferOptDataset
+export AbstractBenchmark, DataSample
 export generate_dataset
 export generate_statistical_model
 export generate_maximizer

src/PortfolioOptimization/portfolio_optimization.jl

@@ -99,7 +99,7 @@ function Utils.generate_dataset(
     maximizer = Utils.generate_maximizer(bench)
     solutions = maximizer.(costs)
 
-    return InferOptDataset(; features, costs, solutions)
+    return [DataSample(; x=x, θ=θ, y=y) for (x, θ, y) in zip(features, costs, solutions)]
 end
 
 """

src/ShortestPath/shortest_paths.jl

@@ -116,7 +116,7 @@ function Utils.generate_dataset(
 
     # Label solutions
     solutions = shortest_path_maximizer.(.-costs)
-    return InferOptDataset(; features, costs, solutions)
+    return [DataSample(; x=x, θ=θ, y=y) for (x, θ, y) in zip(features, costs, solutions)]
 end
 
 """

src/SubsetSelection/subset_selection.jl

@@ -58,7 +58,7 @@ function Utils.generate_dataset(
     features = [randn(rng, Float32, n) for _ in 1:dataset_size]
     costs = copy(features)  # we assume that the cost is the same as the feature
     solutions = top_k.(features, k)
-    return InferOptDataset(; features, solutions, costs)
+    return [DataSample(; x=x, θ=θ, y=y) for (x, θ, y) in zip(features, costs, solutions)]
 end
 
 """

src/Utils/Utils.jl

@@ -5,12 +5,12 @@ using Flux: softplus
 using LinearAlgebra: dot
 using SimpleWeightedGraphs: SimpleWeightedDiGraph
 
-include("dataset.jl")
+include("data_sample.jl")
 include("interface.jl")
 include("grid_graph.jl")
 include("misc.jl")
 
-export InferOptDataset
+export DataSample
 
 export AbstractBenchmark
 export generate_dataset,

src/Utils/data_sample.jl

@@ -0,0 +1,18 @@
+"""
+$TYPEDEF
+
+Data sample data structure.
+
+# Fields
+$TYPEDFIELDS
+"""
+@kwdef struct DataSample{F,S,C,I}
+    "features"
+    x::F
+    "costs"
+    θ::C = nothing
+    "solution"
+    y::S = nothing
+    "instance"
+    instance::I = nothing
+end

src/Utils/dataset.jl

@@ -67,17 +67,14 @@ $TYPEDSIGNATURES
 Default behaviour of `compute_gap` for a benchmark problem where `features`, `solutions` and `costs` are all defined.
 """
 function compute_gap(
-    bench::AbstractBenchmark,
-    dataset::InferOptDataset{<:AbstractArray,<:AbstractArray,<:AbstractArray},
-    statistical_model,
-    maximizer,
+    bench::AbstractBenchmark, dataset::Vector{<:DataSample}, statistical_model, maximizer
 )
     res = 0.0
-    X = dataset.features
-    costs = dataset.costs
-    Y = dataset.solutions
 
-    for (x, θ̄, ȳ) in zip(X, costs, Y)
+    for sample in dataset
+        x = sample.x
+        θ̄ = sample.θ
+        ȳ = sample.y
         θ = statistical_model(x)
         y = maximizer(θ)
         target_obj = objective_value(bench, θ̄, ȳ)