Skip to content

Transitions¤

Transitions map allele-count distributions forward by one generation and destination values to conditional expectations. They also provide segregating operations and expected-visit solves.

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

A reusable backend combined with a frequency map or explicit probabilities.

Given source-row probability \(q_i\) and destination gene-copy count \(M\),

\[P_{ij} = \binom{M}{j} q_i^j (1-q_i)^{M-j}.\]

expect applies \(P\) to destination values; advance applies \(P^T\) to source mass. The backend controls how \(P\) is represented and evaluated. For a frequency map, \(q_i=q(i/M_s)\) where \(M_s\) is source_gene_copies.

Segregating operations restrict the kernel to interior allele counts and discard endpoint mass. They require complete, ordered allele-count grids; sampled probability rows are for full-state operations. For a square interior kernel \(Q\), expected_visits solves \((I-Q^T)g=e_1\), including the initial generation. Finite visits require eventual endpoint hitting; mutation does not restore discarded trajectories.

probabilities property ¤

Current float64 source probabilities, shape (source_gene_copies + 1,).

Evaluates the bound map on each access, or returns the explicit rows. Raises ValueError if the map returns an incompatible shape.

__init__(self, backend: AbstractTransitionBackend, frequency_map: drifft._frequency.AbstractFrequencyMap | None = None, *, source_gene_copies: int | None = None, probabilities: jax.Array | None = None) -> None ¤

Bind a frequency map or explicit probabilities to prepared geometry.

Arguments:

  • backend: Prepared destination backend.
  • frequency_map: AbstractFrequencyMap; mutually exclusive with probabilities.
  • source_gene_copies: Source size, at least two for a map. Defaults to the backend size, or len(probabilities) - 1 for explicit rows.
  • probabilities: Optional nonempty vector of source probabilities in [0, 1]; converted to float64.

Returns:

None. Reuses the backend without rebuilding geometry.

Raises:

ValueError for ambiguous configuration or invalid source dimensions or probability shape; TypeError for an invalid backend or frequency-map type.

Finiteness and probability ranges are caller preconditions, not checked here. Enable JAX x64 before preparing the backend.

from_probabilities(backend: AbstractTransitionBackend, probabilities: Array) -> Transition classmethod ¤

Construct a transition from explicit source-row probabilities.

Arguments:

  • backend: Prepared destination backend.
  • probabilities: Nonempty vector of probabilities in [0, 1]. Rows may represent sampled probabilities rather than a complete allele-count grid.

Returns:

A Transition with float64 probability leaves and source_gene_copies = len(probabilities) - 1.

Raises:

ValueError if the probability array is empty or not one-dimensional. TypeError for an invalid backend. Probability values are not range-checked.

expect(self, values: Array) -> Array ¤

Compute destination-value expectations at each source state.

Arguments:

  • values: Finite real vector of length backend.gene_copies + 1, converted to float64. Signed entries are supported.

Returns:

A float64 vector of length source_gene_copies + 1.

Raises:

ValueError for incompatible input or map-output shapes.

advance(self, distribution: Array) -> Array ¤

Propagate source mass by one generation.

Arguments:

  • distribution: Finite real vector of length source_gene_copies + 1, converted to float64. Signed entries are supported.

Returns:

A float64 vector of length backend.gene_copies + 1.

Raises:

ValueError for incompatible input or map-output shapes.

expect_segregating(self, values: Array) -> Array ¤

Compute interior expectations with endpoint values set to zero.

Arguments:

  • values: Finite real vector of length backend.gene_copies - 1, converted to float64. Signed entries are supported.

Returns:

A float64 vector of length source_gene_copies - 1. No conditioning on survival.

Raises:

ValueError for incompatible input or map-output shapes.

advance_segregating(self, distribution: Array) -> Array ¤

Propagate interior mass, discarding endpoint hits.

Arguments:

  • distribution: Finite real vector of length source_gene_copies - 1, converted to float64. Signed entries are supported.

Returns:

A float64 vector of length backend.gene_copies - 1. No renormalization or re-entry from endpoints.

Raises:

ValueError for incompatible input or map-output shapes.

as_lineax(self, *, segregating: bool = False) -> FunctionLinearOperator ¤

Expose the expectation operation as a Lineax operator.

Arguments:

  • segregating: Exclude endpoint states when true; defaults to false.

Returns:

A rectangular lineax.FunctionLinearOperator from destination values to source expectations. Its transpose advances mass. No symmetry or definiteness tags are attached.

Raises:

ValueError for incompatible dimensions or map-output shape.

expected_visits(self, *, solver: typing.Literal['gmres', 'bicgstab'] = 'gmres', rtol: float = 1e-09, atol: float = 1e-11, max_steps: int = 2000, restart: int = 50) -> ExpectedVisitsResult ¤

Compute interior visits from one copy, including the initial generation.

Arguments:

  • solver: "gmres" (default) or "bicgstab".
  • rtol: Nonnegative relative tolerance; default 1e-9.
  • atol: Nonnegative absolute tolerance; default 1e-11. Both tolerances must be finite and at least one must be positive.
  • max_steps: Positive iteration limit; default 2000 (outer steps for GMRES).
  • restart: Positive GMRES subspace size; default 50, capped at the number of interior states. Validated but unused for BiCGStab.

Returns:

An ExpectedVisitsResult with float64 visits of length backend.gene_copies - 1, plus convergence diagnostics.

Raises:

ValueError for a nonsquare transition, invalid solver controls, or map shape.

Numerical limitations

The internal solve uses throw=False; numerical failure is returned. Check converged before using visits; under vmap, check each result. Visits stop at the first endpoint hit, even with mutation. Global differentiability of solves is not certified.

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

Occupation counts and convergence diagnostics for the approximate operator.

The candidate occupation vector solves \((I-Q^T)g=e_1\), including generation zero. The convergence flag combines solver status, finiteness, and a residual check against the approximate operator.

Fields:

  • visits: Float64 vector of expected interior visits, including generation zero. Entry zero corresponds to one mutant copy.
  • status, stats: Unmodified Lineax result code and iteration diagnostics.
  • residual_norm: Scalar infinity norm of \((I-Q^T)g-e_1\).
  • converged: Scalar boolean combining solver status, finite visits and the explicit residual test. Check this rather than solver status alone.
__new__(_cls, visits: Array, status: RESULTS, stats: dict[str, jax.Array | int | None], residual_norm: Array, converged: Array) staticmethod ¤

Construct the result container without numerical validation.

Arguments:

  • visits: Candidate float64 expected visits of length gene_copies-1, including generation zero. Index zero represents one mutant copy.
  • status: Unmodified Lineax result code.
  • stats: Unmodified Lineax iteration diagnostics.
  • residual_norm: Infinity norm of \((I-Q^T)g-e_1\).
  • converged: Successful status, finite candidate and acceptable residual.

Returns:

An immutable result with the supplied fields.

Usage preconditions

Use visits only when converged is true. Inspect status, stats and the residual otherwise; increase iteration limits or revisit tolerances before retrying. Even a small residual is not an exact-model error certificate.