Utilities

Utilities#

Bessel Functions#

Log modified Bessel function of the second kind.

One centered whole-line Gauss–Legendre kernel (S10) implements both log_kv and log_kv_moments(). Geometry is frozen under stop_gradient; autodiff of log_kv yields cumulants of that discrete measure. No regimes, no lax.cond, no custom_jvp, no finite differences.

backend=’jax’ (default): JAX arrays, JIT-able, differentiable. backend=’cpu’ : the same sums in NumPy.

class normix.utils.bessel.BesselMoments(log_k, u0, mean, cov, d_arg, d2_arg, d2_order_arg)[source]#

Frozen quadrature moments of \((x, e^{-x}-1, e^{x}-1)\).

The discrete exponential family lives on the centered coordinate \(x = u - u_0\) with \(u_0 = \operatorname{asinh}(\nu/z)\). Geometry is constant under AD (see log_kv_moments()); mean and cov are the mean and centered Gram of \(s = (x, \operatorname{expm1}(-x), \operatorname{expm1}(x))\). cov is PSD to rounding (relative-ε entries). Argument-jet views are projections of \(w = 2\sinh(u_0 + x/2)\sinh(x/2)\), not reconstructions from cov.

Parameters:

  • log_k (jax.Array) – \(\log K_\nu(z)\), broadcast shape of (v, z).

  • u0 (jax.Array) – Mode \(\operatorname{asinh}(\nu/z)\).

  • mean (jax.Array) – \(E[s]\), shape (..., 3).

  • cov (jax.Array) – \(\mathrm{Cov}[s]\), shape (..., 3, 3).

  • d_arg (jax.Array) – \(\partial_z \log K_\nu = -(\cosh u_0 + E[w])\).

  • d2_arg (jax.Array) – \(\partial_{zz} \log K_\nu = \mathrm{Var}(w)\).

  • d2_order_arg (jax.Array) – \(\partial_{\nu z} \log K_\nu = -\mathrm{Cov}(x, w)\).

log_k: Array#
u0: Array#
mean: Array#
cov: Array#
d_arg: Array#
d2_arg: Array#
d2_order_arg: Array#
property d_order: Array#

\(\partial_\nu \log K_\nu = E[u] = u_0 + E[x]\).

property d2_order: Array#

\(\partial_{\nu\nu} \log K_\nu = \mathrm{Var}(u) > 0\).

normix.utils.bessel.log_kv_moments(v, z, backend='jax')[source]#

Log \(K_\nu(z)\) and the moment bundle of the quadrature measure.

One centered whole-line Gauss–Legendre kernel; backend selects the array library (JAX or NumPy), not a different approximation. Geometry is frozen under jax.lax.stop_gradient(), so the log-sum-exp identity is exact for any fixed nodes and

\[\partial_\nu \log K = E[u], \qquad \partial_{\nu\nu} \log K = \mathrm{Var}(u) > 0.\]

Stated numerical domain: \(z \ge 10^{-10}\) (GIG never asks below GIG_DEGEN_THRESHOLD), \(|\nu| \le 300\) tested. The point \((\nu,z)=(0, 10^{-300})\) is unresolved at 192 nodes.

Parameters:

  • v (scalar or array) – Order (any real; \(\nu=0\) is an ordinary point).

  • z (scalar or array) – Argument (must be \(> 0\)).

  • backend ({'jax', 'cpu'}, optional) – Array backend. 'jax' is JIT-able and differentiable; 'cpu' is the same sums in NumPy.

Returns:

log_k, u0, mean (shape (..., 3)), cov (shape (..., 3, 3)). Jet views: d_order, d_arg, d2_order, d2_order_arg, d2_arg. Argument jets (d_arg, d2_arg, d2_order_arg) are stored projections of \(w=2\sinh(u_0+x/2)\sinh(x/2)\), not affine images of cov.

Return type:

BesselMoments

Notes

jax.grad(lambda z: log_kv(v, z))(z) equals d_arg; jax.grad with respect to v equals d_order. Both are derivatives of the frozen quadrature, not finite differences. GIG uses this bundle so \(\eta\) and the Fisher come from one pass.

normix.utils.bessel.log_kv(v, z, backend='jax')[source]#

\(\log K_v(z)\) — log modified Bessel function of the second kind.

Parameters:

  • v (scalar or array) – Order (any real; \(K_v = K_{-v}\); \(\nu=0\) is ordinary).

  • z (scalar or array) – Argument (must be \(> 0\)). Stated domain \(z \ge 10^{-10}\).

  • backend (str, optional) – 'jax' (default) or 'cpu'. Both use the same quadrature; 'jax' is JIT-able and differentiable.

Returns:

Same broadcast shape as (v, z).

Return type:

jax.Array

Notes

A 192-point Gauss–Legendre sum of the whole-line integrand for \(K_\nu\). jax.grad with respect to z equals log_kv_moments() d_arg; with respect to v it equals d_order. Neither is a finite difference.

Examples

Evaluate at a single point (JAX backend, JIT-able):

>>> import jax.numpy as jnp
>>> from normix import log_kv
>>> float(log_kv(v=0.5, z=1.0))
-0.774...

CPU backend (same kernel in NumPy):

>>> float(log_kv(v=0.5, z=1.0, backend='cpu'))
-0.774...

Symmetry \(K_v(z) = K_{-v}(z)\):

>>> abs(float(log_kv(0.5, 2.0)) - float(log_kv(-0.5, 2.0))) < 1e-10
True

Differentiable via JAX:

>>> import jax
>>> dlogkv_dz = jax.grad(lambda z: log_kv(0.5, z))(jnp.array(1.0))
>>> float(dlogkv_dz) < 0   # K_v decreases with z
True

Constants#

Shared numerical constants for normix.

Incomplete Gamma#

Pure-JAX inverse of the regularised incomplete gamma function.

JAX exposes jax.scipy.special.gammainc() (the regularised lower incomplete gamma \(P(a, x) = \gamma(a, x)/\Gamma(a)\)) but does not ship its inverse. gammaincinv() solves \(P(a, x) = q\) by Newton iteration with a Wilson–Hilferty starting guess — fully JIT- and vmap-compatible. This is the JAX analogue of scipy.special.gammaincinv().

normix.utils.gammainc.gammaincinv(a, q, max_iter=20, tol=1e-12)[source]#

Solve \(P(a, x) = q\) for \(x\) by Newton iteration.

JAX equivalent of scipy.special.gammaincinv().

Parameters:

  • a (Array) – Shape parameter, \(a > 0\).

  • q (Array) – Probability, \(q \in (0, 1)\).

  • max_iter (int) – Maximum Newton iterations (typically converges in <10 steps).

  • tol (float) – Absolute residual tolerance on \(|P(a, x) - q|\).

Return type:

Array

Notes

Starts from the Wilson–Hilferty cube-root normal approximation and iterates \(x \leftarrow x - (P(a, x) - q) / p(a, x)\), where \(p(a, x) = x^{a-1} e^{-x} / \Gamma(a)\) is the Gamma density. The density is evaluated in log space to avoid overflow for large \(a\).

PINV Sampling#

Generic RVS utilities for univariate distributions.

build_pinv_table() builds a quantile table from any univariate log-kernel in pure JAX (trapezoidal CDF on a \(w\)-grid). Distributions supply log_kernel(w) from their own log_prob (plus a Jacobian when working in \(w = \log x\)). QuantileTable holds the grids for amortised cdf / ppf / rvs; rvs_pinv() samples via inverse lookup.

class normix.utils.rvs.QuantileTable(u_grid, x_grid)[source]#

Frozen PINV quantile table (a pytree — jit/vmap/scan-safe).

Built once via build_pinv_table() (or a distribution’s quantile_table()); reuse across repeated cdf / ppf / rvs calls so the 4000-point grid is not rebuilt each time.

Parameters:
u_grid: Array#
x_grid: Array#
cdf(x)[source]#

CDF lookup \(F(x)\) via linear interpolation on the table.

Parameters:

x (Array)

Return type:

Array

ppf(q)[source]#

Quantile lookup \(F^{-1}(q)\) via linear interpolation.

Parameters:

q (Array)

Return type:

Array

rvs(n, seed=42)[source]#

Inverse-CDF sample of size n from this table.

Parameters:
Return type:

Array

normix.utils.rvs.build_pinv_table(log_kernel, mode, *, x_of_w=None, n_grid=4000, tail_eps=1e-14)[source]#

Build a PINV quantile table in pure JAX.

Parameters:

  • log_kernel (Callable[[Array], Array]) – Callable w -> log f(w) for a univariate density on the internal \(w\)-axis. For support \((0, \infty)\) with \(w = \log x\), pass log_kernel(w) = log_prob(exp(w)) + w. For support \(\mathbb{R}\) with \(w = x\), pass log_kernel(w) = log_prob(w).

  • mode (Array) – Mode of the density on the \(w\)-axis (starting point for tail bisection).

  • x_of_w (Callable[[Array], Array] | None) – Map internal \(w\) to the observation axis (default identity). Use jnp.exp when \(w = \log x\).

  • n_grid (int) – Number of grid points for the trapezoidal CDF.

  • tail_eps (float) – Tail mass below which bisection stops.

Returns:

JAX arrays of shape (n_grid,) — trapezoidal CDF values and corresponding \(x\) values for jnp.interp in ppf / cdf.

Return type:

u_grid, x_grid

normix.utils.rvs.rvs_pinv(key, u_grid, x_grid, n)[source]#

Sample n observations via numerical inverse CDF.

Parameters:
Returns:

Array of shape (n,).

Return type:

samples

Plotting#

Requires the plotting extra (uv sync --extra plotting); used by every executable tutorial for consistent figure styling.

Plotting utilities for normix notebooks.

Requires the plotting extra: pip install normix[plotting].

normix.utils.plotting.set_theme(*, scale=1.0)[source]#

Apply the normix Matplotlib theme (Kami-derived visual tokens).

Shared colour/typography tokens so tutorial figures match the documentation site aesthetic.

Parameters:

scale (float)

Return type:

None

normix.utils.plotting.style_axes(axes, *, grid_axis='y', legend=True)[source]#

Apply final styling to one or more Matplotlib axes.

Parameters:
Return type:

Axes | list[Axes]

normix.utils.plotting.savefig(fig, path, **kwargs)[source]#

Save a figure with normix theme defaults.

Parameters:

  • fig (Figure)

  • path (str)

Return type:

str

normix.utils.plotting.plot_pdf_cdf_comparison(configs, x, xlabel='x', title='')[source]#

Plot PDF and CDF comparing normix vs scipy for multiple distributions.

Parameters:

  • configs (list of dicts with keys {label, dist, scipy})

  • x (1-D evaluation grid)

  • xlabel (str)

  • title (str)

Return type:

Figure

normix.utils.plotting.plot_sample_histograms(configs, ncols=2)[source]#

Grid of histograms with theoretical PDF overlay.

Parameters:

  • configs (list of dicts with keys {label, dist, samples, x_plot (optional)})

  • ncols (int)

Return type:

Figure

normix.utils.plotting.plot_mle_fit(data, fit_results, xlabel='x', title='Maximum Likelihood Estimation')[source]#

Histogram of data with multiple PDF overlays for MLE comparison.

Parameters:

  • data (1-D sample array)

  • fit_results (list of dicts {label, dist, ls, color})

  • xlabel (str)

  • title (str)

Return type:

Figure

normix.utils.plotting.plot_joint_1d(joint_dist, n_samples=5000, seed=42, title='')[source]#

Visualize a 1-D joint distribution f(x, y).

Left: scatter X vs Y. Right: marginal histogram of X.

Parameters:
Return type:

Figure

normix.utils.plotting.plot_marginal_2d(marginal_dist, n_samples=5000, seed=42, title='')[source]#

Visualize a 2-D marginal distribution: scatter + marginal histograms.

Parameters:
Return type:

Figure

normix.utils.plotting.plot_em_convergence(log_likelihoods, title='EM Convergence', true_ll=None)[source]#

Plot EM log-likelihood convergence curve.

Parameters:
Return type:

Figure

Validation#

Moment-validation and parameter-printing helpers used in notebooks and tutorials.

Moment validation and parameter printing utilities for normix notebooks.

normix.utils.validation.validate_moments(dist, n_samples=20000, seed=42, is_joint=True)[source]#

Validate E[X] and E[Y] by comparing samples vs analytical values.

Uses dist.mean() for theoretical E[X] and subordinator().mean() for E[Y].

Parameters:
Return type:

Dict[str, Any]

normix.utils.validation.print_moment_validation(results, title='')[source]#

Print moment validation results.

Parameters:
Return type:

None

normix.utils.validation.print_exp_family_params(dist, label='')[source]#

Print natural and expectation parameters.

Parameters:

label (str)

Return type:

None