62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import numpy as np
|
|
import torch
|
|
import scipy.stats as scst
|
|
from torch.distributions import Poisson
|
|
|
|
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.ptdm_pmf_tensor import PtdmPmfTensor
|
|
|
|
|
|
class TestPtdfPmfTensor(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance."""
|
|
self.test_cases = {
|
|
"Poisson": [
|
|
(5, 3.0, scst.poisson.pmf(3.0, mu=5)),
|
|
(5, [3.0, 4.0], scst.poisson.pmf(np.array([3.0, 4.0], dtype=np.float32), mu=5))
|
|
],
|
|
}
|
|
|
|
def test_distributions(self):
|
|
"""Test pmf."""
|
|
for dist_name, cases in self.test_cases.items(): # For each distribution
|
|
for case in cases: # For each distribution's test cases
|
|
with self.subTest(distribution=dist_name, params=case):
|
|
if dist_name == "Poisson":
|
|
distribution = Poisson(torch.tensor(case[0]))
|
|
k_tensor = torch.tensor(case[1], dtype=torch.float32)
|
|
else:
|
|
raise ValueError("Only Poisson is supported in this unit test.")
|
|
|
|
pmf = PtdmPmfTensor().f(distribution, k_tensor)[0]
|
|
|
|
# Ensure return values are tensors
|
|
self.assertTrue(isinstance(pmf, torch.Tensor), "Probability output is not a tensor")
|
|
|
|
# Convert tensors to floats for comparison
|
|
if isinstance(case[1], float):
|
|
pmf = pmf.item()
|
|
|
|
# Check if computed values are approximately equal to expected values
|
|
self.assertAlmostEqual(pmf, case[-1], places=4,
|
|
msg=f"{dist_name}: Expected {case[-1]}, got {pmf}")
|
|
|
|
elif isinstance(case[1], list):
|
|
pmf = pmf.numpy()
|
|
self.assertTrue(
|
|
np.allclose(pmf, case[-1], atol=1e-5)
|
|
)
|
|
|
|
else:
|
|
raise ValueError("Unsupported data type.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|