From 09b6f0588314150fadbe6509a06f5944ede20767 Mon Sep 17 00:00:00 2001
From: Sathvik Bhagavan <sathvik.bhagavan@juliahub.com>
Date: Tue, 23 Jul 2024 07:40:21 +0000
Subject: [PATCH] test: update tests for SurrogatesAbstractGPs

---
 lib/SurrogatesAbstractGPs/test/runtests.jl | 90 ++++++++--------------
 1 file changed, 32 insertions(+), 58 deletions(-)

diff --git a/lib/SurrogatesAbstractGPs/test/runtests.jl b/lib/SurrogatesAbstractGPs/test/runtests.jl
index fb553f0e..f8a291e1 100644
--- a/lib/SurrogatesAbstractGPs/test/runtests.jl
+++ b/lib/SurrogatesAbstractGPs/test/runtests.jl
@@ -1,11 +1,10 @@
 using SafeTestsets, Test
-using Surrogates: sample, SobolSample
 
 @safetestset "AbstractGPSurrogate" begin
-    using Surrogates
     using SurrogatesAbstractGPs
     using AbstractGPs
     using Zygote
+    using Surrogates: sample, SobolSample
 
     @testset "1D -> 1D" begin
         lb = 0.0
@@ -16,7 +15,7 @@ using Surrogates: sample, SobolSample
         agp1D = AbstractGPSurrogate(x, y, gp = GP(SqExponentialKernel()), Σy = 0.05)
         x_new = 2.5
         y_actual = f.(x_new)
-        y_predicted = agp1D([x_new])
+        y_predicted = agp1D([x_new])[1]
         @test isapprox(y_predicted, y_actual, atol = 0.1)
     end
 
@@ -31,9 +30,9 @@ using Surrogates: sample, SobolSample
         x_new = 2.5
         y_actual = f.(x_new)
         for i in 2:length(x_points)
-            add_point!(agp1D, x_points[i], y_points[i])
+            update!(agp1D, [x_points[i]], [y_points[i]])
         end
-        y_predicted = agp1D([x_new])
+        y_predicted = agp1D(x_new)
         @test isapprox(y_predicted, y_actual, atol = 0.1)
     end
 
@@ -60,7 +59,7 @@ using Surrogates: sample, SobolSample
         logpdf_vals = []
         push!(logpdf_vals, logpdf_surrogate(agp_2D))
         for i in 2:length(x)
-            add_point!(agp_2D, x[i], y[i])
+            update!(agp_2D, [x[i]], [y[i]])
             push!(logpdf_vals, logpdf_surrogate(agp_2D))
         end
         @test first(logpdf_vals) < last(logpdf_vals) #as more points are added log marginal posterior predictive probability increases
@@ -77,31 +76,6 @@ using Surrogates: sample, SobolSample
         @test agpND(x_new)≈f(x_new) atol=0.2
     end
 
-    @testset "Optimization 1D" begin
-        objective_function = x -> 2 * x + 1
-        lb = 0.0
-        ub = 6.0
-        x = [2.0, 4.0, 6.0]
-        y = [5.0, 9.0, 13.0]
-        p = 2
-        a = 2
-        b = 6
-        my_k_EI1 = AbstractGPSurrogate(x, y)
-        surrogate_optimize(objective_function, EI(), a, b, my_k_EI1, RandomSample(),
-            maxiters = 200, num_new_samples = 155)
-    end
-
-    @testset "Optimization ND" begin
-        objective_function_ND = z -> 3 * hypot(z...) + 1
-        x = [(1.2, 3.0), (3.0, 3.5), (5.2, 5.7)]
-        y = objective_function_ND.(x)
-        theta = [2.0, 2.0]
-        lb = [1.0, 1.0]
-        ub = [6.0, 6.0]
-        my_k_E1N = AbstractGPSurrogate(x, y)
-        surrogate_optimize(objective_function_ND, EI(), lb, ub, my_k_E1N, RandomSample())
-    end
-
     @testset "check working of logpdf_surrogate 1D" begin
         lb = 0.0
         ub = 3.0
@@ -122,32 +96,32 @@ using Surrogates: sample, SobolSample
         logpdf_surrogate(agpND)
     end
 
-    lb = 0.0
-    ub = 3.0
-    n = 10
-    x = sample(n, lb, ub, SobolSample())
-    f = x -> x^2
-    y = f.(x)
-
-    #AbstractGP 1D
-    @testset "AbstractGP 1D" begin
-        agp1D = AbstractGPSurrogate(x, y, gp = GP(SqExponentialKernel()), Σy = 0.05)
-        g = x -> agp1D'(x)
-        g([2.0])
-    end
-
-    lb = [0.0, 0.0]
-    ub = [10.0, 10.0]
-    n = 5
-    x = sample(n, lb, ub, SobolSample())
-    f = x -> x[1] * x[2]
-    y = f.(x)
-
-    # AbstractGP ND
-    @testset "AbstractGPSurrogate ND" begin
-        my_agp = AbstractGPSurrogate(x, y, gp = GP(SqExponentialKernel()), Σy = 0.05)
-        g = x -> Zygote.gradient(my_agp, x)
-        #g([(2.0,5.0)])
-        g((2.0, 5.0))
+    @testset "Gradients" begin
+        @testset "1D" begin
+            lb = 0.0
+            ub = 3.0
+            n = 100
+            x = sample(n, lb, ub, SobolSample())
+            f = x -> x^2
+            y = f.(x)
+            agp1D = AbstractGPSurrogate(x, y, gp = GP(SqExponentialKernel()), Σy = 0.05)
+            g = x -> Zygote.gradient(agp1D, x)
+            x_val = 2.0
+            @test g(x_val)[1]≈2 * x_val rtol=1e-1
+        end
+        @testset "ND" begin
+            lb = [0.0, 0.0]
+            ub = [10.0, 10.0]
+            n = 100
+            x = sample(n, lb, ub, SobolSample())
+            f = x -> x[1] * x[2]
+            y = f.(x)
+            my_agp = AbstractGPSurrogate(x, y, gp = GP(SqExponentialKernel()), Σy = 0.05)
+            g = x -> Zygote.gradient(my_agp, x)
+            x_val = (2.0, 5.0)
+            g_val = g(x_val)[1]
+            @test g_val[1]≈x_val[2] rtol=1e-1
+            @test g_val[2]≈x_val[1] rtol=1e-1
+        end
     end
 end