Update SurrogatesPolyChaos to use SurrogatesBase

lib/SurrogatesPolyChaos/Project.toml

@@ -6,10 +6,14 @@ version = "0.1.0"
 [deps]
 PolyChaos = "8d666b04-775d-5f6e-b778-5ac7c70f65a3"
 Surrogates = "6fc51010-71bc-11e9-0e15-a3fcc6593c49"
+SurrogatesBase = "89f642e6-4179-4274-8202-c11f4bd9a72c"
 
 [compat]
 PolyChaos = "0.2"
-Surrogates = "6"
+SafeTestsets = "0.1"
+Surrogates = "6.9"
+SurrogatesBase = "1.1"
+Zygote = "0.6"
 julia = "1.10"
 
 [extras]

lib/SurrogatesPolyChaos/src/SurrogatesPolyChaos.jl

@@ -1,50 +1,73 @@
 module SurrogatesPolyChaos
 
-import Surrogates: AbstractSurrogate, add_point!, _check_dimension
-export PolynomialChaosSurrogate
-
+using SurrogatesBase
 using PolyChaos
 
-mutable struct PolynomialChaosSurrogate{X, Y, L, U, C, O, N} <: AbstractSurrogate
+export PolynomialChaosSurrogate, update!
+
+mutable struct PolynomialChaosSurrogate{X, Y, L, U, C, O, N} <:
+               AbstractDeterministicSurrogate
     x::X
     y::Y
     lb::L
     ub::U
     coeff::C
-    ortopolys::O
+    orthopolys::O
     num_of_multi_indexes::N
 end
 
-function _calculatepce_coeff(x, y, num_of_multi_indexes, op::AbstractCanonicalOrthoPoly)
+function PolynomialChaosSurrogate(x, y, lb::Number, ub::Number;
+        orthopolys::AbstractCanonicalOrthoPoly = GaussOrthoPoly(2))
     n = length(x)
-    A = zeros(eltype(x), n, num_of_multi_indexes)
-    for i in 1:n
-        A[i, :] = PolyChaos.evaluate(x[i], op)
+    poly_degree = orthopolys.deg
+    num_of_multi_indexes = 1 + poly_degree
+    if n < 2 + 3 * num_of_multi_indexes
+        error("To avoid numerical problems, it's strongly suggested to have at least $(2+3*num_of_multi_indexes) samples")
     end
-    return (A' * A) \ (A' * y)
+    coeff = _calculatepce_coeff(x, y, num_of_multi_indexes, orthopolys)
+    return PolynomialChaosSurrogate(x, y, lb, ub, coeff, orthopolys, num_of_multi_indexes)
 end
 
-function PolynomialChaosSurrogate(x, y, lb::Number, ub::Number;
-        op::AbstractCanonicalOrthoPoly = GaussOrthoPoly(2))
+function PolynomialChaosSurrogate(x, y, lb, ub;
+        orthopolys = MultiOrthoPoly([GaussOrthoPoly(2) for j in 1:length(lb)], 2))
     n = length(x)
-    poly_degree = op.deg
-    num_of_multi_indexes = 1 + poly_degree
+    d = length(lb)
+    poly_degree = orthopolys.deg
+    num_of_multi_indexes = binomial(d + poly_degree, poly_degree)
     if n < 2 + 3 * num_of_multi_indexes
-        throw("To avoid numerical problems, it's strongly suggested to have at least $(2+3*num_of_multi_indexes) samples")
+        error("To avoid numerical problems, it's strongly suggested to have at least $(2+3*num_of_multi_indexes) samples")
     end
-    coeff = _calculatepce_coeff(x, y, num_of_multi_indexes, op)
-    return PolynomialChaosSurrogate(x, y, lb, ub, coeff, op, num_of_multi_indexes)
+    coeff = _calculatepce_coeff(x, y, num_of_multi_indexes, orthopolys)
+    return PolynomialChaosSurrogate(x, y, lb, ub, coeff, orthopolys, num_of_multi_indexes)
 end
 
 function (pc::PolynomialChaosSurrogate)(val::Number)
-    # Check to make sure dimensions of input matches expected dimension of surrogate
-    _check_dimension(pc, val)
-
-    return sum([pc.coeff[i] * PolyChaos.evaluate(val, pc.ortopolys)[i]
+    return sum([pc.coeff[i] * PolyChaos.evaluate(val, pc.orthopolys)[i]
                 for i in 1:(pc.num_of_multi_indexes)])
 end
 
-function _calculatepce_coeff(x, y, num_of_multi_indexes, op::MultiOrthoPoly)
+function (pcND::PolynomialChaosSurrogate)(val)
+    sum = zero(eltype(val))
+    for i in 1:(pcND.num_of_multi_indexes)
+        sum = sum +
+              pcND.coeff[i] *
+              first(PolyChaos.evaluate(pcND.orthopolys.ind[i, :], collect(val),
+            pcND.orthopolys))
+    end
+    return sum
+end
+
+function _calculatepce_coeff(
+        x, y, num_of_multi_indexes, orthopolys::AbstractCanonicalOrthoPoly)
+    n = length(x)
+    A = zeros(eltype(x), n, num_of_multi_indexes)
+    for i in 1:n
+        A[i, :] = PolyChaos.evaluate(x[i], orthopolys)
+    end
+    return (A' * A) \ (A' * y)
+end
+
+function _calculatepce_coeff(x, y, num_of_multi_indexes, orthopolys::MultiOrthoPoly)
     n = length(x)
     d = length(x[1])
     A = zeros(eltype(x[1]), n, num_of_multi_indexes)
@@ -53,53 +76,16 @@ function _calculatepce_coeff(x, y, num_of_multi_indexes, op::MultiOrthoPoly)
         for j in 1:d
             xi[j] = x[i][j]
         end
-        A[i, :] = PolyChaos.evaluate(xi, op)
+        A[i, :] = PolyChaos.evaluate(xi, orthopolys)
     end
     return (A' * A) \ (A' * y)
 end
 
-function PolynomialChaosSurrogate(x, y, lb, ub;
-        op::MultiOrthoPoly = MultiOrthoPoly([GaussOrthoPoly(2)
-                                             for j in 1:length(lb)],
-            2))
-    n = length(x)
-    d = length(lb)
-    poly_degree = op.deg
-    num_of_multi_indexes = binomial(d + poly_degree, poly_degree)
-    if n < 2 + 3 * num_of_multi_indexes
-        throw("To avoid numerical problems, it's strongly suggested to have at least $(2+3*num_of_multi_indexes) samples")
-    end
-    coeff = _calculatepce_coeff(x, y, num_of_multi_indexes, op)
-    return PolynomialChaosSurrogate(x, y, lb, ub, coeff, op, num_of_multi_indexes)
-end
-
-function (pcND::PolynomialChaosSurrogate)(val)
-    # Check to make sure dimensions of input matches expected dimension of surrogate
-    _check_dimension(pcND, val)
-
-    sum = zero(eltype(val[1]))
-    for i in 1:(pcND.num_of_multi_indexes)
-        sum = sum +
-              pcND.coeff[i] *
-              first(PolyChaos.evaluate(pcND.ortopolys.ind[i, :], collect(val),
-            pcND.ortopolys))
-    end
-    return sum
-end
-
-function add_point!(polych::PolynomialChaosSurrogate, x_new, y_new)
-    if length(polych.lb) == 1
-        #1D
-        polych.x = vcat(polych.x, x_new)
-        polych.y = vcat(polych.y, y_new)
-        polych.coeff = _calculatepce_coeff(polych.x, polych.y, polych.num_of_multi_indexes,
-            polych.ortopolys)
-    else
-        polych.x = vcat(polych.x, x_new)
-        polych.y = vcat(polych.y, y_new)
-        polych.coeff = _calculatepce_coeff(polych.x, polych.y, polych.num_of_multi_indexes,
-            polych.ortopolys)
-    end
+function SurrogatesBase.update!(polych::PolynomialChaosSurrogate, x_new, y_new)
+    polych.x = vcat(polych.x, x_new)
+    polych.y = vcat(polych.y, y_new)
+    polych.coeff = _calculatepce_coeff(polych.x, polych.y, polych.num_of_multi_indexes,
+        polych.orthopolys)
     nothing
 end
 

lib/SurrogatesPolyChaos/test/runtests.jl

@@ -1,119 +1,69 @@
-using SafeTestsets
+using SafeTestsets, Test
 
 @safetestset "PolynomialChaosSurrogates" begin
     using Surrogates
     using PolyChaos
-    using Surrogates: sample, SobolSample
     using SurrogatesPolyChaos
     using Zygote
 
-    #1D
-    n = 20
-    lb = 0.0
-    ub = 4.0
-    f = x -> 2 * x
-    x = sample(n, lb, ub, SobolSample())
-    y = f.(x)
-
-    my_pce = PolynomialChaosSurrogate(x, y, lb, ub)
-    val = my_pce(2.0)
-    add_point!(my_pce, 3.0, 6.0)
-    my_pce_changed = PolynomialChaosSurrogate(x, y, lb, ub, op = Uniform01OrthoPoly(1))
-
-    #ND
-    n = 60
-    lb = [0.0, 0.0]
-    ub = [5.0, 5.0]
-    f = x -> x[1] * x[2]
-    x = sample(n, lb, ub, SobolSample())
-    y = f.(x)
-
-    my_pce = PolynomialChaosSurrogate(x, y, lb, ub)
-    val = my_pce((2.0, 2.0))
-    add_point!(my_pce, (2.0, 3.0), 6.0)
-
-    op1 = Uniform01OrthoPoly(1)
-    op2 = Beta01OrthoPoly(2, 2, 1.2)
-    ops = [op1, op2]
-    multi_poly = MultiOrthoPoly(ops, min(1, 2))
-    my_pce_changed = PolynomialChaosSurrogate(x, y, lb, ub, op = multi_poly)
-
-    # Surrogate optimization test
-    lb = 0.0
-    ub = 15.0
-    p = 1.99
-    a = 2
-    b = 6
-    objective_function = x -> 2 * x + 1
-    x = sample(20, lb, ub, SobolSample())
-    y = objective_function.(x)
-    my_poly1d = PolynomialChaosSurrogate(x, y, lb, ub)
-    @test_broken surrogate_optimize(objective_function, SRBF(), a, b, my_poly1d,
-        LowDiscrepancySample(; base = 2))
-
-    lb = [0.0, 0.0]
-    ub = [10.0, 10.0]
-    obj_ND = x -> log(x[1]) * exp(x[2])
-    x = sample(40, lb, ub, RandomSample())
-    y = obj_ND.(x)
-    my_polyND = PolynomialChaosSurrogate(x, y, lb, ub)
-    surrogate_optimize(obj_ND, SRBF(), lb, ub, my_polyND, SobolSample(), maxiters = 15)
-
-    # AD Compatibility
+    @testset "Scalar Inputs" begin
+        n = 20
+        lb = 0.0
+        ub = 4.0
+        f = x -> 2 * x
+        x = sample(n, lb, ub, SobolSample())
+        y = f.(x)
+        my_pce = PolynomialChaosSurrogate(x, y, lb, ub)
+        x_val = 1.2
+        @test my_pce(x_val) ≈ f(x_val)
+        update!(my_pce, [3.0], [6.0])
+        my_pce_changed = PolynomialChaosSurrogate(
+            x, y, lb, ub; orthopolys = Uniform01OrthoPoly(1))
+        @test my_pce_changed(x_val) ≈ f(x_val)
+    end
 
-    lb = 0.0
-    ub = 3.0
-    n = 10
-    x = sample(n, lb, ub, SobolSample())
-    f = x -> x^2
-    y = f.(x)
+    @testset "Vector Inputs" begin
+        n = 60
+        lb = [0.0, 0.0]
+        ub = [5.0, 5.0]
+        f = x -> x[1] * x[2]
+        x = collect.(sample(n, lb, ub, SobolSample()))
+        y = f.(x)
+        my_pce = PolynomialChaosSurrogate(x, y, lb, ub)
+        x_val = [1.2, 1.4]
+        @test my_pce(x_val) ≈ f(x_val)
+        update!(my_pce, [[2.0, 3.0]], [6.0])
+        @test my_pce(x_val) ≈ f(x_val)
+        op1 = Uniform01OrthoPoly(1)
+        op2 = Beta01OrthoPoly(2, 2, 1.2)
+        ops = [op1, op2]
+        multi_poly = MultiOrthoPoly(ops, min(1, 2))
+        my_pce_changed = PolynomialChaosSurrogate(x, y, lb, ub, orthopolys = multi_poly)
+    end
 
-    # #Polynomialchaos
-    @testset "Polynomial Chaos" begin
+    @testset "Derivative" begin
+        lb = 0.0
+        ub = 3.0
         f = x -> x^2
         n = 50
-        x = sample(n, lb, ub, SobolSample())
+        x = collect(sample(n, lb, ub, SobolSample()))
         y = f.(x)
         my_poli = PolynomialChaosSurrogate(x, y, lb, ub)
         g = x -> my_poli'(x)
-        g(3.0)
+        x_val = 3.0
+        @test g(x_val) ≈ 2 * x_val
     end
 
-    # #PolynomialChaos
-    @testset "Polynomial Chaos ND" begin
+    @testset "Gradient" begin
         n = 50
         lb = [0.0, 0.0]
         ub = [10.0, 10.0]
-        x = sample(n, lb, ub, SobolSample())
+        x = collect.(sample(n, lb, ub, SobolSample()))
         f = x -> x[1] * x[2]
         y = f.(x)
         my_poli_ND = PolynomialChaosSurrogate(x, y, lb, ub)
-        g = x -> Zygote.gradient(my_poli_ND, x)
-        g((1.0, 1.0))
-
-        n = 10
-        d = 2
-        lb = [0.0, 0.0]
-        ub = [5.0, 5.0]
-        x = sample(n, lb, ub, SobolSample())
-        f = x -> x[1]^2 + x[2]^2
-        y1 = f.(x)
-        grad1 = x -> 2 * x[1]
-        grad2 = x -> 2 * x[2]
-        function create_grads(n, d, grad1, grad2, y)
-            c = 0
-            y2 = zeros(eltype(y[1]), n * d)
-            for i in 1:n
-                y2[i + c] = grad1(x[i])
-                y2[i + c + 1] = grad2(x[i])
-                c = c + 1
-            end
-            return y2
-        end
-        y2 = create_grads(n, d, grad1, grad2, y)
-        y = vcat(y1, y2)
-        my_gek_ND = GEK(x, y, lb, ub)
-        g = x -> Zygote.gradient(my_gek_ND, x)
-        @test_broken g((2.0, 5.0)) #breaks after Zygote version 0.6.43
+        g = x -> Zygote.gradient(my_poli_ND, x)[1]
+        x_val = [1.0, 2.0]
+        @test g(x_val) ≈ [x_val[2], x_val[1]]
     end
 end