bayinx 0.3.12__py3-none-any.whl → 0.3.13__py3-none-any.whl

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/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

{bayinx-0.3.12.dist-info → bayinx-0.3.13.dist-info}/METADATA

@@ -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.dist-info → bayinx-0.3.13.dist-info}/RECORD

@@ -14,12 +14,12 @@ bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,100
 bayinx/dists/gamma2.py,sha256=MuFudL2UTfk8HgWVofNaR36JTmUpmtxvg1Mifu98MvM,1567
 bayinx/dists/normal.py,sha256=Yc2X8F7JoLYwprtK8bA2BPva1tAY7MEs3oSk5pMortI,3822
 bayinx/dists/posnormal.py,sha256=w9plA1EctXwXOiY0doc4ZndjnwptbEZBHHCGdc4gviY,7292
-bayinx/dists/uniform.py,sha256=7XgVvOrzINEFA6HJTYUOFwlWhEtrQQQ1aPJ_ZLOzLEc,2365
+bayinx/dists/uniform.py,sha256=2ZQxEfAX5TFgSPuQ8joFDuFbd_NfmQ1GvmGGjusqvNQ,3461
 bayinx/dists/censored/__init__.py,sha256=UVihMbQgAzCoOk_Zt5wrumPv5-acuTzV3TYMB-U1gOc,49
 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=r8QfEThaaqsqqrr6PHWiY0EdfnjQcbKdTAHN2GmDVzI,3428
+bayinx/dists/censored/posnormal/r.py,sha256=wMDt2Am1TD376ms8B-o6PFCJZXmUJd2-aBC-t9kidH4,3456
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/opt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/mhx/vi/__init__.py,sha256=3T1dEpiiRge4tW-vpS0xBob_RbO1iVFnL3fVCRUawCM,205
@@ -31,7 +31,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.12.dist-info/METADATA,sha256=RvwZfeFmarMW4SlQhurD5TLnAS0sge9lVNcHPUNjIIM,3080
-bayinx-0.3.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.12.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.12.dist-info/RECORD,,
+bayinx-0.3.13.dist-info/METADATA,sha256=1mcHMTXrzMGPcibMy_vHdBJMNPtT0VUF83eVBS_MJlg,3080
+bayinx-0.3.13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.13.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.13.dist-info/RECORD,,