Merge pull request JuliaGaussianProcesses#102 from devmotion/pairwise!

Author: devmotion

src/basekernels/fbm.jl

@@ -45,8 +45,8 @@ end
 
 function kernelmatrix!(K::AbstractMatrix, κ::FBMKernel, x::AbstractVector)
     modx = _mod(x)
-    modxx = pairwise(SqEuclidean(sqroundoff), x)
-    K .= _fbm.(modx, modx', modxx, κ.h)
+    pairwise!(K, SqEuclidean(sqroundoff), x)
+    K .= _fbm.(modx, modx', K, κ.h)
     return K
 end
 
@@ -58,10 +58,10 @@ end
 function kernelmatrix!(
     K::AbstractMatrix,
     κ::FBMKernel,
-    X::AbstractVector,
-    Y::AbstractVector,
+    x::AbstractVector,
+    y::AbstractVector,
 )
-    modxy = pairwise(SqEuclidean(sqroundoff), X, Y,dims = obsdim)
-    K .= _fbm.(_mod(x), _mod(y)', modxy, κ.h)
+    pairwise!(K, SqEuclidean(sqroundoff), x, y)
+    K .= _fbm.(_mod(x), _mod(y)', K, κ.h)
     return K
 end

src/generic.jl

@@ -33,3 +33,16 @@ function Distances.pairwise(
 )
     return pairwise(d, reshape(x, :, 1), reshape(y, :, 1); dims=1)
 end
+
+function Distances.pairwise!(out::AbstractMatrix, d::PreMetric, x::AbstractVector{<:Real})
+    return pairwise!(out, d, reshape(x, :, 1); dims=1)
+end
+
+function Distances.pairwise!(
+    out::AbstractMatrix,
+    d::PreMetric,
+    x::AbstractVector{<:Real},
+    y::AbstractVector{<:Real},
+)
+    return pairwise!(out, d, reshape(x, :, 1), reshape(y, :, 1); dims=1)
+end

src/matrix/kernelmatrix.jl

@@ -69,7 +69,8 @@ kerneldiagmatrix(κ::Kernel, x::AbstractVector) = map(x -> κ(x, x), x)
 
 function kernelmatrix!(K::AbstractMatrix, κ::SimpleKernel, x::AbstractVector)
     validate_inplace_dims(K, x)
-    return map!(d -> kappa(κ, d), K, pairwise(metric(κ), x))
+    pairwise!(K, metric(κ), x)
+    return map!(d -> kappa(κ, d), K, K)
 end
 
 function kernelmatrix!(
@@ -79,7 +80,8 @@ function kernelmatrix!(
     y::AbstractVector,
 )
     validate_inplace_dims(K, x, y)
-    return map!(d -> kappa(κ, d), K, pairwise(metric(κ), x, y))
+    pairwise!(K, metric(κ), x, y)
+    return map!(d -> kappa(κ, d), K, K)
 end
 
 function kernelmatrix(κ::SimpleKernel, x::AbstractVector)

src/utils.jl

@@ -45,7 +45,12 @@ dim(x::ColVecs) = size(x.X, 1)
 
 Distances.pairwise(d::PreMetric, x::ColVecs) = pairwise(d, x.X; dims=2)
 Distances.pairwise(d::PreMetric, x::ColVecs, y::ColVecs) = pairwise(d, x.X, y.X; dims=2)
-
+function Distances.pairwise!(out::AbstractMatrix, d::PreMetric, x::ColVecs)
+    return pairwise!(out, d, x.X; dims=2)
+end
+function Distances.pairwise!(out::AbstractMatrix, d::PreMetric, x::ColVecs, y::ColVecs)
+    return pairwise!(out, d, x.X, y.X; dims=2)
+end
 
 """
     RowVecs(X::AbstractMatrix)
@@ -70,8 +75,12 @@ dim(x::RowVecs) = size(x.X, 2)
 
 Distances.pairwise(d::PreMetric, x::RowVecs) = pairwise(d, x.X; dims=1)
 Distances.pairwise(d::PreMetric, x::RowVecs, y::RowVecs) = pairwise(d, x.X, y.X; dims=1)
-
-
+function Distances.pairwise!(out::AbstractMatrix, d::PreMetric, x::RowVecs)
+    return pairwise!(out, d, x.X; dims=1)
+end
+function Distances.pairwise!(out::AbstractMatrix, d::PreMetric, x::RowVecs, y::RowVecs)
+    return pairwise!(out, d, x.X, y.X; dims=1)
+end
 
 """
 Will be implemented at some point

test/basekernels/fbm.jl

@@ -7,8 +7,14 @@
     # kernelmatrix tests
     m1 = rand(3,3)
     m2 = rand(3,3)
-    @test kernelmatrix(k, m1, m1) ≈ kernelmatrix(k, m1) atol=1e-5
-
+    Kref = kernelmatrix(k, m1, m1)
+    @test kernelmatrix(k, m1) ≈ Kref atol=1e-5
+    K = zeros(3, 3)
+    kernelmatrix!(K, k, m1, m1)
+    @test K ≈ Kref atol=1e-5
+    fill!(K, 0)
+    kernelmatrix!(K, k, m1)
+    @test K ≈ Kref atol=1e-5
 
     x1 = rand()
     x2 = rand()

test/utils.jl

@@ -24,6 +24,12 @@
         DY = ColVecs(Y)
         @test pairwise(SqEuclidean(), DX) ≈ pairwise(SqEuclidean(), X; dims=2)
         @test pairwise(SqEuclidean(), DX, DY) ≈ pairwise(SqEuclidean(), X, Y; dims=2)
+        K = zeros(N, N)
+        pairwise!(K, SqEuclidean(), DX)
+        @test K ≈ pairwise(SqEuclidean(), X; dims=2)
+        K = zeros(N, N + 1)
+        pairwise!(K, SqEuclidean(), DX, DY)
+        @test K ≈ pairwise(SqEuclidean(), X, Y; dims=2)
 
         let
             @test Zygote.pullback(ColVecs, X)[1] == DX
@@ -52,6 +58,12 @@
         DY = RowVecs(Y)
         @test pairwise(SqEuclidean(), DX) ≈ pairwise(SqEuclidean(), X; dims=1)
         @test pairwise(SqEuclidean(), DX, DY) ≈ pairwise(SqEuclidean(), X, Y; dims=1)
+        K = zeros(D, D)
+        pairwise!(K, SqEuclidean(), DX)
+        @test K ≈ pairwise(SqEuclidean(), X; dims=1)
+        K = zeros(D, D + 1)
+        pairwise!(K, SqEuclidean(), DX, DY)
+        @test K ≈ pairwise(SqEuclidean(), X, Y; dims=1)
 
         let
             @test Zygote.pullback(RowVecs, X)[1] == DX