Two-Locus Simulations w. Serial Samples under Constant Population Size

In [1]:
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

import sys
import os
sys.path.append('../../src/')
from coal_cov import *
from plot_utils import *
from tqdm import tqdm

%load_ext autoreload
%autoreload 2
%matplotlib inline
In [2]:
plt.rcParams['font.sans-serif'] = "Arial"
plt.rcParams['figure.facecolor'] = "w"
plt.rcParams['figure.autolayout'] = True
plt.rcParams['pdf.fonttype'] = 3

# Making the relevant figure directories that we want
main_figdir = '../../plots/two_locus_stats/'
supp_figdir = '../../plots/supp_figs/two_locus_stats/'
os.makedirs(main_figdir, exist_ok=True)
os.makedirs(supp_figdir, exist_ok=True)

Two-Loci, Two Haplotypes

In [3]:
%%time
from mpl_toolkits.axes_grid1 import make_axes_locatable

# Define Probability of different states as a function of rho and time
rhos=10**np.linspace(-2, 2., num=100)
t = np.linspace(0.0, 1., num=100)
nrho = rhos.shape[0]
nt = t.shape[0]

prob1_mat = np.zeros(shape=(nrho, nt))
prob2_mat = np.zeros(shape=(nrho, nt))
for i in range(nrho):
    r = rhos[i]
    # Calculation of 
    cur_prob1 = r*(1.  - np.exp(-t*(r/2 + 1)))/(r + 2)
    cur_prob2 = 1. - (r*(1.  - np.exp(-t*(r/2 + 1)))/(r + 2))
    prob1_mat[i] = cur_prob1
    prob2_mat[i] = cur_prob2


# Plotting Values
f, ax = plt.subplots(1,2,figsize=(6,3), sharey=True)
ax[0].set(yscale="log")
ax02 = ax[0].twinx().twiny()
im1 = ax02.imshow(np.flip(prob1_mat, axis=0), vmin=0, vmax=1.0)

ax[0].set_ylim(np.min(rhos), np.max(rhos))
ax02.set_yticks([])
ax02.set_xticks([])
f.colorbar(im1, ax=ax02)

ax12 = ax[1].twinx().twiny()
ax[1].set(yscale="log")
im2 = ax12.imshow(np.flip(prob2_mat, axis=0), vmin=0, vmax=1.0)

ax[1].set_ylim(np.min(rhos), np.max(rhos))
ax12.set_yticks([])
ax12.set_xticks([])
f.colorbar(im2, ax=ax12)


# Tick Labels
idx = np.linspace(0, nt-1, 10, dtype=np.int32)
ax[0].set_xticks(idx);
ax[1].set_xticks(idx);
ax[0].set_xticklabels(np.round(t[idx],decimals=1), {'fontsize': 10})
ax[1].set_xticklabels(np.round(t[idx],decimals=1),{'fontsize': 10})

ax[0].set_ylabel(r'$\rho$', fontsize=16)
ax[0].set_xlabel(r'$t_a$', fontsize=16)
ax[1].set_xlabel(r'$t_a$', fontsize=16)
ax[0].set_title(r'$\mathbb{P}$(Uncoupled)', fontsize=14)
ax[1].set_title(r'$\mathbb{P}$(Coupled)', fontsize=14)
f.tight_layout()

# plt.savefig(supp_figdir + 'fig_S1.pdf', dpi=300, bbox_inches='tight')
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<timed exec> in <module>

~/.pyenv/versions/3.7.5/lib/python3.7/site-packages/matplotlib/axes/_base.py in wrapper(self, *args, **kwargs)
     73
     74         def wrapper(self, *args, **kwargs):
---> 75             return get_method(self)(*args, **kwargs)
     76
     77         wrapper.__module__ = owner.__module__

TypeError: _set_ticklabels() takes 2 positional arguments but 3 were given
In [4]:
%%time
from mpl_toolkits.axes_grid1 import make_axes_locatable

# Define Probability of different states as a function of rho and time
rhos=np.logspace(0, 2., num=100)
t = np.linspace(0, 1., num=100)
nrho = rhos.shape[0]
nt = t.shape[0]

prob1_mat = np.zeros(shape=(nrho, nt))
for i in range(nrho):
    r = rhos[i]
    # Calculation of 
    cur_prob1 = r*(1.  - np.exp(-t*(r/2 + 1)))/(r + 2)
    prob1_mat[i] = cur_prob1

# Plotting Values
f, ax = plt.subplots(1,1,figsize=(4,4), sharey=True)
im1 = ax.imshow(prob1_mat, vmin=0, vmax=1.0, aspect='auto')
ax.set(yscale="log")
ax.set_ylim(np.min(rhos), np.max(rhos))
cbar = f.colorbar(im1, ax=ax)
# Tick Labels
idx = np.linspace(0, nt-1, 10, dtype=np.int32)
ax.set_xticks(idx);
ax.set_xticks(idx);
ax.set_xticklabels(np.round(t[idx],decimals=1), {'fontsize': 10})
ax.set_xticklabels(np.round(t[idx],decimals=1),{'fontsize': 10})

ax.set_ylabel(r'$\rho$', fontsize=16)
# ax.set_xlabel(r'$t_a$', fontsize=16)
ax.set_xlabel(r'$t_a$', fontsize=16)
ax.set_title(r'$\mathbb{P}$(Uncoupled)', fontsize=16)
f.tight_layout()

# plt.savefig(main_figdir + 'fig_2A.pdf', dpi=300, bbox_inches='tight')
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<timed exec> in <module>

~/.pyenv/versions/3.7.5/lib/python3.7/site-packages/matplotlib/axes/_base.py in wrapper(self, *args, **kwargs)
     73
     74         def wrapper(self, *args, **kwargs):
---> 75             return get_method(self)(*args, **kwargs)
     76
     77         wrapper.__module__ = owner.__module__

TypeError: _set_ticklabels() takes 2 positional arguments but 3 were given

Correlation in Branch Length

In [5]:
def se_cov(x,y, unbiased=False):
    """ compute the standard error of a covariance (using Richardsons formula) """
    assert(x.shape[0] == y.shape[0])
    n = x.shape[0]
    mu_x = np.mean(x)
    mu_y = np.mean(y)
    mean_diff_x = (x - mu_x)
    mean_diff_y = (y - mu_y)
    mu_22 = np.mean((mean_diff_x**2) * (mean_diff_y**2))
    mu_11 = np.mean(mean_diff_x * mean_diff_y)
    mu_20 = np.var(x)
    mu_02 = np.var(y)
    var_cov_xy = (n-1)**2/(n**3)*(mu_22 - mu_11**2) + (n-1)/(n**3)*(mu_02 * mu_20 - mu_11**2)
    if unbiased:
        var_cov_xy  = var_cov_xy * ((n/(n-1))**2)
    return(var_cov_xy)

def se_corr(x,y):
    """Calculate standard error of the correlation."""
    assert(x.shape[0] == y.shape[0])
    r = np.corrcoef(x,y)[0,1]
    N = x.shape[0]
    se_r = np.sqrt((1. - r**2)/(N-2))
    return(se_r)
In [6]:
# Showing the different decay rates!
rhos2_t = 10**np.linspace(-1,4, 1000)
u002 = lambda rho : (rho**2 + 14*rho + 36) / (rho**2 + 13*rho + 18)

f, ax = plt.subplots(1,1, figsize=(4,4))
ta = 0.01
Ne = 1
ax.loglog(rhos2_t, TwoLocusTheoryConstant._corrLALB(rhos2_t, 0.0), linestyle='--',
        lw=2, color='black', label=r'$t_a = 0$')
ax.loglog(rhos2_t, TwoLocusTheoryConstant._corrLALB(rhos2_t, ta),
          lw=2, label=r'$t_a = %.2f$' % (ta * Ne))
ax.loglog(rhos2_t, TwoLocusTheoryConstant._corrLALB(rhos2_t, ta*10),
          lw=2, label=r'$t_a = %.2f$' % (ta*10*Ne))
ax.loglog(rhos2_t, TwoLocusTheoryConstant._corrLALB(rhos2_t, ta*100),
        lw=2, label=r'$t_a = %.2f$' % (ta*100*Ne))


# Showing the asymptotic behavior...
ax.loglog(rhos2_t, 1./rhos2_t, lw=1, linestyle='--', color='red')
ax.loglog(rhos2_t, 8.0/(rhos2_t**2), lw=1, linestyle='--', color='red')

ax.set_xlabel(r'$\rho$', fontsize=14)
ax.set_ylabel(r'$Corr(L_A, L_B)$', fontsize=14)
ax.legend(fontsize=10, labelspacing=-0.25)
debox(ax);