def rate(partition): b, k_vec, s = computeSignature(partition) partition_multiset = re.partitionToMultiset(partition) mergerRate = mergerRatesCoalescent(b, k_vec, s) factorSizes = np.prod(map(mp.fac, partition)) factorCounts = np.prod(map(mp.fac, partition_multiset)) # print factorSizes, factorCounts, mergerRate return (mp.fac(n) / factorSizes*factorCounts ) * mergerRate
def p_recursions(N, coalescentType, args): if coalescentType == 'xi_beta': alpha = args[0] def lambdaMergerRate(b, k): return betaMergerRate(b, k, alpha) def fourwayMergerRate(b, k_vec, s): return fourWayMergerRate(b, k_vec, s, lambdaMergerRate) L = [transitionRatesNumberPartitions(n, fourwayMergerRate, 4) for n in range(1, N+1)] Q = Q_matrix_Xi([[]]+L) P = P_matrix(Q) g = g_matrix(P, Q) p = np.zeros((N+1, N+1, N+1), dtype=mp.mpf) p[1, 1, 1] = mp.mpf('1') # print L, "\n" for n, jumpsFrom_n in enumerate(L[1:], start=2): p[n, n, 1] = mp.mpf('1') for n1, jumpsFrom_n_to_n1 in enumerate(jumpsFrom_n[1:-1], start=1): # print n, n1, jumpsFrom_n_to_n1 for lam_multi, rate in [( re.partitionToMultiset(x[0]), x[2] ) for x in jumpsFrom_n_to_n1]: jumpProb = -rate/Q[n, n] for b1 in range(1, n1): for lamSub, b in re.subpartitionsMultiset(lam_multi, b1): lamSubFactor = re.subpartitionProb(lam_multi, lamSub, n1, b1) for k in range(2, min(n1-b1+1, n-b+1) +1): kFactor = p[n1, k, b1] * g[n1, k]/g[n, k] p[n, k, b] += jumpProb * kFactor * lamSubFactor print n,n1,k,b,b1,'\n',lam_multi,'\n',lamSub,'\n' # print '\n' #for testing if things add up # for n, l1 in enumerate(p): # for k, l2 in enumerate(l1): # if sum(l2)>0: # print (n, k, sum(l2)) return p, g, Q, P,L