import popgen.Util as utl
sns.set_style("whitegrid", {"grid.color": "1", 'axes.linewidth': .5, "grid.linewidth": ".09"})
s = 0.05;
sel = pd.read_pickle(utl.outpath + 'msms/' + ('neutral', 'selection')[s != 0] + '.df')
sel.loc[sel.nu == 0, 'nu'] = 1
sel = sel.groupby(level=[1, 2]).mean().reset_index()
s=0;neut=pd.read_pickle(utl.outpath + 'msms/'+('neutral','selection')[s!=0]+'.df').groupby(level=[1,2]).mean().reset_index()
import popgen.Plots as pplt
def D(nu, n, W0, Pi0): return -np.log(1 - nu) * W0 / n - Pi0 * nu ** 2
fontsize = 4
nu = utl.forward(t=1000, s=0.05, x0=0.05)
plt.figure(figsize=(6, 3), dpi=300)
plt.subplot(3, 1, 1);
sns.tsplot(sel, time='gen', value='nu', unit='exp', color='red', ci=99);
sns.tsplot(neut, time='gen', value='nu', unit='exp', ci=99);
nu.plot(color='k', linewidth=1, linestyle='--')
pplt.setSize(plt.gca(), fontsize)
plt.ylabel(r'$\nu_t$', fontsize=fontsize + 2);
plt.ylim([0, 1.05])
plt.title(r'(A)', fontsize=fontsize + 2);
plt.tick_params(
axis='x', # changes apply to the x-axis
which='both', # both major and minor ticks are affected
