Enforce info field to be a NamedTuple

Author: BatyLeo

src/Argmax2D/Argmax2D.jl

@@ -62,7 +62,7 @@ function Utils.generate_sample(bench::Argmax2DBenchmark, rng::AbstractRNG)
     θ_true ./= 2 * norm(θ_true)
     instance = build_polytope(rand(rng, polytope_vertex_range); shift=rand(rng))
     y_true = maximizer(θ_true; instance)
-    return DataSample(; x=x, θ=θ_true, y=y_true, info=instance)
+    return DataSample(; x=x, θ=θ_true, y=y_true, instance=instance)
 end
 
 """
@@ -88,11 +88,11 @@ function Utils.generate_statistical_model(
     return model
 end
 
-function Utils.plot_data(::Argmax2DBenchmark; info, θ, kwargs...)
+function Utils.plot_data(::Argmax2DBenchmark; instance, θ, kwargs...)
     pl = init_plot()
-    plot_polytope!(pl, info)
+    plot_polytope!(pl, instance)
     plot_objective!(pl, θ)
-    return plot_maximizer!(pl, θ, info, maximizer)
+    return plot_maximizer!(pl, θ, instance, maximizer)
 end
 
 """
@@ -101,9 +101,13 @@ $TYPEDSIGNATURES
 Plot the data sample for the [`Argmax2DBenchmark`](@ref).
 """
 function Utils.plot_data(
-    bench::Argmax2DBenchmark, sample::DataSample; info=sample.info, θ=sample.θ, kwargs...
+    bench::Argmax2DBenchmark,
+    sample::DataSample;
+    instance=sample.instance,
+    θ=sample.θ,
+    kwargs...,
 )
-    return Utils.plot_data(bench; info, θ, kwargs...)
+    return Utils.plot_data(bench; instance, θ, kwargs...)
 end
 
 export Argmax2DBenchmark

src/DecisionFocusedLearningBenchmarks.jl

@@ -70,7 +70,7 @@ export generate_sample, generate_dataset, generate_environments, generate_enviro
 export generate_scenario
 export generate_policies
 export generate_statistical_model
-export generate_maximizer, maximizer_kwargs
+export generate_maximizer
 export generate_anticipative_solution
 export is_exogenous, is_endogenous
 

src/DynamicAssortment/DynamicAssortment.jl

@@ -83,7 +83,7 @@ Outputs a data sample containing an [`Instance`](@ref).
 function Utils.generate_sample(
     b::DynamicAssortmentBenchmark, rng::AbstractRNG=MersenneTwister(0)
 )
-    return DataSample(; info=Instance(b, rng))
+    return DataSample(; instance=Instance(b, rng))
 end
 
 """

src/DynamicVehicleScheduling/DynamicVehicleScheduling.jl

@@ -70,7 +70,7 @@ function Utils.generate_dataset(b::DynamicVehicleSchedulingBenchmark, dataset_si
     dataset_size = min(dataset_size, length(files))
     return [
         DataSample(;
-            info=Instance(
+            instance=Instance(
                 read_vsp_instance(files[i]);
                 max_requests_per_epoch,
                 Δ_dispatch,

src/DynamicVehicleScheduling/anticipative_solver.jl

@@ -222,7 +222,7 @@ function anticipative_solver(
             compute_features(state, env.instance)
         end
 
-        return DataSample(; info=(; state, reward), y=y_true, x)
+        return DataSample(; y=y_true, x, state, reward)
     end
 
     return obj, dataset

src/StochasticVehicleScheduling/StochasticVehicleScheduling.jl

@@ -67,7 +67,7 @@ end
 function Utils.objective_value(
     ::StochasticVehicleSchedulingBenchmark, sample::DataSample, y::BitVector
 )
-    return evaluate_solution(y, sample.info)
+    return evaluate_solution(y, sample.instance)
 end
 
 """
@@ -98,7 +98,7 @@ function Utils.generate_sample(
     else
         nothing
     end
-    return DataSample(; x, info=instance, y=y_true)
+    return DataSample(; x, instance, y=y_true)
 end
 
 """
@@ -145,11 +145,12 @@ end
 $TYPEDSIGNATURES
 """
 function plot_instance(
-    ::StochasticVehicleSchedulingBenchmark, sample::DataSample{<:Instance{City}}; kwargs...
+    ::StochasticVehicleSchedulingBenchmark, sample::DataSample; kwargs...
 )
-    (; tasks, district_width, width) = sample.info.city
+    @assert hasproperty(sample.instance, :city) "Sample does not contain city information."
+    (; tasks, district_width, width) = sample.instance.city
     ticks = 0:district_width:width
-    max_time = maximum(t.end_time for t in sample.info.city.tasks[1:(end - 1)])
+    max_time = maximum(t.end_time for t in sample.instance.city.tasks[1:(end - 1)])
     fig = plot(;
         xlabel="x",
         ylabel="y",
@@ -204,11 +205,12 @@ end
 $TYPEDSIGNATURES
 """
 function plot_solution(
-    ::StochasticVehicleSchedulingBenchmark, sample::DataSample{<:Instance{City}}; kwargs...
+    ::StochasticVehicleSchedulingBenchmark, sample::DataSample; kwargs...
 )
-    (; tasks, district_width, width) = sample.info.city
+    @assert hasproperty(sample.instance, :city) "Sample does not contain city information."
+    (; tasks, district_width, width) = sample.instance.city
     ticks = 0:district_width:width
-    solution = Solution(sample.y, sample.info)
+    solution = Solution(sample.y, sample.instance)
     path_list = compute_path_list(solution)
     fig = plot(;
         xlabel="x",

src/Utils/Utils.jl

@@ -35,7 +35,6 @@ export generate_policies
 export generate_anticipative_solution
 
 export plot_data, compute_gap
-export maximizer_kwargs
 export grid_graph, get_path, path_to_matrix
 export neg_tensor, squeeze_last_dims, average_tensor
 export scip_model, highs_model

src/Utils/data_sample.jl

@@ -2,39 +2,99 @@
 $TYPEDEF
 
 Data sample data structure.
+Its main purpose is to store datasets generated by the benchmarks.
+It has 3 main fields: features `x`, cost parameters `θ` and solution `y`.
+Additionally, it has an `info` field to store any additional information as a `NamedTuple`, usually the instance, but can be used for anything else.
 
 # Fields
 $TYPEDFIELDS
 """
-@kwdef struct DataSample{
-    I,
+struct DataSample{
+    I<:NamedTuple,
     F<:Union{AbstractArray,Nothing},
     S<:Union{AbstractArray,Nothing},
     C<:Union{AbstractArray,Nothing},
 }
     "input features (optional)"
-    x::F = nothing
+    x::F
     "intermediate cost parameters (optional)"
-    θ::C = nothing
+    θ::C
     "output solution (optional)"
-    y::S = nothing
+    y::S
     "additional information, usually the instance (optional)"
-    info::I = nothing
+    info::I
 end
 
+"""
+$TYPEDSIGNATURES
+
+Constructor for `DataSample` with keyword arguments.
+
+Additional keyword arguments beyond `x`, `θ`, and `y` are stored in the `info` field
+and can be accessed directly (e.g., `data.instance` instead of `data.info.instance`).
+
+# Examples
+```julia
+d = DataSample(x=[1,2,3], θ=[4,5,6], y=[7,8,9], instance="my_instance")
+d.instance  # "my_instance"
+```
+"""
+function DataSample(; x=nothing, θ=nothing, y=nothing, kwargs...)
+    info = (; kwargs...)
+    return DataSample(x, θ, y, info)
+end
+
+"""
+$TYPEDSIGNATURES
+
+Extended property access for `DataSample`.
+
+Allows accessing `info` fields directly as properties (e.g., `d.instance` instead of `d.info.instance`).
+"""
+function Base.getproperty(d::DataSample, name::Symbol)
+    if name in (:x, :θ, :y, :info)
+        return getfield(d, name)
+    else
+        return getproperty(getfield(d, :info), name)
+    end
+end
+
+"""
+$TYPEDSIGNATURES
+
+Return all property names of a `DataSample`, including both struct fields and `info` fields.
+
+This enables tab completion for all available properties, including those stored in `info`.
+"""
+function Base.propertynames(d::DataSample, private::Bool=false)
+    return (fieldnames(DataSample)..., propertynames(getfield(d, :info), private)...)
+end
+
+"""
+$TYPEDSIGNATURES
+
+Display a `DataSample` with truncated array representations for readability.
+
+Large arrays are automatically truncated with ellipsis (`...`), similar to standard Julia array printing.
+"""
 function Base.show(io::IO, d::DataSample)
     fields = String[]
+    io_limited = IOContext(io, :limit => true, :compact => true)
     if !isnothing(d.x)
-        push!(fields, "x=$(d.x)")
+        x_str = sprint(show, d.x; context=io_limited)
+        push!(fields, "x=$x_str")
     end
     if !isnothing(d.θ)
-        push!(fields, "θ_true=$(d.θ)")
+        θ_str = sprint(show, d.θ; context=io_limited)
+        push!(fields, "θ_true=$θ_str")
     end
     if !isnothing(d.y)
-        push!(fields, "y_true=$(d.y)")
+        y_str = sprint(show, d.y; context=io_limited)
+        push!(fields, "y_true=$y_str")
     end
-    if !isnothing(d.info)
-        push!(fields, "instance=$(d.info)")
+    for (key, value) in pairs(d.info)
+        value_str = sprint(show, value; context=io_limited)
+        push!(fields, "$key=$value_str")
     end
     return print(io, "DataSample(", join(fields, ", "), ")")
 end
@@ -57,7 +117,7 @@ Transform the features in the dataset.
 function StatsBase.transform(t, dataset::AbstractVector{<:DataSample})
     return map(dataset) do d
         (; info, x, θ, y) = d
-        DataSample(; info, x=StatsBase.transform(t, x), θ, y)
+        DataSample(StatsBase.transform(t, x), θ, y, info)
     end
 end
 
@@ -80,7 +140,7 @@ Reconstruct the features in the dataset.
 function StatsBase.reconstruct(t, dataset::AbstractVector{<:DataSample})
     return map(dataset) do d
         (; info, x, θ, y) = d
-        DataSample(; info, x=StatsBase.reconstruct(t, x), θ, y)
+        DataSample(StatsBase.reconstruct(t, x), θ, y, info)
     end
 end
 

src/Utils/interface.jl

@@ -99,26 +99,6 @@ function compute_gap end
 """
 $TYPEDSIGNATURES
 
-For simple benchmarks where there is no instance object, maximizer does not need any keyword arguments.
-"""
-function maximizer_kwargs(
-    ::AbstractBenchmark, sample::DataSample{Nothing,F,S,C}
-) where {F,S,C}
-    return NamedTuple()
-end
-
-"""
-$TYPEDSIGNATURES
-
-For benchmarks where there is an instance object, maximizer needs the instance object as a keyword argument.
-"""
-function maximizer_kwargs(::AbstractBenchmark, sample::DataSample)
-    return (; instance=sample.info)
-end
-
-"""
-$TYPEDSIGNATURES
-
 Default behaviour of `objective_value`.
 """
 function objective_value(::AbstractBenchmark, θ::AbstractArray, y::AbstractArray)
@@ -175,7 +155,7 @@ function compute_gap(
             target_obj = objective_value(bench, sample)
             x = sample.x
             θ = statistical_model(x)
-            y = maximizer(θ; maximizer_kwargs(bench, sample)...)
+            y = maximizer(θ; sample.info...)
             obj = objective_value(bench, sample, y)
             Δ = check ? obj - target_obj : target_obj - obj
             return Δ / abs(target_obj)
@@ -234,7 +214,7 @@ Uses the info field of the sample as the instance.
 function generate_environment(
     bench::AbstractDynamicBenchmark, sample::DataSample, rng::AbstractRNG; kwargs...
 )
-    return generate_environment(bench, sample.info, rng; kwargs...)
+    return generate_environment(bench, sample.instance, rng; kwargs...)
 end
 
 """
@@ -250,7 +230,7 @@ function generate_environments(
     kwargs...,
 )
     Random.seed!(rng, seed)
-    return map(dataset) do instance
-        generate_environment(bench, instance, rng; kwargs...)
+    return map(dataset) do sample
+        generate_environment(bench, sample, rng; kwargs...)
     end
 end

src/Utils/policy.jl

@@ -44,7 +44,7 @@ function evaluate_policy!(
         features, state = observe(env)
         state_copy = deepcopy(state)  # To avoid mutation issues
         reward = step!(env, y)
-        sample = DataSample(; x=features, y=y, info=(; state=state_copy, reward))
+        sample = DataSample(; x=features, y=y, state=state_copy, reward=reward)
         if @isdefined labeled_dataset
             push!(labeled_dataset, sample)
         else