def phi_3D_admix_2_and_3_into_1(phi, f2,f3, xx,yy,zz):
"""
Admix populations 2 and 3 into population 1.
Alters phi in place and returns the new version.
phi: phi corresponding to original 3 populations.
f2: Fraction of updated population 1 to be derived from population 2.
f3: Fraction of updated population 1 to be derived from population 3.
A fraction (1-f2-f3) will be derived from the original pop 1.
xx,yy,zz: Mapping of points in phi to frequencies in populations 1,2 and 3.
"""
lower_w_index, upper_w_index, frac_lower, frac_upper, norm \
= _three_pop_admixture_intermediates(phi, 1-f2-f3,f2, xx,yy,zz, xx)
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
# Basically, we're splitting into a fourth ww population, then integrating
# over yy to be left with the two populations we care about.
idx_i = numpy.arange(phi.shape[0])
for jj in xrange(phi.shape[1]):
for kk in xrange(phi.shape[2]):
phi_int = numpy.zeros((phi.shape[0], phi.shape[0]))
phi_int[idx_i, lower_w_index[:,jj,kk]] = lower_cont[:,jj,kk]
phi_int[idx_i, upper_w_index[:,jj,kk]] = upper_cont[:,jj,kk]
phi[:,jj,kk] = Numerics.trapz(phi_int, xx, axis=0)
return phi
def phi_2D_admix_2_into_1(phi, f, xx,yy):
"""
Admix population 2 into population 1.
Alters phi in place and returns the new version.
phi: phi corresponding to original 2 populations
f: Fraction of updated population 1 to be derived from population 2.
(A fraction 1-f will be derived from the original population 1.)
xx,yy: Mapping of points in phi to frequencies in populations 1 and 2.
"""
# Note that it's 1-f here since f now denotes the fraction coming from
# population 2.
lower_z_index, upper_z_index, frac_lower, frac_upper, norm \
= _two_pop_admixture_intermediates(phi, 1-f, xx,yy,xx)
idx_i = numpy.arange(phi.shape[0])
lower_cont = frac_lower*norm
upper_cont = frac_upper*norm
for jj in xrange(len(yy)):
phi_int = numpy.zeros((len(xx), len(xx)))
phi_int[idx_i, lower_z_index[:,jj]] = lower_cont[:,jj]
phi_int[idx_i, upper_z_index[:,jj]] = upper_cont[:,jj]
phi[:,jj] = Numerics.trapz(phi_int, xx, axis=0)
return phi
def phi_3D_admix_1_and_2_into_3(phi, f1,f2, xx,yy,zz):
"""
Admix populations 1 and 2 into population 3.
Alters phi in place and returns the new version.
phi: phi corresponding to original 3 populations.
f1: Fraction of updated population 3 to be derived from population 1.
f2: Fraction of updated population 3 to be derived from population 2.
A fraction (1-f1-f2) will be derived from the original pop 3.
xx,yy,zz: Mapping of points in phi to frequencies in populations 1,2 and 3.
"""
lower_w_index, upper_w_index, frac_lower, frac_upper, norm \
= _three_pop_admixture_intermediates(phi, f1,f2, xx,yy,zz, zz)
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
# Basically, we're splitting into a fourth ww population, then integrating
# over zz to be left with the two populations we care about.
idx_k = numpy.arange(phi.shape[2])
for ii in xrange(phi.shape[0]):
for jj in xrange(phi.shape[1]):
phi_int = numpy.zeros((phi.shape[2], phi.shape[2]))
phi_int[idx_k, lower_w_index[ii,jj]] = lower_cont[ii,jj]
phi_int[idx_k, upper_w_index[ii,jj]] = upper_cont[ii,jj]
phi[ii,jj] = Numerics.trapz(phi_int, zz, axis=0)
return phi
def phi_4D_admix_into_2(phi, f1, f3, f4, xx, yy, zz, aa):
"""
Admix populations 1,3, and 4 into population 2.
Alters phi in place and returns the new version.
phi: phi corresponding to original 4 populations.
f1: Fraction of updated population 2 to be derived from population 1.
f3: Fraction of updated population 2 to be derived from population 3.
f4: Fraction of updated population 2 to be derived from population 4.
A fraction (1-f1-f3-f4) will be derived from the original pop 2.
xx,yy,zz,aa: Mapping of points in phi to frequencies in populations 1,2,3 and 4.
"""
lower_w_index, upper_w_index, frac_lower, frac_upper, norm \
= _four_pop_admixture_intermediates(phi, f1, 1 - f1 - f3 - f4, f3, xx, yy, zz, aa, yy)
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
idx_j = numpy.arange(phi.shape[1])
for ii in range(phi.shape[0]):
for kk in range(phi.shape[2]):
for ll in range(phi.shape[3]):
phi_int = numpy.zeros((phi.shape[1], phi.shape[1]))
phi_int[idx_j, lower_w_index[ii, :, kk,
ll]] = lower_cont[ii, :, kk, ll]
phi_int[idx_j, upper_w_index[ii, :, kk,
ll]] = upper_cont[ii, :, kk, ll]
phi[ii, :, kk, ll] = Numerics.trapz(phi_int, yy, axis=0)
return phi
def phi_2D_admix_1_into_2(phi, f, xx,yy):
"""
Admix population 1 into population 2.
Alters phi in place and returns the new version.
phi: phi corresponding to original 2 populations
f: Fraction of updated population 2 to be derived from population 1.
(A fraction 1-f will be derived from the original population 2.)
xx,yy: Mapping of points in phi to frequencies in populations 1 and 2.
"""
# This is just like the the split_admix situation, but we're splitting into
# a population with zz=yy. We could do this by creating a xx by yy by yy
# array, then integrating out the second population. That's a big waste of
# memory, however.
lower_z_index, upper_z_index, frac_lower, frac_upper, norm \
= _two_pop_admixture_intermediates(phi, f, xx,yy,yy)
# Basically, we're splitting into a third zz population, then integrating
# over yy to be left with the two populations we care about.
lower_cont = frac_lower*norm
upper_cont = frac_upper*norm
idx_j = numpy.arange(phi.shape[1])
for ii in range(phi.shape[0]):
phi_int = numpy.zeros((phi.shape[1], phi.shape[1]))
# Use fancy indexing to avoid the commented out loop.
#for jj in xrange(len(yy)):
# phi_int[jj, upper_z_index[ii,jj]] += frac_upper[ii,jj]*norm[ii,jj]
# phi_int[jj, lower_z_index[ii,jj]] += frac_lower[ii,jj]*norm[ii,jj]
phi_int[idx_j, lower_z_index[ii]] = lower_cont[ii]
phi_int[idx_j, upper_z_index[ii]] += upper_cont[ii]
phi[ii] = Numerics.trapz(phi_int, yy, axis=0)
return phi
def phi_4D_admix_into_1(phi, f2, f3, f4, xx, yy, zz, aa):
"""
Admix populations 2, 3, and 4 into population 1.
Alters phi in place and returns the new version.
phi: phi corresponding to original 4 populations.
f2: Fraction of updated population 1 to be derived from population 2.
f3: Fraction of updated population 1 to be derived from population 3.
f4: Fraction of updated population 1 to be derived from population 4.
A fraction (1-f2-f3-f4) will be derived from the original pop 1.
xx,yy,zz,aa: Mapping of points in phi to frequencies in populations 1,2,3, and 4.
"""
lower_w_index, upper_w_index, frac_lower, frac_upper, norm \
= _four_pop_admixture_intermediates(phi, 1 - f2 - f3 - f4, f2, f3, xx, yy, zz, aa, xx)
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
# Basically, we're splitting into a fifth bb population, then integrating
# over xx to be left with the two populations we care about.
idx_i = numpy.arange(phi.shape[0])
for jj in range(phi.shape[1]):
for kk in range(phi.shape[2]):
for ll in range(phi.shape[3]):
phi_int = numpy.zeros((phi.shape[0], phi.shape[0]))
phi_int[idx_i, lower_w_index[:, jj, kk,
ll]] = lower_cont[:, jj, kk, ll]
phi_int[idx_i, upper_w_index[:, jj, kk,
ll]] = upper_cont[:, jj, kk, ll]
phi[:, jj, kk, ll] = Numerics.trapz(phi_int, xx, axis=0)
return phi
def phi_3D_admix_2_and_3_into_1(phi, f2, f3, xx, yy, zz):
"""
Admix populations 2 and 3 into population 1.
Alters phi in place and returns the new version.
phi: phi corresponding to original 3 populations.
f2: Fraction of updated population 1 to be derived from population 2.
f3: Fraction of updated population 1 to be derived from population 3.
A fraction (1-f2-f3) will be derived from the original pop 1.
xx,yy,zz: Mapping of points in phi to frequencies in populations 1,2 and 3.
"""
lower_w_index, upper_w_index, frac_lower, frac_upper, norm \
= _three_pop_admixture_intermediates(phi, 1 - f2 - f3, f2, xx, yy, zz, xx)
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
# Basically, we're splitting into a fourth ww population, then integrating
# over yy to be left with the two populations we care about.
idx_i = numpy.arange(phi.shape[0])
for jj in range(phi.shape[1]):
for kk in range(phi.shape[2]):
phi_int = numpy.zeros((phi.shape[0], phi.shape[0]))
phi_int[idx_i, lower_w_index[:, jj, kk]] = lower_cont[:, jj, kk]
phi_int[idx_i, upper_w_index[:, jj, kk]] = upper_cont[:, jj, kk]
phi[:, jj, kk] = Numerics.trapz(phi_int, xx, axis=0)
return phi
def phi_3D_admix_1_and_2_into_3(phi, f1, f2, xx, yy, zz):
"""
Admix populations 1 and 2 into population 3.
Alters phi in place and returns the new version.
phi: phi corresponding to original 3 populations.
f1: Fraction of updated population 3 to be derived from population 1.
f2: Fraction of updated population 3 to be derived from population 2.
A fraction (1-f1-f2) will be derived from the original pop 3.
xx,yy,zz: Mapping of points in phi to frequencies in populations 1,2 and 3.
"""
lower_w_index, upper_w_index, frac_lower, frac_upper, norm \
= _three_pop_admixture_intermediates(phi, f1, f2, xx, yy, zz, zz)
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
# Basically, we're splitting into a fourth ww population, then integrating
# over zz to be left with the two populations we care about.
idx_k = numpy.arange(phi.shape[2])
for ii in range(phi.shape[0]):
for jj in range(phi.shape[1]):
phi_int = numpy.zeros((phi.shape[2], phi.shape[2]))
phi_int[idx_k, lower_w_index[ii, jj]] = lower_cont[ii, jj]
phi_int[idx_k, upper_w_index[ii, jj]] = upper_cont[ii, jj]
phi[ii, jj] = Numerics.trapz(phi_int, zz, axis=0)
return phi
def phi_2D_admix_2_into_1(phi, f, xx, yy):
"""
Admix population 2 into population 1.
Alters phi in place and returns the new version.
phi: phi corresponding to original 2 populations
f: Fraction of updated population 1 to be derived from population 2.
(A fraction 1-f will be derived from the original population 1.)
xx,yy: Mapping of points in phi to frequencies in populations 1 and 2.
"""
# Note that it's 1-f here since f now denotes the fraction coming from
# population 2.
lower_z_index, upper_z_index, frac_lower, frac_upper, norm \
= _two_pop_admixture_intermediates(phi, 1 - f, xx, yy, xx)
idx_i = numpy.arange(phi.shape[0])
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
for jj in range(len(yy)):
phi_int = numpy.zeros((len(xx), len(xx)))
phi_int[idx_i, lower_z_index[:, jj]] = lower_cont[:, jj]
phi_int[idx_i, upper_z_index[:, jj]] = upper_cont[:, jj]
phi[:, jj] = Numerics.trapz(phi_int, xx, axis=0)
return phi
def phi_2D_admix_1_into_2(phi, f, xx, yy):
"""
Admix population 1 into population 2.
Alters phi in place and returns the new version.
phi: phi corresponding to original 2 populations
f: Fraction of updated population 2 to be derived from population 1.
(A fraction 1-f will be derived from the original population 2.)
xx,yy: Mapping of points in phi to frequencies in populations 1 and 2.
"""
# This is just like the the split_admix situation, but we're splitting into
# a population with zz=yy. We could do this by creating a xx by yy by yy
# array, then integrating out the second population. That's a big waste of
# memory, however.
lower_z_index, upper_z_index, frac_lower, frac_upper, norm \
= _two_pop_admixture_intermediates(phi, f, xx, yy, yy)
# Basically, we're splitting into a third zz population, then integrating
# over yy to be left with the two populations we care about.
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
idx_j = numpy.arange(phi.shape[1])
for ii in range(phi.shape[0]):
phi_int = numpy.zeros((phi.shape[1], phi.shape[1]))
# Use fancy indexing to avoid the commented out loop.
# for jj in range(len(yy)):
# phi_int[jj, upper_z_index[ii,jj]] += frac_upper[ii,jj]*norm[ii,jj]
# phi_int[jj, lower_z_index[ii,jj]] += frac_lower[ii,jj]*norm[ii,jj]
phi_int[idx_j, lower_z_index[ii]] = lower_cont[ii]
phi_int[idx_j, upper_z_index[ii]] += upper_cont[ii]
phi[ii] = Numerics.trapz(phi_int, yy, axis=0)
return phi
def integrate_norm(self, params, sel_dist, theta):
"""
"""
#need to include tuple() here to make this function play nice
#with numpy arrays
#compute weights for each fs
sel_args = (self.gammas, ) + tuple(params)
weights = sel_dist(*sel_args)
#compute weight for the effectively neutral portion. not using
#CDF function because I want this to be able to compute weight
#for arbitrary mass functions
weight_neu, err_neu = scipy.integrate.quad(sel_dist,
self.gammas[-1],
0,
args=tuple(params))
#function's adaptable for demographic models from 1-3
#populations but this assumes the selection coefficient is the
#same in both populations
pops = len(self.neu_spec.shape)
if pops == 1:
integrated = self.neu_spec * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis] * self.spectra, self.gammas, axis=0)
elif pops == 2:
integrated = self.neu_spec * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis, numpy.newaxis] * self.spectra,
self.gammas,
axis=0)
elif pops == 3:
integrated = self.neu_spec * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis, numpy.newaxis, numpy.newaxis] *
self.spectra,
self.gammas,
axis=0)
else:
raise IndexError("Must have one to three populations")
integrated_fs = Spectrum(integrated, extrap_x=self.extrap_x)
#normalization
dist_int = Numerics.trapz(weights, self.gammas) + weight_neu
return integrated_fs / dist_int * theta
def integrate(self, params, sel_dist, theta):
"""
integration without re-normalizing the DFE. This assumes the
portion of the DFE that is not integrated is not seen in your
sample.
"""
#need to include tuple() here to make this function play nice
#with numpy arrays
sel_args = (self.gammas, ) + tuple(params)
#compute weights for each fs
weights = sel_dist(*sel_args)
#compute weight for the effectively neutral portion. not using
#CDF function because I want this to be able to compute weight
#for arbitrary mass functions
weight_neu, err_neu = scipy.integrate.quad(sel_dist,
self.gammas[-1],
0,
args=tuple(params))
#function's adaptable for demographic models from 1-3 populations
pops = len(self.neu_spec.shape)
if pops == 1:
integrated = self.neu_spec * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis] * self.spectra, self.gammas, axis=0)
elif pops == 2:
integrated = self.neu_spec * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis, numpy.newaxis] * self.spectra,
self.gammas,
axis=0)
elif pops == 3:
integrated = self.neu_spec * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis, numpy.newaxis, numpy.newaxis] *
self.spectra,
self.gammas,
axis=0)
else:
raise IndexError("Must have one to three populations")
integrated_fs = Spectrum(integrated, extrap_x=self.extrap_x)
#no normalization, allow lethal mutations to fall out
return integrated_fs * theta
def remove_pop(phi, xx, popnum):
"""
Remove a population from phi.
Returns new phi with one fewer population.
phi: phi corresponding to original populations
xx: Mapping of points in phi to frequencies in population to be removed
popnum: Population number to remove, numbering from 1.
"""
return Numerics.trapz(phi, xx, axis=popnum-1)
def remove_pop(phi, xx, popnum):
"""
Remove a population from phi.
Returns new phi with one fewer population.
phi: phi corresponding to original populations
xx: Mapping of points in phi to frequencies in population to be removed
popnum: Population number to remove, numbering from 1.
"""
return Numerics.trapz(phi, xx, axis=popnum - 1)
def phi_5D_admix_into_5(phi, f1, f2, f3, f4, xx, yy, zz, aa, bb):
"""
Admix populations 1, 2, 3, and 4 into population 5.
Alters phi in place and returns the new version.
phi: phi corresponding to original 5 populations.
f1: Fraction of updated population 5 to be derived from population 1.
f2: Fraction of updated population 5 to be derived from population 2.
f3: Fraction of updated population 5 to be derived from population 3.
f4: Fraction of updated population 5 to be derived from population 3.
A fraction (1-f1-f2-f3-f4) will be derived from the original pop 5.
xx,yy,zz,aa,bb: Mapping of points in phi to frequencies in populations 1,2,3,4, and 5.
"""
lower_w_index, upper_w_index, frac_lower, frac_upper, norm \
= _five_pop_admixture_intermediates(phi, f1, f2, f3, f4, xx, yy, zz, aa, bb, xx)
lower_cont = frac_lower * norm
upper_cont = frac_upper * norm
idx_m = numpy.arange(phi.shape[4])
for ii in range(phi.shape[0]):
for jj in range(phi.shape[1]):
for kk in range(phi.shape[2]):
for ll in range(phi.shape[3]):
phi_int = numpy.zeros((phi.shape[4], phi.shape[4]))
phi_int[idx_m,
lower_w_index[ii, jj, kk,
ll, :]] = lower_cont[ii, jj, kk,
ll, :]
phi_int[idx_m,
upper_w_index[ii, jj, kk,
ll, :]] = upper_cont[ii, jj, kk,
ll, :]
phi[ii, jj, kk, ll, :] = Numerics.trapz(phi_int,
xx,
axis=0)
return phi
def demo_selection_distINV(params, ns, sel_dist, theta, cache):
"""
sel_dist should be a function that is evaluated PDF(X) = func(x, param1, param2 ..., paramn)
theta is required for now, just going to use Poisson ll
"""
#load saved objects
#spectra_obj = pickle.load(open('{0}spectra.obj'.format(note),'rb'))
spectra_obj = cache
# Note that first and last entry of SFS are meaningless!
# The last two entries of params are now h_intercept and h_rate
params_DFE = params[:-2]
params_h = params[-2:]
hvalues = comp_h_from_s_INV(spectra_obj['gammas'], *params_h)
# Choose the closest SFS that correspond to the respective hvalues:
hlist_idx = [find_nearest_idx(spectra_obj['hlist'], h) for h in hvalues]
spectra_interp = numpy.array([
spectra_obj['spectra'][hlist_idx[i], i, :]
for i in range(spectra_obj['spectra'].shape[1])
])
# For some reason, some SFS just contain nan when 2Neas*h is large...
# For now, set them to 0 and deal with it later... they should only contain small values anyway
# Update: THis is fixed, the problem where negative values in the SFS
# spectra_interp = [numpy.nan_to_num(sfs) for sfs in spectra_interp] # replaces nan with 0
#compute weights for each SFS
sel_args = (spectra_obj['gammas'], ) + tuple(params_DFE)
weights = sel_dist(*sel_args)
#compute weight for the effectively neutral portion. not using CDF function because
#I want this to be able to compute weight for an arbitrary mass functions
weight_neu, err_neu = scipy.integrate.quad(sel_dist,
spectra_obj['gammas'][-1],
0,
args=tuple(params_DFE))
weight_lethal, err_lethal = scipy.integrate.quad(sel_dist,
-numpy.inf,
spectra_obj['gammas'][0],
args=tuple(params_DFE))
#function's adaptable for demographic models from 1-3 populations
pops = len(spectra_obj['neu_spec'].shape)
if pops == 1:
integrated = spectra_obj['neu_spec'] * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis] * spectra_interp,
spectra_obj['gammas'],
axis=0) + spectra_interp[0] * weight_lethal
elif pops == 2:
integrated = spectra_obj['neu_spec'] * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis, numpy.newaxis] * spectra_interp,
spectra_obj['gammas'],
axis=0) + spectra_interp[0] * weight_lethal
elif pops == 3:
integrated = spectra_obj['neu_spec'] * weight_neu + Numerics.trapz(
weights[:, numpy.newaxis, numpy.newaxis, numpy.newaxis] *
spectra_interp,
spectra_obj['gammas'],
axis=0) + spectra_interp[0] * weight_lethal
else:
raise IndexError("Must have one to three populations")
integrated_fs = Spectrum(integrated, extrap_x=spectra_obj['extrap_x'])
# Changed this:
# Lethal mutations now don't fall out. All lethal mutations contribute the most deleterious SFS.
# This assumes that the range of gamma goes from small to lethal!
return integrated_fs * theta