diff --git a/R/data_curation.R b/R/data_curation.R index f056c37..5a664c0 100644 --- a/R/data_curation.R +++ b/R/data_curation.R @@ -36,17 +36,26 @@ .ftpes_download_one <- function(genomeID, out_dir, connect_timeout = 10L, max_time = 30L, - speed_time = 30L, # end if speed_time - speed_limit = 2048L, # B/s + speed_time = 30L, + speed_limit = 2048L, min_bytes = 100L) { dir.create(out_dir, recursive = TRUE, showWarnings = FALSE) + # Adding this prior dot function helper into this loop for parallel workers + is_complete <- function(dir, genomeID, min_bytes = 100L) { + fna <- file.path(dir, paste0(genomeID, ".fna")) + faa <- file.path(dir, paste0(genomeID, ".PATRIC.faa")) + gff <- file.path(dir, paste0(genomeID, ".PATRIC.gff")) + paths <- c(fna, faa, gff) + all(file.exists(paths)) && + all(vapply(paths, function(x) file.info(x)$size, numeric(1)) > min_bytes) + } + exts <- c(".fna", ".PATRIC.faa", ".PATRIC.gff") for (ext in exts) { dest <- file.path(out_dir, paste0(genomeID, ext)) dest_tmp <- paste0(dest, ".tmp") - # Skip if we already completed these if (file.exists(dest) && file.info(dest)$size > min_bytes) next if (file.exists(dest) && file.info(dest)$size == 0) try(unlink(dest), silent = TRUE) if (file.exists(dest_tmp)) try(unlink(dest_tmp), silent = TRUE) @@ -59,7 +68,7 @@ "--max-time", as.character(max_time), "--speed-time", as.character(speed_time), "--speed-limit", as.character(speed_limit), - "--ftp-ssl", # AUTH TLS on port 21 works in testing + "--ftp-ssl", "--ftp-pasv", "--disable-epsv", "--ipv4", "--user", "anonymous:", "-o", shQuote(dest_tmp), shQuote(url) @@ -69,22 +78,19 @@ status <- attr(res, "status") if (is.null(status)) status <- 0L - # Avoiding 0B files cluttering up results -- atomic rename on successful tmp DL ok <- (status == 0L && file.exists(dest_tmp) && file.info(dest_tmp)$size > min_bytes) if (ok) { if (!file.rename(dest_tmp, dest)) { - # If rename fails for some reason, copy + unlink file.copy(dest_tmp, dest, overwrite = TRUE) unlink(dest_tmp) } } else { try(unlink(dest_tmp), silent = TRUE) - return(FALSE) # Abort early, don't accept partial set + return(FALSE) } } - # Do we have all 3 present? - .is_complete_set(out_dir, genomeID, min_bytes = min_bytes) + is_complete(out_dir, genomeID, min_bytes = min_bytes) } #' Helps manage FTPS downloading from BV-BRC, tryng a quick download first, and @@ -92,8 +98,8 @@ #' hiccup. If 2nd pass fails, log and give up on that file. #' @keywords internal .ftpes_download_two_pass <- function(genome_ids, out_dir, - workers_first = 4L, - workers_second = 4L, + workers_first = 8L, + workers_second = 8L, log_file = NULL) { genome_ids <- unique(as.character(genome_ids)) if (!length(genome_ids)) { @@ -107,13 +113,16 @@ ) } - # Pass 1: 30s per-file cap + old_plan <- future::plan() + on.exit(future::plan(old_plan), add = TRUE) + message("FTPS pass 1 (45s timeout)") - future::plan(future::multisession, workers = max(1, workers_first)) + future::plan(future::multisession, workers = max(1L, workers_first)) res1 <- future.apply::future_lapply( genome_ids, function(gid) { - ok <- .ftpes_download_one(gid, out_dir, + ok <- .ftpes_download_one( + gid, out_dir, connect_timeout = 10L, max_time = 45L, speed_time = 30L, speed_limit = 2048L ) @@ -121,9 +130,11 @@ }, future.seed = TRUE ) + ok1 <- vapply(res1, `[[`, logical(1), "ok") ok_ids_1 <- genome_ids[ok1] fail_ids <- genome_ids[!ok1] + message(sprintf("Pass 1: ok=%d, fail=%d", length(ok_ids_1), length(fail_ids))) if (!is.null(log_file)) { cat(sprintf("[%s] Pass1 ok=%d fail=%d\n", Sys.time(), length(ok_ids_1), length(fail_ids)), @@ -140,13 +151,13 @@ return(ok_ids_1) } - # Pass 2: 60s per-file cap where we retry any failures message("FTPS pass 2 (120s timeout) for failed genomes") - future::plan(future::multisession, workers = max(1, workers_second)) + future::plan(future::multisession, workers = max(1L, workers_second)) res2 <- future.apply::future_lapply( fail_ids, function(gid) { - ok <- .ftpes_download_one(gid, out_dir, + ok <- .ftpes_download_one( + gid, out_dir, connect_timeout = 10L, max_time = 120L, speed_time = 30L, speed_limit = 2048L ) @@ -154,9 +165,11 @@ }, future.seed = TRUE ) + ok2 <- vapply(res2, `[[`, logical(1), "ok") ok_ids_2 <- fail_ids[ok2] still_fail <- setdiff(fail_ids, ok_ids_2) + message(sprintf("Pass 2: ok=%d, still_fail=%d", length(ok_ids_2), length(still_fail))) if (!is.null(log_file)) { cat(sprintf("[%s] Pass2 ok=%d still_fail=%d\n", Sys.time(), length(ok_ids_2), length(still_fail)), @@ -234,6 +247,286 @@ return(df) } +#' Parse BV-BRC TSV output +#' +#' Safely cleans blank links and weird formatting for a safe merge of the BV-BRC +#' TSV data into a tibble. +#' +#' @param x Character vector of output lines from a BV-BRC CLI command. +#' +#' @return A tibble with minimal name repair. +#' @keywords internal +.parse_bvbrc_tsv <- function(x) { + if (!length(x)) { + return(tibble::tibble()) + } + + x <- as.character(x) + x <- x[nzchar(trimws(x))] + + # Keep only lines that look like TSV records/headers + x <- x[grepl("\t", x, fixed = TRUE)] + + if (!length(x)) { + return(tibble::tibble()) + } + + txt <- paste(x, collapse = "\n") + df <- utils::read.table( + text = txt, + sep = "\t", + header = TRUE, + fill = TRUE, + quote = "", + check.names = FALSE, + comment.char = "", + colClasses = "character" + ) + + tibble::as_tibble(df, .name_repair = "minimal") +} + + +#' Apply QC filters to BV-BRC genome metadata +#' +#' Quality control step with default values for contamination and completeness. +#' Optional filters for other genome stats to find genomes that are X deviations +#' away from the median values for those stats. This shouldn't be applied to +#' jobs that use more than one diverse taxon at a time though! Requires a fairly +#' normal distribution like you'd see for (most) single species +#' Returns the cleaned table, retained genome IDs, and a rejection log. +#' +#' CheckM contamination and completeness are treated as the primary QC gate. +#' Deviation-based filters are optional and are best suited to single-taxon runs. +#' +#' @param genome_tbl A tibble of genome metadata containing BV-BRC genome columns. +#' @param checkm_contam Numeric scalar. Maximum allowed CheckM contamination (%). +#' @param checkm_complete Numeric scalar. Minimum allowed CheckM completeness (%). +#' @param gc_deviations Optional numeric scalar. Maximum SDs from the median GC content. +#' @param length_deviations Optional numeric scalar. Maximum SDs from the median genome length. +#' @param cds_deviations Optional numeric scalar. Maximum SDs from the median CDS count. +#' +#' @return A list with components: +#' \describe{ +#' \item{qc_tbl}{QC-annotated genome table} +#' \item{keep_ids}{Character vector of retained genome IDs} +#' \item{rejections}{Tibble describing dropped genomes and reasons} +#' } +#' @keywords internal +.apply_metadata_qc <- function(genome_tbl, + checkm_contam = 10, + checkm_complete = 90, + gc_deviations = NULL, + length_deviations = NULL, + cds_deviations = NULL) { + numify <- function(x) suppressWarnings(as.numeric(x)) + + safe_median <- function(x) { + x <- x[is.finite(x)] + if (!length(x)) { + return(NA_real_) + } + stats::median(x) + } + + safe_sd <- function(x) { + x <- x[is.finite(x)] + if (length(x) < 2L) { + return(NA_real_) + } + stats::sd(x) + } + + qc <- tibble::as_tibble(genome_tbl) |> + dplyr::mutate( + genome_length_qc = numify(.data$`genome.genome_length`), + gc_content_qc = numify(.data$`genome.gc_content`), + cds_qc = numify(.data$`genome.cds`), + checkm_completeness_qc = numify(.data$`genome.checkm_completeness`), + checkm_contamination_qc = numify(.data$`genome.checkm_contamination`) + ) + + med_len <- safe_median(qc$genome_length_qc) + sd_len <- safe_sd(qc$genome_length_qc) + med_gc <- safe_median(qc$gc_content_qc) + sd_gc <- safe_sd(qc$gc_content_qc) + med_cds <- safe_median(qc$cds_qc) + sd_cds <- safe_sd(qc$cds_qc) + + z_len <- if (!is.null(length_deviations) && is.finite(sd_len) && sd_len > 0) { + abs(qc$genome_length_qc - med_len) / sd_len + } else { + rep(NA_real_, nrow(qc)) + } + + z_gc <- if (!is.null(gc_deviations) && is.finite(sd_gc) && sd_gc > 0) { + abs(qc$gc_content_qc - med_gc) / sd_gc + } else { + rep(NA_real_, nrow(qc)) + } + + z_cds <- if (!is.null(cds_deviations) && is.finite(sd_cds) && sd_cds > 0) { + abs(qc$cds_qc - med_cds) / sd_cds + } else { + rep(NA_real_, nrow(qc)) + } + + qc <- qc |> + dplyr::mutate( + flag_checkm_missing = is.na(.data$checkm_completeness_qc) | is.na(.data$checkm_contamination_qc), + flag_checkm_contam = !is.na(.data$checkm_contamination_qc) & .data$checkm_contamination_qc > checkm_contam, + flag_checkm_complete = !is.na(.data$checkm_completeness_qc) & .data$checkm_completeness_qc < checkm_complete, + flag_checkm = .data$flag_checkm_missing | .data$flag_checkm_contam | .data$flag_checkm_complete, + flag_length = if (!is.null(length_deviations) && is.finite(sd_len) && sd_len > 0) { + !is.na(z_len) & z_len > length_deviations + } else { + FALSE + }, + flag_gc = if (!is.null(gc_deviations) && is.finite(sd_gc) && sd_gc > 0) { + !is.na(z_gc) & z_gc > gc_deviations + } else { + FALSE + }, + flag_cds = if (!is.null(cds_deviations) && is.finite(sd_cds) && sd_cds > 0) { + !is.na(z_cds) & z_cds > cds_deviations + } else { + FALSE + }, + qc_keep = !(.data$flag_checkm | .data$flag_length | .data$flag_gc | .data$flag_cds) + ) + + make_rej <- function(flag, rule, observed, threshold, comparator) { + idx <- which(flag %in% TRUE) + if (!length(idx)) { + return(NULL) + } + tibble::tibble( + genome_id = qc$`genome.genome_id`[idx], + genome_name = qc$`genome.genome_name`[idx], + species = qc$`genome.species`[idx], + failed_rule = rule, + observed = as.character(observed[idx]), + threshold = as.character(threshold), + comparator = comparator, + action_taken = "dropped" + ) + } + + rej_list <- list( + make_rej( + qc$flag_checkm_missing, + "checkm_missing", + ifelse(qc$flag_checkm_missing, "missing", NA_character_), + "present", + "is present" + ), + make_rej( + qc$flag_checkm_contam, + "checkm_contamination", + qc$checkm_contamination_qc, + checkm_contam, + ">" + ), + make_rej( + qc$flag_checkm_complete, + "checkm_completeness", + qc$checkm_completeness_qc, + checkm_complete, + "<" + ) + ) + + if (!is.null(gc_deviations)) { + rej_list[[length(rej_list) + 1L]] <- make_rej( + qc$flag_gc, + "gc_deviation_sd", + z_gc, + gc_deviations, + ">" + ) + } + + if (!is.null(length_deviations)) { + rej_list[[length(rej_list) + 1L]] <- make_rej( + qc$flag_length, + "length_deviation_sd", + z_len, + length_deviations, + ">" + ) + } + + if (!is.null(cds_deviations)) { + rej_list[[length(rej_list) + 1L]] <- make_rej( + qc$flag_cds, + "cds_deviation_sd", + z_cds, + cds_deviations, + ">" + ) + } + + rejections <- dplyr::bind_rows(rej_list) + if (nrow(rejections) == 0L) { + rejections <- tibble::tibble( + genome_id = character(), + genome_name = character(), + species = character(), + failed_rule = character(), + observed = character(), + threshold = character(), + comparator = character(), + action_taken = character() + ) + } + + list( + qc_tbl = qc, + keep_ids = unique(qc$`genome.genome_id`[qc$qc_keep]), + rejections = rejections + ) +} + +# Get rid of genomes that are missing files +.purge_genome_files <- function(genome_path, genome_ids, log_file = NULL) { + genome_ids <- unique(as.character(genome_ids)) + if (!length(genome_ids)) { + return(invisible(0L)) + } + + exts <- c( + ".fna", + ".PATRIC.faa", + ".PATRIC.gff", + ".gto", + ".orig2id.tsv", + ".fna.tmp", + ".PATRIC.faa.tmp", + ".PATRIC.gff.tmp" + ) + + n_removed <- 0L + for (gid in genome_ids) { + paths <- file.path(genome_path, paste0(gid, exts)) + existing <- paths[file.exists(paths)] + if (length(existing)) { + unlink(existing, force = TRUE) + n_removed <- n_removed + length(existing) + } + } + + if (!is.null(log_file) && length(genome_ids)) { + cat( + sprintf( + "[%s] Purged incomplete genomes: %s\n", + Sys.time(), paste(head(genome_ids, 50), collapse = ", ") + ), + file = log_file, append = TRUE + ) + } + + invisible(n_removed) +} # Make sure the BV-BRC metadata live where they're supposed to .ensure_bvbrc_cache <- function(base_dir = ".", @@ -910,7 +1203,13 @@ #' @param abx Character or vector. Antibiotic filter. "All" for all antibiotics, else names. #' @param overwrite Logical. If FALSE and DuckDB exists already, abort. Default FALSE. #' @param image Character. Docker image. Default "danylmb/bvbrc:5.3". -#' @param verbose Logical. If TRUE, prints concise messages. +#' @param checkm_contam Numeric scalar. Maximum allowed CheckM contamination (%). +#' @param checkm_complete Numeric scalar. Minimum allowed CheckM completeness (%). +#' @param gc_deviations Optional numeric scalar. Maximum SDs from the median GC content. +#' @param length_deviations Optional numeric scalar. Maximum SDs from the median genome length. +#' @param cds_deviations Optional numeric scalar. Maximum SDs from the median CDS count. +#' @param debug Logical. If TRUE, retain `metadata_full` and QC columns for inspection. +#' @param verbose Logical. If TRUE, print progress messages. #' #' @return A list with: #' - duckdbConnection: live DBI connection to the created DuckDB @@ -923,27 +1222,27 @@ retrieveMetadata <- function(user_bacs, abx = "All", overwrite = FALSE, image = "danylmb/bvbrc:5.3", + checkm_contam = 10, + checkm_complete = 90, + gc_deviations = NULL, + length_deviations = NULL, + cds_deviations = NULL, + debug = FALSE, verbose = TRUE) { base_dir <- normalizePath(base_dir, mustWork = FALSE) - # ------------------------------- - # GENOME ID RESOLUTION (UPDATED) - # ------------------------------- if (!is.null(genome_id_file)) { if (!file.exists(genome_id_file)) { stop("Provided genome_id_file does not exist.") } - if (isTRUE(verbose)) { message("Using genome IDs from file: ", genome_id_file) } - genome_ids <- readLines(genome_id_file, warn = FALSE) genome_ids <- trimws(genome_ids) genome_ids <- genome_ids[genome_ids != ""] } else { if (isTRUE(verbose)) message("Resolving genome IDs for user inputs.") - genome_ids <- .retrieveQueryIDs( base_dir = base_dir, user_bacs = user_bacs, @@ -953,13 +1252,11 @@ retrieveMetadata <- function(user_bacs, } genome_ids <- unique(as.character(genome_ids)) - - if (length(genome_ids) == 0) { + if (length(genome_ids) == 0L) { message("No genome IDs available for the specified inputs.") return(NULL) } - # Desired fields from BV-BRC drug_fields <- paste0( "antibiotic,computational_method,", "evidence,genome_name,id,", @@ -970,11 +1267,13 @@ retrieveMetadata <- function(user_bacs, "pmid,resistant_phenotype,", "source,taxon_id,testing_standard" ) + abx_filter <- if (identical(abx, "All")) { "--required antibiotic" } else { paste0("--in antibiotic,", paste(abx, collapse = ",")) } + amr_fields <- paste0( "assembly_accession,assembly_method,", "bioproject_accession,biosample_accession,", @@ -991,6 +1290,7 @@ retrieveMetadata <- function(user_bacs, "refseq_accessions,", "refseq_project_id,sra_accession,species,taxon_id" ) + microtrait_fields <- paste0( amr_fields, ",genome_length,gram_stain,", "habitat,", @@ -1009,70 +1309,123 @@ retrieveMetadata <- function(user_bacs, "trna" ) - # Batching downloads and parallel implementation + qc_fields <- c( + "genome_length", + "gc_content", + "contigs", + "cds", + "rrna", + "trna", + "checkm_completeness", + "checkm_contamination" + ) + + amr_fields <- paste( + unique(c(strsplit(amr_fields, ",", fixed = TRUE)[[1]], qc_fields)), + collapse = "," + ) + microtrait_fields <- paste( + unique(c(strsplit(microtrait_fields, ",", fixed = TRUE)[[1]], qc_fields)), + collapse = "," + ) + batch_size <- 500L - genome_ids <- as.character(genome_ids) genome_batches <- split(genome_ids, ceiling(seq_along(genome_ids) / batch_size)) n_cores <- max(1L, parallel::detectCores(logical = TRUE) - 1L) - param <- BiocParallel::SnowParam(workers = n_cores) - # Pull AMR metadata + old_plan <- future::plan() + on.exit(future::plan(old_plan), add = TRUE) + future::plan(future::multisession, workers = n_cores) + if (isTRUE(verbose)) message("Retrieving AMR phenotype data in batches.") - batch_drug_data <- BiocParallel::bplapply(genome_batches, function(batch) { - .extractAMRtable( - base_dir = base_dir, - batch_genome_IDs = batch, - abx_filter = abx_filter, - drug_fields = drug_fields, - image = image, - verbose = FALSE - ) - }, BPPARAM = param) - combined_drug_data <- unlist(batch_drug_data, use.names = FALSE) - if (length(combined_drug_data) == 0) { + batch_drug_data <- future.apply::future_lapply( + genome_batches, + function(batch) { + raw <- .extractAMRtable( + base_dir = base_dir, + batch_genome_IDs = batch, + abx_filter = abx_filter, + drug_fields = drug_fields, + image = image, + verbose = FALSE + ) + .parse_bvbrc_tsv(raw) + }, + future.seed = TRUE + ) + + combined_drug_data_tbl <- dplyr::bind_rows(batch_drug_data) |> + dplyr::mutate(dplyr::across(dplyr::everything(), ~ iconv(.x, from = "", to = "UTF-8", sub = ""))) + + if (nrow(combined_drug_data_tbl) == 0L) { message("No drug data returned.") return(NULL) } - combined_drug_data_tbl <- tibble::as_tibble(utils::read.table( - text = combined_drug_data, - sep = "\t", header = TRUE, fill = TRUE, - quote = "", check.names = FALSE, comment.char = "", colClasses = "character" - )) |> - dplyr::mutate(dplyr::across(dplyr::everything(), ~ iconv(.x, from = "", to = "UTF-8", sub = ""))) |> - dplyr::mutate(`genome_drug.genome_id` = as.character(`genome_drug.genome_id`)) - # Pull genome metadata if (isTRUE(verbose)) message("Retrieving genome metadata in batches.") - batch_genome_data <- BiocParallel::bplapply(genome_batches, function(batch) { - .extractGenomeData( - base_dir = base_dir, - batch_genome_IDs = batch, - filter_type = filter_type, - amr_fields = amr_fields, - microtrait_fields = microtrait_fields, - image = image, - verbose = FALSE - ) - }, BPPARAM = param) - combined_genome_data <- unlist(batch_genome_data, use.names = FALSE) - if (length(combined_genome_data) == 0) { + batch_genome_data <- future.apply::future_lapply( + genome_batches, + function(batch) { + raw <- .extractGenomeData( + base_dir = base_dir, + batch_genome_IDs = batch, + filter_type = filter_type, + amr_fields = amr_fields, + microtrait_fields = microtrait_fields, + image = image, + verbose = FALSE + ) + .parse_bvbrc_tsv(raw) + }, + future.seed = TRUE + ) + + combined_genome_data_tbl <- dplyr::bind_rows(batch_genome_data) |> + dplyr::mutate(dplyr::across(dplyr::everything(), ~ iconv(.x, from = "", to = "UTF-8", sub = ""))) + + if (nrow(combined_genome_data_tbl) == 0L) { message("No genome data returned.") return(NULL) } - combined_genome_data_tbl <- tibble::as_tibble(utils::read.table( - text = combined_genome_data, - sep = "\t", header = TRUE, fill = TRUE, - quote = "", check.names = FALSE, comment.char = "", colClasses = "character" - )) |> - dplyr::mutate(dplyr::across(dplyr::everything(), ~ iconv(.x, from = "", to = "UTF-8", sub = ""))) |> - dplyr::mutate(`genome.genome_id` = as.character(`genome.genome_id`)) - # Write & join + if ("genome.genome_id" %in% names(combined_genome_data_tbl)) { + combined_genome_data_tbl <- combined_genome_data_tbl |> + dplyr::filter(grepl("^[0-9]+\\.[0-9]+$", .data$`genome.genome_id`)) |> + dplyr::distinct(`genome.genome_id`, .keep_all = TRUE) + } + + if ("genome_drug.genome_id" %in% names(combined_drug_data_tbl)) { + combined_drug_data_tbl <- combined_drug_data_tbl |> + dplyr::filter(grepl("^[0-9]+\\.[0-9]+$", .data$`genome_drug.genome_id`)) + } + + qc_out <- .apply_metadata_qc( + genome_tbl = combined_genome_data_tbl, + checkm_contam = checkm_contam, + checkm_complete = checkm_complete, + gc_deviations = gc_deviations, + length_deviations = length_deviations, + cds_deviations = cds_deviations + ) + + keep_ids <- qc_out$keep_ids + if (!length(keep_ids)) { + message("No genomes passed QC.") + return(NULL) + } + + combined_genome_data_tbl <- combined_genome_data_tbl |> + dplyr::filter(`genome.genome_id` %in% keep_ids) + + combined_drug_data_tbl <- combined_drug_data_tbl |> + dplyr::filter(`genome_drug.genome_id` %in% keep_ids) + paths <- .buildDBpath(base_dir = base_dir, user_bacs = user_bacs, overwrite = overwrite) db_path <- paths$db_path logs_dir <- file.path(base_dir, "data", "logs") dir.create(logs_dir, recursive = TRUE, showWarnings = FALSE) + cat(sprintf("[%s] Writing metadata DuckDB: %s\n", Sys.time(), db_path), file = file.path(logs_dir, "bvbrc.log"), append = TRUE ) @@ -1082,48 +1435,88 @@ retrieveMetadata <- function(user_bacs, DBI::dbWriteTable(con, "amr_phenotype", combined_drug_data_tbl, overwrite = TRUE) DBI::dbWriteTable(con, "genome_data", combined_genome_data_tbl, overwrite = TRUE) + DBI::dbWriteTable(con, "metadata_qc", qc_out$qc_tbl, overwrite = TRUE) + DBI::dbWriteTable(con, "metadata_qc_rejections", qc_out$rejections, overwrite = TRUE) if (isTRUE(verbose)) message("Joining AMR phenotype and genome metadata.") DBI::dbExecute(con, ' - CREATE OR REPLACE TABLE metadata AS + CREATE OR REPLACE TABLE metadata_full AS SELECT * FROM amr_phenotype INNER JOIN genome_data ON amr_phenotype."genome_drug.genome_id" = genome_data."genome.genome_id" ') - # Debug summary after writes + qc_drop_cols <- c( + "genome.genome_length", + "genome.gc_content", + "genome.contigs", + "genome.cds", + "genome.rrna", + "genome.trna", + "genome.checkm_completeness", + "genome.checkm_contamination" + ) + + if (isTRUE(debug)) { + DBI::dbExecute(con, "CREATE OR REPLACE TABLE metadata AS SELECT * FROM metadata_full") + } else { + keep_cols <- setdiff(DBI::dbListFields(con, "metadata_full"), qc_drop_cols) + select_sql <- paste(DBI::dbQuoteIdentifier(con, keep_cols), collapse = ", ") + DBI::dbExecute( + con, + paste0("CREATE OR REPLACE TABLE metadata AS SELECT ", select_sql, " FROM metadata_full") + ) + DBI::dbExecute(con, "DROP TABLE metadata_full") + } + n_targets <- length(genome_ids) - n_amr_ids <- DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome_drug.genome_id") AS n FROM amr_phenotype')$n - n_gmeta_ids <- DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM genome_data')$n + n_amr_ids <- DBI::dbGetQuery(con, ' + SELECT COUNT(DISTINCT "genome_drug.genome_id") AS n FROM amr_phenotype + ')$n + n_gmeta_ids <- DBI::dbGetQuery(con, ' + SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM genome_data + ')$n + if (isTRUE(verbose)) { message( "Initial summary: targets=", n_targets, " | AMR genomes=", n_amr_ids, " | genome_data genomes=", n_gmeta_ids ) + message( + "QC summary: kept=", length(keep_ids), + " | dropped=", nrow(qc_out$rejections) + ) } - # Tagged view .create_amr_tagged_view(con) - # Final debug summary n_bac <- if ("bac_data" %in% DBI::dbListTables(con)) { - DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM bac_data')$n + DBI::dbGetQuery(con, ' + SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM bac_data + ')$n } else { NA_integer_ } - n_amr_ids <- DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome_drug.genome_id") AS n FROM amr_phenotype')$n - n_gmeta_ids <- DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM genome_data')$n + + n_amr_ids <- DBI::dbGetQuery(con, ' + SELECT COUNT(DISTINCT "genome_drug.genome_id") AS n FROM amr_phenotype + ')$n + n_gmeta_ids <- DBI::dbGetQuery(con, ' + SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM genome_data + ')$n ids_zero_amr <- DBI::dbGetQuery( con, - 'WITH sel AS (SELECT DISTINCT "genome.genome_id" AS gid FROM genome_data), - got AS (SELECT DISTINCT "genome_drug.genome_id" AS gid FROM amr_phenotype) - SELECT sel.gid - FROM sel LEFT JOIN got USING (gid) - WHERE got.gid IS NULL - ORDER BY sel.gid' + ' + WITH sel AS (SELECT DISTINCT "genome.genome_id" AS gid FROM genome_data), + got AS (SELECT DISTINCT "genome_drug.genome_id" AS gid FROM amr_phenotype) + SELECT sel.gid + FROM sel LEFT JOIN got USING (gid) + WHERE got.gid IS NULL + ORDER BY sel.gid + ' )$gid if (isTRUE(verbose)) { @@ -1546,8 +1939,8 @@ retrieveMetadata <- function(user_bacs, #' @param image Docker image for CLI path (default "danylmb/bvbrc:5.3"). #' @param skip_existing Logical; if TRUE, do not re-download genomes already complete. Default TRUE. #' @param ftp_workers Parallel workers for FTP path (default 8). -#' @param cli_fasta_workers Parallel chunk containers for FASTA+GTO (default 4). -#' @param cli_gff_workers Parallel chunk containers for GFF export (default 4). +#' @param cli_fasta_workers Parallel chunk containers for FASTA+GTO (default 8). +#' @param cli_gff_workers Parallel chunk containers for GFF export (default 8). #' @param chunk_size Genomes per chunk container (default 50). #' @param verbose Verbose messages. #' @return Character vector of genome IDs with complete file sets on disk. @@ -1557,9 +1950,9 @@ retrieveGenomes <- function(base_dir = ".", method = c("ftp", "cli"), image = "danylmb/bvbrc:5.3", skip_existing = TRUE, - ftp_workers = 4L, - cli_fasta_workers = 4L, - cli_gff_workers = 4L, + ftp_workers = 8L, + cli_fasta_workers = 8L, + cli_gff_workers = 8L, chunk_size = 50L, evidence_mode = c("lab_only", "lab_or_comp", "comp_only", "any"), # NEW verbose = TRUE) { @@ -1573,11 +1966,15 @@ retrieveGenomes <- function(base_dir = ".", db_path <- paths$db_path con0 <- DBI::dbConnect(duckdb::duckdb(), dbdir = db_path) has_filtered <- "filtered" %in% DBI::dbListTables(con0) + if (has_filtered) { if (isTRUE(verbose)) message("Using existing 'filtered' table (skipping re-filter).") con <- con0 tbl <- "filtered" + on.exit(try(DBI::dbDisconnect(con0, shutdown = TRUE), silent = TRUE), add = TRUE) } else { + DBI::dbDisconnect(con0, shutdown = TRUE) + if (isTRUE(verbose)) message("No 'filtered' table found; filtering now.") f_out <- .filterGenomes( base_dir = base_dir, @@ -1587,6 +1984,7 @@ retrieveGenomes <- function(base_dir = ".", ) con <- f_out$duckdbConnection tbl <- f_out$table_name + on.exit(try(DBI::dbDisconnect(con, shutdown = TRUE), silent = TRUE), add = TRUE) } ids <- tibble::as_tibble(DBI::dbReadTable(con, tbl)) |> @@ -1607,56 +2005,73 @@ retrieveGenomes <- function(base_dir = ".", if (length(ids) == 0L) { if (isTRUE(verbose)) message("All genomes already complete.") - all_complete <- .list_complete( - genome_path, - tibble::as_tibble(DBI::dbReadTable(con, tbl)) |> - dplyr::distinct(`genome.genome_id`) |> - dplyr::pull(`genome.genome_id`) - ) - return(all_complete) + all_ids <- tibble::as_tibble(DBI::dbReadTable(con, tbl)) |> + dplyr::distinct(`genome.genome_id`) |> + dplyr::pull(`genome.genome_id`) + return(.list_complete(genome_path, all_ids)) } if (identical(method, "ftp")) { if (isTRUE(verbose)) message("Trying FTPS download. Workers=", ftp_workers) - ftp_param <- BiocParallel::SnowParam(workers = max(1L, ftp_workers)) - ft_ok <- BiocParallel::bplapply(ids, function(gid) .ftp_download_one(gid, genome_path), - BPPARAM = ftp_param + + ok_ids <- .ftpes_download_two_pass( + genome_ids = ids, + out_dir = genome_path, + workers_first = ftp_workers, + workers_second = ftp_workers, + log_file = file.path(logs_dir, "ftp_download.log") ) - ok_ids <- ids[unlist(ft_ok)] - if (isTRUE(verbose)) message("Complete file sets for ", length(ok_ids), " genomes (FTP).") - return(c(ok_ids, .list_complete(genome_path, setdiff(ids, ok_ids)))) + + dropped_ids <- setdiff(ids, ok_ids) + if (length(dropped_ids)) { + .purge_genome_files( + genome_path = genome_path, + genome_ids = dropped_ids, + log_file = file.path(logs_dir, "ftp_download.log") + ) + } + + if (isTRUE(verbose)) { + message("Complete file sets for ", length(ok_ids), " genomes (FTP).") + if (length(dropped_ids)) { + message(length(dropped_ids), " genomes were excluded because BV-BRC did not provide a complete file set.") + } + } + return(ok_ids) } # CLI for FASTA, FAA, and GTO, then GFF from GTO chunks <- split(ids, ceiling(seq_along(ids) / chunk_size)) - # Parallel chunk containers if (isTRUE(verbose)) { - message( - "CLI being run in parallel for ", length(chunks), - " data chunks." - ) + message("CLI being run in parallel for ", length(chunks), " data chunks.") } - fasta_param <- BiocParallel::SnowParam(workers = max(1L, cli_fasta_workers)) - fa_res <- BiocParallel::bpmapply( + + old_plan <- future::plan() + on.exit(future::plan(old_plan), add = TRUE) + future::plan(future::multisession, workers = max(1L, cli_fasta_workers)) + + fa_res <- future.apply::future_mapply( FUN = function(vec, tag) .cli_dump_fastas_gto_chunk(image, genome_path, vec, tag), - vec = chunks, tag = paste0("fa", seq_along(chunks)), - SIMPLIFY = TRUE, BPPARAM = fasta_param + vec = chunks, + tag = paste0("fa", seq_along(chunks)), + SIMPLIFY = TRUE, + future.seed = TRUE ) if (!all(fa_res) && isTRUE(verbose)) warning(sum(!fa_res), " data chunks failed.") - # GFF extraction in parallel containers if (isTRUE(verbose)) { - message( - "GFF extraction being run in parallel for ", - length(chunks), " data chunks." - ) + message("GFF extraction being run in parallel for ", length(chunks), " data chunks.") } - gff_param <- BiocParallel::SnowParam(workers = max(1L, cli_gff_workers)) - g_res <- BiocParallel::bpmapply( + + future::plan(future::multisession, workers = max(1L, cli_gff_workers)) + + g_res <- future.apply::future_mapply( FUN = function(vec, tag) .cli_export_gff_chunk(image, genome_path, vec, tag), - vec = chunks, tag = paste0("gff", seq_along(chunks)), - SIMPLIFY = TRUE, BPPARAM = gff_param + vec = chunks, + tag = paste0("gff", seq_along(chunks)), + SIMPLIFY = TRUE, + future.seed = TRUE ) if (!all(g_res) && isTRUE(verbose)) warning(sum(!g_res), " GFF chunks had failures.") @@ -1685,6 +2100,7 @@ retrieveGenomes <- function(base_dir = ".", #' @return A list with duckdbConnection and table_name = "files". genomeList <- function(base_dir = ".", user_bacs, + expected_ids = NULL, verbose = TRUE) { base_dir <- normalizePath(base_dir, mustWork = FALSE) paths <- .buildDBpath(base_dir = base_dir, user_bacs = user_bacs) @@ -1706,6 +2122,11 @@ genomeList <- function(base_dir = ".", genome_ids <- unique(c(gff_ids, fna_ids, faa_ids)) + # Make sure all the files we want are present + if (!is.null(expected_ids)) { + genome_ids <- intersect(genome_ids, unique(as.character(expected_ids))) + } + list_of_files <- purrr::map(genome_ids, function(genomeID) { gff_path <- file.path(genome_path, paste0(genomeID, ".PATRIC.gff")) fna_path <- file.path(genome_path, paste0(genomeID, ".fna")) @@ -1763,7 +2184,7 @@ genomeList <- function(base_dir = ".", #' /data//.duckdb #' #' @param user_bacs Character vector. Species and/or taxon IDs (e.g. -#' `c("Shigella flexneri", "623")`). +#' `c("Shigella sonnei", "624")`). #' @param genome_id_file Character or NULL. Optional path to a file listing genome #' IDs (one per line), passed through to `retrieveMetadata()`. If provided, the #' metadata step is restricted to these genome IDs instead of resolving them from @@ -1776,6 +2197,12 @@ genomeList <- function(base_dir = ".", #' @param evidence_mode Character. Sets what types of AMR evidence is acceptable. #' Default `lab_only`. `any` will not require AMR data for downloads. This will #' return very large download lists for many species! +#' @param checkm_contam Numeric scalar. Maximum allowed CheckM contamination (%). +#' @param checkm_complete Numeric scalar. Minimum allowed CheckM completeness (%). +#' @param gc_deviations Optional numeric scalar. Maximum SDs from the median GC content. +#' @param length_deviations Optional numeric scalar. Maximum SDs from the median genome length. +#' @param cds_deviations Optional numeric scalar. Maximum SDs from the median CDS count. +#' @param debug Logical. If TRUE, keep QC columns in metadata tables. #' @param verbose Logical. Print progress messages. Default TRUE. #' #' @return A list (the output of `genomeList()`), containing: @@ -1789,6 +2216,12 @@ prepareGenomes <- function(user_bacs, method = c("ftp", "cli"), overwrite = FALSE, evidence_mode = c("lab_only", "lab_or_comp", "comp_only", "any"), + checkm_contam = 5, + checkm_complete = 95, + gc_deviations = NULL, + length_deviations = NULL, + cds_deviations = NULL, + debug = FALSE, verbose = TRUE) { method <- match.arg(method) evidence_mode <- match.arg(evidence_mode) @@ -1804,6 +2237,12 @@ prepareGenomes <- function(user_bacs, base_dir = base_dir, abx = "All", overwrite = overwrite, + checkm_contam = checkm_contam, + checkm_complete = checkm_complete, + gc_deviations = gc_deviations, + length_deviations = length_deviations, + cds_deviations = cds_deviations, + debug = debug, verbose = verbose )) @@ -1820,13 +2259,13 @@ prepareGenomes <- function(user_bacs, return(NULL) } - # A little summary of what's left after filtering (or not) paths <- .buildDBpath(base_dir = base_dir, user_bacs = user_bacs) con <- DBI::dbConnect(duckdb::duckdb(), dbdir = paths$db_path, read_only = TRUE) n_filtered <- DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM filtered')$n n_meta <- if ("genome_data" %in% DBI::dbListTables(con)) DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome.genome_id") AS n FROM genome_data')$n else NA n_amr <- if ("amr_phenotype" %in% DBI::dbListTables(con)) DBI::dbGetQuery(con, 'SELECT COUNT(DISTINCT "genome_drug.genome_id") AS n FROM amr_phenotype')$n else NA DBI::dbDisconnect(con, shutdown = TRUE) + if (isTRUE(verbose)) { message(sprintf( "Evidence filter summary: filtered=%d | genomes with AMR=%s | genomes with genome_data=%s", @@ -1836,12 +2275,12 @@ prepareGenomes <- function(user_bacs, if (isTRUE(verbose)) message("Step 2: Downloading genomes from BV-BRC (", method, ")") ids <- retrieveGenomes( - base_dir = base_dir, - user_bacs = user_bacs, - method = method, + base_dir = base_dir, + user_bacs = user_bacs, + method = method, skip_existing = !overwrite, evidence_mode = evidence_mode, - verbose = verbose + verbose = verbose ) if (length(ids) == 0L) { message("No genomes downloaded.") @@ -1850,13 +2289,17 @@ prepareGenomes <- function(user_bacs, if (isTRUE(verbose)) message("Step 3: Formatting data into a database for further processing") out <- genomeList( - base_dir = base_dir, + base_dir = base_dir, user_bacs = user_bacs, - verbose = verbose + expected_ids = ids, + verbose = verbose ) if (isTRUE(verbose)) { - message("Done. Files are ready! Continue with downstream processing with runDataProcessing().") + message("Done. Files are ready!") + message("") + message("Continue with downstream processing using:") + message('runDataProcessing("', normalizePath(paths$db_path), '")') } out } diff --git a/R/data_processing.R b/R/data_processing.R index 17b7682..6e1fc3e 100644 --- a/R/data_processing.R +++ b/R/data_processing.R @@ -27,6 +27,30 @@ NULL sub(pattern, container_root, x_unix) } +#' Set a future plan for parallel processing +#' +#' Sets a `future` plan for the duration of a block and restores the +#' previous plan on exit. +#' +#' @param workers Integer. Number of workers to use; if <= 1, uses sequential mode. +#' @param plan Character. Either "multisession" or "sequential". +#' +#' @return Invisibly returns TRUE. +#' @keywords internal +.with_future_plan <- function(workers, plan = c("multisession", "sequential")) { + plan <- match.arg(plan) + old_plan <- future::plan() + on.exit(future::plan(old_plan), add = TRUE) + + if (is.null(workers) || workers <= 1L || identical(plan, "sequential")) { + future::plan(future::sequential) + } else { + future::plan(future::multisession, workers = workers) + } + + invisible(TRUE) +} + # Launch Panaroo to build a pangenome (per batch) #' processPanaroo() #' @@ -203,21 +227,20 @@ NULL # Ensure sum of per-job CPUs does not exceed `threads` panaroo_threads_per_job <- max(1L, floor(threads / n_jobs)) - param <- BiocParallel::SnowParam(workers = max(1L, n_jobs)) - batch_panaroo_run <- BiocParallel::bplapply( + .with_future_plan(workers = n_jobs) + + batch_panaroo_run <- furrr::future_map( panaroo_batches, - function(batch) { - .processPanaroo( - batch_input = batch, - output_path = output_path, - core_threshold = core_threshold, - len_dif_percent = len_dif_percent, - cluster_threshold = cluster_threshold, - family_seq_identity = family_seq_identity, - panaroo_threads_per_job = panaroo_threads_per_job - ) - }, - BPPARAM = param + ~ .processPanaroo( + batch_input = .x, + output_path = output_path, + core_threshold = core_threshold, + len_dif_percent = len_dif_percent, + cluster_threshold = cluster_threshold, + family_seq_identity = family_seq_identity, + panaroo_threads_per_job = panaroo_threads_per_job + ), + .options = furrr::furrr_options(seed = TRUE) ) invisible(batch_panaroo_run) @@ -1205,7 +1228,6 @@ domainFromIPR <- function(duckdb_path, chunks <- list(sequences_df) # Force 1 chunk for RAM limits # Forcing 1 container operation for RAM limits - workers <- 1 cpu_per_container <- threads message(sprintf( @@ -1213,27 +1235,33 @@ domainFromIPR <- function(duckdb_path, cpu_per_container )) - results <- BiocParallel::bplapply(seq_along(chunks), function(i) { - res <- try( - .process_chunk( - chunk = chunks[[i]], - path = path, - ipr_data_path = ipr_data_path, - out_file_base = out_file_base, - appl = appl, - chunk_id = i, - threads = cpu_per_container, - file_format = file_format, - docker_image = ipr_image - ), - silent = TRUE - ) - if (inherits(res, "try-error")) { - message(sprintf("Chunk %d failed: %s", i, as.character(res))) - return(NULL) - } - res - }, BPPARAM = BiocParallel::SerialParam()) + .with_future_plan(workers = 1) + + results <- future.apply::future_lapply( + seq_along(chunks), + function(i) { + res <- try( + .process_chunk( + chunk = chunks[[i]], + path = path, + ipr_data_path = ipr_data_path, + out_file_base = out_file_base, + appl = appl, + chunk_id = i, + threads = cpu_per_container, + file_format = file_format, + docker_image = ipr_image + ), + silent = TRUE + ) + if (inherits(res, "try-error")) { + message(sprintf("Chunk %d failed: %s", i, as.character(res))) + return(NULL) + } + res + }, + future.seed = TRUE + ) # Combine results tsvs <- Filter(function(x) !is.null(x) && file.exists(x), results) @@ -1771,7 +1799,7 @@ runDataProcessing <- function(duckdb_path, ) # 4) Clean metadata and export Parquet + Parquet-backed DuckDB - if (missing(ref_file_path) || is.null(ref_file_path)) { + if (is.null(ref_file_path) || !nzchar(ref_file_path)) { stop("`ref_file_path` (directory with reference TSVs) must be provided to cleanData().") } if (isTRUE(verbose)) message("Cleaning metadata and exporting Parquet-backed views.")