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:

• Univariate implemented distributions
• Multivariate implemented distributions
• Copula implemented distributions

Testing

Tests are comprehensive, but some suites can be slow. A practical workflow is:

1. Run only affected tests first.
2. Run individual test functions while iterating.
3. 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 jupyter notebooks containing example code for using univariate, multivariate and copula distribution objects.