Utilities

Utilities#

Bessel Functions #

JAX-compatible log modified Bessel function of the second kind.

log_kv(v, z) = log K_v(z), fully pure-JAX with zero scipy callbacks.

Regime-specific methods, selected via lax.cond (only one branch executes):

Hankel asymptotic (DLMF 10.40.2) — large z
Olver uniform expansion (DLMF 10.41.4) — large v
Small-z leading asymptotic (DLMF 10.30.2) — z → 0
Gauss-Legendre quadrature (Takekawa 2022) — moderate z, v

Using lax.cond (not jnp.where) means only the selected branch executes at runtime. lax.cond requires scalar conditions, so the core scalar function _log_kv_scalar is vmapped over array inputs.

Custom JVP for full autodiff (defjvp(..., symbolic_zeros=True)):

∂/∂z : exact recurrence K’_v = −(K_{v−1}+K_{v+1})/2; skipped when the z-tangent is a symbolic zero (v-only differentiation)
∂/∂v : central FD on log_kv itself (ε = BESSEL_EPS_V); skipped when the v-tangent is a symbolic zero (z-only differentiation)

backend=’jax’ (default): pure-JAX, JIT-able, differentiable. backend=’cpu’ : scipy.special.kve, fully vectorized numpy.

Not JIT-able. Fast for EM hot path.

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}\)).
z (scalar or array) – Argument (must be > 0).
backend (str, optional) –
'jax' (default) or 'cpu'.
- 'jax': pure-JAX, lax.cond regime selection, custom JVP. JIT-able, differentiable. Default for log_prob, pdf, etc.
- 'cpu': scipy.special.kve, fully vectorised NumPy. Not JIT-able. Fast for EM hot path.

Returns:

Same broadcast shape as (v, z).

Return type:

jax.Array

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.112...

CPU backend (uses scipy, faster for EM hot-paths):

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

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

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)

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:

n (int)
seed (int)

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:

key (Array) – JAX PRNG key.
u_grid (Array) – Arrays returned by build_pinv_table().
x_grid (Array) – Arrays returned by build_pinv_table().
n (int) – Sample size.

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:

axes (Axes | list[Axes] | tuple[Axes, ...])
grid_axis (str)
legend (bool)

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:

n_samples (int)
seed (int)
title (str)

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:

n_samples (int)
seed (int)
title (str)

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:

log_likelihoods (List[float])
title (str)
true_ll (float | None)

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:

n_samples (int)
seed (int)
is_joint (bool)

Return type:

Dict[str, Any]

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

Print moment validation results.

Parameters:

results (Dict[str, Any])
title (str)

Return type:

None

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

Print natural and expectation parameters.

Parameters:: label (str)
Return type:: None

Utilities

Contents

Utilities#

Bessel Functions #

Constants #

Incomplete Gamma #

PINV Sampling #

Plotting #

Validation #

Utilities

Contents

Utilities#

Bessel Functions#

Constants#

Incomplete Gamma#

PINV Sampling#

Plotting#

Validation#

Bessel Functions #

Constants #

Incomplete Gamma #

PINV Sampling #

Plotting #

Validation #