22 lines
991 B
Plaintext
22 lines
991 B
Plaintext
module horse/prob/dist
|
|
|
|
import std/num/decimal
|
|
|
|
tail fun pb-step(pn: list<decimal>, pi: decimal, pmfkm1: decimal, pmf: list<decimal>, next: ctx<list<decimal>>): list<decimal>
|
|
match pn
|
|
Nil -> next ++. Nil // final step overall
|
|
Cons(_, pp) -> match pmf
|
|
Cons(pmfk, pmf') ->
|
|
val next' = next ++ ctx Cons(pi * pmfkm1 + (1.decimal - pi) * pmfk, hole)
|
|
pb-step(pp, pi, pmfk, pmf', next')
|
|
Nil -> next ++. Cons(pi * pmfkm1, Nil) // last step of this iteration
|
|
|
|
// Given `n` different Bernoulli processes with respective probabilities in `pn`,
|
|
// find the distribution of `k` successes for `k` ranging from 0 to `n` inclusive.
|
|
// The index in the result list corresponds to `k`.
|
|
pub fun pmf/poisson-binomial(pn: list<decimal>): list<decimal>
|
|
pn.foldl([1.decimal]) fn(pmf, pi)
|
|
match pmf
|
|
Cons(pmf0, pmf') -> pb-step(pn, pi, pmf0, pmf', ctx Cons((1.decimal - pi) * pmf0, hole))
|
|
Nil -> impossible("fold started with non-empty pmf but got empty pmf")
|