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): """ Integrate a 2-dimensional phi foward. phi: Initial 2-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: Population sizes gamma1,gamma2: Selection coefficients on *all* segregating alleles h1,h2: Dominance coefficients. h = 0.5 corresponds to genic selection. m12,m21: Migration rates. Note that m12 is the rate *into 1 from 2*. theta0: Propotional 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.) frozen1,frozen2: If True, the corresponding population is "frozen" in time (no new mutations and no drift), so the resulting spectrum will correspond to an ancient DNA sample from that population. nomut1,nomut2: If True, no new mutations will be introduced into the given population. 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. """ phi = phi.copy() 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 or frozen2) and (m12 != 0 or m21 != 0): raise ValueError('Population cannot be frozen and have non-zero ' '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]): # 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): return _two_pops_const_params(phi, xx, T, nu1, nu2, m12, m21, gamma1, gamma2, h1, h2, theta0, initial_t, frozen1, frozen2, nomut1, nomut2) yy = xx nu1_f = Misc.ensure_1arg_func(nu1) nu2_f = Misc.ensure_1arg_func(nu2) m12_f = Misc.ensure_1arg_func(m12) m21_f = Misc.ensure_1arg_func(m21) gamma1_f = Misc.ensure_1arg_func(gamma1) gamma2_f = Misc.ensure_1arg_func(gamma2) h1_f = Misc.ensure_1arg_func(h1) h2_f = Misc.ensure_1arg_func(h2) theta0_f = Misc.ensure_1arg_func(theta0) if cuda_enabled: import dadi.cuda phi = dadi.cuda.Integration._two_pops_temporal_params( phi, xx, T, initial_t, nu1_f, nu2_f, m12_f, m21_f, gamma1_f, gamma2_f, h1_f, h2_f, theta0_f, frozen1, frozen2, nomut1, nomut2) return phi current_t = initial_t nu1, nu2 = nu1_f(current_t), nu2_f(current_t) m12, m21 = m12_f(current_t), m21_f(current_t) gamma1, gamma2 = gamma1_f(current_t), gamma2_f(current_t) h1, h2 = h1_f(current_t), h2_f(current_t) dx, dy = numpy.diff(xx), numpy.diff(yy) while current_t < T: dt = min(_compute_dt(dx, nu1, [m12], gamma1, h1), _compute_dt(dy, nu2, [m21], gamma2, h2)) this_dt = min(dt, T - current_t) next_t = current_t + this_dt nu1, nu2 = nu1_f(next_t), nu2_f(next_t) m12, m21 = m12_f(next_t), m21_f(next_t) gamma1, gamma2 = gamma1_f(next_t), gamma2_f(next_t) h1, h2 = h1_f(next_t), h2_f(next_t) theta0 = theta0_f(next_t) if numpy.any(numpy.less([T, nu1, nu2, m12, m21, 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], 0)): raise ValueError('A population size is 0. Has the model been ' 'mis-specified?') _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) current_t = next_t return phi
def four_pops(phi, xx, T, nu1=1, nu2=1, nu3=1, nu4=1, m12=0, m13=0, m14=0, m21=0, m23=0, m24=0, m31=0, m32=0, m34=0, m41=0, m42=0, m43=0, gamma1=0, gamma2=0, gamma3=0, gamma4=0, h1=0.5, h2=0.5, h3=0.5, h4=0.5, theta0=1, initial_t=0, frozen1=False, frozen2=False, frozen3=False, frozen4=False, enable_cuda_const=False): """ Integrate a 4-dimensional phi foward. phi: Initial 4-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: Population sizes gamma1,gamma2,gamma3,gamma4: Selection coefficients on *all* segregating alleles h1,h2,h3,h4: 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 (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)): raise ValueError('Population cannot be frozen and have non-zero ' 'migration to or from it.') 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) 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) 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) m12_f, m13_f, m14_f = Misc.ensure_1arg_func(m12), Misc.ensure_1arg_func( m13), Misc.ensure_1arg_func(m14) m21_f, m23_f, m24_f = Misc.ensure_1arg_func(m21), Misc.ensure_1arg_func( m23), Misc.ensure_1arg_func(m24) m31_f, m32_f, m34_f = Misc.ensure_1arg_func(m31), Misc.ensure_1arg_func( m32), Misc.ensure_1arg_func(m34) m41_f, m42_f, m43_f = Misc.ensure_1arg_func(m41), Misc.ensure_1arg_func( m42), Misc.ensure_1arg_func(m43) 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 = nu1_f(current_t), nu2_f(current_t), nu3_f( current_t), nu4_f(current_t) gamma1, gamma2, gamma3, gamma4 = gamma1_f(current_t), gamma2_f( current_t), gamma3_f(current_t), gamma4_f(current_t) h1, h2, h3, h4 = h1_f(current_t), h2_f(current_t), h3_f(current_t), h4_f( current_t) m12, m13, m14 = m12_f(current_t), m13_f(current_t), m14_f(current_t) m21, m23, m24 = m21_f(current_t), m23_f(current_t), m24_f(current_t) m31, m32, m34 = m31_f(current_t), m32_f(current_t), m34_f(current_t) m41, m42, m43 = m41_f(current_t), m42_f(current_t), m43_f(current_t) dx, dy, dz, da = numpy.diff(xx), numpy.diff(yy), numpy.diff( zz), numpy.diff(aa) while current_t < T: dt = min(_compute_dt(dx, nu1, [m12, m13, m14], gamma1, h1), _compute_dt(dy, nu2, [m21, m23, m24], gamma2, h2), _compute_dt(dz, nu3, [m31, m32, m34], gamma3, h3), _compute_dt(da, nu4, [m41, m42, m43], gamma4, h4)) this_dt = min(dt, T - current_t) next_t = current_t + this_dt nu1, nu2, nu3, nu4 = nu1_f(next_t), nu2_f(next_t), nu3_f( next_t), nu4_f(next_t) gamma1, gamma2, gamma3, gamma4 = gamma1_f(next_t), gamma2_f( next_t), gamma3_f(next_t), gamma4_f(next_t) h1, h2, h3, h4 = h1_f(next_t), h2_f(next_t), h3_f(next_t), h4_f(next_t) m12, m13, m14 = m12_f(next_t), m13_f(next_t), m14_f(next_t) m21, m23, m24 = m21_f(next_t), m23_f(next_t), m24_f(next_t) m31, m32, m34 = m31_f(next_t), m32_f(next_t), m34_f(next_t) m41, m42, m43 = m41_f(next_t), m42_f(next_t), m43_f(next_t) theta0 = theta0_f(next_t) if numpy.any( numpy.less([ T, nu1, nu2, nu3, nu4, m12, m13, m14, m21, m23, m24, m31, m32, m34, m41, m42, m43, 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], 0)): raise ValueError('A population size is 0. Has the model been ' 'mis-specified?') _inject_mutations_4D(phi, this_dt, xx, yy, zz, aa, theta0, frozen1, frozen2, frozen3, frozen4) if not frozen1: phi = int_c.implicit_4Dx(phi, xx, yy, zz, aa, nu1, m12, m13, m14, gamma1, h1, this_dt, use_delj_trick) if not frozen2: phi = int_c.implicit_4Dy(phi, xx, yy, zz, aa, nu2, m21, m23, m24, gamma2, h2, this_dt, use_delj_trick) if not frozen3: phi = int_c.implicit_4Dz(phi, xx, yy, zz, aa, nu3, m31, m32, m34, gamma3, h3, this_dt, use_delj_trick) if not frozen4: phi = int_c.implicit_4Da(phi, xx, yy, zz, aa, nu4, m41, m42, m43, gamma4, h4, this_dt, use_delj_trick) current_t = next_t return phi
def one_pop(phi, xx, T, nu=1, gamma=0, h=0.5, theta0=1.0, initial_t=0, frozen=False, beta=1): """ Integrate a 1-dimensional phi forward. phi: Initial 1-dimensional phi xx: Grid upon (0,1) overwhich phi is defined. nu, gamma, and theta0 may be functions of time. nu: Population size gamma: Selection coefficient on *all* segregating alleles h: Dominance coefficient. h = 0.5 corresponds to genic selection. q Heterozygotes have fitness 1+2sh and homozygotes have fitness 1+2s. theta0: Propotional to ancestral size. Typically constant. beta: Breeding ratio, beta=Nf/Nm. 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.) frozen: If True, population is 'frozen' so that it does not change. In the one_pop case, this is equivalent to not running the integration at all. """ phi = phi.copy() # For a one population integration, freezing means just not integrating. if frozen: return phi 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)) vars_to_check = (nu, gamma, h, theta0, beta) if numpy.all([numpy.isscalar(var) for var in vars_to_check]): return _one_pop_const_params(phi, xx, T, nu, gamma, h, theta0, initial_t, beta) nu_f = Misc.ensure_1arg_func(nu) gamma_f = Misc.ensure_1arg_func(gamma) h_f = Misc.ensure_1arg_func(h) theta0_f = Misc.ensure_1arg_func(theta0) beta_f = Misc.ensure_1arg_func(beta) current_t = initial_t nu, gamma, h = nu_f(current_t), gamma_f(current_t), h_f(current_t) beta = beta_f(current_t) dx = numpy.diff(xx) while current_t < T: dt = _compute_dt(dx, nu, [0], gamma, h) this_dt = min(dt, T - current_t) # Because this is an implicit method, I need the *next* time's params. # So there's a little inconsistency here, in that I'm estimating dt # using the last timepoints nu,gamma,h. next_t = current_t + this_dt nu, gamma, h = nu_f(next_t), gamma_f(next_t), h_f(next_t) beta = beta_f(next_t) theta0 = theta0_f(next_t) if numpy.any(numpy.less([T, nu, 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([nu], 0)): raise ValueError('A population size is 0. Has the model been ' 'mis-specified?') _inject_mutations_1D(phi, this_dt, xx, theta0) # Do each step in C, since it will be faster to compute the a,b,c # matrices there. phi = int_c.implicit_1Dx(phi, xx, nu, gamma, h, beta, this_dt, use_delj_trick=use_delj_trick) current_t = next_t return phi