|
| 1 | +module Argmax2D |
| 2 | + |
| 3 | +using ..Utils |
| 4 | +using Colors: Colors |
| 5 | +using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES |
| 6 | +using Flux: Chain, Dense |
| 7 | +using LaTeXStrings: @L_str |
| 8 | +using LinearAlgebra: dot, norm |
| 9 | +using Plots: Plots |
| 10 | +using Random: Random, MersenneTwister |
| 11 | + |
| 12 | +include("polytope.jl") |
| 13 | + |
| 14 | +""" |
| 15 | +$TYPEDEF |
| 16 | +
|
| 17 | +Argmax becnhmark on a 2d polytope. |
| 18 | +
|
| 19 | +# Fields |
| 20 | +$TYPEDFIELDS |
| 21 | +""" |
| 22 | +struct Argmax2DBenchmark{E,R} <: AbstractBenchmark |
| 23 | + "number of features" |
| 24 | + nb_features::Int |
| 25 | + "true mapping between features and costs" |
| 26 | + encoder::E |
| 27 | + "" |
| 28 | + polytope_vertex_range::R |
| 29 | +end |
| 30 | + |
| 31 | +function Base.show(io::IO, bench::Argmax2DBenchmark) |
| 32 | + (; nb_features) = bench |
| 33 | + return print(io, "Argmax2DBenchmark(nb_features=$nb_features)") |
| 34 | +end |
| 35 | + |
| 36 | +""" |
| 37 | +$TYPEDSIGNATURES |
| 38 | +
|
| 39 | +Custom constructor for [`Argmax2DBenchmark`](@ref). |
| 40 | +""" |
| 41 | +function Argmax2DBenchmark(; nb_features::Int=5, seed=nothing, polytope_vertex_range=[6]) |
| 42 | + Random.seed!(seed) |
| 43 | + model = Dense(nb_features => 2; bias=false) |
| 44 | + return Argmax2DBenchmark(nb_features, model, polytope_vertex_range) |
| 45 | +end |
| 46 | + |
| 47 | +function Utils.is_minimization_problem(::Argmax2DBenchmark) |
| 48 | + return false |
| 49 | +end |
| 50 | + |
| 51 | +maximizer(θ; instance, kwargs...) = instance[argmax(dot(θ, v) for v in instance)] |
| 52 | + |
| 53 | +""" |
| 54 | +$TYPEDSIGNATURES |
| 55 | +
|
| 56 | +Generate a dataset for the [`Argmax2DBenchmark`](@ref). |
| 57 | +""" |
| 58 | +function Utils.generate_dataset( |
| 59 | + bench::Argmax2DBenchmark, dataset_size=10; seed=nothing, rng=MersenneTwister(seed) |
| 60 | +) |
| 61 | + (; nb_features, encoder, polytope_vertex_range) = bench |
| 62 | + return map(1:dataset_size) do _ |
| 63 | + x = randn(rng, Float32, nb_features) |
| 64 | + θ_true = encoder(x) |
| 65 | + θ_true ./= 2 * norm(θ_true) |
| 66 | + instance = build_polytope(rand(rng, polytope_vertex_range); shift=rand(rng)) |
| 67 | + y_true = maximizer(θ_true; instance) |
| 68 | + return DataSample(; x=x, θ_true=θ_true, y_true=y_true, instance=instance) |
| 69 | + end |
| 70 | +end |
| 71 | + |
| 72 | +""" |
| 73 | +$TYPEDSIGNATURES |
| 74 | +
|
| 75 | +Maximizer for the [`Argmax2DBenchmark`](@ref). |
| 76 | +""" |
| 77 | +function Utils.generate_maximizer(::Argmax2DBenchmark) |
| 78 | + return maximizer |
| 79 | +end |
| 80 | + |
| 81 | +""" |
| 82 | +$TYPEDSIGNATURES |
| 83 | +
|
| 84 | +Generate a statistical model for the [`Argmax2DBenchmark`](@ref). |
| 85 | +""" |
| 86 | +function Utils.generate_statistical_model( |
| 87 | + bench::Argmax2DBenchmark; seed=nothing, rng=MersenneTwister(seed) |
| 88 | +) |
| 89 | + Random.seed!(rng, seed) |
| 90 | + (; nb_features) = bench |
| 91 | + model = Dense(nb_features => 2; bias=false) |
| 92 | + return model |
| 93 | +end |
| 94 | + |
| 95 | +function Utils.plot_data(::Argmax2DBenchmark; instance, θ, kwargs...) |
| 96 | + pl = init_plot() |
| 97 | + plot_polytope!(pl, instance) |
| 98 | + plot_objective!(pl, θ) |
| 99 | + return plot_maximizer!(pl, θ, instance, maximizer) |
| 100 | +end |
| 101 | + |
| 102 | +""" |
| 103 | +$TYPEDSIGNATURES |
| 104 | +
|
| 105 | +Plot the data sample for the [`Argmax2DBenchmark`](@ref). |
| 106 | +""" |
| 107 | +function Utils.plot_data( |
| 108 | + bench::Argmax2DBenchmark, |
| 109 | + sample::DataSample; |
| 110 | + instance=sample.instance, |
| 111 | + θ=sample.θ_true, |
| 112 | + kwargs..., |
| 113 | +) |
| 114 | + return Utils.plot_data(bench; instance, θ, kwargs...) |
| 115 | +end |
| 116 | + |
| 117 | +export Argmax2DBenchmark |
| 118 | + |
| 119 | +end |
0 commit comments