bayinx 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

bayinx/dists/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.dists import normal, posnormal
+
+__all__ = ['normal', 'posnormal']

bayinx/dists/censored/__init__.py

@@ -0,0 +1,3 @@
+from . import gamma2, posnormal
+
+__all__ = ['gamma2', 'posnormal']

bayinx/dists/censored/gamma2/__init__.py

@@ -0,0 +1,3 @@
+from . import r
+
+__all__ = ['r']

bayinx/dists/censored/gamma2/r.py

@@ -19,9 +19,10 @@ def prob(
     - `x`:  Value(s) at which to evaluate the PMF/PDF.
     - `mu`: The positive mean.
     - `nu`: The positive inverse dispersion.
+    - `censor`: The positive censor value.
 
     # Returns
-    The PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    The PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `nu`, and `censor`.
     """
     evals: Array = jnp.zeros_like(x * 1.0) # ensure float dtype
 
@@ -29,7 +30,7 @@ def prob(
     uncensored: Array = jnp.array(jnp.logical_and(0.0 < x, x < censor)) # pyright: ignore
     censored: Array = jnp.array(x == censor) # pyright: ignore
 
-    # Evaluate mixed probability (?) function
+    # Evaluate probability mass/density function
     evals = jnp.where(uncensored, gamma2.prob(x, mu, nu), evals)
     evals = jnp.where(censored, gammaincc(nu, x * nu / mu), evals) # pyright: ignore
 
@@ -49,9 +50,10 @@ def logprob(
     - `x`:      Value(s) at which to evaluate the log PMF/PDF.
     - `mu`:     The positive mean/location.
     - `nu`:     The positive inverse dispersion.
+    - `censor`: The positive censor value.
 
     # Returns
-    The log PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    The log PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `nu`, and `censor`.
     """
     evals: Array = jnp.full_like(x * 1.0, -jnp.inf) # ensure float dtype
 
@@ -59,6 +61,7 @@ def logprob(
     uncensored: Array = jnp.array(jnp.logical_and(0.0 < x, x < censor)) # pyright: ignore
     censored: Array = jnp.array(x == censor) # pyright: ignore
 
+    # Evaluate log probability mass/density function
     evals = jnp.where(uncensored, gamma2.logprob(x, mu, nu), evals)
     evals = jnp.where(censored, lax.log(gammaincc(nu, x * nu / mu)), evals) # pyright: ignore
 

bayinx/dists/censored/posnormal/__init__.py

@@ -0,0 +1,3 @@
+from . import r
+
+__all__ = ['r']

bayinx/dists/censored/posnormal/r.py

@@ -0,0 +1,68 @@
+import jax.numpy as jnp
+from jaxtyping import Array, ArrayLike, Float
+
+from bayinx.dists import posnormal
+
+
+def prob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+    censor: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The mixed probability mass/density function (PMF/PDF) for a censored positive Normal distribution.
+
+    # Parameters
+    - `x`:  Value(s) at which to evaluate the PMF/PDF.
+    - `mu`: The mean.
+    - `sigma`: The positive standard deviation.
+    - `censor`: The positive censor value.
+
+    # Returns
+    The PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `sigma`, and `censor`.
+    """
+    # Cast to Array
+    x, mu, sigma, censor = jnp.array(x), jnp.array(mu), jnp.array(sigma), jnp.array(censor)
+
+    # Construct boolean masks
+    uncensored: Array = jnp.logical_and(0.0 < x, x < censor)
+    censored: Array = x == censor
+
+    # Evaluate probability mass/density function
+    evals = jnp.where(uncensored, posnormal.prob(x, mu, sigma), 0.0)
+    evals = jnp.where(censored, posnormal.ccdf(x,mu,sigma), evals)
+
+    return evals
+
+
+def logprob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+    censor: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log-transformed mixed probability mass/density function (log PMF/PDF) for a censored positive Normal distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log PMF/PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+    - `censor`: The censor.
+
+    # Returns
+    The log PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, `sigma`, and `censor`.
+    """
+    # Cast to Array
+    x, mu, sigma, censor = jnp.array(x), jnp.array(mu), jnp.array(sigma), jnp.array(censor)
+
+    # Construct boolean masks
+    uncensored: Array = jnp.logical_and(jnp.array(0.0) < x, x < censor)
+    censored: Array = x == censor
+
+    # Evaluate log probability mass/density function
+    evals = jnp.where(uncensored, posnormal.logprob(x, mu, sigma), -jnp.inf)
+    evals = jnp.where(censored, posnormal.logccdf(x, mu, sigma), evals)
+
+    return evals

bayinx/dists/normal.py

@@ -1,4 +1,6 @@
 import jax.lax as lax
+import jax.numpy as jnp
+import jax.scipy.special as jss
 from jaxtyping import Array, ArrayLike, Float
 
 __PI = 3.141592653589793
@@ -18,8 +20,10 @@ def prob(
     # Returns
     The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
 
-    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / (  # pyright: ignore
+    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / (
         sigma * lax.sqrt(2.0 * __PI)
     )
 
@@ -38,9 +42,11 @@ def logprob(
     # Returns
     The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
 
     return -lax.log(sigma * lax.sqrt(2.0 * __PI)) - 0.5 * lax.square(
-        (x - mu) / sigma  # pyright: ignore
+        (x - mu) / sigma
     )
 
 
@@ -58,8 +64,10 @@ def uprob(
     # Returns
     The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
 
-    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / sigma  # pyright: ignore
+    return lax.exp(-0.5 * lax.square((x - mu) / sigma)) / sigma
 
 
 def ulogprob(
@@ -76,5 +84,39 @@ def ulogprob(
     # Returns
     The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
 
-    return -lax.log(sigma) - 0.5 * lax.square((x - mu) / sigma)  # pyright: ignore
+    return -lax.log(sigma) - 0.5 * lax.square((x - mu) / sigma)
+
+def cdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    return jss.ndtr((x - mu) / sigma)
+
+def logcdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    return jss.log_ndtr((x - mu) / sigma)
+
+def ccdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    return jss.ndtr((mu - x) / sigma)
+
+def logccdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    return jss.log_ndtr((mu - x) / sigma)

bayinx/dists/posnormal.py

@@ -0,0 +1,258 @@
+import jax.numpy as jnp
+from jaxtyping import Array, ArrayLike, Float
+
+from bayinx.dists import normal
+
+
+def prob(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The probability density function (PDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.array(0.0) <= x
+
+    # Evaluate PDF
+    evals = jnp.where(
+        non_negative,
+        normal.prob(x, mu, sigma) / normal.cdf(mu/sigma, 0.0, 1.0),
+        jnp.array(0.0))
+
+    return evals
+
+
+def logprob(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log of the probability density function (log PDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the log PDF.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    The log PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.array(0.0) <= x
+
+    # Evaluate log PDF
+    evals = jnp.where(
+        non_negative,
+        normal.logprob(x, mu, sigma) - normal.logcdf(mu/sigma, 0.0, 1.0),
+        -jnp.inf)
+
+    return evals
+
+
+def uprob(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The unnormalized probability density function (uPDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the uPDF.
+    - `mu`: The mean/location parameter(s).
+    - `sigma`: The positive standard deviation parameter(s).
+
+    # Returns
+    The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.array(0.0) <= x
+
+    # Evaluate PDF
+    evals = jnp.where(
+        non_negative,
+        normal.prob(x, mu, sigma),
+        jnp.array(0.0))
+
+    return evals
+
+
+def ulogprob(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log of the unnormalized probability density function (log uPDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the log uPDF.
+    - `mu`: The mean/location parameter(s).
+    - `sigma`: The non-negative standard deviation parameter(s).
+
+    # Returns
+    The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.array(0.0) <= x
+
+    # Evaluate log PDF
+    evals = jnp.where(
+        non_negative,
+        normal.logprob(x, mu, sigma),
+        -jnp.inf)
+
+    return evals
+
+
+def cdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The cumulative density function (CDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the log uPDF.
+    - `mu`: The mean/location parameter(s).
+    - `sigma`: The non-negative standard deviation parameter(s).
+
+    # Returns
+    The CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+
+    # Notes
+    Not numerically stable for small `x`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.array(0.0) <= x
+
+    # Compute intermediates
+    A: Array = normal.cdf(x, mu, sigma)
+    B: Array = normal.cdf(- mu / sigma, 0.0, 1.0)
+    C: Array = normal.cdf(mu / sigma, 0.0, 1.0)
+
+    # Evaluate CDF
+    evals = jnp.where(
+        non_negative,
+        (A - B) / C,
+        jnp.array(0.0))
+
+    return evals
+
+# TODO: make numerically stable
+def logcdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log-transformed cumulative density function (log CDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the log uPDF.
+    - `mu`: The mean/location parameter(s).
+    - `sigma`: The non-negative standard deviation parameter(s).
+
+    # Returns
+    The log CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+
+    # Notes
+    Not numerically stable for small `x`.
+    """
+    # Cast to Array
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = jnp.array(0.0) <= x
+
+    A: Array = normal.logcdf(x, mu, sigma)
+    B: Array = normal.logcdf(- mu/sigma, 0.0, 1.0)
+    C: Array = normal.logcdf(mu/sigma, 0.0, 1.0)
+
+    # Evaluate log CDF
+    evals = jnp.where(
+        non_negative,
+        A + jnp.log1p(-jnp.exp(B - A)) - C,
+        -jnp.inf)
+
+    return evals
+
+def ccdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The complementary cumulative density function (cCDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the log uPDF.
+    - `mu`: The mean/location parameter(s).
+    - `sigma`: The non-negative standard deviation parameter(s).
+
+    # Returns
+    The cCDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+
+    # Notes
+    Not numerically stable for small `x`.
+    """
+    # Cast to arrays
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = 0.0 <= x
+
+    # Compute intermediates
+    A: Array = normal.cdf(-x, -mu, sigma)
+    B: Array = normal.cdf(mu/sigma, 0.0, 1.0)
+
+    # Evaluate cCDF
+    evals = jnp.where(non_negative, A / B, jnp.array(1.0))
+
+    return evals
+
+
+def logccdf(
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log-transformed complementary cumulative density function (log cCDF) for a positive Normal distribution.
+
+    # Parameters
+    - `x`: Value(s) at which to evaluate the log uPDF.
+    - `mu`: The mean/location parameter(s).
+    - `sigma`: The non-negative standard deviation parameter(s).
+
+    # Returns
+    The log cCDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+
+    # Notes
+    Not numerically stable for small `x`.
+    """
+    # Cast to arrays
+    x, mu, sigma = jnp.array(x), jnp.array(mu), jnp.array(sigma)
+
+    # Construct boolean mask for non-negative elements
+    non_negative: Array = 0.0 <= x
+
+    # Compute intermediates
+    A: Array = normal.logcdf(-x, -mu, sigma)
+    B: Array = normal.logcdf(mu/sigma, 0.0, 1.0)
+
+    # Evaluate log cCDF
+    evals = jnp.where(non_negative, A - B, jnp.array(0.0))
+
+    return evals

{bayinx-0.3.3.dist-info → bayinx-0.3.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.3.3
+Version: 0.3.5
 Summary: Bayesian Inference with JAX
 Requires-Python: >=3.12
 Requires-Dist: equinox>=0.11.12

{bayinx-0.3.3.dist-info → bayinx-0.3.5.dist-info}/RECORD

@@ -8,13 +8,17 @@ bayinx/core/flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
 bayinx/core/model.py,sha256=1vQPVjE0ebCdW7mLuabgQcCTi95o8n8CC6GuzJdNL1s,2956
 bayinx/core/parameter.py,sha256=eECqvfMNWSU8_CkGYaAfOCneMMQGZI21kF0mErsh2Rc,1080
 bayinx/core/variational.py,sha256=lqENISRrKY8ODLtl0D-D7TAA2gD7HGh37BnROM7p5hI,4783
-bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/__init__.py,sha256=qPQrl5vkS9K56GzIaHZXkSUP07YAu4lVB8K2yQ1m3SY,78
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/gamma2.py,sha256=8XYaOtcYJCrr5q1yHWfZaMJmASpLOrfyhrH_J06ksj8,1333
-bayinx/dists/normal.py,sha256=mvm6EoAlORy-yivuhMcExYCZUo0vJzMKMOWH-9iQBZU,2634
+bayinx/dists/normal.py,sha256=BLlp7hGMAxUbroROvzA5ChH5YLXgadeK4VOuBtjjdjs,3978
+bayinx/dists/posnormal.py,sha256=NNr5OHv1fWCxYvc6hwUMIGXX31UAg0sEnc4tsxHLjUg,7726
 bayinx/dists/uniform.py,sha256=7XgVvOrzINEFA6HJTYUOFwlWhEtrQQQ1aPJ_ZLOzLEc,2365
-bayinx/dists/censored/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/dists/censored/gamma2/r.py,sha256=3brRCKhE-74mRXyIyPcnyaWY2OJv8CZyUWPP9T1t09Y,2274
+bayinx/dists/censored/__init__.py,sha256=p8T03TenD-_8YNiOgB_RKksq8hFNFejA5bnoK4JJ8Ms,67
+bayinx/dists/censored/gamma2/__init__.py,sha256=qqm0n2hfid617PvyFRHAOMAp3AvpOlt5v3ns8HgTD7E,33
+bayinx/dists/censored/gamma2/r.py,sha256=dE0MNTAl0E6npQhFONv341U7XbomBB-fNzQhgRjxYpk,2436
+bayinx/dists/censored/posnormal/__init__.py,sha256=qqm0n2hfid617PvyFRHAOMAp3AvpOlt5v3ns8HgTD7E,33
+bayinx/dists/censored/posnormal/r.py,sha256=4MfFkQ2klzOZJNjxS9g4zz1bdoJ6ehBxZQi6QkmPGgE,2232
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/vi/__init__.py,sha256=YfkXKsqo9Dk_AmQGjZKm4vfG8eLer2ez92G-cOExphs,193
 bayinx/mhx/vi/meanfield.py,sha256=M4QrOuHaIMLTuQSD6JNF9vELnTm370tXV68JPB7B67M,3652
@@ -25,6 +29,6 @@ bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvAD
 bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
 bayinx/mhx/vi/flows/radial.py,sha256=e0GfuO-CL8SVr3YnEfsxStpyKcJlFpzMyjMp3sa38hg,2503
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.3.3.dist-info/METADATA,sha256=5tZVPxDvYVajnoDrRIF-KskSMWBix7Zq4h_glyJa-_M,3057
-bayinx-0.3.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.3.dist-info/RECORD,,
+bayinx-0.3.5.dist-info/METADATA,sha256=Hj8GWJef3kfJ6umsHGIFWovYXXtPegAlcsopunoHFFs,3057
+bayinx-0.3.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.5.dist-info/RECORD,,