nodes/HowToSD-ComfyUI-Pt-Wrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added distribution-related nodes.

Browse Source
This commit is contained in:
HowToSD
2025-03-18 01:11:46 -07:00
parent b7d2ea5346
commit 39421a4f8b
51 changed files with 2196 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+48
View File
@@ -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",
docs/reference/node_reference.md
+19
View File
@@ -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 |
| --- | --- |
docs/reference/ptd_bernoulli.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_beta.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_binomial.md
+24
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_categorical.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_chi2.md
+20
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_exponential.md
+20
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_gamma.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_normal.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_poisson.md
+20
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_student_t.md
+24
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptd_uniform.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptdm_cdf.md
+25
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptdm_log_prob.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptdm_pdf.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptdm_pmf.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptdm_sample.md
+22
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Distribution
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
modules/pytorch_wrapper/ptd_bernoulli.py
+59
View File
@@ -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,)
modules/pytorch_wrapper/ptd_beta.py
+52
View File
@@ -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,)
modules/pytorch_wrapper/ptd_binomial.py
+69
View File
@@ -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,)
modules/pytorch_wrapper/ptd_categorical.py
+60
View File
@@ -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,)
modules/pytorch_wrapper/ptd_chi2.py
+50
View File
@@ -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,)
modules/pytorch_wrapper/ptd_exponential.py
+48
View File
@@ -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,)
modules/pytorch_wrapper/ptd_gamma.py
+53
View File
@@ -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,)
modules/pytorch_wrapper/ptd_normal.py
+53
View File
@@ -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,)
modules/pytorch_wrapper/ptd_poisson.py
+48
View File
@@ -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,)
modules/pytorch_wrapper/ptd_student_t.py
+61
View File
@@ -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,)
modules/pytorch_wrapper/ptd_uniform.py
+52
View File
@@ -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,)
modules/pytorch_wrapper/ptdm_cdf.py
+60
View File
@@ -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),)
modules/pytorch_wrapper/ptdm_log_prob.py
+61
View File
@@ -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),)
modules/pytorch_wrapper/ptdm_pdf.py
+57
View File
@@ -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)),)
modules/pytorch_wrapper/ptdm_pmf.py
+52
View File
@@ -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)),)
modules/pytorch_wrapper/ptdm_sample.py
+52
View File
@@ -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),)
pyproject.toml
+1 -1
View File
@@ -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]
tests/pytorch_wrapper/ptd_bernoulli_test.py
+50
View File
@@ -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()
tests/pytorch_wrapper/ptd_beta_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptd_binomial_test.py
+59
View File
@@ -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()
tests/pytorch_wrapper/ptd_categorical_test.py
+29
View File
@@ -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()
tests/pytorch_wrapper/ptd_chi2_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptd_exponential_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptd_gamma_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptd_normal_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptd_poisson_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptd_student_t_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptd_uniform_test.py
+34
View File
@@ -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()
tests/pytorch_wrapper/ptdm_cdf_test.py
+117
View File
@@ -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()
tests/pytorch_wrapper/ptdm_log_prob_test.py
+122
View File
@@ -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()
tests/pytorch_wrapper/ptdm_pdf_test.py
+113
View File
@@ -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()
tests/pytorch_wrapper/ptdm_pmf_test.py
+84
View File
@@ -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()
tests/pytorch_wrapper/ptdm_sample_test.py
+42
View File
@@ -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()