From ae8a6cfb902c2206ab05aacd4812978289d2e3c6 Mon Sep 17 00:00:00 2001 From: HowToSD Date: Wed, 19 Mar 2025 13:40:49 -0700 Subject: [PATCH] Implemented icdf for distribution nodes. --- __init__.py | 6 ++ docs/reference/node_reference.md | 2 + docs/reference/ptdm_icdf.md | 25 +++++ docs/reference/ptdm_icdf_tensor.md | 25 +++++ .../normal_distribution_cdf_icdf.json | 1 + modules/pytorch_wrapper/ptd_bernoulli.py | 17 ++- modules/pytorch_wrapper/ptd_beta.py | 16 +++ modules/pytorch_wrapper/ptd_binomial.py | 19 +++- modules/pytorch_wrapper/ptd_chi2.py | 25 ++++- modules/pytorch_wrapper/ptd_gamma.py | 28 ++++- modules/pytorch_wrapper/ptd_poisson.py | 17 ++- modules/pytorch_wrapper/ptd_student_t.py | 17 +++ modules/pytorch_wrapper/ptdm_icdf.py | 57 ++++++++++ modules/pytorch_wrapper/ptdm_icdf_tensor.py | 49 +++++++++ pyproject.toml | 2 +- .../ptdm_cdf_tensor_beta_test.py | 71 +++++++++++++ .../ptdm_icdf_tensor_bernoulli_test.py | 67 ++++++++++++ .../ptdm_icdf_tensor_beta_test.py | 71 +++++++++++++ .../ptdm_icdf_tensor_binomial_test.py | 66 ++++++++++++ .../ptdm_icdf_tensor_chi2_test.py | 63 +++++++++++ .../ptdm_icdf_tensor_gamma_test.py | 62 +++++++++++ .../ptdm_icdf_tensor_poisson_test.py | 68 ++++++++++++ .../ptdm_icdf_tensor_student_t_test.py | 71 +++++++++++++ .../pytorch_wrapper/ptdm_icdf_tensor_test.py | 47 ++++++++ tests/pytorch_wrapper/ptdm_icdf_test.py | 100 ++++++++++++++++++ 25 files changed, 986 insertions(+), 6 deletions(-) create mode 100644 docs/reference/ptdm_icdf.md create mode 100644 docs/reference/ptdm_icdf_tensor.md create mode 100644 examples/workflows/normal_distribution_cdf_icdf.json create mode 100644 modules/pytorch_wrapper/ptdm_icdf.py create mode 100644 modules/pytorch_wrapper/ptdm_icdf_tensor.py create mode 100644 tests/pytorch_wrapper/ptdm_cdf_tensor_beta_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_bernoulli_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_beta_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_binomial_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_chi2_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_gamma_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_poisson_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_student_t_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_tensor_test.py create mode 100644 tests/pytorch_wrapper/ptdm_icdf_test.py diff --git a/__init__.py b/__init__.py index 463a93f..c7bf831 100644 --- a/__init__.py +++ b/__init__.py @@ -127,6 +127,8 @@ from .modules.pytorch_wrapper.ptd_student_t import PtdStudentT from .modules.pytorch_wrapper.ptd_uniform import PtdUniform from .modules.pytorch_wrapper.ptdm_cdf import PtdmCdf from .modules.pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor +from .modules.pytorch_wrapper.ptdm_icdf import PtdmIcdf +from .modules.pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor from .modules.pytorch_wrapper.ptdm_log_prob import PtdmLogProb from .modules.pytorch_wrapper.ptdm_log_prob_tensor import PtdmLogProbTensor from .modules.pytorch_wrapper.ptdm_pdf import PtdmPdf @@ -303,6 +305,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtdUniform": PtdUniform, "PtdmCdf": PtdmCdf, "PtdmCdfTensor": PtdmCdfTensor, + "PtdmIcdf": PtdmIcdf, + "PtdmIcdfTensor": PtdmIcdfTensor, "PtdmLogProb": PtdmLogProb, "PtdmLogProbTensor": PtdmLogProbTensor, "PtdmPdf": PtdmPdf, @@ -479,6 +483,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtdUniform": "Ptd Uniform", "PtdmCdf": "Ptdm Cdf", "PtdmCdfTensor": "Ptdm Cdf Tensor", + "PtdmIcdf": "Ptdm Icdf", + "PtdmIcdfTensor": "Ptdm Icdf Tensor", "PtdmLogProb": "Ptdm Log Prob", "PtdmLogProbTensor": "Ptdm Log Prob Tensor", "PtdmPdf": "Ptdm Pdf", diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index f0e7537..b572151 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -76,6 +76,8 @@ | [Ptd Uniform](ptd_uniform.md) | Instantiates a Uniform distribution object. | | [Ptdm Cdf](ptdm_cdf.md) | Computes the cumulative distribution function for the input distribution. | | [Ptdm Cdf Tensor](ptdm_cdf_tensor.md) | Computes the cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. | +| [Ptdm Icdf](ptdm_icdf.md) | Computes the inverse of the cumulative distribution function for the input distribution. | +| [Ptdm Icdf Tensor](ptdm_icdf_tensor.md) | Computes the inverse of cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. | | [Ptdm Log Prob](ptdm_log_prob.md) | Computes the log of probability for the input distribution. | | [Ptdm Log Prob Tensor](ptdm_log_prob_tensor.md) | Computes the log of probability for the input distribution. This nodes accepts a tensor so it can be used to compute log of probability for multiple values contained in a tensor. | | [Ptdm Pdf](ptdm_pdf.md) | Computes the probability density for the input distribution. | diff --git a/docs/reference/ptdm_icdf.md b/docs/reference/ptdm_icdf.md new file mode 100644 index 0000000..2e3ddd3 --- /dev/null +++ b/docs/reference/ptdm_icdf.md @@ -0,0 +1,25 @@ +# Ptdm Icdf +Computes the inverse of the cumulative distribution function for the input distribution. + +**Note** +Icdf is not supported for all distributions in PyTorch. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + q (float): Value + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| q | 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_icdf_tensor.md b/docs/reference/ptdm_icdf_tensor.md new file mode 100644 index 0000000..8f31fe2 --- /dev/null +++ b/docs/reference/ptdm_icdf_tensor.md @@ -0,0 +1,25 @@ +# Ptdm Icdf Tensor +Computes the inverse of cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. + +**Note** +Icdf is not supported for all distributions in PyTorch. + +Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): q in Tensor. + +## Input +| Name | Data type | +|---|---| +| distribution | Ptdistribution | +| tens | Tensor | + +## Output +| Data type | +|---| +| Tensor | + +
+Category: PyTorch wrapper - Distribution + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/normal_distribution_cdf_icdf.json b/examples/workflows/normal_distribution_cdf_icdf.json new file mode 100644 index 0000000..6c25fb2 --- /dev/null +++ b/examples/workflows/normal_distribution_cdf_icdf.json @@ -0,0 +1 @@ +{"last_node_id":233,"last_link_id":209,"nodes":[{"id":27,"type":"PtShowText","pos":[666.5792236328125,2204.9599609375],"size":[210,88],"flags":{},"order":5,"mode":0,"inputs":[{"name":"tens","localized_name":"tens","type":"TENSOR","link":163}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PtShowText"},"widgets_values":["tensor(0.8413)"]},{"id":233,"type":"PtStack","pos":[972.0777587890625,2692.753173828125],"size":[315,78],"flags":{},"order":7,"mode":0,"inputs":[{"name":"tens_a","localized_name":"tens_a","type":"TENSOR","link":208},{"name":"tens_b","localized_name":"tens_b","type":"TENSOR","link":207}],"outputs":[{"name":"TENSOR","localized_name":"TENSOR","type":"TENSOR","links":[209],"slot_index":0}],"properties":{"Node name for S&R":"PtStack"},"widgets_values":[1]},{"id":224,"type":"PtShowText","pos":[1405.767822265625,2692.753173828125],"size":[210.00001525878906,573.8150024414062],"flags":{},"order":8,"mode":0,"inputs":[{"name":"tens","localized_name":"tens","type":"TENSOR","link":209}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PtShowText"},"widgets_values":["tensor([[0.8000, 0.8416],\n [0.8050, 0.8596],\n [0.8100, 0.8779],\n [0.8150, 0.8965],\n [0.8200, 0.9154],\n [0.8250, 0.9346],\n [0.8300, 0.9542],\n [0.8350, 0.9741],\n [0.8400, 0.9945],\n [0.8450, 1.0152],\n [0.8500, 1.0364],\n [0.8550, 1.0581],\n [0.8600, 1.0803],\n [0.8650, 1.1031],\n [0.8700, 1.1264],\n [0.8750, 1.1503],\n [0.8800, 1.1750],\n [0.8850, 1.2004],\n [0.8900, 1.2265],\n [0.8950, 1.2536],\n [0.9000, 1.2816],\n [0.9050, 1.3106],\n [0.9100, 1.3408],\n [0.9150, 1.3722],\n [0.9200, 1.4051],\n [0.9250, 1.4395],\n [0.9300, 1.4758],\n [0.9350, 1.5141],\n [0.9400, 1.5548],\n [0.9450, 1.5982],\n [0.9500, 1.6449],\n [0.9550, 1.6954],\n [0.9600, 1.7507],\n [0.9650, 1.8119],\n [0.9700, 1.8808],\n [0.9750, 1.9600],\n [0.9800, 2.0537],\n [0.9850, 2.1701],\n [0.9900, 2.3263],\n [0.9950, 2.5758]])"]},{"id":228,"type":"PtdmIcdf","pos":[980.1187133789062,2488.673828125],"size":[210,58],"flags":{},"order":3,"mode":0,"inputs":[{"name":"distribution","localized_name":"distribution","type":"PTDISTRIBUTION","link":205}],"outputs":[{"name":"TENSOR","localized_name":"TENSOR","type":"TENSOR","links":[204],"slot_index":0}],"properties":{"Node name for S&R":"PtdmIcdf"},"widgets_values":["0.975"]},{"id":191,"type":"PtdmCdf","pos":[666.5792236328125,2093.677978515625],"size":[210,58],"flags":{},"order":2,"mode":0,"inputs":[{"name":"distribution","localized_name":"distribution","type":"PTDISTRIBUTION","link":162}],"outputs":[{"name":"TENSOR","localized_name":"TENSOR","type":"TENSOR","links":[163],"slot_index":0}],"properties":{"Node name for S&R":"PtdmCdf"},"widgets_values":["1"]},{"id":223,"type":"PtdmIcdfTensor","pos":[666.5792236328125,2692.753173828125],"size":[177.26657104492188,55.07500076293945],"flags":{},"order":4,"mode":0,"inputs":[{"name":"distribution","localized_name":"distribution","type":"PTDISTRIBUTION","link":200},{"name":"tens","localized_name":"tens","type":"TENSOR","link":206}],"outputs":[{"name":"TENSOR","localized_name":"TENSOR","type":"TENSOR","links":[207],"slot_index":0}],"properties":{"Node name for S&R":"PtdmIcdfTensor"},"widgets_values":[]},{"id":232,"type":"PtArange","pos":[358.91302490234375,2488.673828125],"size":[213.96755981445312,130],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"TENSOR","localized_name":"TENSOR","type":"TENSOR","links":[206,208],"slot_index":0}],"properties":{"Node name for S&R":"PtArange"},"widgets_values":["0.80","1.00","0.005","float32"]},{"id":28,"type":"PtdNormal","pos":[358.91302490234375,2098.1396484375],"size":[210,83.6666488647461],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTDISTRIBUTION","localized_name":"PTDISTRIBUTION","type":"PTDISTRIBUTION","links":[162,200,205],"slot_index":0}],"properties":{"Node name for S&R":"PtdNormal"},"widgets_values":["0","1"]},{"id":227,"type":"PtShowText","pos":[1405.767822265625,2488.673828125],"size":[210,88],"flags":{},"order":6,"mode":0,"inputs":[{"name":"tens","localized_name":"tens","type":"TENSOR","link":204}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PtShowText"},"widgets_values":["tensor(1.9600)"]}],"links":[[162,28,0,191,0,"PTDISTRIBUTION"],[163,191,0,27,0,"TENSOR"],[200,28,0,223,0,"PTDISTRIBUTION"],[204,228,0,227,0,"TENSOR"],[205,28,0,228,0,"PTDISTRIBUTION"],[206,232,0,223,1,"TENSOR"],[207,223,0,233,1,"TENSOR"],[208,232,0,233,0,"TENSOR"],[209,233,0,224,0,"TENSOR"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.6351091783959867,"offset":[253.49515136244105,-2007.0438933201194]},"node_versions":{"pt-wrapper":"8be93739ad263e040f5fb5091b67b833f5b0fbcc"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptd_bernoulli.py b/modules/pytorch_wrapper/ptd_bernoulli.py index d52f465..7e88136 100644 --- a/modules/pytorch_wrapper/ptd_bernoulli.py +++ b/modules/pytorch_wrapper/ptd_bernoulli.py @@ -24,7 +24,22 @@ class BernoulliEx(Bernoulli): a = x.detach().cpu().numpy() probs = self.probs.detach().cpu().numpy() outputs = scst.bernoulli.cdf(a, probs) - return torch.tensor(outputs, dtype=torch.float32) + return torch.tensor(outputs, dtype=torch.float32).to(x.device) + + def icdf(self, q: torch.Tensor): + """ + Computes the inverse of cumulative distribution function (CDF) of the Bernoulli distribution. + + Args: + q (torch.Tensor): Input tensor containing q. + + Returns: + torch.Tensor: Tensor containing the inverse of cumulative probabilities. + """ + a = q.detach().cpu().numpy() + probs = self.probs.detach().cpu().numpy() + outputs = scst.bernoulli.ppf(a, probs) + return torch.tensor(outputs, dtype=torch.float32).to(q.device) class PtdBernoulli: diff --git a/modules/pytorch_wrapper/ptd_beta.py b/modules/pytorch_wrapper/ptd_beta.py index d25d057..7650ff3 100644 --- a/modules/pytorch_wrapper/ptd_beta.py +++ b/modules/pytorch_wrapper/ptd_beta.py @@ -26,6 +26,22 @@ class BetaEx(Beta): outputs = scst.beta.cdf(x2, a, b) return torch.tensor(outputs, dtype=torch.float32).to(x.device) + def icdf(self, q: torch.Tensor): + """ + Computes the inverse of cumulative distribution function (CDF) of the Beta distribution. + + Args: + q (torch.Tensor): Input tensor containing values where the ICDF is evaluated. + + Returns: + torch.Tensor: Tensor containing the inverse of cumulative probabilities. + """ + x2 = q.detach().cpu().numpy() + a = self.concentration1.detach().cpu().item() + b = self.concentration0.detach().cpu().item() + outputs = scst.beta.ppf(x2, a, b) + return torch.tensor(outputs, dtype=torch.float32).to(q.device) + class PtdBeta: """ diff --git a/modules/pytorch_wrapper/ptd_binomial.py b/modules/pytorch_wrapper/ptd_binomial.py index 171743f..416ad50 100644 --- a/modules/pytorch_wrapper/ptd_binomial.py +++ b/modules/pytorch_wrapper/ptd_binomial.py @@ -24,7 +24,24 @@ class BinomialEx(Binomial): n = self.total_count.detach().cpu().numpy() probs = self.probs.detach().cpu().numpy() outputs = scst.binom.cdf(a, n, probs) - return torch.tensor(outputs, dtype=torch.float32) + return torch.tensor(outputs, dtype=torch.float32).to(x.device) + + def icdf(self, q: torch.Tensor): + """ + Computes the inverse of cumulative distribution function (CDF) of the Binomial distribution. + + Args: + q (torch.Tensor): Input tensor containing values where the ICDF is evaluated. + + Returns: + torch.Tensor: Tensor containing the inverse of cumulative probabilities. + """ + a = q.detach().cpu().numpy() + n = self.total_count.detach().cpu().numpy() + probs = self.probs.detach().cpu().numpy() + outputs = scst.binom.ppf(a, n, probs) + return torch.tensor(outputs, dtype=torch.float32).to(q.device) + class PtdBinomial: diff --git a/modules/pytorch_wrapper/ptd_chi2.py b/modules/pytorch_wrapper/ptd_chi2.py index 789f4dc..2f9d589 100644 --- a/modules/pytorch_wrapper/ptd_chi2.py +++ b/modules/pytorch_wrapper/ptd_chi2.py @@ -2,6 +2,29 @@ from typing import Any, Dict import ast import torch from torch.distributions import Chi2 +import scipy.stats as scst + +class Chi2Ex(Chi2): + """ + A class to extend Chi2 to add missing functionality. + + pragma: skip_doc + """ + def icdf(self, q: torch.Tensor): + """ + Computes the inverse of cumulative distribution function (CDF) of the Chi-squared distribution. + + Args: + q (torch.Tensor): Input tensor containing values where the ICDF is evaluated. + + Returns: + torch.Tensor: Tensor containing the inverse of cumulative probabilities. + """ + a = q.detach().cpu().numpy() + df = self.df.detach().cpu().item() + outputs = scst.chi2.ppf(a, df) + return torch.tensor(outputs, dtype=torch.float32).to(q.device) + class PtdChi2: """ @@ -45,6 +68,6 @@ class PtdChi2: """ d = ast.literal_eval(df) - dist = Chi2( + dist = Chi2Ex( df=torch.tensor(d, dtype=torch.float32)) return (dist,) diff --git a/modules/pytorch_wrapper/ptd_gamma.py b/modules/pytorch_wrapper/ptd_gamma.py index f2e4fbe..54b8a64 100644 --- a/modules/pytorch_wrapper/ptd_gamma.py +++ b/modules/pytorch_wrapper/ptd_gamma.py @@ -2,6 +2,32 @@ from typing import Any, Dict import ast import torch from torch.distributions import Gamma +import scipy.stats as scst + + +class GammaEx(Gamma): + """ + A class to extend Gamma to add missing functionality. + + pragma: skip_doc + """ + + def icdf(self, q: torch.Tensor): + """ + Computes the inverse of cumulative distribution function (CDF) of the Gamma distribution. + + Args: + q (torch.Tensor): Input tensor containing values where the ICDF is evaluated. + + Returns: + torch.Tensor: Tensor containing the inverse of cumulative probabilities. + """ + x2 = q.detach().cpu().numpy() + a = self.concentration.detach().cpu().item() + b = self.rate.detach().cpu().item() + outputs = scst.gamma.ppf(x2, a, 1/b) # Note that you need to pass the inverse of b. + return torch.tensor(outputs, dtype=torch.float32).to(q.device) + class PtdGamma: """ @@ -47,7 +73,7 @@ class PtdGamma: c = ast.literal_eval(alpha) r = ast.literal_eval(beta) - dist = Gamma( + dist = GammaEx( concentration=torch.tensor(c, dtype=torch.float32), rate=torch.tensor(r, dtype=torch.float32)) return (dist,) diff --git a/modules/pytorch_wrapper/ptd_poisson.py b/modules/pytorch_wrapper/ptd_poisson.py index 9006f27..97fe1ad 100644 --- a/modules/pytorch_wrapper/ptd_poisson.py +++ b/modules/pytorch_wrapper/ptd_poisson.py @@ -23,7 +23,22 @@ class PoissonEx(Poisson): a = x.detach().cpu().numpy() rate = self.rate.detach().cpu().numpy() outputs = scst.poisson.cdf(a, rate) - return torch.tensor(outputs, dtype=torch.float32) + return torch.tensor(outputs, dtype=torch.float32).to(x.device) + + def icdf(self, q: torch.Tensor): + """ + Computes the inverse of cumulative distribution function (CDF) of the Poisson distribution. + + Args: + q (torch.Tensor): Input tensor containing values where the ICDF is evaluated. + + Returns: + torch.Tensor: Tensor containing the inverse of cumulative probabilities. + """ + a = q.detach().cpu().numpy() + rate = self.rate.detach().cpu().numpy() + outputs = scst.poisson.ppf(a, rate) + return torch.tensor(outputs, dtype=torch.float32).to(q.device) class PtdPoisson: diff --git a/modules/pytorch_wrapper/ptd_student_t.py b/modules/pytorch_wrapper/ptd_student_t.py index adefe36..3c4f34b 100644 --- a/modules/pytorch_wrapper/ptd_student_t.py +++ b/modules/pytorch_wrapper/ptd_student_t.py @@ -27,6 +27,23 @@ class StudentTEx(StudentT): outputs = scst.t.cdf((a - loc)/scale, df) return torch.tensor(outputs, dtype=torch.float32).to(x.device) + def icdf(self, q: torch.Tensor): + """ + Computes the inverse of cumulative distribution function (CDF) of the StudentT distribution. + + Args: + q (torch.Tensor): Input tensor containing values where the ICDF is evaluated. + + Returns: + torch.Tensor: Tensor containing the inverse of cumulative probabilities. + """ + a = q.detach().cpu().numpy() + df = self.df.detach().cpu().item() + loc = self.loc.detach().cpu().numpy() + scale = self.scale.detach().cpu().numpy() + outputs = scst.t.ppf(a, df) * scale + loc + return torch.tensor(outputs, dtype=torch.float32).to(q.device) + class PtdStudentT: """ diff --git a/modules/pytorch_wrapper/ptdm_icdf.py b/modules/pytorch_wrapper/ptdm_icdf.py new file mode 100644 index 0000000..bcc17bc --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_icdf.py @@ -0,0 +1,57 @@ +from typing import Any, Dict +import torch +from torch.distributions import * +import ast + +class PtdmIcdf: + """ + Computes the inverse of the cumulative distribution function for the input distribution. + + **Note** + Icdf is not supported for all distributions in PyTorch. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + q (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", {}), + "q": ("STRING", {"default": ""}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, q: str) -> tuple: + """ + Computes the inverse of the cumulative distribution function for the input distribution. + + Args: + distribution (torch.distributions.Distribution): Distribution. + q (str): Value + + Returns: + tuple containing the probability + """ + v2 = ast.literal_eval(q) + 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.") + + tens = torch.tensor(v2, dtype=torch.float32) + + return (distribution.icdf(tens),) diff --git a/modules/pytorch_wrapper/ptdm_icdf_tensor.py b/modules/pytorch_wrapper/ptdm_icdf_tensor.py new file mode 100644 index 0000000..b2298d7 --- /dev/null +++ b/modules/pytorch_wrapper/ptdm_icdf_tensor.py @@ -0,0 +1,49 @@ +from typing import Any, Dict +import torch + + +class PtdmIcdfTensor: + """ + Computes the inverse of cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. + + **Note** + Icdf is not supported for all distributions in PyTorch. + + Args: + distribution (torch.distributions.distribution.Distribution): Distribution. + tens (torch.Tensor): q in Tensor. + + category: PyTorch wrapper - Distribution + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the input types for the function. + + Returns: + Dict[str, Any]: A dictionary specifying the required input types. + """ + return { + "required": { + "distribution": ("PTDISTRIBUTION", {}), + "tens": ("TENSOR", {}) + } + } + + RETURN_TYPES: tuple = ("TENSOR",) + FUNCTION: str = "f" + CATEGORY: str = "Distribution" + + def f(self, distribution: torch.distributions.Distribution, tens: torch.Tensor) -> tuple: + """ + Computes the inverse of cumulative distribution function for the input distribution. This nodes accepts a tensor so it can be used to compute cdf for multiple values contained in a tensor. + + Args: + distribution (torch.distributions.Distribution): Distribution. + tens (torch.Tensor): q in Tensor. + + Returns: + tuple containing the probability + """ + return (distribution.icdf(tens),) diff --git a/pyproject.toml b/pyproject.toml index 2b75f44..6c551c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "pt-wrapper" description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training." -version = "3.1.2" +version = "3.1.3" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/ptdm_cdf_tensor_beta_test.py b/tests/pytorch_wrapper/ptdm_cdf_tensor_beta_test.py new file mode 100644 index 0000000..f373dc3 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_cdf_tensor_beta_test.py @@ -0,0 +1,71 @@ +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_beta import BetaEx +from pytorch_wrapper.ptd_beta import PtdBeta +from pytorch_wrapper.ptdm_cdf_tensor import PtdmCdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfCdfBeta(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Beta": [ + # str alpha, str beta, num x, expected + ("2", "8", 0.5, scst.beta.cdf(0.5, 2, 8)), + ("2", "8", 0, scst.beta.cdf(0, 2, 8)), + ("2", "8", [0, 0.5, 1], scst.beta.cdf([0, 0.5, 1], 2, 8)) # 1d array for x + ] + } + + def test_custom_class_1(self): + d = BetaEx(concentration1=2, concentration0=8) + actual = d.cdf(torch.tensor(0.5)).numpy() + expected = torch.tensor(0.98046875).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = BetaEx(concentration1=2, concentration0=8) + actual = d.cdf(torch.tensor(0.0)).numpy() + expected = torch.tensor(0.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_3(self): + d = BetaEx(concentration1=2, concentration0=8) + actual = d.cdf(torch.tensor(1.0)).numpy() + expected = torch.tensor(1.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdBeta().f( + alpha=case[0], + beta=case[1])[0] + cdf = PtdmCdfTensor().f( + distribution, + torch.tensor(case[2]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(cdf, torch.Tensor), "Probability output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_bernoulli_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_bernoulli_test.py new file mode 100644 index 0000000..16f5eaf --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_bernoulli_test.py @@ -0,0 +1,67 @@ +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_bernoulli import BernoulliEx +from pytorch_wrapper.ptd_bernoulli import PtdBernoulli +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfBernoulli(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Bernoulli": [ + (0.4, 0.6, scst.bernoulli.ppf(0.6, 0.4)), # P(0) = 0.6 + (0.4, 1, scst.bernoulli.ppf(1, 0.4)), + (0.4, [0.6, 1], scst.bernoulli.ppf([0.6, 1], 0.4)) # 1d array for x + ] + } + + def test_custom_class_1(self): + d = BernoulliEx(probs=0.4) + actual = d.icdf(torch.tensor(0.6)).numpy() + expected = torch.tensor(0.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = BernoulliEx(probs=0.4) + actual = d.icdf(torch.tensor(1.0)).numpy() + expected = torch.tensor(1.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_3(self): + d = BernoulliEx(probs=0.4) + actual = d.icdf(torch.tensor([0.6, 1.0])).numpy() + expected = torch.tensor([0.0, 1.0]).numpy() + self.assertTrue(np.allclose(actual, expected, atol=1e-5), + msg=f"Expected {expected}, got {actual}") + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdBernoulli().f(str(case[0]), "")[0] + cdf = PtdmIcdfTensor().f( + distribution, + torch.tensor(case[1]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(cdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_beta_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_beta_test.py new file mode 100644 index 0000000..66d37b2 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_beta_test.py @@ -0,0 +1,71 @@ +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_beta import BetaEx +from pytorch_wrapper.ptd_beta import PtdBeta +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfBeta(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Beta": [ + # str alpha, str beta, num q, expected + ("2", "8", 0.5, scst.beta.ppf(0.5, 2, 8)), + ("2", "8", 0, scst.beta.ppf(0, 2, 8)), + ("2", "8", [0, 0.5, 1], scst.beta.ppf([0, 0.5, 1], 2, 8)) # 1d array for q + ] + } + + def test_custom_class_1(self): + d = BetaEx(concentration1=2, concentration0=8) + actual = d.icdf(torch.tensor(0.98046875)).numpy() + expected = torch.tensor(0.5).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = BetaEx(concentration1=2, concentration0=8) + actual = d.icdf(torch.tensor(0.0)).numpy() + expected = torch.tensor(0.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_3(self): + d = BetaEx(concentration1=2, concentration0=8) + actual = d.icdf(torch.tensor(1.0)).numpy() + expected = torch.tensor(1.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdBeta().f( + alpha=case[0], + beta=case[1])[0] + icdf = PtdmIcdfTensor().f( + distribution, + torch.tensor(case[2]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(icdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, icdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_binomial_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_binomial_test.py new file mode 100644 index 0000000..88b7bad --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_binomial_test.py @@ -0,0 +1,66 @@ +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 BinomialEx +from pytorch_wrapper.ptd_binomial import PtdBinomial +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfBinomial(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Binomial": [ + (10.0, 0.5, 0.4, scst.binom.ppf(0.5, 10, 0.4)), + (10.0, [0.5, 0.6], 0.4, scst.binom.ppf([0.5, 0.6], 10, 0.4)) # 1d array for x + ] + } + + def test_custom_class_1(self): + d = BinomialEx(total_count=10, probs=0.4) + actual = d.icdf(torch.tensor(0.3822806015999999)).numpy() + expected = torch.tensor(3.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = BinomialEx(total_count=10, probs=0.4) + actual = d.icdf(torch.tensor(0.04635740159999998)).numpy() + expected = torch.tensor(1.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_3(self): + d = BinomialEx(total_count=10, probs=0.4) + actual = d.icdf(torch.tensor([0.04635740159999998, 0.3822806015999999])).numpy() + expected = torch.tensor([1.0, 3.0]).numpy() + self.assertTrue(np.allclose(actual, expected, atol=1e-5), + msg=f"Expected {expected}, got {actual}") + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdBinomial().f(str(case[0]), str(case[2]), "")[0] # n, p + cdf = PtdmIcdfTensor().f( + distribution, + torch.tensor(case[1]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(cdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_chi2_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_chi2_test.py new file mode 100644 index 0000000..0d173e3 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_chi2_test.py @@ -0,0 +1,63 @@ +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_chi2 import Chi2Ex +from pytorch_wrapper.ptd_chi2 import PtdChi2 +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfChi2(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Chi2": [ + # str df, num q, expected + ("40.0", 0.5, scst.chi2.ppf(0.5, 40)), + ("40.0", 0.9, scst.chi2.ppf(0.9, 40)), + ("40.0", [0.8, 0.9], scst.chi2.ppf([0.8, 0.9], 40)) # 1d array for q + ] + } + + def test_custom_class_1(self): + d = Chi2Ex(df=1) + actual = d.icdf(torch.tensor(0.01)).numpy() + expected = torch.tensor(0.00015708785790970184).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = Chi2Ex(df=24) + actual = d.icdf(torch.tensor(0.6)).numpy() + expected = torch.tensor(25.10634821892835).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdChi2().f( + df=case[0])[0] + icdf = PtdmIcdfTensor().f( + distribution, + torch.tensor(case[1]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(icdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, icdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_gamma_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_gamma_test.py new file mode 100644 index 0000000..d0dd64d --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_gamma_test.py @@ -0,0 +1,62 @@ +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_gamma import GammaEx +from pytorch_wrapper.ptd_gamma import PtdGamma +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfGamma(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Gamma": [ + # str alpha, str gamma, num q, expected + ("2", "8", 0.5, scst.gamma.ppf(0.5, 2, 1/8)), + ("2", "8", 0, scst.gamma.ppf(0, 2, 1/8)), + ("2", "8", [0, 0.5, 1], scst.gamma.ppf([0, 0.5, 1], 2, 1/8)) # 1d array for q + ] + } + + def test_custom_class_1(self): + d = GammaEx(concentration=2, rate=8) + actual = d.icdf(torch.tensor(0.5)).numpy() + expected = torch.tensor(1.8033469900166612).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = GammaEx(concentration=2, rate=8) + actual = d.icdf(torch.tensor([0.5, 0.9])) + expected = torch.tensor([1.8033469900166612, 4.014720169867429]).numpy() + assert_equal_for_tensor_and_list_or_scalar(self, actual, expected) + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdGamma().f( + alpha=case[0], + beta=case[1])[0] + icdf = PtdmIcdfTensor().f( + distribution, + torch.tensor(case[2]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(icdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, icdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_poisson_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_poisson_test.py new file mode 100644 index 0000000..e1f07e8 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_poisson_test.py @@ -0,0 +1,68 @@ +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_poisson import PoissonEx +from pytorch_wrapper.ptd_poisson import PtdPoisson +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfPoisson(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Poisson": [ + # rate, q, expected + (3, 0.4, scst.poisson.ppf(0.4, 3)), + (3, 0.7, scst.poisson.ppf(0.7, 3)), + (3, [0.4, 0.7], scst.poisson.ppf([0.4, 0.7], 3)) # 1d array for q + ] + } + + def test_custom_class_1(self): + d = PoissonEx(rate=3) + actual = d.icdf(torch.tensor(0.42319008112684364)).numpy() + expected = torch.tensor(2).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = PoissonEx(rate=3) + actual = d.icdf(torch.tensor(0.9160820579686966)).numpy() + expected = torch.tensor(5.0).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_3(self): + d = PoissonEx(rate=3) + actual = d.icdf(torch.tensor([0.42319008112684364, 0.9160820579686966])).numpy() + expected = torch.tensor([2.0, 5.0]).numpy() + self.assertTrue(np.allclose(actual, expected, atol=1e-5), + msg=f"Expected {expected}, got {actual}") + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdPoisson().f(str(case[0]))[0] + cdf = PtdmIcdfTensor().f( + distribution, + torch.tensor(case[1]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(cdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_student_t_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_student_t_test.py new file mode 100644 index 0000000..b10d1f6 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_student_t_test.py @@ -0,0 +1,71 @@ +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_student_t import StudentTEx +from pytorch_wrapper.ptd_student_t import PtdStudentT +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfStudentT(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "StudentT": [ + # str df, str loc, str scale, num q, expected + ("40.0", "0", "1", 0.5, scst.t.ppf(0.5, 40)), + ("40.0", "170", "20", 0.9, scst.t.ppf(0.9, 40) * 20 + 170), # adjust loc and scale + ("40.0", "170", "20", [0.8, 0.9], scst.t.ppf([0.8, 0.9], 40) * 20 + 170) # 1d array for q + ] + } + + def test_custom_class_1(self): + d = StudentTEx(df=40, loc=0, scale=1) + actual = d.icdf(torch.tensor(0.690092632383276)).numpy() + expected = torch.tensor(0.5).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_2(self): + d = StudentTEx(df=40, loc=0, scale=1) + actual = d.icdf(torch.tensor(0.813246561322997)).numpy() + expected = torch.tensor(0.9).numpy() + self.assertAlmostEqual(actual, expected, places=4, + msg=f"Expected {expected}, got {actual}") + + def test_custom_class_3(self): + d = StudentTEx(df=40, loc=0, scale=1) + actual = d.icdf(torch.tensor([0.690092632383276, 0.813246561322997])).numpy() + expected = torch.tensor([0.5, 0.9]).numpy() + self.assertTrue(np.allclose(actual, expected, atol=1e-5), + msg=f"Expected {expected}, got {actual}") + + def test_distributions(self): + """Test probability mass and density functions for the target distribution.""" + 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): + distribution = PtdStudentT().f( + df=case[0], + loc=case[1], + scale=case[2])[0] + cdf = PtdmIcdfTensor().f( + distribution, + torch.tensor(case[3]))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(cdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_tensor_test.py b/tests/pytorch_wrapper/ptdm_icdf_tensor_test.py new file mode 100644 index 0000000..4f6b652 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_tensor_test.py @@ -0,0 +1,47 @@ +import os +import sys +import unittest +import numpy as np +import torch +import scipy.stats as scst +from torch.distributions import Exponential + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from pytorch_wrapper.ptdm_icdf_tensor import PtdmIcdfTensor +from .utils import assert_equal_for_tensor_and_list_or_scalar + +class TestPtdfIcdfTensor(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.test_cases = { + "Exponential": [ + (1.0, 0.9, scst.expon.ppf(0.9, scale=1)), + (1.0, [0.8, 0.9], scst.expon.ppf(np.array([0.8, 0.9], dtype=np.float32), scale=1)) + ], + } + + def test_distributions(self): + """Test CDF.""" + for dist_name, cases in self.test_cases.items(): # For each distribution + for case in cases: # For each distribution's test cases + with self.subTest(distribution=dist_name, params=case): + if dist_name == "Exponential": + distribution = Exponential(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + else: + raise ValueError("Only exponential is supported in this unit test.") + + cdf = PtdmIcdfTensor().f(distribution, k_tensor)[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(cdf, torch.Tensor), "Output is not a tensor") + + # Check values are equal + assert_equal_for_tensor_and_list_or_scalar(self, cdf, case[-1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptdm_icdf_test.py b/tests/pytorch_wrapper/ptdm_icdf_test.py new file mode 100644 index 0000000..e92d822 --- /dev/null +++ b/tests/pytorch_wrapper/ptdm_icdf_test.py @@ -0,0 +1,100 @@ +import os +import sys +import unittest +import torch +import scipy.stats as scst +import numpy as np + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) +from pytorch_wrapper.ptdm_icdf import PtdmIcdf +from torch.distributions import ( + Categorical, NegativeBinomial, Multinomial, + Gamma, Normal, Laplace, Exponential, LogNormal, + Pareto, Weibull, Dirichlet, VonMises, Uniform +) + + +class TestPtdfCdf(unittest.TestCase): + def setUp(self): + """Set up test instance. + + Following distributions are tested in separate unit test files: + * Bernoulli + * Binomial + * Poisson + * Chi2 + * Gamma + * StudentT + * Beta + """ + self.test_cases = { + "Uniform": [(0.0, 1.0, 0.5, scst.uniform.ppf(0.5, loc=0, scale=1))], + "Normal": [(0.0, 1.0, 0.95, scst.norm.ppf(0.95, 0, 1))], + "Exponential": [(1.0, 0.9, scst.expon.ppf(0.9, 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): + # Instantiate PyTorch distribution directly if not in wrapper + if dist_name == "Categorical": + distribution = Categorical(case[0]) + k_tensor = torch.tensor(case[1], dtype=torch.int64) + elif dist_name == "NegativeBinomial": + distribution = NegativeBinomial(torch.tensor(case[0]), torch.tensor(case[2])) + k_tensor = torch.tensor(case[1], dtype=torch.int64) + elif dist_name == "Gamma": + distribution = Gamma(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + elif dist_name == "Normal": + distribution = Normal(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + elif dist_name == "Laplace": + distribution = Laplace(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + elif dist_name == "Exponential": + distribution = Exponential(torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + elif dist_name == "LogNormal": + distribution = LogNormal(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + elif dist_name == "Pareto": + distribution = Pareto(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + elif dist_name == "Weibull": + distribution = Weibull(torch.tensor(1.0), torch.tensor(case[0])) + k_tensor = torch.tensor(case[1], dtype=torch.float32) + elif dist_name == "Dirichlet": + distribution = Dirichlet(case[0]) + k_tensor = case[1] # Dirichlet expects tensor input + elif dist_name == "VonMises": + distribution = VonMises(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + elif dist_name == "Uniform": + distribution = Uniform(torch.tensor(case[0]), torch.tensor(case[1])) + k_tensor = torch.tensor(case[2], dtype=torch.float32) + else: + self.fail(f"Unhandled distribution: {dist_name}") + + # Call wrapper log-prob method + if dist_name == "Dirichlet": + icdf = PtdmIcdf().f(distribution, str(k_tensor.tolist()))[0] + else: + icdf = PtdmIcdf().f(distribution, str(k_tensor.item()))[0] + + # Ensure return values are tensors + self.assertTrue(isinstance(icdf, torch.Tensor), "Probability output is not a tensor") + + # Convert tensors to floats for comparison + icdf = icdf.item() + + # Check if computed values are approximately equal to expected values + self.assertAlmostEqual(icdf, case[-1], places=4, + msg=f"{dist_name}: Expected {case[-1]}, got {icdf}") + +if __name__ == "__main__": + unittest.main()