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

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

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

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

@@ -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.9.dist-info/METADATA,sha256=yn3AC0J8_xLnLYEEocDZyNRlyyuN2Ys8C--i9AFFYeY,3079
-bayinx-0.3.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.9.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.9.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,,