bayinx 0.2.29__py3-none-any.whl → 0.2.30__py3-none-any.whl

bayinx/{core/constraints.py → constraints/lower.py}

@@ -1,30 +1,9 @@
-from abc import abstractmethod
 from typing import Tuple
 
-import equinox as eqx
 import jax.numpy as jnp
 from jaxtyping import Array, ArrayLike, Scalar, ScalarLike
 
-
-class Constraint(eqx.Module):
-    """
-    Abstract base class for defining parameter constraints.
-    """
-
-    @abstractmethod
-    def constrain(self, x: ArrayLike) -> Tuple[Array, Scalar]:
-        """
-        Applies the constraining transformation to an unconstrained input and computes the log-absolute-jacobian of the transformation.
-
-        # Parameters
-        - `x`: The unconstrained JAX Array-like input.
-
-        # Returns
-        A tuple containing:
-            - The constrained JAX Array.
-            - A scalar JAX Array representing the laj of the transformation.
-        """
-        pass
+from bayinx.core.constraint import Constraint
 
 
 class LowerBound(Constraint):

bayinx/core/constraint.py

@@ -0,0 +1,26 @@
+from abc import abstractmethod
+from typing import Tuple
+
+import equinox as eqx
+from jaxtyping import Array, ArrayLike, Scalar
+
+
+class Constraint(eqx.Module):
+    """
+    Abstract base class for defining parameter constraints.
+    """
+
+    @abstractmethod
+    def constrain(self, x: ArrayLike) -> Tuple[Array, Scalar]:
+        """
+        Applies the constraining transformation to an unconstrained input and computes the log-absolute-jacobian of the transformation.
+
+        # Parameters
+        - `x`: The unconstrained JAX Array-like input.
+
+        # Returns
+        A tuple containing:
+            - The constrained JAX Array.
+            - A scalar JAX Array representing the laj of the transformation.
+        """
+        pass

bayinx/core/model.py

@@ -3,10 +3,10 @@ from typing import Any, Dict, Tuple
 
 import equinox as eqx
 import jax.numpy as jnp
-import jax.tree_util as jtu
+import jax.tree as jt
 from jaxtyping import Array, Scalar
 
-from bayinx.core.constraints import Constraint
+from bayinx.core.constraint import Constraint
 
 
 class Model(eqx.Module):
@@ -31,13 +31,13 @@ class Model(eqx.Module):
         Generates a filter specification to subset relevant parameters for the model.
         """
         # Generate empty specification
-        filter_spec = jtu.tree_map(lambda _: False, self)
+        filter_spec = jt.map(lambda _: False, self)
 
         # Specify JAX Array parameters
         filter_spec = eqx.tree_at(
             lambda model: model.params,
             filter_spec,
-            replace=jtu.tree_map(eqx.is_array, self.params),
+            replace=jt.map(eqx.is_array, self.params),
         )
 
         return filter_spec

bayinx/core/variational.py

@@ -19,8 +19,8 @@ class Variational(eqx.Module):
     An abstract base class used to define variational methods.
 
     # Attributes
-    - `_unflatten`: A static function to transform draws from the variational distribution back to a `Model`.
-    - `_constraints`: A static partitioned `Model` with the constraints of the `Model` used to initialize the `Variational` object.
+    - `_unflatten`: A function to transform draws from the variational distribution back to a `Model`.
+    - `_constraints`: The static component of a partitioned `Model` used to initialize the `Variational` object.
     """
 
     _unflatten: Callable[[Array], Model]

bayinx/dists/censored/gamma2/r.py

@@ -0,0 +1,65 @@
+import jax.lax as lax
+import jax.numpy as jnp
+from jax.scipy.special import gammaincc
+from jaxtyping import Array, ArrayLike, Float
+
+from bayinx.dists import gamma2
+
+
+def prob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    nu: Float[ArrayLike, "..."],
+    censor: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The mixed probability mass/density function (PMF/PDF) for a (mean-inverse dispersion parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:  Value(s) at which to evaluate the PMF/PDF.
+    - `mu`: The positive mean.
+    - `nu`: The positive inverse dispersion.
+
+    # Returns
+    The PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+    evals: Array = jnp.zeros_like(x * 1.0) # ensure float dtype
+
+    # Construct boolean masks
+    uncensored: Array = jnp.array(jnp.logical_and(0.0 < x, x < censor)) # pyright: ignore
+    censored: Array = jnp.array(x == censor) # pyright: ignore
+
+    # Evaluate mixed probability (?) function
+    evals = jnp.where(uncensored, gamma2.prob(x, mu, nu), evals)
+    evals = jnp.where(censored, gammaincc(nu, x * nu / mu), evals) # pyright: ignore
+
+    return evals
+
+
+def logprob(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    nu: Float[ArrayLike, "..."],
+    censor: Float[ArrayLike, "..."]
+) -> Float[Array, "..."]:
+    """
+    The log-transformed mixed probability mass/density function (log PMF/PDF) for a (mean-inverse dispersion parameterized) Gamma distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log PMF/PDF.
+    - `mu`:     The positive mean/location.
+    - `nu`:     The positive inverse dispersion.
+
+    # Returns
+    The log PMF/PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `nu`.
+    """
+    evals: Array = jnp.full_like(x * 1.0, -jnp.inf) # ensure float dtype
+
+    # Construct boolean masks
+    uncensored: Array = jnp.array(jnp.logical_and(0.0 < x, x < censor)) # pyright: ignore
+    censored: Array = jnp.array(x == censor) # pyright: ignore
+
+    evals = jnp.where(uncensored, gamma2.logprob(x, mu, nu), evals)
+    evals = jnp.where(censored, lax.log(gammaincc(nu, x * nu / mu)), evals) # pyright: ignore
+
+    return evals

bayinx/dists/gamma2.py

@@ -1,17 +1,17 @@
 import jax.lax as lax
 from jax.scipy.special import gammaln
-from jaxtyping import Array, ArrayLike, Float, Real
+from jaxtyping import Array, ArrayLike, Float
 
 
 def prob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], nu: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], nu: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The probability density function (PDF) for a (mean-precision parameterized) Gamma distribution.
 
     # Parameters
     - `x`:      Value(s) at which to evaluate the PDF.
-    - `mu`:     The mean.
+    - `mu`:     The positive mean.
     - `nu`:     The positive inverse dispersion.
 
     # Returns
@@ -22,14 +22,14 @@ def prob(
 
 
 def logprob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], nu: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], nu: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log-transformed probability density function (log PDF) for a (mean-precision parameterized) Gamma distribution.
 
     # Parameters
     - `x`:      Value(s) at which to evaluate the log PDF.
-    - `mu`:     The mean/location.
+    - `mu`:     The positive mean/location.
     - `nu`:     The positive inverse dispersion.
 
     # Returns

bayinx/dists/normal.py

@@ -1,11 +1,11 @@
 import jax.lax as lax
-from jaxtyping import Array, ArrayLike, Float, Real
+from jaxtyping import Array, ArrayLike, Float
 
 __PI = 3.141592653589793
 
 
 def prob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The probability density function (PDF) for a Normal distribution.
@@ -25,7 +25,7 @@ def prob(
 
 
 def logprob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the probability density function (log PDF) for a Normal distribution.
@@ -45,7 +45,7 @@ def logprob(
 
 
 def uprob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The unnormalized probability density function (uPDF) for a Normal distribution.
@@ -63,7 +63,7 @@ def uprob(
 
 
 def ulogprob(
-    x: Real[ArrayLike, "..."], mu: Real[ArrayLike, "..."], sigma: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], mu: Float[ArrayLike, "..."], sigma: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the unnormalized probability density function (log uPDF) for a Normal distribution.

bayinx/dists/uniform.py

@@ -1,10 +1,10 @@
 import jax.lax as _lax
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Float, Real
+from jaxtyping import Array, ArrayLike, Float
 
 
 def prob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The probability density function (PDF) for a Uniform distribution.
@@ -22,7 +22,7 @@ def prob(
 
 
 def logprob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the probability density function (log PDF) for a Uniform distribution.
@@ -40,7 +40,7 @@ def logprob(
 
 
 def uprob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The unnormalized probability density function (uPDF) for a Uniform distribution.
@@ -58,7 +58,7 @@ def uprob(
 
 
 def ulogprob(
-    x: Real[ArrayLike, "..."], lb: Real[ArrayLike, "..."], ub: Real[ArrayLike, "..."]
+    x: Float[ArrayLike, "..."], lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."]
 ) -> Float[Array, "..."]:
     """
     The log of the unnormalized probability density function (log uPDF) for a Uniform distribution.

bayinx/mhx/vi/meanfield.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Self
+from typing import Any, Dict, Self
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -20,8 +20,6 @@ class MeanField(Variational):
     """
 
     var_params: Dict[str, Float[Array, "..."]]
-    _unflatten: Callable[[Float[Array, "..."]], Model]
-    _constraints: Model
 
     def __init__(self, model: Model):
         """

{bayinx-0.2.29.dist-info → bayinx-0.2.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.2.29
+Version: 0.2.30
 Summary: Bayesian Inference with JAX
 Requires-Python: >=3.12
 Requires-Dist: equinox>=0.11.12

{bayinx-0.2.29.dist-info → bayinx-0.2.30.dist-info}/RECORD

@@ -1,18 +1,22 @@
 bayinx/__init__.py,sha256=l20JdkSsE_XGZlZFNEtySXf4NIlbjrao14vXPB-H6aQ,45
 bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/constraints/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/constraints/lower.py,sha256=O37qJ6ojRKbKGJlnQ7Vv7P2VGARcnKrOifwyUNBHma8,912
 bayinx/core/__init__.py,sha256=7vW2F8t3K4TWlSu5nZrYCdUrz5N9FMIfQQBn3IoeH6o,150
-bayinx/core/constraints.py,sha256=lbVs2-xjGRue17YRPGHz3s_mJ0ZiunpYowbD0QvcD-I,1525
+bayinx/core/constraint.py,sha256=T1QP3WV9RU1nPW-9H7loc6DyBzElHVyDp1c4mmZpnmI,733
 bayinx/core/flow.py,sha256=A5Vw5t76LPasnMgghjw6ulBkIm5L2jBprusVt-tuwko,2296
-bayinx/core/model.py,sha256=Z_HaFr0_-keMjG5tg3xxP3hGML7aDFIcCI8Y5dGrtM4,2145
-bayinx/core/variational.py,sha256=W0747jfVJFAtMZqL3mpbtl2wfnARHln-dVBag4xZ23Y,4813
+bayinx/core/model.py,sha256=vfEnqBpHE2MtuJPIDgKvVYIv5n53E2e-KAAXEtqEy0c,2126
+bayinx/core/variational.py,sha256=2stsYKZDri1rLP7mrz7X2GWehBXNESdlWtmF2N9CEas,4787
 bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
-bayinx/dists/gamma2.py,sha256=jc9jiNuuIAv3tdWghQ2Y2ANpJbDscPJKOvs6dOJVJD0,1315
-bayinx/dists/normal.py,sha256=AXsf3Xe2BfCyzfcQ5i8J9AD92LedQjgqviUMQd697D8,2628
-bayinx/dists/uniform.py,sha256=mogFe8VuDelM9KXE6RxGek0-tuZYFrwmo_oMOPHXleA,2359
+bayinx/dists/gamma2.py,sha256=8XYaOtcYJCrr5q1yHWfZaMJmASpLOrfyhrH_J06ksj8,1333
+bayinx/dists/normal.py,sha256=mvm6EoAlORy-yivuhMcExYCZUo0vJzMKMOWH-9iQBZU,2634
+bayinx/dists/uniform.py,sha256=7XgVvOrzINEFA6HJTYUOFwlWhEtrQQQ1aPJ_ZLOzLEc,2365
+bayinx/dists/censored/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/censored/gamma2/r.py,sha256=3brRCKhE-74mRXyIyPcnyaWY2OJv8CZyUWPP9T1t09Y,2274
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/vi/__init__.py,sha256=YfkXKsqo9Dk_AmQGjZKm4vfG8eLer2ez92G-cOExphs,193
-bayinx/mhx/vi/meanfield.py,sha256=8hM1KZ52TpRPLwiQcowsJLlQ-5nJzUEcKrtDiGrFoSs,3732
+bayinx/mhx/vi/meanfield.py,sha256=BobfTagVGA5R-dclv-E0jSA80KZg1X6GGjiw7XR61vE,3643
 bayinx/mhx/vi/normalizing_flow.py,sha256=FvxDtqGRtaEeeF-bXCYnIEAvOOXVHKUK0oCTF9ma02Y,4622
 bayinx/mhx/vi/standard.py,sha256=HaJsIz70Qo1Ql2hMQ-GQhcnfWiOGtyxgkOsm_yQaDKI,1718
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
@@ -20,6 +24,6 @@ bayinx/mhx/vi/flows/fullaffine.py,sha256=s-fxgzv84BEqNMnmLt6vtvwJqBzixCC2OwWXLz3
 bayinx/mhx/vi/flows/planar.py,sha256=WVj-oxcRctuoRA6KJjU63ek1ZgKNG2vI-TLN0QqjtKA,1916
 bayinx/mhx/vi/flows/radial.py,sha256=Obj3SraliawIHmP14F9wRpWt34y3kscY--Izy24eCvM,2499
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.2.29.dist-info/METADATA,sha256=kfq898UxqnChe-87HPe46_icfS2btMbJPp3INv3IXo0,3058
-bayinx-0.2.29.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.2.29.dist-info/RECORD,,
+bayinx-0.2.30.dist-info/METADATA,sha256=VVDldLMllMzTkx_tphb6k2n4u-PSwBbmFOGN-PG1BWc,3058
+bayinx-0.2.30.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.2.30.dist-info/RECORD,,