Merge pull request #104 from devmotion/_scale

Optimization of ScaleTransform for SimpleKernel

Author: theogf

src/generic.jl

@@ -5,11 +5,6 @@ Base.iterate(k::Kernel, ::Any) = nothing
 
 printshifted(io::IO, o, shift::Int) = print(io, o)
 
-# See https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/issues/96
-_scale(t::ScaleTransform, metric::Euclidean, x, y) =  first(t.s) * evaluate(metric, x, y)
-_scale(t::ScaleTransform, metric::Union{SqEuclidean,DotProduct}, x, y) =  first(t.s)^2 * evaluate(metric, x, y)
-_scale(t::ScaleTransform, metric, x, y) = evaluate(metric, apply(t, x), apply(t, y))
-
 ### Syntactic sugar for creating matrices and using kernel functions
 function concretetypes(k, ktypes::Vector)
     isempty(subtypes(k)) ? push!(ktypes, k) : concretetypes.(subtypes(k), Ref(ktypes))

src/kernels/transformedkernel.jl

@@ -11,6 +11,24 @@ end
 
 (k::TransformedKernel)(x, y) = k.kernel(k.transform(x), k.transform(y))
 
+# Optimizations for scale transforms of simple kernels to save allocations:
+# Instead of a multiplying every element of the inputs before evaluating the metric,
+# we perform a scalar multiplcation of the distance of the original inputs, if possible.
+function (k::TransformedKernel{<:SimpleKernel,<:ScaleTransform})(
+    x::AbstractVector{<:Real},
+    y::AbstractVector{<:Real},
+)
+    return kappa(k.kernel, _scale(k.transform, metric(k.kernel), x, y))
+end
+
+function _scale(t::ScaleTransform, metric::Euclidean, x, y)
+    return first(t.s) * evaluate(metric, x, y)
+end
+function _scale(t::ScaleTransform, metric::Union{SqEuclidean,DotProduct}, x, y)
+    return first(t.s)^2 * evaluate(metric, x, y)
+end
+_scale(t::ScaleTransform, metric, x, y) = evaluate(metric, t(x), t(y))
+
 """
 ```julia
     transform(k::BaseKernel, t::Transform) (1)