Improve argmax benchmarks

Author: BatyLeo

src/Argmax/Argmax.jl

@@ -40,6 +40,8 @@ function ArgmaxBenchmark(; instance_dim::Int=10, nb_features::Int=5, seed=nothin
     return ArgmaxBenchmark(instance_dim, nb_features, model)
 end
 
+Utils.is_minimization_problem(::ArgmaxBenchmark) = false
+
 """
 $TYPEDSIGNATURES
 

src/Argmax2D/Argmax2D.jl

@@ -40,11 +40,21 @@ Custom constructor for [`Argmax2DBenchmark`](@ref).
 """
 function Argmax2DBenchmark(; nb_features::Int=5, seed=nothing, polytope_vertex_range=[6])
     Random.seed!(seed)
+<<<<<<< Updated upstream
     model = Chain(Dense(nb_features => 2; bias=false), vec)
     return Argmax2DBenchmark(nb_features, model, polytope_vertex_range)
 end
 
 maximizer(θ; instance) = instance[argmax(dot(θ, v) for v in instance)]
+=======
+    model = Dense(nb_features => 2; bias=false)
+    return Argmax2DBenchmark(nb_features, model, polytope_vertex_range)
+end
+
+Utils.is_minimization_problem(::Argmax2DBenchmark) = false
+
+maximizer(θ; instance, kwargs...) = instance[argmax(dot(θ, v) for v in instance)]
+>>>>>>> Stashed changes
 
 """
 $TYPEDSIGNATURES
@@ -56,7 +66,11 @@ function Utils.generate_dataset(
 )
     (; nb_features, encoder, polytope_vertex_range) = bench
     return map(1:dataset_size) do _
+<<<<<<< Updated upstream
         x = randn(rng, nb_features)
+=======
+        x = randn(rng, Float32, nb_features)
+>>>>>>> Stashed changes
         θ_true = encoder(x)
         θ_true ./= 2 * norm(θ_true)
         instance = build_polytope(rand(rng, polytope_vertex_range); shift=rand(rng))
@@ -84,23 +98,50 @@ function Utils.generate_statistical_model(
 )
     Random.seed!(rng, seed)
     (; nb_features) = bench
+<<<<<<< Updated upstream
     model = Chain(Dense(nb_features => 2; bias=false), vec)
     return model
 end
 
+=======
+    model = Dense(nb_features => 2; bias=false)
+    return model
+end
+
+function Utils.plot_data(::Argmax2DBenchmark; instance, θ, kwargs...)
+    pl = init_plot()
+    plot_polytope!(pl, instance)
+    plot_objective!(pl, θ)
+    return plot_maximizer!(pl, θ, instance, maximizer)
+end
+
+>>>>>>> Stashed changes
 """
 $TYPEDSIGNATURES
 
 Plot the data sample for the [`Argmax2DBenchmark`](@ref).
 """
 function Utils.plot_data(
+<<<<<<< Updated upstream
     ::Argmax2DBenchmark, sample::DataSample; θ_true=sample.θ_true, kwargs...
 )
     (; instance) = sample
     pl = init_plot()
     plot_polytope!(pl, instance)
     plot_objective!(pl, θ_true)
     return plot_maximizer!(pl, θ_true, instance, maximizer)
+=======
+    ::Argmax2DBenchmark,
+    sample::DataSample;
+    instance=sample.instance,
+    θ=sample.θ_true,
+    kwargs...,
+)
+    pl = init_plot()
+    plot_polytope!(pl, instance)
+    plot_objective!(pl, θ)
+    return plot_maximizer!(pl, θ, instance, maximizer)
+>>>>>>> Stashed changes
 end
 
 export Argmax2DBenchmark

src/Argmax2D/polytope.jl

@@ -21,18 +21,29 @@ function plot_polytope!(pl, vertices)
         fillcolor=:gray,
         fillalpha=0.2,
         linecolor=:black,
+<<<<<<< Updated upstream
         label=L"\mathrm{conv}(\mathcal{V})",
     )
 end;
 
 const logocolors = Colors.JULIA_LOGO_COLORS
 
+=======
+        label=L"\mathrm{conv}(\mathcal{Y}(x))",
+    )
+end;
+
+>>>>>>> Stashed changes
 function plot_objective!(pl, θ)
     Plots.plot!(
         pl,
         [0.0, θ[1]],
         [0.0, θ[2]];
+<<<<<<< Updated upstream
         color=logocolors.purple,
+=======
+        color=Colors.JULIA_LOGO_COLORS.purple,
+>>>>>>> Stashed changes
         arrow=true,
         lw=2,
         label=nothing,
@@ -47,7 +58,11 @@ function plot_maximizer!(pl, θ, instance, maximizer)
         pl,
         [ŷ[1]],
         [ŷ[2]];
+<<<<<<< Updated upstream
         color=logocolors.red,
+=======
+        color=Colors.JULIA_LOGO_COLORS.red,
+>>>>>>> Stashed changes
         markersize=9,
         markershape=:square,
         label=L"f(\theta)",
@@ -76,15 +91,23 @@ end;
 #         fillcolor=:blue,
 #         fillalpha=0.1,
 #         linestyle=:dash,
+<<<<<<< Updated upstream
 #         linecolor=logocolors.blue,
+=======
+#         linecolor=Colors.JULIA_LOGO_COLORS.blue,
+>>>>>>> Stashed changes
 #         label=L"\mathrm{conv}(\hat{p}(\theta))",
 #     )
 #     return Plots.scatter!(
 #         pl,
 #         map(first, A),
 #         map(last, A);
 #         markersize=25 .* p .^ 0.5,
+<<<<<<< Updated upstream
 #         markercolor=logocolors.blue,
+=======
+#         markercolor=Colors.JULIA_LOGO_COLORS.blue,
+>>>>>>> Stashed changes
 #         markerstrokewidth=0,
 #         markeralpha=0.4,
 #         label=L"\hat{p}(\theta)",
@@ -97,15 +120,23 @@ end;
 #         pl,
 #         [ŷΩ[1]],
 #         [ŷΩ[2]];
+<<<<<<< Updated upstream
 #         color=logocolors.blue,
+=======
+#         color=Colors.JULIA_LOGO_COLORS.blue,
+>>>>>>> Stashed changes
 #         markersize=6,
 #         markershape=:hexagon,
 #         label=L"\hat{f}(\theta)",
 #     )
 # end;
 
 # function compress_distribution!(
+<<<<<<< Updated upstream
 #     probadist::FixedAtomsProbabilityDistribution{A,W}; atol=0
+=======
+#     probadist::DifferentiableExpectations.FixedAtomsProbabilityDistribution{A,W}; atol=0
+>>>>>>> Stashed changes
 # ) where {A,W}
 #     (; atoms, weights) = probadist
 #     to_delete = Int[]