numerax 0.1.0__py3-none-any.whl

numerax/__init__.py

@@ -0,0 +1,37 @@
+"""
+Statistical and numerical computation functions for JAX, focusing on tools
+not available in the main JAX API.
+
+## Overview
+
+This package provides JAX-compatible implementations of specialized numerical
+functions with full differentiability support. All functions are designed to
+work seamlessly with JAX's transformations (JIT, grad, vmap, etc.) and follow
+JAX's functional programming paradigms.
+
+### Special Functions (`numerax.special`)
+
+Mathematical special functions with custom derivative implementations.
+Functions use numerically stable algorithms and provide exact gradients
+through custom JVP rules where standard automatic differentiation would
+be inefficient or unstable.
+
+### Statistical Methods (`numerax.stats`)
+
+Advanced statistical computation tools for inference problems. Implements
+efficient algorithms for complex statistical models, with particular focus
+on optimization-based methods that benefit from JAX's compilation and
+differentiation capabilities.
+
+### Utilities (`numerax.utils`)
+
+Development utilities for creating JAX-compatible functions with proper
+documentation support. Includes decorators and helpers for preserving
+function metadata when using JAX's advanced features like custom derivatives.
+"""
+
+from . import special, stats, utils
+
+__version__ = "0.1.0"
+
+__all__ = ["special", "stats", "utils"]

numerax/special/__init__.py

@@ -0,0 +1,3 @@
+from .gamma import gammap_inverse
+
+__all__ = ["gammap_inverse"]

numerax/special/gamma.py

@@ -0,0 +1,217 @@
+import jax
+import jax.numpy as jnp
+import jax.scipy.special as special
+from jaxtyping import ArrayLike
+
+# Global constants for numerical stability - adapt to JAX precision setting
+_DTYPE = jnp.float64 if jax.config.jax_enable_x64 else jnp.float32
+TINY = jnp.finfo(_DTYPE).smallest_normal  # For preventing underflow
+EPS = jnp.finfo(_DTYPE).eps  # For convergence tolerance (machine epsilon)
+
+
+@jax.custom_jvp
+def gammap_inverse(p: ArrayLike, a: float) -> ArrayLike:
+    r"""
+    Inverse of the regularized incomplete gamma function.
+
+    ## Overview
+
+    Computes the inverse of the regularized incomplete gamma function, finding
+    $x$ such that $P(a, x) = p$, where $P(a, x)$ is the regularized incomplete
+    gamma function. This is equivalent to computing quantiles of the
+    $\text{Gamma}(a, 1)$ distribution. The general strategy and the initial
+    guess are based on the methods described in
+    Numerical Recipes (Press et al., 2007).
+
+    ## Mathematical Background
+
+    The regularized incomplete gamma function is defined as:
+
+    $$P(a, x) = \frac{\gamma(a, x)}{\Gamma(a)} = \frac{1}{\Gamma(a)}
+    \int_0^x t^{a-1} e^{-t} dt$$
+
+    This function solves the inverse problem:
+
+    $$x = P^{-1}(a, p) \quad \text{such that} \quad P(a, x) = p$$
+
+    For a random variable $X \sim \text{Gamma}(a, 1)$, this gives:
+
+    $$x = F^{-1}(p) \quad \text{where} \quad F(x) = P(\Gamma(a), x)$$
+
+    ## Numerical Method
+
+    Uses Halley's method for fast quadratic convergence:
+
+    $$x_{n+1} = x_n - \frac{2f(x_n)f'(x_n)}{2[f'(x_n)]^2 - f(x_n)f''(x_n)}$$
+
+    where $f(x) = P(a, x) - p$.
+
+    **Initial guess** based on Numerical Recipes (Press et al., 2007):
+    - For $a > 1$: Wilson-Hilferty approximation
+    - For $a \leq 1$: Asymptotic expansions
+
+    ## Args
+
+    - **p**: Probability values in $[0, 1]$. Can be scalar or array.
+    - **a**: Shape parameter (must be positive). Scalar value.
+
+    ## Returns
+
+    Quantiles $x$ where $P(a, x) = p$. Same shape as input `p`.
+
+    ## Example
+
+    ```python
+    import jax.numpy as jnp
+    import numerax
+
+    # Single quantile
+    x = numerax.special.gammap_inverse(0.5, 2.0)  # Median of Gamma(2, 1)
+
+    # Multiple quantiles
+    p_vals = jnp.array([0.1, 0.25, 0.5, 0.75, 0.9])
+    x_vals = numerax.special.gammap_inverse(p_vals, 3.0)
+
+    # Verify inverse relationship
+    from jax.scipy.special import gammainc
+
+    p_recovered = gammainc(2.0, x)  # Should equal original p
+
+    # Differentiable for optimization
+    grad_fn = jax.grad(numerax.special.gammap_inverse)
+    sensitivity = grad_fn(0.5, 2.0)  # ∂x/∂p at median
+    ```
+
+    ## Notes
+
+    - **Convergence**: Typically converges in 3-8 iterations
+    - **Differentiable**: Custom JVP implementation using implicit function
+      theorem
+    - **Numerical stability**: Handles edge cases near 0 and 1
+    - **Performance**: JIT-compiled with adaptive precision
+    - **Domain**: $p \in [0, 1]$ and $a > 0$
+
+    ## References
+
+    Press, W. H., Teukolsky, S. A., Vetterling, W. T., & Flannery, B. P.
+    (2007).
+    *Numerical Recipes: The Art of Scientific Computing* (3rd ed.).
+    Cambridge University Press.
+    """
+
+    def objective(x):
+        """F(x) = gammainc(a, x) - p"""
+        return special.gammainc(a, x) - p
+
+    # Initial guess from Numerical Recipes
+    def initial_guess(u_val, a_val):
+        # a = dof/2 for chi-squared
+
+        def large_a_guess():
+            # For a > 1: use Wilson-Hilferty approximation
+            pp = jnp.where(u_val < 0.5, u_val, 1.0 - u_val)
+            t = jnp.sqrt(-2.0 * jnp.log(pp))
+            x = (2.30753 + t * 0.27061) / (
+                1.0 + t * (0.99229 + t * 0.04481)
+            ) - t
+            x = jnp.where(u_val < 0.5, -x, x)
+            return jnp.fmax(
+                1e-3,
+                a_val
+                * (1.0 - 1.0 / (9.0 * a_val) - x / (3.0 * jnp.sqrt(a_val)))
+                ** 3,
+            )
+
+        def small_a_guess():
+            # For a <= 1: use equations (6.2.8) and (6.2.9)
+            t = 1.0 - a_val * (0.253 + a_val * 0.12)
+            return jnp.where(
+                u_val < t,
+                (u_val / t) ** (1.0 / a_val),
+                1.0 - jnp.log(1.0 - (u_val - t) / (1.0 - t)),
+            )
+
+        return jnp.real(
+            jnp.where(a_val > 1.0, large_a_guess(), small_a_guess())
+        )
+
+    # Derivatives for Halley's method
+    f = objective
+    df_dx = jax.grad(objective)
+    d2f_dx2 = jax.grad(df_dx)
+
+    x = initial_guess(p, a)
+
+    # Use while_loop for dynamic convergence
+    def cond_fn(state):
+        x, step, iteration = state
+        # Continue while step is large and we haven't exceeded max iterations
+        return (jnp.abs(step) > EPS * jnp.abs(x)) & (iteration < 12)
+
+    def body_fn(state):
+        x, _, iteration = state
+
+        f_val = f(x)
+        df_val = df_dx(x)
+        d2f_val = d2f_dx2(x)
+
+        # Halley's method: x_{n+1} = x_n - 2*f*f' / (2*f'^2 - f*f'')
+        numerator = 2 * f_val * df_val
+        denominator = 2 * df_val**2 - f_val * d2f_val
+
+        # Avoid division by zero and ensure step is reasonable
+        denominator = jnp.where(
+            jnp.abs(denominator) < TINY,
+            jnp.sign(denominator) * TINY,
+            denominator,
+        )
+
+        step = numerator / denominator
+        x_new = x - step
+
+        # Ensure x stays positive
+        x_new = jnp.fmax(x_new, TINY)
+
+        return (x_new, step, iteration + 1)
+
+    # Initial state: (x, step, iteration)
+    initial_state = (x, jnp.inf, 0)
+    final_state = jax.lax.while_loop(cond_fn, body_fn, initial_state)
+    x = final_state[0]
+
+    return x
+
+
+@gammap_inverse.defjvp
+def gammap_inverse_jvp(primals, tangents):
+    """
+    Custom JVP for gammap_inverse using implicit function theorem.
+
+    For F(x, p) = gammainc(a, x) - p = 0:
+    dx/dp = -∂F/∂p / ∂F/∂x = 1 / (∂/∂x gammainc(a, x))
+    """
+    p, a = primals
+    p_dot, _ = tangents
+
+    # Forward pass
+    x = gammap_inverse(p, a)
+
+    # Compute derivative: dx/dp = 1 / (d/dx gammainc(a, x))
+    def gammainc_x(x_val):
+        return special.gammainc(a, x_val)
+
+    dgammainc_dx = jax.grad(gammainc_x)(x)
+
+    # Avoid division by zero
+    dgammainc_dx = jnp.where(
+        jnp.abs(dgammainc_dx) < TINY,
+        jnp.sign(dgammainc_dx) * TINY,
+        dgammainc_dx,
+    )
+
+    dx_dp = 1.0 / dgammainc_dx
+
+    # For now, ignore a derivatives (could be added if needed)
+    x_dot = dx_dp * p_dot
+
+    return x, x_dot

numerax/stats/__init__.py

@@ -0,0 +1,5 @@
+"""Statistics submodule for numerax."""
+
+from .profile import make_profile_llh
+
+__all__ = ["make_profile_llh"]

numerax/stats/profile.py

@@ -0,0 +1,165 @@
+"""Profile likelihood functions for statistical inference."""
+
+from collections.abc import Callable
+
+import jax
+import jax.numpy as jnp
+import optax
+from optax import lbfgs
+
+
+def make_profile_llh(
+    llh_fn: Callable,
+    is_nuisance: list[bool] | jnp.ndarray,
+    get_initial_nuisance: Callable,
+    tol: float = 1e-6,
+    initial_value: float = 1e-9,
+    initial_diff: float = 1e9,
+) -> Callable:
+    r"""
+    Factory function for creating profile likelihood functions.
+
+    ## Overview
+
+    Profile likelihood is a statistical technique used when dealing with
+    nuisance parameters that are not of primary interest but are necessary
+    for the model. This function creates an optimized profile likelihood
+    that maximizes over nuisance parameters while keeping inference
+    parameters fixed.
+
+    ## Mathematical Background
+
+    Given a likelihood function $L(\boldsymbol{\theta}, \boldsymbol{\lambda})$
+    where $\boldsymbol{\theta}$ are parameters of interest and
+    $\boldsymbol{\lambda}$ are nuisance parameters, the profile likelihood is:
+
+    $$L_p(\boldsymbol{\theta}) = \max_{\boldsymbol{\lambda}}
+    L(\boldsymbol{\theta}, \boldsymbol{\lambda})$$
+
+    In practice, we work with the log-likelihood
+    $\ell(\boldsymbol{\theta}, \boldsymbol{\lambda}) =
+    \log L(\boldsymbol{\theta}, \boldsymbol{\lambda})$:
+
+    $$\ell_p(\boldsymbol{\theta}) = \max_{\boldsymbol{\lambda}}
+    \ell(\boldsymbol{\theta}, \boldsymbol{\lambda})$$
+
+    This function uses L-BFGS optimization to find the maximum likelihood
+    estimates of nuisance parameters for each fixed value of inference
+    parameters.
+
+    ## Args
+
+    - **llh_fn**: Log likelihood function taking (params, *args) and
+      returning scalar log likelihood value
+    - **is_nuisance**: Boolean array where True indicates nuisance
+      parameters and False indicates inference parameters
+    - **get_initial_nuisance**: Function taking (*args) and returning
+      initial values for nuisance parameters
+    - **tol**: Convergence tolerance for the optimization (default: 1e-6)
+    - **initial_value**: Initial objective value for convergence tracking
+      (default: 1e-9)
+    - **initial_diff**: Initial difference for convergence tracking
+      (default: 1e9)
+
+    ## Returns
+
+    Profile likelihood function with signature:
+    `(inference_values, *args) -> (profile_llh_value, optimal_nuisance,
+    convergence_diff, num_iterations)`
+
+    ## Example
+
+    Consider fitting a normal distribution where we want to infer the mean
+    $\mu$ but treat the variance $\sigma^2$ as a nuisance parameter:
+
+    ```python
+    import jax.numpy as jnp
+    import numerax
+
+    # Sample data
+    data = jnp.array([1.2, 0.8, 1.5, 0.9, 1.1, 1.3, 0.7, 1.4])
+
+
+    # Log likelihood for normal distribution
+    def normal_llh(params, data):
+        mu, log_sigma = params  # Use log(sigma) for numerical stability
+        sigma = jnp.exp(log_sigma)
+        return jnp.sum(
+            -0.5 * jnp.log(2 * jnp.pi)
+            - log_sigma
+            - 0.5 * ((data - mu) / sigma) ** 2
+        )
+
+
+    # Profile over log_sigma (nuisance), infer mu
+    is_nuisance = [False, True]  # mu=inference, log_sigma=nuisance
+
+
+    def get_initial_log_sigma(data):
+        # Initialize with log of sample standard deviation
+        return jnp.array([jnp.log(jnp.std(data))])
+
+
+    profile_llh = numerax.stats.make_profile_llh(
+        normal_llh, is_nuisance, get_initial_log_sigma
+    )
+
+    # Evaluate profile likelihood at different mu values
+    mu_test = 1.0
+    llh_val, opt_log_sigma, diff, n_iter = profile_llh(
+        jnp.array([mu_test]), data
+    )
+    ```
+
+    ## Notes
+
+    - The function is JIT-compiled for performance
+    - Uses L-BFGS optimization which is well-suited for smooth likelihood
+      surfaces
+    - Returns convergence information for diagnostics
+    - Handles parameter masking automatically
+    - Consider using log-parameterization for positive parameters
+      (e.g., $\log \sigma$) to improve numerical stability
+    """
+    nuisance_mask = jnp.array(is_nuisance)
+    inference_mask = ~nuisance_mask
+
+    @jax.jit
+    def profile_llh(inference_values, *args):
+        solver = lbfgs()
+        initial_nuisance = get_initial_nuisance(*args)
+        opt_state = solver.init(initial_nuisance)
+
+        def objective(nuisance_params):
+            # Reconstruct full parameter vector
+            full_params = jnp.zeros(len(nuisance_mask))
+            full_params = full_params.at[inference_mask].set(inference_values)
+            full_params = full_params.at[nuisance_mask].set(nuisance_params)
+            return -llh_fn(full_params, *args)
+
+        value_and_grad = optax.value_and_grad_from_state(objective)
+
+        def profile_llh_loopfun(var):
+            params, last_value, opt_state, _, n = var
+            value, grad = value_and_grad(params, state=opt_state)
+            updates, opt_state = solver.update(
+                grad,
+                opt_state,
+                params,
+                value=value,
+                grad=grad,
+                value_fn=objective,
+            )
+            params = optax.apply_updates(params, updates)
+            diff = last_value - value
+            return params, value, opt_state, diff, n + 1
+
+        params, value, opt_state, diff, n = jax.lax.while_loop(
+            lambda var: jnp.abs(var[-2]) > jnp.abs(var[1] * tol),
+            profile_llh_loopfun,
+            (initial_nuisance, initial_value, opt_state, initial_diff, 0),
+        )
+
+        return -value, params, diff, n
+
+    return profile_llh

numerax/utils.py

@@ -0,0 +1,48 @@
+"""Utility functions for the numerax package."""
+
+import functools
+from collections.abc import Callable
+from typing import TypeVar
+
+F = TypeVar("F", bound=Callable)
+
+
+def preserve_metadata(decorator):
+    """
+    Wrapper that ensures a decorator preserves function metadata for
+    documentation tools.
+
+    ## Overview
+
+    This is particularly useful for JAX decorators like `@custom_jvp` that
+    create special objects which may not preserve `__doc__` and other metadata
+    properly for documentation generators like pdoc.
+
+    ## Args
+
+    - **decorator**: The decorator function to wrap
+
+    ## Returns
+
+    A new decorator that preserves metadata
+
+    ## Example
+
+    ```python
+    import jax
+    from numerax.utils import preserve_metadata
+
+    @preserve_metadata(jax.custom_jvp)
+    def my_function(x):
+        \"\"\"This docstring will be preserved for pdoc.\"\"\"
+        return x
+    ```
+    """
+
+    def metadata_preserving_decorator(func: F) -> F:
+        # Apply the original decorator
+        decorated = decorator(func)
+        # Ensure metadata is preserved using functools.wraps pattern
+        return functools.wraps(func)(decorated)
+
+    return metadata_preserving_decorator

numerax-0.1.0.dist-info/METADATA

@@ -0,0 +1,110 @@
+Metadata-Version: 2.4
+Name: numerax
+Version: 0.1.0
+Project-URL: Documentation, https://github.com/juehang/numerax#readme
+Project-URL: Issues, https://github.com/juehang/numerax/issues
+Project-URL: Source, https://github.com/juehang/numerax
+Author: Juehang Qin
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 4 - Beta
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Python: >=3.12
+Requires-Dist: jax
+Requires-Dist: jaxtyping
+Requires-Dist: optax
+Description-Content-Type: text/markdown
+
+# numerax
+
+[![tests](https://github.com/juehang/numerax/actions/workflows/test.yml/badge.svg)](https://github.com/juehang/numerax/actions/workflows/test.yml)
+[![docs](https://github.com/juehang/numerax/actions/workflows/docs.yml/badge.svg)](https://juehang.github.io/numerax/)
+
+Statistical and numerical computation functions for JAX, focusing on tools not available in the main JAX API.
+
+**[📖 Documentation](https://juehang.github.io/numerax/)**
+
+## Installation
+
+```bash
+pip install numerax
+```
+
+## Features
+
+### Special Functions
+
+Inverse regularized incomplete gamma function with differentiability support:
+
+```python
+import jax.numpy as jnp
+import numerax
+
+# Compute gamma quantiles (inverse CDF)
+p = jnp.array([0.1, 0.5, 0.9])  # Probabilities
+a = 2.0  # Shape parameter
+
+x = numerax.special.gammap_inverse(p, a)
+# Returns quantiles where gammainc(a, x) = p
+
+# Fully differentiable with custom JVP
+grad_fn = jax.grad(numerax.special.gammap_inverse)
+dx_dp = grad_fn(0.5, 2.0)  # Gradient with respect to probability
+```
+
+**Key features:**
+- Halley's method for fast convergence
+- Custom JVP implementation for exact gradients
+- Numerical stability with adaptive precision
+- Equivalent to gamma distribution inverse CDF
+
+### Profile Likelihood
+
+Efficient profile likelihood computation for statistical inference with nuisance parameters:
+
+```python
+import jax.numpy as jnp
+import numerax
+
+# Example: Normal distribution with mean inference, variance profiling
+def normal_llh(params, data):
+    mu, log_sigma = params
+    sigma = jnp.exp(log_sigma)
+    return jnp.sum(-0.5 * jnp.log(2 * jnp.pi) - log_sigma 
+                   - 0.5 * ((data - mu) / sigma) ** 2)
+
+# Profile over log_sigma, infer mu
+is_nuisance = [False, True]  # mu=inference, log_sigma=nuisance
+
+def get_initial_log_sigma(data):
+    return jnp.array([jnp.log(jnp.std(data))])
+
+profile_llh = numerax.stats.make_profile_llh(
+    normal_llh, is_nuisance, get_initial_log_sigma
+)
+
+# Evaluate profile likelihood
+data = jnp.array([1.2, 0.8, 1.5, 0.9, 1.1])
+llh_val, opt_nuisance, diff, n_iter = profile_llh(jnp.array([1.0]), data)
+```
+
+**Key features:**
+- JIT-compiled for performance
+- L-BFGS optimization with convergence diagnostics
+- Configurable tolerance and initial values
+- Handles parameter masking automatically
+
+### Utilities
+
+Development utilities for creating JAX functions with custom derivatives while ensuring proper documentation support. Includes decorators for preserving function metadata when using JAX's advanced features.
+
+## Requirements
+
+- Python ≥ 3.12
+- JAX
+- jaxtyping
+- optax

numerax-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+numerax/__init__.py,sha256=y9LIfr6Fo2w7s9y1Y9QmpyNRfHEaGQuEY6OTs4jKbkw,1340
+numerax/utils.py,sha256=oDEjLGKNJv7_lJQ5JVAZ-pXkYnqzrBjYYirLyj1CSOM,1197
+numerax/special/__init__.py,sha256=sG_dsAOoNynwF1DI_Jn_JlYcA3hWDxMNvgSsE7M0pBs,64
+numerax/special/gamma.py,sha256=SdA3Bb_Pe-NHl02rpDGO2D0B436eGJCSLbJ2lWR7pmY,6443
+numerax/stats/__init__.py,sha256=0FyAqMs_xPWXTlQO-LI4y2RZ0wFQNq00uI_8ZhM_XzE,111
+numerax/stats/profile.py,sha256=KfYG0XyRCguiARrTv3C2N-u-p9GMSS9tIPPA5C9YYj4,5540
+numerax-0.1.0.dist-info/METADATA,sha256=BjwCUYfgtWWQZDDuUX_41zDR9PtZVnSxFptiAGgir2Y,3341
+numerax-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+numerax-0.1.0.dist-info/licenses/LICENSE,sha256=q_748ZfuhHCem3AhXv-C9g5u17674OyDlE74DVDu_Ec,1068
+numerax-0.1.0.dist-info/RECORD,,

numerax-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

numerax-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Juehang Qin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.