<jupyter_start><jupyter_text>Diffusion Theory<jupyter_code># Our numerical workhorses
import numpy as np
import pandas as pd

# Import matplotlib stuff for plotting
import matplotlib.pyplot as plt
import matplotlib.cm as cm

# Seaborn, useful for graphics
import seaborn as sns

# favorite Seaborn settings for notebooks
rc={'lines.linewidth': 2, 
    'axes.facecolor' : 'F4F3F6',
    'axes.edgecolor' : '000000',
    'axes.linewidth' : 1.2,
    'grid.color' : 'a6a6a6',
    'lines.linewidth': 2,
    'axes.labelsize': 18,
    'axes.titlesize': 20,
    'xtick.major' : 13,
    'xtick.labelsize': 'large',
    'ytick.labelsize': 13,
    'font.family': 'Lucida Sans Unicode',
    'grid.linestyle': ':',
    'grid.linewidth': 1.5,
    'mathtext.fontset': 'stixsans',
    'mathtext.sf': 'sans',
    'legend.frameon': True,
    'legend.fontsize': 13}

plt.rc('text.latex', preamble=r'\usepackage{sfmath}')
plt.rc('mathtext', fontset='stixsans', sf='sans')
sns.set_style('darkgrid', rc=rc)
sns.set_palette("colorblind", color_codes=True)
sns.set_context('notebook', rc=rc)

# Magic function to make matplotlib inline; other style specs must come AFTER
%matplotlib inline

# This enables SVG graphics inline (only use with static plots (non-Bokeh))
%config InlineBackend.figure_format = 'svg'<jupyter_output><empty_output><jupyter_text>Kolmogorov Forward Equation.  Mutation and Selection.The steady state distribution for mutation and selection is of the form$$P(f) \propto \left[ f(1-f) \right]^{2N\mu - 1},\tag{1}$$where $f$ is the allele frequency $\mu$ is the mutation rate, and $N$ is the population size (effective population size can be).Let's plot this for different values of $\mu$<jupyter_code>def mut_drift(f, N_mu):
    prob = (f * (1 - f))**(2 * N_mu - 1)
    idx = np.logical_and(prob >=0, prob!=np.inf) # to avoid indeterminations
    return f[idx], prob[idx]
# Define parameters
N_mu = [0.1, 1 / 2, 10, 100]
freq = np.linspace(0, 1, 200)

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for nm in N_mu:
    f, prob = mut_drift(freq, nm)
    plt.plot(f, prob / np.sum(prob),
             label=str(nm))
ax.legend(title='$N \mu$')
ax.set_xlabel('allele frequency $f$')
ax.set_ylabel('$\propto$ probability')
ax.set_title('mutation-drift')
ax.axes.get_yaxis().set_ticks([])
ax.margins(0.02)
plt.tight_layout()
plt.savefig('fig/mutation_drift.png')<jupyter_output>/Users/razo/anaconda3/lib/python3.6/site-packages/ipykernel/__main__.py:2: RuntimeWarning: divide by zero encountered in power
  from ipykernel import kernelapp as app<jupyter_text>Selection and Drift Haploid organismsThe steady state distribution of allele frequency for an allele $A$ with fitness $\omega_A = 1$ that is in a population with another allele $a$ with fitness $\omega_a = 1 - s$ is of the form$$P(f) \propto \frac{e^{-2Ns(1-f)}}{f(1-f)}\tag{2}$$Let's plot it for different values of $s$.<jupyter_code>def sel_drift_hap(f, N_s):
    prob = np.exp(-2 * N_s * (1 - f)) / (f * (1 - f))
    idx = np.logical_and(prob >=0, prob!=np.inf) # to avoid indeterminations
    return f[idx], prob[idx]
# Define parameters
N_s = [0, 0.1, 1, 10]
freq = np.linspace(0, 1, 200)

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for ns in N_s:
    f, prob = sel_drift_hap(freq, ns)
    plt.plot(f, prob, #/ np.sum(prob),
             label=str(ns))
ax.legend(loc='upper center', title='$N s$')
ax.set_xlabel('allele frequency $f$')
ax.set_ylabel('$\propto$ probability')
ax.set_title('selection-drift (haploids)')
ax.axes.get_yaxis().set_ticks([])
ax.margins(0.02)
plt.tight_layout()
plt.savefig('fig/sel_drift_haploid.png')<jupyter_output>/Users/razo/anaconda3/lib/python3.6/site-packages/ipykernel/__main__.py:2: RuntimeWarning: divide by zero encountered in true_divide
  from ipykernel import kernelapp as app<jupyter_text>Diploid organisms Heterozygote advantage.For the case of a diploid organism where the genotypes have fitness values\begin{align}\omega_{AA} = 1,\\\omega_{Aa} = 1 + s,\\\omega_{aa} = 1,\end{align}we have that the steady state distribution of alleles is of the form$$P(f) \propto \frac{e^{4Ns f(1 - f)}}{f(1 - f)}\tag{3}$$Let's look at some distributions.<jupyter_code>def sel_drift_dip_hetero_advantage(f, N_s):
    prob = np.exp(4 * N_s * f * (1 - f)) / (f * (1 - f))
    idx = np.logical_and(prob >=0, prob!=np.inf) # to avoid indeterminations
    return f[idx], prob[idx]
# Define parameters
N_s = [1, 10, 50]
freq = np.linspace(0, 1, 200)

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for ns in N_s:
    f, prob = sel_drift_dip_hetero_advantage(freq, ns)
    plt.plot(f, prob / np.sum(prob),
             label=str(ns))
ax.legend(loc=0, title='$N s$')
ax.set_xlabel('allele frequency $f$')
ax.set_ylabel('$\propto$ probability')
ax.set_title('selection-drift (diploids heterozygote advantage)', fontsize=14)
ax.axes.get_yaxis().set_ticks([])
ax.margins(0.02)
plt.tight_layout()
plt.savefig('fig/sel_drift_diploid_hetero_advantage.png')<jupyter_output>/Users/razo/anaconda3/lib/python3.6/site-packages/ipykernel/__main__.py:2: RuntimeWarning: divide by zero encountered in true_divide
  from ipykernel import kernelapp as app<jupyter_text>Heterozygote Intermediate.For the case of a diploid organism where the genotypes have fitness values\begin{align}\omega_{AA} = 1 + 2s,\\\omega_{Aa} = 1 + s,\\\omega_{aa} = 1,\end{align}we have that the steady state distribution of alleles is of the form$$P(f) \propto \frac{e^{4Ns f}}{f(1 - f)}\tag{3}$$Let's look at some distributions.<jupyter_code>def sel_drift_dip_hetero_intermediate(f, N_s):
    prob = np.exp(4 * N_s * f ) / (f * (1 - f))
    idx = np.logical_and(prob >=0, prob!=np.inf) # to avoid indeterminations
    return f[idx], prob[idx]
# Define parameters
N_s = [0.1, 1, 5]
freq = np.linspace(0, 0.95, 200)

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for ns in N_s:
    f, prob = sel_drift_dip_hetero_intermediate(freq, ns)
    plt.plot(f, prob / np.sum(prob),
             label=str(ns))
ax.legend(loc=0, title='$N s$')
ax.set_xlabel('allele frequency $f$')
ax.set_ylabel('$\propto$ probability')
ax.set_title('selection-drift (diploids heterozygote intermediate)', fontsize=14)
ax.axes.get_yaxis().set_ticks([])
ax.set_xlim([0, 1])
ax.margins(0.02)
plt.tight_layout()
plt.savefig('fig/sel_drift_diploid_hetero_intermediate.png')<jupyter_output>/Users/razo/anaconda3/lib/python3.6/site-packages/ipykernel/__main__.py:2: RuntimeWarning: divide by zero encountered in true_divide
  from ipykernel import kernelapp as app<jupyter_text>Mutation-Selection-Drift.For the case of a diploid organism where the genotypes have fitness values\begin{align}\omega_{AA} = 1 - s,\\\omega_{Aa} = 1,\\\omega_{aa} = 1,\end{align}And the mutation rates are given by\begin{align}\mu_{A \rightarrow a} = 0\\\mu_{a \rightarrow A} = \mu\\\end{align}we have that the steady state distribution of alleles is of the form$$P(f) \propto \frac{e^{-2Ns f^2} (1 - f)^{2N \mu}}{f(1 - f)}\tag{3}$$Let's look at some distributions.<jupyter_code>def mut_sel_drift(f, N, s, mu):
    prob = np.exp(- 2 * N * s * f**2) * (1 - f)**(2 * N * mu) / (f * (1 - f))
    idx = np.logical_and(prob >=0, prob!=np.inf) # to avoid indeterminations
    return f[idx], prob[idx]
# Define parameters
N = [10**2, 10**4]
s = 0.0001
mu = [10**-1, 10**-6]
freq = np.linspace(0, 1, 200)

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for n in N:
    for m in mu:
        f, prob = mut_sel_drift(freq, n, s, m)
        plt.plot(f, prob / np.sum(prob),
                label='$N = {0:d}, \mu = {1:f}$'.format(n, m))
ax.legend(loc=0)
ax.set_xlabel('allele frequency $f$')
ax.set_ylabel('$\propto$ probability')
ax.set_title('selection-drift (diploids heterozygote intermediate)', fontsize=14)
ax.axes.get_yaxis().set_ticks([])
ax.set_xlim([0, 1])
ax.margins(0.02)
plt.tight_layout()
# plt.savefig('fig/sel_drift_diploid_hetero_intermediate.png')
freq = np.linspace(1E-3, 0.02, 200)

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for n in N:
    for m in mu:
        f, prob = mut_sel_drift(freq, n, s, m)
        plt.plot(f, prob / np.sum(prob),
                label='$N = {0:d}, \mu = {1:f}$'.format(n, m))
ax.legend(loc=0)
ax.set_xlabel('allele frequency $f$')
ax.set_ylabel('$\propto$ probability')
ax.set_title('selection-drift (diploids heterozygote intermediate)', fontsize=14)
ax.axes.get_yaxis().set_ticks([])
# ax.set_xlim([0, 1])
ax.margins(0.02)
plt.tight_layout()
# plt.savefig('fig/sel_drift_diploid_hetero_intermediate.png')<jupyter_output><empty_output><jupyter_text>Kolmogorov Backwards Equation.  Probability of fixation (Haploid)Using the Kolmogorov backwards equation we derived that the probability of an allele being fixed is of the form$$P(1, \infty \mid f_o) = \frac{1 - e^{-2N_e s f_o}}{1 - e^{-2N_e s}},\tag{4}$$where $f_o$ is the initial frequency of the allele, $N_e$ is the effective population size, and $s$ is the selection coefficient.Let's define a function to compute this quantity.<jupyter_code>def fix_prob_haploid(N, s, fo):
    return (1 - np.exp(-2 * N * s * fo)) / (1 - np.exp(-2 * N * s))
N_array = 10**np.array([2, 3, 4])
s_array = np.linspace(-0.005, 0.005, 100)

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for n in N_array:
    prob = fix_prob_haploid(n, s_array, 1 / n)
    plt.plot(s_array, prob,
            label='$N = {0:d}$'.format(n))
ax.legend(loc=0)
ax.set_xlabel('selection coefficient $s$')
ax.set_ylabel('fixation probability')
ax.set_title('haploids', fontsize=14)
ax.margins(0.02)
plt.tight_layout()
plt.savefig('fig/fix_prob_haploid.png')<jupyter_output><empty_output><jupyter_text>Probability of fixation (diploids)For the diploid case we have a very similar equation for the probability of fixation of an allele$$P(1, \infty \mid f_o) = \frac{1 - e^{-4N_e h s f_o}}{1 - e^{-2N_e h s}},\tag{5}$$where $f_o$ is the initial frequency of the allele, $N_e$ is the effective population size, $h$ is the dominance, and $s$ is the selection coefficient.Let's define a function to compute this quantity.<jupyter_code>def fix_prob_diploid(N, h, s, fo):
    return (1 - np.exp(-4 * N * h * s * fo)) / (1 - np.exp(-4 * N * h * s))
h_array = [0.05, 0.5, 1]
s_array = np.linspace(-0.005, 0.005, 100)
N = 1E3

fig, ax = plt.subplots(1, 1)
# Loop through mutation rates and plot the distribution
for h in h_array:
    prob = fix_prob_diploid(N, h, s_array, 1 / N)
    plt.plot(s_array, prob,
            label='$h = {0:.1f}$'.format(h))
ax.legend(loc=0)
ax.set_xlabel('selection coefficient $s$')
ax.set_ylabel('fixation probability')
ax.set_title('diploids', fontsize=14)
ax.margins(0.02)
plt.tight_layout()
plt.savefig('fig/fix_prob_diploids.png')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15585966201322557, 'fraction_numerical': 0.027828802351212344, 'mean_word_length': 4.46911099949774, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '156493', 'license': 'MIT', 'n_tokens_mistral': 4431, 'n_tokens_neox': 4087, 'n_words': 1059, 'path': 'code/classic_diffusion/diffusion_theory.ipynb', 'repo_name': 'mrazomej/stat_gen', 'size': '10888'}