from ipywidgets import interactive, IntSlider, FloatSlider, FloatLogSlider, Dropdown, fixed
import numpy as np
from matplotlib import pyplot as plt
import numba
# auxiliary function to facilitate periodic boundary conditions
@numba.njit
def pbc(i, ls):
if i == ls:
return 0
elif i == -1:
return ls-1
else:
return i
def steppingstone(numdemes = 20, N = 50, m = 0.5, relw = 1.0, mu12 = 0.0, mu21 = 0.0,
initialcond = 'uniform', numsteps = 1000):
w1 = relw * 1.0
w2 = 1.0
if initialcond == 'uniform':
n1s = N//2 * np.ones((numdemes,numsteps)) # number individuals of allele 1
elif initialcond == 'nonuniform':
n1s = N * np.zeros((numdemes,numsteps)) # number individuals of allele 1
n1s[int(np.round(numdemes*0.4)):int(np.round(numdemes*0.6)), 0] = int(np.round(N))
for step in range(1, numsteps):
n1x = np.copy(n1s[:, step - 1]) # number of individuals of allele 1, auxiliary array
# migration (in random order)
for d in np.random.permutation(numdemes):
ld = pbc(d - 1, numdemes)
rd = pbc(d + 1, numdemes)
p = m / 2 / N**2 * np.array([1.0, 1.0, 1.0, 1.0])
# 4 cases, see article referenced above
p[0] = p[0] * n1x[d] * (N - n1x[ld])
p[1] = p[1] * n1x[ld] * (N - n1x[d])
p[2] = p[2] * n1x[d] * (N - n1x[rd])
p[3] = p[3] * n1x[rd] * (N - n1x[d])
# choose event
raux = np.random.rand()
if raux < np.cumsum(p)[0]:
n1x[ld] += 1
n1x[d] -= 1
elif raux < np.cumsum(p)[1]:
n1x[ld] -= 1
n1x[d] += 1
elif raux < np.cumsum(p)[2]:
n1x[d] -= 1
n1x[rd] += 1
elif raux < np.cumsum(p)[3]:
n1x[d] += 1
n1x[rd] -= 1
# birth and death
for d in range(numdemes):
f = n1x[d] / N
allelebirth = np.random.choice([1, 2], 1, p = [w1*f/(w1*f+w2*(1-f)), w2*(1-f)/(w1*f+w2*(1-f))])[0]
alleledeath = np.random.choice([1, 2], 1, p = [f, 1-f])[0]
if (allelebirth == 1 and np.random.rand() >= mu12) or (allelebirth == 2 and np.random.rand() < mu21):
n1x[d] += 1
if alleledeath == 1:
n1x[d] -= 1
n1s[:, step] = np.copy(n1x)
fig, ax = plt.subplots(nrows = 1, ncols = 2, figsize = (20, 4))
X, Y = np.meshgrid(np.arange(numdemes), np.arange(numsteps))
im = ax[0].pcolor(X, Y, np.transpose(n1s) / N, cmap = 'PRGn', vmin = 0.0, vmax = 1.0, shading = 'auto')
ax[0].set_ylabel('time (steps)')
ax[0].set_xlabel('deme')
# https://joseph-long.com/writing/colorbars/
cb1 = fig.colorbar(im, ax=ax[0])
cb1.set_label('frequency allele 1')
cf = ax[1].contour(X, Y, np.transpose(n1s) / N)
cb2 = fig.colorbar(cf, ax=ax[1])
cb2.set_label('frequency allele 1')
ax[1].set_ylabel('time (steps)')
ax[1].set_xlabel('deme')
does it work in principle?
steppingstone()
interactive(steppingstone,
numdemes = IntSlider(value = 20, min = 10, max = 100, continuous_update = False),
N = IntSlider(value = 50, min = 10, max = 200, continuous_update = False),
m = FloatSlider(value = 0.5, min = 0.0, max = 1.0, continuous_update = False),
relw = fixed(1.0),
mu12 = fixed(0.0), mu21 = fixed(0.0),
initialcond = Dropdown(options=['uniform', 'nonuniform'], value = 'uniform', description='initial cond'),
numsteps = IntSlider(value = 100, min = 10, max = 2000, continuous_update = False))
interactive(steppingstone,
numdemes = IntSlider(value = 20, min = 10, max = 100, continuous_update = False),
N = IntSlider(value = 50, min = 10, max = 200, continuous_update = False),
m = FloatSlider(value = 0.5, min = 0.0, max = 1.0, continuous_update = False),
relw = FloatLogSlider(value = 1.0, base = 10, min = -1, max = 1, continuous_update = False),
mu12 = fixed(0.0), mu21 = fixed(0.0),
initialcond = Dropdown(options=['uniform', 'nonuniform'], value = 'nonuniform', description='initial cond', disabled=False),
numsteps = IntSlider(value = 100, min = 10, max = 2000, continuous_update = False))