From 29434e7b4a6bc0162f736a1f0cdca1f4e719f840 Mon Sep 17 00:00:00 2001
From: BatyLeo <leo.baty67@gmail.com>
Date: Wed, 16 Oct 2024 15:59:42 +0200
Subject: [PATCH] cleanup and fix fixed size shortest path

---
 .../FixedSizeShortestPath.jl                  | 21 ++-----------------
 test/fixed_size_shortest_path.jl              |  6 +++++-
 2 files changed, 7 insertions(+), 20 deletions(-)

diff --git a/src/FixedSizeShortestPath/FixedSizeShortestPath.jl b/src/FixedSizeShortestPath/FixedSizeShortestPath.jl
index 4a904a8..874fbe1 100644
--- a/src/FixedSizeShortestPath/FixedSizeShortestPath.jl
+++ b/src/FixedSizeShortestPath/FixedSizeShortestPath.jl
@@ -125,13 +125,13 @@ function Utils.generate_dataset(
         [rand(rng, Uniform{type}(1 - ν, 1 + ν), E) for _ in 1:dataset_size]
     end
     costs = [
-        (1 .+ (3 .+ B * zᵢ ./ type(sqrt(p))) .^ deg) .* ξᵢ for (ξᵢ, zᵢ) in zip(ξ, features)
+        -(1 .+ (3 .+ B * zᵢ ./ type(sqrt(p))) .^ deg) .* ξᵢ for (ξᵢ, zᵢ) in zip(ξ, features)
     ]
 
     shortest_path_maximizer = Utils.generate_maximizer(bench)
 
     # Label solutions
-    solutions = shortest_path_maximizer.(.-costs)
+    solutions = shortest_path_maximizer.(costs)
     return [DataSample(; x=x, θ=θ, y=y) for (x, θ, y) in zip(features, costs, solutions)]
 end
 
@@ -145,23 +145,6 @@ function Utils.generate_statistical_model(bench::FixedSizeShortestPathBenchmark)
     return Chain(Dense(p, ne(graph)))
 end
 
-function objective_value(::FixedSizeShortestPathBenchmark, θ, y)
-    return dot(θ, y)
-end
-
-function Utils.compute_gap(
-    bench::FixedSizeShortestPathBenchmark, model, features, costs, solutions, maximizer
-)
-    res = 0.0
-    for (x, ȳ, θ̄) in zip(features, solutions, costs)
-        θ = model(x)
-        y = maximizer(θ)
-        val = objective_value(bench, θ̄, ȳ)
-        res += (objective_value(bench, θ̄, y) - val) / val
-    end
-    return res / length(features)
-end
-
 export FixedSizeShortestPathBenchmark
 
 end
diff --git a/test/fixed_size_shortest_path.jl b/test/fixed_size_shortest_path.jl
index df43cb4..495655b 100644
--- a/test/fixed_size_shortest_path.jl
+++ b/test/fixed_size_shortest_path.jl
@@ -14,15 +14,19 @@
     model = generate_statistical_model(b)
     maximizer = generate_maximizer(b)
 
+    gap = compute_gap(b, dataset, model, maximizer)
+    @test gap >= 0
+
     for sample in dataset
         x = sample.x
         θ_true = sample.θ
         y_true = sample.y
+        @test all(θ_true .< 0)
         @test size(x) == (p,)
         @test length(θ_true) == A
         @test length(y_true) == A
         @test isnothing(sample.instance)
-        @test all(y_true .== maximizer(-θ_true))
+        @test all(y_true .== maximizer(θ_true))
         θ = model(x)
         @test length(θ) == length(θ_true)
         y = maximizer(θ)