Pipeline

End-to-end pipeline entry points, results serialisation, and the core data classes.

Running the pipeline

run_pipeline_streaming

run_pipeline_streaming(
    native_bam: PathLike,
    native_fastq: PathLike,
    native_blow5: PathLike,
    ivt_bam: PathLike,
    ivt_fastq: PathLike,
    ivt_blow5: PathLike,
    ref_fasta: PathLike,
    *,
    min_depth: int = 15,
    depth_mode: Literal[
        "mean_coverage", "read_count"
    ] = "read_count",
    use_cuda: Optional[bool] = None,
    cuda_devices: Optional[list[int]] = None,
    padding: int = 1,
    output_dir: Optional[PathLike] = None,
    cleanup_temp: bool = True,
    rna: bool = True,
    kmer_model: Optional[str] = None,
    pore: str = DEFAULT_PORE,
    krill_hmm: bool = False,
    min_mapq: int = 20,
    primary_only: bool = True,
    threads: int = 1,
    num_cuda_streams: int = 16,
    run_hmm: bool = True,
    hmm_params: object = None,
    target_contigs: Optional[list[str]] = None,
    keep_intermediate: bool = False,
    gpu_memory_limit: Optional[int] = None,
    subsample: bool = True,
    subsample_n: int = 300,
    legacy_scoring: bool = False,
    mod_threshold: float = 0.9,
    write_bam: bool = True,
    resume: bool = False,
    read_intersection: bool = True
) -> tuple[dict[str, Any], PipelineMetadata]

Memory-efficient streaming pipeline: DTW → HMM → aggregation per contig.

Each contig is processed end-to-end in a worker, which writes its own <contig>.tsv and <contig>.bam slices to <output_dir>/per_contig/ and drops the heavy cmr before returning. The main process then merges slices into final outputs.

Peak memory is bounded by O(single_contig + N_workers) rather than growing with the total number of contigs.

Parameters:

Name	Type	Description	Default
output_dir	Optional[PathLike]	Required. Final outputs are written to <output_dir>/site_results.tsv and (when write_bam) <output_dir>/read_results.bam. If None, a temporary directory is used and cleaned up on exit.	None
target_contigs	Optional[list[str]]	If given, only process these contig(s). Contigs not passing depth filters are silently skipped.	None
keep_intermediate	bool	Save per-contig ContigResult pickles under <output_dir>/intermediate/ and keep <output_dir>/per_contig/ on disk after merging.	False
write_bam	bool	Whether to produce a final read_results.bam (set to False to skip mod-BAM construction entirely).	True
run_hmm	bool	Whether to run HMM smoothing (V3).	True
hmm_params	object	Optional trained HMM parameters.	None

Returns:

Type	Description
tuple[dict[str, Any], PipelineMetadata]	(output_paths, metadata) where output_paths is a dict with keys site_tsv (Path), read_bam (Path or None), per_contig_dir (Path or None), n_total_sites (int), and n_significant (int).

Source code in baleen/eventalign/_pipeline.py

def run_pipeline_streaming(
    native_bam: PathLike,
    native_fastq: PathLike,
    native_blow5: PathLike,
    ivt_bam: PathLike,
    ivt_fastq: PathLike,
    ivt_blow5: PathLike,
    ref_fasta: PathLike,
    *,
    min_depth: int = 15,
    depth_mode: Literal["mean_coverage", "read_count"] = "read_count",
    use_cuda: Optional[bool] = None,
    cuda_devices: Optional[list[int]] = None,
    padding: int = 1,
    output_dir: Optional[PathLike] = None,
    cleanup_temp: bool = True,
    rna: bool = True,
    kmer_model: Optional[str] = None,
    pore: str = _eventalign.DEFAULT_PORE,
    krill_hmm: bool = False,
    min_mapq: int = 20,
    primary_only: bool = True,
    threads: int = 1,
    num_cuda_streams: int = 16,
    run_hmm: bool = True,
    hmm_params: object = None,
    target_contigs: Optional[list[str]] = None,
    keep_intermediate: bool = False,
    gpu_memory_limit: Optional[int] = None,
    subsample: bool = True,
    subsample_n: int = 300,
    legacy_scoring: bool = False,
    mod_threshold: float = 0.9,
    write_bam: bool = True,
    resume: bool = False,
    read_intersection: bool = True,
) -> tuple[dict[str, Any], PipelineMetadata]:
    """Memory-efficient streaming pipeline: DTW → HMM → aggregation per contig.

    Each contig is processed end-to-end in a worker, which writes its
    own ``<contig>.tsv`` and ``<contig>.bam`` slices to
    ``<output_dir>/per_contig/`` and drops the heavy ``cmr`` before
    returning.  The main process then merges slices into final outputs.

    Peak memory is bounded by ``O(single_contig + N_workers)`` rather than
    growing with the total number of contigs.

    Parameters
    ----------
    output_dir
        Required.  Final outputs are written to ``<output_dir>/site_results.tsv``
        and (when *write_bam*) ``<output_dir>/read_results.bam``.  If ``None``,
        a temporary directory is used and cleaned up on exit.
    target_contigs
        If given, only process these contig(s).  Contigs not passing depth
        filters are silently skipped.
    keep_intermediate
        Save per-contig ``ContigResult`` pickles under
        ``<output_dir>/intermediate/`` and keep ``<output_dir>/per_contig/``
        on disk after merging.
    write_bam
        Whether to produce a final ``read_results.bam`` (set to False to
        skip mod-BAM construction entirely).
    run_hmm
        Whether to run HMM smoothing (V3).
    hmm_params
        Optional trained HMM parameters.

    Returns
    -------
    tuple[dict[str, Any], PipelineMetadata]
        ``(output_paths, metadata)`` where ``output_paths`` is a dict with
        keys ``site_tsv`` (Path), ``read_bam`` (Path or None),
        ``per_contig_dir`` (Path or None), ``n_total_sites`` (int), and
        ``n_significant`` (int).
    """
    pipeline_t0 = time.perf_counter()
    logger.info("=" * 60)
    logger.info("Starting baleen streaming pipeline")
    logger.info("  native_bam:   %s", native_bam)
    logger.info("  native_fastq: %s", native_fastq)
    logger.info("  native_blow5: %s", native_blow5)
    logger.info("  ivt_bam:      %s", ivt_bam)
    logger.info("  ivt_fastq:    %s", ivt_fastq)
    logger.info("  ivt_blow5:    %s", ivt_blow5)
    logger.info("  ref_fasta:    %s", ref_fasta)
    logger.info("  min_depth=%d  depth_mode=%s  use_cuda=%s  rna=%s  padding=%d  threads=%d",
                min_depth, depth_mode, use_cuda, rna, padding, threads)
    logger.info("  min_mapq=%d  primary_only=%s  cuda_streams=%d",
                min_mapq, primary_only, num_cuda_streams)
    logger.info("  subsample=%s  subsample_n=%d  gpu_memory_limit=%s",
                subsample, subsample_n, gpu_memory_limit)
    logger.info("  run_hmm=%s  legacy_scoring=%s  mod_threshold=%.2f",
                run_hmm, legacy_scoring, mod_threshold)
    logger.info("  target_contigs=%s  keep_intermediate=%s  cleanup_temp=%s",
                target_contigs, keep_intermediate, cleanup_temp)
    logger.info("  kmer_model=%s  pore=%s  krill_hmm=%s", kmer_model, pore, krill_hmm)
    logger.info("=" * 60)

    if threads < 1:
        raise ValueError(f"threads must be >= 1, got {threads}")

    # Resolve cuda_devices from legacy use_cuda if needed
    if cuda_devices is None and use_cuda is not None:
        if use_cuda is True:
            cuda_devices = None  # auto-detect all GPUs
        elif use_cuda is False:
            cuda_devices = []  # CPU mode
    if cuda_devices is not None:
        use_cuda = len(cuda_devices) > 0 if cuda_devices else False

    native_bam = Path(native_bam)
    native_fastq = Path(native_fastq)
    native_blow5 = Path(native_blow5)
    ivt_bam = Path(ivt_bam)
    ivt_fastq = Path(ivt_fastq)
    ivt_blow5 = Path(ivt_blow5)
    ref_fasta = Path(ref_fasta)

    # ---- Step 1: krill engine check ----
    logger.info("[Step 1/5] Checking krill availability...")
    eventalign_version = _eventalign.check_krill()
    logger.info("[Step 1/5] krill version %s OK", eventalign_version)

    # ---- Step 2: Indexing ----
    logger.info("[Step 2/5] Indexing BLOW5 files...")
    step_t0 = time.perf_counter()
    _eventalign.index_blow5(native_blow5)
    _eventalign.index_blow5(ivt_blow5)
    logger.info("[Step 2/5] Indexing complete (%s)", _fmt_elapsed(time.perf_counter() - step_t0))

    # ---- Step 2.5: Read-ID intersection (BAM ∩ FASTQ ∩ BLOW5) ----
    # eventalign silently drops BAM reads whose UUIDs are not in
    # the BLOW5 signal file; computing the intersection up-front keeps
    # contig stats, ``min_depth`` filtering, and subsampling all in
    # sync with the read set that will actually produce signals.
    #
    # The intersection sets are written to disk and passed by path
    # (not by ``set[str]``) across the worker spawn boundary.
    allowed_native_reads_path: Optional[Path] = None
    allowed_ivt_reads_path: Optional[Path] = None
    allowed_native: Optional[set[str]] = None
    allowed_ivt: Optional[set[str]] = None
    if read_intersection:
        from baleen.eventalign._read_ids import (
            compute_condition_intersection,
            write_read_ids,
        )

        logger.info("[Step 2.5/5] Computing read-id intersection (BAM ∩ FASTQ ∩ BLOW5)...")
        step_t0 = time.perf_counter()
        allowed_native = compute_condition_intersection(
            bam=native_bam,
            fastq=native_fastq,
            blow5=native_blow5,
            primary_only=primary_only,
            min_mapq=min_mapq,
            label="native",
        )
        allowed_ivt = compute_condition_intersection(
            bam=ivt_bam,
            fastq=ivt_fastq,
            blow5=ivt_blow5,
            primary_only=primary_only,
            min_mapq=min_mapq,
            label="ivt",
        )
        logger.info(
            "[Step 2.5/5] Intersection complete: native=%d, ivt=%d reads (%s)",
            len(allowed_native), len(allowed_ivt),
            _fmt_elapsed(time.perf_counter() - step_t0),
        )

        # Persist intersection sets to disk so workers can lazy-load via
        # path instead of receiving multi-MB ``set[str]`` payloads across
        # the spawn pickle boundary (2924 contigs × millions of reads
        # would otherwise dominate dispatch cost).
        _intersection_dir = Path(
            tempfile.mkdtemp(prefix="baleen-intersection-")
        )
        allowed_native_reads_path = write_read_ids(
            allowed_native, _intersection_dir / "allowed_native_reads.txt"
        )
        allowed_ivt_reads_path = write_read_ids(
            allowed_ivt, _intersection_dir / "allowed_ivt_reads.txt"
        )

    # ---- Step 3: BAM validation & contig stats ----
    logger.info("[Step 3/5] Validating BAMs and computing contig statistics...")
    step_t0 = time.perf_counter()
    _bam.validate_bam(native_bam)
    _bam.validate_bam(ivt_bam)
    native_stats = _bam.get_contig_stats(
        native_bam, min_mapq=min_mapq, primary_only=primary_only,
        allowed_reads=allowed_native, _validated=True,
    )
    ivt_stats = _bam.get_contig_stats(
        ivt_bam, min_mapq=min_mapq, primary_only=primary_only,
        allowed_reads=allowed_ivt, _validated=True,
    )
    logger.info("[Step 3/5] BAM stats complete: %d native contigs, %d IVT contigs (%s)",
                len(native_stats), len(ivt_stats), _fmt_elapsed(time.perf_counter() - step_t0))

    # ---- Step 4: Contig filtering ----
    logger.info("[Step 4/5] Filtering contigs (min_depth=%d, depth_mode=%s)...",
                min_depth, depth_mode)
    step_t0 = time.perf_counter()
    passed_contigs, filter_results = _bam.filter_contigs(
        native_stats, ivt_stats, min_depth=float(min_depth), depth_mode=depth_mode,
    )

    # Apply target contig filter
    if target_contigs is not None:
        target_set = set(target_contigs)
        skipped_targets = target_set - set(passed_contigs)
        if skipped_targets:
            logger.warning("  Target contigs not passing filters: %s", sorted(skipped_targets))
        passed_contigs = [c for c in passed_contigs if c in target_set]

    logger.info("[Step 4/5] %d contigs to process (%s)",
                len(passed_contigs), _fmt_elapsed(time.perf_counter() - step_t0))

    metadata = PipelineMetadata(
        eventalign_version=eventalign_version,
        min_depth=min_depth,
        use_cuda=use_cuda,
        padding=padding,
        n_contigs_total=len(filter_results),
        n_contigs_passed_filter=len(passed_contigs),
        n_contigs_skipped=len(filter_results) - len(passed_contigs),
        filter_results=filter_results,
    )

    # Resolve output_dir (use a tempdir if not given, cleaned up on exit).
    output_dir_temp: Optional[Path] = None
    if output_dir is None:
        output_dir_temp = Path(tempfile.mkdtemp(prefix="baleen-output-"))
        output_dir_path = output_dir_temp
        logger.warning(
            "  output_dir not specified; using temporary %s (will be removed)",
            output_dir_path,
        )
    else:
        output_dir_path = Path(output_dir)
    output_dir_path.mkdir(parents=True, exist_ok=True)

    per_contig_dir = output_dir_path / "per_contig"
    # Detect prior state BEFORE we create the directory.  Missing dir or
    # missing fingerprint both mean "no usable prior run" — even with
    # --resume, treat it as a fresh run (nothing to skip, nothing to
    # validate against).  This matters for legacy interrupted runs that
    # pre-date the fingerprint file.
    fp_path_pre = per_contig_dir / _RESUME_PARAMS_FILENAME
    has_prior_fingerprint = per_contig_dir.exists() and fp_path_pre.exists()
    per_contig_dir.mkdir(parents=True, exist_ok=True)

    site_tsv_path = output_dir_path / "site_results.tsv"
    final_bam_path: Optional[Path] = (
        output_dir_path / "read_results.bam" if write_bam else None
    )

    # ---- Resume: scan & validate before dispatching workers ----
    fingerprint = _compute_resume_fingerprint(
        native_bam=native_bam,
        native_fastq=native_fastq,
        native_blow5=native_blow5,
        ivt_bam=ivt_bam,
        ivt_fastq=ivt_fastq,
        ivt_blow5=ivt_blow5,
        ref_fasta=ref_fasta,
        min_depth=min_depth,
        depth_mode=depth_mode,
        padding=padding,
        min_mapq=min_mapq,
        primary_only=primary_only,
        subsample=subsample,
        subsample_n=subsample_n,
        legacy_scoring=legacy_scoring,
        mod_threshold=mod_threshold,
        write_bam=write_bam,
        run_hmm=run_hmm,
        target_contigs=target_contigs,
        read_intersection=read_intersection,
        pore=pore,
        krill_hmm=krill_hmm,
    )
    resumed_summaries: list[ContigSummary] = []
    if resume:
        if has_prior_fingerprint:
            _validate_resume_compatibility(per_contig_dir, fingerprint)
            resumed_map = _scan_completed_contigs(
                per_contig_dir, passed_contigs, write_bam,
            )
            resumed_summaries = list(resumed_map.values())
            if resumed_summaries:
                logger.info(
                    "[Resume] Skipping %d/%d contigs already on disk under %s",
                    len(resumed_summaries), len(passed_contigs), per_contig_dir,
                )
            passed_contigs = [c for c in passed_contigs if c not in resumed_map]
        else:
            logger.info(
                "[Resume] No %s under %s; treating as a fresh run.",
                _RESUME_PARAMS_FILENAME, per_contig_dir,
            )
    _write_resume_fingerprint(per_contig_dir, fingerprint)

    if not passed_contigs and not resumed_summaries:
        logger.warning("[Step 5/5] No contigs to process; returning empty results.")
        # Still write an empty TSV (header only) for consistency, unless we
        # are going to delete the parent tempdir on the way out.
        from baleen.eventalign._aggregation import write_site_tsv
        write_site_tsv([], site_tsv_path)
        # Clean up empty per_contig dir we just created.
        shutil.rmtree(per_contig_dir, ignore_errors=True)
        returned_site_tsv: Optional[Path] = site_tsv_path
        if output_dir_temp is not None:
            shutil.rmtree(output_dir_temp, ignore_errors=True)
            returned_site_tsv = None
        elapsed = _fmt_elapsed(time.perf_counter() - pipeline_t0)
        logger.info("Pipeline finished (no results) in %s", elapsed)
        return (
            {
                "site_tsv": returned_site_tsv,
                "read_bam": None,
                "per_contig_dir": None,
                "n_total_sites": 0,
                "n_significant": 0,
            },
            metadata,
        )

    # ---- Step 5: Per-contig streaming (DTW → HMM → aggregation → flush) ----
    logger.info("[Step 5/5] Processing %d contigs (streaming flush: DTW → HMM → aggregation → disk)...",
                len(passed_contigs))
    step5_t0 = time.perf_counter()
    tmp_root = Path(tempfile.mkdtemp(prefix="baleen-streaming-"))

    intermediate_dir: Optional[Path] = None
    if keep_intermediate:
        intermediate_dir = output_dir_path / "intermediate"

    # Build BAM header once (heavy I/O).  pysam.AlignmentHeader is NOT
    # picklable (cdef-class with non-trivial __cinit__), so we ship its
    # ``to_dict()`` representation across the spawn boundary and rebuild
    # in each worker via ``pysam.AlignmentHeader.from_dict``.
    bam_header_dict: Optional[dict] = None
    if write_bam:
        # Ensure full input BAMs are indexed for per-contig fetch().
        from baleen.eventalign._read_bam import (
            _build_header_from_bam,
            _ensure_bam_indexed,
            merge_contig_bams,
        )
        _ensure_bam_indexed(native_bam)
        _ensure_bam_indexed(ivt_bam)
        bam_header_dict = _build_header_from_bam(native_bam, ref_fasta).to_dict()

    from baleen.eventalign._aggregation import merge_contig_tsvs

    gpu_mems = _get_gpu_memory(cuda_devices) if gpu_memory_limit is None else [gpu_memory_limit]
    gpu_workers, device_for_worker = _gpu_concurrent_workers(threads, gpu_mems, cuda_devices)

    summaries: list[ContigSummary] = list(resumed_summaries)
    failed: list[str] = []

    try:
        worker_kwargs = dict(
            native_bam=native_bam,
            native_fastq=native_fastq,
            native_blow5=native_blow5,
            ivt_bam=ivt_bam,
            ivt_fastq=ivt_fastq,
            ivt_blow5=ivt_blow5,
            ref_fasta=ref_fasta,
            native_stats=native_stats,
            ivt_stats=ivt_stats,
            tmp_root=tmp_root,
            use_cuda=use_cuda,
            padding=padding,
            rna=rna,
            kmer_model=kmer_model,
            pore=pore,
            krill_hmm=krill_hmm,
            min_mapq=min_mapq,
            primary_only=primary_only,
            cleanup_temp=cleanup_temp,
            num_cuda_streams=num_cuda_streams,
            per_contig_dir=per_contig_dir,
            bam_header_dict=bam_header_dict,
            write_bam=write_bam,
            run_hmm=run_hmm,
            hmm_params=hmm_params,
            keep_intermediate=keep_intermediate,
            intermediate_dir=intermediate_dir,
            subsample=subsample,
            subsample_n=subsample_n,
            gpu_memory_bytes=gpu_mems[0] if gpu_mems else 8 * 1024 ** 3,
            legacy_scoring=legacy_scoring,
            num_workers=gpu_workers,
            mod_threshold=mod_threshold,
            show_progress=(threads <= 1),
            allowed_native_reads_path=allowed_native_reads_path,
            allowed_ivt_reads_path=allowed_ivt_reads_path,
        )

        if not passed_contigs:
            logger.info("  All contigs already on disk — no workers dispatched.")
        elif threads > 1:
            logger.info("  Using %d parallel workers (spawn context)", threads)
            ctx = mp.get_context('spawn')
            with ProcessPoolExecutor(max_workers=threads, mp_context=ctx) as executor:
                futures = {
                    executor.submit(
                        _process_contig_streaming,
                        contig=contig,
                        contig_idx=idx,
                        total_contigs=len(passed_contigs),
                        cuda_device=device_for_worker[(idx - 1) % len(device_for_worker)] if device_for_worker else 0,
                        **worker_kwargs,
                    ): contig
                    for idx, contig in enumerate(passed_contigs, 1)
                }
                with tqdm(
                    total=len(passed_contigs),
                    desc="Pipeline",
                    unit="contig",
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} contigs [{elapsed}<{remaining}] {postfix}",
                ) as pbar:
                    for future in as_completed(futures):
                        contig = futures[future]
                        try:
                            summary = future.result()
                        except Exception:
                            logger.exception("Worker failed for contig %s", contig)
                            failed.append(contig)
                            pbar.update(1)
                            continue
                        summaries.append(summary)
                        pbar.set_postfix_str(
                            f"{summary.contig_name} ({summary.n_sites} sites)"
                        )
                        pbar.update(1)
                if failed:
                    logger.error("%d contig(s) failed: %s", len(failed), ", ".join(failed))
        else:
            for contig_idx, contig in enumerate(passed_contigs, 1):
                try:
                    summary = _process_contig_streaming(
                        contig=contig,
                        contig_idx=contig_idx,
                        total_contigs=len(passed_contigs),
                        cuda_device=device_for_worker[0] if device_for_worker else 0,
                        **worker_kwargs,
                    )
                except Exception:
                    logger.exception("Worker failed for contig %s", contig)
                    failed.append(contig)
                    continue
                summaries.append(summary)
            if failed:
                logger.error("%d contig(s) failed: %s", len(failed), ", ".join(failed))
    finally:
        if cleanup_temp and tmp_root.exists():
            shutil.rmtree(tmp_root, ignore_errors=True)
        # Clean up intersection set files (only present when read_intersection=True).
        if allowed_native_reads_path is not None:
            inter_parent = Path(allowed_native_reads_path).parent
            if inter_parent.name.startswith("baleen-intersection-") and inter_parent.exists():
                shutil.rmtree(inter_parent, ignore_errors=True)

    # ---- Final merge step ----
    sorted_summaries = sorted(summaries, key=lambda s: s.contig_name)

    merge_contig_tsvs(
        [s.tsv_path for s in sorted_summaries],
        site_tsv_path,
    )

    if write_bam:
        bam_inputs = [s.bam_path for s in sorted_summaries if s.bam_path is not None]
        if bam_inputs:
            merge_contig_bams(bam_inputs, final_bam_path, threads=max(threads, 1))
        else:
            final_bam_path = None

    n_total_sites = sum(s.n_sites for s in sorted_summaries)
    n_significant = sum(s.n_significant for s in sorted_summaries)
    total_positions = sum(s.n_positions for s in sorted_summaries)

    if not keep_intermediate:
        shutil.rmtree(per_contig_dir, ignore_errors=True)
        per_contig_dir_out: Optional[Path] = None
    else:
        per_contig_dir_out = per_contig_dir

    step5_elapsed = _fmt_elapsed(time.perf_counter() - step5_t0)
    logger.info("[Step 5/5] Streaming complete (%s)", step5_elapsed)
    pipeline_elapsed = _fmt_elapsed(time.perf_counter() - pipeline_t0)
    logger.info("=" * 60)
    logger.info(
        "Streaming pipeline complete: %d contigs, %d positions, %d sites, %s",
        len(sorted_summaries), total_positions, n_total_sites, pipeline_elapsed,
    )
    logger.info("=" * 60)

    output_paths: dict[str, Any] = {
        "site_tsv": site_tsv_path,
        "read_bam": final_bam_path,
        "per_contig_dir": per_contig_dir_out,
        "n_total_sites": n_total_sites,
        "n_significant": n_significant,
    }

    if output_dir_temp is not None:
        # Caller did not pass an output_dir — clean up everything we wrote
        # and null the paths in the returned dict so callers do not chase
        # references to deleted files.  If the caller needs persistence,
        # they should pass output_dir.
        shutil.rmtree(output_dir_temp, ignore_errors=True)
        output_paths["site_tsv"] = None
        output_paths["read_bam"] = None
        output_paths["per_contig_dir"] = None

    return output_paths, metadata

run_pipeline

run_pipeline(
    native_bam: PathLike,
    native_fastq: PathLike,
    native_blow5: PathLike,
    ivt_bam: PathLike,
    ivt_fastq: PathLike,
    ivt_blow5: PathLike,
    ref_fasta: PathLike,
    *,
    min_depth: int = 15,
    depth_mode: Literal[
        "mean_coverage", "read_count"
    ] = "read_count",
    use_cuda: Optional[bool] = None,
    cuda_devices: Optional[list[int]] = None,
    padding: int = 1,
    output_dir: Optional[PathLike] = None,
    cleanup_temp: bool = True,
    rna: bool = True,
    kmer_model: Optional[str] = None,
    pore: str = DEFAULT_PORE,
    krill_hmm: bool = False,
    min_mapq: int = 20,
    primary_only: bool = True,
    threads: int = 1,
    num_cuda_streams: int = 16,
    gpu_memory_limit: Optional[int] = None,
    subsample: bool = True,
    subsample_n: int = 300
) -> tuple[dict[str, ContigResult], PipelineMetadata]

Source code in baleen/eventalign/_pipeline.py

def run_pipeline(
    native_bam: PathLike,
    native_fastq: PathLike,
    native_blow5: PathLike,
    ivt_bam: PathLike,
    ivt_fastq: PathLike,
    ivt_blow5: PathLike,
    ref_fasta: PathLike,
    *,
    min_depth: int = 15,
    depth_mode: Literal["mean_coverage", "read_count"] = "read_count",
    use_cuda: Optional[bool] = None,
    cuda_devices: Optional[list[int]] = None,
    padding: int = 1,
    output_dir: Optional[PathLike] = None,
    cleanup_temp: bool = True,
    rna: bool = True,
    kmer_model: Optional[str] = None,
    pore: str = _eventalign.DEFAULT_PORE,
    krill_hmm: bool = False,
    min_mapq: int = 20,
    primary_only: bool = True,
    threads: int = 1,
    num_cuda_streams: int = 16,
    gpu_memory_limit: Optional[int] = None,
    subsample: bool = True,
    subsample_n: int = 300,
) -> tuple[dict[str, ContigResult], PipelineMetadata]:
    pipeline_t0 = time.perf_counter()
    logger.info("=" * 60)
    logger.info("Starting baleen eventalign pipeline")
    logger.info("  native_bam:   %s", native_bam)
    logger.info("  native_fastq: %s", native_fastq)
    logger.info("  native_blow5: %s", native_blow5)
    logger.info("  ivt_bam:      %s", ivt_bam)
    logger.info("  ivt_fastq:    %s", ivt_fastq)
    logger.info("  ivt_blow5:    %s", ivt_blow5)
    logger.info("  ref_fasta:    %s", ref_fasta)
    logger.info("  min_depth=%d  depth_mode=%s  use_cuda=%s  rna=%s  padding=%d  threads=%d",
                min_depth, depth_mode, use_cuda, rna, padding, threads)
    logger.info("  min_mapq=%d  primary_only=%s  cuda_streams=%d",
                min_mapq, primary_only, num_cuda_streams)
    logger.info("  subsample=%s  subsample_n=%d  gpu_memory_limit=%s",
                subsample, subsample_n, gpu_memory_limit)
    logger.info("  cleanup_temp=%s  kmer_model=%s  pore=%s  krill_hmm=%s",
                cleanup_temp, kmer_model, pore, krill_hmm)
    logger.info("  DTW backend:  %s  (GPU=%s)",
                _dtw.backend(), _dtw.CUDA_AVAILABLE)
    logger.info("=" * 60)

    # Validate threads parameter
    if threads < 1:
        raise ValueError(f"threads must be >= 1, got {threads}")

    # Resolve cuda_devices from legacy use_cuda if needed
    if cuda_devices is None and use_cuda is not None:
        if use_cuda is True:
            cuda_devices = None  # auto-detect all GPUs
        elif use_cuda is False:
            cuda_devices = []  # CPU mode
    # Derive use_cuda bool for backward compat in internal code
    if cuda_devices is not None:
        use_cuda = len(cuda_devices) > 0 if cuda_devices else False
    # else: use_cuda stays None (auto-detect)

    native_bam = Path(native_bam)
    native_fastq = Path(native_fastq)
    native_blow5 = Path(native_blow5)
    ivt_bam = Path(ivt_bam)
    ivt_fastq = Path(ivt_fastq)
    ivt_blow5 = Path(ivt_blow5)
    ref_fasta = Path(ref_fasta)

    # ---- Step 1: krill engine check ----
    logger.info("[Step 1/6] Checking krill availability...")
    eventalign_version = _eventalign.check_krill()
    logger.info("[Step 1/6] krill version %s OK", eventalign_version)

    # ---- Step 2: Indexing ----
    logger.info("[Step 2/6] Indexing BLOW5 files...")
    step_t0 = time.perf_counter()
    logger.info("  Indexing native BLOW5...")
    _eventalign.index_blow5(native_blow5)
    logger.info("  Indexing IVT BLOW5...")
    _eventalign.index_blow5(ivt_blow5)
    logger.info("[Step 2/6] Indexing complete (%s)", _fmt_elapsed(time.perf_counter() - step_t0))

    # ---- Step 3: BAM validation & contig stats ----
    logger.info("[Step 3/6] Validating BAMs and computing contig statistics...")
    step_t0 = time.perf_counter()
    _bam.validate_bam(native_bam)
    _bam.validate_bam(ivt_bam)

    logger.info("  Computing native BAM contig stats...")
    native_stats = _bam.get_contig_stats(
        native_bam,
        min_mapq=min_mapq,
        primary_only=primary_only,
        _validated=True,
    )
    logger.info("  Computing IVT BAM contig stats...")
    ivt_stats = _bam.get_contig_stats(
        ivt_bam,
        min_mapq=min_mapq,
        primary_only=primary_only,
        _validated=True,
    )
    logger.info("[Step 3/6] BAM stats complete: %d native contigs, %d IVT contigs (%s)",
                len(native_stats), len(ivt_stats), _fmt_elapsed(time.perf_counter() - step_t0))

    # ---- Step 4: Contig filtering ----
    logger.info("[Step 4/6] Filtering contigs (min_depth=%d, depth_mode=%s)...",
                min_depth, depth_mode)
    passed_contigs, filter_results = _bam.filter_contigs(
        native_stats,
        ivt_stats,
        min_depth=float(min_depth),
        depth_mode=depth_mode,
    )
    logger.info("[Step 4/6] %d/%d contigs passed filtering",
                len(passed_contigs), len(filter_results))
    for fr in filter_results:
        if fr.passed:
            logger.debug("  PASS: %s (native_depth=%.1f, ivt_depth=%.1f)",
                         fr.contig,
                         fr.native_stats.mean_depth if fr.native_stats else 0,
                         fr.ivt_stats.mean_depth if fr.ivt_stats else 0)
        else:
            logger.info("  SKIP: %s — %s", fr.contig, fr.reason.value)

    metadata = PipelineMetadata(
        eventalign_version=eventalign_version,
        min_depth=min_depth,
        use_cuda=use_cuda,
        padding=padding,
        n_contigs_total=len(filter_results),
        n_contigs_passed_filter=len(passed_contigs),
        n_contigs_skipped=len(filter_results) - len(passed_contigs),
        filter_results=filter_results,
    )

    results: dict[str, ContigResult] = {}

    if not passed_contigs:
        logger.warning("[Step 5/6] No contigs passed filtering; returning empty results.")
        if output_dir is not None:
            _ = save_results(results, metadata, Path(output_dir) / "pipeline_results.pkl")
        elapsed = _fmt_elapsed(time.perf_counter() - pipeline_t0)
        logger.info("Pipeline finished (no results) in %s", elapsed)
        return results, metadata

    # ---- Step 5: Per-contig eventalign + signal extraction + DTW ----
    logger.info("[Step 5/6] Processing %d contigs (eventalign → signals → DTW)...",
                len(passed_contigs))
    tmp_root = Path(tempfile.mkdtemp(prefix="baleen-eventalign-"))
    logger.debug("  Temporary directory: %s", tmp_root)

    gpu_mems = _get_gpu_memory(cuda_devices) if gpu_memory_limit is None else [gpu_memory_limit]
    gpu_workers, device_for_worker = _gpu_concurrent_workers(threads, gpu_mems, cuda_devices)

    try:
        if threads > 1:
            # Parallel processing with multiprocessing
            logger.info("  Using %d parallel workers (spawn context)", threads)
            ctx = mp.get_context('spawn')
            with ProcessPoolExecutor(max_workers=threads, mp_context=ctx) as executor:
                futures = {
                    executor.submit(
                        _process_contig,
                        contig=contig,
                        contig_idx=idx,
                        total_contigs=len(passed_contigs),
                        native_bam=native_bam,
                        native_fastq=native_fastq,
                        native_blow5=native_blow5,
                        ivt_bam=ivt_bam,
                        ivt_fastq=ivt_fastq,
                        ivt_blow5=ivt_blow5,
                        ref_fasta=ref_fasta,
                        native_stats=native_stats,
                        ivt_stats=ivt_stats,
                        tmp_root=tmp_root,
                        use_cuda=use_cuda,
                        padding=padding,
                        rna=rna,
                        kmer_model=kmer_model,
                        pore=pore,
                        krill_hmm=krill_hmm,
                        min_mapq=min_mapq,
                        primary_only=primary_only,
                        cleanup_temp=cleanup_temp,
                        num_cuda_streams=num_cuda_streams,
                        subsample=subsample,
                        subsample_n=subsample_n,
                        gpu_memory_bytes=gpu_mems[0] if gpu_mems else 8 * 1024 ** 3,
                        num_workers=gpu_workers,
                        show_progress=False,
                        cuda_device=device_for_worker[idx - 1] if device_for_worker else 0,
                    ): contig
                    for idx, contig in enumerate(passed_contigs, 1)
                }
                failed = []
                with tqdm(
                    total=len(passed_contigs),
                    desc="Pipeline",
                    unit="contig",
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} contigs [{elapsed}<{remaining}] {postfix}",
                ) as pbar:
                    for future in as_completed(futures):
                        contig = futures[future]
                        try:
                            contig_name, contig_result = future.result()
                        except Exception:
                            logger.exception("Worker failed for contig %s", contig)
                            failed.append(contig)
                            pbar.update(1)
                            continue
                        results[contig_name] = contig_result
                        n_pos = len(contig_result.positions)
                        pbar.set_postfix_str(f"{contig_name} ({n_pos} pos)")
                        pbar.update(1)
                if failed:
                    logger.error("%d contig(s) failed: %s", len(failed), ", ".join(failed))
        else:
            # Sequential processing (original behavior)
            for contig_idx, contig in enumerate(passed_contigs, 1):
                contig_name, contig_result = _process_contig(
                    contig=contig,
                    contig_idx=contig_idx,
                    total_contigs=len(passed_contigs),
                    native_bam=native_bam,
                    native_fastq=native_fastq,
                    native_blow5=native_blow5,
                    ivt_bam=ivt_bam,
                    ivt_fastq=ivt_fastq,
                    ivt_blow5=ivt_blow5,
                    ref_fasta=ref_fasta,
                    native_stats=native_stats,
                    ivt_stats=ivt_stats,
                    tmp_root=tmp_root,
                    use_cuda=use_cuda,
                    padding=padding,
                    rna=rna,
                    kmer_model=kmer_model,
                    pore=pore,
                    krill_hmm=krill_hmm,
                    min_mapq=min_mapq,
                    primary_only=primary_only,
                    cleanup_temp=cleanup_temp,
                    num_cuda_streams=num_cuda_streams,
                    subsample=subsample,
                    subsample_n=subsample_n,
                    gpu_memory_bytes=gpu_mems[0] if gpu_mems else 8 * 1024 ** 3,
                    cuda_device=device_for_worker[0] if device_for_worker else 0,
                )
                results[contig_name] = contig_result
    finally:
        if cleanup_temp and tmp_root.exists():
            shutil.rmtree(tmp_root, ignore_errors=True)

    # ---- Step 6: Save results ----
    if output_dir is not None:
        logger.info("[Step 6/6] Saving results to %s ...", output_dir)
        _ = save_results(results, metadata, Path(output_dir) / "pipeline_results.pkl")
    else:
        logger.info("[Step 6/6] No output_dir specified; results returned in memory only")

    total_positions = sum(len(cr.positions) for cr in results.values())
    pipeline_elapsed = _fmt_elapsed(time.perf_counter() - pipeline_t0)
    logger.info("=" * 60)
    logger.info("Pipeline complete: %d contigs, %d positions, %s",
                len(results), total_positions, pipeline_elapsed)
    logger.info("=" * 60)

    return results, metadata

Results I/O

save_results

save_results(
    results: dict[str, ContigResult],
    metadata: PipelineMetadata,
    output_path: PathLike,
) -> Path

Source code in baleen/eventalign/_pipeline.py

def save_results(
    results: dict[str, ContigResult],
    metadata: PipelineMetadata,
    output_path: PathLike,
) -> Path:
    out_path = Path(output_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("wb") as handle:
        pickle.dump({"results": results, "metadata": metadata}, handle)
    logger.info("Saved pipeline results to %s", out_path)
    return out_path

load_results

load_results(
    output_path: PathLike,
) -> tuple[dict[str, ContigResult], PipelineMetadata]

Source code in baleen/eventalign/_pipeline.py

def load_results(output_path: PathLike) -> tuple[dict[str, ContigResult], PipelineMetadata]:
    in_path = Path(output_path)
    with in_path.open("rb") as handle:
        payload = cast(_SerializedPayload, pickle.load(handle))
    return payload["results"], payload["metadata"]

Data classes

PositionResult dataclass

PositionResult(
    position: int,
    reference_kmer: str,
    n_native_reads: int,
    n_ivt_reads: int,
    native_read_names: list[str],
    ivt_read_names: list[str],
    distance_matrix: NDArray[float64],
)

ContigResult dataclass

ContigResult(
    contig: str,
    native_depth: float,
    ivt_depth: float,
    positions: dict[int, PositionResult],
)

ContigSummary dataclass

ContigSummary(
    contig_name: str,
    n_sites: int,
    n_positions: int,
    n_significant: int,
    tsv_path: Path,
    bam_path: Optional[Path] = None,
)

Lightweight per-contig outcome returned by the streaming worker.

Holds counts and on-disk paths to the per-contig TSV/BAM slices — no per-read arrays — so the main process memory stays O(N_contigs).

PipelineMetadata dataclass

PipelineMetadata(
    eventalign_version: str,
    min_depth: int,
    use_cuda: Optional[bool],
    padding: int,
    n_contigs_total: int,
    n_contigs_passed_filter: int,
    n_contigs_skipped: int,
    filter_results: list[ContigFilterResult],
)

ContigStats dataclass

ContigStats(
    contig: str, mapped_reads: int, mean_depth: float
)

Per-contig mapping statistics.

Parameters:

Name	Type	Description	Default
contig	str	Contig/transcript identifier.	required
mapped_reads	int	Number of mapped reads contributing to the contig.	required
mean_depth	float	Mean base-level depth over the full contig length.	required

ContigFilterResult dataclass

ContigFilterResult(
    contig: str,
    passed: bool,
    reason: FilterReason,
    native_stats: Optional[ContigStats] = None,
    ivt_stats: Optional[ContigStats] = None,
)

Outcome of filtering a contig between native and IVT datasets.

Parameters:

Name	Type	Description	Default
contig	str	Contig/transcript identifier.	required
passed	bool	Whether this contig passed all filter criteria.	required
reason	FilterReason	Detailed reason for pass/fail status.	required
native_stats	ContigStats	Native BAM contig statistics when available.	None
ivt_stats	ContigStats	IVT BAM contig statistics when available.	None

FilterReason

Bases: Enum

Reason describing why a contig passed or failed filtering.

PositionSignals dataclass

PositionSignals(
    contig: str,
    position: int,
    reference_kmer: str,
    read_signals: dict[str, NDArray[float32]],
    read_names: list[str] = list(),
)

ContigModificationResult dataclass

ContigModificationResult(
    contig: str,
    position_stats: dict[int, PositionStats],
    native_trajectories: list[ReadTrajectory],
    ivt_trajectories: list[ReadTrajectory],
    global_mu: float,
    global_sigma: float,
)

Full output of the hierarchical pipeline for one contig.

position_stats instance-attribute

position_stats: dict[int, PositionStats]

Keyed by genomic position.

CoverageClass

Bases: str, Enum

IVT coverage quality at a position.