Commit aa876f51 authored by Ryan Gutenkunst's avatar Ryan Gutenkunst
Browse files

Merge branch 'master' of https://bitbucket.org/gutenkunstlab/dadi

parents 90cecc75 c0466d00
......@@ -118,6 +118,28 @@ def _inject_mutations_4D(phi, dt, xx, yy, zz, aa, theta0, frozen1, frozen2, froz
phi[0,0,0,1] += dt/aa[1] * theta0/2 * 16/((aa[2] - aa[0]) * xx[1] * yy[1] * zz[1])
return phi
def _inject_mutations_5D(phi, dt, xx, yy, zz, aa, theta0, frozen1, frozen2, frozen3, frozen4):
"""
Inject novel mutations for a timestep.
"""
# Population 1
# Normalization based on the multi-dimensional trapezoid rule is
# implemented ************** here ***************
if not frozen1:
phi[1,0,0,0,0] += dt/xx[1] * theta0/2 * 32/((xx[2] - xx[0]) * yy[1] * zz[1] * aa[1] * bb[1])
# Population 2
if not frozen2:
phi[0,1,0,0,0] += dt/yy[1] * theta0/2 * 32/((yy[2] - yy[0]) * xx[1] * zz[1] * aa[1] * bb[1])
# Population 3
if not frozen3:
phi[0,0,1,0,0] += dt/zz[1] * theta0/2 * 32/((zz[2] - zz[0]) * xx[1] * yy[1] * aa[1] * bb[1])
# Population 4
if not frozen4:
phi[0,0,0,1,0] += dt/aa[1] * theta0/2 * 32/((aa[2] - aa[0]) * xx[1] * yy[1] * zz[1] * bb[1])
if not frozen5:
phi[0,0,0,0,1] += dt/aa[1] * theta0/2 * 32/((bb[2] - bb[0]) * xx[1] * yy[1] * zz[1] * aa[1])
return phi
def _compute_dt(dx, nu, ms, gamma, h):
"""
Compute the appropriate timestep given the current demographic params.
......@@ -230,7 +252,8 @@ def one_pop(phi, xx, T, nu=1, gamma=0, h=0.5, theta0=1.0, initial_t=0,
def two_pops(phi, xx, T, nu1=1, nu2=1, m12=0, m21=0, gamma1=0, gamma2=0,
h1=0.5, h2=0.5, theta0=1, initial_t=0, frozen1=False,
frozen2=False, nomut1=False, nomut2=False, enable_cuda_const=False):
frozen2=False, nomut1=False, nomut2=False, enable_cuda_const=False,
transpose=False):
"""
Integrate a 2-dimensional phi foward.
......@@ -278,7 +301,7 @@ def two_pops(phi, xx, T, nu1=1, nu2=1, m12=0, m21=0, gamma1=0, gamma2=0,
'migration to or from it.')
vars_to_check = [nu1,nu2,m12,m21,gamma1,gamma2,h1,h2,theta0]
if numpy.all([numpy.isscalar(var) for var in vars_to_check]):
if False and numpy.all([numpy.isscalar(var) for var in vars_to_check]):
# Constant integration with CUDA turns out to be slower,
# so we only use it in specific circumsances.
if not cuda_enabled or (cuda_enabled and enable_cuda_const):
......@@ -332,12 +355,22 @@ def two_pops(phi, xx, T, nu1=1, nu2=1, m12=0, m21=0, gamma1=0, gamma2=0,
_inject_mutations_2D(phi, this_dt, xx, yy, theta0, frozen1, frozen2,
nomut1, nomut2)
if not frozen1:
phi = int_c.implicit_2Dx(phi, xx, yy, nu1, m12, gamma1, h1,
this_dt, use_delj_trick)
if not frozen2:
phi = int_c.implicit_2Dy(phi, xx, yy, nu2, m21, gamma2, h2,
this_dt, use_delj_trick)
if not transpose:
if not frozen1:
phi = int_c.implicit_2Dx(phi, xx, yy, nu1, m12, gamma1, h1,
this_dt, use_delj_trick)
if not frozen2:
phi = int_c.implicit_2Dy(phi, xx, yy, nu2, m21, gamma2, h2,
this_dt, use_delj_trick)
else:
phi = phi.transpose()
if not frozen1:
phi = int_c.implicit_2Dy(phi, xx, yy, nu1, m12, gamma1, h1,
this_dt, use_delj_trick)
phi = phi.transpose()
if not frozen2:
phi = int_c.implicit_2Dy(phi, xx, yy, nu2, m21, gamma2, h2,
this_dt, use_delj_trick)
current_t = next_t
return phi
......@@ -346,7 +379,7 @@ def three_pops(phi, xx, T, nu1=1, nu2=1, nu3=1,
m12=0, m13=0, m21=0, m23=0, m31=0, m32=0,
gamma1=0, gamma2=0, gamma3=0, h1=0.5, h2=0.5, h3=0.5,
theta0=1, initial_t=0, frozen1=False, frozen2=False,
frozen3=False, enable_cuda_const=False):
frozen3=False, enable_cuda_const=False, transpose=False):
"""
Integrate a 3-dimensional phi foward.
......@@ -391,7 +424,7 @@ def three_pops(phi, xx, T, nu1=1, nu2=1, nu3=1,
vars_to_check = [nu1,nu2,nu3,m12,m13,m21,m23,m31,m32,gamma1,gamma2,
gamma3,h1,h2,h3,theta0]
if numpy.all([numpy.isscalar(var) for var in vars_to_check]):
if False and numpy.all([numpy.isscalar(var) for var in vars_to_check]):
if not cuda_enabled or (cuda_enabled and enable_cuda_const):
return _three_pops_const_params(phi, xx, T, nu1, nu2, nu3,
m12, m13, m21, m23, m31, m32,
......@@ -461,15 +494,29 @@ def three_pops(phi, xx, T, nu1=1, nu2=1, nu3=1,
_inject_mutations_3D(phi, this_dt, xx, yy, zz, theta0,
frozen1, frozen2, frozen3)
if not frozen1:
phi = int_c.implicit_3Dx(phi, xx, yy, zz, nu1, m12, m13,
gamma1, h1, this_dt, use_delj_trick)
if not frozen2:
phi = int_c.implicit_3Dy(phi, xx, yy, zz, nu2, m21, m23,
gamma2, h2, this_dt, use_delj_trick)
if not frozen3:
phi = int_c.implicit_3Dz(phi, xx, yy, zz, nu3, m31, m32,
gamma3, h3, this_dt, use_delj_trick)
if not transpose:
if not frozen1:
phi = int_c.implicit_3Dx(phi, xx, yy, zz, nu1, m12, m13,
gamma1, h1, this_dt, use_delj_trick)
if not frozen2:
phi = int_c.implicit_3Dy(phi, xx, yy, zz, nu2, m21, m23,
gamma2, h2, this_dt, use_delj_trick)
if not frozen3:
phi = int_c.implicit_3Dz(phi, xx, yy, zz, nu3, m31, m32,
gamma3, h3, this_dt, use_delj_trick)
else:
phi = phi.transpose(1,2,0)
if not frozen1:
phi = int_c.implicit_3Dz(phi, xx, yy, zz, nu1, m12, m13,
gamma1, h1, this_dt, use_delj_trick)
phi = phi.transpose(2,1,0)
if not frozen2:
phi = int_c.implicit_3Dz(phi, xx, yy, zz, nu2, m21, m23,
gamma2, h2, this_dt, use_delj_trick)
phi = phi.transpose(0,2,1)
if not frozen3:
phi = int_c.implicit_3Dz(phi, xx, yy, zz, nu3, m31, m32,
gamma3, h3, this_dt, use_delj_trick)
current_t = next_t
return phi
......@@ -517,7 +564,8 @@ def four_pops(phi, xx, T, nu1=1, nu2=1, nu3=1, nu4=1,
if (frozen1 and (m12 != 0 or m21 != 0 or m13 !=0 or m31 != 0 or m41 != 0 or m14 != 0))\
or (frozen2 and (m12 != 0 or m21 != 0 or m23 != 0 or m32 != 0 or m24 != 0 or m42 != 0))\
or (frozen3 and (m13 != 0 or m31 != 0 or m23 !=0 or m32 != 0 or m34 != 0 or m43 != 0)):
or (frozen3 and (m13 != 0 or m31 != 0 or m23 !=0 or m32 != 0 or m34 != 0 or m43 != 0))\
or (frozen4 and (m14 != 0 or m41 != 0 or m24 !=0 or m42 != 0 or m34 != 0 or m43 != 0)):
raise ValueError('Population cannot be frozen and have non-zero '
'migration to or from it.')
aa = zz = yy = xx
......@@ -597,6 +645,144 @@ def four_pops(phi, xx, T, nu1=1, nu2=1, nu3=1, nu4=1,
current_t = next_t
return phi
def five_pops(phi, xx, T, nu1=1, nu2=1, nu3=1, nu4=1, nu5=1,
m12=0, m13=0, m14=0, m15=0, m21=0, m23=0, m24=0, m25=0,
m31=0, m32=0, m34=0, m35=0, m41=0, m42=0, m43=0, m45=0,
m51=0, m52=0, m53=0, m54=0,
gamma1=0, gamma2=0, gamma3=0, gamma4=0, gamma5=0,
h1=0.5, h2=0.5, h3=0.5, h4=0.5, h5=0,
theta0=1, initial_t=0, frozen1=False, frozen2=False,
frozen3=False, frozen4=False, frozen5=0,
enable_cuda_const=False):
"""
Integrate a 5-dimensional phi foward.
phi: Initial 5-dimensional phi
xx: 1-dimensional grid upon (0,1) overwhich phi is defined. It is assumed
that this grid is used in all dimensions.
nu's, gamma's, m's, and theta0 may be functions of time.
nu1,nu2,nu3,nu4,nu5: Population sizes
gamma1,gamma2,gamma3,gamma4,gamma5: Selection coefficients on *all* segregating alleles
h1,h2,h3,h4,h5: Dominance coefficients. h = 0.5 corresponds to genic selection.
m12,m13,m21,m23,m31,m32, ...: Migration rates. Note that m12 is the rate
*into 1 from 2*.
theta0: Proportional to ancestral size. Typically constant.
T: Time at which to halt integration
initial_t: Time at which to start integration. (Note that this only matters
if one of the demographic parameters is a function of time.)
enable_cuda_const: If True, enable CUDA integration with slower constant
parameter method. Likely useful only for benchmarking.
Note: Generalizing to different grids in different phi directions is
straightforward. The tricky part will be later doing the extrapolation
correctly.
"""
if T - initial_t == 0:
return phi
elif T - initial_t < 0:
raise ValueError('Final integration time T (%f) is less than '
'intial_time (%f). Integration cannot be run '
'backwards.' % (T, initial_t))
if (frozen1 and (m12 != 0 or m21 != 0 or m13 !=0 or m31 != 0 or m41 != 0 or m14 != 0 or m15 != 0 or m51 != 0))\
or (frozen2 and (m12 != 0 or m21 != 0 or m23 != 0 or m32 != 0 or m24 != 0 or m42 != 0 or m25 !=0 or m52 != 0))\
or (frozen3 and (m13 != 0 or m31 != 0 or m23 != 0 or m32 != 0 or m34 != 0 or m43 != 0 or m35 != 0 or m53 !=0))\
or (frozen4 and (m14 != 0 or m41 != 0 or m24 != 0 or m42 != 0 or m34 != 0 or m43 != 0 or m45 != or m54 != 0))\
or (frozen5 and (m15 != 0 or m51 != 0 or m25 != 0 or m52 != 0 or m35 != 0 or m53 != 0 or m45 != or m54 != 0)):
raise ValueError('Population cannot be frozen and have non-zero '
'migration to or from it.')
bb = aa = zz = yy = xx
nu1_f, nu2_f = Misc.ensure_1arg_func(nu1), Misc.ensure_1arg_func(nu2)
nu3_f, nu4_f = Misc.ensure_1arg_func(nu3), Misc.ensure_1arg_func(nu4)
nu5_f = Misc.ensure_1arg_func(nu5)
gamma1_f, gamma2_f = Misc.ensure_1arg_func(gamma1), Misc.ensure_1arg_func(gamma2)
gamma3_f, gamma4_f = Misc.ensure_1arg_func(gamma3), Misc.ensure_1arg_func(gamma4)
gamma5_f = Misc.ensure_1arg_func(gamma5)
h1_f, h2_f = Misc.ensure_1arg_func(h1), Misc.ensure_1arg_func(h2)
h3_f, h4_f = Misc.ensure_1arg_func(h3), Misc.ensure_1arg_func(h4)
h5_f = Misc.ensure_1arg_func(h5)
m12_f, m13_f, m14_f, m15_f = Misc.ensure_1arg_func(m12), Misc.ensure_1arg_func(m13), Misc.ensure_1arg_func(m14), Misc.ensure_1arg_func(m15)
m21_f, m23_f, m24_f, m25_f = Misc.ensure_1arg_func(m21), Misc.ensure_1arg_func(m23), Misc.ensure_1arg_func(m24), Misc.ensure_1arg_func(m25)
m31_f, m32_f, m34_f, m35_f = Misc.ensure_1arg_func(m31), Misc.ensure_1arg_func(m32), Misc.ensure_1arg_func(m34), Misc.ensure_1arg_func(m35)
m41_f, m42_f, m43_f, m45_f = Misc.ensure_1arg_func(m41), Misc.ensure_1arg_func(m42), Misc.ensure_1arg_func(m43), Misc.ensure_1arg_func(m45)
m51_f, m52_f, m53_f, m54_f = Misc.ensure_1arg_func(m51), Misc.ensure_1arg_func(m52), Misc.ensure_1arg_func(m53), Misc.ensure_1arg_func(m54)
theta0_f = Misc.ensure_1arg_func(theta0)
#if cuda_enabled:
# import dadi.cuda
# phi = dadi.cuda.Integration._three_pops_temporal_params(phi, xx, T, initial_t,
# nu1_f, nu2_f, nu3_f, m12_f, m13_f, m21_f, m23_f, m31_f, m32_f,
# gamma1_f, gamma2_f, gamma3_f, h1_f, h2_f, h3_f,
# theta0_f, frozen1, frozen2, frozen3)
# return phi
current_t = initial_t
nu1, nu2, nu3, nu4, nu5 = nu1_f(current_t), nu2_f(current_t), nu3_f(current_t), nu4_f(current_t), nu5_f(current_t)
gamma1, gamma2, gamma3, gamma4, gamma5 = gamma1_f(current_t), gamma2_f(current_t), gamma3_f(current_t), gamma4_f(current_t), gamma5_f(current_t)
h1, h2, h3, h4, h5 = h1_f(current_t), h2_f(current_t), h3_f(current_t), h4_f(current_t), h5_f(current_t)
m12, m13, m14, m15 = m12_f(current_t), m13_f(current_t), m14_f(current_t), m15_f(current_t)
m21, m23, m24, m25 = m21_f(current_t), m23_f(current_t), m24_f(current_t), m25_f(current_t)
m31, m32, m34, m35 = m31_f(current_t), m32_f(current_t), m34_f(current_t), m35_f(current_t)
m41, m42, m43, m45 = m41_f(current_t), m42_f(current_t), m43_f(current_t), m45_f(current_t)
m51, m52, m53, m54 = m51_f(current_t), m52_f(current_t), m53_f(current_t), m54_f(current_t)
dx,dy,dz,da,db = numpy.diff(xx),numpy.diff(yy),numpy.diff(zz),numpy.diff(aa),numpy.diff(bb)
while current_t < T:
dt = min(_compute_dt(dx,nu1,[m12,m13,m14,m15],gamma1,h1),
_compute_dt(dy,nu2,[m21,m23,m24,m25],gamma2,h2),
_compute_dt(dz,nu3,[m31,m32,m34,m35],gamma3,h3),
_compute_dt(da,nu4,[m41,m42,m43,m45],gamma4,h4),
_compute_dt(da,nu4,[m51,m52,m53,m54],gamma5,h5))
this_dt = min(dt, T - current_t)
next_t = current_t + this_dt
nu1, nu2, nu3, nu4, nu5 = nu1_f(next_t), nu2_f(next_t), nu3_f(next_t), nu4_f(next_t), nu5_f(next_t)
gamma1, gamma2, gamma3, gamma4, gamma5 = gamma1_f(next_t), gamma2_f(next_t), gamma3_f(next_t), gamma4_f(next_t), gamma5_f(next_t)
h1, h2, h3, h4, h5 = h1_f(next_t), h2_f(next_t), h3_f(next_t), h4_f(next_t), h5_f(next_t)
m12, m13, m14, m15 = m12_f(next_t), m13_f(next_t), m14_f(next_t), m15_f(next_t)
m21, m23, m24, m25 = m21_f(next_t), m23_f(next_t), m24_f(next_t), m25_f(next_t)
m31, m32, m34, m35 = m31_f(next_t), m32_f(next_t), m34_f(next_t), m35_f(next_t)
m41, m42, m43, m45 = m41_f(next_t), m42_f(next_t), m43_f(next_t), m45_f(next_t)
m51, m52, m53, m54 = m51_f(next_t), m52_f(next_t), m53_f(next_t), m54_f(next_t)
theta0 = theta0_f(next_t)
if numpy.any(numpy.less([T,nu1,nu2,nu3,nu4,nu5,m12,m13,m14,m15,m21,
m23,m24,m25, m31,m32,m34,m35, m41,m42,m43,m45,
m51,m52,m53,m54, theta0],
0)):
raise ValueError('A time, population size, migration rate, or '
'theta0 is < 0. Has the model been mis-specified?')
if numpy.any(numpy.equal([nu1,nu2,nu3,nu4,nu5], 0)):
raise ValueError('A population size is 0. Has the model been '
'mis-specified?')
_inject_mutations_5D(phi, this_dt, xx, yy, zz, aa, bb, theta0,
frozen1, frozen2, frozen3, frozen4, frozen5)
if not frozen1:
phi = int_c.implicit_5Dx(phi, xx, yy, zz, aa, bb, nu1, m12, m13, m14, m15,
gamma1, h1, this_dt, use_delj_trick)
if not frozen2:
phi = int_c.implicit_5Dy(phi, xx, yy, zz, aa, bb, nu2, m21, m23, m24, m25,
gamma2, h2, this_dt, use_delj_trick)
if not frozen3:
phi = int_c.implicit_5Dz(phi, xx, yy, zz, aa, bb, nu3, m31, m32, m34, m35,
gamma3, h3, this_dt, use_delj_trick)
if not frozen4:
phi = int_c.implicit_5Da(phi, xx, yy, zz, aa, bb, nu4, m41, m42, m43, m45,
gamma4, h4, this_dt, use_delj_trick)
if not frozen5:
phi = int_c.implicit_5Db(phi, xx, yy, zz, aa, bb, nu5, m51, m52, m53, m54,
gamma5, h5, this_dt, use_delj_trick)
current_t = next_t
return phi
#
# Here are the python versions of the population genetic functions.
#
......
......@@ -109,7 +109,7 @@ def perturb_params(params, fold=1, lower_bound=None, upper_bound=None):
upper_bound: If not None, the resulting parameter set is adjusted to have
all value less than upper_bound.
"""
pnew = params * 2**(fold * (2*numpy.random.uniform(len(params))-1))
pnew = params * 2**(fold * (2*numpy.random.uniform(size=len(params))-1))
if lower_bound is not None:
for ii,bound in enumerate(lower_bound):
if bound is None:
......
......@@ -385,6 +385,40 @@ def phi_3D_to_4D(phi, f1,f2, xx,yy,zz,aa):
return phi_4D
def phi_4D_to_5D(phi, f1,f2,f3, xx,yy,zz,aa,bb):
"""
Create population 5 from populations 1, 2, 3, and 4.
Returns a 5D sfs of shape (len(xx),len(yy),len(zz),len(aa),len(bb))
phi: phi corresponding to original 3 populations
f1: Fraction of population 4 derived from population 1.
f2: Fraction of population 4 derived from population 2.
f4: Fraction of population 4 derived from population 3.
(A fraction 1-f1-f2-f3 will be derived from population 4.)
xx,yy,zz,aa: Mapping of points in phi to frequencies in populations 1, 2, 3, 4.
bb: Frequency mapping that will be used along population bb axis.
"""
lower_z_index, upper_z_index, frac_lower, frac_upper, norm \
= _four_pop_admixture_intermediates(phi, f1,f2,f3, xx,yy,zz,aa, bb)
# Assemble our result.
# This uses numpy's fancy indexing. It is much, much faster than an
# explicit loop.
# See the numpy-discussion post "Numpy Advanced Indexing Question" by
# Robert Kern on July 16, 2008
# http://projects.scipy.org/pipermail/numpy-discussion/2008-July/035776.html
idx_i = numpy.arange(len(xx))[:,nuax,nuax,nuax]
idx_j = numpy.arange(len(yy))[nuax,:,nuax,nuax]
idx_k = numpy.arange(len(zz))[nuax,nuax,:,nuax]
idx_l = numpy.arange(len(aa))[nuax,nuax,nuax,:]
phi_5D = numpy.zeros((len(xx), len(yy), len(zz), len(aa), len(bb)))
phi_5D[idx_i, idx_j, idx_k, idx_l, lower_z_index] = frac_lower*norm
phi_5D[idx_i, idx_j, idx_k, idx_l, upper_z_index] += frac_upper*norm
return phi_5D
def phi_2D_admix_1_into_2(phi, f, xx,yy):
"""
Admix population 1 into population 2.
......
......@@ -1667,6 +1667,65 @@ class Spectrum(numpy.ma.masked_array):
fs = dadi.Spectrum(data, mask_corners=mask_corners)
return fs
@staticmethod
def _from_phi_5D_linalg(nx, ny, nz, na, nb, xx, yy, zz, aa, bb, phi, mask_corners=True):
"""
Compute sample Spectrum from population frequency distribution phi.
This function uses analytic formulae for integrating over a
piecewise-linear approximation to phi.
See from_phi for explanation of arguments.
"""
data = numpy.zeros((nx+1,ny+1,nz+1,na+1,nb+1))
dbeta1_bb, dbeta2_bb = cached_dbeta(nb, bb)
dbeta1_aa, dbeta2_aa = cached_dbeta(na, aa)
dbeta1_zz, dbeta2_zz = cached_dbeta(nz, zz)
dbeta1_yy, dbeta2_yy = cached_dbeta(ny, yy)
dbeta1_xx, dbeta2_xx = cached_dbeta(nx, xx)
s_bb = (phi[:,:,:,:,1:]-phi[:,:,:,:,:-1])/(bb[nuax,nuax,nuax,nuax,1:]-bb[nuax,nuax,nuax,nuax:-1])
c1_bb = (phi[:,:,:,:,:-1] - s_bb*bb[nuax,nuax,nuax,nuax,:-1])/(nb+1)
for mm in range(0, nb+1):
term1 = np.dot(c1_bb, dbeta1_bb[mm])
term2 = np.dot(s_bb, dbeta2_bb[bb])
term2 *= (mm+1)/((nb+1)*(nb+2))
over_b = term1 + term2
s_aa = (over_b[:,:,:,1:]-over_b[:,:,:,:-1])/(aa[nuax,nuax,nuax,1:]-aa[nuax,nuax,nuax:-1])
c1_aa = (over_b[:,:,:,:-1] - s_aa*aa[nuax,nuax,nuax,:-1])/(na+1)
for ll in range(0, na+1):
term1 = np.dot(c1_aa, dbeta1_aa[ll])
term2 = np.dot(s_aa, dbeta2_aa[ll])
term2 *= (ll+1)/((na+1)*(na+2))
over_a = term1 + term2
s_zz = (over_a[:,:,1:]-over_a[:,:,:-1])/(zz[nuax,nuax,1:]-zz[nuax,nuax,:-1])
c1_zz = (over_a[:,:,:-1] - s_zz*zz[nuax,nuax,:-1])/(nz+1)
for kk in range(0, nz+1):
term1 = np.dot(c1_zz, dbeta1_zz[kk])
term2 = np.dot(s_zz, dbeta2_zz[kk])
term2 *= (kk+1)/((nz+1)*(nz+2))
over_z = term1 + term2
s_yy = (over_z[:,1:]-over_z[:,:-1])/(yy[nuax,1:]-yy[nuax,:-1])
c1_yy = (over_z[:,:-1] - s_yy*yy[nuax,:-1])/(ny+1)
term1_yy = np.dot(dbeta1_yy, c1_yy.T)
term2_yy = np.dot(dbeta2_yy, s_yy.T) * (np.arange(1,ny+2)[:,np.newaxis]/((ny+1)*(ny+2)))
over_y_all = term1_yy + term2_yy
s_xx_all = (over_y_all[:,1:]-over_y_all[:,:-1])/(xx[1:]-xx[:-1])
c1_xx_all = (over_y_all[:,:-1] - s_xx_all*xx[:-1])/(nx+1)
term1_all = np.dot(dbeta1_xx, c1_xx_all.T)
term2_all = np.dot(dbeta2_xx, s_xx_all.T) * (np.arange(1,nx+2)[:,np.newaxis]/((nx+1)*(nx+2)))
data[:,:,kk,ll,mm] = term1_all + term2_all
fs = dadi.Spectrum(data, mask_corners=mask_corners)
return fs
@staticmethod
def _from_phi_4D_linalg(nx, ny, nz, na, xx, yy, zz, aa, phi, mask_corners=True):
"""
......@@ -1890,8 +1949,13 @@ class Spectrum(numpy.ma.masked_array):
xxs[0], xxs[1], xxs[2], xxs[3],
phi, mask_corners,
admix_props)
elif phi.ndim == 5:
if not het_ascertained and not admix_props and not force_direct:
fs = Spectrum._from_phi_5D_linalg(ns[0], ns[1], ns[2], ns[3], ns[4],
xxs[0], xxs[1], xxs[2], xxs[3], xxs[4],
phi, mask_corners)
else:
raise ValueError('Only implemented for dimensions 1,2 or 3.')
raise ValueError('Only implemented for dimensions 1-5.')
fs.pop_ids = pop_ids
# Record value to use for extrapolation. This is the first grid point,
# which is where new mutations are introduced. Note that extrapolation
......
......@@ -267,7 +267,6 @@ void implicit_4Da(double *phi, double *xx, double *yy, double *zz, double *aa,
double *b = malloc(O * sizeof(*b));
double *c = malloc(O * sizeof(*c));
double *r = malloc(O * sizeof(*r));
double *temp = malloc(O * sizeof(*temp));
double x, y, z;
......@@ -280,7 +279,7 @@ void implicit_4Da(double *phi, double *xx, double *yy, double *zz, double *aa,
for(ll=0; ll < O-1; ll++)
VInt[ll] = Vfunc(aInt[ll], nu4);
tridiag_malloc(N);
tridiag_malloc(O);
for(ii = 0; ii < L; ii++){
for(jj = 0; jj < M; jj++){
for(kk = 0; kk < N; kk++){
......@@ -294,7 +293,7 @@ void implicit_4Da(double *phi, double *xx, double *yy, double *zz, double *aa,
MInt[ll] = Mfunc4D(aInt[ll], x, y, z, m41, m42, m43, gamma4, h4);
compute_delj(da, MInt, VInt, O, delj, use_delj_trick);
compute_abc_nobc(da, dfactor, delj, MInt, V, dt, N, a, b, c);
compute_abc_nobc(da, dfactor, delj, MInt, V, dt, O, a, b, c);
for(ll = 0; ll < O; ll++)
r[ll] = phi[ii*M*N*O + jj*N*O + kk*O + ll]/dt;
......@@ -320,5 +319,4 @@ void implicit_4Da(double *phi, double *xx, double *yy, double *zz, double *aa,
free(b);
free(c);
free(r);
free(temp);
}
\ No newline at end of file
#include "integration_shared.h"
#include "tridiag.h"
#include <stdio.h>
#include <stdlib.h>
/* Visual C++ does not support the C99 standard, so we can't use
* variable-length arrays such as arr[L][M][N].
*
* See integration2D.c for detail on how this complicates indexing, and what
* tricks are necessary.
*/
void implicit_5Dx(double *phi, double *xx, double *yy, double *zz, double *aa, double *bb,
double nu1, double m12, double m13, double m14, double m15, double gamma1, double h1,
double dt, int L, int M, int N, int O, int P, int use_delj_trick){
int ii,jj,kk,ll,mm;
double *dx = malloc((L-1) * sizeof(*dx));
double *dfactor = malloc(L * sizeof(*dfactor));
double *xInt = malloc((L-1) * sizeof(*xInt));
double Mfirst, Mlast;
double *MInt = malloc((L-1) * sizeof(*MInt));
double *V = malloc(L * sizeof(*V));
double *VInt = malloc((L-1) * sizeof(*VInt));
double *delj = malloc((L-1) * sizeof(*delj));
double *a = malloc(L * sizeof(*a));
double *b = malloc(L * sizeof(*b));
double *c = malloc(L * sizeof(*c));
double *r = malloc(L * sizeof(*r));
double *temp = malloc(L * sizeof(*temp));
double y, z, a_, b_;
compute_dx(xx, L, dx);
compute_dfactor(dx, L, dfactor);
compute_xInt(xx, L, xInt);
for(ii=0; ii < L; ii++)
V[ii] = Vfunc(xx[ii], nu1);
for(ii=0; ii < L-1; ii++)
VInt[ii] = Vfunc(xInt[ii], nu1);
tridiag_malloc(L);
for(jj = 0; jj < M; jj++){
for(kk = 0; kk < N; kk++){
for(ll = 0; ll < O; ll++){
for(mm = 0; mm < 0; mm++){
y = yy[jj];
z = zz[kk];
a_ = aa[ll];
b_ = bb[mm];
Mfirst = Mfunc5D(xx[0], y, z, a_, b_, m12, m13, m14, m15, gamma1, h1);
Mlast = Mfunc5D(xx[L-1], y, z, a_, b_, m12, m13, m14, m15, gamma1, h1);
for(ii=0; ii < L-1; ii++)
MInt[ii] = Mfunc5D(xInt[ii], y, z, a_, b_, m12, m13, m14, m15, gamma1, h1);
compute_delj(dx, MInt, VInt, L, delj, use_delj_trick);
compute_abc_nobc(dx, dfactor, delj, MInt, V, dt, L, a, b, c);
for(ii = 0; ii < L; ii++)
r[ii] = phi[ii*M*N*O*P + jj*N*O*P + kk*O*P + ll*P + mm]/dt;
if((yy[jj]==0) && (zz[kk]==0) && (aa[ll]==0) && (bb[mm]==0) && (Mfirst <= 0))
b[0] += (0.5/nu1 - Mfirst)*2./dx[0];
if((yy[jj]==1) && (zz[kk]==1) && (aa[ll]==0) && (bb[mm]==0) && (Mlast >= 0))
b[L-1] += -(-0.5/nu1 - Mlast)*2./dx[L-2];
tridiag_premalloc(a, b, c, r, temp, L);
for(ii = 0; ii < L; ii++)
phi[ii*M*N*O*P + jj*N*O*P + kk*O*P + ll*P + mm] = temp[ii];
}
}
}
}
tridiag_free();
free(dx);
free(dfactor);
free(xInt);
free(MInt);
free(V);
free(VInt);
free(delj);
free(a);
free(b);
free(c);
free(r);
free(temp);
}
void implicit_5Dy(double *phi, double *xx, double *yy, double *zz, double *aa, double *bb,
double nu2, double m21, double m23, double m24, double m25, double gamma2, double h2,
double dt, int L, int M, int N, int O, int P, int use_delj_trick){
int ii,jj,kk,ll,mm;
double *dy = malloc((M-1) * sizeof(*dy));
double *dfactor = malloc(M * sizeof(*dfactor));
double *yInt = malloc((M-1) * sizeof(*yInt));
double Mfirst, Mlast;
double *MInt = malloc((M-1) * sizeof(*MInt));
double *V = malloc(M * sizeof(*V));
double *VInt = malloc((M-1) * sizeof(*VInt));
double *delj = malloc((M-1) * sizeof(*delj));
double *a = malloc(M * sizeof(*a));
double *b = malloc(M * sizeof(*b));
double *c = malloc(M * sizeof(*c));