bayinx 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

bayinx/core/_variational.py

@@ -163,7 +163,7 @@ class Variational(eqx.Module, Generic[M]):
 
     @eqx.filter_jit
     def posterior_predictive(
-        self, func: Callable[[M], Array], n: int, data: Any = None, key: Key = jr.PRNGKey(0)
+        self, func: Callable[[M, Any], Array], n: int, data: Any = None, key: Key = jr.PRNGKey(0)
     ) -> Array:
         # Sample draws from the variational approximation
         draws: Array = self.sample(n, key)
@@ -176,6 +176,6 @@ class Variational(eqx.Module, Generic[M]):
             model: M = self._unflatten(draw)
 
             # Evaluate
-            return func(model)
+            return func(model, data)
 
-        return evaluate(draws)
+        return evaluate(draws, data)

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.8.dist-info → bayinx-0.3.10.dist-info}/METADATA

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

{bayinx-0.3.8.dist-info → bayinx-0.3.10.dist-info}/RECORD

@@ -7,7 +7,7 @@ bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,76
 bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
 bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
 bayinx/core/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
-bayinx/core/_variational.py,sha256=omZ3jFrYNdjfDVAOmNhjPtOrJcBD6yFlypvFROn6GgU,5467
+bayinx/core/_variational.py,sha256=X8o81b8vyU7vJxw8pZYH_nxc3u990tRUZgRhMNodNI4,5484
 bayinx/dists/__init__.py,sha256=9DdPea7HAnBOzaV_4gM5noPX8YCb_p06d8PJvGfFy3Y,118
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/gamma2.py,sha256=MuFudL2UTfk8HgWVofNaR36JTmUpmtxvg1Mifu98MvM,1567
@@ -18,7 +18,7 @@ bayinx/dists/censored/__init__.py,sha256=UVihMbQgAzCoOk_Zt5wrumPv5-acuTzV3TYMB-U
 bayinx/dists/censored/gamma2/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
 bayinx/dists/censored/gamma2/r.py,sha256=dKAOYstufwgDwibQZHrJxA1d2gawj-7K3IkaCRCzNTg,2446
 bayinx/dists/censored/posnormal/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
-bayinx/dists/censored/posnormal/r.py,sha256=hyuNR3HZY-Tgtso-WwjcZT6Ejxfyax_VKwIvVix44Jc,2362
+bayinx/dists/censored/posnormal/r.py,sha256=Ypi6w_t53pAzRVzjcStx2RhozkAlCDLnQmgKykhpQQ4,3426
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/vi/__init__.py,sha256=2woNB5oZxfs8pZCkOfzriGahRFLzkLdkTj8_keTN0I0,205
 bayinx/mhx/vi/meanfield.py,sha256=Z7kGQAyp5iB8rEdjbwAbVTFH4GwxlTKDZFbdJ-FN5Vs,3739
@@ -29,7 +29,7 @@ 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.8.dist-info/METADATA,sha256=Z6oJm_2TgfJyhle7aG-DeDiW3LqE6Mcwt8iNku5ai54,3079
-bayinx-0.3.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.8.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.8.dist-info/RECORD,,
+bayinx-0.3.10.dist-info/METADATA,sha256=y99n8rNP62ezyDKJSmBQvUQx52WKUEgyHtw3QnnVrvs,3080
+bayinx-0.3.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.10.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.10.dist-info/RECORD,,