Skip to content

Frequency maps¤

Frequency maps convert parental allele frequencies into offspring allele probabilities before binomial sampling. Selection, mutation, and their compositions use the same interface.

Interface and custom maps¤

AbstractFrequencyMap

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

Interface for transforming source frequencies into offspring probabilities.

A map takes allele frequencies \(x\) to offspring probabilities \(q(x)\), preserving the input shape. Concrete maps implement apply; __call__ provides shorthand. Implementations must be pure and JAX-compatible; differentiable parameters are stored as array leaves.

apply(self, frequencies: Array) -> Array ¤

Transform allele frequencies into offspring probabilities.

Arguments:

  • frequencies: Float64 array of allele frequencies in [0, 1].

Returns:

An array of finite probabilities in [0, 1] with the same shape as frequencies.

Input domain

Finiteness and the [0, 1] range are caller preconditions. Maps do not clip or normalize their outputs.

__call__(self, frequencies: Array) -> Array ¤

Call apply using function syntax.

Arguments:

  • frequencies: Float64 array of allele frequencies in [0, 1].

Returns:

The result of self.apply(frequencies).

drifft.FrequencyMap(drifft.AbstractFrequencyMap) [source] ¤

Wrap a custom JAX-compatible function in the frequency-map interface.

The wrapped function must be pure, preserve shape, and return finite probabilities in [0, 1]. Its parameters stay connected to differentiation; its output is returned without conversion, clipping, or normalization.

__init__(self, function: collections.abc.Callable[[jax.Array], jax.Array]) -> None ¤

Initialize FrequencyMap.

Arguments:

  • function: Callable returning one probability per input frequency. An Equinox callable can carry dynamic parameters. For closures, construct the function inside the differentiated calculation to retain parameter links.

Returns:

None.

Raises:

TypeError for non-callable functions. Output shape and probability ranges follow AbstractFrequencyMap's contract; transitions check output shape.

Selection¤

drifft.DiploidSelection(drifft.AbstractFrequencyMap) [source] ¤

Diploid viability selection with independent genotype fitness reductions.

Genotypes 0/0, 0/1, and 1/1 have fitnesses \(1\), \(1-s_{\mathrm{het}}\), and \(1-s_{\mathrm{hom}}\). Under random mating, offspring allele frequency is

\[q(x)=\frac{(1-s_{\mathrm{hom}})x^2+(1-s_{\mathrm{het}})x(1-x)} {(1-x)^2+2(1-s_{\mathrm{het}})x(1-x)+(1-s_{\mathrm{hom}})x^2}.\]

Both coefficients are explicit; there is no default dominance relationship or clipping. Negative coefficients represent fitness advantages. Mutation is a separate map, applied in the order specified by Composition.

__init__(self, s_het: typing.Any, s_hom: typing.Any) -> None ¤

Initialize DiploidSelection.

Arguments:

  • s_het: Heterozygote fitness reduction, finite and at most one.
  • s_hom: Derived-homozygote fitness reduction, finite and strictly below one.

Both parameters become dynamic float64 scalar leaves. Enable JAX x64 before construction. The strict homozygote bound keeps mean fitness positive at fixation; a lethal derived homozygote is outside this map's domain.

Returns:

None.

Raises:

ValueError for nonscalar coefficients.

Parameter domain

Value constraints are caller preconditions, including under JIT and vmap. Outputs and parameters are not clipped. All backends support this map. Jacobi's reduced Krylov assembly requires both coefficients to have absolute value at most its configured max_strength, which is strictly below one.

drifft.DominanceSelection(drifft.AbstractFrequencyMap) [source] ¤

Diploid selection with fitnesses \(1\), \(1-hs\), and \(1-s\).

Uses the DiploidSelection numerical rule with coefficients derived from the current parameter leaves on every call.

__init__(self, s: typing.Any, h: typing.Any) -> None ¤

Initialize DominanceSelection.

Arguments:

  • s: Derived-homozygote fitness reduction, finite and strictly below one.
  • h: Finite dominance coefficient, with finite h * s <= 1.

Parameters become dynamic float64 scalar leaves. Enable JAX x64 before construction. Value constraints are caller preconditions, including under JIT and vmap; no clipping is applied. FFT and NUFFT support this rule. Jacobi checks that both derived selection coefficients have absolute value at most its configured max_strength.

Returns:

None.

Raises:

ValueError for nonscalar parameters.

drifft.AdditiveSelection(drifft.AbstractFrequencyMap) [source] ¤

Additive diploid fitness effects.

Genotype fitnesses are \(1\), \(1-s_{\mathrm{het}}\), and \(1-2s_{\mathrm{het}}\).

Uses the DiploidSelection numerical rule with coefficients derived from the current parameter leaves on every call.

__init__(self, s_het: typing.Any) -> None ¤

Initialize AdditiveSelection.

Arguments:

  • s_het: Heterozygote fitness reduction, finite and strictly below 0.5, with finite derived coefficient 2 * s_het.

Parameters become dynamic float64 scalar leaves. Enable JAX x64 before construction. Value constraints are caller preconditions, including under JIT and vmap; no clipping is applied. FFT and NUFFT support this rule. Jacobi checks that both derived selection coefficients have absolute value at most its configured max_strength.

Returns:

None.

Raises:

ValueError for nonscalar parameters.

drifft.RecessiveSelection(drifft.AbstractFrequencyMap) [source] ¤

Completely recessive fitness effects with fitnesses \(1\), \(1\), and \(1-s\).

Uses the DiploidSelection numerical rule with coefficients derived from the current parameter leaves on every call.

__init__(self, s: typing.Any) -> None ¤

Initialize RecessiveSelection.

Arguments:

  • s: Derived-homozygote fitness reduction, finite and strictly below one.

Parameters become dynamic float64 scalar leaves. Enable JAX x64 before construction. Value constraints are caller preconditions, including under JIT and vmap; no clipping is applied. FFT and NUFFT support this rule. Jacobi checks that both derived selection coefficients have absolute value at most its configured max_strength.

Returns:

None.

Raises:

ValueError for nonscalar parameters.

drifft.DominantSelection(drifft.AbstractFrequencyMap) [source] ¤

Completely dominant fitness effects with fitnesses \(1\), \(1-s\), and \(1-s\).

Uses the DiploidSelection numerical rule with coefficients derived from the current parameter leaves on every call.

__init__(self, s: typing.Any) -> None ¤

Initialize DominantSelection.

Arguments:

  • s: Fitness reduction for carriers, finite and strictly below one.

Parameters become dynamic float64 scalar leaves. Enable JAX x64 before construction. Value constraints are caller preconditions, including under JIT and vmap; no clipping is applied. FFT and NUFFT support this rule. Jacobi checks that both derived selection coefficients have absolute value at most its configured max_strength.

Returns:

None.

Raises:

ValueError for nonscalar parameters.

drifft.Underdominance(drifft.AbstractFrequencyMap) [source] ¤

Symmetric heterozygote disadvantage, callable on allele frequencies.

Genotypes 0/0, 0/1, and 1/1 have fitnesses \(1\), \(1-s\), and \(1\). This is the constrained diploid rule DiploidSelection(s_het=s, s_hom=0); both homozygotes have equal fitness. AdditiveSelection(s_het=s) instead assigns derived homozygotes fitness \(1-2s\).

For allele frequency \(x\) and heterozygote disadvantage \(s\),

\[q_s(x) = \frac{x-sx(1-x)}{1-2sx(1-x)}.\]

Uses the shared diploid numerical function. The named class also selects the specialized underdominance assembly supported by the Jacobi backend.

__init__(self, strength: typing.Any) -> None ¤

Initialize Underdominance.

Arguments:

  • strength: Heterozygote fitness reduction in [0, 1], converted to a float64 JAX leaf. Changing its value does not change model structure.

Returns:

None.

Raises:

ValueError for nonscalar strength.

Parameter domain

Strength must be finite and in [0, 1]. This value constraint is unchecked, including under JIT and vmap. Enable JAX x64 before construction.

drifft.Overdominance(drifft.AbstractFrequencyMap) [source] ¤

Symmetric heterozygote advantage with fitnesses \(1\), \(1+s\), and \(1\).

Uses the DiploidSelection numerical rule with coefficients derived from the current parameter leaves on every call.

__init__(self, strength: typing.Any) -> None ¤

Initialize Overdominance.

Arguments:

  • strength: Finite, nonnegative heterozygote fitness advantage.

Parameters become dynamic float64 scalar leaves. Enable JAX x64 before construction. Value constraints are caller preconditions, including under JIT and vmap; no clipping is applied. FFT and NUFFT support this rule. Jacobi checks that both derived selection coefficients have absolute value at most its configured max_strength.

Returns:

None.

Raises:

ValueError for nonscalar parameters.

drifft.MultiplicativeSelection(drifft.AbstractFrequencyMap) [source] ¤

Multiplicative fitness effects with fitnesses \(1\), \(r\), and \(r^2\).

This is the diploid rule with coefficients \(1-r\) and \(1-r^2\). Cancelling the common factor gives \(q(x)=rx/(1-x+rx)\), avoiding cancellation in the coefficients when fitness is small and overflow from squaring large \(r\).

__init__(self, r: typing.Any) -> None ¤

Initialize MultiplicativeSelection.

Arguments:

  • r: Finite, strictly positive fitness multiplier per derived allele.

Parameters become dynamic float64 scalar leaves. Enable JAX x64 before construction. Value constraints are caller preconditions, including under JIT and vmap; no clipping is applied. FFT and NUFFT support this rule. Jacobi checks that both derived selection coefficients have absolute value at most its configured max_strength.

Returns:

None.

Raises:

ValueError for nonscalar parameters.

drifft.Neutral(drifft.AbstractFrequencyMap) [source] ¤

Identity frequency map for neutral reproduction without mutation.

This is DiploidSelection(0, 0), simplified to \(q(x)=x\) to preserve the input dtype and avoid unnecessary arithmetic.

__init__(self) -> None ¤

Initialize Neutral.

Arguments:

None.

Returns:

None.

Mutation and composition¤

drifft.Mutation(drifft.AbstractFrequencyMap) [source] ¤

Apply forward and back mutation to an allele frequency.

For allele-1 frequency \(p\), forward rate \(u\), and backward rate \(v\),

\[q(p) = p(1-v) + (1-p)u.\]
__init__(self, forward_rate: typing.Any = 0.0, backward_rate: typing.Any = 0.0) -> None ¤

Initialize Mutation.

Arguments:

  • forward_rate: Per-generation probability of mutation from allele 0 to 1.
  • backward_rate: Per-generation probability of mutation from allele 1 to 0.

Both rates default to zero and become dynamic float64 scalar leaves. Each rate lies in [0, 1]; their sum need not be less than one.

Returns:

None.

Raises:

ValueError for nonscalar rates.

Parameter domain

Rates must be finite and in [0, 1]. These value constraints are unchecked, including under JIT. Enable JAX x64 before construction.

drifft.Composition(drifft.AbstractFrequencyMap) [source] ¤

Apply frequency maps in the supplied left-to-right order.

For component maps \(q_1,\ldots,q_k\), the combined map is

\[q(x)=q_k(\cdots q_2(q_1(x))\cdots).\]

Selection and mutation generally do not commute, so their order is part of the model. Each component must preserve shape and return finite probabilities in [0, 1]; component failures propagate without additional validation.

__init__(self, *maps: AbstractFrequencyMap) -> None ¤

Bind the ordered frequency maps.

Arguments:

  • *maps: One or more AbstractFrequencyMap instances, applied left to right.

Returns:

None.

Raises:

ValueError for no maps; TypeError for a non-map component.