bayinx 0.3.9__tar.gz → 0.3.11__tar.gz

{bayinx-0.3.9 → bayinx-0.3.11}/PKG-INFO

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

{bayinx-0.3.9 → bayinx-0.3.11}/pyproject.toml

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

{bayinx-0.3.9 → bayinx-0.3.11}/src/bayinx/core/_variational.py

@@ -171,11 +171,11 @@ class Variational(eqx.Module, Generic[M]):
         # Evaluate posterior predictive
         @jax.jit
         @jax.vmap
-        def evaluate(draw: Array, data: Any = None):
+        def evaluate(draw: Array):
             # Reconstruct model
             model: M = self._unflatten(draw)
 
             # Evaluate
             return func(model, data)
 
-        return evaluate(draws, data)
+        return evaluate(draws)

{bayinx-0.3.9 → bayinx-0.3.11}/src/bayinx/dists/censored/posnormal/r.py

@@ -1,5 +1,6 @@
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Float
+import jax.random as jr
+from jaxtyping import Array, ArrayLike, Float, Key
 
 from bayinx.dists import posnormal
 
@@ -76,3 +77,40 @@ def logprob(
     evals = jnp.where(censored, posnormal.logccdf(x, mu, sigma), evals)
 
     return evals
+
+def sample(
+    n: int,
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+    censor: Float[ArrayLike, "..."] = jnp.inf,
+    key: Key = jr.PRNGKey(0)
+) -> Float[Array, "..."]:
+    """
+    Sample from a right-censored positive Normal distribution.
+
+    # Parameters
+    - `n`: Number of draws to sample per-parameter.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+    - `censor`: The censor.
+
+    # Returns
+    Draws from a right-censored positive Normal distribution. The output will have the shape of (n,) + the broadcasted shapes of `mu`, `sigma`, and `censor`.
+    """
+    # Cast to Array
+    mu, sigma, censor = (
+        jnp.asarray(mu),
+        jnp.asarray(sigma),
+        jnp.asarray(censor),
+    )
+
+    # 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)
+
+    return draws

{bayinx-0.3.9 → bayinx-0.3.11}/tests/test_predictive.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Any, Dict
 
 import jax.numpy as jnp
 from jaxtyping import Array
@@ -32,7 +32,7 @@ def test_predictive():
     vari = Standard(model)
 
     # Posterior predictive function
-    def extract_first_x(model: NormalDist) -> Array:
+    def extract_first_x(model: NormalDist, data: Any = None) -> Array:
         return model.x()[0]
 
     # Sample from posterior