Add with_lengthscale (alternative/extension of #335) (#336)

Co-authored-by: willtebbutt <wct23@cam.ac.uk>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: ST John <st--@users.noreply.github.com>

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.6"
+version = "0.10.7"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/transform.md

@@ -36,3 +36,9 @@ SelectTransform
 ChainTransform
 PeriodicTransform
 ```
+
+## Convenience functions
+
+```@docs
+with_lengthscale
+```

docs/src/userguide.md

@@ -15,6 +15,11 @@ For example, a squared exponential kernel is created by
       k = SqExponentialKernel() ∘ ScaleTransform(2.0)
       k = compose(SqExponentialKernel(), ScaleTransform(2.0))
     ```
+    Alternatively, you can use the convenience function [`with_lengthscale`](@ref):
+    ```julia
+    k = with_lengthscale(SqExponentialKernel(), 0.5)
+    ```
+    [`with_lengthscale`](@ref) also works with vector-valued lengthscales for ARD.
     Check the [Input Transforms](@ref input_transforms) page for more details.
 
 !!! tip "How do I set the kernel variance?"

src/KernelFunctions.jl

@@ -27,6 +27,7 @@ export Transform,
     IdentityTransform,
     FunctionTransform,
     PeriodicTransform
+export with_lengthscale
 
 export NystromFact, nystrom
 
@@ -75,6 +76,7 @@ include(joinpath("transform", "selecttransform.jl"))
 include(joinpath("transform", "chaintransform.jl"))
 include(joinpath("transform", "periodic_transform.jl"))
 include(joinpath("kernels", "transformedkernel.jl"))
+include(joinpath("transform", "with_lengthscale.jl"))
 
 include(joinpath("basekernels", "constant.jl"))
 include(joinpath("basekernels", "cosine.jl"))

src/transform/with_lengthscale.jl

@@ -0,0 +1,43 @@
+"""
+    with_lengthscale(kernel::Kernel, lengthscale::Real)
+
+Construct a transformed kernel with `lengthscale`.
+
+# Examples
+
+```jldoctest
+julia> kernel = with_lengthscale(SqExponentialKernel(), 2.5);
+
+julia> x = rand(2);
+
+julia> y = rand(2);
+
+julia> kernel(x, y) ≈ (SqExponentialKernel() ∘ ScaleTransform(0.4))(x, y)
+true
+```
+"""
+function with_lengthscale(kernel::Kernel, lengthscale::Real)
+    return kernel ∘ ScaleTransform(inv(lengthscale))
+end
+
+"""
+    with_lengthscale(kernel::Kernel, lengthscales::AbstractVector{<:Real})
+
+Construct a transformed "ARD" kernel with different `lengthscales` for each dimension.
+
+# Examples
+
+```jldoctest
+julia> kernel = with_lengthscale(SqExponentialKernel(), [0.5, 2.5]);
+
+julia> x = rand(2);
+
+julia> y = rand(2);
+
+julia> kernel(x, y) ≈ (SqExponentialKernel() ∘ ARDTransform([2, 0.4]))(x, y)
+true
+```
+"""
+function with_lengthscale(kernel::Kernel, lengthscales::AbstractVector{<:Real})
+    return kernel ∘ ARDTransform(map(inv, lengthscales))
+end

test/runtests.jl

@@ -76,6 +76,8 @@ include("test_utils.jl")
             print(" ")
             include(joinpath("transform", "periodic_transform.jl"))
             print(" ")
+            include(joinpath("transform", "with_lengthscale.jl"))
+            print(" ")
         end
         @info "Ran tests on Transform"
     end

test/transform/with_lengthscale.jl

@@ -0,0 +1,17 @@
+@testset "with_lengthscale" begin
+    @testset "ScaleTransform" begin
+        l = exp(rand())
+        kernel = @inferred(with_lengthscale(SqExponentialKernel(), l))
+
+        @test kernel isa TransformedKernel{<:SqExponentialKernel,<:ScaleTransform}
+        @test kernel.transform.s[1] ≈ inv(l)
+    end
+
+    @testset "ARDTransform" begin
+        l = map(exp, rand(5))
+        kernel = @inferred(with_lengthscale(SqExponentialKernel(), l))
+
+        @test kernel isa TransformedKernel{<:SqExponentialKernel,<:ARDTransform}
+        @test kernel.transform.v ≈ map(inv, l)
+    end
+end