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
