normix 0.2.0

JAX package for Generalized Hyperbolic distributions as exponential families

normix

JAX package for Generalized Hyperbolic distributions as exponential families.

Built on Equinox with Float64 precision throughout.

Installation

pip install normix

Install optional plotting helpers with:

pip install "normix[plotting]"

For local development:

uv sync
pip install -e .

Quick Start

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

from normix import GeneralizedHyperbolic
from normix.fitting.em import BatchEMFitter

# Fit GH distribution to data via EM
key = jax.random.PRNGKey(0)
X = jax.random.normal(key, (1000, 3))

model = GeneralizedHyperbolic.from_classical(
    mu=jnp.zeros(3), gamma=jnp.zeros(3),
    sigma=jnp.eye(3), p=-0.5, a=2.0, b=1.0,
)
result = BatchEMFitter(max_iter=100).fit(model, X)

# Evaluate log-density (batched via vmap)
log_p = jax.vmap(result.model.log_prob)(X)   # shape (1000,)

Distributions

Univariate (exponential family)

Class	Parameters	Description
Gamma	alpha, beta	Shape α > 0, rate β > 0
InverseGamma	alpha, beta	Shape α > 0, rate β > 0
InverseGaussian	mu, lam	Mean μ > 0, shape λ > 0
GIG / GeneralizedInverseGaussian	p, a, b	Generalized Inverse Gaussian

Multivariate

Class	Parameters	Description
MultivariateNormal	mu, L_Sigma	Mean μ, Cholesky L_Sigma of Σ

Normal Variance-Mean Mixtures (marginal)

Class	Subordinator	Parameters
VarianceGamma	Gamma	mu, gamma, L_Sigma, alpha, beta
NormalInverseGamma	InverseGamma	mu, gamma, L_Sigma, alpha, beta
NormalInverseGaussian	InverseGaussian	mu, gamma, L_Sigma, mu_ig, lam
GeneralizedHyperbolic	GIG	mu, gamma, L_Sigma, p, a, b

Joint distributions

The Joint* classes (e.g. JointGeneralizedHyperbolic) model the full joint $f(x,y)$ where Y is the mixing variable. They are exponential families and are used internally for the EM E-step.

Exponential Family API

All univariate and joint distributions subclass ExponentialFamily(eqx.Module):

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

from normix import Gamma

X = jnp.array([1.0, 1.5, 2.0, 2.5])
dist = Gamma(alpha=jnp.array(2.0), beta=jnp.array(1.0))

# Log-density (single observation)
dist.log_prob(jnp.array(1.5))

# Three parametrizations
theta = dist.natural_params()       # natural parameters θ
eta   = dist.expectation_params()   # expectation parameters η = E[t(X)]
I     = dist.fisher_information()   # Fisher information I(θ) = ∇²ψ(θ)

# Constructors
dist2 = Gamma.from_natural(theta)
dist3 = Gamma.from_expectation(eta)
dist4 = Gamma.fit_mle(X)           # η̂ = mean t(xᵢ)

EM Algorithm

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

from normix import GeneralizedHyperbolic
from normix.fitting.em import BatchEMFitter

d = 3
X = ...  # (n, d) data array

# Initialise from classical parameters
model = GeneralizedHyperbolic.from_classical(
    mu=jnp.zeros(d), gamma=jnp.zeros(d), sigma=jnp.eye(d),
    p=-0.5, a=2.0, b=1.0,
)

# Fit with hybrid CPU/JAX backend for maximum speed
fitter = BatchEMFitter(max_iter=200, tol=1e-6,
                       e_step_backend='cpu', m_step_backend='cpu')
result = fitter.fit(model, X)
fitted = result.model

Bessel Functions

import jax

jax.config.update("jax_enable_x64", True)

from normix import log_kv        # or: from normix.utils.bessel import log_kv

# JIT-able, differentiable (backend='jax', default)
log_kv(v=0.5, z=2.0)

# Fast CPU path for EM hot path (not JIT-able)
log_kv(v=0.5, z=2.0, backend='cpu')

Package Layout

normix/
├── exponential_family.py         # ExponentialFamily base class
├── distributions/                # All distribution implementations
│   ├── gamma.py
│   ├── inverse_gamma.py
│   ├── inverse_gaussian.py
│   ├── generalized_inverse_gaussian.py
│   ├── normal.py
│   ├── variance_gamma.py
│   ├── normal_inverse_gamma.py
│   ├── normal_inverse_gaussian.py
│   └── generalized_hyperbolic.py
├── mixtures/                     # Joint and marginal base classes
├── fitting/em.py                 # BatchEMFitter, EMResult
└── utils/
    ├── bessel.py                 # log_kv with custom JVP
    ├── constants.py              # Shared numerical constants
    ├── plotting.py               # Notebook helpers
    └── validation.py             # EM validation helpers

Development

uv run pytest tests/              # run tests
uv run jupyter lab                # notebooks
make -C docs html                 # build docs

References

• Barndorff-Nielsen, O. E. (1977). Exponentially decreasing distributions for the logarithm of particle size.
• Eberlein, E., & Keller, U. (1995). Hyperbolic distributions in finance.

License

MIT — see LICENSE.