Merge pull request #261 from hyanwong/cli-variational

Allow cli to work with variational_gamma

Author: hyanwong

tests/test_cli.py

@@ -55,7 +55,7 @@ def test_default_values(self):
         assert args.recombination_rate is None
         assert args.epsilon == 1e-6
         assert args.num_threads is None
-        assert args.probability_space == "logarithmic"
+        assert args.probability_space is None  # Use the defaults
         assert args.method == "inside_outside"
         assert not args.progress
 
@@ -128,16 +128,15 @@ def test_probability_space(self):
         )
         assert args.probability_space == "logarithmic"
 
-    def test_method(self):
+    @pytest.mark.parametrize(
+        "method", ["inside_outside", "maximization", "variational_gamma"]
+    )
+    def test_method(self, method):
         parser = cli.tsdate_cli_parser()
         args = parser.parse_args(
-            ["date", self.infile, self.output, "10000", "--method", "inside_outside"]
-        )
-        assert args.method == "inside_outside"
-        args = parser.parse_args(
-            ["date", self.infile, self.output, "10000", "--method", "maximization"]
+            ["date", self.infile, self.output, "10000", "--method", method]
         )
-        assert args.method == "maximization"
+        assert args.method == method
 
     def test_progress(self):
         parser = cli.tsdate_cli_parser()
@@ -262,7 +261,10 @@ def test_method(self):
         with pytest.raises(ValueError):
             self.verify(input_ts, cmd)
 
-    def test_compare_python_api(self):
+    @pytest.mark.parametrize(
+        "method", ["inside_outside", "maximization", "variational_gamma"]
+    )
+    def test_compare_python_api(self, method):
         input_ts = msprime.simulate(
             100,
             Ne=10000,
@@ -271,12 +273,9 @@ def test_compare_python_api(self):
             length=2e4,
             random_seed=10,
         )
-        cmd = "10000 -m 1e-8 --method inside_outside"
-        self.verify(input_ts, cmd)
-        self.compare_python_api(input_ts, cmd, 10000, 1e-8, "inside_outside")
-        cmd = "10000 -m 1e-8 --method maximization"
+        cmd = f"10000 -m 1e-8 --method {method}"
         self.verify(input_ts, cmd)
-        self.compare_python_api(input_ts, cmd, 10000, 1e-8, "maximization")
+        self.compare_python_api(input_ts, cmd, 10000, 1e-8, method)
 
     def preprocess_compare_python_api(self, input_ts):
         with tempfile.TemporaryDirectory() as tmpdir:

tests/test_functions.py

@@ -46,6 +46,7 @@
 from tsdate.core import Likelihoods
 from tsdate.core import LogLikelihoods
 from tsdate.core import posterior_mean_var
+from tsdate.core import variational_dates
 from tsdate.demography import PopulationSizeHistory
 from tsdate.prior import ConditionalCoalescentTimes
 from tsdate.prior import fill_priors
@@ -1531,6 +1532,25 @@ def test_bad_Ne(self):
             tsdate.build_prior_grid(ts, population_size=-10)
 
 
+class TestCallingErrors:
+    def test_bad_vgamma_probability_space(self):
+        ts = utility_functions.single_tree_ts_n2()
+        with pytest.raises(ValueError, match="Cannot specify"):
+            variational_dates(ts, 1, 1, probability_space=base.LOG)
+
+    def test_bad_vgamma_num_threads(self):
+        # Test can be removed if we specify num_threads in the future
+        ts = utility_functions.single_tree_ts_n2()
+        with pytest.raises(ValueError, match="does not currently"):
+            variational_dates(ts, 1, 1, num_threads=2)
+
+    def test_bad_vgamma_ignore_oldest_root(self):
+        # Test can be removed in the future if this is implemented
+        ts = utility_functions.single_tree_ts_n2()
+        with pytest.raises(ValueError, match="not implemented"):
+            variational_dates(ts, 1, 1, ignore_oldest_root=True)
+
+
 class TestPosteriorMeanVar:
     """
     Test posterior_mean_var works as expected

tsdate/cli.py

@@ -125,26 +125,29 @@ def tsdate_cli_parser():
     parser.add_argument(
         "--probability-space",
         type=str,
-        default="logarithmic",
+        default=None,
         help="Should the internal algorithm save probabilities in \
                         'logarithmic' (slower, less liable to to overflow) or 'linear' \
-                        space (faster, may overflow). Default: 'logarithmic'",
+                        space (faster, may overflow). Not relevant for the \
+                        'variational_gamma' method; default otherwise is `None` \
+                        currently treated as 'logarithmic'",
     )
     parser.add_argument(
         "--method",
         type=str,
         default="inside_outside",
         help="Specify which estimation method to use: can be \
                         'inside_outside' (empirically better, theoretically \
-                        problematic) or 'maximization' (worse empirically, especially \
-                        with a gamma approximated prior, but theoretically robust). \
-                        Default: 'inside_outside'",
+                        problematic), 'maximization' (worse empirically, especially \
+                        with a gamma approximated prior, but theoretically robust), or \
+                        'variational_gamma' (a fast experimental continuous-time \
+                        approximation). Default: 'inside_outside'",
     )
     parser.add_argument(
         "--ignore-oldest",
         action="store_true",
-        help="Ignore the oldest node in the tree sequence, which is \
-                        often of low quality when using empirical data.",
+        help="Ignore the oldest node in the tree sequence: in older tsinfer versions \
+                        this could be of low quality when using empirical data.",
     )
     parser.add_argument(
         "-p", "--progress", action="store_true", help="Show progress bar."
@@ -190,18 +193,18 @@ def run_date(args):
         ts = tskit.load(args.tree_sequence)
     except tskit.FileFormatError as ffe:
         error_exit(f"Error loading '{args.tree_sequence}: {ffe}")
-    dated_ts = tsdate.date(
-        ts,
-        args.mutation_rate,
-        args.population_size,
+    params = dict(
         recombination_rate=args.recombination_rate,
-        probability_space=args.probability_space,
         method=args.method,
         eps=args.epsilon,
+        progress=args.progress,
+        probability_space=args.probability_space,
         num_threads=args.num_threads,
         ignore_oldest_root=args.ignore_oldest,
-        progress=args.progress,
     )
+    # TODO: error out if ignore_oldest_root is set,
+    # see https://github.com/tskit-dev/tsdate/issues/262
+    dated_ts = tsdate.date(ts, args.mutation_rate, args.population_size, **params)
     dated_ts.dump(args.output)
 
 

tsdate/core.py

@@ -1287,7 +1287,7 @@ def get_dates(
     ignore_oldest_root=False,
     progress=False,
     cache_inside=False,
-    probability_space=base.LOG,
+    probability_space=None,
 ):
     """
     Infer dates for the nodes in a tree sequence, returning an array of inferred dates
@@ -1333,6 +1333,9 @@ def get_dates(
             )
         priors = priors
 
+    if probability_space is None:
+        probability_space = base.LOG
+
     if probability_space != base.LOG:
         liklhd = Likelihoods(
             tree_sequence,
@@ -1447,6 +1450,9 @@ def variational_dates(
     global_prior=True,
     eps=1e-6,
     progress=False,
+    num_threads=None,  # Unused, matches get_dates()
+    probability_space=None,  # Can only be None, simply to match get_dates()
+    ignore_oldest_root=False,  # Can only be False, simply to match get_dates()
 ):
     """
     Infer dates for the nodes in a tree sequence using expectation propagation,
@@ -1469,6 +1475,21 @@ def variational_dates(
     if not max_iterations >= 1:
         raise ValueError("Maximum number of iterations must be greater than 0")
 
+    # Parameters below are not used in variational dating, but are here
+    # to match the signature of get_dates(). We may be able to remove some
+    # if we move to specifying some params via a control dictionary
+
+    if probability_space is not None:
+        raise ValueError("Cannot specify a probability space in variational dating")
+
+    if num_threads is not None and num_threads != 1:
+        raise ValueError("Variational dating does not currently use multiple threads")
+
+    if ignore_oldest_root:
+        raise ValueError(
+            "Ignoring the oldes root is not implemented in variational dating"
+        )
+
     # Default to not creating approximate priors unless ts has > 1000 samples
     approx_priors = False
     if tree_sequence.num_samples > 1000: