added symmetric_kl and symmetric_kl_grad
@@ -747,6 +747,69 @@ def ll_dirichlet(data1, data2):
)
+@numba.njit(fastmath=True)
+def symmetric_kl(x, y, z=1e-11):
+ """
+ symmetrized KL divergence between two probability distributions
+
+ ..math::
+ D(x, y) = \frac{D_{KL}\left(x \Vert y\right) + D_{KL}\left(y \Vert x\right)}{2}
+ """
+ n = x.shape[0]
+ x_sum = 0.0
+ y_sum = 0.0
+ kl1 = 0.0
+ kl2 = 0.0
+
+ for i in range(n):
+ x[i] += z
+ x_sum += x[i]
+ y[i] += z
+ y_sum += y[i]
+
+ for i in range(n):
+ x[i] /= x_sum
+ y[i] /= y_sum
+
+ for i in range(n):
+ kl1 += x[i] * np.log(x[i] / y[i])
+ kl2 += y[i] * np.log(y[i] / x[i])
+
+ return (kl1 + kl2) / 2
+
+
+@numba.njit(fastmath=True)
+def symmetric_kl_grad(x, y, z=1e-11):
+ """
+ symmetrized KL divergence and its gradient
+
+ """
+ n = x.shape[0]
+ x_sum = 0.0
+ y_sum = 0.0
+ kl1 = 0.0
+ kl2 = 0.0
+
+ for i in range(n):
+ x[i] += z
+ x_sum += x[i]
+ y[i] += z
+ y_sum += y[i]
+
+ for i in range(n):
+ x[i] /= x_sum
+ y[i] /= y_sum
+
+ for i in range(n):
+ kl1 += x[i] * np.log(x[i] / y[i])
+ kl2 += y[i] * np.log(y[i] / x[i])
+
+ dist = (kl1 + kl2) / 2
+ grad = (np.log(y / x) - (x / y) + 1) / 2
+
+ return dist, grad
+
+
@numba.njit()
def correlation_grad(x, y):
mu_x = 0.0
@@ -1091,6 +1154,7 @@ named_distances = {
"haversine": haversine,
"braycurtis": bray_curtis,
"ll_dirichlet": ll_dirichlet,
+ "symmetric_kl": symmetric_kl,
# Binary distances
"hamming": hamming,
"jaccard": jaccard,
@@ -1135,6 +1199,7 @@ named_distances_with_gradients = {
"hellinger": hellinger_grad,
"haversine": haversine_grad,
"braycurtis": bray_curtis_grad,
+ "symmetric_kl": symmetric_kl_grad,
# Special embeddings
"spherical_gaussian_energy": spherical_gaussian_energy_grad,
"diagonal_gaussian_energy": diagonal_gaussian_energy_grad,
