Preserve chromosome name in separate Scaffold attribute

Author: jgrg

src/tola/assembly/assembly_stats.py

@@ -151,25 +151,18 @@ def get_assembly_scaffold_lengths(
         return scaff_lengths
 
     def chromosome_name_csv(self, haplotype: str, asm: Assembly):
-        prefix = self.autosome_prefix
-        suffix = f"_{haplotype}" if haplotype else None
-        last_orig = None
-        chr_name = None
-
         csv_str = io.StringIO()
         for scffld in asm.scaffolds:
             if scffld.rank in (1, 2):
-                name = scffld.name
-                orig = scffld.original_name
-                if last_orig and orig == last_orig:
-                    localised = "no"
-                else:
-                    localised = "yes"
-                    chr_name = name.replace(prefix, "", 1)
-                    if suffix:
-                        chr_name = chr_name.replace(suffix, "", 1)
-                    last_orig = orig
-                csv_str.write(",".join((name, chr_name, localised)))
+                csv_str.write(
+                    ",".join(
+                        (
+                            scffld.name,
+                            scffld.chr_name,
+                            "yes" if scffld.localised else "no",
+                        )
+                    )
+                )
                 csv_str.write("\n")
 
         return csv_str.getvalue() if csv_str.tell() else None
@@ -192,30 +185,15 @@ def chromosomes_report_csv(self, hap_asm: AssemblyDict):
         )
         head_pos = csv_str.tell()
 
-        prefix = self.autosome_prefix
-        chr_name = None
-        last_orig = None
         for hap, asm in hap_asm.items():
-            suffix = f"_{hap}" if hap else None
             for scffld in asm.scaffolds:
                 if scffld.rank in (1, 2):
-                    name = scffld.name
-                    orig = scffld.original_name
-                    if last_orig and orig == last_orig:
-                        # Unlocs share the same original_name
-                        localised = "false"
-                    else:
-                        localised = "true"
-                        last_orig = orig
-                        chr_name = name.replace(prefix, "", 1)
-                        if suffix:
-                            chr_name = chr_name.replace(suffix, "", 1)
                     csvr.writerow(
                         (
                             scffld.haplotype or hap or "Primary",
-                            name,
-                            chr_name,
-                            localised,
+                            scffld.name,
+                            scffld.chr_name,
+                            "true" if scffld.localised else "false",
                             scffld.original_name,
                             scffld.length,
                             scffld.fragments_length,

src/tola/assembly/build_assembly.py

@@ -333,10 +333,15 @@ def cut_scaffold_if_too_long(self, scffld: Scaffold) -> list[Scaffold]:
             cut_at = whole // div
             gap_i = self.index_of_nearest_gap_to_ideal_cut_site(to_cut, cut_at)
             rows = to_cut.rows
+
             # First part is everything up to, but not including, the gap
-            cut_parts.append(Scaffold(to_cut.name, rows[:gap_i]))
+            cut = to_cut.clone_empty()
+            cut.rows = rows[:gap_i]
+            cut_parts.append(cut)
+
             # Second part is everything after the gap
-            to_cut = Scaffold(to_cut.name, rows[gap_i + 1 :])
+            to_cut = to_cut.clone_empty()
+            to_cut.rows = rows[gap_i + 1 :]
         cut_parts.append(to_cut)
 
         # Add suffix "_1", "_2" etc... to cut scaffolds
@@ -393,6 +398,7 @@ def index_of_nearest_gap_to_ideal_cut_site(self, to_cut: Scaffold, cut_at: int):
                     break
 
             if gap_i_before is None and gap_i_after is None:
+                ### Make a cut here ###
                 msg = (
                     f"Failed to find gap before or after {cut_at:_d} in '{to_cut.name}'"
                 )
@@ -403,47 +409,35 @@ def index_of_nearest_gap_to_ideal_cut_site(self, to_cut: Scaffold, cut_at: int):
             elif gap_i_after is None:
                 ovr_i = gap_i_before
             else:
-                length_before = 0
-                length_after = 0
-                for i, this_row in enumerate(rows):
-                    if i < gap_i_before:
-                        length_before += this_row.length
-
-                    if i < gap_i_after:
-                        length_after += this_row.length
-                    else:
-                        break
+                length_before = (
+                    idx_asm.start_end_of_row(to_cut.name, gap_i_before)[0] - 1
+                )
+                length_after = idx_asm.start_end_of_row(to_cut.name, gap_i_after)[0] - 1
 
                 # Choose the gap before or after, whichever is nearest to the
                 # ideal cut point.
                 ovr_i = (
                     gap_i_before
-                    if abs(cut_at - length_before) < abs(cut_at - length_after)
+                    if abs(cut_at - length_before) < abs(length_after - cut_at)
                     else gap_i_after
                 )
 
         return ovr_i
 
     def scaffolds_fused_by_name(self) -> list[Scaffold]:
         gap = self.default_gap
-        hap_name_scaffold: dict[tuple[str, str], Scaffold] = {}
+        hap_name_scaffold: dict[tuple[str | None, str], Scaffold] = {}
         for scffld in self.scaffolds:
             if not scffld.rows:
                 # discard_overhanging_fragments() may have removed the only
                 # row from an OverlapResult
                 continue
 
-            build_scffld = hap_name_scaffold.setdefault(
-                (scffld.haplotype, scffld.name),
-                Scaffold(
-                    scffld.name,
-                    tag=scffld.tag,
-                    haplotype=scffld.haplotype,
-                    rank=scffld.rank,
-                    original_name=scffld.original_name,
-                    original_tags=scffld.original_tags,
-                ),
-            )
+            idx = scffld.haplotype, scffld.name
+            build_scffld = hap_name_scaffold.get(idx)
+            if not build_scffld:
+                hap_name_scaffold[idx] = build_scffld = scffld.clone_empty()
+
             if isinstance(scffld, OverlapResult):
                 build_scffld.append_scaffold(scffld.to_scaffold(), gap)
             else:

src/tola/assembly/indexed_assembly.py

@@ -51,6 +51,11 @@ def overlap_index_by_name(self, name: str) -> list[int]:
         msg = f"Scaffold '{name}' is not indexed."
         raise ValueError(msg)
 
+    def start_end_of_row(self, name: str, index: int) -> tuple[int, int]:
+        idx = self.overlap_index_by_name(name)
+        start = 1 if index == 0 else 1 + idx[index - 1]
+        return start, idx[index]
+
     def find_overlaps(self, bait: Fragment) -> OverlapResult | None:
         """
         Given a Fragment bait, returns an OverlapResult (a subclass of

src/tola/assembly/naming_utils.py

@@ -150,8 +150,12 @@ def label_scaffold(
             if "Painted" not in scaffold_tags:
                 msg = f"Unloc in unpainted scaffold {original_name!r}: {fragment}"
                 raise TaggingError(msg)
+            scaffold.chr_name = name
             name = self.unloc_name()
             self.unloc_scaffolds[-1].append(scaffold)
+        elif "Painted" in scaffold_tags:
+            scaffold.localised = True
+            scaffold.chr_name = name
 
         scaffold.name = name
         scaffold.haplotype = self.current_haplotype
@@ -231,17 +235,17 @@ def length_of_first_haplotype(self):
         return length
 
     @staticmethod
-    def multi_chr_list(chr_name, multi_count, hap_suffix):
+    def multi_chr_list(chr_name, multi_count):
         """
         Adds the suffix "A", "B", "C" etc... to the supplied chromosome name
         for when there are multiple chromosomes in a group.
         """
         if multi_count == 1:
-            return [chr_name + hap_suffix]
+            return [chr_name]
         else:
             chr_list = []
             for ltr in range(ord("A"), ord("A") + multi_count):
-                chr_list.append(chr_name + chr(ltr) + hap_suffix)
+                chr_list.append(chr_name + chr(ltr))
             return chr_list
 
     def max_hap_set_count(self):
@@ -262,13 +266,15 @@ def name_chromosome(self, chr_prefix, chr_n):
             hap_suffix = (
                 "" if (hap_name is None or hap_name == "Primary") else f"_{hap_name}"
             )
-            chr_names = self.multi_chr_list(
-                chr_prefix + str(chr_n), len(hap_set), hap_suffix
-            )
+            chr_names = self.multi_chr_list(str(chr_n), len(hap_set))
             for orig, scffld_list in hap_set.items():
                 this_chr = chr_names.pop(0)
                 for scffld in scffld_list:
-                    scffld.name = scffld.name.replace(orig, this_chr)
+                    scffld.name = scffld.name.replace(
+                        orig,
+                        chr_prefix + this_chr + hap_suffix,
+                    )
+                    scffld.chr_name = this_chr
 
 
 class ChrNamerError(Exception):
@@ -293,7 +299,7 @@ def __init__(self, chr_prefix="SUPER_"):
         self.chr_prefix = chr_prefix
         self.scaffolds = []
         self.haplotypes_seen = {}
-        self.groups = None
+        self.groups: list[ChrGroup] | None = None
 
     def add_scaffold(self, haplotype, scffld):
         # A dict is used to store the haplotypes seen since order is

src/tola/assembly/overlap_result.py

@@ -9,18 +9,19 @@ class OverlapResult(Scaffold):
     def __init__(
         self,
         bait,
-        rows,
         start,
         end,
+        rows,
         name=None,
-        tag=None,
-        haplotype=None,
-        rank=None,
-        original_name=None,
+        **args,
     ):
         if not name:
             name = f"matches to {bait.name} {bait.start} to {bait.end}"
-        super().__init__(name, rows, tag, haplotype, rank, original_name)
+        super().__init__(
+            name,
+            rows,
+            **args,
+        )
         self.bait = bait
         self.start = start
         self.end = end
@@ -174,12 +175,8 @@ def append_scaffold(self):
         raise NotImplementedError
 
     def to_scaffold(self) -> Scaffold:
-        scffld = Scaffold(
-            self.name,
-            self.rows,
-            original_name=self.original_name,
-            original_tags=self.original_tags,
-        )
+        scffld = self.clone_empty()
+        scffld.rows = self.rows
         if self.bait.strand == -1:
             return scffld.reverse()
         else:

src/tola/assembly/scaffold.py

@@ -16,6 +16,8 @@ def __init__(
         rank=0,
         original_name=None,
         original_tags: set[str] | None = None,
+        localised: bool = False,
+        chr_name: str | None = None,
     ):
         self.name = str(name)
         if rows:
@@ -27,6 +29,24 @@ def __init__(
         self.rank: int = rank
         self.original_name: str | None = original_name
         self.original_tags: set[str] | None = original_tags
+        self.localised = localised
+        self.chr_name = chr_name
+
+    def clone_empty(self) -> 'Scaffold':
+        """
+        Clones a `Scaffold` or derived class, returning a base `Scaffold` with
+        all of its attributes apart from `rows`.
+        """
+        return Scaffold(
+            self.name,
+            tag=self.tag,
+            haplotype=self.haplotype,
+            rank=self.rank,
+            original_name=self.original_name,
+            original_tags=self.original_tags,
+            localised=self.localised,
+            chr_name=self.chr_name,
+        )
 
     def __repr__(self):
         txt = io.StringIO()

tests/build_assembly_test.py

@@ -14,12 +14,8 @@
 
 
 def test_multi_chr_list():
-    assert ChrGroup.multi_chr_list("SUPER_3", 1, "_HAP2") == ["SUPER_3_HAP2"]
-    assert ChrGroup.multi_chr_list("SUPER_3", 3, "_HAP1") == [
-        "SUPER_3A_HAP1",
-        "SUPER_3B_HAP1",
-        "SUPER_3C_HAP1",
-    ]
+    assert ChrGroup.multi_chr_list("3", 1) == ["3"]
+    assert ChrGroup.multi_chr_list("3", 3) == ["3A", "3B", "3C"]
 
 
 def list_chr_naming_tests():