add divergence functions and tests

Functions which compute the divergence between two probability distribution. This commit includes kernel Stein discrepancy as a divergence metric and associated tests.

Author: joannajzou

Diff for: src/Kernels/divergences.jl

@@ -0,0 +1,49 @@
+## Discrepancies
+""" 
+    Divergence
+
+    A struct of abstract type Divergence produces a measure of discrepancy between two probability distributions. Discepancies may take as argument analytical distributions or sets of samples representing empirical distributions.
+"""
+abstract type Divergence end
+
+
+"""
+    KernelSteinDiscrepancy <: Divergence
+        score :: Function
+        knl :: Kernel
+
+    Computes the kernel Stein discrepancy between distributions p (from which samples are provided) and q (for which the score is provided) based on the RKHS defined by kernel k.
+"""
+struct KernelSteinDiscrepancy <: Divergence
+    score :: Function
+    kernel :: Kernel
+end
+
+function KernelSteinDiscrepancy(; score, kernel)
+    return KernelSteinDiscrepancy(score, kernel)
+end
+
+
+function compute_divergence(
+    x :: Vector{T},
+    div :: KernelSteinDiscrepancy,
+) where T <: Union{Real, Vector{<:Real}}
+
+    N = length(x)
+    sq = div.score.(x)
+    k = div.kernel
+
+    ksd = 0.0
+    for i = 1:N
+        for j = i:N
+            m = (i == j) ? 1 : 2
+            sks = sq[i]' * compute_kernel(x[i], x[j], k) * sq[j]
+            sk = sq[i]' * compute_grady_kernel(x[i], x[j], k)
+            ks = compute_gradx_kernel(x[i], x[j], k)' * sq[j]
+            trk = tr(compute_gradxy_kernel(x[i], x[j], k))
+
+            ksd += m * (sks + sk + ks + trk) / (N*(N-1.0))
+        end
+    end
+    return ksd
+end

Diff for: src/Kernels/kernels.jl

@@ -1,29 +1,6 @@
 include("features.jl")
 include("distances.jl")
 
-export
-    Distance,
-    Forstner,
-    compute_distance,
-    compute_gradx_distance,
-    compute_grady_distance,
-    compute_gradxy_distance,
-    Euclidean,
-    Feature,
-    GlobalMean,
-    CorrelationMatrix,
-    compute_feature,
-    compute_features,
-    Kernel,
-    DotProduct,
-    get_parameters,
-    RBF,
-    compute_kernel,
-    compute_gradx_kernel,
-    compute_grady_kernel,
-    compute_gradxy_kernel,
-    KernelMatrix
-
 ###############
 """
     Kernel
@@ -94,7 +71,7 @@ function compute_kernel(
     r::RBF,
 ) where {T<:Union{Vector{<:Real},Symmetric{<:Real,<:Matrix{<:Real}}}}
     d2 = compute_distance(A, B, r.d)
-    r.β * exp(-0.5 * d2 / r.ℓ^2)
+    r.β * exp(-0.5 * d2 / (2*r.ℓ^2))
 end
 
 """
@@ -105,12 +82,12 @@ Compute gradient of the kernel between features A and B using kernel k, with res
 function compute_gradx_kernel(
     A::T,
     B::T,
-    knl::RBF,
+    r::RBF,
     ) where {T<:Vector{<:Real}}
 
-    k = compute_kernel(A, B, knl)
-    ∇d = compute_gradx_distance(A, B, knl.d)
-    return -k * ∇d / (2*knl.ℓ^2)
+    k = compute_kernel(A, B, r)
+    ∇d = compute_gradx_distance(A, B, r.d)
+    return -k * ∇d / (2*r.ℓ^2)
 end
 
 """
@@ -121,12 +98,12 @@ Compute gradient of the kernel between features A and B using kernel k, with res
 function compute_grady_kernel(
     A::T,
     B::T,
-    knl::RBF,
+    r::RBF,
     ) where {T<:Vector{<:Real}}
 
-    k = compute_kernel(A, B, knl)
-    ∇d = compute_grady_distance(A, B, knl.d)
-    return -k * ∇d / (2*knl.ℓ^2)
+    k = compute_kernel(A, B, r)
+    ∇d = compute_grady_distance(A, B, r.d)
+    return -k * ∇d / (2*r.ℓ^2)
 end
 
 """
@@ -137,15 +114,15 @@ Compute the second-order cross derivative of the kernel between features A and B
 function compute_gradxy_kernel(
     A::T,
     B::T,
-    knl::RBF,
+    r::RBF,
     ) where {T<:Vector{<:Real}}
 
-    k = compute_kernel(A, B, knl)
-    ∇xd = compute_gradx_distance(A, B, knl.d)
-    ∇yd = compute_grady_distance(A, B, knl.d)
-    ∇xyd = compute_gradxy_distance(A, B, knl.d)
+    k = compute_kernel(A, B, r)
+    ∇xd = compute_gradx_distance(A, B, r.d)
+    ∇yd = compute_grady_distance(A, B, r.d)
+    ∇xyd = compute_gradxy_distance(A, B, r.d)
 
-    return k .* ( -∇xyd/(2*knl.ℓ^2) .+ ∇xd'*∇yd/(4*knl.ℓ^4) )
+    return k .* ( -∇xyd/(2*r.ℓ^2) .+ ∇xd'*∇yd/(4*r.ℓ^4) )
 
 end
 
@@ -213,3 +190,33 @@ function KernelMatrix(
     F2 = compute_feature.(ds2, (f,); dt = dt)
     KernelMatrix(F1, F2, k)
 end
+
+
+
+
+include("divergences.jl")
+export
+    Distance,
+    Forstner,
+    compute_distance,
+    compute_gradx_distance,
+    compute_grady_distance,
+    compute_gradxy_distance,
+    Euclidean,
+    Feature,
+    GlobalMean,
+    CorrelationMatrix,
+    compute_feature,
+    compute_features,
+    Kernel,
+    DotProduct,
+    get_parameters,
+    RBF,
+    compute_kernel,
+    compute_gradx_kernel,
+    compute_grady_kernel,
+    compute_gradxy_kernel,
+    KernelMatrix,
+    Divergence,
+    KernelSteinDiscrepancy,
+    compute_divergence

Diff for: test/kernels/kernel_tests.jl

@@ -1,6 +1,8 @@
 using AtomsBase
 using Unitful, UnitfulAtomic
 using LinearAlgebra
+using Distributions
+
 # initialize some fake descriptors
 d = 8
 num_atoms = 20
@@ -30,6 +32,9 @@ f_cm = compute_feature.(ld, (cm,))
 @test compute_features(ds, gm) == f_gm
 @test compute_features(ds, cm) == f_cm
 
+mvn = MvNormal(zeros(d), I(d))
+f_mvn = [rand(mvn) for j = 1:100]
+
 ## distances 
 fo = Forstner(1e-16)
 e = Euclidean(d)
@@ -82,3 +87,12 @@ rbf_fo = RBF(fo)
 @test typeof(KernelMatrix(f_gm, rbf_e)) <: Symmetric{Float64,Matrix{Float64}}
 @test typeof(KernelMatrix(f_cm, rbf_e)) <: Symmetric{Float64,Matrix{Float64}}
 @test typeof(KernelMatrix(f_cm, rbf_fo)) <: Symmetric{Float64,Matrix{Float64}}
+
+
+# divergences
+std_gauss_score(x) = -x
+ksd = KernelSteinDiscrepancy(score=std_gauss_score, kernel=rbf_e)
+
+@test typeof(ksd) <: Divergence
+@test compute_divergence(f_mvn, ksd) < compute_divergence(f_gm, ksd)
+