63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
import torch
|
|
import numpy as np
|
|
|
|
from src import metrics
|
|
from src.metrics import divergence, evaluate_histogram
|
|
from sklearn.model_selection import GridSearchCV
|
|
from sklearn.neighbors import KernelDensity
|
|
|
|
def test_kl_divergence():
|
|
p = torch.randn(1000000)
|
|
q = 1.0 + 2.0 * torch.randn(1000000)
|
|
|
|
d1 = divergence(p, q, type='kl', n_components=0)
|
|
print(d1)
|
|
assert isinstance(d1, float)
|
|
d2 = divergence(p, q, type='kl', n_components=1)
|
|
print(d2)
|
|
assert isinstance(d2, float)
|
|
assert np.isclose(d1, d2, atol=1e-5)
|
|
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 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 = 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)
|