Commit c0466d00 by Ryan Gutenkunst

Initial 5D implementation, and transpose tests

parent 4b99d4a4
 ... ... @@ -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. # ... ...
 ... ... @@ -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 #include /* 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)); double *r = malloc(M * sizeof(*r)); double *temp = malloc(M * sizeof(*temp)); double x, z, a_, b_; compute_dx(yy, M, dy); compute_dfactor(dy, M, dfactor); compute_xInt(yy, M, yInt); for(jj=0; jj < M; jj++) V[jj] = Vfunc(yy[jj], nu2);