import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from scipy.stats import pearsonr
import pandas as pd
import sys
sys.path.append('../../src/')
from coal_cov import *
from seg_sites_covar import CorrSegSites
from plot_utils import *
%matplotlib inline
%load_ext autoreload
%autoreload 2
plt.rcParams['font.sans-serif'] = "Arial"
plt.rcParams['figure.facecolor'] = "w"
plt.rcParams['figure.autolayout'] = True
plt.rcParams['pdf.fonttype'] = 3
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.ticker as mticker
import os
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)
# NOTE : growth always starts ~ 1000 Generations in the past ...
np.random.seed(42)
tas = [0, 100, 500]
rs = [1e-3, 1e-2, 5e-2]
rec_rate = 1e-3
total_sims = []
for t in tqdm(tas):
for r in rs:
cur_sim = TwoLocusSerialGrowth(ta=t, r=r, T=500, reps=50000, rec_rate=rec_rate)
ts_reps = cur_sim._simulate()
cur_sim._two_locus_branch_length(ts_reps)
# Calculating the marginal variance and means
mu_LA= np.mean(cur_sim.pair_branch_length[:,0])
var_LA = np.var(cur_sim.pair_branch_length[:,0])
cov_LALB = np.cov(cur_sim.pair_branch_length[:,0], cur_sim.pair_branch_length[:,1])[0,1]
corr_LALB = pearsonr(cur_sim.pair_branch_length[:,0], cur_sim.pair_branch_length[:,1])[0]
total_sims.append([t,r, mu_LA, var_LA, cov_LALB, corr_LALB])
total_sims = np.vstack(total_sims)
growth_df = pd.DataFrame(total_sims)
growth_df.columns = ['$t_a$', '$r$','$\mathbb{E}[L]$','$Var(L)$','$Cov(L_A,L_B)$','$Corr(L_A, L_B)$']
growth_df.to_csv('../../results/two_loci/growth_model.csv', index=False)
growth_df
np.random.seed(42)
tas = [0, 100, 500]
rs = [1e-1, 1e-2, 1e-3]
rec_rate = 1e-3
total_sims_bot = []
for t in tqdm(tas):
for r in rs:
cur_sim = TwoLocusSerialBottleneck(Ne=1e4, ta=t, Tstart=50, Tend=500, Nbot=r*1e4, reps=50000, rec_rate=rec_rate)
ts_reps = cur_sim._simulate()
cur_sim._two_locus_branch_length(ts_reps)
# Calculating the marginal variance and means
mu_LA= np.mean(cur_sim.pair_branch_length[:,0])
var_LA = np.var(cur_sim.pair_branch_length[:,0])
cov_LALB = np.cov(cur_sim.pair_branch_length[:,0], cur_sim.pair_branch_length[:,1])[0,1]
corrLALB = pearsonr(cur_sim.pair_branch_length[:,0], cur_sim.pair_branch_length[:,1])[0]
total_sims_bot.append([t, r, mu_LA, var_LA, cov_LALB, corrLALB])
total_sims_bot = np.vstack(total_sims_bot)
bottleneck_df = pd.DataFrame(total_sims_bot)
bottleneck_df.columns = ['$t_a$', '$\phi$','$\mathbb{E}[L]$','$Var(L)$','$Cov(L_A,L_B)$','$Corr(L_A,L_B)$']
bottleneck_df.to_csv('../../results/two_loci/bottleneck.csv', index=False)
bottleneck_df