From db744dbe81711325fae972588c63042c428b4ca2 Mon Sep 17 00:00:00 2001 From: Sam Nicholls Date: Mon, 18 May 2026 12:29:48 +0000 Subject: [PATCH] Scatter-gather bambu quant [CW-7201] --- bin/workflow_glue/report.py | 77 +- bin/workflow_glue/tests/common/test_report.py | 126 ++ bin/workflow_glue_r/R/bambu.R | 1073 ++++++++++++++--- .../tests/testthat/test_bambu.R | 824 +++++++++++-- modules/local/bambu_chunked.nf | 120 ++ subworkflows/transcriptome.nf | 290 +++-- 6 files changed, 2137 insertions(+), 373 deletions(-) create mode 100644 modules/local/bambu_chunked.nf diff --git a/bin/workflow_glue/report.py b/bin/workflow_glue/report.py index af40af4..cff78e0 100644 --- a/bin/workflow_glue/report.py +++ b/bin/workflow_glue/report.py @@ -1,6 +1,7 @@ """Create workflow report for wf-transcriptomes.""" import json +import math from pathlib import Path from dominate.tags import div, h3, p, pre, strong @@ -25,6 +26,35 @@ def _read_table(path, **kwargs): return pd.read_csv(path, sep="\t", **kwargs) +def _coerce_float(value): + """Return a finite float when possible, otherwise None.""" + if value in (None, "", "N/A", "NA", "nan", "NaN"): + return None + try: + numeric = float(value) + except (TypeError, ValueError): + return None + if not math.isfinite(numeric): + return None + return numeric + + +def _format_count_value(value): + """Format count-like values for the report, tolerating NA-like strings.""" + numeric = _coerce_float(value) + if numeric is None: + return "N/A" + return format(round(numeric), ",") + + +def _format_ratio_value(value): + """Format ratio values for the report, tolerating NA-like strings.""" + numeric = _coerce_float(value) + if numeric is None: + return "N/A" + return f"{numeric:.2f}x" + + def _cohort_summary(cohort_dir): tx_meta = _read_table(Path(cohort_dir) / "transcript_metadata.tsv") if tx_meta is None: @@ -262,12 +292,23 @@ def main(args): stats = stats[0] flagstats = flagstats[0] sample_names = sample_names[0] if sample_names else None - fastcat.SeqSummary( - stats, - flagstat=flagstats, - sample_names=sample_names, - alignment_stats=True, - ) + try: + fastcat.SeqSummary( + stats, + flagstat=flagstats, + sample_names=sample_names, + alignment_stats=True, + ) + except Exception as exc: # pragma: no cover - defensive + logger.warning("Skipping read summary plot: %s", exc) + _create_warning_banner( + ( + "Read summary plots could not be rendered for this run. " + "This can happen for degenerate or extremely " + "small input statistics." + ), + level="info", + ) with report.add_section("Sample metadata", "Samples"): tabs = Tabs() @@ -389,33 +430,30 @@ def main(args): ( "Median library size", "{} reads".format( - format( - bambu_qc.get("median_library_size", 0), - ",", + _format_count_value( + bambu_qc.get("median_library_size", 0) ) ), ), ( "Min library size", "{} reads".format( - format( - bambu_qc.get("min_library_size", 0), - ",", + _format_count_value( + bambu_qc.get("min_library_size", 0) ) ), ), ( "Max library size", "{} reads".format( - format( - bambu_qc.get("max_library_size", 0), - ",", + _format_count_value( + bambu_qc.get("max_library_size", 0) ) ), ), ( "Library size ratio (max/min)", - "{:.2f}x".format( + _format_ratio_value( bambu_qc.get("library_size_ratio", 1.0) ), ), @@ -463,11 +501,14 @@ def main(args): with h3("Per-Sample Library Sizes"): lib_size_data = [] for sample, size in bambu_qc["library_sizes"].items(): + numeric_size = _coerce_float(size) lib_size_data.append( { "Sample": sample, - "Library Size": format(size, ","), - "Reads": size, + "Library Size": _format_count_value(size), + "Reads": ( + numeric_size if numeric_size is not None else -1 + ), } ) lib_df = pd.DataFrame(lib_size_data).sort_values( diff --git a/bin/workflow_glue/tests/common/test_report.py b/bin/workflow_glue/tests/common/test_report.py index 5dfa798..8267b81 100644 --- a/bin/workflow_glue/tests/common/test_report.py +++ b/bin/workflow_glue/tests/common/test_report.py @@ -217,6 +217,132 @@ def test_report_main_accepts_optional_file_sentinels(monkeypatch, tmp_path): assert any("GRCh38" in table.to_string() for table in tables) +def test_report_main_handles_degenerate_bambu_qc_and_read_summary( + monkeypatch, + tmp_path, +): + """Tiny/empty stats should not crash the report rendering path.""" + tables = [] + banners = [] + + monkeypatch.setattr(report.labs, "LabsReport", _FakeReport) + monkeypatch.setattr(report, "Tabs", _FakeTabs) + monkeypatch.setattr(report, "p", lambda *args, **kwargs: None) + monkeypatch.setattr(report, "pre", lambda *args, **kwargs: None) + monkeypatch.setattr( + report.fastcat, + "SeqSummary", + lambda *args, **kwargs: (_ for _ in ()).throw(KeyError(1)), + ) + monkeypatch.setattr( + report, + "_create_warning_banner", + lambda message, level="warning": banners.append((level, message)), + ) + monkeypatch.setattr( + report.DataTable, + "from_pandas", + staticmethod(lambda table, *args, **kwargs: tables.append(table.copy())), + ) + + metadata = _write( + tmp_path / "metadata.json", + json.dumps([{"alias": "sampleA", "has_stats": True}]), + ) + params = _write(tmp_path / "params.json", "{}") + versions = tmp_path / "versions" + versions.mkdir() + _write(versions / "versions.txt", "tool,1.0\n") + + cohort = tmp_path / "cohort" + cohort.mkdir() + reference = cohort / "reference" + reference.mkdir() + _write( + reference / "annotation_reference_summary.json", + json.dumps( + { + "seqname_overlap": ["chr1"], + "only_in_annotation": [], + "only_in_reference": [], + "annotation": { + "kept_records": 10, + "excluded_unstranded_records": 0, + "sanitised_attribute_records": 0, + }, + "warnings": [], + } + ), + ) + _write( + cohort / "bambu_qc_stats.json", + json.dumps( + { + "samples": 1, + "library_sizes": {"sampleA": 0}, + "min_library_size": 0, + "max_library_size": 0, + "median_library_size": 0, + "library_size_ratio": "NA", + "total_transcripts_before_filter": 0, + "total_transcripts_after_filter": 0, + "transcripts_filtered": 0, + "median_transcripts_detected": 0, + "total_genes_after_filter": 0, + "transcriptome_mode": "discover", + "ndr_used": "automatic", + } + ), + ) + + samples = tmp_path / "samples" + samples.mkdir() + (samples / "OPTIONAL_FILE").touch() + + sqanti = tmp_path / "sqanti" + sqanti.mkdir() + (sqanti / "OPTIONAL_FILE").touch() + + alignment_stats = tmp_path / "alignment_stats" + alignment_stats.mkdir() + + out_report = tmp_path / "wf-transcriptomes-report.html" + args = report.argparser().parse_args( + [ + str(out_report), + "--metadata", + str(metadata), + "--alignment_stats_dir", + str(alignment_stats), + "--stats", + str(alignment_stats), + "--cohort_dir", + str(cohort), + "--samples_dir", + str(samples), + "--sqanti_dir", + str(sqanti), + "--versions", + str(versions), + "--params", + str(params), + ] + ) + + report.main(args) + + assert out_report.exists() + assert any( + level == "info" and "Read summary plots could not be rendered" in message + for level, message in banners + ) + assert any( + "Library size ratio (max/min)" in table.to_string() + and "N/A" in table.to_string() + for table in tables + ) + + def test_report_main_renders_statistical_methods_and_warnings( monkeypatch, tmp_path, diff --git a/bin/workflow_glue_r/R/bambu.R b/bin/workflow_glue_r/R/bambu.R index 72da7e3..5e4d3f2 100644 --- a/bin/workflow_glue_r/R/bambu.R +++ b/bin/workflow_glue_r/R/bambu.R @@ -1,18 +1,24 @@ bambu_arg_spec <- function() { list( + list( + name = "mode", + flag = "--mode", + help = "Run mode: discover, quant, collate, or empty.", + type = "character", + required = TRUE, + choices = c("discover", "quant", "collate", "empty") + ), list( name = "bams", flag = "--bams", help = "Comma-separated BAM paths.", - type = "character", - required = TRUE + type = "character" ), list( name = "aliases", flag = "--aliases", help = "Comma-separated aliases for --bams.", - type = "character", - required = TRUE + type = "character" ), list( name = "sample_sheet", @@ -24,15 +30,31 @@ bambu_arg_spec <- function() { name = "annotation", flag = "--annotation", help = "Reference annotation GTF/GFF.", - type = "character", - required = TRUE + type = "character" + ), + list( + name = "chunk_rds", + flag = "--chunk_rds", + help = "Path to a chunked rcFile bundle RDS for quant mode.", + type = "character" + ), + list( + name = "discovered_annotation_rds", + flag = "--discovered_annotation_rds", + help = "Path to a discovered annotation RDS for quant mode.", + type = "character" + ), + list( + name = "chunk_dirs", + flag = "--chunk_dirs", + help = "Comma-separated chunk quantification directories for collate mode.", + type = "character" ), list( name = "genome", flag = "--genome", help = "Reference genome FASTA.", - type = "character", - required = TRUE + type = "character" ), list( name = "transcriptome_mode", @@ -76,9 +98,79 @@ bambu_arg_parser <- function() { ) } -bambu_resolve_inputs <- function(args) { +bambu_default_yield_size <- 250000L +bambu_default_seed <- 42L + +bambu_missing <- function(value) { + is.null(value) || + length(value) == 0 || + (length(value) == 1 && is.na(value)) || + (is.character(value) && length(value) == 1 && !nzchar(value)) +} + +bambu_validate_args <- function(args) { + mode <- args$mode + if (bambu_missing(mode)) { + stop("Missing required arguments: --mode", call. = FALSE) + } + if (!mode %in% c("discover", "quant", "collate", "empty")) { + stop("mode must be one of: discover, quant, collate, empty", call. = FALSE) + } + if (bambu_missing(args$out_dir)) { + stop("Missing required arguments: --out_dir", call. = FALSE) + } + + if (identical(mode, "discover")) { + missing_args <- character(0) + if (bambu_missing(args$bams)) missing_args <- c(missing_args, "--bams") + if (bambu_missing(args$aliases)) missing_args <- c(missing_args, "--aliases") + if (bambu_missing(args$annotation)) missing_args <- c(missing_args, "--annotation") + if (bambu_missing(args$genome)) missing_args <- c(missing_args, "--genome") + if (length(missing_args) > 0) { + stop( + sprintf("Missing required arguments: %s", paste(missing_args, collapse = ", ")), + call. = FALSE + ) + } + } + + if (identical(mode, "quant")) { + missing_args <- character(0) + if (bambu_missing(args$chunk_rds)) missing_args <- c(missing_args, "--chunk_rds") + if (bambu_missing(args$discovered_annotation_rds)) { + missing_args <- c(missing_args, "--discovered_annotation_rds") + } + if (bambu_missing(args$genome)) missing_args <- c(missing_args, "--genome") + if (length(missing_args) > 0) { + stop( + sprintf("Missing required arguments: %s", paste(missing_args, collapse = ", ")), + call. = FALSE + ) + } + } + + if (identical(mode, "collate") && bambu_missing(args$chunk_dirs)) { + stop("Missing required arguments: --chunk_dirs", call. = FALSE) + } + + if (identical(mode, "empty") && bambu_missing(args$aliases)) { + stop("Missing required arguments: --aliases", call. = FALSE) + } + + invisible(args) +} + +bambu_resolve_chunk_dirs <- function(args) { + chunk_dirs <- workflow_glue_r_parse_csv_list(args$chunk_dirs) + if (length(chunk_dirs) < 1) { + stop("No chunk quantification directories were provided in --chunk_dirs.", call. = FALSE) + } + chunk_dirs +} + +bambu_resolve_inputs <- function(args, bamfile_list_ctor = Rsamtools::BamFileList) { sample_df <- NULL - if (!is.null(args$sample_sheet)) { + if (!bambu_missing(args$sample_sheet)) { sample_df <- utils::read.csv( args$sample_sheet, check.names = FALSE, @@ -139,7 +231,10 @@ bambu_resolve_inputs <- function(args) { reads <- if (length(bam_paths) == 1) { bam_paths } else { - Rsamtools::BamFileList(bam_paths, yieldSize = 250000L) + bamfile_list_ctor(bam_paths, yieldSize = bambu_default_yield_size) + } + if (is.list(reads) && length(bam_paths) > 1) { + names(reads) <- aliases } list( @@ -154,28 +249,51 @@ bambu_discovery_enabled <- function(args) { identical(args$transcriptome_mode, "discover") } -bambu_build_args <- function(args, reads, annotation_obj) { +bambu_build_args <- function(args, reads, annotation_obj, discovery, quant) { bambu_args <- list( reads = reads, annotations = annotation_obj, genome = args$genome, - ncore = args$threads, - discovery = bambu_discovery_enabled(args), + ncore = as.integer(args$threads), + discovery = discovery, + quant = quant, lowMemory = TRUE, - yieldSize = 250000L, + yieldSize = bambu_default_yield_size, verbose = TRUE ) - if (bambu_discovery_enabled(args) && !is.null(args$ndr)) { + if (discovery && !is.null(args$ndr)) { bambu_args$NDR <- args$ndr } + if (quant) { + # Chunked quant re-estimates degradation bias per chunk, which changes + # the EM inputs and breaks equivalence with unchunked bambu quant. + bambu_args$opt.em <- list(degradationBias = FALSE) + } + bambu_args } +bambu_message_ndr <- function(args) { + if (!is.null(args$ndr)) { + message(sprintf("Using user-specified NDR = %.3f", args$ndr)) + } else { + message("Using bambu automatic NDR selection.") + } + message("Novel Discovery Rate (NDR) controls transcript discovery stringency:") + message(" Lower NDR (e.g., 0.05) = fewer false positive transcripts, may miss real ones") + message(" Higher NDR (e.g., 0.2) = more sensitive discovery, more false positives") + if (is.null(args$ndr)) { + message(" Current NDR = automatic (selected by bambu from the data)") + } else { + message(sprintf(" Current NDR = %.3f balances precision and recall", args$ndr)) + } +} + bambu_effective_threads <- function(args, bam_count) { # bambu's low-memory mode can have issues with multiple BAMs and - # parallel threads due to BiocFileCache writes, + # parallel threads due to BiocFileCache writes, # so we enforce single-threading in that case. threads <- as.integer(args$threads) if (bam_count > 1 && threads > 1L) { @@ -191,6 +309,459 @@ bambu_effective_threads <- function(args, bam_count) { threads } +bambu_normalise_rc_file_list <- function(rc_files, aliases = NULL) { + if (!is.list(rc_files)) { + rc_files <- list(rc_files) + } + if (!is.null(aliases) && is.null(names(rc_files)) && length(rc_files) == length(aliases)) { + names(rc_files) <- aliases + } + rc_files +} + +bambu_chunk_id_for_seqname <- function(seqname) { + seqname <- as.character(seqname) + seqname <- gsub("[^A-Za-z0-9._-]+", "_", seqname) + seqname <- sub("^_+", "", seqname) + seqname <- sub("_+$", "", seqname) + if (!nzchar(seqname)) { + seqname <- "chunk" + } + seqname +} + +bambu_chunk_rc_files <- function(rc_files, aliases, sample_df) { + rc_files <- bambu_normalise_rc_file_list(rc_files, aliases = aliases) + seqnames <- unique(unlist(lapply(rc_files, function(rcf) { + as.character(SummarizedExperiment::rowData(rcf)$chr.rc) + }))) + seqnames <- seqnames[!is.na(seqnames) & nzchar(seqnames)] + chunk_ids <- make.unique(vapply(seqnames, bambu_chunk_id_for_seqname, character(1)), sep = "_") + + Map(function(seqname, chunk_id) { + rc_chunk <- lapply(rc_files, function(rcf) { + idx <- as.character(SummarizedExperiment::rowData(rcf)$chr.rc) == seqname + rcf[idx, , drop = FALSE] + }) + if (!is.null(names(rc_files))) { + names(rc_chunk) <- names(rc_files) + } + list( + chunk_id = chunk_id, + seqname = seqname, + aliases = aliases, + sample_df = sample_df, + rc_files = rc_chunk + ) + }, seqnames, chunk_ids) +} + +bambu_write_discovery_outputs <- function(out_dir, rc_files, discovered_annotations, chunk_bundles, sample_df) { + saveRDS(rc_files, file.path(out_dir, "bambu_rcfiles.rds")) + saveRDS(discovered_annotations, file.path(out_dir, "bambu_discovered_annotations.rds")) + utils::write.csv( + sample_df, + file.path(out_dir, "samples.csv"), + row.names = FALSE, + quote = FALSE + ) + + annotation_tx_counts <- bambu_annotation_tx_counts_by_seqname(discovered_annotations) + chunks_dir <- file.path(out_dir, "chunks") + dir.create(chunks_dir, showWarnings = FALSE, recursive = TRUE) + + manifest <- do.call(rbind, lapply(chunk_bundles, function(bundle) { + annotation_tx_count <- unname(annotation_tx_counts[bundle$seqname]) + if (length(annotation_tx_count) != 1L || is.na(annotation_tx_count)) { + annotation_tx_count <- 0L + } + bundle$annotation_tx_count <- annotation_tx_count + chunk_path <- file.path(chunks_dir, sprintf("%s.rds", bundle$chunk_id)) + saveRDS(bundle, chunk_path) + data.frame( + chunk_id = bundle$chunk_id, + seqname = bundle$seqname, + annotation_tx_count = annotation_tx_count, + rds_path = chunk_path, + stringsAsFactors = FALSE + ) + })) + + utils::write.table( + manifest, + file = file.path(out_dir, "chunk_manifest.tsv"), + sep = "\t", + quote = FALSE, + row.names = FALSE + ) +} + +bambu_run_discover_mode <- function(args, analysis_fn, prepare_annotations_fn, bamfile_list_ctor) { + inputs <- bambu_resolve_inputs(args, bamfile_list_ctor = bamfile_list_ctor) + args$threads <- bambu_effective_threads(args, length(inputs$bam_paths)) + annotation_obj <- prepare_annotations_fn(args$annotation) + + if (length(inputs$bam_paths) > 1) { + message(sprintf("Using BamFileList yieldSize = %d", bambu_default_yield_size)) + } + message(sprintf("Running bambu discover setup with threads = %d", args$threads)) + message("Generating bambu rcFiles...") + rc_files <- bambu_call_analysis( + analysis_fn, + bambu_build_args( + args, + inputs$reads, + annotation_obj, + discovery = FALSE, + quant = FALSE + ) + ) + rc_files <- bambu_normalise_rc_file_list(rc_files, aliases = inputs$aliases) + + discovered_annotations <- if (bambu_discovery_enabled(args)) { + bambu_message_ndr(args) + message("Running global bambu discovery from rcFiles...") + bambu_call_analysis( + analysis_fn, + bambu_build_args( + args, + rc_files, + annotation_obj, + discovery = TRUE, + quant = FALSE + ) + ) + } else { + message("Fixed annotation mode: using prepared annotation without bambu discovery.") + annotation_obj + } + + chunk_bundles <- bambu_chunk_rc_files(rc_files, inputs$aliases, inputs$sample_df) + bambu_write_discovery_outputs( + args$out_dir, + rc_files, + discovered_annotations, + chunk_bundles, + inputs$sample_df + ) + + invisible( + list( + rc_files = rc_files, + discovered_annotations = discovered_annotations, + sample_df = inputs$sample_df, + chunk_bundles = chunk_bundles + ) + ) +} + +bambu_run_quant_mode <- function(args, analysis_fn) { + chunk_bundle <- readRDS(args$chunk_rds) + discovered_annotations <- readRDS(args$discovered_annotation_rds) + sample_names <- bambu_chunk_sample_names(chunk_bundle) + + annotation_tx_count <- chunk_bundle$annotation_tx_count + if (length(annotation_tx_count) != 1L || is.na(annotation_tx_count)) { + annotation_tx_count <- bambu_annotation_tx_count_for_seqname( + discovered_annotations, + chunk_bundle$seqname + ) + } + + if (!is.na(annotation_tx_count) && annotation_tx_count < 1L) { + warning( + sprintf( + paste( + "Skipping bambu quant for chunk '%s' (%s):", + "the discovered annotation contains no transcripts on this seqname." + ), + chunk_bundle$chunk_id, + chunk_bundle$seqname + ), + call. = FALSE + ) + se <- bambu_empty_quant_se( + discovered_annotations = discovered_annotations, + sample_names = sample_names + ) + } else { + se <- bambu_call_analysis( + analysis_fn, + bambu_build_args( + args, + chunk_bundle$rc_files, + discovered_annotations, + discovery = FALSE, + quant = TRUE + ) + ) + } + + if (length(sample_names) > 0 && ncol(se) == length(sample_names)) { + colnames(se) <- sample_names + } + S4Vectors::metadata(se)$incompatibleCounts <- bambu_normalise_incompatible_counts( + S4Vectors::metadata(se)$incompatibleCounts, + sample_names = colnames(se) + ) + + saveRDS(se, file.path(args$out_dir, "bambu_transcripts.rds")) + utils::write.csv( + chunk_bundle$sample_df, + file.path(args$out_dir, "samples.csv"), + row.names = FALSE, + quote = FALSE + ) + qc_stats <- list( + chunk_id = chunk_bundle$chunk_id, + seqname = chunk_bundle$seqname, + samples = ncol(se), + total_transcripts_before_filter = nrow(se), + total_genes_before_filter = length(unique(SummarizedExperiment::rowData(se)$GENEID)) + ) + jsonlite::write_json( + qc_stats, + file.path(args$out_dir, "bambu_qc_stats.json"), + pretty = TRUE, + auto_unbox = TRUE + ) + + invisible( + list( + se = se, + sample_df = chunk_bundle$sample_df, + qc_stats = qc_stats + ) + ) +} + +bambu_run_collate_mode <- function(args, gene_expression_fn, write_gtf_fn) { + chunk_dirs <- bambu_resolve_chunk_dirs(args) + invisible( + bambu_collate_chunk_outputs( + chunk_dirs, + out_dir = args$out_dir, + transcriptome_mode = args$transcriptome_mode, + ndr = args$ndr, + gene_expression_fn = gene_expression_fn, + write_gtf_fn = write_gtf_fn + ) + ) +} + +bambu_run_empty_mode <- function(args, write_gtf_fn) { + aliases <- workflow_glue_r_parse_csv_list(args$aliases) + if (length(aliases) < 1) { + stop("No sample aliases were provided in --aliases.", call. = FALSE) + } + + invisible( + bambu_write_empty_outputs( + sample_aliases = aliases, + out_dir = args$out_dir, + transcriptome_mode = args$transcriptome_mode, + ndr = args$ndr, + write_gtf_fn = write_gtf_fn + ) + ) +} + +bambu_call_analysis <- function(analysis_fn, analysis_args, seed = bambu_default_seed) { + set.seed(seed) + do.call(analysis_fn, analysis_args) +} + +bambu_chunk_sample_names <- function(chunk_bundle) { + if (!is.null(chunk_bundle$aliases)) { + return(as.character(chunk_bundle$aliases)) + } + if (!is.null(chunk_bundle$sample_df$alias)) { + return(as.character(chunk_bundle$sample_df$alias)) + } + + rc_files <- chunk_bundle$rc_files + if (!is.null(names(rc_files))) { + return(names(rc_files)) + } + + paste0("sample", seq_along(rc_files)) +} + +bambu_annotation_tx_counts_by_seqname <- function(discovered_annotations) { + if (methods::is(discovered_annotations, "GenomicRangesList")) { + unlisted_annotations <- unlist(discovered_annotations, use.names = FALSE) + if (length(unlisted_annotations) < 1L) { + return(integer()) + } + partitioning <- IRanges::PartitioningByEnd(discovered_annotations) + tx_seqnames <- as.character( + GenomeInfoDb::seqnames(unlisted_annotations)[IRanges::start(partitioning)] + ) + } else if (methods::is(discovered_annotations, "GenomicRanges")) { + tx_seqnames <- as.character(GenomeInfoDb::seqnames(discovered_annotations)) + } else { + return(integer()) + } + + tx_seqnames <- tx_seqnames[!is.na(tx_seqnames) & nzchar(tx_seqnames)] + if (length(tx_seqnames) < 1L) { + return(integer()) + } + + table(tx_seqnames) +} + +bambu_annotation_tx_count_for_seqname <- function(discovered_annotations, seqname) { + annotation_tx_counts <- bambu_annotation_tx_counts_by_seqname(discovered_annotations) + annotation_tx_count <- unname(annotation_tx_counts[seqname]) + if (length(annotation_tx_count) != 1L || is.na(annotation_tx_count)) { + return(0L) + } + as.integer(annotation_tx_count) +} + +bambu_empty_quant_se <- function(discovered_annotations, sample_names) { + empty_row_ranges <- if (methods::is(discovered_annotations, "GenomicRangesList")) { + discovered_annotations[0] + } else if (methods::is(discovered_annotations, "GenomicRanges")) { + discovered_annotations[0] + } else { + GenomicRanges::GRanges() + } + + empty_matrix <- function() { + matrix( + numeric(0), + nrow = 0, + ncol = length(sample_names), + dimnames = list(character(0), sample_names) + ) + } + + SummarizedExperiment::SummarizedExperiment( + assays = list( + counts = empty_matrix(), + CPM = empty_matrix(), + fullLengthCounts = empty_matrix(), + uniqueCounts = empty_matrix() + ), + rowRanges = empty_row_ranges, + colData = S4Vectors::DataFrame(row.names = sample_names), + metadata = list( + incompatibleCounts = bambu_empty_incompatible_counts(sample_names), + warnings = character(0) + ) + ) +} + +bambu_empty_gene_se <- function(sample_names) { + empty_matrix <- matrix( + numeric(0), + nrow = 0, + ncol = length(sample_names), + dimnames = list(character(0), sample_names) + ) + + SummarizedExperiment::SummarizedExperiment( + assays = list( + counts = empty_matrix, + CPM = empty_matrix + ), + rowData = S4Vectors::DataFrame(GENEID = character(0)), + colData = S4Vectors::DataFrame(row.names = sample_names) + ) +} + +bambu_write_empty_outputs <- function( + sample_aliases, + out_dir, + transcriptome_mode = "discover", + ndr = NULL, + write_gtf_fn = bambu::writeToGTF +) { + sample_df <- data.frame(alias = sample_aliases, stringsAsFactors = FALSE) + tx_se <- bambu_empty_quant_se( + discovered_annotations = GenomicRanges::GRangesList(), + sample_names = sample_aliases + ) + gene_se <- bambu_empty_gene_se(sample_aliases) + qc_stats <- list( + samples = length(sample_aliases), + total_transcripts_before_filter = 0L, + total_genes_before_filter = 0L, + total_transcripts_after_filter = 0L, + total_genes_after_filter = 0L, + transcripts_filtered = 0L, + chunk_count = 0L, + empty_output = TRUE + ) + qc_stats <- bambu_add_library_qc_stats(tx_se, qc_stats) + + args <- list( + out_dir = out_dir, + transcriptome_mode = transcriptome_mode, + ndr = ndr + ) + bambu_write_outputs( + tx_se, + gene_se, + sample_df, + args, + qc_stats, + write_gtf_fn = write_gtf_fn, + write_rds = TRUE + ) + + invisible( + list( + se = tx_se, + gene_se = gene_se, + sample_df = sample_df, + qc_stats = qc_stats + ) + ) +} + +bambu_empty_incompatible_counts <- function(sample_names) { + cols <- c( + list(GENEID = character(0)), + stats::setNames(rep(list(numeric(0)), length(sample_names)), sample_names) + ) + data.table::as.data.table(cols) +} + +bambu_normalise_incompatible_counts <- function(incompatible_counts, sample_names) { + if (is.null(incompatible_counts)) { + return(bambu_empty_incompatible_counts(sample_names)) + } + + incompatible_counts <- data.table::as.data.table(incompatible_counts) + if (!"GENEID" %in% names(incompatible_counts)) { + return(bambu_empty_incompatible_counts(sample_names)) + } + + value_cols <- setdiff(names(incompatible_counts), c("GENEID", "TXNAME")) + if (length(sample_names) > 0 && + !all(sample_names %in% value_cols) && + length(value_cols) == length(sample_names)) { + data.table::setnames(incompatible_counts, value_cols, sample_names) + value_cols <- sample_names + } + + for (sample_name in sample_names) { + if (!sample_name %in% names(incompatible_counts)) { + data.table::set( + incompatible_counts, + j = sample_name, + value = numeric(nrow(incompatible_counts)) + ) + } + } + + keep_cols <- c("GENEID", sample_names) + incompatible_counts[, keep_cols, with = FALSE] +} + bambu_filter_transcripts <- function(se) { assays <- SummarizedExperiment::assays(se) counts_mat <- assays$counts @@ -226,46 +797,20 @@ bambu_filter_transcripts <- function(se) { # transcriptToGeneExpression() expects incompatibleCounts GENEIDs to be a # subset of rowData(se)$GENEID after filtering. sample_names <- colnames(se) - empty_incompatible_counts <- function() { - cols <- c( - list(GENEID = character(0)), - stats::setNames(rep(list(numeric(0)), length(sample_names)), sample_names) - ) - data.table::as.data.table(cols) - } + incompatible_counts <- bambu_normalise_incompatible_counts( + S4Vectors::metadata(se)$incompatibleCounts, + sample_names = sample_names + ) + kept_genes <- unique(SummarizedExperiment::rowData(se)$GENEID) + incompatible_gene_ids <- incompatible_counts$GENEID - incompatible_counts <- S4Vectors::metadata(se)$incompatibleCounts - if (is.null(incompatible_counts)) { - incompatible_counts <- empty_incompatible_counts() - } else { - if (!"GENEID" %in% names(incompatible_counts)) { - incompatible_counts <- empty_incompatible_counts() - } else { - kept_genes <- unique(SummarizedExperiment::rowData(se)$GENEID) - incompatible_gene_ids <- incompatible_counts$GENEID - - rows_to_keep <- incompatible_gene_ids %in% kept_genes - incompatible_counts <- incompatible_counts[rows_to_keep, , drop = FALSE] - data.table::set( - incompatible_counts, - j = "GENEID", - value = incompatible_gene_ids[rows_to_keep] - ) - - for (sample_name in sample_names) { - if (!sample_name %in% names(incompatible_counts)) { - data.table::set( - incompatible_counts, - j = sample_name, - value = numeric(nrow(incompatible_counts)) - ) - } - } - - keep_cols <- c("GENEID", sample_names) - incompatible_counts <- incompatible_counts[, keep_cols, with = FALSE] - } - } + rows_to_keep <- incompatible_gene_ids %in% kept_genes + incompatible_counts <- incompatible_counts[rows_to_keep, , drop = FALSE] + data.table::set( + incompatible_counts, + j = "GENEID", + value = incompatible_gene_ids[rows_to_keep] + ) S4Vectors::metadata(se)$incompatibleCounts <- incompatible_counts qc_stats$total_transcripts_after_filter <- nrow(se) @@ -276,8 +821,30 @@ bambu_filter_transcripts <- function(se) { bambu_matrix_to_df <- function(se_obj, assay_name, id_col, meta_df) { assay_df <- as.data.frame(SummarizedExperiment::assays(se_obj)[[assay_name]]) - assay_df[[id_col]] <- rownames(se_obj) - assay_df <- assay_df[, c(id_col, setdiff(names(assay_df), id_col)), drop = FALSE] + if (nrow(assay_df) < 1) { + assay_df <- data.frame(stringsAsFactors = FALSE) + assay_df[[id_col]] <- character(0) + for (sample_name in colnames(se_obj)) { + assay_df[[sample_name]] <- numeric(0) + } + } else { + assay_df[[id_col]] <- rownames(se_obj) + assay_df <- assay_df[, c(id_col, setdiff(names(assay_df), id_col)), drop = FALSE] + } + + if (nrow(meta_df) < 1 || nrow(assay_df) < 1) { + output_cols <- unique(c(names(meta_df), names(assay_df))) + out_df <- data.frame(stringsAsFactors = FALSE) + for (col_name in output_cols) { + if (col_name %in% colnames(se_obj)) { + out_df[[col_name]] <- numeric(0) + } else { + out_df[[col_name]] <- character(0) + } + } + return(out_df[, output_cols, drop = FALSE]) + } + merge(meta_df, assay_df, by.x = id_col, by.y = id_col, all.y = TRUE, sort = FALSE) } @@ -293,6 +860,226 @@ bambu_format_count <- function(value) { ) } +bambu_add_library_qc_stats <- function(se, qc_stats = list()) { + lib_sizes <- colSums(SummarizedExperiment::assays(se)$counts) + qc_stats$library_sizes <- as.list(lib_sizes) + qc_stats$min_library_size <- min(lib_sizes) + qc_stats$max_library_size <- max(lib_sizes) + qc_stats$median_library_size <- stats::median(lib_sizes) + + if (length(lib_sizes) > 1 && min(lib_sizes) > 0) { + lib_size_ratio <- max(lib_sizes) / min(lib_sizes) + qc_stats$library_size_ratio <- lib_size_ratio + if (lib_size_ratio > 3) { + warning( + sprintf( + paste0( + "Large library size variation detected (%.1fx difference).\n", + " Min: %d, Max: %d reads.\n", + " CPM normalization may not be appropriate for such variation." + ), + lib_size_ratio, + min(lib_sizes), + max(lib_sizes) + ) + ) + qc_stats$library_size_warning <- sprintf("%.1fx variation (>3x threshold)", lib_size_ratio) + } + } else if (length(lib_sizes) > 1) { + qc_stats$library_size_ratio <- NA_real_ + } + + detected_per_sample <- colSums(SummarizedExperiment::assays(se)$counts > 0) + qc_stats$transcripts_detected_per_sample <- as.list(detected_per_sample) + qc_stats$median_transcripts_detected <- stats::median(detected_per_sample) + qc_stats +} + +bambu_sum_incompatible_counts <- function(tx_ses) { + incompatible_counts <- lapply(tx_ses, function(se) { + bambu_normalise_incompatible_counts( + S4Vectors::metadata(se)$incompatibleCounts, + sample_names = colnames(se) + ) + }) + incompatible_counts <- incompatible_counts[!vapply(incompatible_counts, is.null, logical(1))] + if (length(incompatible_counts) < 1) { + return(NULL) + } + + sample_cols <- setdiff(colnames(incompatible_counts[[1]]), c("GENEID", "TXNAME")) + gene_ids <- unique(unlist(lapply(incompatible_counts, function(ic) { + as.character(ic$GENEID) + }))) + + combined <- data.frame( + GENEID = gene_ids, + stringsAsFactors = FALSE + ) + for (sample_col in sample_cols) { + combined[[sample_col]] <- numeric(length(gene_ids)) + } + + for (ic in incompatible_counts) { + current_sample_cols <- setdiff(colnames(ic), c("GENEID", "TXNAME")) + if (!identical(current_sample_cols, sample_cols)) { + stop( + "Chunk quantification outputs have mismatched incompatible count columns.", + call. = FALSE + ) + } + + idx <- match(as.character(ic$GENEID), combined$GENEID) + for (sample_col in sample_cols) { + values <- ic[[sample_col]] + values[is.na(values)] <- 0 + combined[[sample_col]][idx] <- combined[[sample_col]][idx] + values + } + } + + data.table::as.data.table(combined) +} + +bambu_unique_row_ranges <- function(tx_ses, tx_names) { + combined_row_ranges <- do.call(c, lapply(tx_ses, SummarizedExperiment::rowRanges)) + unique_row_ranges <- combined_row_ranges[!duplicated(names(combined_row_ranges))] + unique_row_ranges[match(tx_names, names(unique_row_ranges))] +} + +bambu_sum_chunk_assay <- function(tx_ses, assay_name, tx_names) { + sample_names <- colnames(tx_ses[[1]]) + combined <- matrix( + 0, + nrow = length(tx_names), + ncol = length(sample_names), + dimnames = list(tx_names, sample_names) + ) + + for (se in tx_ses) { + assay_mat <- SummarizedExperiment::assay(se, assay_name) + collapsed <- rowsum(assay_mat, group = rownames(se), reorder = FALSE) + idx <- match(rownames(collapsed), tx_names) + combined[idx, ] <- combined[idx, , drop = FALSE] + collapsed + } + + combined +} + +bambu_combine_transcript_chunks <- function(tx_ses) { + if (length(tx_ses) < 1) { + stop("No chunk quantification results were provided for collation.", call. = FALSE) + } + if (length(tx_ses) == 1) { + return(tx_ses[[1]]) + } + + assay_names <- SummarizedExperiment::assayNames(tx_ses[[1]]) + sample_names <- colnames(tx_ses[[1]]) + col_data <- SummarizedExperiment::colData(tx_ses[[1]]) + object_metadata <- S4Vectors::metadata(tx_ses[[1]]) + + for (se in tx_ses[-1]) { + if (!identical(SummarizedExperiment::assayNames(se), assay_names)) { + stop("Chunk quantification outputs have mismatched assay sets.", call. = FALSE) + } + if (!identical(colnames(se), sample_names)) { + stop("Chunk quantification outputs have mismatched sample columns.", call. = FALSE) + } + } + + tx_names <- unique(unlist(lapply(tx_ses, rownames))) + combined_assays <- setNames( + lapply(assay_names, function(assay_name) { + bambu_sum_chunk_assay(tx_ses, assay_name, tx_names) + }), + assay_names + ) + if ("counts" %in% assay_names && "CPM" %in% assay_names) { + counts_mat <- combined_assays[["counts"]] + combined_assays[["CPM"]] <- t(t(counts_mat) / pmax(colSums(counts_mat), 1)) * 1e6 + } + + combined_row_ranges <- bambu_unique_row_ranges(tx_ses, tx_names) + + object_metadata$incompatibleCounts <- bambu_sum_incompatible_counts(tx_ses) + + SummarizedExperiment::SummarizedExperiment( + assays = combined_assays, + rowRanges = combined_row_ranges, + colData = col_data, + metadata = object_metadata + ) +} + +bambu_collate_chunk_outputs <- function( + chunk_dirs, + out_dir, + transcriptome_mode = "discover", + ndr = NULL, + gene_expression_fn = bambu::transcriptToGeneExpression, + write_gtf_fn = bambu::writeToGTF +) { + if (length(chunk_dirs) < 1) { + stop("No chunk quantification directories were provided for collation.", call. = FALSE) + } + + tx_ses <- lapply(chunk_dirs, function(chunk_dir) { + readRDS(file.path(chunk_dir, "bambu_transcripts.rds")) + }) + raw_se <- bambu_combine_transcript_chunks(tx_ses) + + sample_df <- utils::read.csv( + file.path(chunk_dirs[[1]], "samples.csv"), + check.names = FALSE, + stringsAsFactors = FALSE + ) + + gene_se <- gene_expression_fn(raw_se) + filtered <- bambu_filter_transcripts(raw_se) + se <- filtered$se + qc_stats <- filtered$qc_stats + message( + sprintf( + "Filtering: keeping %d / %d transcripts", + qc_stats$total_transcripts_after_filter, + qc_stats$total_transcripts_before_filter + ) + ) + + keep_gene_ids <- unique(as.character(SummarizedExperiment::rowData(se)$GENEID)) + gene_keep_idx <- rownames(gene_se) %in% keep_gene_ids + gene_se <- gene_se[gene_keep_idx, ] + + qc_stats$chunk_count <- length(chunk_dirs) + qc_stats <- bambu_add_library_qc_stats(se, qc_stats) + colnames(se) <- sample_df$alias + colnames(gene_se) <- sample_df$alias + + dir.create(out_dir, showWarnings = FALSE, recursive = TRUE) + args <- list( + out_dir = out_dir, + transcriptome_mode = transcriptome_mode, + ndr = ndr + ) + bambu_write_outputs( + se, + gene_se, + sample_df, + args, + qc_stats, + write_gtf_fn = write_gtf_fn + ) + + invisible( + list( + se = se, + gene_se = gene_se, + sample_df = sample_df, + qc_stats = qc_stats + ) + ) +} + bambu_write_matrix_tsv <- function(se_obj, assay_name, id_col, meta_df, output_path) { table_df <- bambu_matrix_to_df(se_obj, assay_name, id_col, meta_df) utils::write.table( @@ -306,14 +1093,36 @@ bambu_write_matrix_tsv <- function(se_obj, assay_name, id_col, meta_df, output_p invisible(gc(verbose = FALSE)) } -bambu_write_outputs <- function(se, gene_se, sample_df, args, qc_stats) { - bambu::writeToGTF( - SummarizedExperiment::rowRanges(se), - file = file.path(args$out_dir, "transcripts.gtf") - ) +bambu_write_outputs <- function( + se, + gene_se, + sample_df, + args, + qc_stats, + write_gtf_fn = bambu::writeToGTF, + write_rds = TRUE +) { + row_ranges <- SummarizedExperiment::rowRanges(se) + if (length(row_ranges) > 0) { + write_gtf_fn( + row_ranges, + file = file.path(args$out_dir, "transcripts.gtf") + ) + } else { + dir.create(args$out_dir, showWarnings = FALSE, recursive = TRUE) + writeLines( + c( + "##gff-version 2", + "# Empty GTF generated by supeRglue bambu" + ), + con = file.path(args$out_dir, "transcripts.gtf") + ) + } - saveRDS(se, file.path(args$out_dir, "bambu_transcripts.rds")) - saveRDS(gene_se, file.path(args$out_dir, "bambu_genes.rds")) + if (isTRUE(write_rds)) { + saveRDS(se, file.path(args$out_dir, "bambu_transcripts.rds")) + saveRDS(gene_se, file.path(args$out_dir, "bambu_genes.rds")) + } utils::write.csv( sample_df, file.path(args$out_dir, "samples.csv"), @@ -326,7 +1135,7 @@ bambu_write_outputs <- function(se, gene_se, sample_df, args, qc_stats) { tx_meta$TXNAME <- rownames(se) } if (!"GENEID" %in% names(tx_meta)) { - tx_meta$GENEID <- NA_character_ + tx_meta$GENEID <- rep(NA_character_, nrow(tx_meta)) } gene_meta <- as.data.frame(SummarizedExperiment::rowData(gene_se)) @@ -382,7 +1191,7 @@ bambu_write_outputs <- function(se, gene_se, sample_df, args, qc_stats) { ) qc_stats$transcriptome_mode <- args$transcriptome_mode - qc_stats$ndr_used <- if (bambu_discovery_enabled(args)) { + qc_stats$ndr_used <- if (identical(args$transcriptome_mode, "discover")) { if (is.null(args$ndr)) "automatic" else args$ndr } else { "N/A" @@ -402,7 +1211,7 @@ bambu_write_outputs <- function(se, gene_se, sample_df, args, qc_stats) { "", sprintf("Timestamp: %s", qc_stats$timestamp), sprintf("Mode: %s", args$transcriptome_mode), - if (bambu_discovery_enabled(args)) { + if (identical(args$transcriptome_mode, "discover")) { if (is.null(args$ndr)) { "NDR: automatic (bambu-selected)" } else { @@ -439,107 +1248,61 @@ bambu_write_outputs <- function(se, gene_se, sample_df, args, qc_stats) { writeLines(capture.output(sessionInfo()), file.path(args$out_dir, "session_info.txt")) } -main_run_bambu <- function(args) { - set.seed(42) +main_run_bambu <- function( + args, + analysis_fn = bambu::bambu, + prepare_annotations_fn = bambu::prepareAnnotations, + gene_expression_fn = bambu::transcriptToGeneExpression, + write_gtf_fn = bambu::writeToGTF, + bamfile_list_ctor = Rsamtools::BamFileList +) { + set.seed(bambu_default_seed) # bambu's parallel worker code may rely on these being attached for generics # such as seqlengths(). - suppressPackageStartupMessages(library(GenomicRanges)) - suppressPackageStartupMessages(library(Rsamtools)) + suppressPackageStartupMessages({ + library(GenomicRanges) + library(Rsamtools) + }) + bambu_validate_args(args) dir.create(args$out_dir, showWarnings = FALSE, recursive = TRUE) - inputs <- bambu_resolve_inputs(args) - args$threads <- bambu_effective_threads(args, length(inputs$bam_paths)) - annotation_obj <- bambu::prepareAnnotations(args$annotation) - - if (!is.null(args$ndr)) { - message(sprintf("Using user-specified NDR = %.3f", args$ndr)) - } else if (bambu_discovery_enabled(args)) { - message("Using bambu automatic NDR selection.") + if (identical(args$mode, "discover")) { + return(invisible(bambu_run_discover_mode( + args, + analysis_fn = analysis_fn, + prepare_annotations_fn = prepare_annotations_fn, + bamfile_list_ctor = bamfile_list_ctor + ))) } - if (bambu_discovery_enabled(args)) { - message("Novel Discovery Rate (NDR) controls transcript discovery stringency:") - message(" Lower NDR (e.g., 0.05) = fewer false positive transcripts, may miss real ones") - message(" Higher NDR (e.g., 0.2) = more sensitive discovery, more false positives") - if (is.null(args$ndr)) { - message(" Current NDR = automatic (selected by bambu from the data)") - } else { - message(sprintf(" Current NDR = %.3f balances precision and recall", args$ndr)) - } - } - if (length(inputs$bam_paths) > 1) { - message("Using BamFileList yieldSize = 250000") - } - message(sprintf("Running bambu with threads = %d", args$threads)) - - message("Running bambu...") - se <- do.call(bambu::bambu, bambu_build_args(args, inputs$reads, annotation_obj)) - message("Bambu completed successfully") - colnames(se) <- inputs$aliases - - filtered <- bambu_filter_transcripts(se) - se <- filtered$se - qc_stats <- filtered$qc_stats - gene_se <- bambu::transcriptToGeneExpression(se) - colnames(gene_se) <- inputs$aliases - message( - sprintf( - "Filtering: keeping %d / %d transcripts", - qc_stats$total_transcripts_after_filter, - qc_stats$total_transcripts_before_filter - ) - ) - - lib_sizes <- colSums(SummarizedExperiment::assays(se)$counts) - qc_stats$library_sizes <- as.list(lib_sizes) - qc_stats$min_library_size <- min(lib_sizes) - qc_stats$max_library_size <- max(lib_sizes) - qc_stats$median_library_size <- stats::median(lib_sizes) - - if (length(lib_sizes) > 1) { - lib_size_ratio <- max(lib_sizes) / min(lib_sizes) - qc_stats$library_size_ratio <- lib_size_ratio - if (lib_size_ratio > 3) { - warning( - sprintf( - paste0( - "Large library size variation detected (%.1fx difference).\n", - " Min: %s, Max: %s reads.\n", - " CPM normalization may not be appropriate for such variation." - ), - lib_size_ratio, - bambu_format_count(min(lib_sizes)), - bambu_format_count(max(lib_sizes)) - ) - ) - qc_stats$library_size_warning <- sprintf("%.1fx variation (>3x threshold)", lib_size_ratio) - } + if (identical(args$mode, "collate")) { + return(invisible(bambu_run_collate_mode( + args, + gene_expression_fn = gene_expression_fn, + write_gtf_fn = write_gtf_fn + ))) } - detected_per_sample <- colSums(SummarizedExperiment::assays(se)$counts > 0) - qc_stats$transcripts_detected_per_sample <- as.list(detected_per_sample) - qc_stats$median_transcripts_detected <- stats::median(detected_per_sample) + if (identical(args$mode, "empty")) { + return(invisible(bambu_run_empty_mode( + args, + write_gtf_fn = write_gtf_fn + ))) + } - bambu_write_outputs( - se, - gene_se, - inputs$sample_df, + invisible(bambu_run_quant_mode( args, - qc_stats - ) - - invisible( - list( - se = se, - gene_se = gene_se, - sample_df = inputs$sample_df, - qc_stats = qc_stats - ) - ) + analysis_fn = analysis_fn + )) } run_bambu_cli <- function(argv = commandArgs(trailingOnly = TRUE)) { + if (length(argv) >= 1 && !startsWith(argv[[1]], "-")) { + if (argv[[1]] %in% c("discover", "quant", "collate", "empty")) { + argv <- c("--mode", argv[[1]], argv[-1]) + } + } parsed <- argparser::parse_args(bambu_arg_parser(), argv = argv) args <- workflow_glue_r_normalise_args(parsed, bambu_arg_spec(), raw_argv = argv) main_run_bambu(args) diff --git a/bin/workflow_glue_r/tests/testthat/test_bambu.R b/bin/workflow_glue_r/tests/testthat/test_bambu.R index 4620b7e..6006ef4 100644 --- a/bin/workflow_glue_r/tests/testthat/test_bambu.R +++ b/bin/workflow_glue_r/tests/testthat/test_bambu.R @@ -1,45 +1,75 @@ #' These tests cover the validation logic owned by supeRglue bambu before bambu -#' itself is invoked: BAM/alias argument checks, sample-sheet alignment, -#' transcriptome mode selection, and NDR handling. +#' itself is invoked: mode-specific inputs, explicit BAM aliases, sample-sheet +#' alignment, transcriptome mode selection, NDR handling, and chunk collation. #' #' NOTE: Annotation/reference preparation is handled by Python #' (bin/workflow_glue/prepare_annotation_reference.py) with pytest coverage. #' These tests focus on bambu-specific validation and integration. -# Workflow requires explicit BAM paths and aliases. -# Fail fast with clear error rather than passing invalid inputs to bambu. -testthat::test_that("BAM inputs required", { - args <- list( - annotation = "annotation.gtf", - genome = "genome.fa", - out_dir = tempfile("bambu-out-"), - bams = NULL, - aliases = NULL, - transcriptome_mode = "discover", - ndr = NULL +# Workflow requires an explicit mode plus the appropriate inputs for that mode. +# Fail fast with clear errors rather than passing invalid inputs to bambu. +testthat::test_that("mode-specific bambu inputs required", { + args <- workflow_glue_r_normalise_args( + list(mode = "discover", out_dir = tempfile("bambu-out-")), + bambu_arg_spec() ) - testthat::expect_error( - workflow_glue_r_normalise_args(args, bambu_arg_spec()), - "Missing required arguments: --bams" + bambu_validate_args(args), + "Missing required arguments: --bams, --aliases, --annotation, --genome" ) args$bams <- "sampleA.bam" testthat::expect_error( - workflow_glue_r_normalise_args(args, bambu_arg_spec()), - "Missing required arguments: --aliases" + bambu_validate_args(args), + "Missing required arguments: --aliases, --annotation, --genome" ) args$aliases <- "sampleA" - normalised <- NULL - testthat::expect_silent(normalised <- workflow_glue_r_normalise_args(args, bambu_arg_spec())) - testthat::expect_identical(normalised$threads, 1L) + args$annotation <- "annotation.gtf" + args$genome <- "genome.fa" + testthat::expect_silent(bambu_validate_args(args)) + + quant_args <- workflow_glue_r_normalise_args( + list(mode = "quant", out_dir = tempfile("bambu-out-")), + bambu_arg_spec() + ) + testthat::expect_error( + bambu_validate_args(quant_args), + "Missing required arguments: --chunk_rds, --discovered_annotation_rds, --genome" + ) + quant_args$chunk_rds <- "chunk.rds" + quant_args$discovered_annotation_rds <- "annotations.rds" + quant_args$genome <- "genome.fa" + testthat::expect_silent(bambu_validate_args(quant_args)) + + collate_args <- workflow_glue_r_normalise_args( + list(mode = "collate", out_dir = tempfile("bambu-out-")), + bambu_arg_spec() + ) + testthat::expect_error( + bambu_validate_args(collate_args), + "Missing required arguments: --chunk_dirs" + ) + collate_args$chunk_dirs <- "chunkA,chunkB" + testthat::expect_silent(bambu_validate_args(collate_args)) + + empty_args <- workflow_glue_r_normalise_args( + list(mode = "empty", out_dir = tempfile("bambu-out-")), + bambu_arg_spec() + ) + testthat::expect_error( + bambu_validate_args(empty_args), + "Missing required arguments: --aliases" + ) + empty_args$aliases <- "sampleA,sampleB" + testthat::expect_silent(bambu_validate_args(empty_args)) }) # transcriptome_mode must be "discover" or "fixed_annotation". # NDR (Novel Discovery Rate) must be between 0 and 1 when provided. testthat::test_that("invalid discovery settings rejected", { args <- list( + mode = "discover", annotation = "annotation.gtf", genome = "genome.fa", out_dir = tempfile("bambu-out-"), @@ -87,7 +117,7 @@ testthat::test_that("empty BAM list rejected", { ) testthat::expect_error( - bambu_resolve_inputs(args), + bambu_resolve_inputs(args, bamfile_list_ctor = function(paths, yieldSize) paths), "No BAM files were provided in --bams" ) }) @@ -101,15 +131,16 @@ testthat::test_that("unique sample aliases required", { sample_sheet = NULL ) testthat::expect_error( - bambu_resolve_inputs(args), + bambu_resolve_inputs(args, bamfile_list_ctor = function(paths, yieldSize) paths), "BAM aliases must be unique" ) args$aliases <- "sampleA" testthat::expect_error( - bambu_resolve_inputs(args), + bambu_resolve_inputs(args, bamfile_list_ctor = function(paths, yieldSize) paths), "Provide one alias per BAM in --bams" ) + missing_alias_sheet <- tempfile(fileext = ".csv") writeLines( paste( @@ -125,7 +156,7 @@ testthat::test_that("unique sample aliases required", { sample_sheet = missing_alias_sheet ) testthat::expect_error( - bambu_resolve_inputs(args), + bambu_resolve_inputs(args, bamfile_list_ctor = function(paths, yieldSize) paths), "Sample sheet must contain an 'alias' column" ) @@ -141,7 +172,7 @@ testthat::test_that("unique sample aliases required", { ) args$sample_sheet <- duplicate_alias_sheet testthat::expect_error( - bambu_resolve_inputs(args), + bambu_resolve_inputs(args, bamfile_list_ctor = function(paths, yieldSize) paths), "Sample sheet aliases must be unique" ) }) @@ -166,11 +197,14 @@ testthat::test_that("sample sheet reordered to match BAMs", { aliases = "sampleA,sampleB", sample_sheet = sample_sheet ) - resolved <- bambu_resolve_inputs(args) + resolved <- bambu_resolve_inputs( + args, + bamfile_list_ctor = function(paths, yieldSize) paths + ) testthat::expect_equal(resolved$aliases, c("sampleA", "sampleB")) testthat::expect_equal(resolved$sample_df$alias, c("sampleA", "sampleB")) - testthat::expect_s4_class(resolved$reads, "BamFileList") + testthat::expect_equal(resolved$sample_df$condition, c("control", "treated")) bad_sheet <- tempfile(fileext = ".csv") writeLines( @@ -184,53 +218,410 @@ testthat::test_that("sample sheet reordered to match BAMs", { args$sample_sheet <- bad_sheet testthat::expect_error( - bambu_resolve_inputs(args), + bambu_resolve_inputs(args, bamfile_list_ctor = function(paths, yieldSize) paths), "Sample sheet is missing alias rows" ) }) -# transcriptome_mode="discover" → bambu(discovery=TRUE, NDR=value) -# transcriptome_mode="fixed_annotation" → bambu(discovery=FALSE, no NDR param) -testthat::test_that("transcriptome mode mapped to bambu args", { +# Explicit discovery/quant flags are passed through to bambu consistently. +# NDR is only passed during discovery and omitted when automatic selection is wanted. +testthat::test_that("bambu args include requested discovery and quant flags", { annotation_obj <- structure(list(annotation = TRUE), class = "mockAnnotation") - discover_args <- list( + args <- list( genome = "genome.fa", - threads = 3, + threads = 3L, transcriptome_mode = "discover", ndr = 0.2 ) - discover <- bambu_build_args(discover_args, reads = "sample.bam", annotation_obj = annotation_obj) + discover <- bambu_build_args( + args, + reads = "sample.bam", + annotation_obj = annotation_obj, + discovery = TRUE, + quant = FALSE + ) testthat::expect_true(discover$discovery) + testthat::expect_false(discover$quant) testthat::expect_equal(discover$NDR, 0.2) testthat::expect_equal(discover$ncore, 3L) testthat::expect_true(discover$lowMemory) testthat::expect_equal(discover$yieldSize, 250000L) - auto_ndr_args <- list( - genome = "genome.fa", - threads = 2, - transcriptome_mode = "discover", - ndr = NULL + args$ndr <- NULL + auto_ndr <- bambu_build_args( + args, + reads = "sample.bam", + annotation_obj = annotation_obj, + discovery = TRUE, + quant = FALSE ) - auto_ndr <- bambu_build_args(auto_ndr_args, reads = "sample.bam", annotation_obj = annotation_obj) - testthat::expect_true(auto_ndr$discovery) testthat::expect_false("NDR" %in% names(auto_ndr)) - testthat::expect_equal(auto_ndr$yieldSize, 250000L) - fixed_args <- list( - genome = "genome.fa", - threads = 1, - transcriptome_mode = "fixed_annotation", - ndr = NULL + quant <- bambu_build_args( + args, + reads = "sample.bam", + annotation_obj = annotation_obj, + discovery = FALSE, + quant = TRUE ) - fixed <- bambu_build_args(fixed_args, reads = "sample.bam", annotation_obj = annotation_obj) - testthat::expect_false(fixed$discovery) - testthat::expect_false("NDR" %in% names(fixed)) - testthat::expect_true(fixed$lowMemory) - testthat::expect_equal(fixed$yieldSize, 250000L) + testthat::expect_false(quant$discovery) + testthat::expect_true(quant$quant) + testthat::expect_false("NDR" %in% names(quant)) }) +# Discover mode should write reusable rcFiles, discovered annotations, and chunk bundles. +# This is the scatter source for later per-chromosome quantification. +testthat::test_that("discover mode writes chunked rc outputs", { + fixture_dir <- tempfile("bambu-discover-mode-") + dir.create(fixture_dir) + bam_dir <- file.path(fixture_dir, "bams") + dir.create(bam_dir) + sample_a <- file.path(bam_dir, "sampleA.aligned.sorted.bam") + sample_b <- file.path(bam_dir, "sampleB.bam") + file.create(sample_b) + file.create(sample_a) + + sample_sheet <- file.path(fixture_dir, "sample_sheet.csv") + writeLines( + paste( + "alias,condition", + "sampleB,treated", + "sampleA,control", + sep = "\n" + ), + sample_sheet + ) + + captured <- new.env(parent = emptyenv()) + fake_bamfile_list <- function(paths, yieldSize) { + captured$bamfile_paths <- paths + captured$yield_size <- yieldSize + structure(paths, names = c("sampleA", "sampleB"), class = "mockBamFileList") + } + make_rc_sample <- function(alias) { + rcf <- make_test_tx_se(sample_names = alias) + S4Vectors::mcols(SummarizedExperiment::rowRanges(rcf))$chr.rc <- c("chr1", "chr1", "chr2", "chr2") + rcf + } + fake_analysis <- function( + reads, + annotations, + genome, + ncore, + discovery, + quant, + lowMemory, + yieldSize, + verbose, + NDR = NULL + ) { + captured$calls <- c(captured$calls, list(list( + reads = reads, + annotations = annotations, + genome = genome, + ncore = ncore, + discovery = discovery, + quant = quant, + lowMemory = lowMemory, + yieldSize = yieldSize, + verbose = verbose, + NDR = NDR + ))) + if (!discovery && !quant) { + return(list( + sampleA = make_rc_sample("sampleA"), + sampleB = make_rc_sample("sampleB") + )) + } + structure(list(discovered = TRUE), class = "mockDiscoveredAnnotation") + } + + args <- workflow_glue_r_normalise_args( + list( + mode = "discover", + annotation = "annotation.gtf", + genome = "genome.fa", + out_dir = file.path(fixture_dir, "out"), + bams = paste(c(sample_a, sample_b), collapse = ","), + aliases = "sampleA,sampleB", + sample_sheet = sample_sheet, + transcriptome_mode = "discover", + ndr = 0.25, + threads = 2 + ), + bambu_arg_spec() + ) + + result <- suppressMessages(main_run_bambu( + args, + analysis_fn = fake_analysis, + prepare_annotations_fn = function(annotation) { + captured$annotation_path <- annotation + structure(list(path = annotation), class = "mockAnnotation") + }, + bamfile_list_ctor = fake_bamfile_list + )) + + testthat::expect_equal(captured$annotation_path, "annotation.gtf") + testthat::expect_equal(captured$yield_size, 250000L) + testthat::expect_equal(captured$bamfile_paths, c(sample_a, sample_b)) + testthat::expect_equal(result$sample_df$alias, c("sampleA", "sampleB")) + testthat::expect_equal(result$sample_df$condition, c("control", "treated")) + testthat::expect_length(captured$calls, 2) + testthat::expect_false(captured$calls[[1]]$discovery) + testthat::expect_false(captured$calls[[1]]$quant) + testthat::expect_true(captured$calls[[1]]$lowMemory) + testthat::expect_equal(captured$calls[[1]]$yieldSize, 250000L) + testthat::expect_equal(captured$calls[[1]]$ncore, 1L) + testthat::expect_true(captured$calls[[2]]$discovery) + testthat::expect_false(captured$calls[[2]]$quant) + testthat::expect_equal(captured$calls[[2]]$NDR, 0.25) + testthat::expect_true(file.exists(file.path(args$out_dir, "bambu_rcfiles.rds"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "bambu_discovered_annotations.rds"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "chunk_manifest.tsv"))) + + manifest <- utils::read.delim( + file.path(args$out_dir, "chunk_manifest.tsv"), + check.names = FALSE, + stringsAsFactors = FALSE + ) + testthat::expect_equal(sort(manifest$seqname), c("chr1", "chr2")) + chunk_bundle <- readRDS(manifest$rds_path[[1]]) + testthat::expect_true("annotation_tx_count" %in% names(manifest)) + testthat::expect_true(all(c( + "chunk_id", "seqname", "aliases", "sample_df", "rc_files", "annotation_tx_count" + ) %in% names(chunk_bundle))) + testthat::expect_equal(manifest$annotation_tx_count, c(0L, 0L)) + testthat::expect_equal(chunk_bundle$aliases, c("sampleA", "sampleB")) +}) + +testthat::test_that("annotation tx counts are computed once per seqname", { + discovered_annotations <- GenomicRanges::GRangesList( + tx1 = GenomicRanges::GRanges(seqnames = "chr1", ranges = IRanges::IRanges(c(1, 10), width = 5)), + tx2 = GenomicRanges::GRanges(seqnames = "chr1", ranges = IRanges::IRanges(20, width = 5)), + tx3 = GenomicRanges::GRanges(seqnames = "chr2", ranges = IRanges::IRanges(c(30, 40), width = 5)) + ) + + counts <- bambu_annotation_tx_counts_by_seqname(discovered_annotations) + + testthat::expect_equal(as.integer(counts[c("chr1", "chr2")]), c(2L, 1L)) + testthat::expect_equal(bambu_annotation_tx_count_for_seqname(discovered_annotations, "chr1"), 2L) + testthat::expect_equal(bambu_annotation_tx_count_for_seqname(discovered_annotations, "chr3"), 0L) +}) + +# Quant mode should consume one chunk bundle and write only raw chunk quantification +# artifacts. Filtering and gene aggregation happen during collate. +testthat::test_that("quant mode writes chunk quantification outputs", { + fixture_dir <- tempfile("bambu-quant-mode-") + dir.create(fixture_dir) + + make_rc_sample <- function(alias) { + rcf <- make_test_tx_se(sample_names = alias) + S4Vectors::mcols(SummarizedExperiment::rowRanges(rcf))$chr.rc <- c("chr1", "chr1", "chr1", "chr1") + rcf + } + + chunk_bundle <- list( + chunk_id = "chr1", + seqname = "chr1", + annotation_tx_count = 1L, + aliases = c("sampleA", "sampleB"), + sample_df = data.frame(alias = c("sampleA", "sampleB"), stringsAsFactors = FALSE), + rc_files = list(sampleA = make_rc_sample("sampleA"), sampleB = make_rc_sample("sampleB")) + ) + chunk_rds <- file.path(fixture_dir, "chr1.rds") + saveRDS(chunk_bundle, chunk_rds) + discovered_annotation_rds <- file.path(fixture_dir, "annotations.rds") + saveRDS(structure(list(discovered = TRUE), class = "mockDiscoveredAnnotation"), discovered_annotation_rds) + + captured <- new.env(parent = emptyenv()) + fake_analysis <- function( + reads, + annotations, + genome, + ncore, + discovery, + quant, + lowMemory, + yieldSize, + verbose, + ... + ) { + captured$reads <- reads + captured$annotations <- annotations + captured$genome <- genome + captured$ncore <- ncore + captured$discovery <- discovery + captured$quant <- quant + captured$lowMemory <- lowMemory + captured$yieldSize <- yieldSize + captured$verbose <- verbose + make_test_tx_se(sample_names = c("sampleA", "sampleB")) + } + + args <- workflow_glue_r_normalise_args( + list( + mode = "quant", + genome = "genome.fa", + out_dir = file.path(fixture_dir, "out"), + chunk_rds = chunk_rds, + discovered_annotation_rds = discovered_annotation_rds, + transcriptome_mode = "discover", + ndr = NULL, + threads = 2 + ), + bambu_arg_spec() + ) + + result <- suppressMessages(main_run_bambu( + args, + analysis_fn = fake_analysis + )) + + testthat::expect_equal(captured$genome, "genome.fa") + testthat::expect_false(captured$discovery) + testthat::expect_true(captured$quant) + testthat::expect_true(captured$lowMemory) + testthat::expect_equal(captured$yieldSize, 250000L) + testthat::expect_equal(result$sample_df$alias, c("sampleA", "sampleB")) + testthat::expect_true(file.exists(file.path(args$out_dir, "bambu_transcripts.rds"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "samples.csv"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "bambu_qc_stats.json"))) + qc <- jsonlite::read_json(file.path(args$out_dir, "bambu_qc_stats.json"), simplifyVector = TRUE) + testthat::expect_equal(qc$chunk_id, "chr1") + testthat::expect_equal(qc$seqname, "chr1") +}) + +testthat::test_that("quant mode skips chunks with no discovered annotations on the seqname", { + fixture_dir <- tempfile("bambu-quant-empty-chunk-") + dir.create(fixture_dir) + + make_rc_sample <- function(alias) { + rcf <- make_test_tx_se(sample_names = alias) + S4Vectors::mcols(SummarizedExperiment::rowRanges(rcf))$chr.rc <- rep( + "chr3_GL000221v1_random", + nrow(rcf) + ) + rcf + } + + chunk_bundle <- list( + chunk_id = "chr3_GL000221v1_random", + seqname = "chr3_GL000221v1_random", + aliases = c("sampleA", "sampleB"), + sample_df = data.frame(alias = c("sampleA", "sampleB"), stringsAsFactors = FALSE), + rc_files = list(sampleA = make_rc_sample("sampleA"), sampleB = make_rc_sample("sampleB")) + ) + chunk_rds <- file.path(fixture_dir, "chr3_GL000221v1_random.rds") + saveRDS(chunk_bundle, chunk_rds) + + discovered_annotation_rds <- file.path(fixture_dir, "annotations.rds") + saveRDS(make_test_bambu_row_ranges(fixture_dir), discovered_annotation_rds) + + analysis_called <- FALSE + fake_analysis <- function(...) { + analysis_called <<- TRUE + stop("analysis_fn should not be called for chunks with no matching annotations") + } + + args <- workflow_glue_r_normalise_args( + list( + mode = "quant", + genome = "genome.fa", + out_dir = file.path(fixture_dir, "out"), + chunk_rds = chunk_rds, + discovered_annotation_rds = discovered_annotation_rds, + transcriptome_mode = "discover", + ndr = NULL, + threads = 2 + ), + bambu_arg_spec() + ) + + result <- testthat::expect_warning( + suppressMessages(main_run_bambu( + args, + analysis_fn = fake_analysis + )), + "contains no transcripts on this seqname" + ) + + testthat::expect_false(analysis_called) + testthat::expect_equal(result$sample_df$alias, c("sampleA", "sampleB")) + + se <- readRDS(file.path(args$out_dir, "bambu_transcripts.rds")) + testthat::expect_equal(nrow(se), 0) + testthat::expect_equal(colnames(se), c("sampleA", "sampleB")) + testthat::expect_equal( + SummarizedExperiment::assayNames(se), + c("counts", "CPM", "fullLengthCounts", "uniqueCounts") + ) + testthat::expect_equal( + colnames(S4Vectors::metadata(se)$incompatibleCounts), + c("GENEID", "sampleA", "sampleB") + ) + + qc <- jsonlite::read_json(file.path(args$out_dir, "bambu_qc_stats.json"), simplifyVector = TRUE) + testthat::expect_equal(qc$chunk_id, "chr3_GL000221v1_random") + testthat::expect_equal(qc$seqname, "chr3_GL000221v1_random") + testthat::expect_equal(qc$total_transcripts_before_filter, 0) + testthat::expect_equal(qc$total_genes_before_filter, 0) +}) + +testthat::test_that("empty mode writes valid empty outputs including bambu rds files", { + fixture_dir <- tempfile("bambu-empty-mode-") + dir.create(fixture_dir) + + args <- workflow_glue_r_normalise_args( + list( + mode = "empty", + aliases = "sampleA,sampleB", + out_dir = file.path(fixture_dir, "out"), + transcriptome_mode = "fixed_annotation" + ), + bambu_arg_spec() + ) + + result <- suppressMessages(main_run_bambu(args)) + + testthat::expect_equal(result$sample_df$alias, c("sampleA", "sampleB")) + testthat::expect_true(file.exists(file.path(args$out_dir, "transcripts.gtf"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "transcript_counts.tsv"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "gene_counts.tsv"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "samples.csv"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "bambu_qc_stats.json"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "bambu_transcripts.rds"))) + testthat::expect_true(file.exists(file.path(args$out_dir, "bambu_genes.rds"))) + + tx_counts <- utils::read.delim( + file.path(args$out_dir, "transcript_counts.tsv"), + check.names = FALSE, + stringsAsFactors = FALSE + ) + gene_counts <- utils::read.delim( + file.path(args$out_dir, "gene_counts.tsv"), + check.names = FALSE, + stringsAsFactors = FALSE + ) + qc <- jsonlite::read_json(file.path(args$out_dir, "bambu_qc_stats.json"), simplifyVector = TRUE) + gtf_lines <- readLines(file.path(args$out_dir, "transcripts.gtf"), warn = FALSE) + tx_rds <- readRDS(file.path(args$out_dir, "bambu_transcripts.rds")) + gene_rds <- readRDS(file.path(args$out_dir, "bambu_genes.rds")) + + testthat::expect_equal(nrow(tx_counts), 0) + testthat::expect_equal(nrow(gene_counts), 0) + testthat::expect_true(all(c("sampleA", "sampleB") %in% names(tx_counts))) + testthat::expect_true(all(c("sampleA", "sampleB") %in% names(gene_counts))) + testthat::expect_equal(gtf_lines[[1]], "##gff-version 2") + testthat::expect_true(any(grepl("^#", gtf_lines))) + testthat::expect_s4_class(tx_rds, "RangedSummarizedExperiment") + testthat::expect_s4_class(gene_rds, "SummarizedExperiment") + testthat::expect_equal(nrow(tx_rds), 0) + testthat::expect_equal(nrow(gene_rds), 0) + testthat::expect_true(isTRUE(qc$empty_output)) + testthat::expect_equal(qc$chunk_count, 0) +}) ### # Output serialization # @@ -256,14 +647,13 @@ testthat::test_that("bambu outputs written correctly", { sample_names <- c("sampleA", "sampleB") base_se <- make_test_tx_se(sample_names = sample_names) row_ranges <- make_test_bambu_row_ranges(out_dir) - se <- SummarizedExperiment::SummarizedExperiment( assays = SummarizedExperiment::assays(base_se), rowRanges = row_ranges ) gene_se <- make_test_gene_se(sample_names = sample_names) sample_df <- data.frame(alias = sample_names, stringsAsFactors = FALSE) - argv <- list(out_dir = out_dir, transcriptome_mode = "discover", ndr = 0.15) + args <- list(out_dir = out_dir, transcriptome_mode = "discover", ndr = 0.15) qc_stats <- list( samples = 2, total_transcripts_before_filter = 4, @@ -281,8 +671,14 @@ testthat::test_that("bambu outputs written correctly", { se, gene_se, sample_df, - argv, - qc_stats + args, + qc_stats, + write_gtf_fn = function(row_ranges, file) { + writeLines( + 'chr1\tsim\texon\t1\t50\t.\t+\t.\tgene_id "gene1"; transcript_id "tx1";', + file + ) + } ) testthat::expect_true(file.exists(file.path(out_dir, "transcripts.gtf"))) @@ -306,6 +702,115 @@ testthat::test_that("bambu outputs written correctly", { testthat::expect_true(all(c("TXNAME", "sampleA", "sampleB") %in% names(tx_counts))) }) +testthat::test_that("collate combines multiple chunk quantification outputs", { + fixture_dir <- tempfile("bambu-collate-multi-") + dir.create(fixture_dir) + + sample_df <- data.frame(alias = c("sampleA", "sampleB"), stringsAsFactors = FALSE) + base_se <- make_test_tx_se(sample_names = sample_df$alias) + tx_se <- SummarizedExperiment::SummarizedExperiment( + assays = SummarizedExperiment::assays(base_se), + rowRanges = make_test_bambu_row_ranges(fixture_dir) + ) + + chunk_dirs <- c(file.path(fixture_dir, "chunk1"), file.path(fixture_dir, "chunk2")) + dir.create(chunk_dirs[[1]]) + dir.create(chunk_dirs[[2]]) + saveRDS(tx_se[1:2, ], file.path(chunk_dirs[[1]], "bambu_transcripts.rds")) + saveRDS(tx_se[3:4, ], file.path(chunk_dirs[[2]], "bambu_transcripts.rds")) + utils::write.csv(sample_df, file.path(chunk_dirs[[1]], "samples.csv"), row.names = FALSE, quote = FALSE) + utils::write.csv(sample_df, file.path(chunk_dirs[[2]], "samples.csv"), row.names = FALSE, quote = FALSE) + + mock_gene_expression <- function(se) { + counts <- rowsum( + SummarizedExperiment::assay(se, "counts"), + group = as.character(SummarizedExperiment::rowData(se)$GENEID), + reorder = FALSE + ) + cpm <- t(t(counts) / colSums(counts)) * 1e6 + SummarizedExperiment::SummarizedExperiment( + assays = list(counts = counts, CPM = cpm), + rowData = S4Vectors::DataFrame(GENEID = rownames(counts)) + ) + } + + out_dir <- file.path(fixture_dir, "collated") + suppressWarnings(suppressMessages( + bambu_collate_chunk_outputs( + chunk_dirs, + out_dir = out_dir, + transcriptome_mode = "fixed_annotation", + gene_expression_fn = mock_gene_expression, + write_gtf_fn = function(row_ranges, file) { + writeLines( + 'chr1\tsim\texon\t1\t50\t.\t+\t.\tgene_id "gene1"; transcript_id "tx1";', + file + ) + } + ) + )) + + collated_tx <- readRDS(file.path(out_dir, "bambu_transcripts.rds")) + collated_gene <- readRDS(file.path(out_dir, "bambu_genes.rds")) + testthat::expect_equal(nrow(collated_tx), 4) + testthat::expect_equal(nrow(collated_gene), 2) + testthat::expect_equal(colnames(collated_tx), sample_df$alias) + testthat::expect_true(file.exists(file.path(out_dir, "transcript_counts.tsv"))) + testthat::expect_true(file.exists(file.path(out_dir, "gene_counts.tsv"))) +}) + +testthat::test_that("chunk combiner sums duplicate transcript rows across chunks", { + sample_df <- data.frame(alias = "sampleA", stringsAsFactors = FALSE) + base_se <- make_test_tx_se(sample_names = sample_df$alias) + fixture_dir <- tempfile("bambu-combine-") + dir.create(fixture_dir) + tx_se <- SummarizedExperiment::SummarizedExperiment( + assays = SummarizedExperiment::assays(base_se), + rowRanges = make_test_bambu_row_ranges(fixture_dir) + ) + + chunk1 <- tx_se + chunk2 <- tx_se + counts1 <- SummarizedExperiment::assay(chunk1, "counts") + counts2 <- SummarizedExperiment::assay(chunk2, "counts") + counts1[3:4, ] <- 0 + counts2[1:2, ] <- 0 + SummarizedExperiment::assay(chunk1, "counts", withDimnames = FALSE) <- counts1 + SummarizedExperiment::assay(chunk2, "counts", withDimnames = FALSE) <- counts2 + + cpm1 <- t(t(counts1) / pmax(colSums(counts1), 1)) * 1e6 + cpm2 <- t(t(counts2) / pmax(colSums(counts2), 1)) * 1e6 + SummarizedExperiment::assay(chunk1, "CPM", withDimnames = FALSE) <- cpm1 + SummarizedExperiment::assay(chunk2, "CPM", withDimnames = FALSE) <- cpm2 + + S4Vectors::metadata(chunk1)$incompatibleCounts <- data.frame( + GENEID = c("gene1", "gene2"), + `01` = c(10, 0), + stringsAsFactors = FALSE + ) + S4Vectors::metadata(chunk2)$incompatibleCounts <- data.frame( + GENEID = c("gene1", "gene2"), + `01` = c(0, 20), + stringsAsFactors = FALSE + ) + + combined <- bambu_combine_transcript_chunks(list(chunk1, chunk2)) + + testthat::expect_equal(nrow(combined), 4) + testthat::expect_equal(rownames(combined), rownames(tx_se)) + testthat::expect_equal( + SummarizedExperiment::assay(combined, "counts"), + SummarizedExperiment::assay(tx_se, "counts") + ) + expected_cpm <- t(t(SummarizedExperiment::assay(tx_se, "counts")) / + colSums(SummarizedExperiment::assay(tx_se, "counts"))) * 1e6 + testthat::expect_equal(SummarizedExperiment::assay(combined, "CPM"), expected_cpm) + + incompatible <- S4Vectors::metadata(combined)$incompatibleCounts + testthat::expect_equal(names(incompatible), c("GENEID", "sampleA")) + testthat::expect_equal(incompatible$sampleA, c(10, 20)) +}) + # Transcript filtering edge cases testthat::test_that("filters on counts when fullLengthCounts disagrees", { se <- make_test_tx_se(sample_names = c("sampleA", "sampleB")) @@ -332,7 +837,7 @@ testthat::test_that("filters on counts when fullLengthCounts disagrees", { testthat::test_that("filters on counts when no fullLengthCounts assay", { se <- make_test_tx_se(sample_names = c("sampleA", "sampleB")) - # Set some transcripts with counts + # Set some transcripts with counts. counts <- SummarizedExperiment::assay(se, "counts") counts[1, ] <- c(10, 8) counts[2, ] <- c(5, 3) @@ -340,7 +845,7 @@ testthat::test_that("filters on counts when no fullLengthCounts assay", { counts[4, ] <- c(0, 0) SummarizedExperiment::assay(se, "counts", withDimnames = FALSE) <- counts - # Remove fullLengthCounts assay entirely so it falls back to counts + # Remove fullLengthCounts assay entirely so it falls back to counts. assay_list <- SummarizedExperiment::assays(se) assay_list[["fullLengthCounts"]] <- NULL SummarizedExperiment::assays(se, withDimnames = FALSE) <- assay_list @@ -353,7 +858,7 @@ testthat::test_that("filters on counts when no fullLengthCounts assay", { testthat::test_that("error when all transcripts filtered", { se <- make_test_tx_se(sample_names = c("sampleA", "sampleB")) - # All transcripts have zero counts + # All transcripts have zero counts. SummarizedExperiment::assay(se, "counts", withDimnames = FALSE) <- matrix(0, nrow = 4, ncol = 2) testthat::expect_error( @@ -364,7 +869,7 @@ testthat::test_that("error when all transcripts filtered", { testthat::test_that("QC stats match filtered results", { se <- make_test_tx_se(sample_names = c("sampleA", "sampleB")) - # Set one transcript with counts, rest without + # Set one transcript with counts, rest without. counts <- matrix(0, nrow = 4, ncol = 2) counts[1, ] <- c(10, 8) SummarizedExperiment::assay(se, "counts", withDimnames = FALSE) <- counts @@ -421,15 +926,30 @@ testthat::test_that("bambu_filter_transcripts contract with transcriptToGeneExpr testthat::expect_error(bambu::transcriptToGeneExpression(filtered$se), NA) }) +testthat::test_that("bambu_filter_transcripts renames generic incompatibleCounts columns", { + se <- make_test_tx_se(sample_names = "sampleA") + S4Vectors::metadata(se)$incompatibleCounts <- data.table::data.table( + GENEID = c("gene1", "gene2"), + `01` = c(4, 9) + ) + + filtered <- bambu_filter_transcripts(se) + incompatible <- S4Vectors::metadata(filtered$se)$incompatibleCounts + + testthat::expect_equal(names(incompatible), c("GENEID", "sampleA")) + testthat::expect_equal(incompatible$sampleA, c(4, 9)) +}) + ### # CLI integration tests # # End-to-end tests with real bambu library (not mocked): # - Build BAMs from committed fixtures (reference.fa, annotation.gtf, reads.fastq) -# - Run `supeRglue bambu` with --bams input -# - Verify output files exist and contain data for downstream workflow steps +# - Run `supeRglue bambu discover` to produce rcFiles/chunks +# - Run `supeRglue bambu quant` on one emitted chunk +# - Run `supeRglue bambu collate` to produce downstream workflow outputs -testthat::test_that("CLI single BAM in bams input with fixed annotation", { +testthat::test_that("CLI discover writes reusable chunk artifacts", { fixture_dir <- tempfile("bambu-cli-") dir.create(fixture_dir) @@ -444,6 +964,7 @@ testthat::test_that("CLI single BAM in bams input with fixed annotation", { "supeRglue", c( "bambu", + "discover", "--bams", bam_path, "--aliases", "sampleA", "--sample_sheet", sample_sheet, @@ -460,74 +981,94 @@ testthat::test_that("CLI single BAM in bams input with fixed annotation", { 0L, info = paste(result$output, collapse = "\n") ) - testthat::expect_true(file.exists(file.path(out_dir, "transcripts.gtf"))) - testthat::expect_true(file.exists(file.path(out_dir, "transcript_counts.tsv"))) - testthat::expect_true(file.exists(file.path(out_dir, "gene_counts.tsv"))) + testthat::expect_true(file.exists(file.path(out_dir, "bambu_rcfiles.rds"))) + testthat::expect_true(file.exists(file.path(out_dir, "bambu_discovered_annotations.rds"))) + testthat::expect_true(file.exists(file.path(out_dir, "chunk_manifest.tsv"))) testthat::expect_true(file.exists(file.path(out_dir, "samples.csv"))) - tx_counts <- utils::read.delim( - file.path(out_dir, "transcript_counts.tsv"), + samples <- utils::read.csv( + file.path(out_dir, "samples.csv"), check.names = FALSE, stringsAsFactors = FALSE ) - gene_counts <- utils::read.delim( - file.path(out_dir, "gene_counts.tsv"), + manifest <- utils::read.delim( + file.path(out_dir, "chunk_manifest.tsv"), check.names = FALSE, stringsAsFactors = FALSE ) + chunk_bundle <- readRDS(manifest$rds_path[[1]]) - testthat::expect_gt(nrow(tx_counts), 0) - testthat::expect_gt(nrow(gene_counts), 0) + testthat::expect_equal(samples$alias, "sampleA") + testthat::expect_gt(nrow(manifest), 0) + testthat::expect_true(all(c("chunk_id", "seqname", "annotation_tx_count", "rds_path") %in% names(manifest))) + testthat::expect_true("annotation_tx_count" %in% names(chunk_bundle)) + testthat::expect_true(all(manifest$annotation_tx_count >= 0)) + testthat::expect_equal(chunk_bundle$aliases, "sampleA") }) -testthat::test_that("CLI bams input preserves sample order", { +testthat::test_that("CLI quant consumes a discover chunk", { fixture_dir <- tempfile("bambu-cli-dir-") dir.create(fixture_dir) - bam_dir <- file.path(fixture_dir, "bams") - dir.create(bam_dir) reference <- workflow_glue_r_fixture("bambu", "reference.fa") annotation <- workflow_glue_r_fixture("bambu", "annotation.gtf") reads <- workflow_glue_r_fixture("bambu", "reads.fastq") - expect_bam_fixture_built(reference, reads, bam_dir, alias = "sampleA") - expect_bam_fixture_built(reference, reads, bam_dir, alias = "sampleB") + bam_path <- expect_bam_fixture_built(reference, reads, fixture_dir, alias = "sampleA") sample_sheet <- file.path(fixture_dir, "sample_sheet.csv") writeLines( paste( "alias", - "sampleB", "sampleA", sep = "\n" ), sample_sheet ) - out_dir <- file.path(fixture_dir, "out") - - result <- run_rscript( + discover_out_dir <- file.path(fixture_dir, "discover") + discover_result <- run_rscript( "supeRglue", c( "bambu", - "--bams", paste( - c( - file.path(bam_dir, "sampleA.aligned.sorted.bam"), - file.path(bam_dir, "sampleB.aligned.sorted.bam") - ), - collapse = "," - ), - "--aliases", "sampleA,sampleB", + "discover", + "--bams", bam_path, + "--aliases", "sampleA", "--sample_sheet", sample_sheet, "--annotation", annotation, "--genome", reference, "--transcriptome_mode", "fixed_annotation", "--threads", "1", + "--out_dir", discover_out_dir + ) + ) + + testthat::expect_equal( + discover_result$status, + 0L, + info = paste(discover_result$output, collapse = "\n") + ) + + manifest <- utils::read.delim( + file.path(discover_out_dir, "chunk_manifest.tsv"), + check.names = FALSE, + stringsAsFactors = FALSE + ) + out_dir <- file.path(fixture_dir, "quant") + quant_result <- run_rscript( + "supeRglue", + c( + "bambu", + "quant", + "--chunk_rds", manifest$rds_path[[1]], + "--discovered_annotation_rds", file.path(discover_out_dir, "bambu_discovered_annotations.rds"), + "--genome", reference, + "--threads", "1", "--out_dir", out_dir ) ) testthat::expect_equal( - result$status, + quant_result$status, 0L, - info = paste(result$output, collapse = "\n") + info = paste(quant_result$output, collapse = "\n") ) samples <- utils::read.csv( @@ -535,15 +1076,104 @@ testthat::test_that("CLI bams input preserves sample order", { check.names = FALSE, stringsAsFactors = FALSE ) - tx_counts <- utils::read.delim( - file.path(out_dir, "transcript_counts.tsv"), + tx_se <- readRDS(file.path(out_dir, "bambu_transcripts.rds")) + qc <- jsonlite::read_json( + file.path(out_dir, "bambu_qc_stats.json"), + simplifyVector = TRUE + ) + + testthat::expect_equal(samples$alias, "sampleA") + testthat::expect_gt(nrow(tx_se), 0) + testthat::expect_equal(qc$chunk_id, manifest$chunk_id[[1]]) + testthat::expect_equal(qc$seqname, manifest$seqname[[1]]) +}) + +testthat::test_that("CLI collate consumes quant chunk directories", { + fixture_dir <- tempfile("bambu-cli-collate-") + dir.create(fixture_dir) + + reference <- workflow_glue_r_fixture("bambu", "reference.fa") + annotation <- workflow_glue_r_fixture("bambu", "annotation.gtf") + reads <- workflow_glue_r_fixture("bambu", "reads.fastq") + bam_path <- expect_bam_fixture_built(reference, reads, fixture_dir, alias = "sampleA") + sample_sheet <- file.path(fixture_dir, "sample_sheet.csv") + writeLines( + paste( + "alias", + "sampleA", + sep = "\n" + ), + sample_sheet + ) + + discover_out_dir <- file.path(fixture_dir, "discover") + discover_result <- run_rscript( + "supeRglue", + c( + "bambu", + "discover", + "--bams", bam_path, + "--aliases", "sampleA", + "--sample_sheet", sample_sheet, + "--annotation", annotation, + "--genome", reference, + "--transcriptome_mode", "fixed_annotation", + "--threads", "1", + "--out_dir", discover_out_dir + ) + ) + testthat::expect_equal( + discover_result$status, + 0L, + info = paste(discover_result$output, collapse = "\n") + ) + + manifest <- utils::read.delim( + file.path(discover_out_dir, "chunk_manifest.tsv"), check.names = FALSE, stringsAsFactors = FALSE ) - testthat::expect_equal(samples$alias, c("sampleA", "sampleB")) - testthat::expect_true(all(c("sampleA", "sampleB") %in% names(tx_counts))) - testthat::expect_gt(nrow(tx_counts), 0) + chunk_out_dir <- file.path(fixture_dir, "chunk1") + quant_result <- run_rscript( + "supeRglue", + c( + "bambu", + "quant", + "--chunk_rds", manifest$rds_path[[1]], + "--discovered_annotation_rds", file.path(discover_out_dir, "bambu_discovered_annotations.rds"), + "--genome", reference, + "--threads", "1", + "--out_dir", chunk_out_dir + ) + ) + testthat::expect_equal( + quant_result$status, + 0L, + info = paste(quant_result$output, collapse = "\n") + ) + + collate_out_dir <- file.path(fixture_dir, "sampleA") + collate_result <- run_rscript( + "supeRglue", + c( + "bambu", + "collate", + "--chunk_dirs", chunk_out_dir, + "--transcriptome_mode", "fixed_annotation", + "--out_dir", collate_out_dir + ) + ) + testthat::expect_equal( + collate_result$status, + 0L, + info = paste(collate_result$output, collapse = "\n") + ) + + testthat::expect_true(file.exists(file.path(collate_out_dir, "bambu_transcripts.rds"))) + testthat::expect_true(file.exists(file.path(collate_out_dir, "bambu_genes.rds"))) + testthat::expect_true(file.exists(file.path(collate_out_dir, "transcript_counts.tsv"))) + testthat::expect_true(file.exists(file.path(collate_out_dir, "gene_counts.tsv"))) }) # Ensembl GTF annotations include version numbers in IDs (e.g., ENST000001.7). diff --git a/modules/local/bambu_chunked.nf b/modules/local/bambu_chunked.nf new file mode 100644 index 0000000..be8a586 --- /dev/null +++ b/modules/local/bambu_chunked.nf @@ -0,0 +1,120 @@ +nextflow.enable.dsl = 2 + +OPTIONAL_FILE = file("$projectDir/data/OPTIONAL_FILE") + + +process bambuDiscover { + label "wf_transcriptomes" + cpus { + int requested = (params.threads ?: 4) as int + int sampleCount = aliases instanceof Collection ? aliases.size() : 1 + sampleCount > 1 ? requested : 1 + } + memory "60 GB" + input: + tuple val(meta), val(aliases), path(bams, stageAs: "bams/??.bam"), path(bais, stageAs: "bams/??.bam.bai"), path(sample_sheet) + path annotation, stageAs: "annotation/*" + path reference, stageAs: "reference/*" + output: + tuple val(meta), path("discover"), emit: dir + script: + def bam_list = bams instanceof Collection ? bams : [bams] + def alias_list = aliases instanceof Collection ? aliases : [aliases] + String bams_arg = "--bams '${bam_list.join(",")}'" + String aliases_arg = "--aliases '${alias_list.join(",")}'" + String sample_sheet_arg = sample_sheet.name == OPTIONAL_FILE.name ? "" : "--sample_sheet '${sample_sheet}'" + String ndr_arg = params.ndr != null ? "--ndr ${params.ndr}" : "" + """ + supeRglue bambu discover \ + ${bams_arg} \ + ${aliases_arg} \ + ${sample_sheet_arg} \ + --annotation "${annotation}" \ + --genome "${reference}" \ + --transcriptome_mode "${params.transcriptome_mode}" \ + --threads ${task.cpus} \ + ${ndr_arg} \ + --out_dir discover + """ +} + + +process bambuQuant { + label "wf_transcriptomes" + cpus { + int requested = (params.threads ?: 4) as int + boolean isJoint = meta instanceof Map && meta.alias == 'cohort' + isJoint ? requested : 1 + } + memory { ["8.GB", "16.GB", "48.GB"][task.attempt - 1] } + maxRetries 2 + errorStrategy 'retry' + input: + tuple val(meta), val(chunk_id), val(annotation_tx_count), path(chunk_rds), path(discovered_annotation) + path reference, stageAs: "reference/*" + output: + tuple val(meta), val(chunk_id), path("${chunk_id}"), emit: dir + script: + """ + supeRglue bambu quant \ + --chunk_rds "${chunk_rds}" \ + --discovered_annotation_rds "${discovered_annotation}" \ + --genome "${reference}" \ + --threads ${task.cpus} \ + --out_dir "${chunk_id}" + """ +} + + +process bambuEmpty { + label "wf_transcriptomes" + cpus 1 + memory "4 GB" + input: + tuple val(meta), val(aliases) + output: + tuple val(meta), path("${meta.alias}"), emit: dir + tuple val(meta), path("${meta.alias}/transcripts.gtf"), emit: gtf + tuple val(meta), path("${meta.alias}/transcript_counts.tsv"), emit: transcript_counts + tuple val(meta), path("${meta.alias}/gene_counts.tsv"), emit: gene_counts + tuple val(meta), path("${meta.alias}/bambu_transcripts.rds"), emit: transcript_rds + tuple val(meta), path("${meta.alias}/bambu_genes.rds"), emit: gene_rds + tuple val(meta), path("${meta.alias}/transcript_metadata.tsv"), emit: transcript_metadata + script: + def alias_list = aliases instanceof Collection ? aliases : [aliases] + String aliases_arg = "--aliases '${alias_list.join(",")}'" + """ + supeRglue bambu empty \ + ${aliases_arg} \ + --transcriptome_mode "${params.transcriptome_mode}" \ + --out_dir "${meta.alias}" + """ +} + + +process collateBambuQuant { + label "wf_transcriptomes" + cpus 1 + memory "16 GB" + input: + tuple val(meta), path(chunk_dirs, stageAs: "chunks/*") + output: + tuple val(meta), path("${meta.alias}"), emit: dir + tuple val(meta), path("${meta.alias}/transcripts.gtf"), emit: gtf + tuple val(meta), path("${meta.alias}/transcript_counts.tsv"), emit: transcript_counts + tuple val(meta), path("${meta.alias}/gene_counts.tsv"), emit: gene_counts + tuple val(meta), path("${meta.alias}/bambu_transcripts.rds"), emit: transcript_rds + tuple val(meta), path("${meta.alias}/bambu_genes.rds"), emit: gene_rds + tuple val(meta), path("${meta.alias}/transcript_metadata.tsv"), emit: transcript_metadata + script: + def chunk_dir_list = chunk_dirs instanceof Collection ? chunk_dirs : [chunk_dirs] + String chunk_dirs_arg = "--chunk_dirs '${chunk_dir_list.join(",")}'" + String ndr_arg = params.ndr != null ? "--ndr ${params.ndr}" : "" + """ + supeRglue bambu collate \ + ${chunk_dirs_arg} \ + --transcriptome_mode "${params.transcriptome_mode}" \ + ${ndr_arg} \ + --out_dir "${meta.alias}" + """ +} diff --git a/subworkflows/transcriptome.nf b/subworkflows/transcriptome.nf index bb578a7..a8d4be7 100644 --- a/subworkflows/transcriptome.nf +++ b/subworkflows/transcriptome.nf @@ -2,6 +2,84 @@ nextflow.enable.dsl = 2 OPTIONAL_FILE = file("$projectDir/data/OPTIONAL_FILE") +// use nextflows classic rename trick +include { + // joint + bambuDiscover as runJointBambuDiscover + bambuQuant as runJointBambuQuant + bambuEmpty as runJointBambuEmpty + collateBambuQuant as collateJointBambuQuant + // persample + bambuDiscover as runPerSampleBambuDiscover + bambuQuant as runPerSampleBambuQuant + bambuEmpty as runPerSampleBambuEmpty + collateBambuQuant as collatePerSampleBambuQuant +} from '../modules/local/bambu_chunked' + + +def bambu_discover_to_quant_inputs(discover_channel) { + discover_channel + .map { meta, discover_dir -> + def sampleAliases = discover_dir.resolve("samples.csv") + .readLines() + .drop(1) + .findAll { it?.trim() } + .collect { it.split(",", 2)[0] } + tuple( + meta, + sampleAliases, + discover_dir.resolve("bambu_discovered_annotations.rds"), + discover_dir.resolve("chunks"), + discover_dir.resolve("chunk_manifest.tsv") + ) + } + .splitCsv(header: true, sep: '\t', elem: 4) + .map { meta, sample_aliases, discovered_annotation, chunks_dir, row -> + def txCountRaw = row.annotation_tx_count?.toString()?.trim() + Integer annotationTxCount = (!txCountRaw || txCountRaw == 'NA') ? null : txCountRaw as Integer + tuple( + meta, + sample_aliases, + row.chunk_id, + annotationTxCount, + chunks_dir.resolve("${row.chunk_id}.rds"), + discovered_annotation + ) + } +} + + +def bambu_filter_quant_inputs_with_warning(quant_inputs) { + quant_inputs.filter { meta, sample_aliases, chunk_id, annotation_tx_count, chunk_rds, discovered_annotation -> + boolean keep = annotation_tx_count == null || annotation_tx_count > 0 + if (!keep) { + log.warn("Dropping bambu quant chunk '${chunk_id}' for '${meta.alias}' because annotation_tx_count=0") + } + keep + } +} + + +def bambu_empty_inputs(quant_inputs) { + quant_inputs + .map { meta, sample_aliases, chunk_id, annotation_tx_count, chunk_rds, discovered_annotation -> + tuple(meta.alias, meta, sample_aliases, annotation_tx_count) + } + .groupTuple() + .filter { alias, metas, sample_aliases_sets, annotation_tx_counts -> + !annotation_tx_counts.any { it == null || it > 0 } + } + .map { alias, metas, sample_aliases_sets, annotation_tx_counts -> + tuple(metas[0], sample_aliases_sets[0]) + } +} + + +def bambu_quant_process_inputs(quant_inputs) { + quant_inputs.map { meta, sample_aliases, chunk_id, annotation_tx_count, chunk_rds, discovered_annotation -> + tuple(meta, chunk_id, annotation_tx_count, chunk_rds, discovered_annotation) + } +} process prepareAnnotationReference { label "wf_transcriptomes" @@ -27,79 +105,6 @@ process prepareAnnotationReference { } -process runJointBambu { - label "wf_transcriptomes" - cpus { params.threads ?: 4 } - memory "32 GB" - input: - tuple val(aliases), path(bams, stageAs: "bams/??.bam"), path(bais, stageAs: "bams/??.bam.bai") - path sample_sheet - path annotation, stageAs: "annotation/*" - path reference, stageAs: "reference/*" - output: - path "cohort", emit: dir - path "cohort/transcripts.gtf", emit: gtf - path "cohort/transcript_counts.tsv", emit: transcript_counts - path "cohort/gene_counts.tsv", emit: gene_counts - path "cohort/bambu_transcripts.rds", emit: transcript_rds - path "cohort/bambu_genes.rds", emit: gene_rds - path "cohort/transcript_metadata.tsv", emit: transcript_metadata - script: - def bam_list = bams instanceof Collection ? bams : [bams] // todo dont run joint on single sample anyway - def alias_list = aliases instanceof Collection ? aliases : [aliases] - String bams_arg = "--bams '${bam_list.join(",")}'" - String aliases_arg = "--aliases '${alias_list.join(",")}'" - String sample_sheet_arg = sample_sheet.name == OPTIONAL_FILE.name ? "" : "--sample_sheet ${sample_sheet}" - String ndr_arg = params.ndr != null ? "--ndr ${params.ndr}" : "" - """ - supeRglue bambu \ - ${bams_arg} \ - ${aliases_arg} \ - ${sample_sheet_arg} \ - --annotation "${annotation}" \ - --genome "${reference}" \ - --transcriptome_mode "${params.transcriptome_mode}" \ - --threads ${task.cpus} \ - ${ndr_arg} \ - --out_dir cohort \ - """ -} - - -process runPerSampleBambu { - label "wf_transcriptomes" - cpus { params.threads ?: 4 } - memory "24 GB" - input: - tuple val(meta), path(bam), path(bai), path(stats) - path annotation, stageAs: "annotation/*" - path reference, stageAs: "reference/*" - output: - tuple val(meta), path("${meta.alias}"), emit: dir - tuple val(meta), path("${meta.alias}/transcripts.gtf"), emit: gtf - tuple val(meta), path("${meta.alias}/transcript_counts.tsv"), emit: transcript_counts - tuple val(meta), path("${meta.alias}/gene_counts.tsv"), emit: gene_counts - tuple val(meta), path("${meta.alias}/bambu_transcripts.rds"), emit: transcript_rds - tuple val(meta), path("${meta.alias}/bambu_genes.rds"), emit: gene_rds - tuple val(meta), path("${meta.alias}/transcript_metadata.tsv"), emit: transcript_metadata - script: - String bams_arg = "--bams '${bam.toString()}'" - String aliases_arg = "--aliases '${meta.alias}'" - String ndr_arg = params.ndr != null ? "--ndr ${params.ndr}" : "" - """ - supeRglue bambu \ - ${bams_arg} \ - ${aliases_arg} \ - --annotation "${annotation}" \ - --genome "${reference}" \ - --transcriptome_mode "${params.transcriptome_mode}" \ - --threads ${task.cpus} \ - ${ndr_arg} \ - --out_dir "${meta.alias}" - """ -} - - process buildCohortTranscriptomeFasta { label "wf_transcriptomes" cpus 1 @@ -214,33 +219,112 @@ workflow transcriptome_analysis { analysis_annotation = prepared_reference_annotation.annotation.first() analysis_reference = prepared_reference_annotation.reference.first() - joint_bambu = runJointBambu( + joint_meta = [alias: "cohort"] + + joint_discover = runJointBambuDiscover( alignments - | collect(flat: false) - | map { rows -> - // transform [meta, bam, bai] to [[alias1...aliasN], [bam1...bamN], [bai1...baiN]] - tuple( - rows.collect { it[0].alias }, - rows.collect { it[1] }, - rows.collect { it[2] } - ) - }, - sample_sheet, + .collect(flat: false) + .map { rows -> + // transform [meta, bam, bai] rows to + // [meta, [alias1...aliasN], [bam1...bamN], [bai1...baiN], sample_sheet] + tuple( + joint_meta, + rows.collect { it[0].alias }, + rows.collect { it[1] }, + rows.collect { it[2] }, + sample_sheet + ) + }, analysis_annotation, analysis_reference ) + joint_quant_inputs_all = bambu_discover_to_quant_inputs(joint_discover.dir) + joint_quant = runJointBambuQuant( + bambu_quant_process_inputs( + bambu_filter_quant_inputs_with_warning(joint_quant_inputs_all) + ), + analysis_reference + ) + joint_bambu_real = collateJointBambuQuant( + joint_quant.dir + .map { meta, chunk_id, chunk_dir -> + tuple(meta.alias, meta, chunk_dir) + } + .groupTuple() + .map { alias, metas, chunk_dirs -> + tuple(metas[0], chunk_dirs) + } + ) + joint_bambu_empty = runJointBambuEmpty( + bambu_empty_inputs(joint_quant_inputs_all) + ) - sample_bambu = runPerSampleBambu(alignments, analysis_annotation, analysis_reference) + joint_bambu_dir = joint_bambu_real.dir.mix(joint_bambu_empty.dir) + joint_bambu_gtf = joint_bambu_real.gtf.mix(joint_bambu_empty.gtf) + joint_bambu_transcript_counts = joint_bambu_real.transcript_counts.mix(joint_bambu_empty.transcript_counts) + joint_bambu_gene_counts = joint_bambu_real.gene_counts.mix(joint_bambu_empty.gene_counts) + joint_bambu_transcript_rds = joint_bambu_real.transcript_rds.mix(joint_bambu_empty.transcript_rds) + joint_bambu_gene_rds = joint_bambu_real.gene_rds.mix(joint_bambu_empty.gene_rds) + joint_bambu_metadata = joint_bambu_real.transcript_metadata.mix(joint_bambu_empty.transcript_metadata) - joint_fasta = buildCohortTranscriptomeFasta(joint_bambu.gtf, analysis_reference) - sample_fastas = buildSampleTranscriptomeFasta(sample_bambu.gtf, analysis_reference) + sample_discover = runPerSampleBambuDiscover( + alignments.map { meta, bam, bai, stats -> + tuple( + meta, + [meta.alias], + bam, + bai, + sample_sheet + ) + }, + analysis_annotation, + analysis_reference + ) + sample_quant_inputs_all = bambu_discover_to_quant_inputs(sample_discover.dir) + sample_quant = runPerSampleBambuQuant( + bambu_quant_process_inputs( + bambu_filter_quant_inputs_with_warning(sample_quant_inputs_all) + ), + analysis_reference + ) + sample_bambu_real = collatePerSampleBambuQuant( + sample_quant.dir + .map { meta, chunk_id, chunk_dir -> + tuple(meta.alias, meta, chunk_dir) + } + .groupTuple() + .map { alias, metas, chunk_dirs -> + tuple(metas[0], chunk_dirs) + } + ) + sample_bambu_empty = runPerSampleBambuEmpty( + bambu_empty_inputs(sample_quant_inputs_all) + ) + + sample_bambu_dirs = sample_bambu_real.dir.mix(sample_bambu_empty.dir) + sample_bambu_gtf = sample_bambu_real.gtf.mix(sample_bambu_empty.gtf) + sample_bambu_transcript_counts = sample_bambu_real.transcript_counts.mix(sample_bambu_empty.transcript_counts) + sample_bambu_gene_counts = sample_bambu_real.gene_counts.mix(sample_bambu_empty.gene_counts) + sample_bambu_transcript_rds = sample_bambu_real.transcript_rds.mix(sample_bambu_empty.transcript_rds) + sample_bambu_gene_rds = sample_bambu_real.gene_rds.mix(sample_bambu_empty.gene_rds) + sample_bambu_metadata = sample_bambu_real.transcript_metadata.mix(sample_bambu_empty.transcript_metadata) + + joint_fasta = buildCohortTranscriptomeFasta( + joint_bambu_real.gtf.map { meta, gtf -> gtf }, + analysis_reference + ) + sample_fastas = buildSampleTranscriptomeFasta(sample_bambu_real.gtf, analysis_reference) if (params.skip_sqanti) { joint_sqanti_dir = Channel.empty() sample_sqanti_dirs = Channel.empty() } else { - joint_sqanti = runJointSqanti(joint_bambu.gtf, analysis_annotation, analysis_reference) - sample_sqanti = runPerSampleSqanti(sample_bambu.gtf, analysis_annotation, analysis_reference) + joint_sqanti = runJointSqanti( + joint_bambu_real.gtf.map { meta, gtf -> gtf }, + analysis_annotation, + analysis_reference + ) + sample_sqanti = runPerSampleSqanti(sample_bambu_real.gtf, analysis_annotation, analysis_reference) joint_sqanti_dir = joint_sqanti.dir sample_sqanti_dirs = sample_sqanti.dir } @@ -249,22 +333,22 @@ workflow transcriptome_analysis { annotation = analysis_annotation annotation_reference_summary = prepared_reference_annotation.summary unstranded_annotation = prepared_reference_annotation.unstranded - joint_dir = joint_bambu.dir - joint_gtf = joint_bambu.gtf + joint_dir = joint_bambu_dir.map { meta, dir -> dir } + joint_gtf = joint_bambu_gtf.map { meta, gtf -> gtf } joint_fasta = joint_fasta.fasta - joint_transcript_counts = joint_bambu.transcript_counts - joint_gene_counts = joint_bambu.gene_counts - joint_transcript_rds = joint_bambu.transcript_rds - joint_gene_rds = joint_bambu.gene_rds - joint_metadata = joint_bambu.transcript_metadata - sample_dirs = sample_bambu.dir - sample_gtf = sample_bambu.gtf + joint_transcript_counts = joint_bambu_transcript_counts.map { meta, counts -> counts } + joint_gene_counts = joint_bambu_gene_counts.map { meta, counts -> counts } + joint_transcript_rds = joint_bambu_transcript_rds.map { meta, rds -> rds } + joint_gene_rds = joint_bambu_gene_rds.map { meta, rds -> rds } + joint_metadata = joint_bambu_metadata.map { meta, metadata -> metadata } + sample_dirs = sample_bambu_dirs + sample_gtf = sample_bambu_gtf sample_fastas = sample_fastas.fasta - sample_transcript_counts = sample_bambu.transcript_counts - sample_gene_counts = sample_bambu.gene_counts - sample_transcript_rds = sample_bambu.transcript_rds - sample_gene_rds = sample_bambu.gene_rds - sample_metadata = sample_bambu.transcript_metadata + sample_transcript_counts = sample_bambu_transcript_counts + sample_gene_counts = sample_bambu_gene_counts + sample_transcript_rds = sample_bambu_transcript_rds + sample_gene_rds = sample_bambu_gene_rds + sample_metadata = sample_bambu_metadata joint_sqanti_dir = joint_sqanti_dir sample_sqanti_dirs = sample_sqanti_dirs }