bayinx 0.4.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

bayinx/constraints.py

@@ -1,9 +1,10 @@
 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 Scalar, ScalarLike
+from jaxtyping import Array, PyTree, Scalar, ScalarLike
 
 from bayinx.core import Constraint, Parameter
 from bayinx.core._parameter import T
@@ -17,7 +18,8 @@ class Lower(Constraint):
     lb: Scalar
 
     def __init__(self, lb: ScalarLike):
-        self.lb = jnp.array(lb)
+        # assert greater than 1
+        self.lb = jnp.asarray(lb)
 
     def constrain(self, param: Parameter[T]) -> Tuple[Parameter[T], Scalar]:
         """
@@ -26,16 +28,16 @@ class Lower(Constraint):
         # Parameters
         - `param`: The unconstrained `Parameter`.
 
-        # Parameters
+        # 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 Array)
+        # Extract relevant parameters(all inexact Arrays)
         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 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)
@@ -43,4 +45,91 @@ class Lower(Constraint):
         # Combine into full parameter object
         param = eqx.combine(dyn, static)
 
-        return param, laj
+        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.dist-info → bayinx-0.4.1.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 bayinx/__init__.py,sha256=8etrxEtEGEzSDmKsW0TB4XoUGLiMPt9wpwNR8CGe1gU,93
-bayinx/constraints.py,sha256=wlPOfcu6yP_s-A4nq0wE3-N7bOggfLQ3w5lRVno8-M8,1308
+bayinx/constraints.py,sha256=2ufHsXR-_bWKR4WKKuR-OTjj3XCc4TkSeHVGWYadwCg,4387
 bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/core/__init__.py,sha256=samkrHp2zYyj8n37k-06tlaVrSqbtcgoa1LO0btAEHc,338
 bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,761
@@ -32,7 +32,7 @@ bayinx/mhx/vi/flows/fullaffine.py,sha256=7KXuukwzVtMRIa8bSK_4pjnnP-lLIzVJBCAuKVy
 bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
 bayinx/mhx/vi/flows/radial.py,sha256=AyaqLJCwn871L6E8lBCU4Y8zZBF9UYZu6KIhzV6Z6wo,2503
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.4.0.dist-info/METADATA,sha256=lE2h9zq-1Jl15zpP8RiPh9rcNZ-JsebkWbqjZPSmqRQ,3021
-bayinx-0.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.4.0.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.4.0.dist-info/RECORD,,
+bayinx-0.4.1.dist-info/METADATA,sha256=7Zw-9hVqUVxj3ncyGBfn72FQzTDomdDaXXH2hOsJM60,2989
+bayinx-0.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.4.1.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.4.1.dist-info/RECORD,,