API Reference

Reference pages, one per subpackage:

• Distributions — the nine GH-family distributions (univariate, multivariate, and their joint/marginal/univariate-wrapper/factor variants)

• Mixtures — the JointNormalMixture / NormalMixture / FactorNormalMixture base classes that every mixture distribution builds on

• Fitting — EM fitters, solvers, and the incremental-EM \(\eta\)-update machinery

• Finance — portfolio projection, risk measures, mean-risk optimization, and transaction-cost QPs

• Utilities — Bessel functions, constants, sampling, and plotting helpers

Base Classes

class normix.exponential_family.ExponentialFamily

Bases: Module

Abstract base class for exponential family distributions.

Concrete subclasses must implement:

_log_partition_from_theta, natural_params, sufficient_statistics, log_base_measure

abstractmethod natural_params()

\(\theta\) from stored classical parameters.

Return type:

Array

abstractmethod static sufficient_statistics(x)

\(t(x)\) for a single unbatched observation.

Parameters:

x (Array)

Return type:

Array

abstractmethod static log_base_measure(x)

\(\log h(x)\) for a single unbatched observation.

Parameters:

x (Array)

Return type:

Array

log_partition()

\(\psi(\theta)\) at current parameters.

Return type:

Array

expectation_params(backend='jax')

\(\eta = \nabla\psi(\theta)\).

Parameters:

backend (str) – 'jax' (default, JIT-able) or 'cpu' (numpy/scipy).

Return type:

Array

fisher_information(backend='jax')

\(I(\theta) = \nabla^2\psi(\theta)\).

Parameters:

backend (str) – 'jax' (default, JIT-able) or 'cpu' (numpy/scipy).

Return type:

Array

log_prob(x)

\(\log p(x\mid\theta) = \log h(x) + \theta^\top t(x) - \psi(\theta)\), single observation.

Parameters:

x (Array)

Return type:

Array

pdf(x)

p(x|θ), single observation. Batch via jax.vmap.

Parameters:

x (Array)

Return type:

Array

log_density_power(alpha)

Log density-power integral \(R(\alpha) = \log\int p(x)^\alpha\,\mu(dx)\).

For an exponential family whose base measure is constant on the support,

\[R(\alpha) = (\alpha - 1)\,b_0 + \psi(\alpha\theta) - \alpha\,\psi(\theta), \qquad b_0 = \log h,\]

because \(p_\theta(x)^\alpha \propto \exp\{\alpha\theta^\top t(x)\}\) stays in the family with natural parameter \(\alpha\theta\). This is the cumulant generator of the information content \(-\log p(X)\): entropy is \(H = -R'(1)\), varentropy is \(V_H = R''(1)\), and the Rényi entropy of order \(\alpha\) is \(R(\alpha)/(1-\alpha)\).

Parameters:

alpha (Array)

Return type:

Array

entropy()

Differential entropy \(H = \mathbb{E}[-\log p(X)] = -R'(1)\).

Obtained by differentiating log_density_power(); for constant base measure this equals \(\psi(\theta) - \theta^\top\eta - b_0\).

Return type:

Array

varentropy()

Varentropy \(V_H = \mathrm{Var}[-\log p(X)] = R''(1)\).

Obtained by differentiating log_density_power() twice. For constant base measure this is the Fisher quadratic form \(\theta^\top I(\theta)\,\theta\). The autodiff route is used deliberately: it flows through the accurate log_kv custom JVP.

Return type:

Array

renyi(alpha)

Rényi entropy \(H_\alpha = (1-\alpha)^{-1}\log\int p(x)^\alpha\,\mu(dx)\).

Defined for order \(\alpha > 0\). At \(\alpha = 1\) the value is the removable-singularity limit, the Shannon entropy(); the first-order expansion \(H_\alpha = H - \tfrac{1}{2}V_H(\alpha-1) + \mathcal{O}((\alpha-1)^2)\) links it to the varentropy.

A jax.custom_jvp() supplies the limiting derivative \(H_\alpha'(1) = -V_H/2\). The primal uses jax.lax.cond() so the varentropy is not evaluated for \(\alpha\) away from 1 (a bare jnp.where() singularity guard both zeroes the gradient at \(\alpha = 1\) and forces the expensive unused branch).

Parameters:

alpha (Array)

Return type:

Array

mean()

E[X]. Subclasses should override with analytical formulas.

Return type:

Array

var()

Var[X]. Subclasses should override with analytical formulas.

Return type:

Array

std()

\(\mathrm{Std}[X] = \sqrt{\mathrm{Var}[X]}\).

Return type:

Array

cdf(x)

CDF F(x). Subclasses should override with analytical formulas.

Parameters:

x (Array)

Return type:

Array

rvs(n, seed=42)

Sample n observations via JAX PRNG (JIT-able).

Parameters:

• n (int)

• seed (int)

Return type:

Array

squared_hellinger(other)

Squared Hellinger distance \(H^2(p, q)\).

Both operands are lifted into their divergence gauge before evaluating the general exponential-family formula via the gauge \(\psi\).

Parameters:

other (ExponentialFamily)

Return type:

Array

kl_divergence(other)

KL divergence \(D_{\mathrm{KL}}(\mathrm{self} \| \mathrm{other})\).

Uses the gauge \(\psi\) after lifting both operands, with \(\eta_p\) from the source family’s closed-form tail split (_divergence_eta()) rather than \(\nabla\psi\) at a boundary \(\theta\).

Parameters:

other (ExponentialFamily)

Return type:

Array

classmethod from_natural(theta)

Construct from natural parameters θ. Subclasses must override.

Parameters:

theta (Array)

Return type:

ExponentialFamily

classmethod bregman_divergence(theta, eta)

Bregman divergence \(\psi(\theta) - \theta\cdot\eta\) (conjugate dual).

Minimising over \(\theta\) yields \(\nabla\psi(\theta^*) = \eta\), i.e. the natural parameters corresponding to expectation parameters \(\eta\).

Parameters:

• theta (Array)

• eta (Array)

Return type:

Array

classmethod from_expectation(eta, *, theta0=None, maxiter=500, tol=1e-10, backend='jax', method='lbfgs', verbose=0)

Construct from expectation parameters \(\eta\) by solving \(\nabla\psi(\theta) = \eta\).

Minimises the Bregman divergence \(\psi(\theta) - \theta\cdot\eta\) via solve_bregman. Subclasses can override for closed-form inverses.

Parameters:

• backend (str) – 'jax' (default, JIT-able) or 'cpu' (scipy, more robust).

• method (str) – 'lbfgs' (default), 'bfgs', or 'newton'.

• verbose (int) – 0 = silent, >= 1 = print solver summary.

• eta (Array)

• theta0 (Array | None)

• maxiter (int)

• tol (float)

Return type:

ExponentialFamily

classmethod fit_mle(X, *, theta0=None, maxiter=500, tol=1e-10, verbose=0)

MLE via exponential family identity: \(\hat\eta = \frac{1}{n}\sum_i t(x_i)\).

Batches over X using jax.vmap, then calls from_expectation(\(\hat\eta\)).

Parameters:

• X (jax.Array) – (n, ...) array of observations.

• theta0 (jax.Array, optional) – Initial natural parameters \(\theta_0\) for the \(\eta\to\theta\) solver.

• maxiter (int) – Maximum iterations for the \(\eta\to\theta\) solver.

• tol (float) – Convergence tolerance for the \(\eta\to\theta\) solver.

• verbose (int) – 0 = silent, >= 1 = print solver summary.

Return type:

ExponentialFamily

fit(X, *, maxiter=500, tol=1e-10, verbose=0, **kwargs)

Fit using self as initialization (warm start).

Computes \(\hat\eta = \frac{1}{n}\sum_i t(x_i)\) and solves from_expectation(\(\hat\eta\)) using self.natural_params() as the initial \(\theta_0\).

Parameters:

• X (jax.Array) – (n, ...) array of observations.

• maxiter (int) – Maximum iterations for the \(\eta\to\theta\) solver.

• tol (float) – Convergence tolerance for the \(\eta\to\theta\) solver.

• verbose (int) – 0 = silent, >= 1 = print solver summary.

Return type:

ExponentialFamily

classmethod default_init(X)

Moment-based initialisation from data.

Computes \(\hat\eta = \frac{1}{n}\sum_i t(x_i)\) and inverts to get an initial model. For distributions with closed-form from_expectation (Gamma, InverseGamma, InverseGaussian), this gives the MLE directly.

Parameters:

X (Array)

Return type:

ExponentialFamily

Divergences

Statistical divergences for exponential family distributions.

Tier 1 — functional core (pure JAX, maximum composability):

\[H^2(p, q) = 1 - \exp\!\left(\psi\!\left(\frac{\theta_p + \theta_q}{2}\right) - \frac{\psi(\theta_p) + \psi(\theta_q)}{2}\right)\]

\[D_{\mathrm{KL}}(p \| q) = \psi(\theta_q) - \psi(\theta_p) - (\theta_q - \theta_p)^\top \nabla\psi(\theta_p)\]

kl_divergence_from_eta is the tail-split sibling used when a boundary moment may be infinite (\(\eta = \eta_{\mathrm{fin}} + m\,v\)).

Tier 3 — module convenience (delegates to Tier 2 instance methods):

squared_hellinger(p, q) and kl_divergence(p, q) accept ExponentialFamily or NormalMixture objects. Tier 2 lifts both operands into their divergence gauge before calling Tier 1.

normix.divergences.squared_hellinger_from_psi(psi, theta_p, theta_q)

Squared Hellinger distance from the log-partition function alone.

\[H^2(p, q) = 1 - \exp\!\left(\psi\!\left(\frac{\theta_p + \theta_q}{2}\right) - \frac{\psi(\theta_p) + \psi(\theta_q)}{2}\right)\]

Parameters:

• psi (callable) – Log-partition function \(\psi(\theta) \to \mathbb{R}\).

• theta_p (jax.Array) – Natural parameter vectors.

• theta_q (jax.Array) – Natural parameter vectors.

Return type:

Array

normix.divergences.kl_divergence_from_psi(psi, grad_psi, theta_p, theta_q)

KL divergence \(D_{\mathrm{KL}}(p \| q)\) as a Bregman divergence of \(\psi\).

\[D_{\mathrm{KL}}(p \| q) = \psi(\theta_q) - \psi(\theta_p) - (\theta_q - \theta_p)^\top \nabla\psi(\theta_p)\]

Parameters:

• psi (callable) – Log-partition function.

• grad_psi (callable) – Gradient \(\nabla\psi(\theta)\).

• theta_p (jax.Array) – Natural parameter vectors.

• theta_q (jax.Array) – Natural parameter vectors.

Return type:

Array

normix.divergences.kl_divergence_from_eta(psi, theta_p, theta_q, eta_fin, m, v)

KL with a tail-split expectation parameter \(\eta = \eta_{\mathrm{fin}} + m\,v\).

\[D_{\mathrm{KL}}(p \| q) = \psi(\theta_q) - \psi(\theta_p) - \Delta\theta^\top \eta_{\mathrm{fin}} + m\,A, \qquad A = \max(-\Delta\theta^\top v,\, 0)\]

The product \(m A\) is masked when \(A = 0\) so that an infinite moment with a zero chord coefficient stays finite (no \(0\cdot\infty\)).

Parameters:

• psi (callable) – Gauge log-partition.

• theta_p (jax.Array) – Natural parameters in gauge coordinates.

• theta_q (jax.Array) – Natural parameters in gauge coordinates.

• eta_fin (jax.Array) – Tail-split of \(\eta_p\) (see _divergence_eta()).

• m (jax.Array) – Tail-split of \(\eta_p\) (see _divergence_eta()).

• v (jax.Array) – Tail-split of \(\eta_p\) (see _divergence_eta()).

Return type:

Array

normix.divergences.squared_hellinger(p, q)

Squared Hellinger distance between two distributions.

For ExponentialFamily objects, calls p.squared_hellinger(q) (Tier 2), which defaults to the general \(\psi\)-based formula and can be overridden by subclasses.

For NormalMixture objects, delegates to the joint distributions as an upper-bound approximation.

Parameters:

• p (ExponentialFamily or NormalMixture) – Must be the same type (or at least share the same log-partition).

• q (ExponentialFamily or NormalMixture) – Must be the same type (or at least share the same log-partition).

Return type:

Array

normix.divergences.kl_divergence(p, q)

KL divergence \(D_{\mathrm{KL}}(p \| q)\).

For ExponentialFamily objects, calls p.kl_divergence(q) (Tier 2).

For NormalMixture objects, delegates to the joint distributions.

Parameters:

• p (ExponentialFamily or NormalMixture) – Must be the same type.

• q (ExponentialFamily or NormalMixture) – Must be the same type.

Return type:

Array