normix

normix#

normix is a JAX package for Generalized Hyperbolic distributions and their relatives, implemented as exponential families. It is built on Equinox with float64 precision throughout. Log-densities, moments, and the JAX log-partition triad are JIT-compiled and differentiable. log_prob takes one observation. vmap applies it to a batch. The supported surface is the capability matrix.

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from normix import GeneralizedHyperbolic

# Fit a Generalized Hyperbolic distribution to data via EM
key = jax.random.PRNGKey(0)
X = jax.random.normal(key, (1000, 3))
model = GeneralizedHyperbolic.default_init(X)
result = model.fit(X, max_iter=100)

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

Key features#

Exponential-family structure — classical parameters and natural \(\theta\) convert exactly. from_expectation inverts \(\eta = \nabla\psi(\theta)\) in closed form or to a solver tolerance.
The full GH family — Gamma, Inverse Gamma, Inverse Gaussian, GIG, Variance Gamma, Normal-Inverse Gamma, Normal-Inverse Gaussian, and Generalized Hyperbolic, plus univariate and factor variants.
EM fitting — batch and incremental EM. Bessel values come from the log_kv kernel in JAX or NumPy. SciPy supplies the CPU optimizers. CPU Gamma uses scipy.special digamma and polygamma. scipy.special.kve is a test oracle.
Closed-form divergences — squared Hellinger and KL with no Monte Carlo.
Finance toolkit — portfolio projection and differentiable CVaR.
Immutable — all distributions are equinox.Module pytrees.