Breaking of long scaffolds, 2 Gbp by default

Author: jgrg

src/tola/assembly/build_assembly.py

@@ -31,6 +31,7 @@ def __init__(
         default_gap=None,
         bp_per_texel=None,
         autosome_prefix=None,
+        max_contig_length=2_000_000_000,
     ):
         super().__init__(name, header, scaffolds, bp_per_texel)
         self.default_gap = default_gap
@@ -40,6 +41,7 @@ def __init__(
         self.assembly_stats: AssemblyStats = AssemblyStats()
         if autosome_prefix:
             self.autosome_prefix = autosome_prefix
+        self.max_contig_length = max_contig_length
 
     @property
     def autosome_prefix(self):
@@ -288,17 +290,21 @@ def __build_name_and_sort_assemblies(
                 asm_key = hap
             else:
                 asm_key = None
-            new_asm = assemblies.setdefault(
-                asm_key, Assembly(self.name, curated=curated)
-            )
-            new_asm.add_scaffold(scffld)
-            if scffld.rank == 1:
-                # Add autosome to the ChrNamer
-                chr_namer.add_scaffold(asm_key, scffld)
-            elif scffld.rank == 2:
-                chr_namer.add_chr_prefix(scffld, hap)
-            elif hap is not None:
-                chr_namer.add_haplotype_prefix(scffld, hap)
+
+            if not (new_asm := assemblies.get(asm_key)):
+                new_asm = Assembly(self.name, curated=curated)
+                assemblies[asm_key] = new_asm
+
+            for cut_scffld in self.cut_scaffold_if_too_long(scffld):
+                new_asm.add_scaffold(cut_scffld)
+
+                if cut_scffld.rank == 1:
+                    # Add autosome to the ChrNamer
+                    chr_namer.add_scaffold(asm_key, cut_scffld)
+                elif cut_scffld.rank == 2:
+                    chr_namer.add_chr_prefix(cut_scffld, hap)
+                elif hap is not None:
+                    chr_namer.add_haplotype_prefix(cut_scffld, hap)
 
         # ChrNamer names autosome chromosomes by size
         chr_namer.name_chromosomes()
@@ -311,6 +317,113 @@ def __build_name_and_sort_assemblies(
 
         return scaffolds, assemblies
 
+    def cut_scaffold_if_too_long(self, scffld: Scaffold) -> list[Scaffold]:
+        whole = scffld.length
+        pieces = math.ceil(whole / self.max_contig_length)
+
+        if pieces == 1:
+            return [scffld]
+
+        # If, for example, we need to cut the scaffold into 3 pieces, this
+        # loop will take the first 1/3 off the scaffold, then 1/2 of what's
+        # remaining.
+        cut_parts = []
+        to_cut = scffld
+        for div in range(pieces, 1, -1):
+            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]))
+            # Second part is everything after the gap
+            to_cut = Scaffold(to_cut.name, rows[gap_i + 1 :])
+        cut_parts.append(to_cut)
+
+        # Add suffix "_1", "_2" etc... to cut scaffolds
+        for i, part in enumerate(cut_parts):
+            part.name = f"{part.name}_{i + 1}"
+
+        # Format report of cuts made
+        whole_str = f"{whole:,d}"
+        wl = len(whole_str)
+        nl = len(scffld.name) + 4
+        log.info(
+            f"Cut {scffld.name:<{nl}}  {whole:{wl},d} bp (including gaps) into:\n"
+            + "".join(
+                [f"    {x.name:<{nl}}  {x.length:{wl},d} bp\n" for x in cut_parts]
+            )
+        )
+
+        return cut_parts
+
+    def index_of_nearest_gap_to_ideal_cut_site(self, to_cut: Scaffold, cut_at: int):
+        # Make a temporary `IndexedAssembly` to use it's code to search for an
+        # row which overlaps out cut coordiante.
+        idx_asm = IndexedAssembly(
+            f"Temporary Assembly for cutting '{to_cut.name}' at {cut_at:_d}",
+            scaffolds=[to_cut],
+        )
+
+        # Find the row which overlaps the ideal cut site
+        ovr_i_j = idx_asm.overlapping_indices_by_scaffold_start_end(
+            to_cut, cut_at, cut_at
+        )
+        if not ovr_i_j:
+            msg = (
+                f"Failed to find an element at {cut_at:_d}"
+                f" within '{to_cut.name}' of length {to_cut.length:_d}"
+            )
+            raise ValueError(msg)
+
+        ovr_i = ovr_i_j[0]
+        rows = to_cut.rows
+        ele = rows[ovr_i]
+        if not isinstance(ele, Gap):
+            # This isn't a gap, so we need to find the nearest
+            gap_i_before = None
+            for i in range(ovr_i, 0, -1):
+                if isinstance(rows[i], Gap):
+                    gap_i_before = i
+                    break
+
+            gap_i_after = None
+            for i in range(ovr_i, len(rows)):
+                if isinstance(rows[i], Gap):
+                    gap_i_after = i
+                    break
+
+            if gap_i_before is None and gap_i_after is None:
+                msg = (
+                    f"Failed to find gap before or after {cut_at:_d} in '{to_cut.name}'"
+                )
+                raise ValueError(msg)
+
+            if gap_i_before is None:
+                ovr_i = gap_i_after
+            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
+
+                # 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)
+                    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] = {}

src/tola/assembly/indexed_assembly.py

@@ -64,9 +64,10 @@ def find_overlaps(self, bait: Fragment) -> OverlapResult | None:
             raise ValueError(msg)
 
         idx = self.overlap_index_by_name(scffld.name)
-        i_ovr, j_ovr = self.get_overlap_fragment_indices(idx, bait.start, bait.end)
-        if i_ovr is None or j_ovr is None:
+        rng = self.__get_overlap_fragment_indices(idx, bait.start, bait.end)
+        if rng is None:
             return None
+        i_ovr, j_ovr = rng
 
         # Walk start and end pointers back to ignore Gaps on the ends
         while isinstance(scffld.rows[i_ovr], Gap):
@@ -87,9 +88,9 @@ def find_overlaps(self, bait: Fragment) -> OverlapResult | None:
             rows=overlaps,
         )
 
-    def get_overlap_fragment_indices(
+    def __get_overlap_fragment_indices(
         self, idx: list[int], start: int, end: int
-    ) -> tuple[int | None, int | None]:
+    ) -> tuple[int, int] | None:
         """
         Perfoms a binary search through the index `idx` (a list of
         `Fragment | Gap` end positions), returning the start and end indices
@@ -116,7 +117,7 @@ def get_overlap_fragment_indices(
                 break
 
         if ovr is None:
-            return None, None
+            return None
 
         # The span of overlapping entries may extend to the left or right
         # of "ovr"
@@ -132,3 +133,12 @@ def get_overlap_fragment_indices(
             j_ovr = j
 
         return i_ovr, j_ovr
+
+    def overlapping_indices_by_scaffold_start_end(
+        self, scaffold: Scaffold, start: int, end: int
+    ) -> tuple[int, int] | None:
+        """
+        Get the start and end index
+        """
+        idx = self.overlap_index_by_name(scaffold.name)
+        return self.__get_overlap_fragment_indices(idx, start, end)

src/tola/assembly/naming_utils.py

@@ -396,6 +396,7 @@ def build_groups(self):
             #           },
             #       }
             #   )
+
             group.add_scaffold_to_haplotype(haplotype, scffld)
             last_haplotype = haplotype
             last_orig = orig

src/tola/assembly/scaffold.py

@@ -10,7 +10,7 @@ class Scaffold:
     def __init__(
         self,
         name,
-        rows=None,
+        rows: list[Fragment | Gap] | None = None,
         tag=None,
         haplotype=None,
         rank=0,
@@ -21,7 +21,7 @@ def __init__(
         if rows:
             self.rows = [*rows]
         else:
-            self.rows = []
+            self.rows: list[Fragment | Gap] = []
         self.tag: str | None = tag
         self.haplotype: str | None = haplotype
         self.rank: int = rank
@@ -57,7 +57,7 @@ def __str__(self):
             txt.write(f"  {row.length:14_d}  {row}\n")
         return txt.getvalue()
 
-    def add_row(self, row):
+    def add_row(self, row: Fragment | Gap):
         self.rows.append(row)
 
     @property

src/tola/assembly/scripts/pretext_to_asm.py

@@ -3,6 +3,7 @@
 import re
 import sys
 from pathlib import Path
+from typing import no_type_check
 
 import click
 import yaml
@@ -187,6 +188,18 @@ def ul(txt):
         to 'map-order' if the '--keep-map-order' option is given.
         """,
 )
+@click.option(
+    "--max-contig-length",
+    "max_contig_length",
+    type=int,
+    help=f"""
+        Maximum length for a single contig.  If a contig exceeds this size, it
+        will be broken at gaps into a number of pieces such that each is less
+        than this number.  The pieces will have the suffixes '_1', '_2'
+        {it("etc...")} added to their names.  [default: 2 Gbp]
+        """,
+    default=2_000_000_000,
+)
 def cli(
     assembly_file,
     pretext_file,
@@ -197,6 +210,7 @@ def cli(
     write_log,
     keep_map_order,
     default_asm_name,
+    max_contig_length,
 ):
     logfile = setup_logging(log_level, output_file, write_log, clobber)
 
@@ -218,6 +232,7 @@ def cli(
         "stdout",
         default_gap=Gap(200, "scaffold"),
         autosome_prefix=autosome_prefix,
+        max_contig_length=max_contig_length,
     )
     build_asm.remap_to_input_assembly(prtxt_asm, input_asm)
 
@@ -306,9 +321,9 @@ def name_assemblies(
     # has been curated. One of the painted chromosomes in the curated
     # haplotype has been tagged with 'Primary'
     if asm_dict.get("Primary"):
-        # <ToLID>.1.primary.curated.fa  <- Sequence from "Hap1" tagged scaffolds)
+        # <ToLID>.1.primary.curated.fa  <- Sequence from "Hap1" tagged scaffolds
         # <ToLID>.1.primary.chromosome.list.csv
-        # <ToLID>.1.all_haplotigs.curated.fa  <- Sequence from "Hap2" tagged scaffolds)
+        # <ToLID>.1.all_haplotigs.curated.fa  <- Sequence from "Hap2" tagged scaffolds
         other_asm = []
         for asm_key, asm in asm_dict.items():
             if asm_key == "Primary":
@@ -332,8 +347,8 @@ def name_assemblies(
     elif asm_dict.get(None):
         # <ToLID>.1.primary.curated.fa
         # <ToLID>.1.primary.chromosome.list.csv
-        # <ToLID>.1.additional_haplotigs.curated.fa <- Sequence from "Haplotig"
-        #                                              tagged scaffolds
+        # <ToLID>.1.additional_haplotigs.curated.fa  <- Sequence from "Haplotig"
+        #                                               tagged scaffolds
         other_asm = []
         for asm_key, asm in asm_dict.items():
             if asm_key is None:
@@ -483,7 +498,12 @@ def write_chr_csv_files(
                 csv_fh.write(chr_names)
 
 
-def write_info_yaml(output_file, stats: AssemblyStats, out_assemblies: AssemblyDict, clobber):
+def write_info_yaml(
+    output_file,
+    stats: AssemblyStats,
+    out_assemblies: AssemblyDict,
+    clobber,
+):
     asm_stats = stats.per_assembly_stats
     info = {"assemblies": asm_stats}
     if len(asm_stats) > 1: