Merge pull request #62 from JuliaDecisionFocusedLearning/cleanup

Cleanup and reinforce the interface

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/Argmax/Argmax.jl

@@ -5,6 +5,8 @@ using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES
 using Flux: Chain, Dense
 using Random
 
+using LinearAlgebra: dot
+
 """
 $TYPEDEF
 
@@ -29,6 +31,8 @@ function Base.show(io::IO, bench::ArgmaxBenchmark)
     )
 end
 
+Utils.objective_value(::ArgmaxBenchmark, sample::DataSample, y) = dot(sample.θ, y)
+
 """
 $TYPEDSIGNATURES
 

src/Argmax2D/Argmax2D.jl

@@ -30,6 +30,8 @@ function Base.show(io::IO, bench::Argmax2DBenchmark)
     return print(io, "Argmax2DBenchmark(nb_features=$nb_features)")
 end
 
+Utils.objective_value(::Argmax2DBenchmark, sample::DataSample, y) = dot(sample.θ, y)
+
 """
 $TYPEDSIGNATURES
 

src/ContextualStochasticArgmax/ContextualStochasticArgmax.jl

@@ -3,6 +3,7 @@ module ContextualStochasticArgmax
 using ..Utils
 using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES
 using Flux: Dense
+using LinearAlgebra: dot
 using Random: Random, AbstractRNG, MersenneTwister
 using Statistics: mean
 
@@ -11,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(ξ)`.
@@ -44,25 +45,42 @@ end
 Utils.is_minimization_problem(::ContextualStochasticArgmaxBenchmark) = false
 Utils.generate_maximizer(::ContextualStochasticArgmaxBenchmark) = one_hot_argmax
 
+function Utils.objective_value(
+    ::ContextualStochasticArgmaxBenchmark, sample::DataSample, y, scenario
+)
+    return dot(scenario, y)
+end
+
+function Utils.objective_value(
+    bench::ContextualStochasticArgmaxBenchmark, sample::DataSample, y
+)
+    if hasproperty(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.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,
@@ -71,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/DecisionFocusedLearningBenchmarks.jl

@@ -79,6 +79,7 @@ export generate_statistical_model
 export generate_maximizer
 export generate_anticipative_solver, generate_parametric_anticipative_solver
 export is_exogenous, is_endogenous
+export is_minimization_problem
 
 export objective_value
 export has_visualization, plot_instance, plot_solution, plot_trajectory, animate_trajectory

src/FixedSizeShortestPath/FixedSizeShortestPath.jl

@@ -55,9 +55,9 @@ function FixedSizeShortestPathBenchmark(;
 end
 
 function Utils.objective_value(
-    ::FixedSizeShortestPathBenchmark, θ::AbstractArray, y::AbstractArray
+    ::FixedSizeShortestPathBenchmark, sample::DataSample, y::AbstractArray
 )
-    return -dot(θ, y)
+    return -dot(sample.θ, y)
 end
 
 """

src/PortfolioOptimization/PortfolioOptimization.jl

@@ -6,7 +6,7 @@ using Distributions: Uniform, Bernoulli
 using Flux: Chain, Dense
 using Ipopt: Ipopt
 using JuMP: @variable, @objective, @constraint, optimize!, value, Model, set_silent
-using LinearAlgebra: I
+using LinearAlgebra: I, dot
 using Random: Random, AbstractRNG, MersenneTwister
 
 """
@@ -38,6 +38,10 @@ struct PortfolioOptimizationBenchmark <: AbstractBenchmark
     f::Vector{Float32}
 end
 
+function Utils.objective_value(::PortfolioOptimizationBenchmark, sample::DataSample, y)
+    return dot(sample.θ, y)
+end
+
 """
 $TYPEDSIGNATURES
 

src/Ranking/Ranking.jl

@@ -5,6 +5,8 @@ using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES
 using Flux: Chain, Dense
 using Random
 
+using LinearAlgebra: dot
+
 """
 $TYPEDEF
 
@@ -29,6 +31,9 @@ function Base.show(io::IO, bench::RankingBenchmark)
     )
 end
 
+Utils.objective_value(::RankingBenchmark, sample::DataSample, y) = dot(sample.θ, y)
+Utils.is_minimization_problem(::RankingBenchmark) = false
+
 """
 $TYPEDSIGNATURES