def test_prob_coal_counts(self):
self.assertAlmostEqual(
coal.prob_coal_counts(2, 2, 100, 1000),
0.904837418036)
self.assertAlmostEqual(
coal.prob_coal_counts(5, 2, 100, 1000),
0.0184034834527)
def test_coal_counts(self):
b = 1
t = 1000.0
n = 1000
for a in xrange(1, 100):
i = coal.prob_coal_counts(a, b, t, n)
j = coal.cdf_mrca(t, a, n)
fequal(i, j)
for a in xrange(1, 10):
i = sum(coal.prob_coal_counts(a, b, t, n)
for b in xrange(1, a+1))
fequal(i, 1.0)
def test_coal_counts2(self):
b = 3
t = 1000.0
n = 1000
for b in xrange(1, 10):
for a in xrange(b, 10):
i = coal.prob_coal_counts(a, b, t, n)
j = coal.prob_coal_counts_slow(a, b, t, n)
fequal(i, j)
def test_cdf_mrca2(self):
n = 1000
k = 6
step = 10
x = list(frange(0, 5000, step))
y = [coal.prob_mrca(i, k, n) * step for i in x]
y2 = cumsum(y)
y3 = [coal.cdf_mrca(t, k, n) for t in x]
y4 = [coal.prob_coal_counts(k, 1, t, n) for t in x]
fequals(y2, y3, eabs=.01)
fequals(y3, y4)
def sample_coal_cond_counts_cdf(a, b, t, n, p=None):
"""
Samples the next coalescent between 'a' lineages in a population size of
'n', conditioned on there being 'b' lineages at time 't'.
"""
# this code solves this equation for t
# cdf(t) - p = 0
# where p ~ U(0, 1)
import scipy.optimize
if p is None:
p = random.random()
# compute constants
lama = -a*(a-1)/2.0/n
C0 = stats.prod((b+y)*(a-1-y)/(a-1+y) for y in xrange(b))
c = -b*(b-1)/2.0/n
d = 1.0/stats.factorial(b) * (-lama) / coal.prob_coal_counts(a, b, t, n)
# CDF(t) - p
def f(x):
if x <= 0:
return x - p
if x >= t:
return 1.0 - p + (x - t)
C = C0
s = exp(c*t) * (exp((lama-c)*x)-1.0) / (lama-c) * C
for k in xrange(b+1, a):
k1 = k - 1
lam = -k*k1/2.0/n
C = (b+k1)*(a-1-k1)/(a-1+k1)/(b-k) * C
s += exp(lam*t) * (exp((lama-lam)*x) - 1.0) / (lama - lam) \
* (2*k-1) / (k1+b) * C
print "f", p, x, s * d - p
return s * d - p
return f
