numerax 0.1.0__tar.gz

numerax-0.1.0/.claude/settings.local.json

@@ -0,0 +1,13 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(find:*)",
+      "Bash(ruff format:*)",
+      "mcp__vscode-tools__get_symbol_definition_code",
+      "mcp__vscode-tools__search_symbols_code",
+      "Bash(hatch run lint:*)",
+      "mcp__vscode-tools__get_document_symbols_code"
+    ],
+    "deny": []
+  }
+}

numerax-0.1.0/.github/workflows/docs.yml

@@ -0,0 +1,59 @@
+name: documentation
+
+# build the documentation whenever there are new commits on main
+on:
+  push:
+    branches:
+      - main
+    # Alternative: only build for tags.
+    # tags:
+    #   - '*'
+
+# security: restrict permissions for CI jobs.
+permissions:
+  contents: read
+
+jobs:
+  # Build the documentation and upload the static HTML files as an artifact.
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      
+      # Extract Python version from pyproject.toml
+      - name: Get Python version from pyproject.toml
+        id: python-version
+        run: |
+          python_version=$(grep -E "requires-python.*=" pyproject.toml | sed -E 's/.*">=([0-9.]+)".*/\1/')
+          echo "version=$python_version" >> $GITHUB_OUTPUT
+      
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ steps.python-version.outputs.version }}
+
+      # Install Hatch
+      - run: pip install hatch
+      
+      # Build documentation with MkDocs using dev environment
+      - run: hatch run mkdocs build
+
+      - uses: actions/upload-pages-artifact@v3
+        with:
+          path: site/
+
+  # Deploy the artifact to GitHub pages.
+  # This is a separate job so that only actions/deploy-pages has the necessary permissions.
+  deploy:
+    needs: build
+    runs-on: ubuntu-latest
+    permissions:
+      pages: write
+      id-token: write
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - id: deployment
+        uses: actions/deploy-pages@v4

numerax-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,84 @@
+name: publish
+
+# Publish to PyPI when a new release is created
+on:
+  release:
+    types: [published]
+
+# Security: restrict permissions for CI jobs
+permissions:
+  contents: read
+
+jobs:
+  # Build and test before publishing
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+      
+      # Install hatch for environment management
+      - name: Install Hatch
+        uses: pypa/hatch@install
+      
+      # Run code quality checks using hatch
+      - name: Run linting and formatting checks
+        run: hatch run lint:check
+      
+      # Run tests using hatch
+      - name: Run tests
+        run: hatch run test
+
+  # Build the package
+  build:
+    needs: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+      
+      # Install hatch for building
+      - name: Install Hatch
+        uses: pypa/hatch@install
+      
+      # Build the package
+      - name: Build package
+        run: hatch build
+      
+      # Upload build artifacts
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist-artifacts
+          path: dist/
+
+  # Publish to PyPI
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write  # Required for trusted publishing
+    environment:
+      name: pypi
+      url: https://pypi.org/p/numerax
+    steps:
+      - name: Download build artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: dist-artifacts
+          path: dist/
+      
+      # Publish to PyPI using trusted publishing
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

numerax-0.1.0/.github/workflows/test.yml

@@ -0,0 +1,39 @@
+name: tests
+
+# Run tests on pushes to main and pull requests
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+# Security: restrict permissions for CI jobs
+permissions:
+  contents: read
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+      
+      # Install hatch for environment management
+      - name: Install Hatch
+        uses: pypa/hatch@install
+      
+      # Run code quality checks using hatch
+      - name: Run linting and formatting checks
+        run: hatch run lint:check
+      
+      # Run tests using hatch
+      - name: Run tests
+        run: hatch run test

numerax-0.1.0/.gitignore

@@ -0,0 +1,5 @@
+.ruff_cache/*
+CLAUDE.md
+*.pyc
+__pycache__/
+*.pyo

numerax-0.1.0/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.

numerax-0.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,89 @@
+# 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/docs/api/index.md

@@ -0,0 +1,9 @@
+# API Reference
+
+Comprehensive API documentation for all Numerax modules.
+
+## Modules
+
+- **[Special Functions](special.md)** - Mathematical special functions with custom derivatives
+- **[Statistics](stats.md)** - Advanced statistical computation tools  
+- **[Utilities](utils.md)** - Development utilities for JAX functions

numerax-0.1.0/docs/api/special.md

@@ -0,0 +1,6 @@
+# Special Functions
+
+::: numerax.special
+    options:
+      show_source: false
+      heading_level: 2

numerax-0.1.0/docs/api/stats.md

@@ -0,0 +1,6 @@
+# Statistics
+
+::: numerax.stats
+    options:
+      show_source: false
+      heading_level: 2

numerax-0.1.0/docs/api/utils.md

@@ -0,0 +1,6 @@
+# Utilities
+
+::: numerax.utils
+    options:
+      show_source: false
+      heading_level: 2

numerax-0.1.0/docs/index.md

@@ -0,0 +1,6 @@
+# Numerax
+
+::: numerax
+    options:
+      show_source: false
+      heading_level: 1

numerax-0.1.0/docs/stylesheets/extra.css

@@ -0,0 +1,18 @@
+/* Custom styles for MkDocs Material theme */
+
+/* Improve math rendering */
+.arithmatex {
+  padding: 0.5em;
+  margin: 1em 0;
+}
+
+/* Better code block styling */
+.highlight pre {
+  padding: 1em;
+  overflow-x: auto;
+}
+
+/* Improve navigation */
+.md-nav__item .md-nav__link--active {
+  font-weight: bold;
+}

numerax-0.1.0/mkdocs.yml

@@ -0,0 +1,73 @@
+site_name: Numerax
+site_description: Statistical and numerical computation functions for JAX
+site_url: https://juehang.github.io/numerax/
+repo_url: https://github.com/juehang/numerax
+repo_name: juehang/numerax
+
+theme:
+  name: material
+  features:
+    - navigation.tabs
+    - navigation.sections
+    - navigation.expand
+    - navigation.top
+    - search.highlight
+    - search.share
+    - content.code.copy
+  palette:
+    - scheme: default
+      primary: blue
+      accent: blue
+      toggle:
+        icon: material/brightness-7
+        name: Switch to dark mode
+    - scheme: slate
+      primary: blue
+      accent: blue
+      toggle:
+        icon: material/brightness-4
+        name: Switch to light mode
+
+plugins:
+  - search
+  - mkdocstrings:
+      handlers:
+        python:
+          options:
+            docstring_style: google
+            show_source: false
+            show_root_heading: true
+            show_root_toc_entry: false
+            merge_init_into_class: true
+            show_signature_annotations: true
+            separate_signature: true
+
+markdown_extensions:
+  - pymdownx.arithmatex:
+      generic: true
+  - pymdownx.highlight:
+      anchor_linenums: true
+  - pymdownx.inlinehilite
+  - pymdownx.snippets
+  - pymdownx.superfences
+  - admonition
+  - pymdownx.details
+  - pymdownx.tabbed:
+      alternate_style: true
+  - toc:
+      permalink: true
+
+extra_javascript:
+  - https://polyfill.io/v3/polyfill.min.js?features=es6
+  - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js
+
+extra_css:
+  - stylesheets/extra.css
+
+nav:
+  - Home: index.md
+  - API Reference:
+    - Overview: api/index.md
+    - Special Functions: api/special.md
+    - Statistics: api/stats.md
+    - Utilities: api/utils.md

numerax-0.1.0/pyproject.toml

@@ -0,0 +1,133 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "numerax"
+dynamic = ["version"]
+description = ""
+readme = "README.md"
+requires-python = ">=3.12"
+license = "MIT"
+keywords = []
+authors = [
+  { name = "Juehang Qin"},
+]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: Implementation :: CPython",
+  "Programming Language :: Python :: Implementation :: PyPy",
+]
+dependencies = [
+  "jax",
+  "jaxtyping",
+  "optax"
+]
+
+[project.urls]
+Documentation = "https://github.com/juehang/numerax#readme"
+Issues = "https://github.com/juehang/numerax/issues"
+Source = "https://github.com/juehang/numerax"
+
+[tool.hatch.version]
+path = "src/numerax/__init__.py"
+
+[tool.hatch.envs.default]
+dependencies = [
+  "coverage[toml]>=6.5",
+  "pytest",
+  "ruff",
+  "mkdocs-material",
+  "mkdocstrings[python]"
+]
+[tool.hatch.envs.default.scripts]
+test = "pytest {args:tests}"
+test-cov = "coverage run -m pytest {args:tests}"
+cov-report = [
+  "- coverage combine",
+  "coverage report",
+]
+cov = [
+  "test-cov",
+  "cov-report",
+]
+
+[[tool.hatch.envs.all.matrix]]
+python = ["3.12"]
+
+[tool.hatch.envs.lint]
+detached = true
+dependencies = [
+  "ruff>=0.0.243",
+]
+[tool.hatch.envs.lint.scripts]
+check = [
+  "ruff check {args:.}",
+  "ruff format --check --diff {args:.}",
+]
+fmt = [
+  "ruff format {args:.}",
+  "ruff check --fix {args:.}",
+  "check",
+]
+
+[tool.coverage.run]
+source_pkgs = ["numerax", "tests"]
+branch = true
+parallel = true
+omit = [
+  "src/numerax/__about__.py",
+]
+
+[tool.coverage.paths]
+numerax = ["src/numerax", "*/numerax/src/numerax"]
+tests = ["tests", "*/numerax/tests"]
+
+[tool.coverage.report]
+exclude_lines = [
+  "no cov",
+  "if __name__ == .__main__.:",
+  "if TYPE_CHECKING:",
+]
+
+[tool.ruff]
+line-length = 79
+target-version = "py312"
+extend-exclude = ["old_code"]
+
+[tool.ruff.lint]
+select = [
+    "E",  # pycodestyle errors
+    "F",  # pyflakes
+    "UP", # pyupgrade
+    "B",  # flake8-bugbear
+    "SIM", # flake8-simplify
+    "I",  # isort
+    "ASYNC", # flake8-async
+    "S",  # flake8-bandit
+    "ARG", # flake8-unused-arguments
+    "Q",  # flake8-quotes
+    "SIM", # flake8-simplify
+    "NPY", # numpy-specific rules
+    "PD",  # pandas-specific rules
+    "N",  # pep8-naming
+    "W",  # warning
+    "PLC", # pylint convention
+    "PLE", # pylint error
+    "PLW", # pylint warning
+    "RUF", # ruff-specific rules
+]
+ignore = ["COM812"]  # Ignore trailing comma rule that conflicts with formatter
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**" = ["S101"]  # Allow assert statements in tests
+
+# Enable Ruff's formatter
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+line-ending = "auto"
+docstring-code-format = true

numerax-0.1.0/src/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-0.1.0/src/numerax/special/__init__.py

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

numerax-0.1.0/src/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-0.1.0/src/numerax/stats/__init__.py

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

numerax-0.1.0/src/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-0.1.0/src/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/tests/__init__.py

@@ -0,0 +1 @@
+"""Test configuration for numerax."""

numerax-0.1.0/tests/test_gammap_inverse.py

@@ -0,0 +1,50 @@
+"""
+Test suite for gamma special functions.
+"""
+
+import jax
+import jax.numpy as jnp
+import pytest
+from jax.scipy import special
+
+from numerax.special import gammap_inverse
+
+
+@pytest.mark.parametrize(
+    "p,a",
+    [
+        (0.1, 0.5),
+        (0.3, 1.0),
+        (0.5, 1.5),
+        (0.8, 2.0),
+        (0.95, 3.0),
+    ],
+)
+def test_gamma_inverse_correctness(p, a):
+    """Test that gammap_inverse correctly inverts gammainc."""
+    result = special.gammainc(a, gammap_inverse(p, a))
+    assert jnp.abs(result - p) < 1e-6
+
+
+@pytest.mark.parametrize(
+    "p,a",
+    [
+        (0.2, 0.5),
+        (0.4, 1.0),
+        (0.6, 1.5),
+        (0.8, 2.0),
+        (0.9, 2.5),
+    ],
+)
+def test_gamma_inverse_gradients(p, a):
+    """Test that gradients of gammap_inverse are computed correctly."""
+    # Compute x such that gammainc(a, x) = p
+    x = gammap_inverse(p, a)
+
+    # Manual gradient: 1 / d/dx gammainc(a, x)
+    manual_grad = 1 / jax.grad(lambda x_val: special.gammainc(a, x_val))(x)
+
+    # Autodiff gradient
+    auto_grad = jax.grad(gammap_inverse)(p, a)
+
+    assert jnp.abs(manual_grad - auto_grad) < 1e-6

numerax-0.1.0/tests/test_profile_gaussian.py

@@ -0,0 +1,116 @@
+"""Test profile likelihood with Gaussian data and analytical validation."""
+
+import jax
+import jax.numpy as jnp
+
+from numerax.stats import make_profile_llh
+
+
+def test_gaussian_profile_analytical_validation():
+    r"""Test profile likelihood against analytical Gaussian MLE solution.
+
+    For Gaussian data with fixed $\mu$, the MLE of $\sigma$ is:
+    $\hat{\sigma}(\mu) = \sqrt{\Sigma(x_i-\mu)^2/N}$
+    This test validates that our profile likelihood optimization finds the same
+    result as this analytical solution.
+    """
+    # Set random seed for reproducible test
+    key = jax.random.PRNGKey(42)
+
+    # Generate synthetic Gaussian data
+    true_mu = 2.0
+    true_sigma = 1.5
+    n_samples = 50
+
+    data = jax.random.normal(key, (n_samples,)) * true_sigma + true_mu
+
+    # Define normal log likelihood function
+    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
+        )
+
+    # Set up profile likelihood: mu=inference, log_sigma=nuisance
+    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))])
+
+    # Create profile likelihood function
+    profile_llh = make_profile_llh(
+        normal_llh, is_nuisance, get_initial_log_sigma
+    )
+
+    # Test at a fixed mu value
+    test_mu = 1.8
+    llh_val, opt_log_sigma, diff, n_iter = profile_llh(
+        jnp.array([test_mu]), data
+    )
+
+    # Extract optimized sigma
+    sigma_opt = jnp.exp(opt_log_sigma[0])
+
+    # Calculate analytical solution: sigma_hat(mu) = sqrt(sum((x_i-mu)^2)/N)
+    sigma_analytical = jnp.sqrt(jnp.mean((data - test_mu) ** 2))
+
+    # Validate convergence
+    assert n_iter > 0, "Optimization should have run at least one iteration"
+    assert jnp.abs(diff) <= jnp.abs(llh_val * 1e-6), (
+        f"Should have converged, but diff={diff}"
+    )
+
+    # Validate against analytical solution
+    assert jnp.allclose(sigma_opt, sigma_analytical, atol=1e-3), (
+        f"Optimized sigma={sigma_opt:.6f} should match "
+        f"analytical sigma={sigma_analytical:.6f}"
+    )
+
+    # Validate likelihood value is finite and reasonable
+    assert jnp.isfinite(llh_val), "Profile likelihood value should be finite"
+    assert llh_val < 0, "Log likelihood should be negative"
+
+
+def test_gaussian_profile_at_true_mean():
+    """Test profile likelihood when mu equals the true data-generating mean."""
+    key = jax.random.PRNGKey(123)
+
+    # Generate data
+    true_mu = 3.0
+    true_sigma = 0.8
+    n_samples = 30
+
+    data = jax.random.normal(key, (n_samples,)) * true_sigma + true_mu
+
+    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
+        )
+
+    is_nuisance = [False, True]
+
+    def get_initial_log_sigma(data):
+        return jnp.array([jnp.log(jnp.std(data))])
+
+    profile_llh = make_profile_llh(
+        normal_llh, is_nuisance, get_initial_log_sigma
+    )
+
+    # Test at sample mean (should be close to optimal)
+    sample_mean = jnp.mean(data)
+    llh_val, opt_log_sigma, _, _ = profile_llh(jnp.array([sample_mean]), data)
+
+    sigma_opt = jnp.exp(opt_log_sigma[0])
+    sigma_analytical = jnp.sqrt(jnp.mean((data - sample_mean) ** 2))
+
+    # Should converge quickly and match analytical solution well
+    assert jnp.allclose(sigma_opt, sigma_analytical, atol=1e-4)
+    assert jnp.isfinite(llh_val)