diff --git a/__init__.py b/__init__.py index 6ced78c..6388e25 100644 --- a/__init__.py +++ b/__init__.py @@ -124,9 +124,13 @@ from .modules.pytorch_wrapper.ptd_poisson import PtdPoisson from .modules.pytorch_wrapper.ptd_student_t import PtdStudentT from .modules.pytorch_wrapper.ptd_uniform import PtdUniform from .modules.pytorch_wrapper.ptdm_cdf import PtdmCdf +from .modules.pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor from .modules.pytorch_wrapper.ptdm_log_prob import PtdmLogProb +from .modules.pytorch_wrapper.ptdm_log_prob_tensor import PtdmLogProbTensor from .modules.pytorch_wrapper.ptdm_pdf import PtdmPdf +from .modules.pytorch_wrapper.ptdm_pdf_tensor import PtdmPdfTensor from .modules.pytorch_wrapper.ptdm_pmf import PtdmPmf +from .modules.pytorch_wrapper.ptdm_pmf_tensor import PtdmPmfTensor from .modules.pytorch_wrapper.ptdm_sample import PtdmSample from .modules.pytorch_wrapper.ptf_gelu import PtfGELU from .modules.pytorch_wrapper.ptf_leaky_relu import PtfLeakyReLU @@ -294,9 +298,13 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtdStudentT": PtdStudentT, "PtdUniform": PtdUniform, "PtdmCdf": PtdmCdf, + "PtdmCdfTensor": PtdmCdfTensor, "PtdmLogProb": PtdmLogProb, + "PtdmLogProbTensor": PtdmLogProbTensor, "PtdmPdf": PtdmPdf, + "PtdmPdfTensor": PtdmPdfTensor, "PtdmPmf": PtdmPmf, + "PtdmPmfTensor": PtdmPmfTensor, "PtdmSample": PtdmSample, "PtfGELU": PtfGELU, "PtfLeakyReLU": PtfLeakyReLU, @@ -464,9 +472,13 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtdStudentT": "Ptd Student T", "PtdUniform": "Ptd Uniform", "PtdmCdf": "Ptdm Cdf", + "PtdmCdfTensor": "Ptdm Cdf Tensor", "PtdmLogProb": "Ptdm Log Prob", + "PtdmLogProbTensor": "Ptdm Log Prob Tensor", "PtdmPdf": "Ptdm Pdf", + "PtdmPdfTensor": "Ptdm Pdf Tensor", "PtdmPmf": "Ptdm Pmf", + "PtdmPmfTensor": "Ptdm Pmf Tensor", "PtdmSample": "Ptdm Sample", "PtfGELU": "Ptf GELU", "PtfLeakyReLU": "Ptf Leaky ReLU", diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index ddc55a7..41e480d 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -75,9 +75,13 @@ | [Ptd Student T](ptd_student_t.md) | Instantiates a StudentT distribution object. | | [Ptd Uniform](ptd_uniform.md) | Instantiates a Uniform distribution object. | | [Ptdm Cdf](ptdm_cdf.md) | Computes the cumulative distribution function for the input distribution. | +| [Ptdm Cdf Tensor](ptdm_cdf_tensor.md) | Computes the cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. | | [Ptdm Log Prob](ptdm_log_prob.md) | Computes the log of probability for the input distribution. | +| [Ptdm Log Prob Tensor](ptdm_log_prob_tensor.md) | Computes the log of probability for the input distribution. This nodes accepts a tensor so it can be used to compute log of probability for multiple values contained in a tensor. | | [Ptdm Pdf](ptdm_pdf.md) | Computes the probability density for the input distribution. | +| [Ptdm Pdf Tensor](ptdm_pdf_tensor.md) | Computes the probability density for the input distribution. This nodes accepts a tensor so it can be used to compute pdf for multiple values contained in a tensor. | | [Ptdm Pmf](ptdm_pmf.md) | Computes the probability for the input distribution. | +| [Ptdm Pmf Tensor](ptdm_pmf_tensor.md) | Computes the probability for the input distribution. This nodes accepts a tensor so it can be used to compute pmf for multiple values contained in a tensor. | | [Ptdm Sample](ptdm_sample.md) | Samples from the input distribution. | ## PyTorch wrapper - Image processing | Node | Description | diff --git a/docs/reference/ptdm_cdf.md b/docs/reference/ptdm_cdf.md index 819ff73..f8f496f 100644 --- a/docs/reference/ptdm_cdf.md +++ b/docs/reference/ptdm_cdf.md @@ -6,13 +6,13 @@ Cdf is not supported for all distributions in PyTorch. Args: distribution (torch.distributions.distribution.Distribution): Distribution. - v (float): Value + x (float): Value ## Input | Name | Data type | |---|---| | distribution | Ptdistribution | -| v | String | +| x | String | ## Output | Data type | diff --git a/docs/reference/ptdm_cdf_tensor.md b/docs/reference/ptdm_cdf_tensor.md new file mode 100644 index 0000000..ba70f28 --- /dev/null +++ b/docs/reference/ptdm_cdf_tensor.md @@ -0,0 +1,25 @@ +# Ptdm Cdf Tensor +Computes the cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. + +**Note** +Cdf is not supported for all distributions in PyTorch. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| tens | Tensor | + +## Output +| Data type | +|---| +| Tensor | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptdm_log_prob.md b/docs/reference/ptdm_log_prob.md index c98aa7c..61b1b29 100644 --- a/docs/reference/ptdm_log_prob.md +++ b/docs/reference/ptdm_log_prob.md @@ -9,7 +9,7 @@ Args: | Name | Data type | |---|---| | distribution | Ptdistribution | -| value | String | +| x | String | ## Output | Data type | diff --git a/docs/reference/ptdm_log_prob_tensor.md b/docs/reference/ptdm_log_prob_tensor.md new file mode 100644 index 0000000..b097405 --- /dev/null +++ b/docs/reference/ptdm_log_prob_tensor.md @@ -0,0 +1,22 @@ +# Ptdm Log Prob Tensor +Computes the log of probability for the input distribution. This nodes accepts a tensor so it can be used to compute log of probability for multiple values contained in a tensor. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| tens | Tensor | + +## Output +| Data type | +|---| +| Tensor | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptdm_pdf.md b/docs/reference/ptdm_pdf.md index eef9f7a..2baf2ac 100644 --- a/docs/reference/ptdm_pdf.md +++ b/docs/reference/ptdm_pdf.md @@ -3,13 +3,13 @@ Computes the probability density for the input distribution. Args: distribution (torch.distributions.distribution.Distribution): Distribution. - v (float): Value + x (float): Value ## Input | Name | Data type | |---|---| | distribution | Ptdistribution | -| v | String | +| x | String | ## Output | Data type | diff --git a/docs/reference/ptdm_pdf_tensor.md b/docs/reference/ptdm_pdf_tensor.md new file mode 100644 index 0000000..e732829 --- /dev/null +++ b/docs/reference/ptdm_pdf_tensor.md @@ -0,0 +1,22 @@ +# Ptdm Pdf Tensor +Computes the probability density for the input distribution. This nodes accepts a tensor so it can be used to compute pdf for multiple values contained in a tensor. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| tens | Tensor | + +## Output +| Data type | +|---| +| Tensor | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptdm_pmf_tensor.md b/docs/reference/ptdm_pmf_tensor.md new file mode 100644 index 0000000..59cd787 --- /dev/null +++ b/docs/reference/ptdm_pmf_tensor.md @@ -0,0 +1,22 @@ +# Ptdm Pmf Tensor +Computes the probability for the input distribution. This nodes accepts a tensor so it can be used to compute pmf for multiple values contained in a tensor. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): k(s) in Tensor. + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| tens | Tensor | + +## Output +| Data type | +|---| +| Tensor | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/modules/pytorch_wrapper/ptdm_cdf.py b/modules/pytorch_wrapper/ptdm_cdf.py index 348aeb8..cf00d42 100644 --- a/modules/pytorch_wrapper/ptdm_cdf.py +++ b/modules/pytorch_wrapper/ptdm_cdf.py @@ -12,7 +12,7 @@ class PtdmCdf: Args: distribution (torch.distributions.distribution.Distribution): Distribution. - v (float): Value + x (float): Value category: PyTorch wrapper - Distribution """ @@ -28,7 +28,7 @@ class PtdmCdf: return { "required": { "distribution": ("PTDISTRIBUTION", {}), - "v": ("STRING", {"default": ""}) + "x": ("STRING", {"default": ""}) } } @@ -36,18 +36,18 @@ class PtdmCdf: FUNCTION: str = "f" CATEGORY: str = "Distribution" - def f(self, distribution: torch.distributions.Distribution, v: str) -> tuple: + def f(self, distribution: torch.distributions.Distribution, x: str) -> tuple: """ Computes the cumulative distribution function for the input distribution. Args: distribution (torch.distributions.Distribution): Distribution. - v (str): Value + x (str): Value Returns: tuple containing the probability """ - v2 = ast.literal_eval(v) + v2 = ast.literal_eval(x) if isinstance(v2, int) is False and isinstance(v2, float) is False : if isinstance(distribution, Dirichlet) is False: raise ValueError("You need to specify a float or an int.") diff --git a/modules/pytorch_wrapper/ptdm_cdf_tensor.py b/modules/pytorch_wrapper/ptdm_cdf_tensor.py new file mode 100644 index 0000000..815a83b --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_cdf_tensor.py @@ -0,0 +1,49 @@ +from typing import Any, Dict +import torch + + +class PtdmCdfTensor: + """ + Computes the cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. + + **Note** + Cdf is not supported for all distributions in PyTorch. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + + category: PyTorch wrapper - Distribution + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying the required input types. + """ + return { + "required": { + "distribution": ("PTDISTRIBUTION", {}), + "tens": ("TENSOR", {}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, tens: torch.Tensor) -> tuple: + """ + Computes the cumulative distribution function for the input distribution. + + Args: + distribution (torch.distributions.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + + Returns: + tuple containing the probability + """ + return (distribution.cdf(tens),) diff --git a/modules/pytorch_wrapper/ptdm_log_prob.py b/modules/pytorch_wrapper/ptdm_log_prob.py index 73826e0..4ff21c9 100644 --- a/modules/pytorch_wrapper/ptdm_log_prob.py +++ b/modules/pytorch_wrapper/ptdm_log_prob.py @@ -25,7 +25,7 @@ class PtdmLogProb: return { "required": { "distribution": ("PTDISTRIBUTION", {}), - "value": ("STRING", {"default": ""}) + "x": ("STRING", {"default": ""}) } } @@ -35,27 +35,27 @@ class PtdmLogProb: def f(self, distribution:torch.distributions.distribution.Distribution, - value: str) -> tuple: + x: str) -> tuple: """ Computes the probability for the input distribution. Args: distribution (torch.distributions.distribution.Distribution): Distribution. - value (Union[int,float])): Value or number of successes. + x (Union[int,float])): Value or number of successes. Returns: tuple containing the log of probability """ - k = ast.literal_eval(value) - if isinstance(k, int) is False and isinstance(k, float) is False : + v = ast.literal_eval(x) + if isinstance(v, int) is False and isinstance(v, float) is False : if isinstance(distribution, Dirichlet) is False: raise ValueError("You need to specify a float or an int.") if isinstance(distribution, (Poisson, Categorical, Multinomial, Binomial, NegativeBinomial)): - tens = torch.tensor(k, dtype=torch.int64) + tens = torch.tensor(v, dtype=torch.int64) else: - tens = torch.tensor(k, dtype=torch.float32) + tens = torch.tensor(v, dtype=torch.float32) return (distribution.log_prob(tens),) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptdm_log_prob_tensor.py b/modules/pytorch_wrapper/ptdm_log_prob_tensor.py new file mode 100644 index 0000000..08fb98b --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_log_prob_tensor.py @@ -0,0 +1,49 @@ +from typing import Any, Dict +import torch + + +class PtdmLogProbTensor: + """ + Computes the log of probability for the input distribution. This nodes accepts a tensor so it can be used to compute log of probability for multiple values contained in a tensor. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + + category: PyTorch wrapper - Distribution + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying the required input types. + """ + return { + "required": { + "distribution": ("PTDISTRIBUTION", {}), + "tens": ("TENSOR", {}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, + distribution:torch.distributions.distribution.Distribution, + tens: torch.Tensor) -> tuple: + """ + Computes the probability for the input distribution. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + + Returns: + tuple containing the log of probability + """ + return (distribution.log_prob(tens),) + \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptdm_pdf.py b/modules/pytorch_wrapper/ptdm_pdf.py index 11b3c60..72e8eb8 100644 --- a/modules/pytorch_wrapper/ptdm_pdf.py +++ b/modules/pytorch_wrapper/ptdm_pdf.py @@ -9,7 +9,7 @@ class PtdmPdf: Args: distribution (torch.distributions.distribution.Distribution): Distribution. - v (float): Value + x (float): Value category: PyTorch wrapper - Distribution """ @@ -25,7 +25,7 @@ class PtdmPdf: return { "required": { "distribution": ("PTDISTRIBUTION", {}), - "v": ("STRING", {"default": ""}) + "x": ("STRING", {"default": ""}) } } @@ -33,18 +33,18 @@ class PtdmPdf: FUNCTION: str = "f" CATEGORY: str = "Distribution" - def f(self, distribution: torch.distributions.Distribution, v: str) -> tuple: + def f(self, distribution: torch.distributions.Distribution, x: str) -> tuple: """ Computes the probability for the input distribution. Args: distribution (torch.distributions.Distribution): Distribution. - v (str): Value + x (str): Value Returns: tuple containing the probability """ - v2 = ast.literal_eval(v) + v2 = ast.literal_eval(x) if isinstance(v2, int) is False and isinstance(v2, float) is False : if isinstance(distribution, Dirichlet) is False: raise ValueError("You need to specify a float or an int.") diff --git a/modules/pytorch_wrapper/ptdm_pdf_tensor.py b/modules/pytorch_wrapper/ptdm_pdf_tensor.py new file mode 100644 index 0000000..88ab28f --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_pdf_tensor.py @@ -0,0 +1,46 @@ +from typing import Any, Dict +import torch + + +class PtdmPdfTensor: + """ + Computes the probability density for the input distribution. This nodes accepts a tensor so it can be used to compute pdf for multiple values contained in a tensor. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + + category: PyTorch wrapper - Distribution + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying the required input types. + """ + return { + "required": { + "distribution": ("PTDISTRIBUTION", {}), + "tens": ("TENSOR", {}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, tens: torch.Tensor) -> tuple: + """ + Computes the probability for the input distribution. + + Args: + distribution (torch.distributions.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + + Returns: + tuple containing the probability + """ + return (torch.exp(distribution.log_prob(tens)),) diff --git a/modules/pytorch_wrapper/ptdm_pmf_tensor.py b/modules/pytorch_wrapper/ptdm_pmf_tensor.py new file mode 100644 index 0000000..c963619 --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_pmf_tensor.py @@ -0,0 +1,46 @@ +from typing import Any, Dict +import torch + + +class PtdmPmfTensor: + """ + Computes the probability for the input distribution. This nodes accepts a tensor so it can be used to compute pmf for multiple values contained in a tensor. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): k(s) in Tensor. + + category: PyTorch wrapper - Distribution + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying the required input types. + """ + return { + "required": { + "distribution": ("PTDISTRIBUTION", {}), + "tens": ("TENSOR", {}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, tens: torch.Tensor) -> tuple: + """ + Computes the probability for the input distribution. + + Args: + distribution (torch.distributions.Distribution): Distribution. + tens (torch.Tensor): Value(s) in Tensor. + + Returns: + tuple containing the probability + """ + return (torch.exp(distribution.log_prob(tens)),) diff --git a/pyproject.toml b/pyproject.toml index 7a2b6e7..f758b08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "pt-wrapper" description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training." -version = "3.1.0" +version = "3.1.1" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptdm_cdf_tensor_test.py b/tests/pytorch_wrapper/ptdm_cdf_tensor_test.py new file mode 100644 index 0000000..cbec628 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_cdf_tensor_test.py @@ -0,0 +1,61 @@ +import os +import sys +import unittest +import numpy as np +import torch +import scipy.stats as scst +from torch.distributions import Exponential + +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_tensor import PtdmCdfTensor + + +class TestPtdfCdfTensor(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Exponential": [ + (1.0, 3.0, scst.expon.cdf(3.0, scale=1)), + (1.0, [3.0, 4.0], scst.expon.cdf(np.array([3.0, 4.0], dtype=np.float32), scale=1)) + ], + } + + def test_distributions(self): + """Test CDF.""" + 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 == "Exponential": + distribution = Exponential(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + else: + raise ValueError("Only exponential is supported in this unit test.") + + cdf = PtdmCdfTensor().f(distribution, k_tensor)[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor") + + # Convert tensors to floats for comparison + if isinstance(case[1], float): + 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}") + + elif isinstance(case[1], list): + cdf = cdf.numpy() + self.assertTrue( + np.allclose(cdf, case[-1], atol=1e-5) + ) + + else: + raise ValueError("Unsupported data type.") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_log_prob_tensor_test.py b/tests/pytorch_wrapper/ptdm_log_prob_tensor_test.py new file mode 100644 index 0000000..497888d --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_log_prob_tensor_test.py @@ -0,0 +1,61 @@ +import os +import sys +import unittest +import numpy as np +import torch +import scipy.stats as scst +from torch.distributions import Exponential + +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_log_prob_tensor import PtdmLogProbTensor + + +class TestPtdfLogProbTensor(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Exponential": [ + (1.0, 3.0, scst.expon.logpdf(3.0, scale=1)), + (1.0, [3.0, 4.0], scst.expon.logpdf(np.array([3.0, 4.0], dtype=np.float32), scale=1)) + ], + } + + def test_distributions(self): + """Test log prob.""" + 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 == "Exponential": + distribution = Exponential(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + else: + raise ValueError("Only exponential is supported in this unit test.") + + log_prob = PtdmLogProbTensor().f(distribution, k_tensor)[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(log_prob, torch.Tensor), "Probability output is not a tensor") + + # Convert tensors to floats for comparison + if isinstance(case[1], float): + log_prob = log_prob.item() + + # Check if computed values are approximately equal to expected values + self.assertAlmostEqual(log_prob, case[-1], places=4, + msg=f"{dist_name}: Expected {case[-1]}, got {log_prob}") + + elif isinstance(case[1], list): + log_prob = log_prob.numpy() + self.assertTrue( + np.allclose(log_prob, case[-1], atol=1e-5) + ) + + else: + raise ValueError("Unsupported data type.") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_pdf_tensor_test.py b/tests/pytorch_wrapper/ptdm_pdf_tensor_test.py new file mode 100644 index 0000000..b7365ff --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_pdf_tensor_test.py @@ -0,0 +1,61 @@ +import os +import sys +import unittest +import numpy as np +import torch +import scipy.stats as scst +from torch.distributions import Exponential + +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_pdf_tensor import PtdmPdfTensor + + +class TestPtdfPdfTensor(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Exponential": [ + (1.0, 3.0, scst.expon.pdf(3.0, scale=1)), + (1.0, [3.0, 4.0], scst.expon.pdf(np.array([3.0, 4.0], dtype=np.float32), scale=1)) + ], + } + + def test_distributions(self): + """Test pdf.""" + 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 == "Exponential": + distribution = Exponential(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + else: + raise ValueError("Only exponential is supported in this unit test.") + + pdf = PtdmPdfTensor().f(distribution, k_tensor)[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(pdf, torch.Tensor), "Probability output is not a tensor") + + # Convert tensors to floats for comparison + if isinstance(case[1], float): + pdf = pdf.item() + + # Check if computed values are approximately equal to expected values + self.assertAlmostEqual(pdf, case[-1], places=4, + msg=f"{dist_name}: Expected {case[-1]}, got {pdf}") + + elif isinstance(case[1], list): + pdf = pdf.numpy() + self.assertTrue( + np.allclose(pdf, case[-1], atol=1e-5) + ) + + else: + raise ValueError("Unsupported data type.") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_pmf_tensor_test.py b/tests/pytorch_wrapper/ptdm_pmf_tensor_test.py new file mode 100644 index 0000000..1cdaa69 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_pmf_tensor_test.py @@ -0,0 +1,61 @@ +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()