ENH: better docking for repetitive elements

Author: davidmeijer

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "versalign"
-version = "2.1.0"
+version = "2.2.0"
 description = "Naive alignment for lists of arbitrary objects"
 readme = "README.md"
 requires-python = ">=3.10"

src/versalign/docking.py

@@ -10,10 +10,7 @@
 from versalign.config import DEFAULT_GAP_REPR
 from versalign.helpers import seq_to_arr, arr_to_seq
 from versalign.blocking import get_symbol_score_lookup
-from versalign.pairwise import pairwise_alignment, pairwise_alignment_score
-
-
-DockStrategy = Literal["nonoverlap"]
+from versalign.pairwise import pairwise_alignment
 
 
 @dataclass(frozen=True)
@@ -44,10 +41,10 @@ class DockingResult:
     """
     Represents the result of docking multiple blocks against a target sequence.
 
-    :var placements: chosen placements, reordered by center coordinate
-    :var unused_blocks: per-block indices that were not used
-    :var docked_row: center-anchored row for visualization
-    :var total_score: total score of all placements
+    :param placements: chosen placements, reordered by center coordinate
+    :param unused_blocks: per-block indices that were not used
+    :param docked_row: center-anchored row for visualization
+    :param total_score: total score of all placements
     """
 
     placements: list[DockPlacement]
@@ -60,8 +57,8 @@ def _alignment_ctx(aligner: Aligner, gap_repr: str) -> tuple[PairwiseAligner, li
     """
     Prepare alignment context for docking.
 
-    :var aligner: Aligner object    
-    :var gap_repr: representation of gaps in sequences
+    :param aligner: Aligner object    
+    :param gap_repr: representation of gaps in sequences
     :return: tuple containing the underlying PairwiseAligner, alphabet list, gap index,
     """
     aligner_obj: PairwiseAligner = aligner.aligner
@@ -75,9 +72,9 @@ def _extract_center_interval(center_aln: list[str], block_aln: list[str], gap_re
     """
     Extract the start and end indices of the block alignment within the center alignment.
 
-    :var center_aln: aligned center sequence segment
-    :var block_aln: aligned block sequence segment
-    :var gap_repr: representation of gaps in sequences
+    :param center_aln: aligned center sequence segment
+    :param block_aln: aligned block sequence segment
+    :param gap_repr: representation of gaps in sequences
     :return: tuple of (start, end) indices in the center sequence, or None if no coverage
     """
     # Map candidate columns to target positions
@@ -106,9 +103,9 @@ def _make_center_anchored_row(
     """
     Create a center-anchored row representing the docked blocks.
 
-    :var center: the target sequence
-    :var placements: list of DockPlacement objects
-    :var gap_repr: representation of gaps in sequences
+    :param center: the target sequence
+    :param placements: list of DockPlacement objects
+    :param gap_repr: representation of gaps in sequences
     :return: list of symbols representing the docked row
     """
     # Simple mapping view: one symbol per center position; no insertion
@@ -173,103 +170,23 @@ def score_docked_region(
     return total
 
 
-def dock_against_target(
-    aligner: Aligner,
-    target: Sequence[str],
-    candidates: Sequence[Sequence[str]],
-    gap_repr: str = DEFAULT_GAP_REPR,
-    allow_block_reverse: bool = False,
-    strategy: DockStrategy = "nonoverlap",
-) -> DockingResult:
-    """
-    Dock multiple blocks against a target sequence.
-
-    :var aligner: Aligner object
-    :var target: target sequence as a list of symbols
-    :var candidates: list of block sequences to dock
-    :var gap_repr: representation of gaps in sequences
-    :var allow_block_reverse: whether to consider reversed blocks
-    :var strategy: docking strategy to use (currently only "nonoverlap" supported)
-    :return: DockingResult object containing placements and summary information
+def _select_nonoverlapping(placements: list[DockPlacement]) -> list[DockPlacement]:
     """
-    if not target:
-        raise ValueError("target sequence must not be empty")
-    if not candidates:
-        return DockingResult(
-            placements=[],
-            unused_blocks=[],
-            docked_row=[gap_repr] * len(target),
-            total_score=0.0,
-        )
+    Weighted interval scheduling on DockPlacement.start/end with weight=score.
     
-    aligner_obj, alphabet, gap_idx, label_fn = _alignment_ctx(aligner, gap_repr)
-    idx, mat = get_symbol_score_lookup(aligner)
-
-    target_int = seq_to_arr(list(target), alphabet, label_fn).astype(np.int32)
-
-    # Generate placements for each block (fwd and optionally rev)
-    placements: list[DockPlacement] = []
-    for bi, blk in enumerate(candidates):
-        blk_fwd = seq_to_arr(list(blk), alphabet, label_fn).astype(np.int32)
-        orientations = [(False, blk_fwd)]
-        if allow_block_reverse:
-            blk_rev = seq_to_arr(list(reversed(blk)), alphabet, label_fn).astype(np.int32)
-            orientations.append((True, blk_rev))
-
-        best_for_block: DockPlacement | None = None
-
-        for is_rev, blk_int in orientations:
-            _, tgt_aln_int, blk_aln_int = pairwise_alignment(aligner_obj, target_int, blk_int, gap_repr=gap_idx)
-            tgt_aln = arr_to_seq(tgt_aln_int, alphabet)
-            blk_aln = arr_to_seq(blk_aln_int, alphabet)
-
-            interval = _extract_center_interval(tgt_aln, blk_aln, gap_repr)
-            if interval is None:
-                continue
-            start, end = interval
-
-            region_score = score_docked_region(
-                center_aln=tgt_aln,
-                block_aln=blk_aln,
-                start=start,
-                end=end,
-                gap_repr=gap_repr,
-                idx=idx,
-                mat=mat,
-            )
-
-            cand = DockPlacement(
-                block_idx=bi,
-                reversed=is_rev,
-                start=start,
-                end=end,
-                score=float(region_score),
-                center_aln=tgt_aln,
-                block_aln=blk_aln,
-            )
-            if best_for_block is None or cand.score > best_for_block.score:
-                best_for_block = cand
-
-        if best_for_block is not None and best_for_block.score > 0.0:
-            placements.append(best_for_block)
-
+    :param placements: list of DockPlacement objects
+    :return: list of chosen non-overlapping DockPlacement objects with maximum total score
+    """
     if not placements:
-        return DockingResult(
-            placements=[],
-            unused_blocks=list(range(len(candidates))),
-            docked_row=[gap_repr] * len(target),
-            total_score=0.0,
-        )
-    
-    # Weighted interval scheduling (non-overlapping)
-    placements.sort(key=lambda p: (p.end, p.start))
-    ends = [p.end for p in placements]
+        return []
+
+    placements = sorted(placements, key=lambda p: (p.end, p.start))
 
     def prev_nonoverlap(i: int) -> int:
         """
         Find the index of the last placement that does not overlap with placement i.
         
-        :var i: index of the current placement
+        :param i: index of the current placement
         :return: index of the last non-overlapping placement, or -1 if none exists
         """
         lo, hi = 0, i - 1
@@ -283,7 +200,6 @@ def prev_nonoverlap(i: int) -> int:
                 hi = mid - 1
         return ans
 
-    # Dynamic programming to find optimal set of non-overlapping placements
     n = len(placements)
     pidx = [prev_nonoverlap(i) for i in range(n)]
     dp = [0.0] * n
@@ -298,7 +214,7 @@ def prev_nonoverlap(i: int) -> int:
             take[i] = True
         else:
             dp[i] = excl
-    
+
     chosen: list[DockPlacement] = []
     i = n - 1
     while i >= 0:
@@ -308,19 +224,132 @@ def prev_nonoverlap(i: int) -> int:
         else:
             i -= 1
     chosen.reverse()
+    return chosen
 
-    used_blocks = {p.block_idx for p in chosen}
-    unused_blocks = [i for i in range(len(candidates)) if i not in used_blocks]
 
-    docker_row = _make_center_anchored_row(list(target), chosen, gap_repr)
-    total_score = float(sum(p.score for p in chosen))
+def dock_against_target(
+    aligner: Aligner,
+    target: Sequence[str],
+    candidates: Sequence[Sequence[str]],
+    gap_repr: str,
+    mask_repr: str,
+    allow_block_reverse: bool = False,
+    max_passes: int = 3,
+) -> DockingResult:
+    """
+    Dock multiple blocks against a target sequence.
+
+    :param aligner: Aligner object
+    :param target: target sequence as a list of symbols
+    :param candidates: list of block sequences to dock
+    :param gap_repr: representation of gaps in sequences
+    :param allow_block_reverse: whether to consider reversed blocks
+    :param strategy: docking strategy to use (currently only "nonoverlap" supported)
+    :param max_passes: maximum number of passes for iterative docking
+    :return: DockingResult object containing placements and summary information
+    """
+    if not target:
+        raise ValueError("target sequence must not be empty")
+    if not candidates:
+        return DockingResult(
+            placements=[],
+            unused_blocks=[],
+            docked_row=[gap_repr] * len(target),
+            total_score=0.0,
+        )
+    
+    aligner_obj, alphabet, gap_idx, label_fn = _alignment_ctx(aligner, gap_repr)
+    idx, mat = get_symbol_score_lookup(aligner)
+
+    target_int = seq_to_arr(list(target), alphabet, label_fn).astype(np.int32)
+    mask_idx = alphabet.index(mask_repr)
+
+    chosen_all: list[DockPlacement] = []
+    occupied = np.zeros(len(target), dtype=bool)
+    
+    remaining = list(range(len(candidates)))
+
+    for _pass in range(max_passes):
+        if not remaining:
+            break
+
+        # Build masked target for this pass
+        masked_target_int = target_int.copy()
+        masked_target_int[occupied] = mask_idx
+
+        # Propose onse best placement per remaining block
+        proposed: list[DockPlacement] = []
+        
+        for bi in remaining:
+            blk = candidates[bi]
+            blk_fwd = seq_to_arr(list(blk), alphabet, label_fn).astype(np.int32)
+            orientations = [(False, blk_fwd)]
+            if allow_block_reverse:
+                blk_rev = seq_to_arr(list(reversed(blk)), alphabet, label_fn).astype(np.int32)
+                orientations.append((True, blk_rev))
+
+            best_for_block: DockPlacement | None = None
+
+            for is_rev, blk_int in orientations:
+                _, tgt_aln_int, blk_aln_int = pairwise_alignment(aligner_obj, masked_target_int, blk_int, gap_repr=gap_idx)
+                tgt_aln = arr_to_seq(tgt_aln_int, alphabet)
+                blk_aln = arr_to_seq(blk_aln_int, alphabet)
+
+                interval = _extract_center_interval(tgt_aln, blk_aln, gap_repr)
+                if interval is None:
+                    continue
+                start, end = interval
+
+                region_score = score_docked_region(
+                    center_aln=tgt_aln,
+                    block_aln=blk_aln,
+                    start=start,
+                    end=end,
+                    gap_repr=gap_repr,
+                    idx=idx,
+                    mat=mat,
+                )
+
+                cand = DockPlacement(
+                    block_idx=bi,
+                    reversed=is_rev,
+                    start=start,
+                    end=end,
+                    score=float(region_score),
+                    center_aln=tgt_aln,
+                    block_aln=blk_aln,
+                )
+                if best_for_block is None or cand.score > best_for_block.score:
+                    best_for_block = cand
+
+            if best_for_block is not None and best_for_block.score > 0.0:
+                proposed.append(best_for_block)
+
+        # Choose a non-overlappin subset among the newly proposed placements
+        chosen_this_pass = _select_nonoverlapping(proposed)
+
+        if not chosen_this_pass:
+            break  # no progress => stop
+
+        # Accept them and update occupied mask
+        chosen_all.extend(chosen_this_pass)
+        for p in chosen_this_pass:
+            occupied[p.start : p.end + 1] = True
+
+        used_now = {p.block_idx for p in chosen_all}
+        remaining = [i for i in range(len(candidates)) if i not in used_now]
+
+    # Finalize
+    chosen_sorted = sorted(chosen_all, key=lambda p: (p.start, p.end))
+    used_blocks = {p.block_idx for p in chosen_sorted}
+    unused_blocks = [i for i in range(len(candidates)) if i not in used_blocks]
 
-    # Reorder by target coordinate for readability
-    chosen_sorted = sorted(chosen, key=lambda p: (p.start, p.end))
+    docked_row = _make_center_anchored_row(list(target), chosen_sorted, gap_repr)
+    total_score = float(sum(p.score for p in chosen_sorted))
 
     return DockingResult(
         placements=chosen_sorted,
         unused_blocks=unused_blocks,
-        docked_row=docker_row,
+        docked_row=docked_row,
         total_score=total_score,
     )