From 1ee46214a71c46747a475835723db24b386cd17e Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Tue, 3 Aug 2021 17:03:37 +0200 Subject: [PATCH] Implement jenson shannon divergence --- src/metrics.py | 58 ++++++++++++++++++++++++++++++++----------- tests/test_metrics.py | 21 ++++++++++++++-- 2 files changed, 62 insertions(+), 17 deletions(-) diff --git a/src/metrics.py b/src/metrics.py index 56a157e..8819df7 100644 --- a/src/metrics.py +++ b/src/metrics.py @@ -96,13 +96,15 @@ def average_velocity(states): arg_v = nanmean(vehicle_avg_v) return arg_v -def divergence(p, q, type='kl', n_components=0): +def divergence(p, q, type='js', n_components=0): """ Calculate a divergence between p and q Args: p (torch.tensor): (n) samples from p q (torch.tensor): (m) samples from q type (str): divergence to use + 'kl': Kullback-Leibler divergence KL(p||q) + 'js': Jensen-Shannon divergence (symmetric KLD) n_components (int): method to use to compute kl divergence n_components < 0: approximate samples with histogram density n_components == 0: approximate samples by Gaussian distributions and compute analytically @@ -110,22 +112,23 @@ def divergence(p, q, type='kl', n_components=0): Returns: d (float): approximate divergence """ - if type == 'kl': + if type == 'js': + # Use histogram binning to discretize sampled distributions + p_hist = np.histogram(p, bins='auto', density=True) + q_hist = np.histogram(q, bins='auto', density=True) + m = torch.cat([p, q], dim=0) + m_weights = torch.cat([torch.full_like(p, 1./len(p)), torch.full_like(q, 1./len(q))], dim=0) + m_bins = np.sort(np.concatenate([p_hist[1], q_hist[1]])) + m_hist = np.histogram(m, bins=m_bins, density=True, weights=m_weights) + d = .5 * kl_histogram(p, p_hist, m_hist) + .5 * kl_histogram(q, q_hist, m_hist) + return d + + elif type == 'kl': if n_components < 0: # Use histogram binning to discretize sampled distributions - 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) - qx = evaluate_histogram(p, q_hist, q_edges) - p_supp = ~np.isclose(px, 0.0) - q_supp = ~np.isclose(qx, 0.0) - if np.any(np.logical_and(p_supp, ~q_supp)): - # if not support(p) subset support(q) - return np.nan - elif ~np.any(p_supp): - # if p is zero everywhere - return 0. - d = np.mean(np.log(px[p_supp] / qx[p_supp])) + p_hist = np.histogram(p, bins='auto', density=True) + q_hist = np.histogram(q, bins='auto', density=True) + d = kl_histogram(p, p_hist, q_hist) return d elif n_components == 0: # Assume p and q to be Gaussian @@ -148,6 +151,31 @@ def divergence(p, q, type='kl', n_components=0): else: raise NotImplementedError("Please implement divergence for type '{}'".format(type)) +def kl_histogram(p_sample, p_hist, q_hist): + """ + Calculate the kl divergence between p and q based on a histogram representation + Args: + p_sample (torch.tensor): (n) samples from p + p_hist (tuple): result of np.histogram(density=True) for samples from p + q_hist (tuple): result of np.histogram(density=True) for samples from q + Returns: + d (float): approximate KL divergence + """ + p_density, p_edges = p_hist + q_density, q_edges = q_hist + px = evaluate_histogram(p_sample, p_density, p_edges) + qx = evaluate_histogram(p_sample, q_density, q_edges) + p_supp = ~np.isclose(px, 0.0) + q_supp = ~np.isclose(qx, 0.0) + if np.any(np.logical_and(p_supp, ~q_supp)): + # if not support(p) subset support(q) + return np.inf + elif ~np.any(p_supp): + # if p is zero everywhere + return 0. + d = np.mean(np.log(px[p_supp] / qx[p_supp])) + return d + def kl_normal(pm, pv, qm, qv): """ diff --git a/tests/test_metrics.py b/tests/test_metrics.py index bef990d..0806e97 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -25,10 +25,10 @@ def test_kl_divergence(): d5 = divergence(q, p, type='kl', n_components=-1) assert isinstance(d5, float) - p = 1000000 * torch.randn(1000) + p = 1000 * torch.randn(1000) q = torch.randn(1000) d6 = divergence(p, q, type='kl', n_components=-1) - assert np.isnan(d6) + assert ~np.isfinite(d6) def test_evaluate_histogram(): N = 10000 @@ -47,6 +47,23 @@ def test_evaluate_histogram(): assert qx.shape == q.shape +def test_js_divergence(): + N = 1000 + p = torch.randn(N) + q = 1.0 + 2.0 * torch.randn(2*N) + + d1 = divergence(p, q, type='js') + d2 = divergence(q, p, type='js') + assert d1 == d2 + + p = 1000 * torch.randn(1000) + q = torch.randn(1000) + d1 = divergence(p, q, type='js') + d2 = divergence(q, p, type='js') + assert d1 == d2 + assert np.isfinite(d1) + + if __name__ == '__main__': p = torch.randn(10) q = 1.0 + 2.0 * torch.randn(10)