fix: pass marginal flag to all infer_ancestral_sequences calls

test/test_treetime.py

@@ -305,3 +305,59 @@ def real_lh():
     print(abs(ref.max() - real))
     # joint chooses the most likely realization of the tree
     assert abs(ref.max() - real) < 1e-10
+
+
+def test_marginal_mode_used_in_all_iterations(root_dir=None):
+    """Verify that marginal reconstruction is used in every call to
+    infer_ancestral_sequences when branch_length_mode='marginal'.
+
+    Regression test for https://github.com/neherlab/treetime/issues/601:
+    later iterations passed marginal_sequences via **seq_kwargs but not as the
+    explicit `marginal` parameter, causing infer_ancestral_sequences to default
+    to joint reconstruction.
+    """
+    import os
+    from unittest.mock import patch
+    from treetime import TreeTime
+    from treetime.utils import parse_dates
+
+    if root_dir is None:
+        root_dir = os.path.dirname(os.path.realpath(__file__))
+
+    fasta = root_dir + '/treetime_examples/data/h3n2_na/h3n2_na_20.fasta'
+    nwk = root_dir + '/treetime_examples/data/h3n2_na/h3n2_na_20.nwk'
+    dates = parse_dates(root_dir + '/treetime_examples/data/h3n2_na/h3n2_na_20.metadata.csv')
+
+    myTree = TreeTime(
+        gtr='Jukes-Cantor',
+        tree=nwk,
+        use_fft=False,
+        aln=fasta,
+        verbose=1,
+        dates=dates,
+        precision=3,
+        debug=True,
+        rng_seed=1234,
+    )
+
+    marginal_args_log = []
+    original = myTree.infer_ancestral_sequences.__func__
+
+    def tracking_wrapper(self, *args, marginal=False, **kwargs):
+        marginal_args_log.append(marginal)
+        return original(self, *args, marginal=marginal, **kwargs)
+
+    with patch.object(type(myTree), 'infer_ancestral_sequences', tracking_wrapper):
+        myTree.run(
+            branch_length_mode='marginal',
+            infer_gtr=False,
+            max_iter=2,
+            time_marginal=False,
+        )
+
+    assert len(marginal_args_log) > 0, "infer_ancestral_sequences was never called"
+    for i, was_marginal in enumerate(marginal_args_log):
+        assert was_marginal, (
+            f"Call #{i} to infer_ancestral_sequences used marginal=False (joint) "
+            f"when marginal reconstruction was requested"
+        )

treetime/treetime.py

@@ -261,7 +261,7 @@ def _run(
 
         if self.branch_length_mode == 'input':
             if self.aln:
-                self.infer_ancestral_sequences(**seq_kwargs)
+                self.infer_ancestral_sequences(marginal=seq_kwargs['marginal_sequences'], **seq_kwargs)
         else:
             self.optimize_tree(max_iter=1, method_anc=method_anc, **seq_kwargs)
 
@@ -342,10 +342,14 @@ def _run(
             if need_new_time_tree:
                 self.make_time_tree(**tt_kwargs)
                 if self.aln:
-                    ndiff = self.infer_ancestral_sequences('ml', **seq_kwargs)
+                    ndiff = self.infer_ancestral_sequences(
+                        'ml', marginal=seq_kwargs['marginal_sequences'], **seq_kwargs
+                    )
             else:  # no refinements, just iterate
                 if self.aln:
-                    ndiff = self.infer_ancestral_sequences('ml', **seq_kwargs)
+                    ndiff = self.infer_ancestral_sequences(
+                        'ml', marginal=seq_kwargs['marginal_sequences'], **seq_kwargs
+                    )
                 self.make_time_tree(**tt_kwargs)
 
             self.tree.coalescent_joint_LH = self.merger_model.total_LH() if Tc else 0.0