Add max_assembly_sequences circuit breaker to iterative_overlap_assembly

When the number of variant sequences after substring collapse exceeds
max_assembly_sequences (default 1000), skip the O(n^2)-per-round greedy
merge and return sequences sorted by read support. This prevents silent
hangs on high-coverage loci.

Three new tests: breaker triggers above threshold, disabled with None,
normal merge below threshold.

Fixes #151

Author: iskandr

isovar/__init__.py

@@ -10,7 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "1.4.12"
+__version__ = "1.4.13"
 
 
 from .allele_read import AlleleRead

isovar/assembly.py

@@ -130,23 +130,34 @@ def collapse_substrings(variant_sequences):
     ]
 
 
+DEFAULT_MAX_ASSEMBLY_SEQUENCES = 1000
+
+
 def iterative_overlap_assembly(
         variant_sequences,
-        min_overlap_size=MIN_VARIANT_SEQUENCE_ASSEMBLY_OVERLAP_SIZE):
+        min_overlap_size=MIN_VARIANT_SEQUENCE_ASSEMBLY_OVERLAP_SIZE,
+        max_assembly_sequences=DEFAULT_MAX_ASSEMBLY_SEQUENCES):
     """
     Assembles longer sequences from reads centered on a variant by
-    between merging all pairs of overlapping sequences and collapsing
+    merging all pairs of overlapping sequences and collapsing
     shorter sequences onto every longer sequence which contains them.
 
+    Parameters
+    ----------
+    variant_sequences : list of VariantSequence
+
+    min_overlap_size : int
+
+    max_assembly_sequences : int or None
+        If the number of input sequences exceeds this threshold after
+        substring collapse, skip the O(n^2)-per-round greedy merge and
+        return sequences sorted by read support. Set to None to disable.
+
     Returns a list of variant sequences, sorted by decreasing read support.
     """
     if len(variant_sequences) <= 1:
-        # if we don't have at least two sequences to start with then
-        # skip the whole mess below
         return variant_sequences
 
-    # reduce the number of inputs to the merge algorithm by first collapsing
-    # shorter sequences onto the longer sequences which contain them
     n_before_collapse = len(variant_sequences)
     variant_sequences = collapse_substrings(variant_sequences)
     n_after_collapse = len(variant_sequences)
@@ -155,7 +166,16 @@ def iterative_overlap_assembly(
         n_before_collapse,
         n_after_collapse)
 
-    merged_variant_sequences = greedy_merge(variant_sequences, min_overlap_size)
+    if (max_assembly_sequences is not None
+            and n_after_collapse > max_assembly_sequences):
+        logger.warning(
+            "Too many variant sequences (%d > %d) for greedy assembly; "
+            "skipping merge and returning sequences sorted by read support",
+            n_after_collapse,
+            max_assembly_sequences)
+    else:
+        variant_sequences = greedy_merge(variant_sequences, min_overlap_size)
+
     return list(sorted(
-        merged_variant_sequences,
+        variant_sequences,
         key=lambda seq: -len(seq.reads)))

tests/test_assembly.py

@@ -422,3 +422,61 @@ def test_iterative_assembly_deterministic_across_input_orders():
         result_sequences = sorted(r.sequence for r in result)
         assert result_sequences == reference_sequences, \
             "Trial %d: different assembly from shuffled input" % trial
+
+
+def test_assembly_circuit_breaker_skips_merge():
+    """
+    When the number of sequences after collapse exceeds
+    max_assembly_sequences, the greedy merge is skipped and sequences
+    are returned sorted by read support.
+    """
+    # Use distinct prefixes so collapse_substrings can't fold them
+    prefixes = ["GAT", "CTA", "AGC", "TCA", "GCG"]
+    variant_sequences = [
+        VariantSequence(
+            prefix=p, alt="X", suffix="TTT",
+            reads={"%s_%d" % (p, k) for k in range(i + 1)})
+        for i, p in enumerate(prefixes)
+    ]
+
+    result = iterative_overlap_assembly(
+        variant_sequences,
+        min_overlap_size=1,
+        max_assembly_sequences=3)
+
+    eq_(len(result), 5, "All 5 sequences should be returned (no merge)")
+    read_counts = [len(r.reads) for r in result]
+    assert read_counts == sorted(read_counts, reverse=True), \
+        "Results should be sorted by decreasing read count, got %s" % read_counts
+
+
+def test_assembly_circuit_breaker_disabled_with_none():
+    """max_assembly_sequences=None disables the circuit breaker."""
+    variant_sequences = [
+        VariantSequence(prefix="A" * 5, alt="CC", suffix="T" * 5, reads={"r1"}),
+        VariantSequence(prefix="A" * 5, alt="CC", suffix="T" * 5 + "G", reads={"r2"}),
+    ]
+    result = iterative_overlap_assembly(
+        variant_sequences,
+        min_overlap_size=1,
+        max_assembly_sequences=None)
+    # Should merge normally despite any threshold
+    eq_(len(result), 1)
+
+
+def test_assembly_circuit_breaker_not_triggered_under_threshold():
+    """
+    When the sequence count is under the threshold, greedy merge proceeds
+    normally (same as no circuit breaker).
+    """
+    variant_sequences = [
+        VariantSequence(prefix="AAAA", alt="CC", suffix="TTTT", reads={"r1"}),
+        VariantSequence(prefix="AA", alt="CC", suffix="TTTTGG", reads={"r2"}),
+    ]
+    result_with_breaker = iterative_overlap_assembly(
+        variant_sequences, min_overlap_size=1, max_assembly_sequences=100)
+    result_without_breaker = iterative_overlap_assembly(
+        variant_sequences, min_overlap_size=1, max_assembly_sequences=None)
+    eq_(len(result_with_breaker), len(result_without_breaker))
+    for r1, r2 in zip(result_with_breaker, result_without_breaker):
+        eq_(r1.sequence, r2.sequence)