diff --git a/src/plexus/blast/blast_runner.py b/src/plexus/blast/blast_runner.py index ccbba9b..7c07bb1 100644 --- a/src/plexus/blast/blast_runner.py +++ b/src/plexus/blast/blast_runner.py @@ -24,6 +24,23 @@ def _check_blast_tools() -> None: # https://github.com/JasonAHendry/multiply/blob/master/src/multiply/blast/runner.py +_BLAST_DTYPES = { + "pident": "float32", + "length": "int32", + "mismatch": "int32", + "gapopen": "int32", + "qstart": "int32", + "qend": "int32", + "sstart": "int32", # human genome coords max ~250M, fits int32 (max 2.1B) + "send": "int32", + "evalue": "float32", + "bitscore": "float32", + "qlen": "int16", +} + +_BLAST_CATEGORICAL_COLS = ("qseqid", "sseqid", "sstrand") + + class BlastRunner: BLAST_COLS = "qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore sstrand qlen" @@ -205,8 +222,13 @@ def _load_as_dataframe(self): # Load as a dataframe self.blast_df = pd.read_csv( - self.output_table, sep="\t", names=self.BLAST_COLS.split(" ") + self.output_table, + sep="\t", + names=self.BLAST_COLS.split(" "), + dtype=_BLAST_DTYPES, ) + for col in _BLAST_CATEGORICAL_COLS: + self.blast_df[col] = self.blast_df[col].astype("category") def get_dataframe(self): """ diff --git a/src/plexus/blast/offtarget_finder.py b/src/plexus/blast/offtarget_finder.py index 947ed31..204115d 100644 --- a/src/plexus/blast/offtarget_finder.py +++ b/src/plexus/blast/offtarget_finder.py @@ -146,33 +146,49 @@ def find_amplicons(self, max_size_bp=6000): "R_target": r_targets, "F_primer": np.full(n_matches, f_qseqid), "R_primer": matched_r_qseqids, - "F_start": np.full(n_matches, f_start, dtype=int), - "R_start": matched_r_starts, - "product_bp": matched_r_starts - f_start + 1, + "F_start": np.full(n_matches, f_start, dtype=np.int32), + "R_start": matched_r_starts.astype(np.int32), + "product_bp": (matched_r_starts - f_start + 1).astype(np.int32), "F_pident": np.full( - n_matches, f_pident[i] if f_pident is not None else None + n_matches, + f_pident[i] if f_pident is not None else np.nan, + dtype=np.float32, + ), + "R_pident": ( + r_pident[lo:hi].astype(np.float32) + if r_pident is not None + else np.full(n_matches, np.nan, dtype=np.float32) ), - "R_pident": r_pident[lo:hi] - if r_pident is not None - else np.full(n_matches, None), "F_mismatch": np.full( - n_matches, f_mismatch[i] if f_mismatch is not None else None + n_matches, + f_mismatch[i] if f_mismatch is not None else -1, + dtype=np.int32, + ), + "R_mismatch": ( + r_mismatch[lo:hi].astype(np.int32) + if r_mismatch is not None + else np.full(n_matches, -1, dtype=np.int32) ), - "R_mismatch": r_mismatch[lo:hi] - if r_mismatch is not None - else np.full(n_matches, None), "F_align_len": np.full( - n_matches, f_length[i] if f_length is not None else None + n_matches, + f_length[i] if f_length is not None else -1, + dtype=np.int32, + ), + "R_align_len": ( + r_length[lo:hi].astype(np.int32) + if r_length is not None + else np.full(n_matches, -1, dtype=np.int32) ), - "R_align_len": r_length[lo:hi] - if r_length is not None - else np.full(n_matches, None), "F_evalue": np.full( - n_matches, f_evalue[i] if f_evalue is not None else None + n_matches, + f_evalue[i] if f_evalue is not None else np.nan, + dtype=np.float32, + ), + "R_evalue": ( + r_evalue[lo:hi].astype(np.float32) + if r_evalue is not None + else np.full(n_matches, np.nan, dtype=np.float32) ), - "R_evalue": r_evalue[lo:hi] - if r_evalue is not None - else np.full(n_matches, None), } result_chunks.append(pd.DataFrame(chunk, columns=amplicon_columns)) @@ -180,5 +196,137 @@ def find_amplicons(self, max_size_bp=6000): if result_chunks: self.amplicon_df = pd.concat(result_chunks, ignore_index=True) del result_chunks + for col in ("chrom", "F_target", "R_target", "F_primer", "R_primer"): + self.amplicon_df[col] = self.amplicon_df[col].astype("category") else: self.amplicon_df = pd.DataFrame(columns=amplicon_columns) + + def find_amplicons_by_chrom(self, max_size_bp=6000): + """Yield (chrom, amplicon_df) tuples, one per chromosome. + + Same logic as find_amplicons() but yields per-chromosome results + instead of accumulating into self.amplicon_df. This bounds peak + memory to the largest single chromosome's amplicon set. + """ + amplicon_columns = [ + "chrom", + "F_target", + "R_target", + "F_primer", + "R_primer", + "F_start", + "R_start", + "product_bp", + "F_pident", + "R_pident", + "F_mismatch", + "R_mismatch", + "F_align_len", + "R_align_len", + "F_evalue", + "R_evalue", + ] + target_map = self.target_map + + for chrom, chrom_df in self.bound_df.groupby("sseqid"): + fwd = chrom_df[chrom_df["sstrand"] == "plus"] + rev = chrom_df[chrom_df["sstrand"] == "minus"] + + if fwd.empty or rev.empty: + continue + + # Sort reverse hits by sstart for searchsorted + rev = rev.sort_values("sstart") + r_starts = rev["sstart"].values + r_qseqids = rev["qseqid"].values + r_pident = self._get_col_or_none(rev, "pident") + r_mismatch = self._get_col_or_none(rev, "mismatch") + r_length = self._get_col_or_none(rev, "length") + r_evalue = self._get_col_or_none(rev, "evalue") + + f_starts = fwd["sstart"].values + f_qseqids = fwd["qseqid"].values + f_pident = self._get_col_or_none(fwd, "pident") + f_mismatch = self._get_col_or_none(fwd, "mismatch") + f_length = self._get_col_or_none(fwd, "length") + f_evalue = self._get_col_or_none(fwd, "evalue") + + result_chunks = [] + + for i in range(len(f_starts)): + f_start = f_starts[i] + lo = np.searchsorted(r_starts, f_start, side="right") + hi = np.searchsorted(r_starts, f_start + max_size_bp, side="left") + + if lo >= hi: + continue + + n_matches = hi - lo + f_qseqid = f_qseqids[i] + f_target = target_map.get(f_qseqid, f_qseqid.split("_")[0]) + + matched_r_starts = r_starts[lo:hi] + matched_r_qseqids = r_qseqids[lo:hi] + + r_targets = [ + target_map.get(rq, rq.split("_")[0]) for rq in matched_r_qseqids + ] + + chunk = { + "chrom": np.full(n_matches, chrom), + "F_target": np.full(n_matches, f_target), + "R_target": r_targets, + "F_primer": np.full(n_matches, f_qseqid), + "R_primer": matched_r_qseqids, + "F_start": np.full(n_matches, f_start, dtype=np.int32), + "R_start": matched_r_starts.astype(np.int32), + "product_bp": (matched_r_starts - f_start + 1).astype(np.int32), + "F_pident": np.full( + n_matches, + f_pident[i] if f_pident is not None else np.nan, + dtype=np.float32, + ), + "R_pident": ( + r_pident[lo:hi].astype(np.float32) + if r_pident is not None + else np.full(n_matches, np.nan, dtype=np.float32) + ), + "F_mismatch": np.full( + n_matches, + f_mismatch[i] if f_mismatch is not None else -1, + dtype=np.int32, + ), + "R_mismatch": ( + r_mismatch[lo:hi].astype(np.int32) + if r_mismatch is not None + else np.full(n_matches, -1, dtype=np.int32) + ), + "F_align_len": np.full( + n_matches, + f_length[i] if f_length is not None else -1, + dtype=np.int32, + ), + "R_align_len": ( + r_length[lo:hi].astype(np.int32) + if r_length is not None + else np.full(n_matches, -1, dtype=np.int32) + ), + "F_evalue": np.full( + n_matches, + f_evalue[i] if f_evalue is not None else np.nan, + dtype=np.float32, + ), + "R_evalue": ( + r_evalue[lo:hi].astype(np.float32) + if r_evalue is not None + else np.full(n_matches, np.nan, dtype=np.float32) + ), + } + result_chunks.append(pd.DataFrame(chunk, columns=amplicon_columns)) + + if result_chunks: + chrom_amplicon_df = pd.concat(result_chunks, ignore_index=True) + del result_chunks + for col in ("chrom", "F_target", "R_target", "F_primer", "R_primer"): + chrom_amplicon_df[col] = chrom_amplicon_df[col].astype("category") + yield chrom, chrom_amplicon_df diff --git a/src/plexus/blast/specificity.py b/src/plexus/blast/specificity.py index d900166..ec522fe 100644 --- a/src/plexus/blast/specificity.py +++ b/src/plexus/blast/specificity.py @@ -46,6 +46,7 @@ def run_specificity_check( blast_penalty: int = -1, blast_max_hsps: int = 100, blast_dust: str = "yes", + max_bound_per_primer: int | None = 10000, ): """ Run BLAST on all candidate primers in the panel to check for specificity @@ -123,87 +124,92 @@ def run_specificity_check( bound_df = annotator.get_predicted_bound() _log_df_memory("bound_df", bound_df) + if max_bound_per_primer is not None: + pre_cap = len(bound_df) + bound_df = ( + bound_df.sort_values("evalue") + .groupby("qseqid", observed=True) + .head(max_bound_per_primer) + .reset_index(drop=True) + ) + if len(bound_df) < pre_cap: + n_capped = ( + bound_df.groupby("qseqid", observed=True) + .size() + .eq(max_bound_per_primer) + .sum() + ) + logger.info( + f"Per-primer cap ({max_bound_per_primer}): " + f"{pre_cap:,} → {len(bound_df):,} bound hits " + f"({n_capped} primer(s) capped)" + ) + _log_df_memory("bound_df (after cap)", bound_df) + # Free the full BLAST results — only bound_df is needed from here on del blast_df, runner, annotator _log_process_memory() - finder = AmpliconFinder(bound_df, target_map=target_map) - finder.find_amplicons(max_size_bp=max_amplicon_size) - _log_df_memory( - "amplicon_df", - finder.amplicon_df if finder.amplicon_df is not None else pd.DataFrame(), - ) - _log_process_memory() - - all_amplicons_df = finder.amplicon_df - - if all_amplicons_df is None or all_amplicons_df.empty: - logger.info("No amplicons found (on- or off-target).") - return - - # 6. Map results back to PrimerPairs - amplicon_groups = all_amplicons_df.groupby(["F_primer", "R_primer"]) - _empty = pd.DataFrame() - - # Use the panel's unique map to look up IDs + # 6. Build reverse lookup: (f_id, r_id) -> [(junction, pair), ...] seq_to_id = panel.unique_primer_map - + id_to_pairs: dict[tuple[str, str], list[tuple]] = {} for junction in panel.junctions: - n_checked = 0 - n_missing_on_target = 0 - for pair in junction.primer_pairs: - f_seq = pair.forward.seq - r_seq = pair.reverse.seq - - f_id = seq_to_id.get(f_seq) - r_id = seq_to_id.get(r_seq) - + f_id = seq_to_id.get(pair.forward.seq) + r_id = seq_to_id.get(pair.reverse.seq) if not f_id or not r_id: continue - pair.specificity_checked = True - n_checked += 1 - - fwd = ( - amplicon_groups.get_group((f_id, r_id)) - if (f_id, r_id) in amplicon_groups.groups - else _empty - ) - rev = ( - amplicon_groups.get_group((r_id, f_id)) - if (r_id, f_id) in amplicon_groups.groups - else _empty - ) - potential_products = pd.concat([fwd, rev]) if not rev.empty else fwd + pair.off_target_products = [] + pair.on_target_detected = False + id_to_pairs.setdefault((f_id, r_id), []).append((junction, pair)) + id_to_pairs.setdefault((r_id, f_id), []).append((junction, pair)) - if potential_products.empty: - pair.off_target_products = [] - pair.on_target_detected = False - else: + # 7. Process amplicons per chromosome to bound peak memory + finder = AmpliconFinder(bound_df, target_map=target_map) + total_amplicons = 0 + for _chrom, chrom_amp_df in finder.find_amplicons_by_chrom( + max_size_bp=max_amplicon_size + ): + total_amplicons += len(chrom_amp_df) + for (f_id, r_id), group_df in chrom_amp_df.groupby( + ["F_primer", "R_primer"], observed=True + ): + for junction, pair in id_to_pairs.get((f_id, r_id), []): on_mask = _is_on_target_vec( - potential_products, junction, pair, tolerance=ontarget_tolerance - ) - pair.on_target_detected = bool(on_mask.any()) - off_df = potential_products[~on_mask] - pair.off_target_products = ( - off_df.to_dict("records") if not off_df.empty else [] + group_df, junction, pair, tolerance=ontarget_tolerance ) + if on_mask.any(): + pair.on_target_detected = True + off_df = group_df[~on_mask] + if not off_df.empty: + pair.off_target_products.extend(off_df.to_dict("records")) + if total_amplicons == 0: + logger.info("No amplicons found (on- or off-target).") + else: + logger.info(f"Processed {total_amplicons:,} amplicons across all chromosomes") + + # Post-loop: log per-pair debug info and per-junction warnings + for junction in panel.junctions: + n_checked = sum(1 for p in junction.primer_pairs if p.specificity_checked) + n_missing = 0 + for pair in junction.primer_pairs: + if not pair.specificity_checked: + continue if not pair.on_target_detected: - n_missing_on_target += 1 + n_missing += 1 logger.debug( f"Pair {pair.pair_id}: on-target amplicon not detected by BLAST." ) - if pair.off_target_products: logger.debug( - f"Pair {pair.pair_id} has {len(pair.off_target_products)} off-target products." + f"Pair {pair.pair_id} has {len(pair.off_target_products)} " + "off-target products." ) - - if n_missing_on_target > 0: + if n_missing > 0: logger.warning( - f"Junction {junction.name}: {n_missing_on_target}/{n_checked} pairs " + f"Junction {junction.name}: {n_missing}/{n_checked} pairs " "have no on-target amplicon detected by BLAST. " "Check BLAST sensitivity or junction coordinates." ) diff --git a/src/plexus/config.py b/src/plexus/config.py index abf1674..0c3e735 100644 --- a/src/plexus/config.py +++ b/src/plexus/config.py @@ -352,6 +352,18 @@ class BlastParameters(BaseModel): "off-target detection." ), ) + max_bound_per_primer: int | None = Field( + default=10000, + ge=100, + description=( + "Maximum predicted binding sites to retain per primer, ranked by " + "lowest e-value. Acts as a safety net against primers in highly " + "repetitive regions (e.g. Alu elements) that can cause combinatorial " + "explosion in amplicon finding. Normal primers have 700-7000 bound " + "sites; only extreme outliers (>10000) are affected. " + "Set to None to disable." + ), + ) max_amplicon_size: int = Field( default=2000, ge=100, diff --git a/src/plexus/pipeline.py b/src/plexus/pipeline.py index 609433e..6b3717d 100644 --- a/src/plexus/pipeline.py +++ b/src/plexus/pipeline.py @@ -678,6 +678,7 @@ def advance_step(label=None): blast_penalty=blast_config.blast_penalty, blast_max_hsps=blast_config.blast_max_hsps, blast_dust=blast_config.blast_dust, + max_bound_per_primer=blast_config.max_bound_per_primer, ) result.steps_completed.append("specificity_checked") logger.info("Specificity check complete") diff --git a/src/plexus/version.py b/src/plexus/version.py index 72f26f5..0b2f79d 100644 --- a/src/plexus/version.py +++ b/src/plexus/version.py @@ -1 +1 @@ -__version__ = "1.1.2" +__version__ = "1.1.3" diff --git a/tests/test_blast_runner.py b/tests/test_blast_runner.py index 7fdc00d..9d6c658 100644 --- a/tests/test_blast_runner.py +++ b/tests/test_blast_runner.py @@ -255,6 +255,32 @@ def test_get_dataframe(runner, tmp_path): assert df.iloc[0]["qseqid"] == "Q1" +def test_get_dataframe_uses_compact_dtypes(runner, tmp_path): + """DataFrame uses categorical strings and downcasted numerics.""" + output_table = tmp_path / "output.txt" + runner.output_table = str(output_table) + + dummy_data = ["Q1", "S1", 100, 20, 0, 0, 1, 20, 100, 119, 0.0, 40.0, "plus", 20] + df_content = " ".join(map(str, dummy_data)) + output_table.write_text(df_content) + + df = runner.get_dataframe() + + # String columns should be categorical + for col in ("qseqid", "sseqid", "sstrand"): + assert df[col].dtype.name == "category", f"{col} should be categorical" + + # Numeric columns should be downcasted + assert df["sstart"].dtype == "int32" + assert df["send"].dtype == "int32" + assert df["length"].dtype == "int32" + assert df["mismatch"].dtype == "int32" + assert df["evalue"].dtype == "float32" + assert df["pident"].dtype == "float32" + assert df["bitscore"].dtype == "float32" + assert df["qlen"].dtype == "int16" + + def test_run_direct_tabular_output(runner, tmp_path): """output_table= produces outfmt 6 and sets self.output_table.""" output_table = str(tmp_path / "table.txt") diff --git a/tests/test_blast_specificity.py b/tests/test_blast_specificity.py index 536be95..1786310 100644 --- a/tests/test_blast_specificity.py +++ b/tests/test_blast_specificity.py @@ -43,6 +43,28 @@ def mock_panel(): return panel +def _mock_finder_by_chrom(MockFinder, amplicon_df): + """Configure a MockFinder to yield per-chromosome amplicon DataFrames. + + If amplicon_df is empty, the generator yields nothing. + Otherwise it groups by 'chrom' (if present) or yields a single + ('unknown', amplicon_df) tuple. + """ + finder_instance = MockFinder.return_value + if amplicon_df.empty: + finder_instance.find_amplicons_by_chrom.return_value = iter([]) + elif "chrom" in amplicon_df.columns: + chunks = [] + for chrom, group_df in amplicon_df.groupby("chrom"): + chunks.append((chrom, group_df.reset_index(drop=True))) + finder_instance.find_amplicons_by_chrom.return_value = iter(chunks) + else: + finder_instance.find_amplicons_by_chrom.return_value = iter( + [("unknown", amplicon_df)] + ) + return finder_instance + + def test_direct_tabular_no_archive(mock_panel, tmp_path): """Specificity check uses output_table directly, no archive file is created.""" with ( @@ -55,11 +77,10 @@ def test_direct_tabular_no_archive(mock_panel, tmp_path): annotator_instance = MockAnnotator.return_value annotator_instance.get_predicted_bound.return_value = pd.DataFrame( - {"dummy_bound": [1]} + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} ) - finder_instance = MockFinder.return_value - finder_instance.amplicon_df = pd.DataFrame() + _mock_finder_by_chrom(MockFinder, pd.DataFrame()) run_specificity_check(mock_panel, str(tmp_path), "fake_genome.fa") @@ -88,12 +109,11 @@ def test_run_specificity_check_integration(mock_panel, tmp_path): # 2. Mock Annotator behavior annotator_instance = MockAnnotator.return_value annotator_instance.get_predicted_bound.return_value = pd.DataFrame( - {"dummy_bound": [1]} + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} ) # 3. Mock Finder behavior — off-target at wrong coordinates - finder_instance = MockFinder.return_value - finder_instance.amplicon_df = pd.DataFrame( + amplicon_df = pd.DataFrame( [ { "chrom": "chr7", @@ -105,6 +125,7 @@ def test_run_specificity_check_integration(mock_panel, tmp_path): } ] ) + _mock_finder_by_chrom(MockFinder, amplicon_df) # Run run_specificity_check(mock_panel, str(tmp_path), "fake_genome.fa") @@ -138,11 +159,10 @@ def test_run_specificity_check_forwards_num_threads(mock_panel, tmp_path): annotator_instance = MockAnnotator.return_value annotator_instance.get_predicted_bound.return_value = pd.DataFrame( - {"dummy_bound": [1]} + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} ) - finder_instance = MockFinder.return_value - finder_instance.amplicon_df = pd.DataFrame() + _mock_finder_by_chrom(MockFinder, pd.DataFrame()) run_specificity_check(mock_panel, str(tmp_path), "genome.fa", num_threads=6) @@ -164,11 +184,10 @@ def test_run_specificity_check_forwards_blast_parameters(mock_panel, tmp_path): annotator_instance = MockAnnotator.return_value annotator_instance.get_predicted_bound.return_value = pd.DataFrame( - {"dummy_bound": [1]} + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} ) - finder_instance = MockFinder.return_value - finder_instance.amplicon_df = pd.DataFrame() + finder_instance = _mock_finder_by_chrom(MockFinder, pd.DataFrame()) run_specificity_check( mock_panel, @@ -195,8 +214,10 @@ def test_run_specificity_check_forwards_blast_parameters(mock_panel, tmp_path): three_prime_tolerance=3, ) - # Verify finder received custom max amplicon size - finder_instance.find_amplicons.assert_called_once_with(max_size_bp=5000) + # Verify finder used find_amplicons_by_chrom with custom max amplicon size + finder_instance.find_amplicons_by_chrom.assert_called_once_with( + max_size_bp=5000 + ) # Verify blast parameters were forwarded to runner.run() _, run_kwargs = runner_instance.run.call_args @@ -248,13 +269,12 @@ def test_run_specificity_check_swapped_orientation_off_target(mock_panel, tmp_pa annotator_instance = MockAnnotator.return_value annotator_instance.get_predicted_bound.return_value = pd.DataFrame( - {"dummy_bound": [1]} + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} ) # Swapped orientation: reverse primer on plus strand (F_primer), # forward primer on minus strand (R_primer). - finder_instance = MockFinder.return_value - finder_instance.amplicon_df = pd.DataFrame( + amplicon_df = pd.DataFrame( [ { "chrom": "chr12", @@ -266,6 +286,7 @@ def test_run_specificity_check_swapped_orientation_off_target(mock_panel, tmp_pa } ] ) + _mock_finder_by_chrom(MockFinder, amplicon_df) run_specificity_check(mock_panel, str(tmp_path), "fake_genome.fa") @@ -289,14 +310,13 @@ def test_run_specificity_check_on_target_detected(mock_panel, tmp_path): annotator_instance = MockAnnotator.return_value annotator_instance.get_predicted_bound.return_value = pd.DataFrame( - {"dummy_bound": [1]} + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} ) # Amplicon at the CORRECT coordinates: # F_start = design_start + fwd_start = 1000 + 10 = 1010 # R_start = design_start + rev_start + rev_length - 1 = 1000 + 180 + 22 - 1 = 1201 - finder_instance = MockFinder.return_value - finder_instance.amplicon_df = pd.DataFrame( + amplicon_df = pd.DataFrame( [ { "chrom": "chr7", @@ -308,6 +328,7 @@ def test_run_specificity_check_on_target_detected(mock_panel, tmp_path): } ] ) + _mock_finder_by_chrom(MockFinder, amplicon_df) run_specificity_check(mock_panel, str(tmp_path), "fake_genome.fa") @@ -330,12 +351,11 @@ def test_run_specificity_check_on_target_not_detected(mock_panel, tmp_path): annotator_instance = MockAnnotator.return_value annotator_instance.get_predicted_bound.return_value = pd.DataFrame( - {"dummy_bound": [1]} + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} ) # BLAST only found an off-target hit (pair maps somewhere, but not the intended locus) - finder_instance = MockFinder.return_value - finder_instance.amplicon_df = pd.DataFrame( + amplicon_df = pd.DataFrame( [ { "chrom": "chr7", @@ -347,6 +367,7 @@ def test_run_specificity_check_on_target_not_detected(mock_panel, tmp_path): } ] ) + _mock_finder_by_chrom(MockFinder, amplicon_df) run_specificity_check(mock_panel, str(tmp_path), "fake_genome.fa") @@ -356,6 +377,164 @@ def test_run_specificity_check_on_target_not_detected(mock_panel, tmp_path): assert len(pair.off_target_products) == 1 +def test_run_specificity_check_accepts_max_bound_per_primer(mock_panel, tmp_path): + """The max_bound_per_primer kwarg is accepted without error.""" + with ( + patch("plexus.blast.specificity.BlastRunner") as MockRunner, + patch("plexus.blast.specificity.BlastResultsAnnotator") as MockAnnotator, + patch("plexus.blast.specificity.AmpliconFinder") as MockFinder, + patch("os.makedirs"), + ): + runner_instance = MockRunner.return_value + runner_instance.get_dataframe.return_value = pd.DataFrame({"dummy": [1]}) + + annotator_instance = MockAnnotator.return_value + annotator_instance.get_predicted_bound.return_value = pd.DataFrame( + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} + ) + + _mock_finder_by_chrom(MockFinder, pd.DataFrame()) + + # Should not raise + run_specificity_check( + mock_panel, + str(tmp_path), + "genome.fa", + max_bound_per_primer=500, + ) + + +def test_bound_cap_keeps_lowest_evalue(): + """Cap keeps the rows with the lowest e-values per primer.""" + bound_df = pd.DataFrame( + { + "qseqid": pd.Categorical(["P1"] * 20), + "evalue": list(range(20, 0, -1)), # 20, 19, ..., 1 + "sseqid": ["chr1"] * 20, + "sstart": list(range(100, 2100, 100)), + "sstrand": ["plus"] * 20, + } + ) + cap = 5 + pre_cap = len(bound_df) + result = ( + bound_df.sort_values("evalue") + .groupby("qseqid", observed=True) + .head(cap) + .reset_index(drop=True) + ) + assert len(result) == cap + assert result["evalue"].max() == 5 # kept lowest 5: 1, 2, 3, 4, 5 + assert len(result) < pre_cap + + +def test_bound_cap_none_disables(): + """When max_bound_per_primer is None, bound_df is unchanged.""" + bound_df = pd.DataFrame( + { + "qseqid": pd.Categorical(["P1"] * 20), + "evalue": list(range(20)), + "sseqid": ["chr1"] * 20, + "sstart": list(range(100, 2100, 100)), + "sstrand": ["plus"] * 20, + } + ) + # Simulating what the code does when max_bound_per_primer is None: nothing + result = bound_df.copy() + assert len(result) == 20 + + +def test_per_chrom_multi_chromosome_accumulation(mock_panel, tmp_path): + """Amplicons on chr7 and chr12 both contribute off-targets to the same pair.""" + with ( + patch("plexus.blast.specificity.BlastRunner") as MockRunner, + patch("plexus.blast.specificity.BlastResultsAnnotator") as MockAnnotator, + patch("plexus.blast.specificity.AmpliconFinder") as MockFinder, + patch("os.makedirs"), + ): + runner_instance = MockRunner.return_value + runner_instance.get_dataframe.return_value = pd.DataFrame({"dummy": [1]}) + + annotator_instance = MockAnnotator.return_value + annotator_instance.get_predicted_bound.return_value = pd.DataFrame( + {"qseqid": pd.Categorical(["P1_F"]), "evalue": [1.0]} + ) + + # Two chromosomes, both with off-target hits for the same primer pair + amplicon_df = pd.DataFrame( + [ + { + "chrom": "chr7", + "F_primer": "P1_F", + "R_primer": "P1_R", + "product_bp": 500, + "F_start": 5000, + "R_start": 5500, + }, + { + "chrom": "chr12", + "F_primer": "P1_F", + "R_primer": "P1_R", + "product_bp": 300, + "F_start": 8000, + "R_start": 8300, + }, + ] + ) + # Mock yields two separate chromosome chunks + finder_instance = MockFinder.return_value + chr7_df = amplicon_df[amplicon_df["chrom"] == "chr7"].reset_index(drop=True) + chr12_df = amplicon_df[amplicon_df["chrom"] == "chr12"].reset_index(drop=True) + finder_instance.find_amplicons_by_chrom.return_value = iter( + [("chr7", chr7_df), ("chr12", chr12_df)] + ) + + run_specificity_check(mock_panel, str(tmp_path), "fake_genome.fa") + + pair = mock_panel.junctions[0].primer_pairs[0] + assert pair.specificity_checked is True + # Both off-target hits accumulated from separate chromosomes + assert len(pair.off_target_products) == 2 + chroms = {p["chrom"] for p in pair.off_target_products} + assert chroms == {"chr7", "chr12"} + + +# --------------------------------------------------------------------------- +# AmpliconFinder.find_amplicons_by_chrom unit test +# --------------------------------------------------------------------------- + + +def test_find_amplicons_by_chrom_yields_per_chromosome(): + """find_amplicons_by_chrom yields one tuple per chromosome with amplicons.""" + from plexus.blast.offtarget_finder import AmpliconFinder + + bound_df = pd.DataFrame( + { + "qseqid": ["F1", "F1", "R1", "R1"], + "sseqid": ["chr1", "chr7", "chr1", "chr7"], + "sstart": [100, 200, 200, 400], + "send": [120, 220, 180, 380], + "sstrand": ["plus", "plus", "minus", "minus"], + "qlen": [20, 20, 20, 20], + "qend": [20, 20, 20, 20], + "predicted_bound": [True, True, True, True], + } + ) + + finder = AmpliconFinder(bound_df, target_map={"F1": "T1", "R1": "T1"}) + results = list(finder.find_amplicons_by_chrom(max_size_bp=6000)) + + # Should yield results for both chromosomes + chroms = [chrom for chrom, _ in results] + assert "chr1" in chroms + assert "chr7" in chroms + + for chrom, df in results: + assert not df.empty + # All rows should be for this chromosome + assert (df["chrom"] == chrom).all() + + # --------------------------------------------------------------------------- # _is_on_target coordinate-based classification # ---------------------------------------------------------------------------