vastly simplifying primer preparation with less heavy handed handling by nextflow, please yet more tidying around the codebase

Author: nrminor

Containerfile

@@ -76,5 +76,5 @@ COPY Cargo.lock $HOME/Cargo.lock
 COPY bin/find_and_trim_amplicons.rs $HOME/bin/find_and_trim_amplicons.rs
 RUN cd $HOME && \
     export PATH="$HOME/.pixi/envs/default/bin:$PATH" && \
-    cargo build --release && \
+    RUSTFLAGS="-C target-cpu=native" cargo build --release && \
     cp $HOME/target/release/find_and_trim_amplicons $HOME/.pixi/envs/default/bin/

bin/create_amplicon_tsv.py

@@ -6,128 +6,200 @@
 # ]
 # ///
 
-"""Summarize amplicon coverage from stats files and BED file."""
+"""
+Summarize amplicon coverage from stats files and primer position data.
+
+Supports two input modes for position data:
+  1. --primer-tsv: Use primer_pairs.tsv from prepare_primers.py (preferred)
+  2. --bed: Use BED file with primer coordinates (legacy)
+
+The primer TSV approach is cleaner as it already contains pre-computed
+amplicon positions and works regardless of primer naming conventions.
+"""
 
 from __future__ import annotations
 
 import argparse
+import sys
+from pathlib import Path
 
 import polars as pl
 
 
-def main(bed_file: str, output_file: str, stats_pattern: str) -> None:
+def _load_positions_from_tsv(tsv_path: str) -> pl.LazyFrame:
+    """
+    Load amplicon positions from primer_pairs.tsv.
+
+    The TSV has columns: amplicon_name, fwd_sequence, rev_sequence, chrom,
+    amplicon_start, amplicon_end
+
+    Returns a LazyFrame with: amplicon_name, start_pos, end_pos
+    """
+    return pl.scan_csv(tsv_path, separator="\t").select(
+        pl.col("amplicon_name"),
+        pl.col("amplicon_start").alias("start_pos"),
+        pl.col("amplicon_end").alias("end_pos"),
+    )
+
+
+def _load_positions_from_bed(
+    bed_path: str, fwd_suffix: str, rev_suffix: str
+) -> pl.LazyFrame:
     """
-    Create amplicon summary TSV from stats files and BED file.
+    Load amplicon positions from BED file (legacy approach).
 
-    Reads all stats files matching the glob pattern, extracts sample and amplicon
-    names from the 'file' column, joins with BED file to get amplicon positions,
-    and writes a summary TSV.
+    Extracts base amplicon name from primer names and computes positions
+    from LEFT/RIGHT primer coordinates.
 
-    Parameters:
-        bed_file: Path to BED file with primer positions
-        output_file: Path to output TSV file
-        stats_pattern: Glob pattern for stats files (e.g., "stats_*.tsv")
+    Returns a LazyFrame with: amplicon_name, start_pos, end_pos
     """
-    # Read all stats files with glob, extract info from the 'file' column
-    # (which contains filenames like "SAMPLE.QIAseq_X-Y.no_downsampling.fasta.gz")
-    stats = (
+    bed = (
         pl.scan_csv(
-            stats_pattern,
+            bed_path,
             separator="\t",
-            glob=True,
+            has_header=False,
+            new_columns=["chrom", "start", "end", "name", "score", "strand"],
         )
         .with_columns(
-            # Normalize .QIAseq_ to _QIAseq_ for consistent parsing
-            pl.col("file")
-            .str.replace(r"\.QIAseq_", "_QIAseq_")
-            .alias("normalized_file")
+            # Extract base amplicon name by removing suffixes and indices
+            pl.col("name")
+            .str.replace(fwd_suffix, "")
+            .str.replace(rev_suffix, "")
+            .str.replace(r"_splice\d+$", "")
+            .str.replace(r"-\d+$", "")
+            .alias("amplicon_name"),
+            # Flag primer direction
+            pl.col("name").str.contains(fwd_suffix).alias("is_fwd"),
+            pl.col("name").str.contains(rev_suffix).alias("is_rev"),
         )
+        .filter(pl.col("amplicon_name").is_not_null())
+    )
+
+    # Aggregate to get amplicon span
+    return bed.group_by("amplicon_name").agg(
+        pl.col("start").filter(pl.col("is_fwd")).min().alias("start_pos"),
+        pl.col("end").filter(pl.col("is_rev")).max().alias("end_pos"),
+    )
+
+
+def _parse_stats_files(stats_pattern: str) -> pl.LazyFrame:
+    """
+    Parse amplicon stats files and extract sample/amplicon info.
+
+    Stats files have columns including 'file' and 'num_seqs'.
+    The 'file' column contains filenames like:
+      "SAMPLE.amplicon_name.no_downsampling.fasta.gz"
+
+    Returns a LazyFrame with: sample_name, amplicon_name, reads
+    """
+    return (
+        pl.scan_csv(stats_pattern, separator="\t", glob=True)
         .with_columns(
-            # Extract sample name: everything before _QIAseq_, or first dot-segment
-            pl.when(pl.col("normalized_file").str.contains("_QIAseq_"))
-            .then(
-                pl.col("normalized_file").str.extract(
-                    r"^([^_]+(?:_[^_]+)*?)_QIAseq_", group_index=1
-                )
-            )
-            .otherwise(
-                pl.col("normalized_file").str.extract(r"^([^.]+)", group_index=1)
-            )
-            .alias("sample_name"),
-            # Extract amplicon name: QIAseq_XXX (including any suffix like -1)
+            # Extract sample name: first segment before '.'
+            pl.col("file").str.extract(r"^([^.]+)", group_index=1).alias("sample_name"),
+            # Extract amplicon name: second segment (between first and second '.')
             pl.col("file")
-            .str.extract(r"(QIAseq_[^.]+)", group_index=1)
+            .str.extract(r"^[^.]+\.([^.]+)", group_index=1)
             .alias("amplicon_name"),
-            # Extract base amplicon for joining: QIAseq_N (just the number, no suffix)
-            pl.col("file")
-            .str.extract(r"(QIAseq_\d+)", group_index=1)
-            .alias("base_amplicon"),
         )
         .select(
             "sample_name",
             "amplicon_name",
-            "base_amplicon",
             pl.col("num_seqs").alias("reads"),
         )
     )
 
-    # Read BED file, compute amplicon start/end positions
-    bed = (
-        pl.scan_csv(
-            bed_file,
-            separator="\t",
-            has_header=False,
-            new_columns=["chrom", "start", "end", "name", "score", "strand"],
-        )
-        .with_columns(
-            # Extract base amplicon from primer name (e.g., QIAseq_2_LEFT -> QIAseq_2)
-            pl.col("name")
-            .str.extract(r"(QIAseq_\d+)", group_index=1)
-            .alias("base_amplicon"),
-            # Flag LEFT vs RIGHT primers
-            pl.col("name").str.contains("_LEFT").alias("is_left"),
-            pl.col("name").str.contains("_RIGHT").alias("is_right"),
-        )
-        .filter(pl.col("base_amplicon").is_not_null())
-    )
-
-    # Aggregate to get min(start) for LEFT primers, max(end) for RIGHT primers
-    amplicon_positions = bed.group_by("base_amplicon").agg(
-        pl.col("start").filter(pl.col("is_left")).min().alias("start_pos"),
-        pl.col("end").filter(pl.col("is_right")).max().alias("end_pos"),
-    )
 
-    # Join stats with positions and format output
+def _create_summary(
+    stats_lf: pl.LazyFrame,
+    positions_lf: pl.LazyFrame,
+    output_path: str,
+) -> None:
+    """
+    Join stats with positions and write summary TSV.
+    """
     result = (
-        stats.join(amplicon_positions, on="base_amplicon", how="left")
+        stats_lf.join(positions_lf, on="amplicon_name", how="left")
         .select(
             "sample_name",
             "amplicon_name",
             pl.col("start_pos").cast(pl.String).fill_null("NA"),
             pl.col("end_pos").cast(pl.String).fill_null("NA"),
             "reads",
         )
+        .sort("sample_name", "amplicon_name")
         .collect()
     )
 
-    result.write_csv(output_file, separator="\t")
+    result.write_csv(output_path, separator="\t")
+    print(f"Wrote {len(result)} rows to {output_path}", file=sys.stderr)
 
 
-if __name__ == "__main__":
+def main() -> None:
     parser = argparse.ArgumentParser(
-        description="Summarize amplicon coverage from stats files and BED file.",
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    parser.add_argument("--bed", required=True, help="Path to BED file")
-    parser.add_argument(
-        "--output",
-        default="amplicon_summary.tsv",
-        help="Output TSV file",
+
+    # Position data input (mutually exclusive)
+    pos_group = parser.add_mutually_exclusive_group(required=True)
+    pos_group.add_argument(
+        "--primer-tsv",
+        type=Path,
+        help="Primer pairs TSV from prepare_primers.py (preferred)",
     )
+    pos_group.add_argument(
+        "--bed",
+        type=Path,
+        help="BED file with primer coordinates (legacy)",
+    )
+
+    # Stats input
     parser.add_argument(
         "--pattern",
         default="stats_*.tsv",
         help="Glob pattern for stats files (default: stats_*.tsv)",
     )
 
+    # Output
+    parser.add_argument(
+        "--output",
+        default="amplicon_summary.tsv",
+        help="Output TSV file (default: amplicon_summary.tsv)",
+    )
+
+    # BED-specific options
+    parser.add_argument(
+        "--fwd-suffix",
+        default="_LEFT",
+        help="Forward primer suffix for BED parsing (default: _LEFT)",
+    )
+    parser.add_argument(
+        "--rev-suffix",
+        default="_RIGHT",
+        help="Reverse primer suffix for BED parsing (default: _RIGHT)",
+    )
+
     args = parser.parse_args()
-    main(args.bed, args.output, args.pattern)
+
+    # Load position data
+    if args.primer_tsv:
+        assert args.primer_tsv.is_file(), f"Primer TSV not found: {args.primer_tsv}"
+        positions_lf = _load_positions_from_tsv(str(args.primer_tsv))
+    else:
+        assert args.bed.is_file(), f"BED file not found: {args.bed}"
+        positions_lf = _load_positions_from_bed(
+            str(args.bed),
+            args.fwd_suffix,
+            args.rev_suffix,
+        )
+
+    # Parse stats files
+    stats_lf = _parse_stats_files(args.pattern)
+
+    # Create and write summary
+    _create_summary(stats_lf, positions_lf, args.output)
+
+
+if __name__ == "__main__":
+    main()