helpers

helpers#

Functions#

`count_min_instids`(→ int)	Minimum number of `instid` groups needed for `M` exposure windows
`assign_instids`(→ tinygp.helpers.JAXArray)	Assign each exposure window to one of `num_insts` groups such that no
`assign_min_instids`(→ tuple[tinygp.helpers.JAXArray, int])	Compute the minimum number of non-overlapping `instid` groups for a
`block_view`(A, b)
`Q_from_VanLoan`(→ tinygp.helpers.JAXArray)	Compute the process noise covariance via the Van Loan method.
`Phibar_from_VanLoan`(→ tinygp.helpers.JAXArray)	Compute the integrated transition matrix via the Van Loan method.
`robust_sqrt`(→ tinygp.helpers.JAXArray)	Symmetric-PSD matrix square root via eigendecomposition.
`transition_sequence`(→ tuple[tinygp.helpers.JAXArray, ...)	Per-step transition matrices and process noise, for precomputing.
`kalman_gain`(→ tinygp.helpers.JAXArray)	The Kalman gain
`smoothing_gain`(→ tinygp.helpers.JAXArray)	Computes the RTS smoothing gain \(G_k\) from
`VanLoan`(→ dict[str, tinygp.helpers.JAXArray])	Compute all submatrices of the Van Loan matrix exponential.

Module Contents#

smolgp.helpers.count_min_instids(t: tinygp.helpers.JAXArray, delta: tinygp.helpers.JAXArray) → int[source]#

Minimum number of instid groups needed for M exposure windows \((t_i - \delta_i/2,\; t_i + \delta_i/2)\) with arbitrary overlap.

This is the chromatic number of an interval graph, which equals its maximum clique. See https://en.wikipedia.org/wiki/Interval_graph.

Note for exposure windows, - \(b_i = a_j\) is not an overlap and can use the same group id. - A zero-width window (delta == 0) strictly inside another window’s

span must conflict, since its readout would otherwise corrupt the enclosing exposure’s running integra. However, two coincident zero-width windows do not conflict with each other.

So instead of a sweep, count for each window j the windows still open at \(a_j\), itself included, and take the largest.

Cost is \(O(M \log M)\), dominated by the sorts.

Deliberately numpy rather than jnp: the result sizes the augmented state (see merge_exposure_test_coords()), so it must be a concrete Python int and cannot be produced under jit.

Parameters:

t – Exposure midpoints, length M.
delta – Exposure widths, length M (must be >= 0).

Returns:

The minimum number of groups, i.e. int(jnp.max(instid)) + 1 for the assignment assign_instids() produces.

smolgp.helpers.assign_instids(t: tinygp.helpers.JAXArray, delta: tinygp.helpers.JAXArray, num_insts: int) → tinygp.helpers.JAXArray[source]#

Assign each exposure window to one of num_insts groups such that no two conflicting windows share a group.

Unlike count_min_instids() this is fully jittable, provided num_insts is a static Python int.

Since the count is already known, simply sweep the windows in order of start time and give each one any currently-free group. Optimal and vectorizes.

Cost is \(O(M \log M)\) for the sort plus \(O(M \cdot n)\), and beats the \(O(M \log M)\) for an eager heap because the heap’s cost is Python interpreter overhead rather than its asymptotics, except only for very large M, and the heap is not jittable.

Parameters:

t – Exposure midpoints, length M.
delta – Exposure widths, length M (must be >= 0).
num_insts – Number of groups to assign into; must be static, and at least count_min_instids() or the assignment is not valid.

Returns:

Length M integer array of group assignments, 0 <= instid[i] < num_insts.

smolgp.helpers.assign_min_instids(t: tinygp.helpers.JAXArray, delta: tinygp.helpers.JAXArray) → tuple[tinygp.helpers.JAXArray, int][source]#

Compute the minimum number of non-overlapping instid groups for a set of M exposure windows \((t_i - \delta_i/2,\; t_i + \delta_i/2)\) with arbitrary overlap.

Main use is in merge_exposure_test_coords(), to optimally reduce the dimensionality of the augmented state. Because sample draws are from a joint multivariate Gaussian, they cannot be drawn independently for overlapping exposures, so we must track them with separate instrument indices. The sampling algorithm scales with O(n^3), n the number of instruments, so we want to minimize that number.

This is the “minimum number of meeting rooms” problem, which is optimally solved with the standard “reuse whichever group’s window finished earliest, if it’s already finished” greedy sweep. Cost is \(O(M \log M)\).

For calls inside jit, first call count_min_instids() once outside the trace and then call assign_instids() directly.

Returns:: Length M integer array of group assignments. num_insts: The number of distinct groups used.
Return type:: instid

smolgp.helpers.block_view(A, b)[source]#

smolgp.helpers.Q_from_VanLoan(F: tinygp.helpers.JAXArray, L: tinygp.helpers.JAXArray, Qc: tinygp.helpers.JAXArray, dt: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Compute the process noise covariance via the Van Loan method.

Evaluates

\[Q_k = \int_0^{\Delta t} e^{F(\Delta t - s)}\, L\, Q_c\, L^T\, e^{F^T(\Delta t - s)}\, ds\]

See Van Loan (1978), “Computing Integrals Involving the Matrix Exponential” (PDF).

Parameters:

F – Feedback (design) matrix \(F\) from design_matrix().
L – Noise effect matrix \(L\) from noise_effect_matrix().
Qc – Spectral density \(Q_c\) from noise().
dt – Time step \(\Delta t = X_2 - X_1\).

Returns:

Process noise covariance matrix \(Q_k\) over time step \(\Delta t\).

smolgp.helpers.Phibar_from_VanLoan(F: tinygp.helpers.JAXArray, dt: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Compute the integrated transition matrix via the Van Loan method.

Evaluates

\[\bar{\Phi} = \int_0^{\Delta t} e^{F s}\, ds\]

See Van Loan (1978), “Computing Integrals Involving the Matrix Exponential” (PDF).

Parameters:

F – Feedback (design) matrix \(F\) from design_matrix().
dt – Time step \(\Delta t = X_2 - X_1\).

Returns:

Integrated transition matrix \(\bar{\Phi}\) over time step \(\Delta t\).

smolgp.helpers.robust_sqrt(M: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Symmetric-PSD matrix square root via eigendecomposition.

Returns \(S\) such that \(S S^T = M\), via \(M = V \mathrm{diag}(w) V^T\), \(S = V \mathrm{diag}(\sqrt{\max(w,0)})\).

In cases where \(M\) is numerically singular, jnp.linalg.cholesky() will fail. This method, while slightly slower, is robust to singularity and returns a valid square root. Needed for sampling (which uses a square root of the covariance) when either 1. Q_k is exactly zero, either from a zero-length transition (such as the first step,

or two states at the same instant), or for kernels with Q_k=0 everwhere (e.g. Cosine)

Q_k is numerically singular due to multiple instruments resetting at the same instant, which produces perfectly correlated integral states (see docstring of get_smoothing_gain()).

smolgp.helpers.transition_sequence(A, Q, t: tinygp.helpers.JAXArray) → tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray][source]#

Per-step transition matrices and process noise, for precomputing.

Returns A(0, Delta_k) and Q(0, Delta_k) for every step k, with \(\Delta_0 = 0\) for the first step (transition from the prior) and \(\Delta_k = t_k - t_{k-1}\) thereafter.

Building these with jax.vmap() ahead of the scan, and streaming them in as scan xs, is much faster than calling A/Q inside the scan body even though the arithmetic is identical: both are matrix exponentials (or Van Loan blocks) of a fixed generator, so as a scan body they are a serial chain of small unfusable kernels, whereas vmapped they become one batched kernel. The tradeoff is carrying two extra arrays of shape (N, dim, dim) in memory, which is not a significant addition.

smolgp.helpers.kalman_gain(S: tinygp.helpers.JAXArray, PHt: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

The Kalman gain

\[K_k = \mathbf{P}_k^- \mathbf{H}_k^T \mathbf{S}_k^{-1}.\]

Parameters:

S – The \(D \times D\) innovation covariance \(\mathbf{H}_k \mathbf{P}_k^- \mathbf{H}_k^T + \mathbf{R}_k\)
PHt – The product \(\mathbf{P}_k^- \mathbf{H}_k^T\), shape (dim, D).

For a scalar observation (D == 1) the inverse is just division, and D = S.shape[0] is a static shape known at trace time, so we can boost performance for this common case by avoiding solve.

smolgp.helpers.smoothing_gain(P_pred_next: tinygp.helpers.JAXArray, PAt: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Computes the RTS smoothing gain \(G_k\) from

\[G_k = \mathbf{P}_k \mathbf{A}_k^T \left[\mathbf{P}_{k+1}^{-}\right]^{-1}.\]

Parameters:

P_pred_next – The predicted covariance \(\mathbf{P}_{k+1}^{-}\) for the next step.
PAt – The product \(\mathbf{P}_k \mathbf{A}_k^T\) for the current step.

Returns:

The RTS smoothing gain \(G_k\).

TLDR; usage converts: G_k = jnp.linalg.solve(P_pred_next.T, (P_k @ A_k.T).T).T
to: G_k = get_smoothing_gain(P_pred_next, P_k @ A_k.T)

Guards against \(\mathbf{P}_{\mathrm{pred,next}}\) being exactly singular. This arises when two states in an exposure-aware model occupy the same instant in time, which produces a singular covariance in two different ways:

An exposure-start reset zeroes a row/column of the covariance. If the following transition has zero duration, that zeroed row/column passes through unregularized.

At nonzero transition lengths, whenever two or more instruments are reset at the exact same instant, their integral states become perfectly (diagonal and row identical) correlated, since they share the same driving process noise. This can persist through several further transitions before it clears.

Checking Delta == 0 only catches the first case, so we instead directly check for singularity in \(\mathbf{P}_{\mathrm{pred,next}}\) itself.

The detection is via P_pred_next’s (scale-normalized) log-determinant: dividing by trace(P_pred_next) / n before taking slogdet makes the threshold independent of the kernel’s overall amplitude (a plain absolute threshold on the raw determinant would not be, since determinant scales as amplitudeⁿ). A genuinely singular matrix here shows up many orders of magnitude below this threshold; typically -inf to around -40, whereas well-conditioned states have around -2 to -24. The threshold set here is -30.

Both branches compute the correct smoothing gain by inverting the predicted covariance, but with different methods: - The common (non-singular) case uses jnp.linalg.solve() (LU-based, cheap), which assumes invertibility. - The degenerate case uses jnp.linalg.lstsq() (SVD-based), which is more expensive but handles singular matrices correctly.

smolgp.helpers.VanLoan(F: tinygp.helpers.JAXArray, L: tinygp.helpers.JAXArray, Qc: tinygp.helpers.JAXArray, dt: tinygp.helpers.JAXArray) → dict[str, tinygp.helpers.JAXArray][source]#

Compute all submatrices of the Van Loan matrix exponential.

Assembles the block matrix \(C\) and returns its matrix exponential, partitioned into the submatrices F1-F4, G1-G3, H1-H2, K1 (see Van Loan 1978 for notation), from which various integrals such as Q_from_VanLoan() and Phibar_from_VanLoan() can be derived.

See Van Loan (1978), “Computing Integrals Involving the Matrix Exponential” (PDF).

Parameters:

F – Feedback (design) matrix \(F\).
L – Noise effect matrix \(L\).
Qc – Spectral density \(Q_c\).
dt – Time step \(\Delta t = X_2 - X_1\).

Returns:

Dictionary of named submatrices of the Van Loan exponential.