def test_zip_1_gate(rate):
# if gate is 1 ZIP is Delta(0)
zip1 = ZeroInflatedPoisson(torch.tensor(rate), gate=torch.ones(1))
zip2 = ZeroInflatedPoisson(torch.tensor(rate),
gate_logits=torch.tensor(math.inf))
delta = Delta(torch.zeros(1))
s = torch.tensor([0.0, 1.0])
zip1_prob = zip1.log_prob(s)
zip2_prob = zip2.log_prob(s)
delta_prob = delta.log_prob(s)
assert_close(zip1_prob, delta_prob)
assert_close(zip2_prob, delta_prob)
def test_zip_0_gate(rate):
# if gate is 0 ZIP is Poisson
zip1 = ZeroInflatedPoisson(torch.tensor(rate), gate=torch.zeros(1))
zip2 = ZeroInflatedPoisson(torch.tensor(rate),
gate_logits=torch.tensor(-99.9))
pois = Poisson(torch.tensor(rate))
s = pois.sample((20, ))
zip1_prob = zip1.log_prob(s)
zip2_prob = zip2.log_prob(s)
pois_prob = pois.log_prob(s)
assert_close(zip1_prob, pois_prob)
assert_close(zip2_prob, pois_prob)
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)