Accuracy and differentiation¤

A calculation combines a transition approximation with a method for repeated evolution or a linear solve. Tightening one part does not remove errors from the others. Validate the final statistic and its gradients for your workload.

Check the quantity you will use¤

Compare the same calculation with tighter numerical settings. Use a nonlinear observable that exercises the transition approximation: a mean frequency alone can agree because low-degree polynomials are reproduced accurately.

This example checks one-generation loss probability and its selection derivative. It also compares with the direct binomial formula for this particular event:

import jax
import jax.numpy as jnp
import drifft

jax.config.update("jax_enable_x64", True)
M = 256
coarse = drifft.NUFFT(M, accuracy=1e-6)
fine = drifft.NUFFT(M, accuracy=1e-10, derivative_accuracy=True)
loss_indicator = jnp.zeros(M + 1).at[0].set(1.0)

# The allele starts with four copies.
def loss_probability(backend, s_het):
    transition = drifft.Transition(
        backend, drifft.AdditiveSelection(s_het=s_het)
    )
    return transition.expect(loss_indicator)[4]

def exact_loss(s_het):
    q = drifft.AdditiveSelection(s_het=s_het)(jnp.asarray(4 / M))
    return (1 - q) ** M

strength = jnp.asarray(0.01)
exact_value = exact_loss(strength)
exact_gradient = jax.grad(exact_loss)(strength)
for label, backend in [("coarse", coarse), ("fine", fine)]:
    value = loss_probability(backend, strength)
    gradient = jax.grad(lambda s: loss_probability(backend, s))(strength)
    print(f"{label}: value error={abs(value - exact_value):.3e}, "
          f"gradient error={abs(gradient - exact_gradient):.3e}")

Inspect both errors against the precision required by your application. This reference uses exact binomial reproduction and the same selection map; it checks the transition approximation, not the selection formula itself. For general quantities, use Dense on a small population as an independent binomial reference. For a history or an expected-visits solve, also compare the final statistic under tighter settings: one-step agreement is not sufficient.

Which settings control which errors?¤

Choose controls for the component whose accuracy you want to test. Tightening a power or solve tolerance does not improve the transition representation.

For reduced Jacobi evolution, prepare max_strength to cover the absolute values of both diploid selection coefficients over your parameter range. See the Jacobi API for assembly restrictions.

With Krylov, inspect converged and status when using throw=False. Its error_estimate assesses the power approximation; the Krylov chapter defines its acceptance criterion. Expected-visits results instead report a solve residual.

Tune the representation separately from the method¤

For FFT, tighten tail truncation and increase interpolation order. For NUFFT, lower accuracy and compare the final quantity again. Setting derivative_accuracy=True increases polynomial degree by 50%, capped at exact degree; it is a heuristic for derivative accuracy. Filtered encoding can reduce the sample count needed by its spectral bound, but must still be checked for your statistic and gradient.

For Jacobi, increase rank; for Krylov, increase dimension separately. If a shifted solve fails, increasing the outer Krylov dimension does not fix it. See the backend API and method API for parameter defaults and NUFFT transforms for encoding bounds.

Check mass, negative entries, and the final observable. The implementations do not clip or renormalize transition outputs; tiny negative entries can arise from the numerical approximation.

Differentiating biological parameters¤

Differentiate biological parameters while holding numerical settings and prepared geometry fixed. Gradient accuracy needs its own checks, even when the forward result has converged.

Selection and mutation parameters, initial mass, and immigration can be array-valued inputs to a differentiated function. Population sizes, durations, and numerical dimensions are fixed configuration. Construct backends outside the function, then create or update the frequency map inside it; the population-history example shows this pattern with jax.grad.

Krylov operator sensitivities support first derivatives; use Sequential for higher-order calculations. Sensitivity failures raise even with throw=False. See the autodiff accuracy contract for the derivative targets and limitations.

For a complete finite-time calculation with a derivative plot and finite-difference checks, see sensitivity to selection.

Memory and timing¤

Differentiation can retain more intermediates than forward evaluation, especially across long intervals. NUFFT's batch_size controls temporary binomial-sampling storage; native FINUFFT workspace is additional to JAX memory reports. Krylov basis storage grows with working-state size and dimension.

For timing, separate preparation, first-call compilation, and warm execution, and synchronize JAX outputs. Use validation and benchmarks for the repository's comparison commands and recorded results.