Allow OffsetArrays and integer arrays in CosineDist #164

Project.toml

@@ -1,6 +1,6 @@
 name = "Distances"
 uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
-version = "0.8.2"
+version = "0.9.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -10,8 +10,9 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 julia = "1"
 
 [extras]
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Random", "Test"]
+test = ["OffsetArrays", "Random", "Test"]

src/common.jl

@@ -99,16 +99,26 @@ function sqrt!(a::AbstractArray)
 end
 
 function sumsq_percol(a::AbstractMatrix{T}) where {T}
-    m = size(a, 1)
     n = size(a, 2)
     r = Vector{T}(undef, n)
-    for j = 1:n
+    @simd for j in 1:n
         aj = view(a, :, j)
         r[j] = dot(aj, aj)
     end
     return r
 end
 
+function norm_percol(a::AbstractMatrix{T}) where {T}
+    n = size(a, 2)
+    √T = typeof(sqrt(oneunit(T)))
+    r = Vector{√T}(undef, n)
+    @simd for j in 1:n
+        aj = view(a, :, j)
+        r[j] = sqrt(dot(aj, aj))
+    end
+    return r
+end
+
 function wsumsq_percol(w::AbstractArray, a::AbstractMatrix)
     m = size(a, 1)
     n = size(a, 2)

src/metrics.jl

@@ -236,7 +236,7 @@ Base.@propagate_inbounds function _evaluate(d::UnionMetrics, a::AbstractArray, b
     end
     @inbounds begin
         s = eval_start(d, a, b)
-        if IndexStyle(a, b) === IndexLinear() || size(a) == size(b)
+        if (IndexStyle(a, b) === IndexLinear() && eachindex(a) == eachindex(b)) || axes(a) == axes(b)
             @simd for I in eachindex(a, b)
                 ai = a[I]
                 bi = b[I]
@@ -265,8 +265,9 @@ Base.@propagate_inbounds function _evaluate(d::UnionMetrics, a::AbstractArray, b
     end
     @inbounds begin
         s = eval_start(d, a, b)
-        if IndexStyle(a, b, p) === IndexLinear() || size(a) == size(b)
-            @simd for I in eachindex(a, b)
+        if (IndexStyle(a, b, p) === IndexLinear() && eachindex(a) == eachindex(b) == eachindex(p)) ||
+                axes(a) == axes(b) == axes(p)
+            @simd for I in eachindex(a, b, p)
                 ai = a[I]
                 bi = b[I]
                 pi = p[I]
@@ -801,8 +802,8 @@ function _pairwise!(r::AbstractMatrix, dist::CosineDist,
                     a::AbstractMatrix, b::AbstractMatrix)
     m, na, nb = get_pairwise_dims(r, a, b)
     mul!(r, a', b)
-    ra = sqrt!(sumsq_percol(a))
-    rb = sqrt!(sumsq_percol(b))
+    ra = norm_percol(a)
+    rb = norm_percol(b)
     for j = 1:nb
         @simd for i = 1:na
             @inbounds r[i, j] = max(1 - r[i, j] / (ra[i] * rb[j]), 0)
@@ -813,7 +814,7 @@ end
 function _pairwise!(r::AbstractMatrix, dist::CosineDist, a::AbstractMatrix)
     m, n = get_pairwise_dims(r, a)
     mul!(r, a', a)
-    ra = sqrt!(sumsq_percol(a))
+    ra = norm_percol(a)
     @inbounds for j = 1:n
         @simd for i = j + 1:n
             r[i, j] = max(1 - r[i, j] / (ra[i] * ra[j]), 0)

test/runtests.jl

@@ -2,6 +2,7 @@ using Distances
 
 using Test
 using LinearAlgebra
+using OffsetArrays
 using Random
 using Statistics
 

test/test_dists.jl

@@ -170,6 +170,7 @@ end
             x_int, y_int = Int64.(x), Int64.(y)
             @test cosine_dist(x_int, y_int) == (1.0 - 112.0 / sqrt(19530.0))
             @test corr_dist(x, y) ≈ cosine_dist(x .- mean(x), vec(y) .- mean(y))
+            @test corr_dist(OffsetVector(x, -1:length(x)-2), y) == corr_dist(x, y)
             @test chisq_dist(x, y) == sum((x - vec(y)).^2 ./ (x + vec(y)))
             @test spannorm_dist(x, y) == maximum(x - vec(y)) - minimum(x - vec(y))
 
@@ -504,8 +505,8 @@ function test_pairwise(dist, x, y, T)
             rxx[i, j] = dist(x[:, i], x[:, j])
         end
         # As earlier, we have small rounding errors in accumulations
-        @test pairwise(dist, x, y) ≈ rxy
-        @test pairwise(dist, x) ≈ rxx
+        @test pairwise(dist, x, y, dims=2) ≈ rxy
+        @test pairwise(dist, x, dims=2) ≈ rxx
         @test pairwise(dist, x, y, dims=2) ≈ rxy
         @test pairwise(dist, x, dims=2) ≈ rxx
         @test pairwise(dist, permutedims(x), permutedims(y), dims=1) ≈ rxy
@@ -536,6 +537,7 @@ end
     test_pairwise(Hamming(), A, B, T)
 
     test_pairwise(CosineDist(), X, Y, T)
+    test_pairwise(CosineDist(), A, B, T)
     test_pairwise(CorrDist(), X, Y, T)
 
     test_pairwise(ChiSqDist(), X, Y, T)
@@ -570,16 +572,16 @@ end
 
 @testset "Euclidean precision" begin
     X = [0.1 0.2; 0.3 0.4; -0.1 -0.1]
-    pd = pairwise(Euclidean(1e-12), X, X)
+    pd = pairwise(Euclidean(1e-12), X, X, dims=2)
     @test pd[1, 1] == 0
     @test pd[2, 2] == 0
-    pd = pairwise(Euclidean(1e-12), X)
+    pd = pairwise(Euclidean(1e-12), X, dims=2)
     @test pd[1, 1] == 0
     @test pd[2, 2] == 0
-    pd = pairwise(SqEuclidean(1e-12), X, X)
+    pd = pairwise(SqEuclidean(1e-12), X, X, dims=2)
     @test pd[1, 1] == 0
     @test pd[2, 2] == 0
-    pd = pairwise(SqEuclidean(1e-12), X)
+    pd = pairwise(SqEuclidean(1e-12), X, dims=2)
     @test pd[1, 1] == 0
     @test pd[2, 2] == 0
 end