Repository: tskit-dev/pyslim (language: python)

Issue:
add method to compute discretized population density
Make a method to compute the population size in intervals of x, y, and time from a tree sequence. It should take a tree sequence, bins in the x and y directions, and bins in time and return an easy to plot 3-d array of the number of individuals alive in each x and y bin averaged across the times in each time bin.

```
def population_size(ts, x_bins = 1, y_bins = 1, time_bins = 1)
```

If `x_bins`, `y_bins`, or `time_bins` is a single number, the breakpoints will be `np.linspace(0, max(x), xbins + 1)`. If any of the bins is an array it will be the breakpoints.

Returns a 3-d array with dimensions (number of x breakpoints - 1) by (number of y breakpoints - 1) by (number of time breakpoints - 1)

The goal is to make it easy to make plots like this:
```
import matplotlib.pyplot as plt
x_bins = [0, 10, 20]
time_bins = np.linspace(0, 1000, 101)
popsize = pyslim.population_size(ts, x_bins = x_bins, time_bins = time_bins)
plt.pcolormesh(time_bins, x_bins, popsize)
```
![image](https://user-images.githubusercontent.com/32884936/126401926-22cac0e1-dbb2-47cf-8227-dcb948368d05.png)



Discussion / hints:

Fix patch:
diff --git a/pyslim/__init__.py b/pyslim/__init__.py
index 96afe502..7188e7ad 100644
--- a/pyslim/__init__.py
+++ b/pyslim/__init__.py
@@ -5,4 +5,5 @@
 from pyslim.slim_tree_sequence import *  # NOQA
 from pyslim.provenance import *          # NOQA
 from pyslim.methods import *             # NOQA
+from pyslim.spatial import *             # NOQA
 from pyslim._version import pyslim_version as __version__
diff --git a/pyslim/spatial.py b/pyslim/spatial.py
new file mode 100644
index 00000000..0fdd5b7a
--- /dev/null
+++ b/pyslim/spatial.py
@@ -0,0 +1,166 @@
+import msprime
+import tskit
+import warnings
+import numpy as np
+
+from .slim_tree_sequence import *
+
+
+def _in_location_bin(locations, x0, x1, y0, y1):
+    '''
+    Takes the locations of individuals and tests if each is within [x0, x1) and
+    [y0, y1). Returns a list the same length as the number of individuals with
+    True if the corresponding individual is within the bounds and False
+    otherwise.  This is a helper function for `population_size`, and may change
+    or be removed in the future.
+
+    :param array locations: An n x 3 array with rows individuals, the first
+        column the x coordinates, and the second column the y coordinates.
+    :param float x0: The lower x coordinate boundary.
+    :param float x1: The upper x coordinate boundary.
+    :param float y0: The lower y coordinate boundary.
+    :param float y1: The upper y coordinate boundary.
+    '''
+    return np.logical_and(
+            np.logical_and(
+                locations[:,0] < x1,
+                locations[:,0] >= x0),
+            np.logical_and(
+                locations[:,1] < y1,
+                locations[:,1] >= y0
+            )
+        )
+
+
+def _average_time_alive(birth_times, death_times, t0, t1):
+    '''
+    Calculates the average population size in the time interval [`t0`, `t1`)
+    from birth times and death times of individuals. This is a helper function for
+    `population_size`, and may change or be removed in the future.
+
+    'Average population size' is the sum over all individuals of the length of
+    time each individual was alive in [`t0`, `t1`), divided by `t1 - t0`. For
+    integer t0 and t1, this is equivalent to the average over all integer time
+    steps in [t0, t1) of the number of individuals alive at that time step;
+    however, this is not restricted to integer times.
+
+    Time goes bacwards, so t1 is further in the past than t0 and birth times
+    are greater than death times. Birth times and death times should already
+    include an offset for the stage that individuals are remembered at. See
+    `population_size` code.
+
+    For example, suppose t0 = 1, t1 = 4, and there are four individuals with
+    birth times [3, 1, 5, 4] and death times [0, 1, 2, 4]. The first individual
+    is alive for three time steps in [t0, t1): 1, 2, and 3; the second
+    individual for one time step: 1; the third individual for two time steps: 2
+    and 3, and the fourth individual for zero (since the time interval is open
+    on the right and it was born at t1). So the average population size is (3 +
+    1 + 2 + 0)/(4-1) = 2.
+
+    :param array birth_times: Birth times of individuals.
+    :param array death_times: Death times of individuals. Should be the same
+        length as `birth_times`. Death times are less than or equal to birth times.
+    :param float t0: Lower time endpoint.
+    :param float t1: Upper time endpoint. t1 is greater than t0.
+
+    '''
+    # length of the segment of time in [t0, t1) that the individual
+    # was alive for, i.e., of the intersection with [birth, death].
+    a = np.maximum(0, np.minimum(t1, 1 + birth_times) -  np.maximum(t0, death_times))
+    return sum(a)/(t1-t0)
+
+
+def population_size(ts, x_bins, y_bins, time_bins, stage='late', remembered_stage=None):
+    '''
+    Calculates the population size in each of the spatial bins defined by grid lines
+    at ``x_bins`` and ``y_bins``, averaged over each of the time intervals separated by
+    ``time_bins``. To obtain actual (census) sizes, the tree sequence must contain
+    all individuals alive, e.g., from a SLiM simulation with all individuals
+    permanently remembered.
+
+    With ``nx``, ``ny`` and ``nt`` the number of bins in the ``x``, ``y`` and time directions
+    (so, ``nx`` is one less than the length of ``x_bins``),
+    returns a 3-d array with dimensions ``(nx, ny, nt)``.
+    The ``(i,j,k)``th element of the array is the average number of
+    individuals with ``x`` coordinate in the half-open interval ``[x_bins[i], x_bins[i + 1])``
+    and ``y`` coordinate in the half-open interval ``[y_bins[i], y_bins[i + 1])``,
+    averaged across all times in the half-open interval ``[time_bins[k], time_bins[k + 1])``.
+
+    For integer endpoints of the time bins, this average is equivalent to
+    recording the number of indivduals that are alive at each time in the time
+    interval and have location in the relevant location bin, then taking the
+    mean of these recorded population sizes.
+
+    :param TreeSequence ts: The tree sequence to calculate population size from.
+    :param array x_bins: The x-coordinates of the boundaries of the location bins.
+    :param array y_bins: The y-coordinates of the boundaries of the location bins.
+    :param array time_bins: The endpoints of the time bins.
+    :param str stage: The stage in the SLiM life cycle that the endpoints of
+        the time bins refer to (either "early" or "late"; defaults to "late").
+    :param str remembered_stage: The stage in the SLiM life cycle during which
+        individuals were Remembered (defaults to the stage the tree sequence was
+        recorded at, stored in metadata).
+    '''
+
+    # TODO: make a helper function to do this
+    if stage not in ("late", "early"):
+        raise ValueError(f"Unknown stage '{stage}': "
+                          "should be either 'early' or 'late'.")
+
+    if remembered_stage is None:
+        remembered_stage = ts.metadata['SLiM']['stage']
+
+    if remembered_stage not in ("late", "early"):
+        raise ValueError(f"Unknown remembered_stage '{remembered_stage}': "
+                          "should be either 'early' or 'late'.")
+    if remembered_stage != ts.metadata['SLiM']['stage']:
+        warnings.warn(f"Provided remembered_stage '{remembered_stage}' does not"
+                      " match the stage at which the tree sequence was saved"
+                      f" ('{ts.metadata['SLiM']['stage']}'). This is not necessarily"
+                      " an error, but mismatched stages will lead to inconsistencies:"
+                      " make sure you know what you're doing.")
+
+    # see individuals_alive_at for explanation
+    if stage == "early" and ts.metadata['SLiM']['model_type'] == "WF":
+        birth_offset = 1
+    else:
+        birth_offset = 0
+    birth_times = ts.individual_times - birth_offset
+    if (ts.metadata['SLiM']['model_type'] == "WF"
+            or stage == remembered_stage):
+        age_offset = 0
+    else:
+        if (remembered_stage == "early"
+                and stage == "late"):
+            age_offset = -1
+        else:
+            age_offset = 1
+    ages = ts.individual_ages + age_offset
+    death_times = birth_times - ages
+
+    time_breaks = time_bins
+    x_breaks = x_bins
+    y_breaks = y_bins
+
+    nxbins = len(x_breaks) - 1
+    nybins = len(y_breaks) - 1
+    ntbins = len(time_breaks) - 1
+    popsize = np.empty((nxbins, nybins, ntbins))
+
+    locations = ts.individual_locations
+
+    for i in np.arange(nxbins):
+        for j in np.arange(nybins):
+            x0, x1 = x_breaks[i], x_breaks[i + 1]
+            y0, y1 = y_breaks[j], y_breaks[j + 1]
+            in_bin = _in_location_bin(locations, x0, x1, y0, y1)
+            for k in np.arange(ntbins):
+                t0, t1 = time_breaks[k], time_breaks[k + 1]
+                popsize[i, j, k] = _average_time_alive(
+                        birth_times[in_bin],
+                        death_times[in_bin],
+                        t0,
+                        t1,
+                )
+
+    return popsize


Test patch:
diff --git a/tests/test_spatial.py b/tests/test_spatial.py
new file mode 100644
index 00000000..32ec83fc
--- /dev/null
+++ b/tests/test_spatial.py
@@ -0,0 +1,209 @@
+"""
+Test cases for tree sequences.
+"""
+import pickle
+import random
+import os
+
+import numpy as np
+import pytest
+import tskit
+import msprime
+import pyslim
+
+import tests
+
+from .recipe_specs import recipe_eq
+
+class TestPopulationSize(tests.PyslimTestCase):
+
+    def population_size_simple(self, ts, x_bins, y_bins, time_bins, **kwargs):
+        '''
+        Calculates population size in each location bin averaged over each time_bin.   
+        '''
+        
+        # Want to return
+        # [[[px0y0t0, px0y0t1, ...], [px0y1t0, px0y1t1, ...], ...],
+        #  [[px1y0t0, px0y0t1, ...], [px1y1t0, px1y1t1, ...], ...],
+        #  [[px2y0t0, px0y0t1, ...], [px2y1t0, px2y1t1, ...], ...]]
+        # where pxiyjtk is the population size in [x_breaks[i], x_breaks[i + 1]) and [y_breaks[j], x_breaks[j + 1]),
+        # averaged over each time point in [t[k], t[k + 1])
+        # different offsets for different stages, from pyslim.individuals_alive_at code
+               
+        time_breaks = time_bins
+        x_breaks = x_bins
+        y_breaks = y_bins
+
+        # Initialize array to store popoluationsize
+        nxbins = len(x_breaks) - 1
+        nybins = len(y_breaks) - 1
+        ntbins = len(time_breaks) - 1
+        popsize = np.empty((nxbins, nybins, ntbins))
+        #print(np.shape(popsize))
+
+        # Location, times, and ages of individuals
+        locations = ts.individual_locations
+        times = ts.individual_times
+
+        # Iterate through location bins and time bins
+        for i in np.arange(nxbins):
+            for j in np.arange(nybins):
+                # Endpoints of bins
+                x0, x1 = x_breaks[i], x_breaks[i + 1]
+                y0, y1 = y_breaks[j], y_breaks[j + 1]
+                for k in np.arange(ntbins):
+                    #print(i, j, k)
+                    # Endpoints of bins
+                    t0, t1 = time_breaks[k], time_breaks[k + 1]
+                    alive = 0
+                    for t in np.arange(np.ceil(t0), t1):
+                        for ind_id in ts.individuals_alive_at(t, **kwargs):
+                            ind = ts.individual(ind_id)
+                            if (ind.location[0] < x1 and ind.location[0] >= x0 and ind.location[1] < y1 and ind.location[1] >= y0):
+                                alive += 1/(t1-t0)
+                    popsize[i, j, k] = alive     
+        return(popsize)
+
+    def make_bins(self, ts):
+        for nx in np.arange(1, 20, 10):
+            for ny in np.arange(1, 20, 10):
+                for nt in np.arange(1, 60, 30):
+                    yield [np.linspace(0, round(max(ts.individual_locations[:,0])), nx + 1), 
+                           np.linspace(0, round(max(ts.individual_locations[:,1])), ny + 1),
+                           np.round(np.linspace(0, max(nt, max(ts.individual_times)), nt + 1))]  
+
+    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))
+
+    @pytest.mark.parametrize('recipe', [next(recipe_eq("everyone"))], indirect=True)
+    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)
+
+    @pytest.mark.parametrize('recipe', [next(recipe_eq("pedigree", "WF"))], indirect=True)
+    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")
+
+    @pytest.mark.parametrize('recipe', recipe_eq("everyone"), indirect=True)
+    def test_population_size(self, recipe):
+        # compare output to the right answer
+        ts = recipe["ts"]
+        remembered_stage = 'early' if 'remembered_early' in recipe else 'late'
+        self.verify(ts, remembered_stage=remembered_stage)
+
+    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)
+        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]]]
+        )