bayinx 0.2.21__tar.gz → 0.2.22__tar.gz

{bayinx-0.2.21 → bayinx-0.2.22}/PKG-INFO

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

{bayinx-0.2.21 → bayinx-0.2.22}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.21"
+version = "0.2.22"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -19,7 +19,7 @@ build-backend = "hatchling.build"
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
 
 [tool.bumpversion]
-current_version = "0.2.21"
+current_version = "0.2.22"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{bayinx-0.2.21 → bayinx-0.2.22}/src/bayinx/dists/normal.py

@@ -1,7 +1,7 @@
 import jax.lax as _lax
 from jaxtyping import Array, ArrayLike, Float, Real
 
-_PI = 3.141592653589793
+__PI = 3.141592653589793
 
 
 def prob(
@@ -20,7 +20,7 @@ def prob(
     """
 
     return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / ( # pyright: ignore
-        sigma * _lax.sqrt(2.0 * _PI)
+        sigma * _lax.sqrt(2.0 * __PI)
     )
 
 
@@ -39,7 +39,7 @@ def logprob(
     The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
     """
 
-    return -_lax.log(sigma * _lax.sqrt(2.0 * _PI)) - 0.5 * _lax.square((x - mu) / sigma) # pyright: ignore
+    return -_lax.log(sigma * _lax.sqrt(2.0 * __PI)) - 0.5 * _lax.square((x - mu) / sigma) # pyright: ignore
 
 
 def uprob(

bayinx-0.2.22/src/bayinx/dists/uniform.py

@@ -0,0 +1,75 @@
+import jax.lax as _lax
+import jax.numpy as jnp
+from jaxtyping import Array, ArrayLike, Float, Real
+
+
+def prob(
+    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The probability density function (PDF) for a Uniform distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `lb`:     The lower bound parameter(s).
+    - `ub`:     The upper bound parameter(s).
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
+    """
+
+    return 1.0 / (ub - lb) # pyright: ignore
+
+
+def logprob(
+    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log of the probability density function (log PDF) for a Uniform distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `lb`:     The lower bound parameter(s).
+    - `ub`:     The upper bound parameter(s).
+
+    # Returns
+    The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
+    """
+
+    return _lax.log(1.0) - _lax.log(ub - lb) # pyright: ignore
+
+
+def uprob(
+    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The unnormalized probability density function (uPDF) for a Uniform distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `lb`:     The lower bound parameter(s).
+    - `ub`:     The upper bound parameter(s).
+
+    # Returns
+    The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
+    """
+
+    return jnp.ones(jnp.broadcast_arrays(x,lb,ub))
+
+
+def ulogprob(
+    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log of the unnormalized probability density function (log uPDF) for a Uniform distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `lb`:     The lower bound parameter(s).
+    - `ub`:     The upper bound parameter(s).
+
+    # Returns
+    The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
+    """
+
+    return jnp.zeros(jnp.broadcast_arrays(x,lb,ub))