Add: Jensen-Shannon Divergence-based distance

ashvardanian · ashvardanian · commit 799284c9da14 · 2023-10-24T21:24:26.000Z
diff --git a/c/lib.cpp b/c/lib.cpp
@@ -22,6 +22,7 @@ metric_kind_t metric_kind_to_cpp(usearch_metric_kind_t kind) {
     case usearch_metric_l2sq_k: return metric_kind_t::l2sq_k;
     case usearch_metric_cos_k: return metric_kind_t::cos_k;
     case usearch_metric_haversine_k: return metric_kind_t::haversine_k;
+    case usearch_metric_divergence_k: return metric_kind_t::divergence_k;
     case usearch_metric_pearson_k: return metric_kind_t::pearson_k;
     case usearch_metric_jaccard_k: return metric_kind_t::jaccard_k;
     case usearch_metric_hamming_k: return metric_kind_t::hamming_k;
@@ -37,6 +38,7 @@ usearch_metric_kind_t metric_kind_to_c(metric_kind_t kind) {
     case metric_kind_t::l2sq_k: return usearch_metric_l2sq_k;
     case metric_kind_t::cos_k: return usearch_metric_cos_k;
     case metric_kind_t::haversine_k: return usearch_metric_haversine_k;
+    case metric_kind_t::divergence_k: return usearch_metric_divergence_k;
     case metric_kind_t::pearson_k: return usearch_metric_pearson_k;
     case metric_kind_t::jaccard_k: return usearch_metric_jaccard_k;
     case metric_kind_t::hamming_k: return usearch_metric_hamming_k;
diff --git a/c/usearch.h b/c/usearch.h
@@ -43,6 +43,7 @@ USEARCH_EXPORT typedef enum usearch_metric_kind_t {
     usearch_metric_ip_k,
     usearch_metric_l2sq_k,
     usearch_metric_haversine_k,
+    usearch_metric_divergence_k,
     usearch_metric_pearson_k,
     usearch_metric_jaccard_k,
     usearch_metric_hamming_k,
diff --git a/cpp/README.md b/cpp/README.md
@@ -76,6 +76,7 @@ The following distances are pre-packaged:
 - `metric_sorensen_gt<scalar_t>` for "Dice-Sorensen" coefficient for bit-strings.
 - `metric_pearson_gt<scalar_t>` for "Pearson" correlation between probability distributions.
 - `metric_haversine_gt<scalar_t>` for "Haversine" or "Great Circle" distance between coordinates used in GIS applications.
+- `metric_divergence_gt<scalar_t>` for the "Jensen Shannon" similarity between probability distributions.
 
 ## Multi-Threading
 
diff --git a/cpp/bench.cpp b/cpp/bench.cpp
@@ -447,6 +447,7 @@ struct args_t {
     bool metric_l2 = false;
     bool metric_cos = false;
     bool metric_haversine = false;
+    bool metric_divergence = false;
     bool metric_hamming = false;
     bool metric_tanimoto = false;
     bool metric_sorensen = false;
@@ -458,6 +459,8 @@ struct args_t {
             return metric_kind_t::cos_k;
         if (metric_haversine)
             return metric_kind_t::haversine_k;
+        if (metric_divergence)
+            return metric_kind_t::divergence_k;
         if (metric_hamming)
             return metric_kind_t::hamming_k;
         if (metric_tanimoto)
diff --git a/golang/lib.go b/golang/lib.go
@@ -21,6 +21,7 @@ const (
 	Cosine
 	L2sq
 	Haversine
+	Divergence
 	Pearson
 	Hamming
 	Tanimoto
@@ -37,6 +38,8 @@ func (m Metric) String() string {
 		return "cos"
 	case Haversine:
 		return "haversine"
+	case Divergence:
+		return "divergence"
 	case Pearson:
 		return "pearson"
 	case Hamming:
diff --git a/include/usearch/index_plugins.hpp b/include/usearch/index_plugins.hpp
@@ -112,6 +112,7 @@ enum class metric_kind_t : std::uint8_t {
     // Custom:
     pearson_k = 'p',
     haversine_k = 'h',
+    divergence_k = 'd',
 
     // Sets:
     jaccard_k = 'j',
@@ -239,6 +240,7 @@ inline char const* metric_kind_name(metric_kind_t metric) noexcept {
     case metric_kind_t::l2sq_k: return "l2sq";
     case metric_kind_t::pearson_k: return "pearson";
     case metric_kind_t::haversine_k: return "haversine";
+    case metric_kind_t::divergence_k: return "divergence";
     case metric_kind_t::jaccard_k: return "jaccard";
     case metric_kind_t::hamming_k: return "hamming";
     case metric_kind_t::tanimoto_k: return "tanimoto";
@@ -275,6 +277,8 @@ inline expected_gt<metric_kind_t> metric_from_name(char const* name, std::size_t
         parsed.result = metric_kind_t::cos_k;
     } else if (str_equals(name, len, "haversine")) {
         parsed.result = metric_kind_t::haversine_k;
+    } else if (str_equals(name, len, "divergence")) {
+        parsed.result = metric_kind_t::divergence_k;
     } else if (str_equals(name, len, "pearson")) {
         parsed.result = metric_kind_t::pearson_k;
     } else if (str_equals(name, len, "hamming")) {
@@ -284,8 +288,8 @@ inline expected_gt<metric_kind_t> metric_from_name(char const* name, std::size_t
     } else if (str_equals(name, len, "sorensen")) {
         parsed.result = metric_kind_t::sorensen_k;
     } else
-        parsed.failed(
-            "Unknown distance, choose: l2sq, ip, cos, haversine, jaccard, pearson, hamming, tanimoto, sorensen");
+        parsed.failed("Unknown distance, choose: l2sq, ip, cos, haversine, divergence, jaccard, pearson, hamming, "
+                      "tanimoto, sorensen");
     return parsed;
 }
 
@@ -1180,6 +1184,35 @@ template <typename scalar_at = float, typename result_at = float> struct metric_
     }
 };
 
+#include <cmath>
+#include <vector>
+
+/**
+ *  @brief  Measures Jensen-Shannon Divergence between two probability distributions.
+ */
+template <typename scalar_t = float, typename result_t = float> struct metric_divergence_gt {
+    using scalar_t = scalar_t;
+    using result_t = result_t;
+
+    inline result_t operator()(scalar_t const* p, scalar_t const* q, std::size_t dim) const noexcept {
+        result_t kld_pm{}, kld_qm{};
+        scalar_t epsilon = std::numeric_limits<scalar_t>::epsilon();
+#if USEARCH_USE_OPENMP
+#pragma omp simd reduction(+ : kld_pm, kld_qm)
+#elif defined(USEARCH_DEFINED_CLANG)
+#pragma clang loop vectorize(enable)
+#elif defined(USEARCH_DEFINED_GCC)
+#pragma GCC ivdep
+#endif
+        for (std::size_t i = 0; i != dim; ++i) {
+            scalar_t mi = (p[i] + q[i]) / 2 + epsilon;
+            kld_pm += p[i] * std::log((p[i] + epsilon) / mi);
+            kld_qm += q[i] * std::log((q[i] + epsilon) / mi);
+        }
+        return (kld_pm + kld_qm) / 2;
+    }
+};
+
 struct cos_i8_t {
     using scalar_t = i8_t;
     using result_t = f32_t;
@@ -1452,6 +1485,17 @@ class metric_punned_t {
             }
             break;
         }
+        case metric_kind_t::divergence_k: {
+            switch (scalar_kind_) {
+            case scalar_kind_t::f16_k:
+                raw_ptr_ = (punned_ptr_t)&equidimensional_<metric_divergence_gt<f16_t, f32_t>>;
+                break;
+            case scalar_kind_t::f32_k: raw_ptr_ = (punned_ptr_t)&equidimensional_<metric_divergence_gt<f32_t>>; break;
+            case scalar_kind_t::f64_k: raw_ptr_ = (punned_ptr_t)&equidimensional_<metric_divergence_gt<f64_t>>; break;
+            default: raw_ptr_ = nullptr; break;
+            }
+            break;
+        }
         case metric_kind_t::jaccard_k: // Equivalent to Tanimoto
         case metric_kind_t::tanimoto_k:
             raw_ptr_ = (punned_ptr_t)&equidimensional_<metric_tanimoto_gt<b1x8_t>>,
diff --git a/javascript/usearch.js b/javascript/usearch.js
@@ -11,6 +11,7 @@ const MetricKind = {
     IP: 'ip',
     L2sq: 'l2sq',
     Haversine: 'haversine',
+    Divergence: 'divergence',
     Pearson: 'pearson',
     Jaccard: 'jaccard',
     Hamming: 'hamming',
diff --git a/objc/USearchObjective.mm b/objc/USearchObjective.mm
@@ -33,6 +33,9 @@ metric_kind_t to_native_metric(USearchMetric m) {
         case USearchMetricHaversine:
             return metric_kind_t::haversine_k;
 
+        case USearchMetricDivergence:
+            return metric_kind_t::divergence_k;
+
         case USearchMetricJaccard:
             return metric_kind_t::jaccard_k;
 
diff --git a/python/lib.cpp b/python/lib.cpp
@@ -864,6 +864,7 @@ PYBIND11_MODULE(compiled, m) {
         .value("L2sq", metric_kind_t::l2sq_k)
 
         .value("Haversine", metric_kind_t::haversine_k)
+        .value("Divergence", metric_kind_t::divergence_k)
         .value("Pearson", metric_kind_t::pearson_k)
         .value("Jaccard", metric_kind_t::jaccard_k)
         .value("Hamming", metric_kind_t::hamming_k)
diff --git a/python/usearch/index.py b/python/usearch/index.py
@@ -151,6 +151,7 @@ def _normalize_metric(metric) -> MetricKind:
             "l2sq": MetricKind.L2sq,
             "l2_sq": MetricKind.L2sq,
             "haversine": MetricKind.Haversine,
+            "divergence": MetricKind.Divergence,
             "pearson": MetricKind.Pearson,
             "hamming": MetricKind.Hamming,
             "tanimoto": MetricKind.Tanimoto,
diff --git a/rust/lib.cpp b/rust/lib.cpp
@@ -107,6 +107,7 @@ metric_kind_t rust_to_cpp_metric(MetricKind value) {
     case MetricKind::Cos: return metric_kind_t::cos_k;
     case MetricKind::Pearson: return metric_kind_t::pearson_k;
     case MetricKind::Haversine: return metric_kind_t::haversine_k;
+    case MetricKind::Divergence: return metric_kind_t::divergence_k;
     case MetricKind::Hamming: return metric_kind_t::hamming_k;
     case MetricKind::Tanimoto: return metric_kind_t::tanimoto_k;
     case MetricKind::Sorensen: return metric_kind_t::sorensen_k;
diff --git a/simsimd b/simsimd
@@ -1 +1 @@
-Subproject commit 8f2c8881e440a55cfea246996984662623b4d5dd
+Subproject commit cdf2c1fb3c84f2a7bd71543eb13d6c0ccb46e9b2
