106 lines
5.5 KiB
Python
106 lines
5.5 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch
|
|
import scipy.stats as scst
|
|
import numpy as np
|
|
|
|
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_cdf import PtdmCdf
|
|
from torch.distributions import (
|
|
Categorical, Poisson, NegativeBinomial, Multinomial,
|
|
Bernoulli, Beta, Gamma, Normal, Laplace, Exponential, LogNormal,
|
|
Pareto, Weibull, Dirichlet, VonMises, StudentT, Uniform, Chi2
|
|
)
|
|
|
|
|
|
class TestPtdfCdf(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance.
|
|
|
|
Following distributions are tested in separate unit test files:
|
|
* Bernoulli
|
|
* Binomial
|
|
* Poisson
|
|
* StudentT
|
|
* Beta
|
|
"""
|
|
self.test_cases = {
|
|
"Uniform": [(0.0, 1.0, 0.5, scst.uniform.cdf(0.5, loc=0, scale=1))],
|
|
"Normal": [(0.0, 1.0, -1.5, scst.norm.cdf(-1.5, 0, 1))],
|
|
"Chi2": [(1.0, 2.0, scst.chi2.cdf(2.0, 1))], # df=1, value=2.0
|
|
"Exponential": [(1.0, 3.0, scst.expon.cdf(3.0, scale=1))],
|
|
"Gamma": [(2.0, 1.0, 3.0, scst.gamma.cdf(3.0, 2, scale=1))],
|
|
# "Dirichlet": [(torch.tensor([1.0, 1.0, 1.0]), torch.tensor([0.2, 0.5, 0.3]), scst.dirichlet.cdf([0.2, 0.5, 0.3], [1.0, 1.0, 1.0]))],
|
|
# "VonMises": [(0.0, 1.0, 1.0, scst.vonmises.cdf(1.0, kappa=1.0, loc=0))],
|
|
}
|
|
|
|
def test_distributions(self):
|
|
"""Test probability mass and density functions for all PyTorch distributions."""
|
|
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):
|
|
# Instantiate PyTorch distribution directly if not in wrapper
|
|
if dist_name == "Categorical":
|
|
distribution = Categorical(case[0])
|
|
k_tensor = torch.tensor(case[1], dtype=torch.int64)
|
|
elif dist_name == "NegativeBinomial":
|
|
distribution = NegativeBinomial(torch.tensor(case[0]), torch.tensor(case[2]))
|
|
k_tensor = torch.tensor(case[1], dtype=torch.int64)
|
|
elif dist_name == "Gamma":
|
|
distribution = Gamma(torch.tensor(case[0]), torch.tensor(case[1]))
|
|
k_tensor = torch.tensor(case[2], dtype=torch.float32)
|
|
elif dist_name == "Normal":
|
|
distribution = Normal(torch.tensor(case[0]), torch.tensor(case[1]))
|
|
k_tensor = torch.tensor(case[2], dtype=torch.float32)
|
|
elif dist_name == "Laplace":
|
|
distribution = Laplace(torch.tensor(case[0]), torch.tensor(case[1]))
|
|
k_tensor = torch.tensor(case[2], dtype=torch.float32)
|
|
elif dist_name == "Exponential":
|
|
distribution = Exponential(torch.tensor(case[0]))
|
|
k_tensor = torch.tensor(case[1], dtype=torch.float32)
|
|
elif dist_name == "LogNormal":
|
|
distribution = LogNormal(torch.tensor(case[0]), torch.tensor(case[1]))
|
|
k_tensor = torch.tensor(case[2], dtype=torch.float32)
|
|
elif dist_name == "Pareto":
|
|
distribution = Pareto(torch.tensor(case[0]), torch.tensor(case[1]))
|
|
k_tensor = torch.tensor(case[2], dtype=torch.float32)
|
|
elif dist_name == "Weibull":
|
|
distribution = Weibull(torch.tensor(1.0), torch.tensor(case[0]))
|
|
k_tensor = torch.tensor(case[1], dtype=torch.float32)
|
|
elif dist_name == "Dirichlet":
|
|
distribution = Dirichlet(case[0])
|
|
k_tensor = case[1] # Dirichlet expects tensor input
|
|
elif dist_name == "VonMises":
|
|
distribution = VonMises(torch.tensor(case[0]), torch.tensor(case[1]))
|
|
k_tensor = torch.tensor(case[2], dtype=torch.float32)
|
|
elif dist_name == "Uniform":
|
|
distribution = Uniform(torch.tensor(case[0]), torch.tensor(case[1]))
|
|
k_tensor = torch.tensor(case[2], dtype=torch.float32)
|
|
elif dist_name == "Chi2":
|
|
distribution = Chi2(torch.tensor(case[0])) # df
|
|
k_tensor = torch.tensor(case[1], dtype=torch.float32)
|
|
else:
|
|
self.fail(f"Unhandled distribution: {dist_name}")
|
|
|
|
# Call wrapper log-prob method
|
|
if dist_name == "Dirichlet":
|
|
cdf = PtdmCdf().f(distribution, str(k_tensor.tolist()))[0]
|
|
else:
|
|
cdf = PtdmCdf().f(distribution, str(k_tensor.item()))[0]
|
|
|
|
# Ensure return values are tensors
|
|
self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor")
|
|
|
|
# Convert tensors to floats for comparison
|
|
cdf = cdf.item()
|
|
|
|
# Check if computed values are approximately equal to expected values
|
|
self.assertAlmostEqual(cdf, case[-1], places=4,
|
|
msg=f"{dist_name}: Expected {case[-1]}, got {cdf}")
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|