def test_errors(self, recipe):
ts = recipe["ts"]
x_bins = [0, 1.0]
y_bins = [0, 1.0]
time_bins = [0, 10.0]
for stage in ['abcd', 10, []]:
with pytest.raises(ValueError):
pyslim.population_size(ts,
x_bins,
y_bins,
time_bins,
stage=stage)
with pytest.raises(ValueError):
pyslim.population_size(ts,
x_bins,
y_bins,
time_bins,
remembered_stage=stage)
def test_mismatched_remembered_stage(self, recipe):
ts = recipe["ts"]
x_bins = [0, 1.0]
y_bins = [0, 1.0]
time_bins = [0, 10.0]
info = recipe["info"]
if "remembered_early" in recipe:
with pytest.warns(UserWarning):
pyslim.population_size(ts,
x_bins,
y_bins,
time_bins,
remembered_stage="late")
else:
with pytest.warns(UserWarning):
pyslim.population_size(ts,
x_bins,
y_bins,
time_bins,
remembered_stage="early")
def verify(self, ts, remembered_stage):
for bins in self.make_bins(ts):
for stage in ('early', 'late'):
x_bins, y_bins, time_bins = bins
# as computed by pyslim
popsize0 = pyslim.population_size(
ts,
x_bins,
y_bins,
time_bins,
stage=stage,
remembered_stage=remembered_stage)
# as computed in a simple way
popsize1 = self.population_size_simple(
ts,
x_bins,
y_bins,
time_bins,
stage=stage,
remembered_stage=remembered_stage)
assert (np.allclose(popsize1, popsize0))
params = ts.metadata['SLiM']['user_metadata']
for n in params:
if len(params[n]) == 1:
params[n] = params[n][0]
dx = 0.8
T = params['RUNTIME']
dt = params["DT"]
max_gen = ts.metadata["SLiM"]["generation"] - 1
y_bins = [0, params['HEIGHT']]
x_bins = np.arange(0, params['WIDTH'] + dx, dx)
x_mids = x_bins[1:] - np.diff(x_bins) / 2
t_bins = np.arange(0, max_gen + 1)
# Population sizes for each time step within each x bin
popsize = pyslim.population_size(ts, x_bins, y_bins, t_bins)
# "Observed" values of u
uhat = np.sum(popsize, axis=1) / (params["K"] * dx)
assert uhat.shape[0] == len(x_bins) - 1
assert uhat.shape[1] == len(t_bins) - 1
def plot_times(ax, times):
cm = matplotlib.cm.cool
steps = np.searchsorted(t_bins, max_gen - times / dt) - 1
legend_steps = list(range(0, len(steps) + 1, int((len(steps) + 1) / 4)))
for m, (j, t) in enumerate(zip(steps, times)):
ax.plot(x_mids,
uhat[:, j],
c=cm(m / len(steps)),
outbase = ".".join(treefile.split(".")[:-1])
ts = pyslim.load(treefile)
num_gens = ts.metadata['SLiM']['generation']
params = ts.metadata['SLiM']['user_metadata']
try:
dt = params['DT'][0]
except KeyError:
dt = 1.0
width = params['WIDTH'][0]
height = params['HEIGHT'][0]
popsize = pyslim.population_size(
ts,
x_bins=np.linspace(0, width, 101),
y_bins=(0, height),
time_bins=np.arange(num_gens),
)[:, 0, :]
today = ts.individuals_alive_at(0)
has_parents = ts.has_individual_parents()
max_time_ago = np.max(ts.individual_times[has_parents])
if len(today) < num_indivs:
raise ValueError(f"Not enough individuals: only {len(today)} alive today!")
size = (figwidth, figheight)
fig, ax = plt.subplots(figsize=size)
ax.set_xlabel("(forwards) time")
ax.set_ylabel("geographic position")
ax.imshow(
def test_known_answer(self):
# a simple example to make sure we've got the edge cases right
tables = tskit.TableCollection(sequence_length=1)
pyslim.set_tree_sequence_metadata(tables,
model_type='nonWF',
generation=0)
pyslim.set_metadata_schemas(tables)
locs = [
[0, 0], # alive at 0, 1
[0, 1], # alive at 0, 1, 2
[2, 0], # alive at 0, 1, 2
[1, 1], # alive at 0, 1, 2
[0, 0], # alive at 1
[0.5, 1], # alive at 1
[2, 2], # alive at 1
[3, 2] # alive at 0, 1, 2, 3
]
births = [1, 2, 2, 2, 1, 1, 1, 3]
ages = [1, 2, 2, 2, 0, 0, 0, 3]
x_bins = [0, 1, 3]
for xy, a in zip(locs, ages):
md = pyslim.default_slim_metadata('individual')
md['age'] = a
tables.individuals.add_row(
location=xy + [np.nan],
metadata=md,
)
for j, b in enumerate(births):
tables.nodes.add_row(time=b, individual=j)
ts = pyslim.SlimTreeSequence(tables.tree_sequence())
# check we've got this right
for k, n in enumerate([[0, 1, 2, 3, 7], [0, 1, 2, 3, 4, 5, 6, 7],
[1, 2, 3, 7], [7]]):
np.testing.assert_array_equal(n, ts.individuals_alive_at(k))
# no-one
counts = pyslim.population_size(
ts,
x_bins=np.arange(10),
y_bins=np.arange(10),
time_bins=[100, 200, 300],
)
np.testing.assert_array_equal(
counts,
np.zeros((9, 9, 2)),
)
# everyone at the start
counts = pyslim.population_size(
ts,
x_bins=[0, 10],
y_bins=[0, 10],
time_bins=[0, 1],
)
np.testing.assert_array_equal(counts, [[[5]]])
# should omit the last one
counts = pyslim.population_size(ts,
x_bins=[0, 3],
y_bins=[0, 3],
time_bins=[0, 1])
np.testing.assert_array_equal(counts, [[[4]]])
# now should omit the ones at the boundaries
counts = pyslim.population_size(ts,
x_bins=[0, 1, 2],
y_bins=[0, 1, 2],
time_bins=[0, 1, 2, 5])
np.testing.assert_array_equal(
counts, [[[1, 2, 0], [1, 2, 1 / 3]], [[0, 0, 0], [1, 1, 1 / 3]]])