bayinx 0.3.12__tar.gz → 0.3.13__tar.gz

{bayinx-0.3.12 → bayinx-0.3.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.3.12
+Version: 0.3.13
 Summary: Bayesian Inference with JAX
 License-File: LICENSE
 Requires-Python: >=3.12

{bayinx-0.3.12 → bayinx-0.3.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.3.12"
+version = "0.3.13"
 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.3.12"
+current_version = "0.3.13"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{bayinx-0.3.12 → bayinx-0.3.13}/src/bayinx/dists/censored/posnormal/r.py

@@ -1,8 +1,9 @@
 import jax.numpy as jnp
 import jax.random as jr
+from jax.scipy.special import ndtri
 from jaxtyping import Array, ArrayLike, Float, Key
 
-from bayinx.dists import posnormal
+from bayinx.dists import normal, posnormal
 
 
 def prob(
@@ -108,10 +109,8 @@ def sample(
     # Derive shape
     shape = (n,) + jnp.broadcast_shapes(mu.shape, sigma.shape, censor.shape)
 
-    # Draw from positive normal
-    draws = jr.truncated_normal(key, 0.0, jnp.inf, shape) * sigma + mu
-
-    # Censor values
-    draws = jnp.where(censor <= draws, censor, draws)
+    # Construct draws
+    draws = jr.uniform(key, shape)
+    draws = mu + sigma * ndtri(normal.cdf(-mu/sigma, 0.0, 1.0) + draws * normal.cdf(mu/sigma, 0.0, 1.0))
 
     return draws

{bayinx-0.3.12 → bayinx-0.3.13}/src/bayinx/dists/uniform.py

@@ -1,6 +1,7 @@
 import jax.lax as _lax
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Float
+import jax.random as jr
+from jaxtyping import Array, ArrayLike, Float, Key
 
 
 def prob(
@@ -17,8 +18,10 @@ def prob(
     # Returns
     The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
+    # Cast to Array
+    x, lb, ub = jnp.asarray(x), jnp.asarray(lb), jnp.asarray(ub)
 
-    return 1.0 / (ub - lb)  # pyright: ignore
+    return 1.0 / (ub - lb)
 
 
 def logprob(
@@ -35,8 +38,10 @@ def logprob(
     # Returns
     The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
+    # Cast to Array
+    x, lb, ub = jnp.asarray(x), jnp.asarray(lb), jnp.asarray(ub)
 
-    return _lax.log(1.0) - _lax.log(ub - lb)  # pyright: ignore
+    return _lax.log(1.0) - _lax.log(ub - lb)
 
 
 def uprob(
@@ -53,8 +58,10 @@ def uprob(
     # Returns
     The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
+    # Cast to Array
+    x, lb, ub = jnp.asarray(x), jnp.asarray(lb), jnp.asarray(ub)
 
-    return jnp.ones(jnp.broadcast_arrays(x, lb, ub))
+    return jnp.ones(jnp.broadcast_shapes(x.shape, lb.shape, ub.shape))
 
 
 def ulogprob(
@@ -71,5 +78,32 @@ def ulogprob(
     # Returns
     The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `lb`, and `ub`.
     """
+    # Cast to Array
+    x, lb, ub = jnp.asarray(x), jnp.asarray(lb), jnp.asarray(ub)
 
-    return jnp.zeros(jnp.broadcast_arrays(x, lb, ub))
+    return jnp.zeros(jnp.broadcast_shapes(x.shape, lb.shape, ub.shape))
+
+def sample(
+    n: int, lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."], key: Key = jr.PRNGKey(0),
+) -> Float[Array, "..."]:
+    """
+    Sample from a Uniform distribution.
+
+    # Parameters
+    - `n`: Number of draws to sample per-parameter.
+    - `lb`: The lower bound parameter(s).
+    - `ub`: The upper bound parameter(s).
+
+    # Returns
+    Draws from a Uniform distribution. The output will have the shape of (n,) + the broadcasted shapes of `lb` and `ub`.
+    """
+    # Cast to Array
+    lb, ub = jnp.asarray(lb), jnp.asarray(ub)
+
+    # Derive shape
+    shape = (n,) + jnp.broadcast_shapes(lb.shape, ub.shape)
+
+    # Construct draws
+    draws = jr.uniform(key, shape, minval = lb, maxval = ub)
+
+    return draws