def test_zip_1_gate(rate):
# if gate is 1 ZIP is Delta(0)
zip_ = ZeroInflatedPoisson(torch.ones(1), torch.tensor(rate))
delta = Delta(torch.zeros(1))
s = torch.tensor([0., 1.])
zip_prob = zip_.log_prob(s)
delta_prob = delta.log_prob(s)
assert_tensors_equal(zip_prob, delta_prob)
def test_zip_0_gate(rate):
# if gate is 0 ZIP is Poisson
zip_ = ZeroInflatedPoisson(torch.zeros(1), torch.tensor(rate))
pois = Poisson(torch.tensor(rate))
s = pois.sample((20, ))
zip_prob = zip_.log_prob(s)
pois_prob = pois.log_prob(s)
assert_tensors_equal(zip_prob, pois_prob)
def test_zip_mean_variance(gate, rate):
num_samples = 1000000
zip_ = ZeroInflatedPoisson(torch.tensor(gate), torch.tensor(rate))
s = zip_.sample((num_samples, ))
expected_mean = zip_.mean
estimated_mean = s.mean()
expected_std = zip_.stddev
estimated_std = s.std()
assert_tensors_equal(expected_mean, estimated_mean, prec=1e-02)
assert_tensors_equal(expected_std, estimated_std, prec=1e-02)
def test_log_prob_d2(eta):
dist = LKJCorrCholesky(2, torch.tensor([eta]))
test_dist = TransformedDistribution(Beta(eta, eta),
AffineTransform(loc=-1., scale=2.0))
samples = dist.sample(torch.Size([100]))
lp = dist.log_prob(samples)
x = samples[..., 1, 0]
tst = test_dist.log_prob(x)
assert_tensors_equal(lp, tst, prec=1e-6)
def test_log_prob_d2(concentration):
dist = LKJCholesky(2, torch.tensor([concentration]))
test_dist = TransformedDistribution(Beta(concentration, concentration),
AffineTransform(loc=-1., scale=2.0))
samples = dist.sample(torch.Size([100]))
lp = dist.log_prob(samples)
x = samples[..., 1, 0]
tst = test_dist.log_prob(x)
# LKJ prevents inf values in log_prob
lp[tst == math.inf] = math.inf # substitute inf for comparison
assert_tensors_equal(lp, tst, prec=1e-3)
def test_corr_cholesky_transform(x_shape, mapping):
transform = mapping(constraints.corr_cholesky)
x = torch.randn(x_shape, requires_grad=True).clamp(-2, 2)
y = transform(x)
# test codomain
assert (transform.codomain.check(y) == 1).all()
# test inv
z = transform.inv(y)
assert_tensors_equal(x, z, prec=1e-4)
# test domain
assert (transform.domain.check(z) == 1).all()
# test log_abs_det_jacobian
log_det = transform.log_abs_det_jacobian(x, y)
assert log_det.shape == x_shape[:-1]
def test_unconstrained_to_corr_cholesky_transform(y_shape):
transform = transforms.CorrLCholeskyTransform()
y = torch.empty(y_shape).uniform_(-4, 4).requires_grad_()
x = transform(y)
# test codomain
assert (transform.codomain.check(x) == 1).all()
# test inv
y_prime = transform.inv(x)
assert_tensors_equal(y, y_prime, prec=1e-4)
# test domain
assert (transform.domain.check(y_prime) == 1).all()
# test log_abs_det_jacobian
log_det = transform.log_abs_det_jacobian(y, x)
assert log_det.shape == y_shape[:-1]
if len(y_shape) == 1:
triu_index = x.new_ones(x.shape).triu(diagonal=1).to(torch.bool)
x_tril_vector = x.t()[triu_index]
assert_tensors_equal(_autograd_log_det(x_tril_vector, y),
log_det,
prec=1e-4)
x_tril_vector = x_tril_vector.detach().requires_grad_()
x = x.new_zeros(x.shape)
x[triu_index] = x_tril_vector
x = x.t()
z = transform.inv(x)
assert_tensors_equal(_autograd_log_det(z, x_tril_vector),
-log_det,
prec=1e-4)
