def three_epoch(params,
pts,
ns,
rho=0.0,
dt=0.005,
gammaA=0.0,
gammaB=0.0,
hA=0.5,
hB=0.5,
thetaA=1.0,
thetaB=1.0):
nu1, nu2, T1, T2 = params
x = numerics.grid(pts)
yA, yB, phi = equilibrium_phi(pts, rho, dt, gammaA, gammaB, hA, hB, thetaA,
thetaB)
yA, yB, phi = integration.advance(phi,
x,
T1,
yA,
yB,
nu=nu1,
gammaA=gammaA,
gammaB=gammaB,
hA=hA,
hB=hB,
rho=rho,
thetaA=thetaA,
thetaB=thetaB,
dt=dt)
yA, yB, phi = integration.advance(phi,
x,
T2,
yA,
yB,
nu=nu2,
gammaA=gammaA,
gammaB=gammaB,
hA=hA,
hB=hB,
rho=rho,
thetaA=thetaA,
thetaB=thetaB,
dt=dt)
fs = numerics.sample_cached(phi, ns, x)
return fs
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 equilibrium_phi(pts,
rho,
dt,
gammaA=0,
gammaB=0,
hA=0,
hB=0,
thetaA=1.0,
thetaB=1.0):
rho = float(rho)
gammaA = float(gammaA)
gammaB = float(gammaB)
hA = float(hA)
hB = float(hB)
thetaA = float(thetaA)
thetaB = float(thetaB)
fname_template = 'phis_pts{0}_dt{1}_rho{2}_gammaA{3}_gammaB{4}_hA{5}_hB{6}_thetaA{7}_thetaB{8}.npz'
fname = fname_template.format(pts, rho, dt, gammaA, gammaB, hA, hB, thetaA,
thetaB)
fname = os.path.join(cache_path, fname)
try:
fid = np.load(fname)
yA, yB, phi = fid['arr_0'], fid['arr_1'], fid['arr_2']
fid.close()
except IOError:
print('Calculating equilibrium phi from scratch.')
x = numerics.grid(pts)
yA = dadi.PhiManip.phi_1D(x, gamma=gammaA, h=hA, theta0=thetaA)
yB = dadi.PhiManip.phi_1D(x, gamma=gammaB, h=hB, theta0=thetaB)
phi = np.zeros((len(x), len(x), len(x)))
yA, yB, phi = integration.advance(phi, x, 20., yA, yB, 1.0, gammaA,
gammaB, hA, hB, rho, thetaA, thetaB,
dt)
np.savez(fname, yA, yB, phi)
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