Implement histogram based kl divergence computation

tests/test_metrics.py

@@ -2,7 +2,7 @@ import torch
 import numpy as np
 
 from src import metrics
-from src.metrics import divergence
+from src.metrics import divergence, evaluate_histogram
 from sklearn.model_selection import GridSearchCV
 from sklearn.neighbors import KernelDensity
 
@@ -20,39 +20,43 @@ def test_kl_divergence():
     d3 = divergence(p, q, type='kl', n_components=3)
     print(d3)
     assert isinstance(d3, float)
+    d4 = divergence(p, q, type='kl', n_components=-1)
+    assert isinstance(d4, float)
+    d5 = divergence(q, p, type='kl', n_components=-1)
+    assert isinstance(d5, float)
 
+    p = 1000000 * torch.randn(1000)
+    q = torch.randn(1000)
+    d6 = divergence(p, q, type='kl', n_components=-1)
+    assert np.isnan(d6)
 
-def evaluate_histogram(x, hist, bin_edges):
-    idx = np.digitize(x, bin_edges)
-    mask = np.logical_and(np.less(0, idx), np.less(idx, len(bin_edges)))
-    r = np.zeros_like(x)
-    r[mask] = hist[idx[mask] - 1]
+def test_evaluate_histogram():
+    N = 10000
+    p = torch.randn(N)
+    q = 1.0 + 2.0 * torch.randn(2*N)
+    
+    p_hist, p_edges = np.histogram(p.unsqueeze(-1), bins='auto', density=True)
+    q_hist, q_edges = np.histogram(q.unsqueeze(-1), bins='auto', density=True)
+    px = evaluate_histogram(p, p_hist, p_edges)
+    assert px.shape == p.shape
+    qx = evaluate_histogram(q, p_hist, p_edges)
+    assert qx.shape == q.shape
+    px = evaluate_histogram(p, q_hist, q_edges)
+    assert px.shape == p.shape
+    qx = evaluate_histogram(q, q_hist, q_edges)
+    assert qx.shape == q.shape
 
 
 if __name__ == '__main__':
-    # p = torch.randn(10)
-    # q = 1.0 + 2.0 * torch.randn(10)
+    p = torch.randn(10)
+    q = 1.0 + 2.0 * torch.randn(10)
 
-    # p = p.unsqueeze(-1)
-    # q = q.unsqueeze(-1)
-    # p_hist, p_edges = np.histogram(p.unsqueeze(-1), bins='auto', density=True)
-    # q_hist, q_edges = np.histogram(q.unsqueeze(-1), bins='auto', density=True)
-    # # px = p_hist[np.digitize(p, p_edges) - 1]
+    p = p.unsqueeze(-1)
+    q = q.unsqueeze(-1)
+    p_hist, p_edges = np.histogram(p.unsqueeze(-1), bins='auto', density=True)
+    q_hist, q_edges = np.histogram(q.unsqueeze(-1), bins='auto', density=True)
+    # px = p_hist[np.digitize(p, p_edges) - 1]
 
-    # # qx = q_hist[np.digitize(p, q_edges) - 1]
-    # px = evaluate_histogram(p, p_hist, p_edges)
-    # qx = evaluate_histogram(q, p_hist, p_edges)
-
-    # use grid search cross-validation to optimize the bandwidth
-    params = {'bandwidth': np.logspace(-1, 1, 3)}
-    grid = GridSearchCV(KernelDensity(), params)
-    grid.fit(p)
-    p_kde = grid.best_estimator_
-    grid = GridSearchCV(KernelDensity(), params)
-    grid.fit(q)
-    q_kde = grid.best_estimator_
-    px = p_kde.score_samples(p)
-    qx = q_kde.score_samples(p)
-    d = np.mean(px - qx).item()
-    print(d)
-    test_kl_divergence()
+    # qx = q_hist[np.digitize(p, q_edges) - 1]
+    px = evaluate_histogram(p, p_hist, p_edges)
+    qx = evaluate_histogram(q, p_hist, p_edges)