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

bayinx/constraints/__init__.py

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

bayinx/core/__init__.py

@@ -1,7 +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", "Parameter", "Variational"]
+__all__ = [
+    "Constraint",
+    "Flow",
+    "Model",
+    "constrain",
+    "optimize_model",
+    "Parameter",
+    "Variational",
+]

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/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/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)
@@ -171,11 +181,11 @@ class Variational(eqx.Module, Generic[M]):
         # Evaluate posterior predictive
         @jax.jit
         @jax.vmap
-        def evaluate(draw: Array, data: Any = None):
+        def evaluate(draw: Array):
             # Reconstruct model
             model: M = self._unflatten(draw)
 
             # Evaluate
             return func(model, data)
 
-        return evaluate(draws, data)
+        return evaluate(draws)

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/dists/censored/posnormal/r.py

@@ -78,12 +78,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.

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

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

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

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

@@ -1,14 +1,15 @@
 bayinx/__init__.py,sha256=TM-aoRaPX6jSYtCM7Jv59TPV-H6bcDk1-VMttYP1KME,99
 bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/constraints/__init__.py,sha256=PiWXZKi7YdbTMKvw-OE5f-t87jJT893uAFrwWWBfOdg,64
+bayinx/constraints/__init__.py,sha256=027WJxRLkybXZkmusfvR6iZayY2pDid7Tw6TTTeA6ko,64
 bayinx/constraints/lower.py,sha256=30y0l6PF-tbS9LR_tto9AvwmsvXq1ExU-v8DLrJD4g4,1446
-bayinx/core/__init__.py,sha256=bZvQITgW0DWuPKl3wCLKt6WHKogYKx8Zz36g8z9Aung,253
+bayinx/core/__init__.py,sha256=Qmy0EjzqqKwI9F8rjmC9j6J8hiDw6A54yOck2WuQJkY,344
 bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,761
 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=X8o81b8vyU7vJxw8pZYH_nxc3u990tRUZgRhMNodNI4,5484
-bayinx/dists/__init__.py,sha256=9DdPea7HAnBOzaV_4gM5noPX8YCb_p06d8PJvGfFy3Y,118
+bayinx/core/_optimization.py,sha256=cXO07guDG5kd64kWq4_de-gLgxTT6vIOU3IOL3TMl6U,2583
+bayinx/core/_parameter.py,sha256=UfLonbTwxzr1g76Cf3HzAh9u4UBtlcaLByYLXgq-aCQ,1082
+bayinx/core/_variational.py,sha256=bQYiN8c4AGPt4hsNT68zN7J6o0fdsLGwVjbDyl62LnI,5639
+bayinx/dists/__init__.py,sha256=BIrypqMnTLWK3a_zw8fYKMyuEMxP_qGsLfLeScias0o,118
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/gamma2.py,sha256=MuFudL2UTfk8HgWVofNaR36JTmUpmtxvg1Mifu98MvM,1567
 bayinx/dists/normal.py,sha256=Yc2X8F7JoLYwprtK8bA2BPva1tAY7MEs3oSk5pMortI,3822
@@ -18,18 +19,19 @@ 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=Ypi6w_t53pAzRVzjcStx2RhozkAlCDLnQmgKykhpQQ4,3426
+bayinx/dists/censored/posnormal/r.py,sha256=r8QfEThaaqsqqrr6PHWiY0EdfnjQcbKdTAHN2GmDVzI,3428
 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
-bayinx/mhx/vi/normalizing_flow.py,sha256=8pLMDdZPIt5wlgbhHWSFY1ChSWM9pvSD2bQx3zgz1F8,4710
-bayinx/mhx/vi/standard.py,sha256=W-ZvigJkUpqVlREgiFm9io8ansT1XpZwq5AqSmdv--E,1578
+bayinx/mhx/opt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/mhx/vi/__init__.py,sha256=3T1dEpiiRge4tW-vpS0xBob_RbO1iVFnL3fVCRUawCM,205
+bayinx/mhx/vi/meanfield.py,sha256=kx5WeD93-XO8NHxd4L4pZ8V19Y9B6j-yE3Y5OBXMcTk,3743
+bayinx/mhx/vi/normalizing_flow.py,sha256=vzLu5H1G1-pBqhgHWmIZkUTyPE1DxC9vBwpiZeIyu1I,4712
+bayinx/mhx/vi/standard.py,sha256=LYgglaGQMGmXpzFR4eMJnXkl2PhBeggbXMvO5zJpf2c,1578
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
 bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvADwVYQ,1954
 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.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,,
+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,,