I added a few function calls in _approx_rbf that dispatch on ::RadialBasis{F,Q,X,<:AbstractArray{<:Number}} -- a type that restricts the output to just a number.

Doing this lets us elide creating an output buffer, which means RBFs are now allocation-free when the output is a single dimension.

The motivation behind this is specific to my use-case (creating surrogates of old Fortran spectral models to use in MCMC fitting), but I hope others may find this useful too.

Using an example from the docs:

using Surrogates
using BenchmarkTools

function booth(x)
    x1=x[1]
    x2=x[2]
    term1 = (x1 + 2*x2 - 7)^2;
    term2 = (2*x1 + x2 - 5)^2;
    y = term1 + term2;
end

n_samples = 80
lower_bound = [-5.0, 0.0]
upper_bound = [10.0, 15.0]

xys = sample(n_samples, lower_bound, upper_bound, SobolSample())
zs = booth.(xys)

radial_basis = RadialBasis(xys, zs,  lower_bound, upper_bound)

# new sample space
xys = sample(500, lower_bound, upper_bound, SobolSample())

function benchmarktest(radial_basis, xys)
    s = 0.0
    for v in xys
        s += radial_basis(v)
    end
    s
end

@btime benchmarktest($radial_basis, $xys)
# before: 253.122 μs (500 allocations: 31.25 KiB)
# after: 166.757 μs (0 allocations: 0 bytes)

I also added a few @inbound annotations to squeeze a bit more performance, since _check_dimension already ensures things are pretty good:

@btime benchmarktest($radial_basis, $xys)
# after + @inbounds 57.860 μs (0 allocations: 0 bytes)

Overall, the result is a 4-5x speedup, with all the RBF test cases still passing (at least on my machine), and inspecting invocations with @code_warntype shows the methods to be type-stable.

Let me know if this is alright, or if there's anything else I can do! Also any feedback or alternative suggestions on the implementation would be welcomed :) I would have used @generated for this but I know that can be a bit controversial.

Cheers!

Codecov Report

Merging #392 (6a8efb4) into master (c1b8397) will decrease coverage by 0.38%. The diff coverage is 95.65%.

@@            Coverage Diff             @@
##           master     #392      +/-   ##
==========================================
- Coverage   79.53%   79.15%   -0.39%     
==========================================
  Files          16       16              
  Lines        2336     2331       -5     
==========================================
- Hits         1858     1845      -13     
- Misses        478      486       +8

Impacted Files	Coverage Δ
src/Radials.jl	`85.84% <95.65%> (-2.44%)`	:arrow_down:
src/GEKPLS.jl	`94.21% <0.00%> (-0.83%)`	:arrow_down:
src/Optimization.jl	`71.68% <0.00%> (-0.37%)`	:arrow_down:

:mega: Codecov can now indicate which changes are the most critical in Pull Requests. Learn more

SciML / Surrogates.jl

Speed up radial basis functions when output is single number #392

Codecov Report