diff --git a/__init__.py b/__init__.py index b0316d3..6ced78c 100644 --- a/__init__.py +++ b/__init__.py @@ -112,6 +112,22 @@ from .modules.pytorch_wrapper.pt_var import PtVar from .modules.pytorch_wrapper.pt_view import PtView from .modules.pytorch_wrapper.pt_where import PtWhere from .modules.pytorch_wrapper.pt_zeros import PtZeros +from .modules.pytorch_wrapper.ptd_bernoulli import PtdBernoulli +from .modules.pytorch_wrapper.ptd_beta import PtdBeta +from .modules.pytorch_wrapper.ptd_binomial import PtdBinomial +from .modules.pytorch_wrapper.ptd_categorical import PtdCategorical +from .modules.pytorch_wrapper.ptd_chi2 import PtdChi2 +from .modules.pytorch_wrapper.ptd_exponential import PtdExponential +from .modules.pytorch_wrapper.ptd_gamma import PtdGamma +from .modules.pytorch_wrapper.ptd_normal import PtdNormal +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_log_prob import PtdmLogProb +from .modules.pytorch_wrapper.ptdm_pdf import PtdmPdf +from .modules.pytorch_wrapper.ptdm_pmf import PtdmPmf +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 from .modules.pytorch_wrapper.ptf_log_softmax import PtfLogSoftmax @@ -266,6 +282,22 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtView": PtView, "PtWhere": PtWhere, "PtZeros": PtZeros, + "PtdBernoulli": PtdBernoulli, + "PtdBeta": PtdBeta, + "PtdBinomial": PtdBinomial, + "PtdCategorical": PtdCategorical, + "PtdChi2": PtdChi2, + "PtdExponential": PtdExponential, + "PtdGamma": PtdGamma, + "PtdNormal": PtdNormal, + "PtdPoisson": PtdPoisson, + "PtdStudentT": PtdStudentT, + "PtdUniform": PtdUniform, + "PtdmCdf": PtdmCdf, + "PtdmLogProb": PtdmLogProb, + "PtdmPdf": PtdmPdf, + "PtdmPmf": PtdmPmf, + "PtdmSample": PtdmSample, "PtfGELU": PtfGELU, "PtfLeakyReLU": PtfLeakyReLU, "PtfLogSoftmax": PtfLogSoftmax, @@ -420,6 +452,22 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtView": "Pt View", "PtWhere": "Pt Where", "PtZeros": "Pt Zeros", + "PtdBernoulli": "Ptd Bernoulli", + "PtdBeta": "Ptd Beta", + "PtdBinomial": "Ptd Binomial", + "PtdCategorical": "Ptd Categorical", + "PtdChi2": "Ptd Chi2", + "PtdExponential": "Ptd Exponential", + "PtdGamma": "Ptd Gamma", + "PtdNormal": "Ptd Normal", + "PtdPoisson": "Ptd Poisson", + "PtdStudentT": "Ptd Student T", + "PtdUniform": "Ptd Uniform", + "PtdmCdf": "Ptdm Cdf", + "PtdmLogProb": "Ptdm Log Prob", + "PtdmPdf": "Ptdm Pdf", + "PtdmPmf": "Ptdm Pmf", + "PtdmSample": "Ptdm Sample", "PtfGELU": "Ptf GELU", "PtfLeakyReLU": "Ptf Leaky ReLU", "PtfLogSoftmax": "Ptf Log Softmax", diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index d817ba0..ddc55a7 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -60,6 +60,25 @@ | [Pt To Int64](pt_to_int64.md) | Converts the input tensor's data type to int64. | | [Pt To Int8](pt_to_int8.md) | Converts the input tensor's data type to int8. | | [Pt To Uint8](pt_to_uint8.md) | Converts the input tensor's data type to uint8. | +## PyTorch wrapper - Distribution +| Node | Description | +| --- | --- | +| [Ptd Bernoulli](ptd_bernoulli.md) | Instantiates a Bernoulli distribution object. | +| [Ptd Beta](ptd_beta.md) | Instantiates a Beta distribution object. | +| [Ptd Binomial](ptd_binomial.md) | Instantiates a Binomial distribution object. | +| [Ptd Categorical](ptd_categorical.md) | Instantiates a Categorical distribution object from the input probabilities or logits. You have to specify one of them and not both. | +| [Ptd Chi2](ptd_chi2.md) | Instantiates a Chi-squared distribution object. | +| [Ptd Exponential](ptd_exponential.md) | Instantiates a Exponential distribution object. | +| [Ptd Gamma](ptd_gamma.md) | Instantiates a Gamma distribution object. | +| [Ptd Normal](ptd_normal.md) | Instantiates a Normal distribution object. | +| [Ptd Poisson](ptd_poisson.md) | Instantiates a Poisson distribution object. | +| [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 Log Prob](ptdm_log_prob.md) | Computes the log of probability for the input distribution. | +| [Ptdm Pdf](ptdm_pdf.md) | Computes the probability density for the input distribution. | +| [Ptdm Pmf](ptdm_pmf.md) | Computes the probability for the input distribution. | +| [Ptdm Sample](ptdm_sample.md) | Samples from the input distribution. | ## PyTorch wrapper - Image processing | Node | Description | | --- | --- | diff --git a/docs/reference/ptd_bernoulli.md b/docs/reference/ptd_bernoulli.md new file mode 100644 index 0000000..354ec60 --- /dev/null +++ b/docs/reference/ptd_bernoulli.md @@ -0,0 +1,22 @@ +# Ptd Bernoulli +Instantiates a Bernoulli distribution object. + + Args: + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| probs | String | +| logits | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_beta.md b/docs/reference/ptd_beta.md new file mode 100644 index 0000000..051084a --- /dev/null +++ b/docs/reference/ptd_beta.md @@ -0,0 +1,22 @@ +# Ptd Beta +Instantiates a Beta distribution object. + + Args: + alpha: The alpha parameter of the distsribution. You can specify a scalar or a tuple of float. + beta: The beta parameter of the distribution. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| alpha | String | +| beta | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_binomial.md b/docs/reference/ptd_binomial.md new file mode 100644 index 0000000..1934a0c --- /dev/null +++ b/docs/reference/ptd_binomial.md @@ -0,0 +1,24 @@ +# Ptd Binomial +Instantiates a Binomial distribution object. + + Args: + total_count (str): Total count, or n. You can specify a scalar or a tuple of int. + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| total_count | String | +| probs | String | +| logits | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_categorical.md b/docs/reference/ptd_categorical.md new file mode 100644 index 0000000..dd99c66 --- /dev/null +++ b/docs/reference/ptd_categorical.md @@ -0,0 +1,22 @@ +# Ptd Categorical +Instantiates a Categorical distribution object from the input probabilities or logits. You have to specify one of them and not both. + +Args: + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| probs | String | +| logits | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_chi2.md b/docs/reference/ptd_chi2.md new file mode 100644 index 0000000..e09aa4f --- /dev/null +++ b/docs/reference/ptd_chi2.md @@ -0,0 +1,20 @@ +# Ptd Chi2 +Instantiates a Chi-squared distribution object. + + Args: + df (str): Degree of freedom. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| df | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_exponential.md b/docs/reference/ptd_exponential.md new file mode 100644 index 0000000..b76d821 --- /dev/null +++ b/docs/reference/ptd_exponential.md @@ -0,0 +1,20 @@ +# Ptd Exponential +Instantiates a Exponential distribution object. + + Args: + rate (str): Rate. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| rate | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_gamma.md b/docs/reference/ptd_gamma.md new file mode 100644 index 0000000..1dbcb3d --- /dev/null +++ b/docs/reference/ptd_gamma.md @@ -0,0 +1,22 @@ +# Ptd Gamma +Instantiates a Gamma distribution object. + + Args: + alpha: Rate. Mean of the distsribution. You can specify a scalar or a tuple of float. + beta: Rate. Standard deviation of the distribution. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| alpha | String | +| beta | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_normal.md b/docs/reference/ptd_normal.md new file mode 100644 index 0000000..d27e338 --- /dev/null +++ b/docs/reference/ptd_normal.md @@ -0,0 +1,22 @@ +# Ptd Normal +Instantiates a Normal distribution object. + + Args: + loc: Mean of the distsribution. You can specify a scalar or a tuple of float. + scale: Standard deviation of the distribution. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| loc | String | +| scale | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_poisson.md b/docs/reference/ptd_poisson.md new file mode 100644 index 0000000..04d9692 --- /dev/null +++ b/docs/reference/ptd_poisson.md @@ -0,0 +1,20 @@ +# Ptd Poisson +Instantiates a Poisson distribution object. + + Args: + rate (str): Rate. You can specify a scalar or a tuple of int. + +## Input +| Name | Data type | +|---|---| +| rate | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_student_t.md b/docs/reference/ptd_student_t.md new file mode 100644 index 0000000..2c0d737 --- /dev/null +++ b/docs/reference/ptd_student_t.md @@ -0,0 +1,24 @@ +# Ptd Student T +Instantiates a StudentT distribution object. + + Args: + df (str): Degree of freedom. You can specify a scalar or a tuple of float. + loc (str): Rate. Mean of the distsribution. You can specify a scalar or a tuple of float. + scale (str): Rate. Scale of the distribution. You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| df | String | +| loc | String | +| scale | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptd_uniform.md b/docs/reference/ptd_uniform.md new file mode 100644 index 0000000..edb848a --- /dev/null +++ b/docs/reference/ptd_uniform.md @@ -0,0 +1,22 @@ +# Ptd Uniform +Instantiates a Uniform distribution object. + + Args: + low (str): Minimum (inclusive). You can specify a scalar or a tuple of float. + high (str): Maximum (exclusive). You can specify a scalar or a tuple of float. + +## Input +| Name | Data type | +|---|---| +| low | String | +| high | String | + +## Output +| Data type | +|---| +| Ptdistribution | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptdm_cdf.md b/docs/reference/ptdm_cdf.md new file mode 100644 index 0000000..819ff73 --- /dev/null +++ b/docs/reference/ptdm_cdf.md @@ -0,0 +1,25 @@ +# Ptdm Cdf +Computes the cumulative distribution function for the input distribution. + +**Note** +Cdf is not supported for all distributions in PyTorch. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + v (float): Value + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| v | String | + +## 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 new file mode 100644 index 0000000..c98aa7c --- /dev/null +++ b/docs/reference/ptdm_log_prob.md @@ -0,0 +1,22 @@ +# Ptdm Log Prob +Computes the log of probability for the input distribution. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + value (Union[int,float]): Value (Number of successes for PMF). + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| value | String | + +## 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 new file mode 100644 index 0000000..eef9f7a --- /dev/null +++ b/docs/reference/ptdm_pdf.md @@ -0,0 +1,22 @@ +# Ptdm Pdf +Computes the probability density for the input distribution. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + v (float): Value + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| v | String | + +## 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.md b/docs/reference/ptdm_pmf.md new file mode 100644 index 0000000..ccbfdc0 --- /dev/null +++ b/docs/reference/ptdm_pmf.md @@ -0,0 +1,22 @@ +# Ptdm Pmf +Computes the probability for the input distribution. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + k (int): Number of successes. + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| k | Int | + +## 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_sample.md b/docs/reference/ptdm_sample.md new file mode 100644 index 0000000..9153fdc --- /dev/null +++ b/docs/reference/ptdm_sample.md @@ -0,0 +1,22 @@ +# Ptdm Sample +Samples from the input distribution. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + sample_shape (torch.Size): Shape of the sample. + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| sample_shape | String | + +## 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/ptd_bernoulli.py b/modules/pytorch_wrapper/ptd_bernoulli.py new file mode 100644 index 0000000..beb47c0 --- /dev/null +++ b/modules/pytorch_wrapper/ptd_bernoulli.py @@ -0,0 +1,59 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Bernoulli + +class PtdBernoulli: + """ + Instantiates a Bernoulli distribution object. + + Args: + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + + 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": { + "probs": ("STRING", {"default": ""}), + "logits": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, probs: str, logits: str) -> tuple: + """ + Instantiates a Bernoulli distribution object. + + Args: + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + Returns: + tuple: (torch.distributions.bernoulli.Bernoulli) + """ + p = ast.literal_eval(probs) if probs else None + l = ast.literal_eval(logits) if logits else None + + if p is not None and l is not None: # You cannot use if p because p or l can be "0". + raise ValueError("You can specify either probabilities or logits, but not both.") + if p is None and l is None: + raise ValueError("You have to specify either probabilities or logits") + + if p is not None: + dist = Bernoulli(probs=torch.tensor(p, dtype=torch.float32)) + elif l is not None: + dist = Bernoulli(logits=torch.tensor(l, dtype=torch.float32)) + + return (dist,) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptd_beta.py b/modules/pytorch_wrapper/ptd_beta.py new file mode 100644 index 0000000..8fe4eaa --- /dev/null +++ b/modules/pytorch_wrapper/ptd_beta.py @@ -0,0 +1,52 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Beta + +class PtdBeta: + """ + Instantiates a Beta distribution object. + + Args: + alpha: The alpha parameter of the distsribution. You can specify a scalar or a tuple of float. + beta: The beta parameter of the distribution. You can specify a scalar or a tuple of float. + + 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": { + "alpha": ("STRING", {"default": ""}), + "beta": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, alpha: str, beta: str) -> tuple: + """ + Instantiates a Beta distribution object. + + Args: + alpha: The alpha parameter of the distsribution. You can specify a scalar or a tuple of float. + beta: The beta parameter of the distribution. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Beta) + """ + c1 = ast.literal_eval(alpha) + c2 = ast.literal_eval(beta) + + dist = Beta( + concentration1=torch.tensor(c1, dtype=torch.float32), concentration0=torch.tensor(c2, dtype=torch.float32)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_binomial.py b/modules/pytorch_wrapper/ptd_binomial.py new file mode 100644 index 0000000..0f679b0 --- /dev/null +++ b/modules/pytorch_wrapper/ptd_binomial.py @@ -0,0 +1,69 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Binomial + +class PtdBinomial: + """ + Instantiates a Binomial distribution object. + + Args: + total_count (str): Total count, or n. You can specify a scalar or a tuple of int. + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + + 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": { + "total_count": ("STRING", {"default": ""}), + "probs": ("STRING", {"default": ""}), + "logits": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, total_count: str, probs: str, logits: str): + """ + Instantiates a Binomial distribution object. + + Args: + total_count (str): Total count, or n. You can specify a scalar or a tuple of int. + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Binomial) + """ + + n = ast.literal_eval(total_count) + p = ast.literal_eval(probs) if probs else None + l = ast.literal_eval(logits) if logits else None + + if p is not None and l is not None: # You cannot use if p because p or l can be "0". + raise ValueError("You can specify either probabilities or logits, but not both.") + if p is None and l is None: + raise ValueError("You have to specify either probabilities or logits") + + if p is not None: + dist = Binomial( + total_count=torch.tensor(n, dtype=torch.int64), + probs=torch.tensor(p, dtype=torch.float32)) + elif l is not None: + dist = Binomial( + total_count=torch.tensor(n, dtype=torch.int64), + logits=torch.tensor(l, dtype=torch.float32)) + + return(dist,) diff --git a/modules/pytorch_wrapper/ptd_categorical.py b/modules/pytorch_wrapper/ptd_categorical.py new file mode 100644 index 0000000..6d9e4ba --- /dev/null +++ b/modules/pytorch_wrapper/ptd_categorical.py @@ -0,0 +1,60 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Categorical + +class PtdCategorical: + """ + Instantiates a Categorical distribution object from the input probabilities or logits. You have to specify one of them and not both. + + Args: + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + + 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": { + "probs": ("STRING", {"default": ""}), + "logits": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, probs: str, logits: str) -> tuple: + """ + Instantiates a Categorical distribution object from the input probabilities or logits. You have to specify one of them and not both. + + Args: + probs (str): Probabilities of the distsribution. You can specify a scalar or a tuple of float. + logits (str): Logits of the distribution. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Categorical) + """ + p = ast.literal_eval(probs) if probs else None + l = ast.literal_eval(logits) if logits else None + + if p is not None and l is not None: # You cannot use if p because p or l can be "0". + raise ValueError("You can specify either probabilities or logits, but not both.") + if p is None and l is None: + raise ValueError("You have to specify either probabilities or logits") + + if p is not None: + dist = Categorical(probs=torch.tensor(p)) + elif l is not None: + dist = Categorical(logits=torch.tensor(l)) + + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_chi2.py b/modules/pytorch_wrapper/ptd_chi2.py new file mode 100644 index 0000000..789f4dc --- /dev/null +++ b/modules/pytorch_wrapper/ptd_chi2.py @@ -0,0 +1,50 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Chi2 + +class PtdChi2: + """ + Ptd Chi2: + Instantiates a Chi-squared distribution object. + + Args: + df (str): Degree of freedom. You can specify a scalar or a tuple of float. + + 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": { + "df": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, + df: str) -> tuple: + """ + Instantiates a Chi-squared distribution object. + + Args: + df (str): Degree of freedom. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Chi2) + """ + d = ast.literal_eval(df) + + dist = Chi2( + df=torch.tensor(d, dtype=torch.float32)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_exponential.py b/modules/pytorch_wrapper/ptd_exponential.py new file mode 100644 index 0000000..1330008 --- /dev/null +++ b/modules/pytorch_wrapper/ptd_exponential.py @@ -0,0 +1,48 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Exponential + +class PtdExponential: + """ + Instantiates a Exponential distribution object. + + Args: + rate (str): Rate. You can specify a scalar or a tuple of float. + + 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": { + "rate": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, rate: str) -> tuple: + """ + Instantiates a Exponential distribution object. + + Args: + rate (str): Probability. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Exponential) + """ + r = ast.literal_eval(rate) + + dist = Exponential( + rate=torch.tensor(r)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_gamma.py b/modules/pytorch_wrapper/ptd_gamma.py new file mode 100644 index 0000000..14080fe --- /dev/null +++ b/modules/pytorch_wrapper/ptd_gamma.py @@ -0,0 +1,53 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Gamma + +class PtdGamma: + """ + Instantiates a Gamma distribution object. + + Args: + alpha: Rate. Mean of the distsribution. You can specify a scalar or a tuple of float. + beta: Rate. Standard deviation of the distribution. You can specify a scalar or a tuple of float. + + 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": { + "alpha": ("STRING", {"default": ""}), + "beta": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, alpha: str, beta: str) -> tuple: + """ + Instantiates a Gamma distribution object. + + Args: + alpha (str): Concentration You can specify a scalar or a tuple of float. + beta (str): Rate. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Gamma) + """ + c = ast.literal_eval(alpha) + r = ast.literal_eval(beta) + + dist = Gamma( + concentration=torch.tensor(c, dtype=torch.float32), + rate=torch.tensor(r, dtype=torch.float32)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_normal.py b/modules/pytorch_wrapper/ptd_normal.py new file mode 100644 index 0000000..5d644d7 --- /dev/null +++ b/modules/pytorch_wrapper/ptd_normal.py @@ -0,0 +1,53 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Normal + +class PtdNormal: + """ + Instantiates a Normal distribution object. + + Args: + loc: Mean of the distsribution. You can specify a scalar or a tuple of float. + scale: Standard deviation of the distribution. You can specify a scalar or a tuple of float. + + 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": { + "loc": ("STRING", {"default": ""}), + "scale": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, loc: str, scale: str) -> tuple: + """ + Instantiates a Normal distribution object. + + Args: + loc (str): Mean of the distsribution. You can specify a scalar or a tuple of float. + scale (str): Standard deviation of the distribution. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Normal) + """ + mu = ast.literal_eval(loc) + s = ast.literal_eval(scale) + + dist = Normal( + loc=torch.tensor(mu, dtype=torch.float32), + scale=torch.tensor(s, dtype=torch.float32)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_poisson.py b/modules/pytorch_wrapper/ptd_poisson.py new file mode 100644 index 0000000..4017196 --- /dev/null +++ b/modules/pytorch_wrapper/ptd_poisson.py @@ -0,0 +1,48 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Poisson + +class PtdPoisson: + """ + Instantiates a Poisson distribution object. + + Args: + rate (str): Rate. You can specify a scalar or a tuple of int. + + 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": { + "rate": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, rate: str) -> tuple: + """ + Instantiates a Poisson distribution object. + + Args: + rate (str): Probability. You can specify a scalar or a tuple of int. + + Returns: + tuple: (torch.distributions.poisson.Poisson) + """ + r = ast.literal_eval(rate) + + dist = Poisson( + rate=torch.tensor(r, dtype=torch.float32)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_student_t.py b/modules/pytorch_wrapper/ptd_student_t.py new file mode 100644 index 0000000..9accd6d --- /dev/null +++ b/modules/pytorch_wrapper/ptd_student_t.py @@ -0,0 +1,61 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import StudentT + +class PtdStudentT: + """ + Instantiates a StudentT distribution object. + + Args: + df (str): Degree of freedom. You can specify a scalar or a tuple of float. + loc (str): Rate. Mean of the distsribution. You can specify a scalar or a tuple of float. + scale (str): Rate. Scale of the distribution. You can specify a scalar or a tuple of float. + + 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": { + "df": ("STRING", {"default": ""}), + "loc": ("STRING", {"default": ""}), + "scale": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, + df: str, + loc: str, + scale: str) -> tuple: + """ + Instantiates a StudentT distribution object. + + Args: + df (str): Degree of freedom. You can specify a scalar or a tuple of float. + loc (str): Rate. Mean of the distsribution. You can specify a scalar or a tuple of float. + scale (str): Rate. Standard deviation of the distribution. You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.StudentT) + """ + d = ast.literal_eval(df) + mu = ast.literal_eval(loc) + s = ast.literal_eval(scale) + + dist = StudentT( + df=torch.tensor(d, dtype=torch.float32), + loc=torch.tensor(mu, dtype=torch.float32), + scale=torch.tensor(s, dtype=torch.float32)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptd_uniform.py b/modules/pytorch_wrapper/ptd_uniform.py new file mode 100644 index 0000000..2e5e6f8 --- /dev/null +++ b/modules/pytorch_wrapper/ptd_uniform.py @@ -0,0 +1,52 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import Uniform + +class PtdUniform: + """ + Instantiates a Uniform distribution object. + + Args: + low (str): Minimum (inclusive). You can specify a scalar or a tuple of float. + high (str): Maximum (exclusive). You can specify a scalar or a tuple of float. + + 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": { + "low": ("STRING", {"default": ""}), + "high": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("PTDISTRIBUTION",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, low: str, high: str) -> tuple: + """ + Instantiates a Uniform distribution object. + + Args: + low (str): Minimum (inclusive). You can specify a scalar or a tuple of float. + high (str): Maximum (exclusive). You can specify a scalar or a tuple of float. + + Returns: + tuple: (torch.distributions.binomial.Uniform) + """ + l = ast.literal_eval(low) + h = ast.literal_eval(high) + + dist = Uniform( + low=torch.tensor(l, dtype=torch.float32), high=torch.tensor(h, dtype=torch.float32)) + return (dist,) diff --git a/modules/pytorch_wrapper/ptdm_cdf.py b/modules/pytorch_wrapper/ptdm_cdf.py new file mode 100644 index 0000000..348aeb8 --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_cdf.py @@ -0,0 +1,60 @@ +from typing import Any, Dict +import torch +from torch.distributions import * +import ast + +class PtdmCdf: + """ + Computes the cumulative distribution function for the input distribution. + + **Note** + Cdf is not supported for all distributions in PyTorch. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + v (float): Value + + 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", {}), + "v": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, v: str) -> tuple: + """ + Computes the cumulative distribution function for the input distribution. + + Args: + distribution (torch.distributions.Distribution): Distribution. + v (str): Value + + Returns: + tuple containing the probability + """ + v2 = ast.literal_eval(v) + 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.") + + if isinstance(distribution, (Poisson, Categorical, Multinomial, Binomial, NegativeBinomial)): + tens = torch.tensor(v2, dtype=torch.int64) + else: + tens = torch.tensor(v2, dtype=torch.float32) + + return (distribution.cdf(tens),) diff --git a/modules/pytorch_wrapper/ptdm_log_prob.py b/modules/pytorch_wrapper/ptdm_log_prob.py new file mode 100644 index 0000000..73826e0 --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_log_prob.py @@ -0,0 +1,61 @@ +from typing import Any, Dict +import ast +import torch +from torch.distributions import * + +class PtdmLogProb: + """ + Computes the log of probability for the input distribution. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + value (Union[int,float]): Value (Number of successes for PMF). + + 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", {}), + "value": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, + distribution:torch.distributions.distribution.Distribution, + value: 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. + + Returns: + tuple containing the log of probability + """ + + k = ast.literal_eval(value) + if isinstance(k, int) is False and isinstance(k, 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) + else: + tens = torch.tensor(k, dtype=torch.float32) + + 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 new file mode 100644 index 0000000..11b3c60 --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_pdf.py @@ -0,0 +1,57 @@ +from typing import Any, Dict +import torch +from torch.distributions import * +import ast + +class PtdmPdf: + """ + Computes the probability density for the input distribution. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + v (float): Value + + 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", {}), + "v": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, v: str) -> tuple: + """ + Computes the probability for the input distribution. + + Args: + distribution (torch.distributions.Distribution): Distribution. + v (str): Value + + Returns: + tuple containing the probability + """ + v2 = ast.literal_eval(v) + 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.") + + if isinstance(distribution, (Poisson, Categorical, Multinomial, Binomial, NegativeBinomial)): + tens = torch.tensor(v2, dtype=torch.int64) + else: + tens = torch.tensor(v2, dtype=torch.float32) + + return (torch.exp(distribution.log_prob(tens)),) diff --git a/modules/pytorch_wrapper/ptdm_pmf.py b/modules/pytorch_wrapper/ptdm_pmf.py new file mode 100644 index 0000000..7cbf720 --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_pmf.py @@ -0,0 +1,52 @@ +from typing import Any, Dict +import torch +from torch.distributions import * +import ast + +class PtdmPmf: + """ + Computes the probability for the input distribution. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + k (int): Number of successes. + + 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", {}), + "k": ("INT", {"default": 0, "min": -2**31, "max": 2**31}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, k: int) -> tuple: + """ + Computes the probability for the input distribution. + + Args: + distribution (torch.distributions.Distribution): Distribution. + k (int): Number of successes. + + Returns: + tuple containing the probability + """ + if isinstance(distribution, (Poisson, Categorical, Multinomial, Binomial, NegativeBinomial)): + tens = torch.tensor(k, dtype=torch.int64) + else: + tens = torch.tensor(k, dtype=torch.float32) + + return (torch.exp(distribution.log_prob(tens)),) diff --git a/modules/pytorch_wrapper/ptdm_sample.py b/modules/pytorch_wrapper/ptdm_sample.py new file mode 100644 index 0000000..52a561c --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_sample.py @@ -0,0 +1,52 @@ +from typing import Any, Dict +import torch +import ast + +class PtdmSample: + """ + Samples from the input distribution. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + sample_shape (torch.Size): Shape of the sample. + + 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", {}), + "sample_shape": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, + distribution:torch.distributions.distribution.Distribution, + sample_shape: str) -> tuple: + """ + Samples from the input distribution. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + sample_shape (torch.Size): Shape of the sample. + + Returns: + tuple containing the sample + """ + + shape = ast.literal_eval(sample_shape) + if isinstance(shape, int): + shape = (shape,) + return (distribution.sample(shape),) diff --git a/pyproject.toml b/pyproject.toml index dc44806..7a2b6e7 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.0.13" +version = "3.1.0" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptd_bernoulli_test.py b/tests/pytorch_wrapper/ptd_bernoulli_test.py new file mode 100644 index 0000000..6091457 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_bernoulli_test.py @@ -0,0 +1,50 @@ +import os +import sys +import unittest +import torch + +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_bernoulli import PtdBernoulli + +class TestPtdBernoulli(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdBernoulli() + + def test_bernoulli_distribution_probs1(self): + """Test instantiation of bernoulli distribution.""" + dist = self.node.f("0", "")[0] # p=0.5 + self.assertTrue(isinstance(dist, torch.distributions.bernoulli.Bernoulli), "Probability distribution is not Bernoulli") + + def test_bernoulli_distribution_probs2(self): + """Test instantiation of bernoulli distribution.""" + dist = self.node.f("0.5", "")[0] # p=0.5 + self.assertTrue(isinstance(dist, torch.distributions.bernoulli.Bernoulli), "Probability distribution is not Bernoulli") + + def test_bernoulli_distribution_probs3(self): + """Test instantiation of bernoulli distribution.""" + dist = self.node.f("1", "")[0] # p=0.5 + self.assertTrue(isinstance(dist, torch.distributions.bernoulli.Bernoulli), "Probability distribution is not Bernoulli") + + def test_bernoulli_distribution_logits1(self): + """Test instantiation of bernoulli distribution.""" + dist = self.node.f("", "-40")[0] + self.assertTrue(isinstance(dist, torch.distributions.bernoulli.Bernoulli), "Probability distribution is not Bernoulli") + + def test_bernoulli_distribution_logits2(self): + """Test instantiation of bernoulli distribution.""" + dist = self.node.f("", "0")[0] + self.assertTrue(isinstance(dist, torch.distributions.bernoulli.Bernoulli), "Probability distribution is not Bernoulli") + + def test_bernoulli_distribution_logits3(self): + """Test instantiation of bernoulli distribution.""" + dist = self.node.f("", "40")[0] + self.assertTrue(isinstance(dist, torch.distributions.bernoulli.Bernoulli), "Probability distribution is not Bernoulli") + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_beta_test.py b/tests/pytorch_wrapper/ptd_beta_test.py new file mode 100644 index 0000000..ea5eba6 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_beta_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_beta import PtdBeta + +class TestPtdBeta(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdBeta() + + def test_beta_distribution_scalar(self): + """Test instantiation of beta distribution.""" + dist = self.node.f("1.0","9.0")[0] + self.assertTrue(isinstance(dist, torch.distributions.beta.Beta), "Probability distribution is not Beta") + + def test_beta_distribution_1d(self): + """Test instantiation of beta distribution.""" + dist = self.node.f("(1.0,)", "(9.0,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.beta.Beta), "Probability distribution is not Beta") + + def test_beta_distribution_2d(self): + """Test instantiation of beta distribution.""" + dist = self.node.f("(1.0,3.0)","(9.0,10.0)")[0] + self.assertTrue(isinstance(dist, torch.distributions.beta.Beta), "Probability distribution is not Beta") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_binomial_test.py b/tests/pytorch_wrapper/ptd_binomial_test.py new file mode 100644 index 0000000..0c86c78 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_binomial_test.py @@ -0,0 +1,59 @@ +import os +import sys +import unittest +import torch + +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 + +class TestPtdBinomial(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdBinomial() + + def test_binomial_distribution_probs1(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "0", "")[0] # n = 10, p=0 + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + def test_binomial_distribution_probs2(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "0.0", "")[0] # n = 10, p=0 + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + def test_binomial_distribution_probs3(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "0.5", "")[0] # n = 10, p=0.5 + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + def test_binomial_distribution_probs4(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "1", "")[0] # n = 10, p=1 + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + def test_binomial_distribution_probs5(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "1.0", "")[0] # n = 10, p=1 + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + def test_binomial_distribution_logits1(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "", "-40")[0] + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + def test_binomial_distribution_logits2(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "", "0")[0] + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + def test_binomial_distribution_logits3(self): + """Test instantiation of binomial distribution.""" + dist = self.node.f("10", "", "40")[0] + self.assertTrue(isinstance(dist, torch.distributions.binomial.Binomial), "Probability distribution is not Binomial") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_categorical_test.py b/tests/pytorch_wrapper/ptd_categorical_test.py new file mode 100644 index 0000000..3cedf06 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_categorical_test.py @@ -0,0 +1,29 @@ +import os +import sys +import unittest +import torch + +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_categorical import PtdCategorical + +class TestPtdCategorical(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdCategorical() + + def test_categorical_distribution_probs(self): + """Test instantiation of categorical distribution.""" + dist = self.node.f("[0.1,0.2,0.7]", "")[0] + self.assertTrue(isinstance(dist, torch.distributions.categorical.Categorical), "Probability distribution is not Categorical") + + def test_categorical_distribution_logits(self): + """Test instantiation of categorical distribution.""" + dist = self.node.f("", "[-40, -0.40, 0.36797678529459443]")[0] + self.assertTrue(isinstance(dist, torch.distributions.categorical.Categorical), "Probability distribution is not Categorical") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_chi2_test.py b/tests/pytorch_wrapper/ptd_chi2_test.py new file mode 100644 index 0000000..8c2ce34 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_chi2_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_chi2 import PtdChi2 + +class TestPtdChi2(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdChi2() + + def test_chi2_distribution_scalar(self): + """Test instantiation of chi2 distribution.""" + dist = self.node.f("10.0")[0] + self.assertTrue(isinstance(dist, torch.distributions.studentT.Chi2), "Probability distribution is not Chi2") + + def test_chi2_distribution_1d(self): + """Test instantiation of chi2 distribution.""" + dist = self.node.f("(10.0,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.studentT.Chi2), "Probability distribution is not Chi2") + + def test_chi2_distribution_2d(self): + """Test instantiation of chi2 distribution.""" + dist = self.node.f("(10.0, 11.0)")[0] + self.assertTrue(isinstance(dist, torch.distributions.studentT.Chi2), "Probability distribution is not Chi2") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_exponential_test.py b/tests/pytorch_wrapper/ptd_exponential_test.py new file mode 100644 index 0000000..5b98278 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_exponential_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_exponential import PtdExponential + +class TestPtdExponential(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdExponential() + + def test_exponential_distribution_scalar(self): + """Test instantiation of exponential distribution.""" + dist = self.node.f("3")[0] + self.assertTrue(isinstance(dist, torch.distributions.exponential.Exponential), "Probability distribution is not Exponential") + + def test_exponential_distribution_1d(self): + """Test instantiation of exponential distribution.""" + dist = self.node.f("(3,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.exponential.Exponential), "Probability distribution is not Exponential") + + def test_exponential_distribution_2d(self): + """Test instantiation of exponential distribution.""" + dist = self.node.f("(2,3)")[0] + self.assertTrue(isinstance(dist, torch.distributions.exponential.Exponential), "Probability distribution is not Exponential") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_gamma_test.py b/tests/pytorch_wrapper/ptd_gamma_test.py new file mode 100644 index 0000000..892a880 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_gamma_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_gamma import PtdGamma + +class TestPtdGamma(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdGamma() + + def test_gamma_distribution_scalar(self): + """Test instantiation of gamma distribution.""" + dist = self.node.f("1.0","1.0")[0] + self.assertTrue(isinstance(dist, torch.distributions.gamma.Gamma), "Probability distribution is not Gamma") + + def test_gamma_distribution_1d(self): + """Test instantiation of gamma distribution.""" + dist = self.node.f("(1.0,)", "(1.0,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.gamma.Gamma), "Probability distribution is not Gamma") + + def test_gamma_distribution_2d(self): + """Test instantiation of gamma distribution.""" + dist = self.node.f("(1.0,1.0)","(1.0,1.0)")[0] + self.assertTrue(isinstance(dist, torch.distributions.gamma.Gamma), "Probability distribution is not Gamma") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_normal_test.py b/tests/pytorch_wrapper/ptd_normal_test.py new file mode 100644 index 0000000..145f5f7 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_normal_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_normal import PtdNormal + +class TestPtdNormal(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdNormal() + + def test_normal_distribution_scalar(self): + """Test instantiation of normal distribution.""" + dist = self.node.f("0.0","1.0")[0] + self.assertTrue(isinstance(dist, torch.distributions.normal.Normal), "Probability distribution is not Normal") + + def test_normal_distribution_1d(self): + """Test instantiation of normal distribution.""" + dist = self.node.f("(0.0,)", "(1.0,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.normal.Normal), "Probability distribution is not Normal") + + def test_normal_distribution_2d(self): + """Test instantiation of normal distribution.""" + dist = self.node.f("(155.0,170.0)","(8.0,10.0)")[0] + self.assertTrue(isinstance(dist, torch.distributions.normal.Normal), "Probability distribution is not Normal") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_poisson_test.py b/tests/pytorch_wrapper/ptd_poisson_test.py new file mode 100644 index 0000000..49c3fc3 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_poisson_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_poisson import PtdPoisson + +class TestPtdPoisson(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdPoisson() + + def test_poisson_distribution_scalar(self): + """Test instantiation of poisson distribution.""" + dist = self.node.f("3")[0] + self.assertTrue(isinstance(dist, torch.distributions.poisson.Poisson), "Probability distribution is not Poisson") + + def test_poisson_distribution_1d(self): + """Test instantiation of poisson distribution.""" + dist = self.node.f("(3,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.poisson.Poisson), "Probability distribution is not Poisson") + + def test_poisson_distribution_2d(self): + """Test instantiation of poisson distribution.""" + dist = self.node.f("(2,3)")[0] + self.assertTrue(isinstance(dist, torch.distributions.poisson.Poisson), "Probability distribution is not Poisson") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_student_t_test.py b/tests/pytorch_wrapper/ptd_student_t_test.py new file mode 100644 index 0000000..da21c0b --- /dev/null +++ b/tests/pytorch_wrapper/ptd_student_t_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_student_t import PtdStudentT + +class TestPtdStudentT(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdStudentT() + + def test_student_t_distribution_scalar(self): + """Test instantiation of student_t distribution.""" + dist = self.node.f("10.0", "0.0","1.0")[0] + self.assertTrue(isinstance(dist, torch.distributions.studentT.StudentT), "Probability distribution is not StudentT") + + def test_student_t_distribution_1d(self): + """Test instantiation of student_t distribution.""" + dist = self.node.f("(10.0,)", "(0.0,)", "(1.0,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.studentT.StudentT), "Probability distribution is not StudentT") + + def test_student_t_distribution_2d(self): + """Test instantiation of student_t distribution.""" + dist = self.node.f("(10.0, 11.0)", "(155.0,170.0)","(8.0,10.0)")[0] + self.assertTrue(isinstance(dist, torch.distributions.studentT.StudentT), "Probability distribution is not StudentT") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptd_uniform_test.py b/tests/pytorch_wrapper/ptd_uniform_test.py new file mode 100644 index 0000000..fc5f2a3 --- /dev/null +++ b/tests/pytorch_wrapper/ptd_uniform_test.py @@ -0,0 +1,34 @@ +import os +import sys +import unittest +import torch + +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_uniform import PtdUniform + +class TestPtdUniform(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.node = PtdUniform() + + def test_uniform_distribution_scalar(self): + """Test instantiation of uniform distribution.""" + dist = self.node.f("0.0","1.0")[0] + self.assertTrue(isinstance(dist, torch.distributions.uniform.Uniform), "Probability distribution is not Uniform") + + def test_uniform_distribution_1d(self): + """Test instantiation of uniform distribution.""" + dist = self.node.f("(0.0,)", "(1.0,)")[0] + self.assertTrue(isinstance(dist, torch.distributions.uniform.Uniform), "Probability distribution is not Uniform") + + def test_uniform_distribution_2d(self): + """Test instantiation of uniform distribution.""" + dist = self.node.f("(155.0,170.0)","(160.0,175.0)")[0] + self.assertTrue(isinstance(dist, torch.distributions.uniform.Uniform), "Probability distribution is not Uniform") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_cdf_test.py b/tests/pytorch_wrapper/ptdm_cdf_test.py new file mode 100644 index 0000000..7ae9d43 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_cdf_test.py @@ -0,0 +1,117 @@ +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.ptd_binomial import PtdBinomial +from pytorch_wrapper.ptdm_cdf import PtdmCdf +from torch.distributions import ( + Binomial, 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.""" + self.test_cases = { + # "Beta": [(2.0, 5.0, 0.3, scst.beta.cdf(0.3, 2, 5))], + "Gamma": [(2.0, 1.0, 3.0, scst.gamma.cdf(3.0, 2, scale=1))], + "Normal": [(0.0, 1.0, -1.5, scst.norm.cdf(-1.5, 0, 1))], + "Exponential": [(1.0, 3.0, scst.expon.cdf(3.0, 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))], + # "StudentT": [(1.0, 0.0, scst.t.cdf(0.0, 1))], + "Uniform": [(0.0, 1.0, 0.5, scst.uniform.cdf(0.5, loc=0, scale=1))], + "Chi2": [(1.0, 2.0, scst.chi2.cdf(2.0, 1))], # df=1, value=2.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): + if dist_name == "Binomial": + distribution = PtdBinomial().f(str(case[0]), str(case[2]), "")[0] # n, p + cdf = PtdmCdf().f(distribution, str(case[1]))[0] # k + else: + # Instantiate PyTorch distribution directly if not in wrapper + if dist_name == "Poisson": + distribution = Poisson(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.int64) + elif 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 == "Bernoulli": + distribution = Bernoulli(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + elif dist_name == "Beta": + distribution = Beta(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + 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 == "StudentT": + distribution = StudentT(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], 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() diff --git a/tests/pytorch_wrapper/ptdm_log_prob_test.py b/tests/pytorch_wrapper/ptdm_log_prob_test.py new file mode 100644 index 0000000..0db04e6 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_log_prob_test.py @@ -0,0 +1,122 @@ +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.ptd_binomial import PtdBinomial +from pytorch_wrapper.ptdm_log_prob import PtdmLogProb +from torch.distributions import ( + Binomial, Categorical, Poisson, NegativeBinomial, Multinomial, + Bernoulli, Beta, Gamma, Normal, Laplace, Exponential, LogNormal, + Pareto, Weibull, Dirichlet, VonMises, StudentT, Uniform +) + + +class TestPtdfLogProb(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Binomial": [(10, k, 0.5, np.log(scst.binom.pmf(k, 10, 0.5))) for k in range(11)], + "Poisson": [(3.0, k, np.log(scst.poisson.pmf(k, 3.0))) for k in range(10)], + "Categorical": [(torch.tensor([0.2, 0.8]), 0, np.log(0.2)), (torch.tensor([0.2, 0.8]), 1, np.log(0.8))], + "NegativeBinomial": [(5, k, 0.5, np.log(scst.nbinom.pmf(k, 5, 0.5))) for k in range(10)], + "Bernoulli": [(0.5, k, np.log(scst.bernoulli.pmf(k, 0.5))) for k in [0, 1]], + "Beta": [(2.0, 5.0, 0.3, np.log(scst.beta.pdf(0.3, 2, 5)))], + "Gamma": [(2.0, 1.0, 3.0, np.log(scst.gamma.pdf(3.0, 2, scale=1)))], + "Normal": [(0.0, 1.0, -1.5, np.log(scst.norm.pdf(-1.5, 0, 1)))], + "Laplace": [(0.0, 1.0, -1.5, np.log(scst.laplace.pdf(-1.5, 0, 1)))], + "Exponential": [(1.0, 3.0, np.log(scst.expon.pdf(3.0, scale=1)))], + "LogNormal": [(0.0, 1.0, 2.0, np.log(scst.lognorm.pdf(2.0, s=1, scale=np.exp(0))))], + "Pareto": [(2.5, 3.0, 3.5, np.log(scst.pareto.pdf(3.5, b=3.0, scale=2.5)))], + "Weibull": [(2.0, 3.0, np.log(scst.weibull_min.pdf(3.0, c=2.0, scale=1.0)))], + "Dirichlet": [(torch.tensor([1.0, 1.0, 1.0]), torch.tensor([0.2, 0.5, 0.3]), np.log(scst.dirichlet.pdf([0.2, 0.5, 0.3], [1.0, 1.0, 1.0])))], + "VonMises": [(0.0, 1.0, 1.0, np.log(scst.vonmises.pdf(1.0, kappa=1.0, loc=0)))], + "StudentT": [(1.0, 0.0, np.log(scst.t.pdf(0.0, 1)))], + "Uniform": [(0.0, 1.0, 0.5, np.log(scst.uniform.pdf(0.5, loc=0, scale=1)))], + } + + 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): + if dist_name == "Binomial": + distribution = PtdBinomial().f(str(case[0]), str(case[2]), "")[0] # n, p + log_prob = PtdmLogProb().f(distribution, str(case[1]))[0] # k + else: + # Instantiate PyTorch distribution directly if not in wrapper + if dist_name == "Poisson": + distribution = Poisson(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.int64) + elif 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 == "Bernoulli": + distribution = Bernoulli(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + elif dist_name == "Beta": + distribution = Beta(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + 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 == "StudentT": + distribution = StudentT(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], 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) + else: + self.fail(f"Unhandled distribution: {dist_name}") + + # Call wrapper log-prob method + if dist_name == "Dirichlet": + log_prob = PtdmLogProb().f(distribution, str(k_tensor.tolist()))[0] + else: + log_prob = PtdmLogProb().f(distribution, str(k_tensor.item()))[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 + 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}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_pdf_test.py b/tests/pytorch_wrapper/ptdm_pdf_test.py new file mode 100644 index 0000000..50fc671 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_pdf_test.py @@ -0,0 +1,113 @@ +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.ptd_binomial import PtdBinomial +from pytorch_wrapper.ptdm_pdf import PtdmPdf +from torch.distributions import ( + Binomial, Categorical, Poisson, NegativeBinomial, Multinomial, + Bernoulli, Beta, Gamma, Normal, Laplace, Exponential, LogNormal, + Pareto, Weibull, Dirichlet, VonMises, StudentT, Uniform +) + + +class TestPtdfPdf(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Beta": [(2.0, 5.0, 0.3, scst.beta.pdf(0.3, 2, 5))], + "Gamma": [(2.0, 1.0, 3.0, scst.gamma.pdf(3.0, 2, scale=1))], + "Normal": [(0.0, 1.0, -1.5, scst.norm.pdf(-1.5, 0, 1))], + "Exponential": [(1.0, 3.0, scst.expon.pdf(3.0, scale=1))], + "Dirichlet": [(torch.tensor([1.0, 1.0, 1.0]), torch.tensor([0.2, 0.5, 0.3]), scst.dirichlet.pdf([0.2, 0.5, 0.3], [1.0, 1.0, 1.0]))], + "VonMises": [(0.0, 1.0, 1.0, scst.vonmises.pdf(1.0, kappa=1.0, loc=0))], + "StudentT": [(1.0, 0.0, scst.t.pdf(0.0, 1))], + "Uniform": [(0.0, 1.0, 0.5, scst.uniform.pdf(0.5, loc=0, scale=1))], + } + + 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): + if dist_name == "Binomial": + distribution = PtdBinomial().f(str(case[0]), str(case[2]), "")[0] # n, p + pdf = PtdmPdf().f(distribution, str(case[1]))[0] # k + else: + # Instantiate PyTorch distribution directly if not in wrapper + if dist_name == "Poisson": + distribution = Poisson(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.int64) + elif 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 == "Bernoulli": + distribution = Bernoulli(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + elif dist_name == "Beta": + distribution = Beta(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + 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 == "StudentT": + distribution = StudentT(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], 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) + else: + self.fail(f"Unhandled distribution: {dist_name}") + + # Call wrapper log-prob method + if dist_name == "Dirichlet": + pdf = PtdmPdf().f(distribution, str(k_tensor.tolist()))[0] + else: + pdf = PtdmPdf().f(distribution, str(k_tensor.item()))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(pdf, torch.Tensor), "Probability output is not a tensor") + + # Convert tensors to floats for comparison + 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}") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_pmf_test.py b/tests/pytorch_wrapper/ptdm_pmf_test.py new file mode 100644 index 0000000..ed40b8c --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_pmf_test.py @@ -0,0 +1,84 @@ +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() diff --git a/tests/pytorch_wrapper/ptdm_sample_test.py b/tests/pytorch_wrapper/ptdm_sample_test.py new file mode 100644 index 0000000..dee10e8 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_sample_test.py @@ -0,0 +1,42 @@ +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.ptd_binomial import PtdBinomial +from pytorch_wrapper.ptdm_sample import PtdmSample + +class TestPtdfSample(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + + def test_scalar_binom(self): + """Test binom scalar.""" + distribution_node = PtdBinomial().f(str(10), str(0.5), "")[0] + sample = PtdmSample().f(distribution_node, str(5))[0] + + self.assertEqual(sample.size(), (5,)) + + def test_1d_binom(self): + """Test binom 1d.""" + distribution_node = PtdBinomial().f(str(10), str(0.5), "")[0] + sample = PtdmSample().f(distribution_node, "(5,)")[0] + + self.assertEqual(sample.size(), (5,)) + + def test_2d_binom(self): + """Test binom 2d.""" + distribution_node = PtdBinomial().f(str(10), str(0.5), "")[0] + sample = PtdmSample().f(distribution_node, ("(2,3)"))[0] + + self.assertEqual(sample.size(), (2,3)) + + +if __name__ == "__main__": + unittest.main()