Add isolated_as_missing to ts.alignements

Author: benjeffery

python/CHANGELOG.rst

@@ -1,5 +1,5 @@
 ----------------------
-[1.0.0b1] - 2025-09-24
+[1.0.0b2] - 2025-XX-XX
 ----------------------
 
 **Breaking Changes**
@@ -21,6 +21,10 @@
 - ``draw_svg()`` methods now associate tree branches with edge IDs
   (:user:`hyanwong`, :pr:`3193`, :issue:`557`)
 
+- ``TreeSequence.alignments`` now accepts ``isolated_as_missing=False`` so that
+  alignments can be emitted for non-sample nodes (e.g., internal ARG nodes) when
+  missing data are imputed to the ancestral state. (:user:`benjeffery`, :issue:`3293`)
+
 - ``draw_svg()`` methods now allow the y-axis to be placed on the right-hand side
   using ``y_axis="right"`` (:user:`hyanwong`, :pr:`3201`)
 

python/tests/test_genotypes.py

@@ -1286,6 +1286,22 @@ def test_non_sample_samples(self):
         with pytest.raises(tskit.LibraryError, match="MUST_IMPUTE_NON_SAMPLES"):
             list(ts.alignments(samples=[4]))
 
+    def test_internal_nodes_with_imputation(self):
+        ts = self.ts()
+        internal = [u for u in range(ts.num_nodes) if u not in set(ts.samples())]
+        seqs = list(ts.alignments(samples=internal, isolated_as_missing=False))
+        assert seqs == ["NNANNNNNNC", "NNANNNNNNT"]
+
+    def test_internal_and_sample_nodes_with_imputation(self):
+        ts = self.ts()
+        samples = ts.samples()
+        internal = next(u for u in range(ts.num_nodes) if u not in set(samples))
+        seqs = list(
+            ts.alignments(samples=[internal, samples[0]], isolated_as_missing=False)
+        )
+        assert seqs[0] == "NNANNNNNNC"
+        assert seqs[1] == "NNGNNNNNNT"
+
     def test_alignments_missing_data_char(self):
         A = list(self.ts().alignments(missing_data_character="x"))
         assert A[0] == "xxGxxxxxxT"

python/tskit/trees.py

@@ -5596,6 +5596,7 @@ def alignments(
         samples=None,
         left=None,
         right=None,
+        isolated_as_missing=None,
     ):
         """
         Returns an iterator over the full sequence alignments for the defined samples
@@ -5660,7 +5661,9 @@ def alignments(
            currently supported by this method and it will raise a ValueError
            if called on tree sequences containing isolated samples.
            See https://github.com/tskit-dev/tskit/issues/1896 for more
-           information.
+           information. If you wish to include non-sample nodes (e.g., internal
+           ARG nodes) in the output, set ``isolated_as_missing=False`` to opt out
+           of this guard and impute missing data as the ancestral state.
 
         See also the :meth:`.variants` iterator for site-centric access
         to sample genotypes and :meth:`.haplotypes` for access to sample sequences
@@ -5679,6 +5682,11 @@ def alignments(
             (default) alignments start at 0.
         :param int right: Alignments will stop before this genomic position. If ``None``
             (default) alignments will continue until the end of the tree sequence.
+        :param bool isolated_as_missing: If True (default), isolated samples without
+            mutations are treated as missing data and this method raises an error
+            when any are detected. If False, missing data are imputed with the
+            ancestral state, which also permits alignments to be generated for
+            non-sample nodes such as internal ARG nodes.
         :return: An iterator over the alignment strings for specified samples in
             this tree sequence, in the order given in ``samples``.
         :rtype: collections.abc.Iterable
@@ -5693,6 +5701,8 @@ def alignments(
         missing_data_character = (
             "N" if missing_data_character is None else missing_data_character
         )
+        if isolated_as_missing is None:
+            isolated_as_missing = True
 
         L = interval.span
         a = np.empty(L, dtype=np.int8)
@@ -5730,7 +5740,9 @@ def alignments(
         # incorrectly if have a sample isolated over the region (a, b],
         # and if we have sites at each position from a to b, and at
         # each site there is a mutation over the isolated sample.
-        if any(tree._has_isolated_samples() for tree in self.trees()):
+        if isolated_as_missing and any(
+            tree._has_isolated_samples() for tree in self.trees()
+        ):
             raise ValueError(
                 "Missing data not currently supported in alignments; see "
                 "https://github.com/tskit-dev/tskit/issues/1896 for details."
@@ -5741,6 +5753,7 @@ def alignments(
             interval=interval,
             missing_data_character=missing_data_character,
             samples=samples,
+            isolated_as_missing=isolated_as_missing,
         )
         site_pos = self.sites_position.astype(np.int64)[
             first_site_id : last_site_id + 1