Support node_labels in write_nexus via TRANSLATE

Co-authored-by: Jerome Kelleher <jk@well.ox.ac.uk>

Author: kaathewisegit

python/CHANGELOG.rst

@@ -7,6 +7,8 @@ In development
 - Add ``json+struct`` metadata codec that allows storing binary data using a struct
   schema alongside JSON metadata. (:user:`benjeffery`, :pr:`3306`)
 
+- Add `node_labels` parameter to `write_nexus`. (:user:`kaathewisegit`, :pr:`3442`)
+
 --------------------
 [1.0.2] - 2026-03-06
 --------------------

python/tests/test_phylo_formats.py

@@ -474,6 +474,73 @@ def test_as_nexus_precision_1(self):
         )
         assert ts.as_nexus(precision=1) == expected
 
+    def test_as_nexus_labels_basic(self):
+        ts = self.tree().tree_sequence
+        labels = {0: "human", 1: "chimp", 2: "bonobo"}
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=3;
+              TAXLABELS human chimp bonobo;
+            END;
+            BEGIN TREES;
+              TRANSLATE n0 human, n1 chimp, n2 bonobo;
+              TREE t0^1 = [&R] (n0:2,(n1:1,n2:1):1);
+            END;
+            """
+        )
+        assert expected == ts.as_nexus(
+            include_alignments=False,
+            node_labels=labels,
+        )
+
+    def test_as_nexus_labels_partial(self):
+        ts = self.tree().tree_sequence
+        labels = {0: "human", 2: "bonobo"}
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=3;
+              TAXLABELS human n1 bonobo;
+            END;
+            BEGIN TREES;
+              TRANSLATE n0 human, n2 bonobo;
+              TREE t0^1 = [&R] (n0:2,(n1:1,n2:1):1);
+            END;
+            """
+        )
+        assert expected == ts.as_nexus(
+            include_alignments=False,
+            node_labels=labels,
+        )
+
+    def test_as_nexus_labels_empty(self):
+        ts = self.tree().tree_sequence
+
+        with pytest.raises(ValueError, match="`node_labels` cannot be empty"):
+            ts.as_nexus(include_alignments=False, node_labels={})
+
+    def test_as_nexus_labels_none(self):
+        ts = self.tree().tree_sequence
+        expected = textwrap.dedent(
+            """\
+            #NEXUS
+            BEGIN TAXA;
+              DIMENSIONS NTAX=3;
+              TAXLABELS n0 n1 n2;
+            END;
+            BEGIN TREES;
+              TREE t0^1 = [&R] (n0:2,(n1:1,n2:1):1);
+            END;
+            """
+        )
+        assert expected == ts.as_nexus(
+            include_alignments=False,
+            node_labels=None,
+        )
+
 
 class TestFractionalBranchLengths:
     # 0.67┊   4   ┊

python/tskit/text_formats.py

@@ -120,6 +120,7 @@ def write_nexus(
     include_alignments,
     reference_sequence,
     missing_data_character,
+    node_labels,
     isolated_as_missing=None,
 ):
     # See TreeSequence.write_nexus for documentation on parameters.
@@ -134,7 +135,21 @@ def write_nexus(
     print("#NEXUS", file=out)
     print("BEGIN TAXA;", file=out)
     print("", f"DIMENSIONS NTAX={ts.num_samples};", sep=indent, file=out)
-    taxlabels = " ".join(f"n{u}" for u in ts.samples())
+
+    if node_labels is not None:
+        if len(node_labels) == 0:
+            raise ValueError("`node_labels` cannot be empty")
+
+        if max(node_labels) >= ts.num_samples:
+            raise ValueError(
+                "Can only remap samples, got node with id {max(node_labels)}"
+            )
+
+        taxlabels = " ".join(
+            node_labels[u] if u in node_labels else f"n{u}" for u in ts.samples()
+        )
+    else:
+        taxlabels = " ".join(f"n{u}" for u in ts.samples())
     print("", f"TAXLABELS {taxlabels};", sep=indent, file=out)
     print("END;", file=out)
 
@@ -166,6 +181,11 @@ def write_nexus(
     include_trees = True if include_trees is None else include_trees
     if include_trees:
         print("BEGIN TREES;", file=out)
+
+        if node_labels is not None:
+            translations = ", ".join(f"n{u} {name}" for u, name in node_labels.items())
+            print(f"  TRANSLATE {translations};", file=out)
+
         for tree in ts.trees():
             start_interval = "{0:.{1}f}".format(tree.interval.left, pos_precision)
             end_interval = "{0:.{1}f}".format(tree.interval.right, pos_precision)

python/tskit/trees.py

@@ -6797,6 +6797,7 @@ def write_nexus(
         reference_sequence=None,
         missing_data_character=None,
         isolated_as_missing=None,
+        node_labels=None,
     ):
         """
         Returns a `nexus encoding <https://en.wikipedia.org/wiki/Nexus_file>`_
@@ -6896,6 +6897,10 @@ def write_nexus(
         :param str missing_data_character: As for the :meth:`.alignments` method,
             but defaults to "?".
         :param bool isolated_as_missing: As for the :meth:`.alignments` method.
+        :param node_labels: A map of type `{index: name}`.  Samples present in
+            the map will have the given name instead of `n{index}`.  Note that
+            the names must not have whitespace (spaces should be replaced by
+            underscores) or puncuation in them.
         :return: A nexus representation of this :class:`TreeSequence`
         :rtype: str
         """
@@ -6908,6 +6913,7 @@ def write_nexus(
             reference_sequence=reference_sequence,
             missing_data_character=missing_data_character,
             isolated_as_missing=isolated_as_missing,
+            node_labels=node_labels,
         )
 
     def as_nexus(self, **kwargs):