CopulAX is an open-source library for probability distribution fitting, written in JAX with an emphasis on speed and 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

• Table of contents
• Documentation
• Installation
• Quick Start
• Saving and Loading
• 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

A pinned uv.lock is also committed for reproducible development environments — clone the repo and run uv sync to install the exact dependency set used by maintainers.

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
from copulax.preprocessing import DataScaler

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)

# Preprocessing — jittable z-score / min-max / robust / max-abs scaler
scaler, x_scaled = DataScaler("zscore").fit_transform(x_mvt)
x_original = scaler.inverse_transform(x_scaled)

Saving and Loading

Fitted distributions and fitted DataScaler preprocessors can be saved to disk and loaded back in a later session. All distribution types (univariate, multivariate, and copula) are supported, as are DataScaler instances. Files use the .cpx format and are cross-platform.

import copulax

# Save a fitted distribution
fitted_uni.save("my_model.cpx")

# Load it back (in the same or a different session)
loaded = copulax.load("my_model.cpx")
loaded.logpdf(x_uni)  # identical to fitted_uni.logpdf(x_uni)

# DataScaler uses the same save/load entry points
scaler.save("my_scaler.cpx")
loaded_scaler = copulax.load("my_scaler.cpx")

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
• Preprocessing utilities

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.

# Default iteration: exclude slow tests (what CI runs)
pytest copulax/tests/ -v -m "not slow"

# Specific test file (e.g. multivariate copulas)
pytest copulax/tests/test_copulas_mv.py -v -m "not slow"

# Specific test function
pytest copulax/tests/test_copulas_mv.py::TestCopulaFitting::test_fit_returns_valid_params -v

# Append test output to a running log (PowerShell)
pytest copulax/tests/test_copulas_mv.py -v -m "not slow" *>&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.