rename instance_kwargs to maximizer_kwargs, + other small fixes

Author: BatyLeo

docs/src/tutorials/warcraft_tutorial.jl

@@ -30,8 +30,8 @@ x = sample.x
 θ_true = sample.θ
 # `y` correspond to the optimal shortest path, encoded as a binary matrix:
 y_true = sample.y
-# `context` is not used in this benchmark (no solver kwargs needed), so it is empty:
-isempty(sample.instance_kwargs)
+# `maximizer_kwargs` is not used in this benchmark (no solver kwargs needed), so it is empty:
+isempty(sample.maximizer_kwargs)
 
 # For some benchmarks, we provide the following plotting method [`plot_solution`](@ref) to visualize the data:
 plot_solution(b, sample)

ext/DFLBenchmarksPlotsExt.jl

@@ -22,7 +22,7 @@ function plot_solution(bench::AbstractBenchmark, sample::DataSample, y; kwargs..
     return plot_solution(
         bench,
         DataSample(;
-            sample.instance_kwargs..., x=sample.x, θ=sample.θ, y=y, extra=sample.extra
+            sample.maximizer_kwargs..., x=sample.x, θ=sample.θ, y=y, extra=sample.extra
         );
         kwargs...,
     )

src/ContextualStochasticArgmax/ContextualStochasticArgmax.jl

@@ -5,18 +5,12 @@ using DocStringExtensions: TYPEDEF, TYPEDFIELDS
 using Flux: Dense
 using Random: Random, AbstractRNG, MersenneTwister
 
-function one_hot_argmax(z::AbstractVector{R}; kwargs...) where {R<:Real}
-    e = zeros(R, length(z))
-    e[argmax(z)] = one(R)
-    return e
-end
-
 """
 $TYPEDEF
 
 Minimal contextual stochastic argmax benchmark.
 
-Per instance: `c_base ~ U[0,1]^n` (base utility, part of instance kwargs and base features).
+Per instance: `c_base ~ U[0,1]^n` (base utility, stored in `extra` of the instance sample).
 Per context draw: `x_raw ~ N(0, I_d)` (observable context). Features: `x = [c_base; x_raw]`.
 Per scenario: `ξ = c_base + W * x_raw + noise`, `noise ~ N(0, noise_std² I)`.
 The learner sees `x` and must predict `θ̂` so that `argmax(θ̂)` ≈ `argmax(ξ)`.
@@ -49,29 +43,42 @@ end
 Utils.is_minimization_problem(::ContextualStochasticArgmaxBenchmark) = false
 Utils.generate_maximizer(::ContextualStochasticArgmaxBenchmark) = one_hot_argmax
 
-# c_base: base features (in x) and solver kwarg (in instance_kwargs for generate_scenario)
+"""
+    generate_instance(::ContextualStochasticArgmaxBenchmark, rng)
+
+Draw `c_base ~ U[0,1]^n` and store it in `extra`. No solver kwargs are needed
+(the maximizer is `one_hot_argmax`, which takes no kwargs).
+"""
 function Utils.generate_instance(
     bench::ContextualStochasticArgmaxBenchmark, rng::AbstractRNG; kwargs...
 )
     c_base = rand(rng, Float32, bench.n)
-    return DataSample(; x=c_base, c_base=c_base)
+    return DataSample(; extra=(; c_base))
 end
 
-# Enriches instance_sample: x = [c_base; x_raw], x_raw in extra for generate_scenario
+"""
+    generate_context(::ContextualStochasticArgmaxBenchmark, rng, instance_sample)
+
+Draw `x_raw ~ N(0, I_d)` and return a context sample with:
+- `x = [c_base; x_raw]`: full feature vector seen by the ML model.
+- `extra = (; c_base, x_raw)`: latents spread into [`generate_scenario`](@ref).
+"""
 function Utils.generate_context(
     bench::ContextualStochasticArgmaxBenchmark,
     rng::AbstractRNG,
     instance_sample::DataSample,
 )
+    c_base = instance_sample.c_base
     x_raw = randn(rng, Float32, bench.d)
-    return DataSample(;
-        x=vcat(instance_sample.x, x_raw),
-        instance_sample.instance_kwargs...,
-        extra=(; x_raw),
-    )
+    return DataSample(; x=vcat(c_base, x_raw), extra=(; x_raw, c_base))
 end
 
-# ξ = c_base + W * x_raw + noise  (c_base from instance_kwargs, x_raw from ctx.extra)
+"""
+    generate_scenario(::ContextualStochasticArgmaxBenchmark, rng; c_base, x_raw, kwargs...)
+
+Draw `ξ = c_base + W * x_raw + noise`, `noise ~ N(0, noise_std² I)`.
+`c_base` and `x_raw` are spread from `ctx.extra` by the framework.
+"""
 function Utils.generate_scenario(
     bench::ContextualStochasticArgmaxBenchmark,
     rng::AbstractRNG;

src/DecisionFocusedLearningBenchmarks.jl

@@ -74,7 +74,7 @@ export generate_instance,
     generate_sample, generate_dataset, generate_environments, generate_environment
 export generate_scenario, generate_context
 export generate_baseline_policies
-export SAA
+export SampleAverageApproximation
 export generate_statistical_model
 export generate_maximizer
 export generate_anticipative_solver, generate_parametric_anticipative_solver

src/StochasticVehicleScheduling/StochasticVehicleScheduling.jl

@@ -116,7 +116,7 @@ Returns a [`DataSample`](@ref) with features `x` and `instance` set, but `y=noth
 To obtain labeled samples, pass a `target_policy` to [`generate_dataset`](@ref):
 
 ```julia
-policy = sample -> DataSample(; sample.instance_kwargs..., x=sample.x,
+policy = sample -> DataSample(; sample.maximizer_kwargs..., x=sample.x,
                                 y=column_generation_algorithm(sample.instance))
 dataset = generate_dataset(benchmark, N; target_policy=policy)
 ```

src/StochasticVehicleScheduling/policies.jl

@@ -5,12 +5,12 @@ SAA baseline policy: builds a stochastic instance from all K scenarios and solve
 via column generation.
 Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
 """
-function svs_saa_policy(instance_sample, ctx_sample, scenarios)
-    stochastic_inst = build_stochastic_instance(instance_sample.instance, scenarios)
+function svs_saa_policy(ctx_sample, scenarios)
+    stochastic_inst = build_stochastic_instance(ctx_sample.instance, scenarios)
     y = column_generation_algorithm(stochastic_inst)
     return [
         DataSample(;
-            instance_sample.instance_kwargs...,
+            ctx_sample.maximizer_kwargs...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),
@@ -24,13 +24,11 @@ $TYPEDSIGNATURES
 Deterministic baseline policy: solves the deterministic MIP (ignores scenario delays).
 Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
 """
-function svs_deterministic_policy(
-    instance_sample, ctx_sample, scenarios; model_builder=highs_model
-)
-    y = deterministic_mip(instance_sample.instance; model_builder)
+function svs_deterministic_policy(ctx_sample, scenarios; model_builder=highs_model)
+    y = deterministic_mip(ctx_sample.instance; model_builder)
     return [
         DataSample(;
-            instance_sample.instance_kwargs...,
+            ctx_sample.maximizer_kwargs...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),
@@ -45,12 +43,12 @@ Local search baseline policy: builds a stochastic instance from all K scenarios
 solves via local search heuristic.
 Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
 """
-function svs_local_search_policy(instance_sample, ctx_sample, scenarios)
-    stochastic_inst = build_stochastic_instance(instance_sample.instance, scenarios)
+function svs_local_search_policy(ctx_sample, scenarios)
+    stochastic_inst = build_stochastic_instance(ctx_sample.instance, scenarios)
     y = local_search(stochastic_inst)
     return [
         DataSample(;
-            instance_sample.instance_kwargs...,
+            ctx_sample.maximizer_kwargs...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),
@@ -68,13 +66,11 @@ Returns a single labeled [`DataSample`](@ref) with `extra=(; scenarios)`.
 Prefer this over [`svs_saa_policy`](@ref) when an exact solution is needed; requires
 SCIP (default) or Gurobi.
 """
-function svs_saa_mip_policy(
-    instance_sample, ctx_sample, scenarios; model_builder=scip_model
-)
-    y = compact_linearized_mip(instance_sample.instance, scenarios; model_builder)
+function svs_saa_mip_policy(ctx_sample, scenarios; model_builder=scip_model)
+    y = compact_linearized_mip(ctx_sample.instance, scenarios; model_builder)
     return [
         DataSample(;
-            instance_sample.instance_kwargs...,
+            ctx_sample.maximizer_kwargs...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),
@@ -86,7 +82,7 @@ end
 $TYPEDSIGNATURES
 
 Return the named baseline policies for [`StochasticVehicleSchedulingBenchmark`](@ref).
-Each policy has signature `(instance_sample, ctx_sample, scenarios) -> Vector{DataSample}`.
+Each policy has signature `(ctx_sample, scenarios) -> Vector{DataSample}`.
 """
 function svs_generate_baseline_policies(::StochasticVehicleSchedulingBenchmark)
     return (;

src/Utils/Utils.jl

@@ -22,7 +22,7 @@ include("model_builders.jl")
 
 export DataSample, Policy
 export evaluate_policy!
-export TopKMaximizer
+export TopKMaximizer, one_hot_argmax
 
 export AbstractEnvironment, get_seed, is_terminated, observe, reset!, step!
 
@@ -33,7 +33,7 @@ export generate_instance, generate_sample, generate_dataset
 export generate_statistical_model, generate_maximizer
 export generate_scenario, generate_context
 export generate_environment, generate_environments
-export SAA
+export SampleAverageApproximation
 export generate_baseline_policies
 export generate_anticipative_solver, generate_parametric_anticipative_solver
 

src/Utils/data_sample.jl

@@ -4,8 +4,8 @@ $TYPEDEF
 Data sample data structure.
 Its main purpose is to store datasets generated by the benchmarks.
 It has 3 main (optional) fields: features `x`, cost parameters `θ`, and solution `y`.
-Additionally, it has an `instance_kwargs` field (solver kwargs, spread into the maximizer as
-`maximizer(θ; sample.instance_kwargs...)`) and an `extra` field (non-solver data, never passed
+Additionally, it has an `maximizer_kwargs` field (solver kwargs, spread into the maximizer as
+`maximizer(θ; sample.maximizer_kwargs...)`) and an `extra` field (non-solver data, never passed
 to the maximizer).
 
 The separation prevents silent breakage from accidentally passing non-solver data
@@ -15,8 +15,8 @@ The separation prevents silent breakage from accidentally passing non-solver dat
 $TYPEDFIELDS
 """
 struct DataSample{
-    CTX<:NamedTuple,
-    EX<:NamedTuple,
+    K<:NamedTuple,
+    E<:NamedTuple,
     F<:Union{AbstractArray,Nothing},
     S<:Union{AbstractArray,Nothing},
     C<:Union{AbstractArray,Nothing},
@@ -28,63 +28,63 @@ struct DataSample{
     "output solution (optional)"
     y::S
     "solver kwargs, e.g. instance, graph, etc."
-    instance_kwargs::CTX
+    maximizer_kwargs::K
     "additional data, never passed to the maximizer, e.g. scenario, objective value, reward,
     step count, etc. Can be used for any purpose by the user, such as plotting utilities."
-    extra::EX
+    extra::E
 end
 
 """
 $TYPEDSIGNATURES
 
 Constructor for `DataSample` with keyword arguments.
 
-All keyword arguments beyond `x`, `θ`, `y`, and `extra` are collected into the `instance_kwargs`
+All keyword arguments beyond `x`, `θ`, `y`, and `extra` are collected into the `maximizer_kwargs`
 field (solver kwargs). The `extra` keyword accepts a `NamedTuple` of non-solver data.
 
-Fields in `instance_kwargs` and `extra` must be disjoint. An error is thrown if they overlap.
+Fields in `maximizer_kwargs` and `extra` must be disjoint. An error is thrown if they overlap.
 Both can be accessed directly via property forwarding.
 
 # Examples
 ```julia
-# Instance goes in instance_kwargs
+# Instance goes in maximizer_kwargs
 d = DataSample(x=[1,2,3], θ=[4,5,6], y=[7,8,9], instance="my_instance")
-d.instance  # "my_instance" (from instance_kwargs)
+d.instance  # "my_instance" (from maximizer_kwargs)
 
 # Scenario goes in extra
 d = DataSample(x=x, y=y, instance=inst, extra=(; scenario=ξ))
 d.scenario  # ξ (from extra)
 
-# State goes in instance_kwargs, reward in extra
+# State goes in maximizer_kwargs, reward in extra
 d = DataSample(x=x, y=y, instance=state, extra=(; reward=-1.5))
-d.instance  # state (from instance_kwargs)
+d.instance  # state (from maximizer_kwargs)
 d.reward    # -1.5 (from extra)
 ```
 """
 function DataSample(; x=nothing, θ=nothing, y=nothing, extra=NamedTuple(), kwargs...)
-    instance_kwargs = (; kwargs...)
-    overlap = intersect(keys(instance_kwargs), keys(extra))
+    maximizer_kwargs = (; kwargs...)
+    overlap = intersect(keys(maximizer_kwargs), keys(extra))
     if !isempty(overlap)
         error(
-            "Keys $(collect(overlap)) appear in both instance_kwargs and extra of DataSample",
+            "Keys $(collect(overlap)) appear in both maximizer_kwargs and extra of DataSample",
         )
     end
-    return DataSample(x, θ, y, instance_kwargs, extra)
+    return DataSample(x, θ, y, maximizer_kwargs, extra)
 end
 
 """
 $TYPEDSIGNATURES
 
 Extended property access for `DataSample`.
 
-Allows accessing `instance_kwargs` and `extra` fields directly as properties.
-`instance_kwargs` is searched first; if the key is not found there, `extra` is searched.
+Allows accessing `maximizer_kwargs` and `extra` fields directly as properties.
+`maximizer_kwargs` is searched first; if the key is not found there, `extra` is searched.
 """
 function Base.getproperty(d::DataSample, name::Symbol)
-    if name in (:x, :θ, :y, :instance_kwargs, :extra)
+    if name in (:x, :θ, :y, :maximizer_kwargs, :extra)
         return getfield(d, name)
     else
-        ctx = getfield(d, :instance_kwargs)
+        ctx = getfield(d, :maximizer_kwargs)
         if haskey(ctx, name)
             return getproperty(ctx, name)
         end
@@ -96,12 +96,12 @@ end
 $TYPEDSIGNATURES
 
 Return all property names of a `DataSample`, including both struct fields and forwarded
-fields from `instance_kwargs` and `extra`.
+fields from `maximizer_kwargs` and `extra`.
 
 This enables tab completion for all available properties.
 """
 function Base.propertynames(d::DataSample, private::Bool=false)
-    ctx_names = propertynames(getfield(d, :instance_kwargs), private)
+    ctx_names = propertynames(getfield(d, :maximizer_kwargs), private)
     extra_names = propertynames(getfield(d, :extra), private)
     return (fieldnames(DataSample)..., ctx_names..., extra_names...)
 end
@@ -128,7 +128,7 @@ function Base.show(io::IO, d::DataSample)
         y_str = sprint(show, d.y; context=io_limited)
         push!(fields, "y_true=$y_str")
     end
-    for (key, value) in pairs(d.instance_kwargs)
+    for (key, value) in pairs(d.maximizer_kwargs)
         value_str = sprint(show, value; context=io_limited)
         push!(fields, "$key=$value_str")
     end
@@ -156,8 +156,8 @@ Transform the features in the dataset.
 """
 function StatsBase.transform(t, dataset::AbstractVector{<:DataSample})
     return map(dataset) do d
-        (; instance_kwargs, extra, x, θ, y) = d
-        DataSample(StatsBase.transform(t, x), θ, y, instance_kwargs, extra)
+        (; maximizer_kwargs, extra, x, θ, y) = d
+        DataSample(StatsBase.transform(t, x), θ, y, maximizer_kwargs, extra)
     end
 end
 
@@ -179,8 +179,8 @@ Reconstruct the features in the dataset.
 """
 function StatsBase.reconstruct(t, dataset::AbstractVector{<:DataSample})
     return map(dataset) do d
-        (; instance_kwargs, extra, x, θ, y) = d
-        DataSample(StatsBase.reconstruct(t, x), θ, y, instance_kwargs, extra)
+        (; maximizer_kwargs, extra, x, θ, y) = d
+        DataSample(StatsBase.reconstruct(t, x), θ, y, maximizer_kwargs, extra)
     end
 end