85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch
|
|
import scipy.stats as scst
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
|
|
from pytorch_wrapper.ptd_binomial import PtdBinomial
|
|
from pytorch_wrapper.ptd_poisson import PtdPoisson
|
|
from pytorch_wrapper.ptdm_pmf import PtdmPmf
|
|
|
|
|
|
class TestPtdfPmf(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance."""
|
|
# Binomial
|
|
test_cases = []
|
|
n = 10
|
|
p = 0.5
|
|
for k in range(n+1):
|
|
e = (n, k, p, scst.binom.pmf(k, n, p))
|
|
test_cases.append(e)
|
|
|
|
p = 0.01
|
|
for k in range(n+1):
|
|
e = (n, k, p, scst.binom.pmf(k, n, p))
|
|
test_cases.append(e)
|
|
|
|
p = 0.99
|
|
for k in range(n+1):
|
|
e = (n, k, p, scst.binom.pmf(k, n, p))
|
|
test_cases.append(e)
|
|
|
|
self.test_cases_binomial = test_cases
|
|
|
|
# Poisson
|
|
test_cases = []
|
|
lamb = 3 # rate
|
|
n = 10
|
|
for k in range(n+1):
|
|
e = (k, lamb, scst.poisson.pmf(k, lamb))
|
|
test_cases.append(e)
|
|
self.test_cases_poisson = test_cases
|
|
|
|
def test_binomial_distribution(self):
|
|
"""Test binomial probability mass."""
|
|
for n, k, p, expected_mass in self.test_cases_binomial:
|
|
with self.subTest(n=n, k=k, p=p):
|
|
distribution_node = PtdBinomial().f(str(n), str(p), "")[0]
|
|
prob = PtdmPmf().f(distribution_node, k)[0]
|
|
|
|
# Ensure return values are tensors
|
|
self.assertTrue(isinstance(prob, torch.Tensor), "Probability output is not a tensor")
|
|
|
|
# Convert tensors to floats for comparison
|
|
prob = prob.item()
|
|
|
|
# Check if computed values are approximately equal to expected values
|
|
self.assertAlmostEqual(prob, expected_mass, places=5,
|
|
msg=f"Expected probability {expected_mass}, got {prob}")
|
|
|
|
def test_poisson_distribution(self):
|
|
"""Test poisson probability."""
|
|
for k, lamb, expected_mass in self.test_cases_poisson:
|
|
with self.subTest(k=k):
|
|
distribution_node = PtdPoisson().f(str(lamb))[0]
|
|
prob = PtdmPmf().f(distribution_node, k)[0]
|
|
|
|
# Ensure return values are tensors
|
|
self.assertTrue(isinstance(prob, torch.Tensor), "Probability output is not a tensor")
|
|
|
|
# Convert tensors to floats for comparison
|
|
prob = prob.item()
|
|
|
|
# Check if computed values are approximately equal to expected values
|
|
self.assertAlmostEqual(prob, expected_mass, places=5,
|
|
msg=f"Expected probability {expected_mass}, got {prob}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|