refactor(SurrogatesPolyChaos): update it to use SurrogatesBase

lib/SurrogatesPolyChaos/src/SurrogatesPolyChaos.jl

@@ -1,49 +1,70 @@
 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 (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, op::AbstractCanonicalOrthoPoly)
+    n = length(x)
+    A = zeros(eltype(x), n, num_of_multi_indexes)
+    for i in 1:n
+        A[i, :] = PolyChaos.evaluate(x[i], op)
+    end
+    return (A' * A) \ (A' * y)
+end
+
 function _calculatepce_coeff(x, y, num_of_multi_indexes, op::MultiOrthoPoly)
     n = length(x)
     d = length(x[1])
@@ -58,48 +79,11 @@ function _calculatepce_coeff(x, y, num_of_multi_indexes, op::MultiOrthoPoly)
     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