Copulax
JAX based probability modelling
Install / Use
/learn @tfm000/CopulaxREADME
CopulAX is an open-source library for probability distribution fitting, written in JAX with an emphasis on low-dimensional optimization. It is the spiritual successor to SklarPy and provides univariate, multivariate and copula distribution objects with JIT compilation and automatic differentiation support.
This library is designed for use cases ranging from machine learning to finance.
Table of contents
- Documentation
- Installation
- Quick Start
- Low-Dimensional Optimization
- Development Status
- Implemented Distributions
- Testing
- Examples
Documentation
- Read the Docs: https://copulax.readthedocs.io/en/latest/
- API reference: https://copulax.readthedocs.io/en/latest/api/index.html
Installation
CopulAX is available on PyPI and can be installed by running:
pip install copulax
Quick Start
import jax.random as jr
from copulax.univariate import normal, univariate_fitter
from copulax.multivariate import mvt_normal
from copulax.copulas import gaussian_copula
key = jr.PRNGKey(0)
k1, k2, k3 = jr.split(key, 3)
# Univariate fitting
x_uni = jr.normal(k1, shape=(500,))
fitted_uni = normal.fit(x_uni)
best_idx, candidates = univariate_fitter(x_uni)
# Multivariate fitting
x_mvt = jr.normal(k2, shape=(500, 3))
fitted_mvt = mvt_normal.fit(x_mvt)
# Copula fitting
x_cop = jr.normal(k3, shape=(500, 3))
fitted_cop = gaussian_copula.fit(x_cop)
Low-Dimensional Optimization
In many settings, sample sizes are limited. Probabilistic modeling can help generate additional data with similar statistical structure, but multivariate and copula models often require shape/covariance/correlation parameters that grow as O($d^2$). CopulAX reduces this burden where possible by using analytical relationships between these matrices and other parameters. Estimating location and shape robustly outside the optimization loop can materially reduce the number of numerically optimized parameters.
Development Status
CopulAX is under active development. Current coverage includes:
- Continuous univariate distributions.
- Multivariate normal-mixture families.
- Elliptical and Archimedean copulas.
- JIT/autodiff-compatible fitting workflows and utility functions.
Near-term roadmap:
- Additional univariate distributions (including discrete support).
- Additional multivariate and copula families.
- Broader CDF coverage for multivariate and elliptical copula objects.
- Empirical distribution support with multiple fitting methods.
Implemented Distributions
A list of all implemented distributions can be found here:
- <a href="https://github.com/tfm000/copulax/blob/main/copulax/univariate/README.md">Univariate implemented distributions</a>
- <a href="https://github.com/tfm000/copulax/blob/main/copulax/multivariate/README.md">Multivariate implemented distributions</a>
- <a href="https://github.com/tfm000/copulax/blob/main/copulax/copulas/README.md">Copula implemented distributions</a>
Testing
Tests are comprehensive, but some suites can be slow. A practical workflow is:
- Run only affected tests first.
- Run individual test functions while iterating.
- Keep a timestamped test log while debugging.
# Specific function
pytest copulax/tests/copulas/test_copulas.py::TestFitting::test_fit -v
# Affected file/family only
pytest copulax/tests/copulas/test_copulas.py -v
# Append test output to a running log (PowerShell)
pytest copulax/tests/copulas/test_copulas.py::TestFitting::test_fit -v *>&1 `
| Tee-Object -FilePath copula_test_results.txt -Append
Examples
We have provided <a href="https://github.com/tfm000/copulax/tree/main/examples">jupyter notebooks</a> containing example code for using univariate, multivariate and copula distribution objects.
Related Skills
claude-opus-4-5-migration
82.7kMigrate prompts and code from Claude Sonnet 4.0, Sonnet 4.5, or Opus 4.1 to Opus 4.5
model-usage
335.8kUse CodexBar CLI local cost usage to summarize per-model usage for Codex or Claude, including the current (most recent) model or a full model breakdown. Trigger when asked for model-level usage/cost data from codexbar, or when you need a scriptable per-model summary from codexbar cost JSON.
feishu-drive
335.8k|
things-mac
335.8kManage Things 3 via the `things` CLI on macOS (add/update projects+todos via URL scheme; read/search/list from the local Things database)
