Merge pull request #248 from nspope/ne-rescale-time

Priors with piecewise demographic histories

Author: hyanwong

.gitignore

@@ -9,3 +9,5 @@ tsdate.egg-info
 data/prior_1000df.txt
 *.ipynb
 .ipynb*
+build/
+.*.swp

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
     hooks:
       - id: black
         language_version: python3
-  - repo: https://gitlab.com/pycqa/flake8
+  - repo: https://github.com/pycqa/flake8
     rev: 3.9.2
     hooks:
       - id: flake8

CHANGELOG.rst

@@ -12,11 +12,22 @@
   in ``get_dates`` and the internal ``normalize`` terminology is changed to
   ``standardize`` to better reflect the fact that the maximum (not sum) is one.
 
+- The ``Ne`` argument to ``date`` has been deprecated (although it is
+  still in the API for backward compatibility).  The equivalent argument
+  ``population_size`` should be used instead.
+
+**Features**
+
+- Priors may be calculated using a piecewise-constant effective population trajectory,
+  which is implemented in the ``demography.PopulationSizeHistory`` class. The
+  ``population_size`` argument to ``date`` accepts either a single scalar effective
+  population size, or a ``PopulationSizeHistory`` instance.
+
 **Bugfixes**
 
 - The returned posteriors when ``return_posteriors=True`` now return actual
   probabilities (scaled so that they sum to one) rather than standardized
-  "probabilites" whose maximum value is one.
+  "probabilities" whose maximum value is one.
 
 --------------------
 [0.1.5] - 2022-06-07

tests/test_accuracy.py

@@ -104,7 +104,7 @@ def test_basic(
         mu = sim_mutations_parameters["rate"]
 
         dts, posteriors = tsdate.date(
-            ts, Ne=Ne, mutation_rate=mu, return_posteriors=True
+            ts, population_size=Ne, mutation_rate=mu, return_posteriors=True
         )
         # make sure we can read node metadata - old tsdate versions didn't set a schema
         if dts.table_metadata_schemas.node.schema is None:
@@ -139,6 +139,38 @@ def test_scaling(self, Ne):
         Test that we are in the right theoretical ballpark given known Ne
         """
         ts = tskit.Tree.generate_comb(2).tree_sequence
-        dts = tsdate.date(ts, Ne=Ne, mutation_rate=None)
+        dts = tsdate.date(ts, population_size=Ne, mutation_rate=None)
         # Check the date is within 10% of the expected
         assert 0.9 < dts.node(dts.first().root).time / (2 * Ne) < 1.1
+
+    @pytest.mark.parametrize(
+        "bkwd_rate, trio_tmrca",
+        [  # calculated from simulations
+            (-1.0, 0.76),
+            (-0.9, 0.79),
+            (-0.8, 0.82),
+            (-0.7, 0.85),
+            (-0.6, 0.89),
+            (-0.5, 0.94),
+            (-0.4, 0.99),
+            (-0.3, 1.05),
+            (-0.2, 1.12),
+            (-0.1, 1.21),
+            (0.0, 1.32),
+        ],
+    )
+    def test_piecewise_scaling(self, bkwd_rate, trio_tmrca):
+        """
+        Test that we are in the right theoretical ballpark given known Ne,
+        under exponential growth.
+
+        Check coalescence time of a trio instead of a pair, because of
+        https://github.com/tskit-dev/tsdate/issues/230
+        """
+        time = np.linspace(0, 10, 100)
+        ne = 0.5 * np.exp(bkwd_rate * time)
+        ts = tskit.Tree.generate_comb(3).tree_sequence
+        demo = tsdate.demography.PopulationSizeHistory(ne, time[1:])
+        dts = tsdate.date(ts, population_size=demo, mutation_rate=None)
+        # Check the date is within 10% of the expected
+        assert 0.9 < dts.node(dts.first().root).time / trio_tmrca < 1.1

tests/test_cache.py

@@ -27,13 +27,13 @@ def test_cached_prior(self):
         if os.path.isfile(fn):
             raise AssertionError(f"The file {fn} already exists. Delete before testing")
         with self.assertLogs(level="WARNING") as log:
-            priors_approx10 = ConditionalCoalescentTimes(10, Ne=1)
+            priors_approx10 = ConditionalCoalescentTimes(10)
             assert len(log.output) == 1
             assert "user cache" in log.output[0]
         priors_approx10.add(10)
         # Check we have created the prior file
         assert os.path.isfile(fn)
-        priors_approxNone = ConditionalCoalescentTimes(None, Ne=1)
+        priors_approxNone = ConditionalCoalescentTimes(None)
         priors_approxNone.add(10)
         assert np.allclose(priors_approx10[10], priors_approxNone[10], equal_nan=True)
         # Test when using a bigger n that we're using the precalculated version

tests/test_cli.py

@@ -50,7 +50,7 @@ def test_default_values(self):
             args = parser.parse_args(["date", self.infile, self.output, "1"])
         assert args.tree_sequence == self.infile
         assert args.output == self.output
-        assert args.Ne == 1
+        assert args.population_size == 1
         assert args.mutation_rate is None
         assert args.recombination_rate is None
         assert args.epsilon == 1e-6
@@ -216,7 +216,7 @@ def compare_python_api(self, input_ts, cmd, Ne, mutation_rate, method):
             cli.tsdate_main(full_cmd.split())
             output_ts = tskit.load(output_filename)
         dated_ts = tsdate.date(
-            input_ts, Ne=Ne, mutation_rate=mutation_rate, method=method
+            input_ts, population_size=Ne, mutation_rate=mutation_rate, method=method
         )
         print(dated_ts.tables.nodes.time, output_ts.tables.nodes.time)
         assert np.array_equal(dated_ts.tables.nodes.time, output_ts.tables.nodes.time)