bayinx 0.3.11__tar.gz → 0.3.13__tar.gz

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

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

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

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

{bayinx-0.3.11 → bayinx-0.3.13}/src/bayinx/constraints/__init__.py

@@ -1,3 +1,3 @@
 from bayinx.constraints.lower import Lower
 
-__all__ = ['Lower']
+__all__ = ["Lower"]

bayinx-0.3.13/src/bayinx/core/__init__.py

@@ -0,0 +1,16 @@
+from ._constraint import Constraint
+from ._flow import Flow
+from ._model import Model, constrain
+from ._optimization import optimize_model
+from ._parameter import Parameter
+from ._variational import Variational
+
+__all__ = [
+    "Constraint",
+    "Flow",
+    "Model",
+    "constrain",
+    "optimize_model",
+    "Parameter",
+    "Variational",
+]

bayinx-0.3.13/src/bayinx/core/_optimization.py

@@ -0,0 +1,94 @@
+from typing import Any, Tuple
+
+import equinox as eqx
+import jax.lax as lax
+import jax.numpy as jnp
+import optax as opx
+from jaxtyping import PyTree, Scalar
+from optax import GradientTransformation, OptState, Schedule
+
+from ._model import Model
+
+
+@eqx.filter_jit
+def optimize_model(
+    model: Model,
+    max_iters: int,
+    data: Any = None,
+    learning_rate: float = 1,
+    weight_decay: float = 0.0,
+    tolerance: float = 1e-4,
+) -> Model:
+    """
+    Optimize the dynamic parameters of the model.
+
+    # Parameters
+    - `max_iters`: Maximum number of iterations for the optimization loop.
+    - `data`: Data to evaluate the model with.
+    - `learning_rate`: Initial learning rate for optimization.
+    - `weight_decay`: Weight decay for the optimizer.
+    - `tolerance`: Relative tolerance of loss decrease for stopping early (not implemented in the loop).
+    """
+    # Get dynamic and static componts of model
+    dyn, static = eqx.partition(model, model.filter_spec)
+
+    # Derive gradient for posterior
+    @eqx.filter_jit
+    @eqx.filter_grad
+    def eval_grad(dyn: Model):
+        # Reconstruct model
+        model: Model = eqx.combine(dyn, static)
+
+        # Evaluate posterior
+        return model.eval(data)
+
+    # Construct scheduler
+    schedule: Schedule = opx.warmup_cosine_decay_schedule(
+        init_value=1e-16,
+        peak_value=learning_rate,
+        warmup_steps=int(max_iters / 10),
+        decay_steps=max_iters - int(max_iters / 10),
+    )
+
+    optim: GradientTransformation = opx.chain(
+        opx.scale(-1.0), opx.nadamw(schedule, weight_decay=weight_decay)
+    )
+    opt_state: OptState = optim.init(dyn)
+
+    @eqx.filter_jit
+    def condition(state: Tuple[PyTree, OptState, Scalar]):
+        # Unpack iteration state
+        current_opt_dyn, opt_state, i = state
+
+        return i < max_iters
+
+    @eqx.filter_jit
+    def body(state: Tuple[PyTree, OptState, Scalar]):
+        # Unpack iteration state
+        dyn, opt_state, i = state
+
+        # Update iteration
+        i = i + 1
+
+        # Evaluate gradient of posterior
+        updates = eval_grad(dyn)
+
+        # Compute updates
+        updates, opt_state = optim.update(
+            updates, opt_state, eqx.filter(dyn, dyn.filter_spec)
+        )
+
+        # Update model
+        dyn = eqx.apply_updates(dyn, updates)
+
+        return dyn, opt_state, i
+
+    # Run optimization loop
+    dyn = lax.while_loop(
+        cond_fun=condition,
+        body_fun=body,
+        init_val=(dyn, opt_state, jnp.array(0, jnp.uint32)),
+    )[0]
+
+    # Return optimized model
+    return eqx.combine(dyn, static)

{bayinx-0.3.11 → bayinx-0.3.13}/src/bayinx/core/_parameter.py

@@ -5,6 +5,8 @@ import jax.tree as jt
 from jaxtyping import PyTree
 
 T = TypeVar("T", bound=PyTree)
+
+
 class Parameter(eqx.Module, Generic[T]):
     """
     A container for a parameter of a `Model`.

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

@@ -13,7 +13,9 @@ from optax import GradientTransformation, OptState, Schedule
 
 from ._model import Model
 
-M = TypeVar('M', bound=Model)
+M = TypeVar("M", bound=Model)
+
+
 class Variational(eqx.Module, Generic[M]):
     """
     An abstract base class used to define variational methods.
@@ -80,6 +82,7 @@ class Variational(eqx.Module, Generic[M]):
         # Evaluate posterior density
         return model.eval(data)
 
+    # TODO: get rid of this and put it all in each vari's methods, forgot abt discrete parameters :V
     @eqx.filter_jit
     def fit(
         self,
@@ -107,7 +110,10 @@ class Variational(eqx.Module, Generic[M]):
 
         # Construct scheduler
         schedule: Schedule = opx.warmup_cosine_decay_schedule(
-            init_value=1e-16, peak_value=learning_rate, warmup_steps=int(max_iters/10), decay_steps=max_iters-int(max_iters/10)
+            init_value=1e-16,
+            peak_value=learning_rate,
+            warmup_steps=int(max_iters / 10),
+            decay_steps=max_iters - int(max_iters / 10),
         )
 
         # Initialize optimizer
@@ -163,7 +169,11 @@ class Variational(eqx.Module, Generic[M]):
 
     @eqx.filter_jit
     def posterior_predictive(
-        self, func: Callable[[M, Any], 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)

{bayinx-0.3.11 → bayinx-0.3.13}/src/bayinx/dists/__init__.py

@@ -1,3 +1,3 @@
 from bayinx.dists import censored, gamma2, normal, posnormal
 
-__all__ = ['censored', "gamma2", "normal", "posnormal"]
+__all__ = ["censored", "gamma2", "normal", "posnormal"]

{bayinx-0.3.11 → bayinx-0.3.13}/src/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(
@@ -78,12 +79,13 @@ def logprob(
 
     return evals
 
+
 def sample(
     n: int,
     mu: Float[ArrayLike, "..."],
     sigma: Float[ArrayLike, "..."],
     censor: Float[ArrayLike, "..."] = jnp.inf,
-    key: Key = jr.PRNGKey(0)
+    key: Key = jr.PRNGKey(0),
 ) -> Float[Array, "..."]:
     """
     Sample from a right-censored positive Normal distribution.
@@ -107,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-0.3.11 → bayinx-0.3.13}/src/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.11 → bayinx-0.3.13}/src/bayinx/mhx/vi/__init__.py

@@ -2,4 +2,4 @@ from bayinx.mhx.vi.meanfield import MeanField
 from bayinx.mhx.vi.normalizing_flow import NormalizingFlow
 from bayinx.mhx.vi.standard import Standard
 
-__all__ = ['MeanField', 'NormalizingFlow', 'Standard']
+__all__ = ["MeanField", "NormalizingFlow", "Standard"]

{bayinx-0.3.11 → bayinx-0.3.13}/src/bayinx/mhx/vi/meanfield.py

@@ -10,7 +10,9 @@ from jaxtyping import Array, Float, Key, Scalar
 from bayinx.core import Model, Variational
 from bayinx.dists import normal
 
-M = TypeVar('M', bound=Model)
+M = TypeVar("M", bound=Model)
+
+
 class MeanField(Variational, Generic[M]):
     """
     A fully factorized Gaussian approximation to a posterior distribution.
@@ -19,7 +21,7 @@ class MeanField(Variational, Generic[M]):
     - `var_params`: The variational parameters for the approximation.
     """
 
-    var_params: Dict[str, Float[Array, "..."]] #todo: just expand to attributes
+    var_params: Dict[str, Float[Array, "..."]]  # todo: just expand to attributes
 
     def __init__(self, model: M):
         """

{bayinx-0.3.11 → bayinx-0.3.13}/src/bayinx/mhx/vi/normalizing_flow.py

@@ -9,7 +9,9 @@ from jaxtyping import Array, Key, Scalar
 
 from bayinx.core import Flow, Model, Variational
 
-M = TypeVar('M', bound=Model)
+M = TypeVar("M", bound=Model)
+
+
 class NormalizingFlow(Variational, Generic[M]):
     """
     An ordered collection of diffeomorphisms that map a base distribution to a

bayinx-0.3.11/src/bayinx/core/__init__.py

@@ -1,7 +0,0 @@
-from ._constraint import Constraint
-from ._flow import Flow
-from ._model import Model, constrain
-from ._parameter import Parameter
-from ._variational import Variational
-
-__all__ = ["Constraint", "Flow", "Model", "constrain", "Parameter", "Variational"]

{bayinx-0.3.11 → bayinx-0.3.13}/src/bayinx/mhx/vi/standard.py

@@ -1,4 +1,3 @@
-
 import equinox as eqx
 import jax.numpy as jnp
 import jax.random as jr
@@ -17,6 +16,7 @@ class Standard(Variational[M]):
     # Attributes
     - `dim`: Dimension of the parameter space.
     """
+
     dim: int
 
     def __init__(self, model: M):