def equilibrium(params, ns, pts, sig1=0.0, sig2=0.0, theta1=1.0, theta2=1.0, misid=0.0, dt=0.005, folded=False): """ Integrate the density function to equilibrium params = unused sig1, sig2 - population scaled selection coefficients for the two derived alleles theta1, theta2 - population scaled mutation rates misid - ancestral misidentification parameter dt - time step to use for integration folded = True - fold the frequency spectrum (if we assume we don't know the order that derived alleles appeared) """ x = np.linspace(0, 1, pts + 1) sig1, sig2 = np.float(sig1), np.float(sig2) y1 = dadi.PhiManip.phi_1D(x, gamma=sig1) y2 = dadi.PhiManip.phi_1D(x, gamma=sig2) phi = np.zeros((len(x), len(x))) if sig1 == sig2 == 0.0 and theta1 == theta2 == 1: phi = integration.equilibrium_neutral_exact(x) else: phi = integration.equilibrium_neutral_exact(x) phi, y1, y2 = integration.advance(phi, x, 2, y1, y2, nu=1., sig1=sig1, sig2=sig2, theta1=theta1, theta2=theta2, dt=dt) dx = numerics.grid_dx(x) try: ns = int(ns) except TypeError: ns = ns[0] fs = numerics.sample(phi, ns, x) fs.extrap_t = dt if folded == True: fs = fs.fold_major() if misid > 0.0: fs = numerics.misidentification(fs, misid) return fs
def advance_injection_test(phi, x, T, yA, yB, nu=1., gammaA=0.0, gammaB=0.0, rho=0.0, thetaA=1.0, thetaB=1.0, dt=0.001): """ Integrate phi, yA, and yB forward in time, using dadi for the biallelic density functions and numerics methods for phi """ dx = numerics.grid_dx(x) dx3 = numerics.grid_dx3(x, dx) U01 = numerics.domain(x) for ii in range(int(T / dt)): # integrate the biallelic density functions forward in time using dadi yA = dadi.Integration.one_pop(yA, x, dt, nu=nu, gamma=gammaA, theta0=thetaA) yB = dadi.Integration.one_pop(yB, x, dt, nu=nu, gamma=gammaB, theta0=thetaB) # inject new mutations using biallelic density functions phi = numerics.injectA(x, dx, dt, yA, phi, thetaA) phi = numerics.injectB(x, dx, dt, yB, phi, thetaB) return yA, yB, phi
def two_epoch(params, ns, pts, sig1=0.0, sig2=0.0, theta1=1.0, theta2=1.0, misid=0.0, dt=0.005, folded=False): """ Two epoch demography - a single population size change at some point in the past params = [nu,T,sig1,sig2,theta1,theta2,misid,dt] nu - relative poplulation size change to ancestral population size T - time in past that size change occured (scaled by 2N generations) sig1, sig2 - population scaled selection coefficients for the two derived alleles theta1, theta2 - population scaled mutation rates misid - ancestral misidentification parameter dt - time step to use for integration """ nu, T = params x = np.linspace(0, 1, pts + 1) sig1, sig2 = np.float(sig1), np.float(sig2) y1 = dadi.PhiManip.phi_1D(x, gamma=sig1) y2 = dadi.PhiManip.phi_1D(x, gamma=sig2) phi = np.zeros((len(x), len(x))) # integrate to equilibrium first if sig1 == sig2 == 0.0 and theta1 == theta2 == 1: phi = integration.equilibrium_neutral_exact(x) else: phi = integration.equilibrium_neutral_exact(x) phi, y1, y2 = integration.advance(phi, x, 2, y1, y2, nu=1., sig1=sig1, sig2=sig2, theta1=theta1, theta2=theta2, dt=dt) phi, y1, y2 = integration.advance(phi, x, T, y1, y2, nu=nu, sig1=sig1, sig2=sig2, theta1=theta1, theta2=theta2, dt=dt) dx = numerics.grid_dx(x) try: ns = int(ns) except TypeError: ns = ns[0] fs = numerics.sample(phi, ns, x) fs.extrap_t = dt if folded == True: fs = fs.fold_major() if misid > 0.0: fs = numerics.misidentification(fs, misid) return fs
def bottlegrowth(params, ns, pts, sig1=0.0, sig2=0.0, theta1=1.0, theta2=1.0, misid=0.0, dt=0.005, folded=False): """ Three epoch demography - two instantaneous population size changes in the past params = [nu1,nu2,T1,T2] nu1,nu2 - relative poplulation size changes to ancestral population size (nu1 occurs before nu2, historically) T1,T2 - time for which population had relative sizes nu1, nu2 (scaled by 2N generations) sig1, sig2 - population scaled selection coefficients for the two derived alleles theta1, theta2 - population scaled mutation rates misid - ancestral misidentification parameter dt - time step to use for integration """ nuB, nuF, T = params if nuB == nuF: nu = nuB else: nu = lambda t: nuB * np.exp(np.log(nuF / nuB) * t / T) x = np.linspace(0, 1, pts + 1) sig1, sig2 = np.float(sig1), np.float(sig2) y1 = dadi.PhiManip.phi_1D(x, gamma=sig1) y2 = dadi.PhiManip.phi_1D(x, gamma=sig2) phi = np.zeros((len(x), len(x))) # integrate to equilibrium first if sig1 == sig2 == 0.0 and theta1 == theta2 == 1: phi = integration.equilibrium_neutral_exact(x) else: phi = integration.equilibrium_neutral_exact(x) phi, y1, y2 = integration.advance(phi, x, 2, y1, y2, nu=1., sig1=sig1, sig2=sig2, theta1=theta1, theta2=theta2, dt=dt) phi, y1, y2 = integration.advance(phi, x, T, y1, y2, nu, sig1, sig2, theta1, theta2, dt=dt) dx = numerics.grid_dx(x) try: ns = int(ns) except TypeError: ns = ns[0] fs = numerics.sample(phi, ns, x) fs.extrap_t = dt if folded == True: fs = fs.fold_major() if misid > 0.0: fs = numerics.misidentification(fs, misid) return fs
def advance(phi, x, T, y1, y2, nu=1., sig1=0., sig2=0., theta1=1., theta2=1., dt=0.001): """ Integrate phi, y1, and y2 forward in time phi - density function for triallelic sites y1,y2 - density of biallelic background sites, integrated forward alongside phi T - amount of time to integrate, scaled by 2N generations nu - relative size of population to ancestral size sig1,sig2 - selection coefficients for two derived alleles theta1,theta2 - population scaled mutation rates dt - time step for integration lam - proportion of mutations that occur from simulateous mutation model (Hodgkinson/Eyre-Walker 2010) """ dx = numerics.grid_dx(x) U01 = numerics.domain(x) C_base = numerics.transition12(x, dx, U01) V1_base, M1_base = numerics.transition1(x, dx, U01, sig1, sig2) V2_base, M2_base = numerics.transition2(x, dx, U01, sig1, sig2) V1D1_base, M1D1_base = numerics.transition1D(x, dx, sig1) V1D2_base, M1D2_base = numerics.transition1D(x, dx, sig2) sig_line = sig1 - sig2 Vline_base, Mline_base = numerics.transition1D(x, dx, sig_line) if np.isscalar(nu): C = identity(len(x)**2) + dt / nu * C_base P1 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V1_base / nu + M1_base) P2 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V2_base / nu + M2_base) P1D1 = np.eye(len(x)) + dt * (V1D1_base / nu + M1D1_base) P1D2 = np.eye(len(x)) + dt * (V1D2_base / nu + M1D2_base) Pline = np.eye(len(x)) + dt * (Vline_base / nu + Mline_base) P = numerics.remove_diag_density_weights_nonneutral( x, dt, nu, sig1, sig2) for ii in range(int(T / dt)): y1[1] += dt / dx[1] / x[1] / 2 * theta1 y1 = numerics.advance1D(y1, P1D1) y2[1] += dt / dx[1] / x[1] / 2 * theta2 y2 = numerics.advance1D(y2, P1D2) phi = inject_mutations_1(phi, dt, x, dx, y2, theta1) phi = inject_mutations_2(phi, dt, x, dx, y1, theta2) phi = numerics.advance_adi(phi, U01, P1, P2, x, ii) phi = numerics.advance_cov(phi, C, x, dx) #phi *= 1-P # move density to diagonal boundary and integrate it phi = numerics.move_density_to_bdry(x, phi, P) phi = numerics.advance_line(x, phi, Pline) T_elapsed = int(T / dt) * dt if T - T_elapsed > 1e-8: # adjust dt and integrate last time step dt = T - T_elapsed C = identity(len(x)**2) + dt / nu * C_base P1 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V1_base / nu + M1_base) P2 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V2_base / nu + M2_base) P1D1 = np.eye(len(x)) + dt * (V1D1_base / nu + M1D1_base) P1D2 = np.eye(len(x)) + dt * (V1D2_base / nu + M1D2_base) Pline = np.eye(len(x)) + dt * (Vline_base / nu + Mline_base) P = numerics.remove_diag_density_weights_nonneutral( x, dt, nu, sig1, sig2) y1[1] += dt / dx[1] / x[1] / 2 * theta1 y1 = numerics.advance1D(y1, P1D1) y2[1] += dt / dx[1] / x[1] / 2 * theta2 y2 = numerics.advance1D(y2, P1D2) phi = inject_mutations_1(phi, dt, x, dx, y2, theta1) phi = inject_mutations_2(phi, dt, x, dx, y1, theta2) phi = numerics.advance_adi(phi, U01, P1, P2, x, 0) phi = numerics.advance_cov(phi, C, x, dx) #phi *= 1-P # move density to diagonal boundary and integrate it phi = numerics.move_density_to_bdry(x, phi, P) phi = numerics.advance_line(x, phi, Pline) else: Ts = np.concatenate((np.linspace(0, np.floor(T / dt) * dt, np.floor(T / dt) + 1), np.array([T]))) for ii in range(int(T / dt)): nu_current = nu(Ts[ii]) C = identity(len(x)**2) + dt / nu_current * C_base P1 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V1_base / nu_current + M1_base) P2 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V2_base / nu_current + M2_base) P1D1 = np.eye(len(x)) + dt * (V1D1_base / nu_current + M1D1_base) P1D2 = np.eye(len(x)) + dt * (V1D2_base / nu_current + M1D2_base) Pline = np.eye( len(x)) + dt * (Vline_base / nu_current + Mline_base) P = numerics.remove_diag_density_weights_nonneutral( x, dt, nu_current, sig1, sig2) y1[1] += dt / dx[1] / x[1] / 2 * theta1 y1 = numerics.advance1D(y1, P1D1) y2[1] += dt / dx[1] / x[1] / 2 * theta2 y2 = numerics.advance1D(y2, P1D2) phi = inject_mutations_1(phi, dt, x, dx, y2, theta1) phi = inject_mutations_2(phi, dt, x, dx, y1, theta2) phi = numerics.advance_adi(phi, U01, P1, P2, x, ii) phi = numerics.advance_cov(phi, C, x, dx) #phi *= 1-P # move density to diagonal boundary and integrate it phi = numerics.move_density_to_bdry(x, phi, P) phi = numerics.advance_line(x, phi, Pline) T_elapsed = int(T / dt) * dt if T - T_elapsed > 1e-8: # adjust dt and integrate last time step dt = T - T_elapsed nu_current = nu(Ts[-1]) C = identity(len(x)**2) + dt / nu_current * C_base P1 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V1_base / nu_current + M1_base) P2 = np.outer(np.array([0, 1, 0]), np.ones( len(x))) + dt * (V2_base / nu_current + M2_base) P1D1 = np.eye(len(x)) + dt * (V1D1_base / nu_current + M1D1_base) P1D2 = np.eye(len(x)) + dt * (V1D2_base / nu_current + M1D2_base) Pline = np.eye( len(x)) + dt * (Vline_base / nu_current + Mline_base) P = numerics.remove_diag_density_weights_nonneutral( x, dt, nu_current, sig1, sig2) y1[1] += dt / dx[1] / x[1] / 2 * theta1 y1 = numerics.advance1D(y1, P1D1) y2[1] += dt / dx[1] / x[1] / 2 * theta2 y2 = numerics.advance1D(y2, P1D2) phi = inject_mutations_1(phi, dt, x, dx, y2, theta1) phi = inject_mutations_2(phi, dt, x, dx, y1, theta2) phi = numerics.advance_adi(phi, U01, P1, P2, x, 0) phi = numerics.advance_cov(phi, C, x, dx) #phi *= 1-P # move density to diagonal boundary and integrate it phi = numerics.move_density_to_bdry(x, phi, P) phi = numerics.advance_line(x, phi, Pline) return phi, y1, y2
def advance(phi, x, T, yA, yB, nu=1., gammaA=0.0, gammaB=0.0, hA=0.5, hB=0.5, rho=0.0, thetaA=1.0, thetaB=1.0, dt=0.001): """ Integrate phi, yA, and yB forward in time, using dadi for the biallelic density functions and numerics methods for phi """ dx = numerics.grid_dx(x) dx3 = numerics.grid_dx3(x, dx) U01 = numerics.domain(x) P1 = np.outer([0, 1, 0], np.ones(len(x))) + dt * numerics.transition1( x, dx, U01, gammaA, gammaB, rho, nu, hA=hA, hB=hB) P2 = np.outer([0, 1, 0], np.ones(len(x))) + dt * numerics.transition2( x, dx, U01, gammaA, gammaB, rho, nu, hA=hA, hB=hB) P3 = np.outer([0, 1, 0], np.ones(len(x))) + dt * numerics.transition3( x, dx, U01, gammaA, gammaB, rho, nu, hA=hA, hB=hB) C12 = numerics.transition12(x, dx, U01) for kk in range(len(x)): C12[kk] = identity(len(x)**2) + dt / nu * C12[kk] C13 = numerics.transition13(x, dx, U01) for kk in range(len(x)): C13[kk] = identity(len(x)**2) + dt / nu * C13[kk] C23 = numerics.transition23(x, dx, U01) for kk in range(len(x)): C23[kk] = identity(len(x)**2) + dt / nu * C23[kk] Psurf = numerics.move_density_to_surface(x, dx, dt, gammaA, gammaB, nu, hA=hA, hB=hB) # surface transition matrices #### 9/11 still need to incorporate dominance into surface integration U01surf = numerics.domain_surf(x) P1surf = np.outer([0, 1, 0], np.ones( len(x))) + dt * numerics.transition1_surf( x, dx, U01surf, gammaA, gammaB, rho, nu, hA=hA, hB=hB) P2surf = np.outer([0, 1, 0], np.ones( len(x))) + dt * numerics.transition2_surf( x, dx, U01surf, gammaA, gammaB, rho, nu, hA=hA, hB=hB) Csurf = identity(len(x)** 2) + dt / nu * numerics.transition12_surf(x, dx, U01surf) Pline = np.eye(len(x)) P = numerics.move_surf_to_line(x, dx, dt, gammaA, gammaB, nu) #Pline = numerics.transition1D(x, dx, dt, gamma, nu) if np.all( phi == 0 ) and T >= 5: # solving to equilibrium - integrate at first without covariance term so that the surface is smooth first yA, yB, phi = advance_without_cov(phi, x, dx, dt, yA, yB, thetaA, thetaB, U01, P1, P2, P3, Psurf, rho, P1surf, P2surf, Pline, P, U01surf, 1.) for ii in range(int(T / dt)): # integrate the biallelic density functions forward in time using dadi yA = dadi.Integration.one_pop(yA, x, dt, nu=nu, gamma=gammaA, theta0=thetaA) yB = dadi.Integration.one_pop(yB, x, dt, nu=nu, gamma=gammaB, theta0=thetaB) # inject new mutations using biallelic density functions phi = numerics.injectA( x, dx, dt, yB, phi, thetaA) # A is injected onto B background, so need yB phi = numerics.injectB( x, dx, dt, yA, phi, thetaB) # B is injected onto A background, so need yA # advance bulk of phi forward by dt using numerics methods advance_adi and advance_cov phi = numerics.advance_adi(phi, U01, P1, P2, P3, x, ii) phi = numerics.advance_cov(phi, C12, C13, C23, x, ii) # methods for interaction with and integration of non-axis surface phi = numerics.surface_interaction(phi, x, Psurf) phi = numerics.advance_surface(phi, x, P1surf, P2surf, Csurf, Pline, P, U01surf) phi = numerics.surface_recombination( phi, x, rho / 2., dt ) ## changed to rho/2 5/29 - note that this is effectively only "half" of the recombination events. the other half are along the surface. return yA, yB, phi