route/vendor/github.com/aclements/go-moremath/stats/hypergdist.go

102 lines
2.5 KiB
Go

// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package stats
import (
"math"
"github.com/aclements/go-moremath/mathx"
)
// HypergeometicDist is a hypergeometric distribution.
type HypergeometicDist struct {
// N is the size of the population. N >= 0.
N int
// K is the number of successes in the population. 0 <= K <= N.
K int
// Draws is the number of draws from the population. This is
// usually written "n", but is called Draws here because of
// limitations on Go identifier naming. 0 <= Draws <= N.
Draws int
}
// PMF is the probability of getting exactly int(k) successes in
// d.Draws draws with replacement from a population of size d.N that
// contains exactly d.K successes.
func (d HypergeometicDist) PMF(k float64) float64 {
ki := int(math.Floor(k))
l, h := d.bounds()
if ki < l || ki > h {
return 0
}
return d.pmf(ki)
}
func (d HypergeometicDist) pmf(k int) float64 {
return math.Exp(mathx.Lchoose(d.K, k) + mathx.Lchoose(d.N-d.K, d.Draws-k) - mathx.Lchoose(d.N, d.Draws))
}
// CDF is the probability of getting int(k) or fewer successes in
// d.Draws draws with replacement from a population of size d.N that
// contains exactly d.K successes.
func (d HypergeometicDist) CDF(k float64) float64 {
// Based on Klotz, A Computational Approach to Statistics.
ki := int(math.Floor(k))
l, h := d.bounds()
if ki < l {
return 0
} else if ki >= h {
return 1
}
// Use symmetry to compute the smaller sum.
flip := false
if ki > (d.Draws+1)/(d.N+1)*(d.K+1) {
flip = true
ki = d.K - ki - 1
d.Draws = d.N - d.Draws
}
p := d.pmf(ki) * d.sum(ki)
if flip {
p = 1 - p
}
return p
}
func (d HypergeometicDist) sum(k int) float64 {
const epsilon = 1e-14
sum, ak := 1.0, 1.0
L := maxint(0, d.Draws+d.K-d.N)
for dk := 1; dk <= k-L && ak/sum > epsilon; dk++ {
ak *= float64(1+k-dk) / float64(d.Draws-k+dk)
ak *= float64(d.N-d.K-d.Draws+k+1-dk) / float64(d.K-k+dk)
sum += ak
}
return sum
}
func (d HypergeometicDist) bounds() (int, int) {
return maxint(0, d.Draws+d.K-d.N), minint(d.Draws, d.K)
}
func (d HypergeometicDist) Bounds() (float64, float64) {
l, h := d.bounds()
return float64(l), float64(h)
}
func (d HypergeometicDist) Step() float64 {
return 1
}
func (d HypergeometicDist) Mean() float64 {
return float64(d.Draws*d.K) / float64(d.N)
}
func (d HypergeometicDist) Variance() float64 {
return float64(d.Draws*d.K*(d.N-d.K)*(d.N-d.Draws)) /
float64(d.N*d.N*(d.N-1))
}