Docstrings and CLI

tsdate/cli.py

@@ -199,6 +199,34 @@ def tsdate_cli_parser():
             "but does not exactly minimize KL divergence in each EP update."
         ),
     )
+    parser.add_argument(
+        "--max-iterations",
+        type=int,
+        help=(
+            "The number of iterations used in the expectation propagation "
+            "algorithm. Default: 20"
+        ),
+        default=20,
+    )
+    parser.add_argument(
+        "--em-iterations",
+        type=int,
+        help=(
+            "The number of expectation-maximization iterations used to optimize the "
+            "global mixture prior at the end of each expectation propagation step. "
+            "Setting to zero disables optimization. Default: 10"
+        ),
+        default=10,
+    )
+    parser.add_argument(
+        "--global-prior",
+        type=int,
+        help=(
+            "The number of components in the i.i.d. mixture prior for node "
+            "ages. Default: 1"
+        ),
+        default=1,
+    )
     parser.set_defaults(runner=run_date)
 
     parser = subparsers.add_parser(
@@ -253,8 +281,11 @@ def run_date(args):
             method=args.method,
             eps=args.epsilon,
             progress=args.progress,
+            max_iterations=args.max_iterations,
             max_shape=args.max_shape,
             match_central_moments=args.match_central_moments,
+            em_iterations=args.em_iterations,
+            global_prior=args.global_prior,
         )
     else:
         params = dict(

tsdate/core.py

@@ -1478,17 +1478,20 @@ def variational_dates(
     which invokes this method and inserts the resulting node ages into the tree
     sequence.
 
+    An i.i.d. gamma mixture is used as a prior for each node, that is
+    initialized from the conditional coalescent and updated via expectation
+    maximization in each iteration. In addition, node-specific priors may be
+    specified via a grid of shape/rate parameters.
+
     :param ~tskit.TreeSequence tree_sequence: See :func:`date`.
     :param float mutation_rate: See :func:`date`.
     :param float population_size: See :func:`date`.
     :param float recombination_rate: See :func:`date`.
     :param bool progress: See :func:`date`.
     :param ~tsdate.base.NodeGridValues priors: the prior
         parameters for each node-to-be-dated, assuming a gamma prior on node
-        age and using shape/rate parameterization. If ``None`` (default), use
-        an iid prior derived from the conditional coalescent prior, tilted
-        according to population size, and assume the nodes to be dated are all
-        the non-sample nodes in the input tree sequence.
+        age and using shape/rate parameterization. If ``None`` (default), node
+        specific priors are omitted and only a global mixture prior is used.
     :param int max_iterations: The number of iterations used in the expectation
         propagation algorithm. Default: 20.
     :param float max_shape: The maximum value for the shape parameter in the variational
@@ -1500,7 +1503,7 @@ def variational_dates(
         Kullback-Leibler divergence. Default: False.
     :param int global_prior: The number of components in the i.i.d. mixture prior
         for node ages. Default: 1.
-    :param int em_iterations: The number of expectation-maximization iterations used
+    :param int em_iterations: The number of expectation maximization iterations used
         to optimize the global mixture prior. Setting to zero disables optimization.
         Default: 10.