Numerical backends¤

Binomial expectations through local FFT convolutions.

For destination values \(v_j\) and gene-copy count \(M\), approximate \(B_M[v](q)=\sum_{j=0}^M\binom{M}{j}q^j(1-q)^{M-j}v_j\). On each local log-odds grid, the retained binomial kernel factors as

Here \(j_0\) is the reference allele count, \(h\) is the log-odds spacing, and \(a_r,b_j\) are diagonal scaling weights. The sum is a convolution, evaluated by zero-padded FFTs; local polynomial interpolation supplies expectations at arbitrary \(q\). Near zero and one, direct binomial sums replace log-odds interpolation.

For offspring probabilities \(q_i\), Transition uses \((Pv)_i=B_M[v](q_i)\) for expectations and the transpose \(P^Tp\) for forward mass. The approximation comes from tail truncation and interpolation; the Gaussian-shaped convolution above is a factorization of the binomial kernel, not a normal approximation to it.

Binomial expectations through Chebyshev approximation and NUFFT evaluation.

For destination values \(v_j\) and gene-copy count \(M\), approximate \(B_M[v](q)=\sum_{j=0}^M\binom{M}{j}q^j(1-q)^{M-j}v_j\) by

Here \(T_k\) is the degree-\(k\) Chebyshev polynomial and \(L\) is the prepared approximation degree. The coefficients \(c_k\) come from binomial sampling at fixed Chebyshev–Lobatto nodes followed by a DCT-I. A type-2 NUFFT evaluates the cosine series at arbitrary offspring probabilities. encoding="lobatto" interpolates the samples. encoding="filtered" tapers the modes and adds a linear endpoint correction before evaluation. Endpoint expectations are set directly to \(v_0\) and \(v_M\) to suppress evaluation roundoff. degree always means the output polynomial degree; compressed filtered encoding uses degree + 2 sampling nodes.

For offspring probabilities \(q_i\), Transition uses \((Pv)_i=B_M[v](q_i)\) for expectations and the transpose \(P^Tp\) for forward mass. Polynomial truncation, truncated binomial sampling, and NUFFT evaluation each contribute approximation error.

Jacobi geometry for binomial expectations and reduced Krylov evolution.

For gene-copy count \(M\) and destination values \(v_j\), the binomial expectation is

Write \(z=2q-1\). After separating the linear interpolant between \(v_0\) and \(v_M\), approximate the remaining function in the basis \((1-z^2)\phi_k(z)\), \(k=0,\ldots,r-1\). Here \(r\) is rank and \(\phi_k\) is a normalized Jacobi polynomial \(P_k^{(1,1)}\), orthonormal for weight \(3(1-z^2)/4\) on \([-1,1]\). The envelope makes every interior basis function vanish at the count endpoints. Gauss–Jacobi quadrature and binomial stencils map count-space values to these coefficients.

Let \(E\) denote this coefficient map and \(D\) evaluation at the offspring probabilities \(q_i=f(i/M)\). The expectation operator factors as \(P\approx DE\), so forward mass evolves by its transpose. For \(g\geq1\) generations with a fixed frequency map and no immigration,

The reduced operator has \(r\) interior coordinates, plus two endpoint coordinates for full-state evolution. It is assembled using polynomial recurrences and a padded tridiagonal coordinate matrix, without forming the population-sized transition matrix. A separate epoch method computes the power of \(H^T\); representation rank and padding control different approximations from the method's power tolerance.

Transition and Sequential accept arbitrary valid offspring probabilities through the common expectation interface. Reduced Krylov assembly supports neutral and diploid selection, standalone mutation, and one mutation composed with one supported selection rule in either order. Custom maps and longer compositions are not supported by reduced assembly. Recurrent mutation couples the endpoint states to the interior; full-state evolution retains this coupling, while segregating evolution discards endpoint mass each generation. Construct outside JIT with JAX x64 enabled; the prepared geometry is independent of the dynamic process parameters. Reduced epoch domains are checked at runtime, including for zero generations: selection must stay within the configured bound, and mutation rates must be finite and in [0, 1]. In DemographicTransition, size-changing steps use the same backend by default; an explicit FFT or NUFFT boundary_backend can be supplied instead.

Form binomial transition rows for an independent numerical reference.

For destination gene-copy count \(M\) and source probabilities \(q_i\), construct

Each row is built by repeated Bernoulli convolution, starting at count zero. This avoids factorials, logarithms, and divisions by endpoint probabilities. The calculation is polynomial throughout, so ordinary JAX differentiation gives finite endpoint derivatives without custom rules. It shares no numerical kernels with FFT, NUFFT, or Jacobi.

expect explicitly forms the matrix before multiplying by destination values. With \(N\) source rows, this uses \(O(NM)\) matrix storage and \(O(NM^2)\) arithmetic. Reverse-mode differentiation can retain intermediate rows. Use this backend for small-population checks of transitions, demographic evolution, and gradients; it is not intended for large populations. There is no truncation or interpolation, but floating-point roundoff remains.