bayinx 0.4.0__tar.gz → 0.4.1__tar.gz

{bayinx-0.4.0 → bayinx-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.4.0
+Version: 0.4.1
 Summary: Bayesian Inference with JAX
 Author: Todd McCready
 Maintainer: Todd McCready
@@ -12,13 +12,13 @@ Requires-Dist: jaxtyping>=0.2.36
 Requires-Dist: optax>=0.2.4
 Description-Content-Type: text/markdown
 
-# `Bayinx`: <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
+# Bayinx: <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
 
-The original aim of this project was to build a PPL in Python that is similar in feel to `Stan` or `Nimble`(where there is a nice declarative syntax for defining the model) and allows for arbitrary models(e.g., ones with discrete parameters that may not be just integers); most of this goal has been moved to [baycian](https://github.com/toddmccready/baycian) for the foreseeable future.
+The original aim of this project was to build a PPL in Python that is similar in feel to Stan or Nimble(where there is a nice declarative syntax for defining the model) and allows for arbitrary models(e.g., ones with discrete parameters that may not be just integers); most of this goal has been moved to [baycian](https://github.com/toddmccready/baycian) for the foreseeable future.
 
-Part of the reason for this move is that Rust's ability to embed a "nice" DSL is comparitively easier due to [Rust macros](https://doc.rust-lang.org/rust-by-example/macros/dsl.html); I can define syntax similar to `Stan` and parse it to valid Rust code. Additionally, the current state of `bayinx` is relatively functional(plus/minus a few things to clean-up and documentation) and it offers enough functionality for one of my other projects: [disize](https://github.com/toddmccready/disize)! I plan to rewrite `disize` in Python with JAX, and `bayinx` makes it easy to handle constraining transformations, filtering for parameters for gradient calculations, etc.
+Part of the reason for this move is that Rust's ability to embed a "nice" DSL is comparitively easier due to [Rust macros](https://doc.rust-lang.org/rust-by-example/macros/dsl.html); I can define syntax similar to Stan and parse it to valid Rust code. Additionally, the current state of bayinx is relatively functional(plus/minus a few things to clean-up and documentation) and it offers enough for one of my other projects: [disize](https://github.com/toddmccready/disize)! I plan to rewrite disize in Python with JAX, and bayinx makes it easy to handle constraining transformations, filtering for parameters for gradient calculations, etc.
 
-Instead, this project is narrowing on implementing much of `Stan`'s functionality(restricted to continuously parameterized models, point estimation + vi + mcmc, etc) without most of the nice syntax, at least for versions `0.4.#`. Therefore, people will work with `target` directly and return the density like below:
+Instead, this project is narrowing on implementing much of Stan's functionality(restricted to continuously parameterized models, point estimation + vi + mcmc, etc) without most of the nice syntax, at least for versions `0.4.#`. Therefore, people will work with `target` directly and return the density like below:
 
 ```py
 class NormalDist(Model):
@@ -34,7 +34,7 @@ class NormalDist(Model):
         return target
 ```
 
-I have ideas for using a context manager and implementing `Node`: `Observed`/`Stochastic` classes that will try and replicate what `baycian` is trying to do, but that is for the future and versions `0.4.#` will retain the functionality needed for `disize`.
+I have ideas for using a context manager and implementing `Node`: `Observed`/`Stochastic` classes that will try and replicate what `baycian` is trying to do, but that is for the future and versions `0.4.#` will retain the functionality needed for disize.
 
 
 # TODO

{bayinx-0.4.0 → bayinx-0.4.1}/README.md

@@ -1,10 +1,10 @@
-# `Bayinx`: <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
+# Bayinx: <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
 
-The original aim of this project was to build a PPL in Python that is similar in feel to `Stan` or `Nimble`(where there is a nice declarative syntax for defining the model) and allows for arbitrary models(e.g., ones with discrete parameters that may not be just integers); most of this goal has been moved to [baycian](https://github.com/toddmccready/baycian) for the foreseeable future.
+The original aim of this project was to build a PPL in Python that is similar in feel to Stan or Nimble(where there is a nice declarative syntax for defining the model) and allows for arbitrary models(e.g., ones with discrete parameters that may not be just integers); most of this goal has been moved to [baycian](https://github.com/toddmccready/baycian) for the foreseeable future.
 
-Part of the reason for this move is that Rust's ability to embed a "nice" DSL is comparitively easier due to [Rust macros](https://doc.rust-lang.org/rust-by-example/macros/dsl.html); I can define syntax similar to `Stan` and parse it to valid Rust code. Additionally, the current state of `bayinx` is relatively functional(plus/minus a few things to clean-up and documentation) and it offers enough functionality for one of my other projects: [disize](https://github.com/toddmccready/disize)! I plan to rewrite `disize` in Python with JAX, and `bayinx` makes it easy to handle constraining transformations, filtering for parameters for gradient calculations, etc.
+Part of the reason for this move is that Rust's ability to embed a "nice" DSL is comparitively easier due to [Rust macros](https://doc.rust-lang.org/rust-by-example/macros/dsl.html); I can define syntax similar to Stan and parse it to valid Rust code. Additionally, the current state of bayinx is relatively functional(plus/minus a few things to clean-up and documentation) and it offers enough for one of my other projects: [disize](https://github.com/toddmccready/disize)! I plan to rewrite disize in Python with JAX, and bayinx makes it easy to handle constraining transformations, filtering for parameters for gradient calculations, etc.
 
-Instead, this project is narrowing on implementing much of `Stan`'s functionality(restricted to continuously parameterized models, point estimation + vi + mcmc, etc) without most of the nice syntax, at least for versions `0.4.#`. Therefore, people will work with `target` directly and return the density like below:
+Instead, this project is narrowing on implementing much of Stan's functionality(restricted to continuously parameterized models, point estimation + vi + mcmc, etc) without most of the nice syntax, at least for versions `0.4.#`. Therefore, people will work with `target` directly and return the density like below:
 
 ```py
 class NormalDist(Model):
@@ -20,7 +20,7 @@ class NormalDist(Model):
         return target
 ```
 
-I have ideas for using a context manager and implementing `Node`: `Observed`/`Stochastic` classes that will try and replicate what `baycian` is trying to do, but that is for the future and versions `0.4.#` will retain the functionality needed for `disize`.
+I have ideas for using a context manager and implementing `Node`: `Observed`/`Stochastic` classes that will try and replicate what `baycian` is trying to do, but that is for the future and versions `0.4.#` will retain the functionality needed for disize.
 
 
 # TODO

{bayinx-0.4.0 → bayinx-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.4.0"
+version = "0.4.1"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -21,7 +21,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.4.0"
+current_version = "0.4.1"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

bayinx-0.4.1/src/bayinx/constraints.py

@@ -0,0 +1,135 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax.nn as jnn
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import Array, PyTree, Scalar, ScalarLike
+
+from bayinx.core import Constraint, Parameter
+from bayinx.core._parameter import T
+
+
+class Lower(Constraint):
+    """
+    Enforces a lower bound on the parameter.
+    """
+
+    lb: Scalar
+
+    def __init__(self, lb: ScalarLike):
+        # assert greater than 1
+        self.lb = jnp.asarray(lb)
+
+    def constrain(self, param: Parameter[T]) -> Tuple[Parameter[T], Scalar]:
+        """
+        Enforces a lower bound on the parameter and adjusts the posterior density.
+
+        # Parameters
+        - `param`: The unconstrained `Parameter`.
+
+        # Returns
+        A tuple containing:
+            - A modified `Parameter` with relevant leaves satisfying the constraint.
+            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+        """
+        # Extract relevant parameters(all inexact Arrays)
+        dyn, static = eqx.partition(param, param.filter_spec)
+
+        # Compute Jacobian adjustment
+        total_laj: Scalar = jt.reduce(lambda a, b: a + b, jt.map(jnp.sum, dyn))
+
+        # Compute transformation
+        dyn = jt.map(lambda v: jnp.exp(v) + self.lb, dyn)
+
+        # Combine into full parameter object
+        param = eqx.combine(dyn, static)
+
+        return param, total_laj
+
+
+class LogSimplex(Constraint):
+    """
+    Enforces a log-transformed simplex constraint on the parameter.
+
+    # Attributes
+    - `sum`: The total sum of the parameter.
+    """
+
+    sum: Scalar
+
+    def __init__(self, sum_val: ScalarLike = 1.0):
+        """
+        # Parameters
+        - `sum_val`: The target sum for the exponentiated simplex. Defaults to 1.0.
+        """
+        self.sum = jnp.asarray(sum_val)
+
+    def constrain(self, param: Parameter[T]) -> Tuple[Parameter[T], Scalar]:
+        """
+        Enforces a log-transformed simplex constraint on the parameter and adjusts the posterior density.
+
+        # Parameters
+        - `param`: The unconstrained `Parameter`.
+
+        # Returns
+        A tuple containing:
+            - A modified `Parameter` with relevant leaves satisfying the constraint.
+            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+        """
+        # Partition the parameter into dynamic (to be transformed) and static parts
+        dyn, static = eqx.partition(param, param.filter_spec)
+
+        # Map transformation leaf-wise
+        transformed = jt.map(self._transform_leaf, dyn) ## filter spec handles subsetting arrays, is_leaf unnecessary
+
+        # Extract constrained parameters and Jacobian adjustments
+        dyn_constrained: PyTree = jt.map(lambda x: x[0], transformed)
+        lajs: PyTree = jt.map(lambda x: x[1], transformed)
+
+        # Sum to get total Jacobian adjustment
+        total_laj = jt.reduce(lambda a, b: a + b, lajs)
+
+        # Recombine the transformed dynamic parts with the static parts
+        param = eqx.combine(dyn_constrained, static)
+
+        return param, total_laj
+
+    def _transform_leaf(self, x: Array) -> Tuple[Array, Scalar]:
+        """
+        Internal function that applies a log-transformed simplex constraint on a single array.
+        """
+        laj: Scalar = jnp.array(0.0)
+
+        # Save output shape
+        output_shape: tuple[int, ...] = x.shape
+
+        if x.size == 1:
+            return(jnp.full(output_shape, jnp.log(self.sum)), laj)
+        else:
+            # Flatten x
+            x = x.flatten()
+
+            # Subset first K - 1 elements
+            x = x[:-1]
+
+            # Compute shifted cumulative sum
+            zeta: Array = jnp.concat([jnp.zeros(1), x.cumsum()[:-1]])
+
+            # Compute intermediate proportions vector
+            eta: Array = jnn.sigmoid(x - zeta)
+
+            # Compute Jacobian adjustment
+            laj += jnp.sum(jnp.log(eta) + jnp.log(1 - eta)) # TODO: check for correctness
+
+            # Compute log-transformed simplex weights
+            w: Array = jnp.log(eta) + jnp.concatenate([jnp.array([0.0]), jnp.log(jnp.cumprod((1-eta)[:-1]))])
+            w = jnp.concatenate([w, jnp.log(jnp.prod(1 - eta, keepdims=True))])
+
+            # Scale unit simplex on log-scale
+            w = w + jnp.log(self.sum)
+
+            # Reshape for output
+            w = w.reshape(output_shape)
+
+            return (w, laj)

bayinx-0.4.0/src/bayinx/constraints.py

@@ -1,46 +0,0 @@
-from typing import Tuple
-
-import equinox as eqx
-import jax.numpy as jnp
-import jax.tree as jt
-from jaxtyping import Scalar, ScalarLike
-
-from bayinx.core import Constraint, Parameter
-from bayinx.core._parameter import T
-
-
-class Lower(Constraint):
-    """
-    Enforces a lower bound on the parameter.
-    """
-
-    lb: Scalar
-
-    def __init__(self, lb: ScalarLike):
-        self.lb = jnp.array(lb)
-
-    def constrain(self, param: Parameter[T]) -> Tuple[Parameter[T], Scalar]:
-        """
-        Enforces a lower bound on the parameter and adjusts the posterior density.
-
-        # Parameters
-        - `param`: The unconstrained `Parameter`.
-
-        # Parameters
-        A tuple containing:
-            - A modified `Parameter` with relevant leaves satisfying the constraint.
-            - A scalar Array representing the log-absolute-Jacobian of the transformation.
-        """
-        # Extract relevant parameters(all Array)
-        dyn, static = eqx.partition(param, param.filter_spec)
-
-        # Compute density adjustment
-        laj: Scalar = jt.reduce(lambda a, b: a + b, jt.map(jnp.sum, dyn))
-
-        # Compute transformation
-        dyn = jt.map(lambda v: jnp.exp(v) + self.lb, dyn)
-
-        # Combine into full parameter object
-        param = eqx.combine(dyn, static)
-
-        return param, laj