cleanup docstrings and add optional noise in dataset

Author: BatyLeo

src/Argmax/Argmax.jl

@@ -8,13 +8,13 @@ using Random
 """
 $TYPEDEF
 
-Benchmark problem with an argmax as the CO algorithm.
+Basic benchmark problem with an argmax as the CO algorithm.
 
 # Fields
 $TYPEDFIELDS
 """
 struct ArgmaxBenchmark <: AbstractBenchmark
-    "iinstances dimension, total number of classes"
+    "instances dimension, total number of classes"
     instance_dim::Int
     "number of features"
     nb_features::Int
@@ -45,7 +45,7 @@ end
 """
 $TYPEDSIGNATURES
 
-Return a top k maximizer.
+Return an argmax maximizer.
 """
 function Utils.generate_maximizer(bench::ArgmaxBenchmark)
     return one_hot_argmax
@@ -54,19 +54,23 @@ end
 """
 $TYPEDSIGNATURES
 
-Generate a dataset of labeled instances for the subset selection problem.
-The mapping between features and cost is identity.
+Generate a dataset of labeled instances for the argmax problem.
 """
-function Utils.generate_dataset(bench::ArgmaxBenchmark, dataset_size::Int=10; seed::Int=0)
+function Utils.generate_dataset(
+    bench::ArgmaxBenchmark, dataset_size::Int=10; seed::Int=0, noise_std=0.0
+)
     (; instance_dim, nb_features) = bench
     rng = MersenneTwister(seed)
     features = [randn(rng, Float32, nb_features, instance_dim) for _ in 1:dataset_size]
     mapping = Chain(Dense(nb_features => 1; bias=false), vec)
     costs = mapping.(features)
-    solutions = one_hot_argmax.(costs)
+    # solutions = one_hot_argmax.(costs)
+    noisy_solutions = [
+        one_hot_argmax(θ + noise_std * randn(rng, Float32, instance_dim)) for θ in costs
+    ]
     return [
         DataSample(; x, θ_true, y_true) for
-        (x, θ_true, y_true) in zip(features, costs, solutions)
+        (x, θ_true, y_true) in zip(features, costs, noisy_solutions)
     ]
 end
 

src/Ranking/Ranking.jl

@@ -8,13 +8,13 @@ using Random
 """
 $TYPEDEF
 
-Benchmark problem with an argmax as the CO algorithm.
+Basic benchmark problem with ranking as the CO algorithm.
 
 # Fields
 $TYPEDFIELDS
 """
 struct RankingBenchmark <: AbstractBenchmark
-    "iinstances dimension, total number of classes"
+    "instances dimension, total number of classes"
     instance_dim::Int
     "number of features"
     nb_features::Int
@@ -43,7 +43,7 @@ end
 """
 $TYPEDSIGNATURES
 
-Return a top k maximizer.
+Return a ranking maximizer.
 """
 function Utils.generate_maximizer(bench::RankingBenchmark)
     return ranking
@@ -52,19 +52,23 @@ end
 """
 $TYPEDSIGNATURES
 
-Generate a dataset of labeled instances for the subset selection problem.
-The mapping between features and cost is identity.
+Generate a dataset of labeled instances for the ranking problem.
 """
-function Utils.generate_dataset(bench::RankingBenchmark, dataset_size::Int=10; seed::Int=0)
+function Utils.generate_dataset(
+    bench::RankingBenchmark, dataset_size::Int=10; seed::Int=0, noise_std=0.0
+)
     (; instance_dim, nb_features) = bench
     rng = MersenneTwister(seed)
     features = [randn(rng, Float32, nb_features, instance_dim) for _ in 1:dataset_size]
     mapping = Chain(Dense(nb_features => 1; bias=false), vec)
     costs = mapping.(features)
-    solutions = ranking.(costs)
+    # solutions = ranking.(costs)
+    noisy_solutions = [
+        ranking(θ .+ noise_std * randn(rng, Float32, instance_dim)) for θ in costs
+    ]
     return [
         DataSample(; x, θ_true, y_true) for
-        (x, θ_true, y_true) in zip(features, costs, solutions)
+        (x, θ_true, y_true) in zip(features, costs, noisy_solutions)
     ]
 end