Update and fix tests

Author: BatyLeo

src/Argmax2D/Argmax2D.jl

@@ -50,18 +50,14 @@ function Utils.generate_dataset(
     bench::Argmax2DBenchmark, dataset_size=10; seed=nothing, rng=MersenneTwister(seed)
 )
     (; nb_features, encoder, polytope_vertex_range) = bench
-    X = [randn(rng, nb_features) for _ in 1:dataset_size]
-    θs = encoder.(X)
-    θs ./= 2 * norm.(θs)
-    instances = [
-        build_polytope(rand(rng, polytope_vertex_range); shift=rand(rng)) for
-        _ in 1:dataset_size
-    ]
-    Y = [maximizer(θ; instance) for (θ, instance) in zip(θs, instances)]
-    return [
-        DataSample(; x, θ_true, y_true, instance) for
-        (x, θ_true, y_true, instance) in zip(X, θs, Y, instances)
-    ]
+    return map(1:dataset_size) do _
+        x = randn(rng, nb_features)
+        θ_true = encoder(x)
+        θ_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=θ_true, y_true=y_true, instance=instance)
+    end
 end
 
 Utils.generate_maximizer(::Argmax2DBenchmark) = maximizer

test/argmax_2d.jl

@@ -1,31 +1,33 @@
 @testitem "Argmax2D" begin
     using DecisionFocusedLearningBenchmarks
 
-    instance_dim = 10
     nb_features = 5
-
-    b = ArgmaxBenchmark(; instance_dim=instance_dim, nb_features=nb_features)
+    b = Argmax2DBenchmark(; nb_features=nb_features)
 
     io = IOBuffer()
     show(io, b)
-    @test String(take!(io)) == "ArgmaxBenchmark(instance_dim=10, nb_features=5)"
+    @test String(take!(io)) == "Argmax2DBenchmark(nb_features=5)"
 
     dataset = generate_dataset(b, 50)
     model = generate_statistical_model(b)
     maximizer = generate_maximizer(b)
 
     for (i, sample) in enumerate(dataset)
-        (; x, θ_true, y_true) = sample
-        @test size(x) == (nb_features, instance_dim)
-        @test length(θ_true) == instance_dim
-        @test length(y_true) == instance_dim
-        @test isnothing(sample.instance)
-        @test all(y_true .== maximizer(θ_true))
+        (; x, θ_true, y_true, instance) = sample
+        @test length(x) == nb_features
+        @test length(θ_true) == 2  # 2D vectors
+        @test length(y_true) == 2  # 2D point
+        @test !isnothing(sample.instance)  # instance is a polytope
+        @test instance isa Vector{Vector{Float64}}  # polytope is vector of 2D points
+        @test all(length(vertex) == 2 for vertex in instance)  # all vertices are 2D
+        @test y_true in instance  # solution should be a vertex of the polytope
+        @test y_true == maximizer(θ_true; instance=instance)
 
         θ = model(x)
-        @test length(θ) == instance_dim
+        @test length(θ) == 2  # 2D vector
 
-        y = maximizer(θ)
-        @test length(y) == instance_dim
+        y = maximizer(θ; instance=instance)
+        @test length(y) == 2  # 2D point
+        @test y in instance  # solution should be a vertex of the polytope
     end
 end