rename maximizer_kwargs to context

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
-# `maximizer_kwargs` is not used in this benchmark (no solver kwargs needed), so it is empty:
-isempty(sample.maximizer_kwargs)
+# `context` is not used in this benchmark (no solver kwargs needed), so it is empty:
+isempty(sample.context)
 
 # For some benchmarks, we provide the following plotting method [`plot_solution`](@ref) to visualize the data:
 plot_solution(b, sample)

ext/DFLBenchmarksPlotsExt.jl

@@ -21,9 +21,7 @@ Reconstruct a new sample with `y` overridden and delegate to the 2-arg
 function plot_solution(bench::AbstractBenchmark, sample::DataSample, y; kwargs...)
     return plot_solution(
         bench,
-        DataSample(;
-            sample.maximizer_kwargs..., x=sample.x, θ=sample.θ, y=y, extra=sample.extra
-        );
+        DataSample(; sample.context..., x=sample.x, θ=sample.θ, y=y, extra=sample.extra);
         kwargs...,
     )
 end

src/ContextualStochasticArgmax/ContextualStochasticArgmax.jl

@@ -12,7 +12,7 @@ $TYPEDEF
 
 Minimal contextual stochastic argmax benchmark.
 
-Per instance: `c_base ~ U[0,1]^n` (base utility, stored in `extra` of the instance sample).
+Per instance: `c_base ~ U[0,1]^n` (base utility, stored in `context` 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(ξ)`.
@@ -55,34 +55,32 @@ function Utils.objective_value(
     bench::ContextualStochasticArgmaxBenchmark, sample::DataSample, y
 )
     if hasproperty(sample.extra, :scenario)
-        return Utils.objective_value(bench, sample, y, sample.extra.scenario)
+        return Utils.objective_value(bench, sample, y, sample.scenario)
     elseif hasproperty(sample.extra, :scenarios)
-        return mean(
-            Utils.objective_value(bench, sample, y, ξ) for ξ in sample.extra.scenarios
-        )
+        return mean(Utils.objective_value(bench, sample, y, ξ) for ξ in sample.scenarios)
     end
     return error("Sample must have scenario or scenarios")
 end
 
 """
     generate_instance(::ContextualStochasticArgmaxBenchmark, rng)
 
-Draw `c_base ~ U[0,1]^n` and store it in `extra`. No solver kwargs are needed
+Draw `c_base ~ U[0,1]^n` and store it in `context`. 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(; extra=(; c_base))
+    return DataSample(; c_base)
 end
 
 """
     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).
+- `c_base`, `x_raw` in `context`: spread into [`generate_scenario`](@ref).
 """
 function Utils.generate_context(
     bench::ContextualStochasticArgmaxBenchmark,
@@ -91,14 +89,14 @@ function Utils.generate_context(
 )
     c_base = instance_sample.c_base
     x_raw = randn(rng, Float32, bench.d)
-    return DataSample(; x=vcat(c_base, x_raw), extra=(; x_raw, c_base))
+    return DataSample(; x=vcat(c_base, x_raw), c_base, x_raw)
 end
 
 """
     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.
+`c_base` and `x_raw` are spread from `ctx.context` by the framework.
 """
 function Utils.generate_scenario(
     bench::ContextualStochasticArgmaxBenchmark,

src/ContextualStochasticArgmax/policies.jl

@@ -11,7 +11,7 @@ function csa_saa_policy(ctx_sample, scenarios)
     y = one_hot_argmax(mean(scenarios))
     return [
         DataSample(;
-            ctx_sample.maximizer_kwargs...,
+            ctx_sample.context...,
             x=ctx_sample.x,
             y=y,
             extra=(; ctx_sample.extra..., scenarios),

src/StochasticVehicleScheduling/StochasticVehicleScheduling.jl

@@ -131,7 +131,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.maximizer_kwargs..., x=sample.x,
+policy = sample -> DataSample(; sample.context..., x=sample.x,
                                 y=column_generation_algorithm(sample.instance))
 dataset = generate_dataset(benchmark, N; target_policy=policy)
 ```

src/StochasticVehicleScheduling/policies.jl

@@ -10,7 +10,7 @@ function svs_saa_policy(ctx_sample, scenarios)
     y = column_generation_algorithm(stochastic_inst)
     return [
         DataSample(;
-            ctx_sample.maximizer_kwargs...,
+            ctx_sample.context...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),
@@ -28,7 +28,7 @@ function svs_deterministic_policy(ctx_sample, scenarios; model_builder=highs_mod
     y = deterministic_mip(ctx_sample.instance; model_builder)
     return [
         DataSample(;
-            ctx_sample.maximizer_kwargs...,
+            ctx_sample.context...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),
@@ -48,7 +48,7 @@ function svs_local_search_policy(ctx_sample, scenarios)
     y = local_search(stochastic_inst)
     return [
         DataSample(;
-            ctx_sample.maximizer_kwargs...,
+            ctx_sample.context...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),
@@ -70,7 +70,7 @@ function svs_saa_mip_policy(ctx_sample, scenarios; model_builder=scip_model)
     y = compact_linearized_mip(ctx_sample.instance, scenarios; model_builder)
     return [
         DataSample(;
-            ctx_sample.maximizer_kwargs...,
+            ctx_sample.context...,
             x=ctx_sample.x,
             y,
             extra=(; ctx_sample.extra..., scenarios),

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 a `maximizer_kwargs` field (solver kwargs, spread into the maximizer as
-`maximizer(θ; sample.maximizer_kwargs...)`) and an `extra` field (non-solver data, never passed
+Additionally, it has a `context` field (solver and scenario-generation context, spread into the
+maximizer as `maximizer(θ; sample.context...)`) and an `extra` field (non-solver data, never passed
 to the maximizer).
 
 The separation prevents silent breakage from accidentally passing non-solver data
@@ -27,8 +27,8 @@ struct DataSample{
     θ::C
     "output solution (optional)"
     y::S
-    "solver kwargs, e.g. instance, graph, etc."
-    maximizer_kwargs::K
+    "solver and scenario-generation context, e.g. instance, graph, contextual information"
+    context::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::E
@@ -39,65 +39,61 @@ $TYPEDSIGNATURES
 
 Constructor for `DataSample` with keyword arguments.
 
-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.
+All keyword arguments beyond `x`, `θ`, `y`, and `extra` are collected into the `context`
+field (solver and scenario-generation context). The `extra` keyword accepts a `NamedTuple` of non-solver data.
 
-Fields in `maximizer_kwargs` and `extra` must be disjoint. Neither may use a reserved
-struct field name (`x`, `θ`, `y`, `maximizer_kwargs`, `extra`). An error is thrown in
+Fields in `context` and `extra` must be disjoint. Neither may use a reserved
+struct field name (`x`, `θ`, `y`, `context`, `extra`). An error is thrown in
 both cases.
 Both can be accessed directly via property forwarding.
 
 # Examples
 ```julia
-# Instance goes in maximizer_kwargs
+# Instance goes in context
 d = DataSample(x=[1,2,3], θ=[4,5,6], y=[7,8,9], instance="my_instance")
-d.instance  # "my_instance" (from maximizer_kwargs)
+d.instance  # "my_instance" (from context)
 
 # Scenario goes in extra
 d = DataSample(x=x, y=y, instance=inst, extra=(; scenario=ξ))
 d.scenario  # ξ (from extra)
 
-# State goes in maximizer_kwargs, reward in extra
+# State goes in context, reward in extra
 d = DataSample(x=x, y=y, instance=state, extra=(; reward=-1.5))
-d.instance  # state (from maximizer_kwargs)
+d.instance  # state (from context)
 d.reward    # -1.5 (from extra)
 ```
 """
 function DataSample(; x=nothing, θ=nothing, y=nothing, extra=NamedTuple(), kwargs...)
-    maximizer_kwargs = (; kwargs...)
-    overlap = intersect(keys(maximizer_kwargs), keys(extra))
+    context = (; kwargs...)
+    overlap = intersect(keys(context), keys(extra))
     if !isempty(overlap)
-        error(
-            "Keys $(collect(overlap)) appear in both maximizer_kwargs and extra of DataSample",
-        )
+        error("Keys $(collect(overlap)) appear in both context and extra of DataSample")
     end
-    reserved = (:x, :θ, :y, :maximizer_kwargs, :extra)
-    shadowed_ctx = intersect(keys(maximizer_kwargs), reserved)
+    reserved = (:x, :θ, :y, :context, :extra)
+    shadowed_ctx = intersect(keys(context), reserved)
     if !isempty(shadowed_ctx)
-        error(
-            "Keys $(collect(shadowed_ctx)) in maximizer_kwargs shadow DataSample struct fields",
-        )
+        error("Keys $(collect(shadowed_ctx)) in context shadow DataSample struct fields")
     end
     shadowed_extra = intersect(keys(extra), reserved)
     if !isempty(shadowed_extra)
         error("Keys $(collect(shadowed_extra)) in extra shadow DataSample struct fields")
     end
-    return DataSample(x, θ, y, maximizer_kwargs, extra)
+    return DataSample(x, θ, y, context, extra)
 end
 
 """
 $TYPEDSIGNATURES
 
 Extended property access for `DataSample`.
 
-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.
+Allows accessing `context` and `extra` fields directly as properties.
+`context` 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, :maximizer_kwargs, :extra)
+    if name in (:x, :θ, :y, :context, :extra)
         return getfield(d, name)
     else
-        ctx = getfield(d, :maximizer_kwargs)
+        ctx = getfield(d, :context)
         if haskey(ctx, name)
             return getproperty(ctx, name)
         end
@@ -109,12 +105,12 @@ end
 $TYPEDSIGNATURES
 
 Return all property names of a `DataSample`, including both struct fields and forwarded
-fields from `maximizer_kwargs` and `extra`.
+fields from `context` and `extra`.
 
 This enables tab completion for all available properties.
 """
 function Base.propertynames(d::DataSample, private::Bool=false)
-    ctx_names = propertynames(getfield(d, :maximizer_kwargs), private)
+    ctx_names = propertynames(getfield(d, :context), private)
     extra_names = propertynames(getfield(d, :extra), private)
     return (fieldnames(DataSample)..., ctx_names..., extra_names...)
 end
@@ -141,7 +137,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.maximizer_kwargs)
+    for (key, value) in pairs(d.context)
         value_str = sprint(show, value; context=io_limited)
         push!(fields, "$key=$value_str")
     end
@@ -169,8 +165,8 @@ Transform the features in the dataset.
 """
 function StatsBase.transform(t, dataset::AbstractVector{<:DataSample})
     return map(dataset) do d
-        (; maximizer_kwargs, extra, x, θ, y) = d
-        DataSample(StatsBase.transform(t, x), θ, y, maximizer_kwargs, extra)
+        (; context, extra, x, θ, y) = d
+        DataSample(StatsBase.transform(t, x), θ, y, context, extra)
     end
 end
 
@@ -192,8 +188,8 @@ Reconstruct the features in the dataset.
 """
 function StatsBase.reconstruct(t, dataset::AbstractVector{<:DataSample})
     return map(dataset) do d
-        (; maximizer_kwargs, extra, x, θ, y) = d
-        DataSample(StatsBase.reconstruct(t, x), θ, y, maximizer_kwargs, extra)
+        (; context, extra, x, θ, y) = d
+        DataSample(StatsBase.reconstruct(t, x), θ, y, context, extra)
     end
 end