refactor *Mahalanobis to avoid unnecessary symmetry checks (#231)

Author: dkarrasch

src/mahalanobis.jl

@@ -72,8 +72,7 @@ function result_type(d::SqMahalanobis, ::Type{T1}, ::Type{T2}) where {T1,T2}
     return typeof(z * zero(eltype(d.qmat)) * z)
 end
 
-# SqMahalanobis
-
+# TODO: merge the following two once we lift the lower bound for julia (above v1.4?)
 function (dist::SqMahalanobis)(a::AbstractVector, b::AbstractVector)
     if length(a) != length(b)
         throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
@@ -83,24 +82,47 @@ function (dist::SqMahalanobis)(a::AbstractVector, b::AbstractVector)
     z = a - b
     return dot(z, Q * z)
 end
+function (dist::Mahalanobis)(a::AbstractVector, b::AbstractVector)
+    if length(a) != length(b)
+        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
+    end
 
-sqmahalanobis(a::AbstractVector, b::AbstractVector, Q::AbstractMatrix) = SqMahalanobis(Q)(a, b)
-
-function colwise!(r::AbstractArray, dist::SqMahalanobis, a::AbstractMatrix, b::AbstractMatrix)
     Q = dist.qmat
-    get_colwise_dims(size(Q, 1), r, a, b)
     z = a - b
-    dot_percol!(r, Q * z, z)
+    return sqrt(dot(z, Q * z))
 end
 
-function colwise!(r::AbstractArray, dist::SqMahalanobis, a::AbstractVector, b::AbstractMatrix)
+sqmahalanobis(a::AbstractVector, b::AbstractVector, Q::AbstractMatrix) = SqMahalanobis(Q)(a, b)
+mahalanobis(a::AbstractVector, b::AbstractVector, Q::AbstractMatrix) = Mahalanobis(Q)(a, b)
+
+function _colwise!(r, dist, a, b)
     Q = dist.qmat
     get_colwise_dims(size(Q, 1), r, a, b)
     z = a .- b
     dot_percol!(r, Q * z, z)
 end
 
-function _pairwise!(r::AbstractMatrix, dist::SqMahalanobis, a::AbstractMatrix, b::AbstractMatrix)
+function colwise!(r::AbstractArray, dist::SqMahalanobis, a::AbstractMatrix, b::AbstractMatrix)
+    _colwise!(r, dist, a, b)
+end
+function colwise!(r::AbstractArray, dist::SqMahalanobis, a::AbstractVector, b::AbstractMatrix)
+    _colwise!(r, dist, a, b)
+end
+function colwise!(r::AbstractArray, dist::SqMahalanobis, a::AbstractMatrix, b::AbstractVector)
+    _colwise!(r, dist, a, b)
+end
+
+function colwise!(r::AbstractArray, dist::Mahalanobis, a::AbstractMatrix, b::AbstractMatrix)
+    sqrt!(_colwise!(r, dist, a, b))
+end
+function colwise!(r::AbstractArray, dist::Mahalanobis, a::AbstractVector, b::AbstractMatrix)
+    sqrt!(_colwise!(r, dist, a, b))
+end
+function colwise!(r::AbstractArray, dist::Mahalanobis, a::AbstractMatrix, b::AbstractVector)
+    sqrt!(_colwise!(r, dist, a, b))
+end
+
+function _pairwise!(r::AbstractMatrix, dist::Union{SqMahalanobis,Mahalanobis}, a::AbstractMatrix, b::AbstractMatrix)
     Q = dist.qmat
     m, na, nb = get_pairwise_dims(size(Q, 1), r, a, b)
 
@@ -112,13 +134,13 @@ function _pairwise!(r::AbstractMatrix, dist::SqMahalanobis, a::AbstractMatrix, b
 
     for j = 1:nb
         @simd for i = 1:na
-            @inbounds r[i, j] = max(sa2[i] + sb2[j] - 2 * r[i, j], 0)
+            @inbounds r[i, j] = eval_end(dist, max(sa2[i] + sb2[j] - 2 * r[i, j], 0))
         end
     end
     r
 end
 
-function _pairwise!(r::AbstractMatrix, dist::SqMahalanobis, a::AbstractMatrix)
+function _pairwise!(r::AbstractMatrix, dist::Union{SqMahalanobis,Mahalanobis}, a::AbstractMatrix)
     Q = dist.qmat
     m, n = get_pairwise_dims(size(Q, 1), r, a)
 
@@ -132,33 +154,11 @@ function _pairwise!(r::AbstractMatrix, dist::SqMahalanobis, a::AbstractMatrix)
         end
         r[j, j] = 0
         for i = (j + 1):n
-            @inbounds r[i, j] = max(sa2[i] + sa2[j] - 2 * r[i, j], 0)
+            @inbounds r[i, j] = eval_end(dist, max(sa2[i] + sa2[j] - 2 * r[i, j], 0))
         end
     end
     r
 end
 
-
-# Mahalanobis
-
-function (dist::Mahalanobis)(a::AbstractVector, b::AbstractVector)
-    sqrt(SqMahalanobis(dist.qmat, skipchecks = true)(a, b))
-end
-
-mahalanobis(a::AbstractVector, b::AbstractVector, Q::AbstractMatrix) = Mahalanobis(Q)(a, b)
-
-function colwise!(r::AbstractArray, dist::Mahalanobis, a::AbstractMatrix, b::AbstractMatrix)
-    sqrt!(colwise!(r, SqMahalanobis(dist.qmat, skipchecks = true), a, b))
-end
-
-function colwise!(r::AbstractArray, dist::Mahalanobis, a::AbstractVector, b::AbstractMatrix)
-    sqrt!(colwise!(r, SqMahalanobis(dist.qmat, skipchecks = true), a, b))
-end
-
-function _pairwise!(r::AbstractMatrix, dist::Mahalanobis, a::AbstractMatrix, b::AbstractMatrix)
-    sqrt!(_pairwise!(r, SqMahalanobis(dist.qmat, skipchecks = true), a, b))
-end
-
-function _pairwise!(r::AbstractMatrix, dist::Mahalanobis, a::AbstractMatrix)
-    sqrt!(_pairwise!(r, SqMahalanobis(dist.qmat, skipchecks = true), a))
-end
+eval_end(::SqMahalanobis, x) = x
+eval_end(::Mahalanobis, x) = sqrt(x)

test/test_dists.jl

@@ -375,6 +375,8 @@ end # testset
     Q = Q * Q'  # make sure Q is positive-definite
     @test_throws DimensionMismatch mahalanobis(p, q, Q)
     @test_throws DimensionMismatch mahalanobis(q, q, Q)
+    @test_throws DimensionMismatch sqmahalanobis(p, q, Q)
+    @test_throws DimensionMismatch sqmahalanobis(q, q, Q)
     mat23 = [0.3 0.2 0.0; 0.1 0.0 0.4]
     mat22 = [0.3 0.2; 0.1 0.4]
     @test_throws DimensionMismatch colwise!(mat23, Euclidean(), mat23, mat23)