Test skeletons

Author: nspope

tests/test_inference.py

@@ -372,3 +372,41 @@ def test_simple_sim_multi_tree(self):
             assert all(
                 [a == b for a, b in zip(ts.haplotypes(), io_dated_ts.haplotypes())]
             )
+
+
+class TestVariational:
+    """
+    Tests for tsdate with variational algorithm
+    """
+
+    def test_simple_sim_1_tree(self):
+        ts = msprime.simulate(8, mutation_rate=5, random_seed=2)
+        tsdate.date(ts, mutation_rate=5, population_size=1, method="variational_gamma")
+
+    def test_simple_sim_multi_tree(self):
+        ts = msprime.simulate(8, mutation_rate=5, recombination_rate=5, random_seed=2)
+        tsdate.date(ts, mutation_rate=5, population_size=1, method="variational_gamma")
+
+    def test_nonglobal_priors(self):
+        ts = msprime.simulate(8, mutation_rate=5, recombination_rate=5, random_seed=2)
+        priors = tsdate.prior.MixturePrior(ts, prior_distribution="gamma")
+        grid = priors.make_parameter_grid(population_size=1)
+        grid.grid_data[:] = [1.0, 0.0]  # noninformative prior
+        tsdate.date(
+            ts,
+            mutation_rate=5,
+            method="variational_gamma",
+            priors=grid,
+            global_prior=False,
+        )
+
+    def test_bad_arguments(self):
+        ts = utility_functions.two_tree_mutation_ts()
+        with pytest.raises(ValueError, match="Maximum number of iterations"):
+            tsdate.date(
+                ts,
+                mutation_rate=5,
+                population_size=1,
+                method="variational_gamma",
+                max_iterations=-1,
+            )

tsdate/approx.py

@@ -111,8 +111,8 @@ def sufficient_statistics(a_i, b_i, a_j, b_j, y_ij, mu_ij):
 
     :return: normalizing constant, E[t_i], E[log t_i], E[t_j], E[log t_j]
     """
-    assert a_i > 0 and b_i > 0, "Invalid parent parameters"
-    assert a_j > 0 and b_j > 0, "Invalid child parameters"
+    assert a_i > 0 and b_i >= 0, "Invalid parent parameters"
+    assert a_j > 0 and b_j >= 0, "Invalid child parameters"
     assert y_ij >= 0 and mu_ij > 0, "Invalid edge parameters"
 
     a = a_i + a_j + y_ij

tsdate/core.py

@@ -630,17 +630,20 @@ def __init__(self, priors, lik, *, progress=False):
 
     # === Grouped edge iterators ===
 
-    def edges_by_parent_asc(self):
+    def edges_by_parent_asc(self, grouped=True):
         """
         Return an itertools.groupby object of edges grouped by parent in ascending order
         of the time of the parent. Since tree sequence properties guarantee that edges
         are listed in nondecreasing order of parent time
         (https://tskit.readthedocs.io/en/latest/data-model.html#edge-requirements)
         we can simply use the standard edge order
         """
-        return itertools.groupby(self.ts.edges(), operator.attrgetter("parent"))
+        if grouped:
+            return itertools.groupby(self.ts.edges(), operator.attrgetter("parent"))
+        else:
+            return self.ts.edges()
 
-    def edges_by_child_desc(self):
+    def edges_by_child_desc(self, grouped=True):
         """
         Return an itertools.groupby object of edges grouped by child in descending order
         of the time of the child.
@@ -651,9 +654,12 @@ def edges_by_child_desc(self):
                 (self.ts.edges_child, -self.ts.nodes_time[self.ts.edges_child])
             )
         )
-        return itertools.groupby(it, operator.attrgetter("child"))
+        if grouped:
+            return itertools.groupby(it, operator.attrgetter("child"))
+        else:
+            return it
 
-    def edges_by_child_then_parent_desc(self):
+    def edges_by_child_then_parent_desc(self, grouped=True):
         """
         Return an itertools.groupby object of edges grouped by child in descending order
         of the time of the child, then by descending order of age of child
@@ -675,7 +681,10 @@ def edges_by_child_then_parent_desc(self):
                 np.argsort(w, order=("child_age", "child_node", "parent_age"))
             )
         )
-        return itertools.groupby(sorted_child_parent, operator.attrgetter("child"))
+        if grouped:
+            return itertools.groupby(sorted_child_parent, operator.attrgetter("child"))
+        else:
+            return sorted_child_parent
 
     # === MAIN ALGORITHMS ===
 
@@ -1038,8 +1047,8 @@ def iterate(self, *, progress=None, **kwargs):
         Update edge factors from leaves to root then from root to leaves,
         and return approximate log marginal likelihood
         """
-        self.propagate(edges=self.edges_by_parent_asc(), progress=progress)
-        self.propagate(edges=self.edges_by_child_desc(), progress=progress)
+        self.propagate(edges=self.edges_by_parent_asc(grouped=False), progress=progress)
+        self.propagate(edges=self.edges_by_child_desc(grouped=False), progress=progress)
         # TODO
         # marginal_lik = np.sum(self.factor_norm)
         # return marginal_lik

tsdate/hypergeo.py

@@ -341,8 +341,8 @@ def _hyp2f1_dlmf1583(a_i, b_i, a_j, b_j, y, mu):
     """
     DLMF 15.8.3, sum of recurrence and series expansion
     """
-    assert b_i > 0
-    assert 0 < mu <= b_j
+    assert b_i >= 0
+    assert 0 <= mu <= b_j
     assert y >= 0 and y % 1.0 == 0.0
 
     f_1, s_1, da_i_1, db_i_1, da_j_1, db_j_1 = _hyp2f1_dlmf1583_first(
@@ -380,8 +380,8 @@ def _hyp2f1_dlmf1521(a_i, b_i, a_j, b_j, y, mu):
     """
     DLMF 15.2.1, series expansion without transformation
     """
-    assert b_i > 0
-    assert mu >= b_j > 0
+    assert b_i >= 0
+    assert mu >= b_j >= 0
     assert y >= 0 and y % 1 == 0.0
 
     y = int(y)
@@ -409,8 +409,8 @@ def _hyp2f1_dlmf1581(a_i, b_i, a_j, b_j, y, mu):
     """
     DLMF 15.8.1, series expansion with Pfaff transformation
     """
-    assert b_i > 0
-    assert 0 < mu <= b_j
+    assert b_i >= 0
+    assert 0 <= mu <= b_j
     assert y >= 0 and y % 1 == 0.0
 
     y = int(y)
@@ -456,7 +456,7 @@ def _hyp2f1(a_i, b_i, a_j, b_j, y, mu):
     and dividing the gradient by the function value.
     """
     z = (mu - b_j) / (mu + b_i)
-    assert z < 1.0  # TODO: allow z == 1.0 for improper prior
+    assert z < 1.0, "Invalid hypergeometric function argument"
     if 0.0 <= z < 1.0:
         return _hyp2f1_dlmf1521(a_i, b_i, a_j, b_j, y, mu)
     elif -1.0 < z < 0.0: