Getting Started

Installation

CopulAX is available on PyPI:

pip install copulax

To install with documentation build dependencies:

pip install copulax[docs]

Documentation

Requirements

  • Python >= 3.10

  • JAX >= 0.4.38

  • equinox >= 0.13.0

  • optax >= 0.2.4

  • quadax >= 0.2.8

  • interpax >= 0.3.7

  • matplotlib >= 3.9.2

Quick Start

Univariate fitting

import jax.random as jr
from copulax.univariate import univariate_fitter, normal

key = jr.PRNGKey(0)

# Fit a normal distribution
data = jr.normal(key, shape=(500,))
fitted = normal.fit(data)
print(fitted.stats())

# Automatic distribution selection
best_idx, results = univariate_fitter(data)

Multivariate fitting

import jax.random as jr
from copulax.multivariate import mvt_normal

key = jr.PRNGKey(1)
data = jr.normal(key, shape=(500, 3))
fitted = mvt_normal.fit(data)
samples = fitted.rvs(size=50)

Copula fitting

import jax.random as jr
from copulax.copulas import gaussian_copula

key = jr.PRNGKey(2)
data = jr.normal(key, shape=(500, 3))
fitted = gaussian_copula.fit(data)
samples = fitted.rvs(size=50)

Archimedean copulas

from copulax.copulas import clayton_copula
from copulax import get_random_key

key = get_random_key()

# sample from a Clayton copula
params = clayton_copula.example_params(dim=3)
u = clayton_copula.copula_rvs(size=200, params=params, key=key)

# evaluate copula CDF, PDF and log-PDF
cdf = clayton_copula.copula_cdf(u, params=params)
pdf = clayton_copula.copula_pdf(u, params=params)

# fit a copula to uniform data
fitted = clayton_copula.fit_copula(u)

# model selection with AIC / BIC
aic = clayton_copula.aic(u, params=params)
bic = clayton_copula.bic(u, params=params)

Preprocessing — DataScaler

DataScaler is a jittable, autodiff-compatible affine rescaler with support for z-score, min-max, robust, and max-abs methods. It accepts arbitrary n-dimensional input (axis 0 is the sample axis) and is serialisable to .cpx via the same copulax.load entry point used for distributions.

import jax.random as jr
from copulax.multivariate import mvt_normal
from copulax.preprocessing import DataScaler

key = jr.PRNGKey(3)
data = jr.normal(key, shape=(500, 3))

# Fit the scaler and transform in one step
scaler, data_scaled = DataScaler("zscore").fit_transform(data)

# Round-trip
data_back = scaler.inverse_transform(data_scaled)

# Save / load
scaler.save("my_scaler.cpx")
import copulax
loaded_scaler = copulax.load("my_scaler.cpx")

Saving and loading distributions

Fitted distributions can be saved to disk and loaded back in a later session. All distribution types — univariate, multivariate, and copula — are supported. Files use the .cpx format and are cross-platform (Windows, macOS, Linux).

import copulax
from copulax.univariate import normal

# Fit and save
fitted = normal.fit(data)
fitted.save("my_model.cpx")

# Load (same or different session)
loaded = copulax.load("my_model.cpx")
loaded.logpdf(data)  # identical output

The name keyword on load lets you rename the instance on load:

loaded = copulax.load("my_model.cpx", name="production_model")

Copulas (including their fitted marginals) are saved and loaded in exactly the same way:

from copulax.copulas import gaussian_copula

fitted_cop = gaussian_copula.fit(data)
fitted_cop.save("copula_model.cpx")
loaded_cop = copulax.load("copula_model.cpx")

Testing efficiently

The full test suite can take time. During development, run only affected tests and prefer single test functions while iterating.

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

# affected file only
pytest copulax/tests/test_copulas_mv.py -v -m "not slow"

# keep an append-only log while iterating
pytest copulax/tests/test_copulas_mv.py -v -m "not slow" *>&1 `
  | Tee-Object -FilePath copula_test_results.txt -Append