from ipywidgets import interactive, IntSlider, FloatSlider, FloatLogSlider, fixed
import numpy as np
from matplotlib import pyplot as plt
import numba
def travpopwave(numdemes = 100, N = 2000, r0 = 0.01, m = 0.01, numsteps = 4000, rs = 42):
np.random.seed(rs)
n1s = np.zeros((numdemes,numsteps))
n2s = np.zeros((numdemes,numsteps))
n1s[int(np.round(numdemes*0.45)):int(np.round(numdemes*0.55)), 0] = N//2
n2s[int(np.round(numdemes*0.45)):int(np.round(numdemes*0.55)), 0] = N//2
for step in range(1, numsteps):
n1x = np.copy(n1s[:, step - 1])
n2x = np.copy(n2s[:, step - 1])
# migration
n1x = (1 - m) * n1x + m / 2 * (np.roll(n1x, 1) + np.roll(n1x, -1))
n2x = (1 - m) * n2x + m / 2 * (np.roll(n2x, 1) + np.roll(n2x, -1))
# population growth
n1xnew = (1 + r0 * (1 - (n1x + n2x) / N)) * n1x
n2xnew = (1 + r0 * (1 - (n1x + n2x) / N)) * n2x
n1x = np.copy(n1xnew)
n2x = np.copy(n2xnew)
for d in range(numdemes):
f1 = n1x[d] / N
f2 = n2x[d] / N
if f1 + f2 > 1:
f1new = f1 / (f1 + f2)
f2new = f2 / (f1 + f2)
f1 = f1new
f2 = f2new
p = [f1, f2, 1 - (f1 + f2)]
aux = np.random.multinomial(N, p)
n1x[d] = aux[0]
n2x[d] = aux[1]
n1s[:, step] = np.copy(n1x)
n2s[:, step] = np.copy(n2x)
fig, ax = plt.subplots(nrows = 1, ncols = 2, figsize = (20, 4))
X, Y = np.meshgrid(np.arange(numdemes), np.arange(numsteps))
im1 = ax[0].pcolor(X, Y, np.transpose(n1s + n2s), cmap = 'gray', vmin = 0, vmax = N, shading = 'auto')
ax[0].set_ylabel('time (steps)')
ax[0].set_xlabel('deme')
# https://joseph-long.com/writing/colorbars/
cb1 = fig.colorbar(im1, ax=ax[0])
cb1.set_label('population size in deme')
im2 = ax[1].pcolor(X, Y, np.transpose((n1s - n2s) / N), cmap = 'PRGn',
vmin = -1.0, vmax = 1.0, shading = 'auto')
ax[1].set_ylabel('time (steps)')
ax[1].set_xlabel('deme')
cb2 = fig.colorbar(im2, ax=ax[1])
cb2.set_label('deviation from allele 1 and 2 equally abundant')
does it work in principle?
travpopwave()
interactive(travpopwave,
numdemes = IntSlider(value = 100, min = 20, max = 200, continuous_update = False),
N = IntSlider(value = 2000, min = 100, max = 10000, continuous_update = False),
r0 = FloatSlider(value = 0.05, min = 0.0, max = 0.1, step = 0.01, continuous_update = False),
m = FloatSlider(value = 0.05, min = 0.0, max = 0.1, step = 0.01, continuous_update = False),
numsteps = IntSlider(value = 1000, min = 100, max = 2000, continuous_update = False),
rs = IntSlider(value = 42, min = 0, max = 100, continuous_update = False))