File size: 2,728 Bytes
9d901ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """Tests for NB distribution utilities."""
import numpy as np
import pytest
import torch
from scptr.deep._distributions import log_nb_positive
class TestLogNBPositive:
def test_output_shape(self):
x = torch.tensor([[1.0, 2.0, 3.0]])
mu = torch.tensor([[2.0, 2.0, 2.0]])
theta = torch.tensor([5.0, 5.0, 5.0])
ll = log_nb_positive(x, mu, theta)
assert ll.shape == (1, 3)
def test_non_positive(self):
"""Log-probabilities should be <= 0."""
x = torch.randint(0, 20, (100, 50)).float()
mu = torch.rand(100, 50) * 10 + 0.1
theta = torch.rand(50) * 10 + 0.1
ll = log_nb_positive(x, mu, theta)
assert (ll <= 1e-5).all(), "Log probabilities should be non-positive"
def test_peak_at_mean(self):
"""For integer means, likelihood should peak near the mean."""
mu = torch.tensor([[10.0]])
theta = torch.tensor([50.0]) # low dispersion
xs = torch.arange(0, 30).float().unsqueeze(1)
ll = log_nb_positive(xs, mu.expand(30, 1), theta)
peak = ll.argmax().item()
assert abs(peak - 10) <= 2, f"Peak at {peak}, expected near 10"
def test_higher_theta_less_variance(self):
"""Higher theta (less dispersion) should give sharper distribution."""
mu = torch.tensor([[5.0]])
xs = torch.arange(0, 20).float().unsqueeze(1)
theta_low = torch.tensor([1.0])
theta_high = torch.tensor([100.0])
ll_low = log_nb_positive(xs, mu.expand(20, 1), theta_low)
ll_high = log_nb_positive(xs, mu.expand(20, 1), theta_high)
# High theta should have higher peak probability
assert ll_high.max() > ll_low.max()
def test_gradient_flows(self):
"""Ensure gradients flow through all parameters."""
x = torch.tensor([[3.0, 5.0]])
mu = torch.tensor([[2.0, 4.0]], requires_grad=True)
theta = torch.tensor([5.0, 5.0], requires_grad=True)
ll = log_nb_positive(x, mu, theta).sum()
ll.backward()
assert mu.grad is not None
assert theta.grad is not None
assert not torch.isnan(mu.grad).any()
assert not torch.isnan(theta.grad).any()
def test_batch_consistency(self):
"""Batched computation should match individual computation."""
torch.manual_seed(42)
x = torch.randint(0, 10, (5, 3)).float()
mu = torch.rand(5, 3) * 5 + 0.1
theta = torch.rand(3) * 5 + 0.1
ll_batch = log_nb_positive(x, mu, theta)
for i in range(5):
ll_single = log_nb_positive(
x[i : i + 1], mu[i : i + 1], theta
)
assert torch.allclose(ll_batch[i], ll_single[0], atol=1e-5)
|