110 lines
3.8 KiB
TypeScript
110 lines
3.8 KiB
TypeScript
/**
|
|
* Punnett engine: combine two parent genotypes into the offspring genotype
|
|
* distribution, with exact fractional probabilities.
|
|
*
|
|
* Per locus, each parent contributes one of its two alleles with prob 1/2,
|
|
* giving 4 equally-likely gametic combinations that we fold into unordered
|
|
* (canonically-ordered) allele pairs. The full-genotype distribution is the
|
|
* product across all loci.
|
|
*
|
|
* Wildcard ("?") alleles are expanded uniformly over the locus' allele set
|
|
* before combining, so a parent known only by phenotype can still be paired.
|
|
*/
|
|
import { add, frac, multiply, ONE, type Fraction } from './fraction'
|
|
import { LOCI, LOCUS_ORDER, type LocusKey } from './loci'
|
|
import {
|
|
canonicalPair,
|
|
toDisplayString,
|
|
WILDCARD,
|
|
type AllelePair,
|
|
type Genotype,
|
|
} from './genotype'
|
|
|
|
/** A probability distribution over outcomes of type T (keyed by a string). */
|
|
export interface DistEntry<T> {
|
|
readonly value: T
|
|
readonly probability: Fraction
|
|
}
|
|
|
|
/** Expand a (possibly wildcard) parent allele pair into weighted concrete alleles. */
|
|
function parentAlleleWeights(locus: LocusKey, pair: AllelePair): Map<string, Fraction> {
|
|
const weights = new Map<string, Fraction>()
|
|
const addWeight = (allele: string, w: Fraction) => {
|
|
weights.set(allele, add(weights.get(allele) ?? frac(0, 1), w))
|
|
}
|
|
const alleles = LOCI[locus].alleles
|
|
for (const a of pair) {
|
|
if (a === WILDCARD) {
|
|
// Unknown allele: uniform over the locus set, each contributing 1/2 of the gamete.
|
|
const share = frac(1, 2 * alleles.length)
|
|
for (const concrete of alleles) addWeight(concrete, share)
|
|
} else {
|
|
addWeight(a, frac(1, 2))
|
|
}
|
|
}
|
|
return weights
|
|
}
|
|
|
|
/** Offspring allele-pair distribution at a single locus. */
|
|
export function combineLocus(
|
|
locus: LocusKey,
|
|
father: AllelePair,
|
|
mother: AllelePair,
|
|
): Map<string, DistEntry<AllelePair>> {
|
|
const fa = parentAlleleWeights(locus, father)
|
|
const mo = parentAlleleWeights(locus, mother)
|
|
const result = new Map<string, DistEntry<AllelePair>>()
|
|
for (const [fAllele, fw] of fa) {
|
|
for (const [mAllele, mw] of mo) {
|
|
const pair = canonicalPair(locus, fAllele, mAllele)
|
|
const key = pair[0] + pair[1]
|
|
const p = multiply(fw, mw)
|
|
const existing = result.get(key)
|
|
result.set(key, {
|
|
value: pair,
|
|
probability: existing ? add(existing.probability, p) : p,
|
|
})
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
/** Full offspring genotype distribution (Cartesian product across loci). */
|
|
export function combineGenotypes(
|
|
father: Genotype,
|
|
mother: Genotype,
|
|
): DistEntry<Genotype>[] {
|
|
// Start with one certain empty partial genotype, then fold in each locus.
|
|
let dist: DistEntry<Record<LocusKey, AllelePair>>[] = [
|
|
{ value: {} as Record<LocusKey, AllelePair>, probability: ONE },
|
|
]
|
|
for (const locus of LOCUS_ORDER) {
|
|
const locusDist = combineLocus(locus, father[locus], mother[locus])
|
|
const next: DistEntry<Record<LocusKey, AllelePair>>[] = []
|
|
for (const partial of dist) {
|
|
for (const entry of locusDist.values()) {
|
|
next.push({
|
|
value: { ...partial.value, [locus]: entry.value },
|
|
probability: multiply(partial.probability, entry.probability),
|
|
})
|
|
}
|
|
}
|
|
dist = next
|
|
}
|
|
return dist.map((e) => ({ value: e.value as Genotype, probability: e.probability }))
|
|
}
|
|
|
|
/** Merge distribution entries that share the same display string, summing probs. */
|
|
export function mergeByGenotype(dist: DistEntry<Genotype>[]): DistEntry<Genotype>[] {
|
|
const merged = new Map<string, DistEntry<Genotype>>()
|
|
for (const entry of dist) {
|
|
const key = toDisplayString(entry.value)
|
|
const existing = merged.get(key)
|
|
merged.set(key, {
|
|
value: entry.value,
|
|
probability: existing ? add(existing.probability, entry.probability) : entry.probability,
|
|
})
|
|
}
|
|
return Array.from(merged.values())
|
|
}
|