Add Binomial, Bernoulli, Poisson and NegativeBinomial

distributions/__init__.py

@@ -2,5 +2,4 @@
 PyTensor powered distributions.
 """
 
-
-__version__ = "0.1.0"
+__version__ = "0.1.0"

distributions/bernoulli.py

@@ -0,0 +1,80 @@
+import pytensor.tensor as pt
+from pytensor.tensor.xlogx import xlogx
+
+from .helper import ppf_bounds_disc, cdf_bounds
+
+
+def mean(p):
+    return p
+
+
+def mode(p):
+    return median(p)
+
+
+def median(p):
+    return pt.switch(pt.le(p, 0.5), 0, 1)
+
+
+def var(p):
+    return p * (1 - p)
+
+
+def std(p):
+    return pt.sqrt(var(p))
+
+
+def skewness(p):
+    q = 1 - p
+    return (q - p) / pt.sqrt(p * q)
+
+
+def kurtosis(p):
+    return (1 - 6 * p * (1 - p)) / (p * (1 - p))
+
+
+def entropy(p):
+    q = 1 - p
+    return -xlogx(p) - xlogx(q)
+
+
+def rvs(p, size=None, random_state=None):
+    return pt.random.binomial(1, p, size=size, rng=random_state)
+
+
+def cdf(x, p):
+    x = pt.as_tensor_variable(x)
+    return cdf_bounds(pt.switch(pt.lt(x, 1), 1 - p, 1.0), x, 0, 1)
+
+
+def ppf(q, p):
+    q = pt.as_tensor_variable(q)
+    x_val = pt.switch(pt.lt(q, 1 - p), 0, 1)
+    return ppf_bounds_disc(x_val, q, 0, 1)
+
+
+def sf(x, p):
+    return pt.switch(pt.lt(x, 0), 1.0, pt.switch(pt.lt(x, 1), p, 0.0))
+
+
+def isf(q, p):
+    return ppf(1 - q, p)
+
+
+def pdf(x, p):
+    x = pt.as_tensor_variable(x)
+    return pt.switch(pt.eq(x, 1), p, pt.switch(pt.eq(x, 0), 1 - p, 0.0))
+
+
+def logpdf(x, p):
+    x = pt.as_tensor_variable(x)
+    return pt.switch(pt.eq(x, 1), pt.log(p), pt.switch(pt.eq(x, 0), pt.log(1 - p), -pt.inf))
+
+
+def logcdf(x, p):
+    x = pt.as_tensor_variable(x)
+    return pt.switch(pt.lt(x, 0), -pt.inf, pt.switch(pt.lt(x, 1), pt.log(1 - p), 0.0))
+
+
+def logsf(x, mu):
+    return pt.log1p(-cdf(x, mu))

distributions/beta.py

@@ -8,38 +8,38 @@
 def mean(alpha, beta):
     return alpha / (alpha + beta)
 
+
 def mode(alpha, beta):
     alpha_b, beta_b = pt.broadcast_arrays(alpha, beta)
     result = pt.full_like(alpha_b, pt.nan)
 
     result = pt.where(pt.equal(alpha_b, 1) & (beta_b > 1), 0.0, result)
     result = pt.where(pt.equal(beta_b, 1) & (alpha_b > 1), 1.0, result)
-    result = pt.where((alpha_b > 1) & (beta_b > 1),
-                      (alpha_b - 1) / (alpha_b + beta_b - 2),
-                      result)
-
+    result = pt.where((alpha_b > 1) & (beta_b > 1), (alpha_b - 1) / (alpha_b + beta_b - 2), result)
+
     return result
 
+
 def median(alpha, beta):
     return ppf(0.5, alpha, beta)
 
+
 def var(alpha, beta):
-    return (alpha * beta) / (
-        pt.pow(alpha + beta, 2) * (alpha + beta + 1)
-    )
+    return (alpha * beta) / (pt.pow(alpha + beta, 2) * (alpha + beta + 1))
+
 
 def std(alpha, beta):
     return pt.sqrt(var(alpha, beta))
 
+
 def skewness(alpha, beta):
     alpha_b, beta_b = pt.broadcast_arrays(alpha, beta)
-    
+
     psc = alpha_b + beta_b
     result = pt.where(
-        pt.eq(alpha_b, beta_b), 0.0,
-        (2 * (beta_b - alpha_b) * pt.sqrt(psc + 1)) / (
-            (psc + 2) * pt.sqrt(alpha_b * beta_b)
-        )
+        pt.eq(alpha_b, beta_b),
+        0.0,
+        (2 * (beta_b - alpha_b) * pt.sqrt(psc + 1)) / ((psc + 2) * pt.sqrt(alpha_b * beta_b)),
     )
     return result
 
@@ -48,10 +48,14 @@ def kurtosis(alpha, beta):
     alpha_b, beta_b = pt.broadcast_arrays(alpha, beta)
     psc = alpha_b + beta_b
     prod = alpha_b * beta_b
-    result = (6 * (pt.abs(alpha_b - beta_b) ** 2 * (psc + 1) - prod * (psc + 2))
-    / (prod * (psc + 2) * (psc + 3)))
+    result = (
+        6
+        * (pt.abs(alpha_b - beta_b) ** 2 * (psc + 1) - prod * (psc + 2))
+        / (prod * (psc + 2) * (psc + 3))
+    )
     return result
 
+
 def entropy(alpha, beta):
     alpha_b, beta_b = pt.broadcast_arrays(alpha, beta)
     psc = alpha_b + beta_b
@@ -62,23 +66,30 @@ def entropy(alpha, beta):
         + (psc - 2) * pt.psi(psc)
     )
 
+
 def cdf(x, alpha, beta):
     return pt.exp(logcdf(x, alpha, beta))
 
+
 def isf(x, alpha, beta):
     return ppf(1 - x, alpha, beta)
 
+
 def pdf(x, alpha, beta):
     return pt.exp(logpdf(x, alpha, beta))
 
+
 def ppf(q, alpha, beta):
     return ppf_bounds_cont(betaincinv(alpha, beta, q), q, 0.0, 1.0)
 
+
 def sf(x, alpha, beta):
     return pt.exp(logsf(x, alpha, beta))
 
+
 def rvs(alpha, beta, size=None, random_state=None):
-     return pt.random.beta(alpha, beta, rng=random_state, size=size)
+    return pt.random.beta(alpha, beta, rng=random_state, size=size)
+
 
 def logcdf(x, alpha, beta):
     return pt.switch(
@@ -91,6 +102,7 @@ def logcdf(x, alpha, beta):
         ),
     )
 
+
 def logpdf(x, alpha, beta):
     result = (
         pt.switch(pt.eq(alpha, 1.0), 0.0, (alpha - 1.0) * pt.log(x))
@@ -103,31 +115,30 @@ def logpdf(x, alpha, beta):
 def logsf(x, alpha, beta):
     return pt.switch(
         pt.lt(x, 0),
-        0, 
+        0,
         pt.switch(
             pt.lt(x, 1),
-            pt.log(pt.betainc(beta, alpha, 1 - x)),  
+            pt.log(pt.betainc(beta, alpha, 1 - x)),
             -pt.inf,
         ),
     )
 
+
 def from_mu_sigma(mu, sigma):
     nu = mu * (1 - mu) / sigma**2 - 1
     alpha = mu * nu
     beta = (1 - mu) * nu
     return alpha, beta
 
+
 def from_mu_nu(mu, nu):
     alpha = mu * nu
     beta = (1 - mu) * nu
     return alpha, beta
 
+
 def to_mu_sigma(alpha, beta):
     alpha_plus_beta = alpha + beta
     mu = alpha / alpha_plus_beta
     sigma = (alpha * beta) ** 0.5 / alpha_plus_beta / (alpha_plus_beta + 1) ** 0.5
     return mu, sigma
-
-
-
-

distributions/betascaled.py

@@ -1,4 +1,3 @@
-
 import pytensor.tensor as pt
 from pytensor.tensor.special import betaln
 from pytensor.tensor.xlogx import xlogy0
@@ -15,12 +14,14 @@ def mode(alpha, beta, lower, upper):
     alpha_b, beta_b, lower_b, upper_b = pt.broadcast_arrays(alpha, beta, lower, upper)
     result = pt.full_like(alpha_b, pt.nan)
 
-    result = pt.where((alpha_b < 1) & (beta_b < 1), pt.nan, result)    
+    result = pt.where((alpha_b < 1) & (beta_b < 1), pt.nan, result)
     result = pt.where((alpha_b <= 1) & (beta_b > 1), lower_b, result)
     result = pt.where((beta_b <= 1) & (alpha_b > 1), upper_b, result)
-    result = pt.where((alpha_b > 1) & (beta_b > 1),
-                      lower_b + (alpha_b - 1) / (alpha_b + beta_b - 2) * (upper_b - lower_b),
-                      result)
+    result = pt.where(
+        (alpha_b > 1) & (beta_b > 1),
+        lower_b + (alpha_b - 1) / (alpha_b + beta_b - 2) * (upper_b - lower_b),
+        result,
+    )
     return result
 
 
@@ -29,11 +30,7 @@ def median(alpha, beta, lower, upper):
 
 
 def var(alpha, beta, lower, upper):
-    return (
-        (alpha * beta)
-        / ((alpha + beta) ** 2 * (alpha + beta + 1))
-        * (upper - lower) ** 2
-    )
+    return (alpha * beta) / ((alpha + beta) ** 2 * (alpha + beta + 1)) * (upper - lower) ** 2
 
 
 def std(alpha, beta, lower, upper):
@@ -42,13 +39,12 @@ def std(alpha, beta, lower, upper):
 
 def skewness(alpha, beta, lower, upper):
     alpha_b, beta_b = pt.broadcast_arrays(alpha, beta)
-    
+
     psc = alpha_b + beta_b
     result = pt.where(
-        pt.eq(alpha_b, beta_b), 0.0,
-        (2 * (beta_b - alpha_b) * pt.sqrt(psc + 1)) / (
-            (psc + 2) * pt.sqrt(alpha_b * beta_b)
-        )
+        pt.eq(alpha_b, beta_b),
+        0.0,
+        (2 * (beta_b - alpha_b) * pt.sqrt(psc + 1)) / ((psc + 2) * pt.sqrt(alpha_b * beta_b)),
     )
     return result
 
@@ -57,8 +53,11 @@ def kurtosis(alpha, beta, lower, upper):
     alpha_b, beta_b = pt.broadcast_arrays(alpha, beta)
     psc = alpha_b + beta_b
     prod = alpha_b * beta_b
-    result = (6 * (pt.abs(alpha_b - beta_b) ** 2 * (psc + 1) - prod * (psc + 2))
-    / (prod * (psc + 2) * (psc + 3)))
+    result = (
+        6
+        * (pt.abs(alpha_b - beta_b) ** 2 * (psc + 1) - prod * (psc + 2))
+        / (prod * (psc + 2) * (psc + 3))
+    )
     return result
 
 
@@ -105,21 +104,16 @@ def logcdf(x, alpha, beta, lower, upper):
     return pt.switch(
         pt.lt(x, lower),
         -pt.inf,
-        pt.switch(
-            pt.gt(x, upper),
-            0.0,
-            pt.log(pt.betainc(alpha, beta, x_normalized))
-        )
+        pt.switch(pt.gt(x, upper), 0.0, pt.log(pt.betainc(alpha, beta, x_normalized))),
     )
 
 
 def logpdf(x, alpha, beta, lower, upper):
     return pt.switch(
         pt.bitwise_or(pt.lt(x, lower), pt.gt(x, upper)),
         -pt.inf,
-        (xlogy0((alpha - 1), (x - lower)) + xlogy0((beta - 1), (upper - x))) - (
-             xlogy0((alpha + beta - 1), (upper - lower)) + betaln(alpha, beta)
-        )
+        (xlogy0((alpha - 1), (x - lower)) + xlogy0((beta - 1), (upper - x)))
+        - (xlogy0((alpha + beta - 1), (upper - lower)) + betaln(alpha, beta)),
     )
 
 
@@ -128,11 +122,7 @@ def logsf(x, alpha, beta, lower, upper):
     return pt.switch(
         pt.lt(x, lower),
         0.0,
-        pt.switch(
-            pt.gt(x, upper),
-            -pt.inf,
-            pt.log(pt.betainc(beta, alpha, 1 - x_normalized))
-        )
+        pt.switch(pt.gt(x, upper), -pt.inf, pt.log(pt.betainc(beta, alpha, 1 - x_normalized))),
     )
 
 
@@ -142,13 +132,15 @@ def from_mu_sigma(mu, sigma):
     beta = (1 - mu) * nu
     return alpha, beta
 
+
 def from_mu_nu(mu, nu):
     alpha = mu * nu
     beta = (1 - mu) * nu
     return alpha, beta
 
+
 def to_mu_sigma(alpha, beta):
     alpha_plus_beta = alpha + beta
     mu = alpha / alpha_plus_beta
     sigma = (alpha * beta) ** 0.5 / alpha_plus_beta / (alpha_plus_beta + 1) ** 0.5
-    return mu, sigma
+    return mu, sigma