Skip to content

Evolution methods¤

Evolution methods propagate allele-count distributions over a constant-size interval. They apply the prepared transition repeatedly or approximate its power.

AbstractEpochMethod

drifft.AbstractEpochMethod(equinox.Module) [source] ¤

An algorithm for a prepared linear power action, independent of geometry.

Concrete methods declare uses_reduced_representation: true allows backend reduction during preparation; false retains population coordinates. Direct calls and demographic histories use the same declaration.

validate(self, backend: Module, frequency_map: AbstractFrequencyMap) -> None ¤

Reject unsupported representation and process combinations.

Arguments:

  • backend: Prepared transition geometry.
  • frequency_map: Process applied in each generation.

Returns:

None. Checks types and supported combinations, without evaluating dynamic parameter values. Runtime domain checks occur during preparation.

Raises:

TypeError for unsupported backend or map types. Sequential accepts all transition backends; reduced Jacobi powers restrict map compositions.

solve(self, problem: _PowerProblem, *, throw: bool = True) -> EpochResult ¤

Evaluate a power problem prepared internally by drifft.

Arguments:

  • problem: Internal operator, initial state, static exponent and reconstruction map. Use advance to supply population-space inputs.
  • throw: Raise on numerical failure when true; otherwise return status.

Returns:

An EpochResult in population coordinates. Its error estimate covers the power approximation, excluding backend and representation errors.

Raises:

Runtime errors for numerical failure with throw=True. A failed batch member may fail the entire call. Krylov sensitivity failures raise even with throw=False; higher-order operator sensitivities are unsupported.

advance(self, backend: Module, distribution: Array, *, frequency_map: AbstractFrequencyMap, generations: int, segregating: bool = False, immigration: jax.Array | None = None, throw: bool = True) -> EpochResult ¤

Bind geometry and evolve a constant-size interval.

Arguments:

  • backend: Prepared geometry compatible with this method.
  • distribution: Real mass, converted to float64. Length is backend.gene_copies + 1, or backend.gene_copies - 1 if segregating. Signed entries are supported; no normalization is applied.
  • frequency_map: Process shared by every generation.
  • generations: Nonnegative Python integer, static under JIT.
  • segregating: Discard endpoint mass at each generation.
  • immigration: Mass added after every generation, matching the input shape; None means no source. In segregating mode it is added after endpoint mass is discarded.
  • throw: Raise on numerical failure; under vmap this may fail the whole batch.

Returns:

EpochResult; zero generations preserves the input without immigration.

Raises:

TypeError for unsupported combinations; ValueError for static input errors. Runtime domain checks remain active even with throw=False.

drifft.EpochResult(builtins.tuple) [source] ¤

Evolved mass, power-action L1 error estimate, convergence and status.

The estimate excludes one-generation approximation and representation error. With throw=False, use the distribution only after inspecting convergence.

Fields:

  • distribution: Float64 full or interior mass, with the input population shape. Signed inputs are not normalized.
  • error_estimate: Scalar absolute L1 power-error estimate in population coordinates; zero for Sequential and identity actions.
  • converged: Scalar boolean, true exactly when status is successful.
  • status: EvolutionStatus describing the numerical outcome.

drifft.EvolutionStatus(equinox._enum.Enumeration) [source] ¤

Numerical outcomes; successful power checks do not certify backend accuracy.

  • successful: Finite output with accepted power-error and inner-solve checks.
  • nonfinite: Nonfinite output or error estimate.
  • unconverged: Power-error estimate exceeds the configured tolerance.
  • linear_solve_failed: A shifted solve failed its status or residual checks.

Compare a result's status with these members. A failed shifted solve takes precedence over nonfiniteness, which takes precedence over power nonconvergence.


drifft.Sequential(drifft.AbstractEpochMethod) [source] ¤

Repeat one-generation transitions without an additional power approximation.

Construct as Sequential(). Geometry is supplied separately to advance or through DemographicTransition. The reported power-action error is zero.

__init__(self) -> None ¤

Configure repeated one-generation evolution.

Arguments:

No arguments. Supply geometry separately to advance or DemographicTransition.

Returns:

None.

advance(self, backend: Module, distribution: Array, *, frequency_map: AbstractFrequencyMap, generations: int, segregating: bool = False, immigration: jax.Array | None = None, throw: bool = True) -> EpochResult ¤

Repeat the one-generation transition in population coordinates.

Arguments:

  • backend, distribution, frequency_map, generations, segregating, immigration, throw: See AbstractEpochMethod.advance for shapes, source timing and failure behavior.

Returns:

An EpochResult with zero power-approximation error. Backend errors and roundoff remain; Sequential supports higher-order differentiation.


drifft.PolynomialKrylov(drifft.AbstractEpochMethod) [source] ¤

Approximate repeated evolution with matrix-free polynomial Arnoldi.

For a prepared operator \(A\), seed \(b\), and \(g\) steps, build an orthonormal basis \(V\) for \(\mathcal K_m(A,b)=\mathrm{span}(b,Ab,\ldots,A^{m-1}b)\). With \(H=V^TAV\), approximate \(A^g b\approx VH^gV^Tb\). Only operator applications are needed; there are no shifted solves. All transition backends are supported, including Dense for small reference calculations.

The basis uses O(n m) storage for working-state size n and dimension m. Larger dimensions can be necessary for long intervals; no restart or automatic Sequential fallback is performed. Convergence checks exclude backend approximation error. First-order derivatives use a coupled Krylov sensitivity action and a two-sided operator pullback, without basis AD. Sensitivity or pullback nonconvergence raises even with throw=False. Higher-order operator sensitivities are not implemented. This method does not provide a matched expectation transpose of its primal approximation.

__init__(self, dimension: int = 64, rtol: float = 1e-08, atol: float = 1e-12) -> None ¤

Configure polynomial powering independently of geometry.

Arguments:

  • dimension: Maximum basis size, a Python integer at least two. Capped by the working-state size and the number of steps plus one.
  • rtol, atol: Finite nonnegative relative and absolute tolerances for the power-action error estimate; these do not control backend approximation.

Returns:

None. Supply transition geometry separately.

Raises:

ValueError for invalid dimension or tolerances.

advance(self, backend: Module, distribution: Array, *, frequency_map: AbstractFrequencyMap, generations: int, segregating: bool = False, immigration: jax.Array | None = None, throw: bool = True) -> EpochResult ¤

Evolve mass using a checked Krylov power approximation.

Arguments:

  • backend, distribution, frequency_map, generations, segregating, immigration, throw: See AbstractEpochMethod.advance for shapes and source timing.

Returns:

An EpochResult whose power-error estimate excludes representation, quadrature and assembly error. Check converged with throw=False; increase rank and dimension independently to assess both approximations.

Raises:

Static input errors follow AbstractEpochMethod.advance. Runtime Jacobi domain checks and sensitivity convergence checks raise even with throw=False. Higher-order operator sensitivities are unsupported.


drifft.RationalKrylov(drifft.AbstractEpochMethod) [source] ¤

Approximate a prepared power action using a resolvent Krylov basis.

Each basis vector solves \([I+\tau(I-A)]z=v\), with \(\tau=g/8\) for working exponent \(g\). Jacobi uses a reusable dense LU factorization; other transition backends use unpreconditioned Lineax GMRES. Solver choice is internal. Inner failures are reported as linear_solve_failed.

GMRES uses rtol=1e-11, atol=1e-13, restart=50 (capped by state size), and at most 100 restart cycles per solve, with a recomputed residual check. These inner tolerances are independent of the outer power tolerances. FFT/NUFFT speedups depend on how many operator applications GMRES needs.

First-order derivatives approximate the mathematical power sensitivity with a coupled Krylov action and a two-sided operator pullback. Arnoldi bases are not differentiated. Sensitivity or pullback nonconvergence raises even with throw=False. Higher-order operator sensitivities are not implemented. The input-dependent primal approximation has no matched expectation transpose.

__init__(self, dimension: int = 64, rtol: float = 1e-08, atol: float = 1e-12) -> None ¤

Configure the power algorithm without preparing geometry.

Arguments:

  • dimension: Maximum Krylov basis size, an integer at least two; capped by the prepared working-state dimension.
  • rtol, atol: Finite nonnegative relative and absolute tolerances for the reconstructed power-action error estimate, excluding backend error.

Returns:

None. Shifted solves use dense LU for Jacobi and operator-based GMRES for other transition backends. Solver selection is internal.

Raises:

ValueError for invalid dimension or tolerances.

advance(self, backend: Module, distribution: Array, *, frequency_map: AbstractFrequencyMap, generations: int, segregating: bool = False, immigration: jax.Array | None = None, throw: bool = True) -> EpochResult ¤

Evolve mass using a checked Krylov power approximation.

Arguments:

  • backend, distribution, frequency_map, generations, segregating, immigration, throw: See AbstractEpochMethod.advance for shapes and source timing.

Returns:

An EpochResult whose power-error estimate excludes representation, quadrature and assembly error. Check converged with throw=False; increase rank and dimension independently to assess both approximations.

Raises:

Static input errors follow AbstractEpochMethod.advance. Runtime Jacobi domain checks and sensitivity convergence checks raise even with throw=False. Higher-order operator sensitivities are unsupported.