Rework topological constraint func

Fixes #216

Author: hyanwong

tests/test_functions.py

@@ -742,14 +742,16 @@ def test_logsumexp(self):
         assert np.allclose(LogLikelihoods.logsumexp(log_lls), np.log(ll_sum))
 
     def test_zeros_logsumexp(self):
-        lls = np.log(np.concatenate([np.zeros(100), np.random.rand(1000)]))
-        assert np.allclose(LogLikelihoods.logsumexp(lls), self.naive_logsumexp(lls))
+        with np.errstate(divide="ignore"):
+            lls = np.log(np.concatenate([np.zeros(100), np.random.rand(1000)]))
+            assert np.allclose(LogLikelihoods.logsumexp(lls), self.naive_logsumexp(lls))
 
     def test_logsumexp_underflow(self):
         # underflow in the naive case, but not in the LogLikelihoods implementation
-        lls = np.array([-1000, -1001])
-        assert self.naive_logsumexp(lls) == -np.inf
-        assert LogLikelihoods.logsumexp(lls) != -np.inf
+        with np.errstate(divide="ignore"):
+            lls = np.array([-1000, -1001])
+            assert self.naive_logsumexp(lls) == -np.inf
+            assert LogLikelihoods.logsumexp(lls) != -np.inf
 
     def test_log_tri_functions(self):
         ts = utility_functions.two_tree_mutation_ts()
@@ -1047,7 +1049,7 @@ def test_dangling_fails(self):
         print(ts.draw_text())
         print("Samples:", ts.samples())
         Ne = 0.5
-        with pytest.raises(ValueError, match="simplified"):
+        with pytest.raises(ValueError, match="simplify"):
             tsdate.build_prior_grid(ts, Ne, timepoints=np.array([0, 1.2, 2]))
         # mut_rate = 1
         # eps = 1e-6
@@ -1420,7 +1422,7 @@ def test_date_input(self):
 
     def test_sample_as_parent_fails(self):
         ts = utility_functions.single_tree_ts_n3_sample_as_parent()
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError, match="samples at non-zero times"):
             tsdate.date(ts, mutation_rate=None, Ne=1)
 
     def test_recombination_not_implemented(self):

tests/test_inference.py

@@ -61,7 +61,7 @@ def test_bad_Ne(self):
 
     def test_dangling_failure(self):
         ts = utility_functions.single_tree_ts_n2_dangling()
-        with pytest.raises(ValueError, match="simplified"):
+        with pytest.raises(ValueError, match="simplify"):
             tsdate.date(ts, mutation_rate=None, Ne=1)
 
     def test_unary_failure(self):
@@ -271,16 +271,29 @@ def test_fails_multi_root(self):
         with pytest.raises(ValueError):
             tsdate.date(multiroot_ts, Ne=1, mutation_rate=2, priors=good_priors)
 
-    def test_non_contemporaneous(self):
+    def test_non_contemporaneous_warn(self):
         samples = [
             msprime.Sample(population=0, time=0),
             msprime.Sample(population=0, time=0),
             msprime.Sample(population=0, time=0),
             msprime.Sample(population=0, time=1.0),
         ]
         ts = msprime.simulate(samples=samples, Ne=1, mutation_rate=2, random_seed=12)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError, match="samples at non-zero times"):
             tsdate.date(ts, Ne=1, mutation_rate=2)
+        with pytest.raises(ValueError, match="samples at non-zero times"):
+            tsdate.build_prior_grid(ts, Ne=1)
+
+    def test_non_contemporaneous(self):
+        samples = [
+            msprime.Sample(population=0, time=0),
+            msprime.Sample(population=0, time=0),
+            msprime.Sample(population=0, time=0),
+            msprime.Sample(population=0, time=1.0),
+        ]
+        ts = msprime.simulate(samples=samples, Ne=1, mutation_rate=2, random_seed=12)
+        priors = tsdate.build_prior_grid(ts, Ne=1, allow_historical_samples=True)
+        tsdate.date(ts, priors=priors, mutation_rate=2)
 
     def test_no_mutation_times(self):
         ts = msprime.simulate(20, Ne=1, mutation_rate=1, random_seed=12)

tsdate/base.py

@@ -134,6 +134,19 @@ def force_probability_space(self, probability_space):
         else:
             logging.warning("Cannot force", *descr)
 
+    def sum_to_unity(self, arr):
+        """
+        Return an array for a node in which the (untransformed) values
+        sum to unity
+        """
+        if self.probability_space == LIN:
+            return arr / np.sum(arr)
+        elif self.probability_space == LOG:
+            with np.errstate(divide="ignore"):
+                return np.log(np.exp(arr) / np.sum(np.exp(arr)))
+        else:
+            raise RuntimeError("Probability space is not", LIN, "or", LOG)
+
     def normalize(self):
         """
         normalize grid and fixed data so the max is one

tsdate/core.py

@@ -653,7 +653,10 @@ def inside_pass(self, *, normalize=True, cache_inside=False, progress=None):
             self.priors.nonfixed_node_ids(), self.ts.samples()
         )
         for u in nonfixed_samples:
-            inside[u][:] = self.priors[u]
+            # this is in the same probability space as the prior, so we should be
+            # OK just to copy the prior values straight in (but we should check they
+            # are normalised so that they sum to unity)
+            inside[u][:] = self.priors.sum_to_unity(self.priors[u])
 
         if cache_inside:
             g_i = np.full(
@@ -922,34 +925,31 @@ def posterior_mean_var(ts, timepoints, posterior, *, fixed_node_set=None):
     return ts, mn_post, vr_post
 
 
-def constrain_ages_topo(ts, post_mn, eps, nodes_to_date=None, progress=False):
+def constrain_ages_topo(ts, node_times, eps, progress=False):
     """
-    If predicted node times violate topology, restrict node ages so that they
-    must be older than all their children.
+    If node_times violate topology, return increased node_times so that each node is
+    guaranteed to be older than any of its their children.
     """
-    new_mn_post = np.copy(post_mn)
-    if nodes_to_date is None:
-        nodes_to_date = np.arange(ts.num_nodes, dtype=np.uint64)
-        nodes_to_date = nodes_to_date[~np.isin(nodes_to_date, ts.samples())]
-
     tables = ts.tables
-    parents = tables.edges.parent
-    nd_children = tables.edges.child[np.argsort(parents)]
-    parents = sorted(parents)
-    parents_unique = np.unique(parents, return_index=True)
-    parent_indices = parents_unique[1][np.isin(parents_unique[0], nodes_to_date)]
-    for index, nd in tqdm(
-        enumerate(sorted(nodes_to_date)), desc="Constrain Ages", disable=not progress
+    new_node_times = np.copy(node_times)
+    # Traverse through the ARG, ensuring children come before parents.
+    # This can be done by iterating over groups of edges with the same parent
+    new_parent_edge_idx = np.concatenate(
+        (
+            [0],
+            np.where(np.diff(tables.edges.parent) != 0)[0] + 1,
+            [tables.edges.num_rows],
+        )
+    )
+    for edges_start, edges_end in zip(
+        new_parent_edge_idx[:-1], new_parent_edge_idx[1:]
     ):
-        if index + 1 != len(nodes_to_date):
-            children_index = np.arange(parent_indices[index], parent_indices[index + 1])
-        else:
-            children_index = np.arange(parent_indices[index], ts.num_edges)
-        children = nd_children[children_index]
-        time = np.max(new_mn_post[children])
-        if new_mn_post[nd] <= time:
-            new_mn_post[nd] = time + eps
-    return new_mn_post
+        parent = tables.edges.parent[edges_start]
+        child_ids = tables.edges.child[edges_start:edges_end]  # May contain dups
+        oldest_child_time = np.max(new_node_times[child_ids])
+        if oldest_child_time >= new_node_times[parent]:
+            new_node_times[parent] = oldest_child_time + eps
+    return new_node_times
 
 
 def date(
@@ -1040,7 +1040,7 @@ def date(
         progress=progress,
         **kwargs
     )
-    constrained = constrain_ages_topo(tree_sequence, dates, eps, nds, progress)
+    constrained = constrain_ages_topo(tree_sequence, dates, eps, progress)
     tables = tree_sequence.dump_tables()
     tables.time_units = time_units
     tables.nodes.time = constrained