.

Author: DaanJansen94

recmpox/recmpox.py

@@ -1558,24 +1558,14 @@ def _run_one_phylogeny(
             for i in range(0, len(seq), line_len):
                 out.write(seq[i : i + line_len] + "\n")
         for sid, seq in part1_seqs.items():
-            # Tract FASTA headers include coverage suffix like
-            # "<safe_sample_id>_<ref1_label>_tract_HC_<cov>%". Recover the base
-            # sample ID so that phylogeny tract IDs can be mapped back to the
-            # per-genome records by stripping the coverage part.
-            base_id = sid
-            cov_suffix = f"_{ref1_label}_tract_HC_"
-            if cov_suffix in base_id:
-                base_id = base_id.split(cov_suffix, 1)[0]
-            header = _safe_fasta_id(base_id) + f"_{ref1_label}_tracts"
+            # Use tract FASTA headers (which already encode sample, clade and coverage)
+            # but sanitize them for downstream tools (e.g. remove '%' characters).
+            header = _safe_fasta_id(sid)
             out.write(f">{header}\n")
             for i in range(0, len(seq), line_len):
                 out.write(seq[i : i + line_len] + "\n")
         for sid, seq in part2_seqs.items():
-            base_id = sid
-            cov_suffix = f"_{ref2_label}_tract_HC_"
-            if cov_suffix in base_id:
-                base_id = base_id.split(cov_suffix, 1)[0]
-            header = _safe_fasta_id(base_id) + f"_{ref2_label}_tracts"
+            header = _safe_fasta_id(sid)
             out.write(f">{header}\n")
             for i in range(0, len(seq), line_len):
                 out.write(seq[i : i + line_len] + "\n")
@@ -1620,9 +1610,11 @@ def _outbreak_label_from_header(h: str) -> str:
                 return "Ib"
             return "other"
 
-        ref_ids = [hid for hid in aln_dict.keys() if not hid.endswith(f"_{ref1_label}_tracts") and not hid.endswith(f"_{ref2_label}_tracts")]
-        tract1_ids = [hid for hid in aln_dict.keys() if hid.endswith(f"_{ref1_label}_tracts")]
-        tract2_ids = [hid for hid in aln_dict.keys() if hid.endswith(f"_{ref2_label}_tracts")]
+        # Partition alignment IDs into references vs tract sequences for ref1/ref2.
+        tract1_ids = [hid for hid in aln_dict.keys() if f"_{ref1_label}_tract_HC_" in hid]
+        tract2_ids = [hid for hid in aln_dict.keys() if f"_{ref2_label}_tract_HC_" in hid]
+        tract_ids = set(tract1_ids) | set(tract2_ids)
+        ref_ids = [hid for hid in aln_dict.keys() if hid not in tract_ids]
 
         ref_outbreak = {rid: _outbreak_label_from_header(rid) for rid in ref_ids}
 
@@ -1657,8 +1649,10 @@ def _nearest_outbreak(seq_id: str) -> str:
             args._phylogeny_ancestors = {}
 
         def _sample_key_from_tract(hid: str, label: str) -> str:
-            suffix = f"_{label}_tracts"
-            return hid[: -len(suffix)] if hid.endswith(suffix) else hid
+            # Tract IDs contain coverage like "<safe_id>_<label>_tract_HC_<cov>".
+            # Recover the base sample ID by stripping the coverage suffix.
+            cov_suffix = f"_{label}_tract_HC_"
+            return hid.split(cov_suffix, 1)[0] if cov_suffix in hid else hid
 
         for hid in tract1_ids:
             key = _sample_key_from_tract(hid, ref1_label)
@@ -1799,8 +1793,8 @@ def _nearest_outbreak_tree(seq_id: str) -> Optional[str]:
                 return best_label
 
             def _sample_key_from_tract(hid: str, label: str) -> str:
-                suffix = f"_{label}_tracts"
-                return hid[: -len(suffix)] if hid.endswith(suffix) else hid
+                cov_suffix = f"_{label}_tract_HC_"
+                return hid.split(cov_suffix, 1)[0] if cov_suffix in hid else hid
 
             for hid in tract1_ids:
                 key = _sample_key_from_tract(hid, ref1_label)

recmpox/references/root_tree_figure.R

@@ -33,7 +33,47 @@ rooted_path <- file.path(out_dir, "rooted.tree")
 write.tree(tr, rooted_path)
 
 labels <- tr$tip.label
-type <- ifelse(grepl("_tracts", labels, fixed = TRUE), "recombinant", "reference")
+
+# Shorten recombinant tract labels for plotting: keep a 20-char sample ID prefix,
+# the clade label (Ia/Ib/IIa/IIb), and the tract coverage (HC_xx.x) when present.
+shorten_label <- function(x) {
+  if (!grepl("_tract_HC_", x, fixed = TRUE)) {
+    return(x)
+  }
+  clade <- NA_character_
+  split_pat <- NULL
+  if (grepl("_Ia_tract_HC_", x, fixed = TRUE)) {
+    split_pat <- "_Ia_tract_HC_"
+    clade <- "Ia"
+  } else if (grepl("_Ib_tract_HC_", x, fixed = TRUE)) {
+    split_pat <- "_Ib_tract_HC_"
+    clade <- "Ib"
+  } else if (grepl("_IIa_tract_HC_", x, fixed = TRUE)) {
+    split_pat <- "_IIa_tract_HC_"
+    clade <- "IIa"
+  } else if (grepl("_IIb_tract_HC_", x, fixed = TRUE)) {
+    split_pat <- "_IIb_tract_HC_"
+    clade <- "IIb"
+  } else {
+    return(x)
+  }
+  parts <- strsplit(x, split_pat, fixed = TRUE)[[1]]
+  base <- parts[1]
+  cov <- if (length(parts) >= 2) parts[2] else ""
+  base_short <- substr(base, 1, 20)
+  cov_clean <- gsub("[^0-9.]", "", cov)
+  if (nzchar(cov_clean)) {
+    cov_id <- gsub("\\.", "_", cov_clean)
+    return(paste0(base_short, "_", clade, "_tract_HC_", cov_id))
+  } else {
+    return(paste0(base_short, "_", clade, "_tract"))
+  }
+}
+labels <- vapply(labels, shorten_label, character(1))
+# Use shortened labels in the tree so PDF/SVG and any downstream plot show them
+tr$tip.label <- labels
+
+type <- ifelse(grepl("_tract_HC_", labels, fixed = TRUE), "recombinant", "reference")
 clade <- rep("Ia", length(labels))
 clade[type == "recombinant"] <- "Recombinant"
 clade[grepl("sh2024Ib", labels, fixed = TRUE)] <- "sh2024Ib"