augax.ks_sample

augax.ks_sample(u, key)

Draw the Kolmogorov scale lam given a logistic disturbance u, on-device.

A standard logistic \(\varepsilon\) is the normal scale mixture \(\varepsilon \mid \lambda \sim N(0, \lambda^2)\), \(\lambda = 2K\) with \(K\) Kolmogorov distributed. Given \(\varepsilon = u\) this draws \(\lambda\) from

\[p(\lambda \mid u) \propto \lambda^{-1}\,\varphi(u/\lambda)\,f_{2K}(\lambda),\]

exactly: a rejection sampler with an inverse-Gaussian proposal for \(1/\lambda^2\) and the alternating-series squeeze of the Kolmogorov density in both of its forms (Holmes & Held 2006).

Parameters:

• u (array_like) – Shape [n] float64: the disturbances.

• key (jax.Array) – A JAX PRNG key. vmap over a batch of keys draws each element’s batch with its own RNG stream.

Returns:

lam – The scales, with the shape of u; lam**2 is the variance.

Return type:

jax.Array