Merge pull request #773 from mcveanlab/release-0.7.1-final

Final updates for 0.7.1

CHANGELOG.rst

@@ -1,3 +1,26 @@
+********************
+[0.7.1] - 2019-06-08
+********************
+
+**New features**
+
+- Discrete Time Wright-Fisher simulation model (:user:`DomNelson`).
+- SMC/SMC' simulation models (:user:`jeromekelleher`).
+- Mixed simulation models (:user:`jeromekelleher`).
+- Specify ``end_time`` to allow early-finish for simulations (:user:`jeromekelleher`).
+- Calculation of historical coalescence rates in the DemographyDebugger
+  (:user:`jgallowa07`, :user:`petrelharp`).
+- Additional information on population sizes in DemographyDebugger
+  (:user:`andrewkern`).
+- Remove support for Python 2 (:user:`hugovk`).
+- Allow specifying metadata for populations (:user:`jeromekelleher`).
+
+**Bug fixes**:
+
+- Various minor bug and doc fixes from :user:`hyanwong`, :user:`petrelharp`,
+  :user:`brianzhang01`, :user:`mufernando` and :user:`andrewkern`.
+
+
 **********************
 [0.7.1b1] - 2019-05-31
 **********************

msprime/simulations.py

@@ -709,6 +709,12 @@ def get_tree_sequence(self, mutation_generator=None, provenance_record=None):
         tables = tskit.TableCollection.fromdict(self.ll_tables.asdict())
         if provenance_record is not None:
             tables.provenances.add_row(provenance_record)
+        if self.from_ts is None:
+            # Add the populations with metadata
+            assert len(tables.populations) == len(self.population_configurations)
+            tables.populations.clear()
+            for pop_config in self.population_configurations:
+                tables.populations.add_row(metadata=pop_config.encoded_metadata)
         return tables.tree_sequence()
 
     def reset(self):
@@ -908,15 +914,26 @@ class PopulationConfiguration(object):
         population per generation. Growth rates can be negative. This is zero for a
         constant population size, and positive for a population that has been
         growing. Defaults to 0.
+    :param dict metadata: A JSON-encodable dictionary of metadata to associate
+        with the corresponding Population in the output tree sequence.
+        If not specified or None, no metadata is stored (i.e., an empty bytes array).
+        Note that this metadata is ignored when using the ``from_ts`` argument to
+        :func:`simulate`, as the population definitions in the tree sequence that
+        is used as the starting point take precedence.
     """
-    def __init__(self, sample_size=None, initial_size=None, growth_rate=0.0):
+    def __init__(
+            self, sample_size=None, initial_size=None, growth_rate=0.0, metadata=None):
         if initial_size is not None and initial_size <= 0:
             raise ValueError("Population size must be > 0")
         if sample_size is not None and sample_size < 0:
             raise ValueError("Sample size must be >= 0")
         self.sample_size = sample_size
         self.initial_size = initial_size
         self.growth_rate = growth_rate
+        self.metadata = metadata
+        self.encoded_metadata = b''
+        if self.metadata is not None:
+            self.encoded_metadata = json.dumps(self.metadata).encode()
 
     def get_ll_representation(self):
         """
@@ -1379,7 +1396,10 @@ def __init__(
             population_configurations=population_configurations,
             migration_matrix=migration_matrix,
             demographic_events=demographic_events)
-        # TODO implement the model change events here.
+        if len(simulator.model_change_events) > 0:
+            raise ValueError(
+                "Model changes not currently supported by the DemographyDebugger. "
+                "Please open an issue on GitHub if this feature would be useful to you")
         assert len(simulator.model_change_events) == 0
         self._make_epochs(simulator, sorted(demographic_events, key=lambda e: e.time))
         self.simulation_model = simulator.model

tests/test_demography.py

@@ -23,6 +23,7 @@
 import math
 import tempfile
 import unittest
+import json
 
 import numpy as np
 import scipy.linalg
@@ -370,6 +371,16 @@ def test_equal_after(self):
         self.assertEqual(population_configurations[0].sample_size, 10)
         self.assertEqual(population_configurations[1].sample_size, 20)
 
+    def test_model_change_events(self):
+        population_configurations = [
+            msprime.PopulationConfiguration(sample_size=10)]
+        demographic_events = [
+            msprime.SimulationModelChange(1, "hudson")]
+        with self.assertRaises(ValueError):
+            msprime.DemographyDebugger(
+                population_configurations=population_configurations,
+                demographic_events=demographic_events)
+
     def test_one_pop_zero_events(self):
         dd = msprime.DemographyDebugger(
             population_configurations=[msprime.PopulationConfiguration()])
@@ -2183,3 +2194,53 @@ def test_4_populations(self):
                 model='dtwf')
         for node in ts.nodes():
             self.assertEqual(node.time, int(node.time))
+
+
+class TestPopulationMetadata(unittest.TestCase):
+    """
+    Tests for the metadata behaviour on populations.
+    """
+    def test_simple_case(self):
+        md = {"x": "y"}
+        ts = msprime.simulate(
+            population_configurations=[
+                msprime.PopulationConfiguration(2, metadata=md)],
+            random_seed=1)
+        self.assertEqual(ts.num_populations, 1)
+        pop = ts.population(0)
+        self.assertEqual(md, json.loads(pop.metadata.decode()))
+
+    def test_default(self):
+        ts = msprime.simulate(
+            population_configurations=[msprime.PopulationConfiguration(2)],
+            random_seed=1)
+        self.assertEqual(ts.num_populations, 1)
+        pop = ts.population(0)
+        self.assertEqual(b'', pop.metadata)
+
+        ts = msprime.simulate(
+            population_configurations=[
+                msprime.PopulationConfiguration(2, metadata=None)],
+            random_seed=1)
+        self.assertEqual(ts.num_populations, 1)
+        pop = ts.population(0)
+        self.assertEqual(b'', pop.metadata)
+
+    def test_errors(self):
+        for bad_metadata in [b"asdf", Exception]:
+            with self.assertRaises(TypeError):
+                msprime.PopulationConfiguration(2, metadata=bad_metadata)
+
+    def test_multi_population(self):
+        for num_pops in range(1, 10):
+            pop_configs = [
+                msprime.PopulationConfiguration(2, metadata={"x": "x" * j})
+                for j in range(num_pops)]
+            ts = msprime.simulate(
+                population_configurations=pop_configs,
+                random_seed=1, end_time=1)
+            self.assertEqual(ts.num_populations, num_pops)
+            for j in range(num_pops):
+                pop = ts.population(j)
+                self.assertEqual(
+                    pop_configs[j].metadata, json.loads(pop.metadata.decode()))