fix doc

Author: BatyLeo

docs/src/tutorials/warcraft.jl

@@ -17,23 +17,23 @@ b = WarcraftBenchmark()
 # This method takes as input the benchmark object for which the dataset is to be generated, and a second argument specifying the number of samples to generate:
 dataset = generate_dataset(b, 50)
 
-# We obtain an [`InferOptDataset`](@ref) object, which contains all needed data for the problem.
+# We obtain a vector of [`DataSample`](@ref) object, which contains all needed data for the problem.
 # Subdatasets can be created through regular slicing:
 train_dataset, test_dataset = dataset[1:45], dataset[46:50]
 
 # And getting an individual sample will return a NamedTuple with four keys: `features`, `instance`, `costs`, and `solution`:
 sample = test_dataset[1]
 # In the case the the Warcraft benchmark, `features` correspond to the input image:
-x = sample.features
+x = sample.x
 # `costs` correspond to the true unknown terrain weights:
-θ_true = sample.costs
+θ_true = sample.θ
 # `solution` correspond to the true shortest path:
-y_true = sample.solution
+y_true = sample.y
 # `instance` is not used in this benchmark, therefore set to nothing:
 sample.instance
 
 # For some benchmarks, we provide the following plotting method [`plot_data`](@ref) to visualize the data:
-plot_data(b; sample...)
+plot_data(b, sample)
 # We can see here the terrain image, the true terrain weights, and the true shortest path avoiding the high cost cells.
 
 # ## Building a pipeline
@@ -50,7 +50,7 @@ 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(θ)
 # As we can see, currently the pipeline predicts random noise as cell weights, and therefore the maximizer returns a straight line path.
-plot_data(b; features=x, costs=θ, solution=y)
+plot_data(b, DataSample(; x, θ, y))
 # We can evaluate the current pipeline performance using the optimality gap metric:
 starting_gap = compute_gap(b, test_dataset, model, maximizer)
 
@@ -61,9 +61,6 @@ using InferOpt
 using Flux
 using Plots
 
-X_train = train_dataset.features
-Y_train = train_dataset.solutions
-
 perturbed_maximizer = PerturbedMultiplicative(maximizer; ε=0.2, nb_samples=100)
 loss = FenchelYoungLoss(perturbed_maximizer)
 
@@ -73,7 +70,7 @@ opt_state = Flux.setup(Adam(1e-3), model)
 loss_history = Float64[]
 for epoch in 1:50
     val, grads = Flux.withgradient(model) do m
-        sum(loss(m(x), y) for (x, y) in zip(X_train, Y_train)) / length(train_dataset)
+        sum(loss(m(sample.x), sample.y) for sample in train_dataset) / length(train_dataset)
     end
     Flux.update!(opt_state, model, grads[1])
     push!(loss_history, val)
@@ -86,4 +83,6 @@ plot(loss_history; xlabel="Epoch", ylabel="Loss", title="Training loss")
 final_gap = compute_gap(b, test_dataset, model, maximizer)
 
 #
-plot_data(b; sample...)
+θ = model(x)
+y = maximizer(θ)
+plot_data(b, DataSample(; x, θ, y))

src/Utils/interface.jl

@@ -14,10 +14,10 @@ The following methods exist for benchmarks:
 abstract type AbstractBenchmark end
 
 """
-    generate_dataset(::AbstractBenchmark, dataset_size::Int) -> InferOptDataset
+    generate_dataset(::AbstractBenchmark, dataset_size::Int) -> Vector{<:DataSample}
 
-Generate an [`InferOptDataset`](@ref) for given benchmark as a Vector of length `dataset_size`.
-Content of the dataset can be visualized using [`plot_data`](@ref).
+Generate a `Vector` of [`DataSample`](@ref)  of length `dataset_size` for given benchmark.
+Content of the dataset can be visualized using [`plot_data`](@ref), when it applies.
 """
 function generate_dataset end
 
@@ -46,7 +46,7 @@ Check the specific benchmark documentation of `plot_data` for more details on th
 function plot_data end
 
 """
-    compute_gap(::AbstractBenchmark, dataset::InferOptDataset, statistical_model, maximizer) -> Float64
+    compute_gap(::AbstractBenchmark, dataset::Vector{<:DataSample}, statistical_model, maximizer) -> Float64
 
 Compute the average relative optimality gap of the pipeline on the dataset.
 """

src/Warcraft/Warcraft.jl

@@ -28,18 +28,14 @@ $TYPEDSIGNATURES
 Plot the image `im`, the weights `weights`, and the path `path` on the same Figure.
 """
 function Utils.plot_data(
-    ::WarcraftBenchmark;
-    features,
-    solution,
-    costs,
+    ::WarcraftBenchmark,
+    sample::DataSample,
     θ_title="Weights",
     y_title="Path",
-    θ_true=costs,
+    θ_true=sample.θ,
     kwargs...,
 )
-    x = features
-    y = solution
-    θ = costs
+    (; x, y, θ) = sample
     im = dropdims(x; dims=4)
     img = convert_image_for_plot(im)
     p1 = Plots.plot(