Revert "add gate_logits to zero inflated"

This reverts commit 7415a15256.
This commit is contained in:
Dr. Kashif Rasul
2020-10-21 16:44:44 +02:00
parent 7415a15256
commit 149a35a7f8
3 changed files with 45 additions and 122 deletions
+11 -23
View File
@@ -25,7 +25,7 @@ def test_zid_shape(gate_shape, base_shape):
gate = torch.rand(gate_shape)
base_dist = Normal(torch.randn(base_shape), torch.randn(base_shape).exp())
d = ZeroInflatedDistribution(base_dist, gate=gate)
d = ZeroInflatedDistribution(gate, base_dist)
assert d.batch_shape == broadcast_shape(gate_shape, base_shape)
assert d.support == base_dist.support
@@ -36,22 +36,19 @@ def test_zid_shape(gate_shape, base_shape):
@pytest.mark.parametrize("rate", [0.1, 0.5, 0.9, 1.0, 1.1, 2.0, 10.0])
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))
zip_ = ZeroInflatedPoisson(torch.zeros(1), torch.tensor(rate))
pois = Poisson(torch.tensor(rate))
s = pois.sample((20,))
zip1_prob = zip1.log_prob(s)
zip2_prob = zip2.log_prob(s)
zip_prob = zip_.log_prob(s)
pois_prob = pois.log_prob(s)
assert_close(zip1_prob, pois_prob, atol=1e-05)
assert_close(zip2_prob, pois_prob, atol=1e-05)
assert_close(zip_prob, pois_prob, atol=1e-06)
@pytest.mark.parametrize("gate", [0.0, 0.25, 0.5, 0.75, 1.0])
@pytest.mark.parametrize("rate", [0.1, 0.5, 0.9, 1.0, 1.1, 2.0, 10.0])
def test_zip_mean_variance(gate, rate):
num_samples = 1000000
zip_ = ZeroInflatedPoisson(torch.tensor(rate), gate=torch.tensor(gate))
zip_ = ZeroInflatedPoisson(torch.tensor(gate), torch.tensor(rate))
s = zip_.sample((num_samples,))
expected_mean = zip_.mean
estimated_mean = s.mean()
@@ -65,23 +62,14 @@ def test_zip_mean_variance(gate, rate):
@pytest.mark.parametrize("probs", [0.1, 0.5, 0.9])
def test_zinb_0_gate(total_count, probs):
# if gate is 0 ZINB is NegativeBinomial
zinb1 = ZeroInflatedNegativeBinomial(
total_count=torch.tensor(total_count),
gate=torch.zeros(1),
probs=torch.tensor(probs),
)
zinb2 = ZeroInflatedNegativeBinomial(
total_count=torch.tensor(total_count),
gate_logits=torch.tensor(-99.9),
probs=torch.tensor(probs),
zinb_ = ZeroInflatedNegativeBinomial(
torch.zeros(1), total_count=torch.tensor(total_count), probs=torch.tensor(probs)
)
neg_bin = NegativeBinomial(torch.tensor(total_count), probs=torch.tensor(probs))
s = neg_bin.sample((20,))
zinb1_prob = zinb1.log_prob(s)
zinb2_prob = zinb2.log_prob(s)
zinb_prob = zinb_.log_prob(s)
neg_bin_prob = neg_bin.log_prob(s)
assert_close(zinb1_prob, neg_bin_prob, atol=1e-05)
assert_close(zinb2_prob, neg_bin_prob, atol=1e-05)
assert_close(zinb_prob, neg_bin_prob, atol=1e-06)
@pytest.mark.parametrize("gate", [0.0, 0.25, 0.5, 0.75, 1.0])
@@ -90,8 +78,8 @@ def test_zinb_0_gate(total_count, probs):
def test_zinb_mean_variance(gate, total_count, logits):
num_samples = 1000000
zinb_ = ZeroInflatedNegativeBinomial(
torch.tensor(gate),
total_count=torch.tensor(total_count),
gate=torch.tensor(gate),
logits=torch.tensor(logits),
)
s = zinb_.sample((num_samples,))
@@ -100,4 +88,4 @@ def test_zinb_mean_variance(gate, total_count, logits):
expected_std = zinb_.stddev
estimated_std = s.std()
assert_close(expected_mean, estimated_mean, atol=1e-01)
assert_close(expected_std, estimated_std, atol=1e-01)
assert_close(expected_std, estimated_std, atol=1e-1)