R/ancestry_wide.R

In Rduckhts: 'DuckDB' High Throughput Sequencing File Formats Reader Extension

.ancestry_bounded_proportions <- function(
  con, input_table, reference_table, correction_table, input_kind,
  sum_to_one, min_cor, table_name, overwrite, group_labels = NULL
) {
  .ancestry_profile_stage("entry")
  relations <- vapply(list(input_table, reference_table, correction_table),
                      function(x) sql_quote_identifier(con, x), character(1L))
  contract <- .ancestry_check_site_contract(con, reference_table, relations)
  groups <- contract$groups
  if (is.null(group_labels)) {
    group_labels <- groups
  } else if (is.character(group_labels) && length(group_labels) == length(groups) &&
             !anyNA(group_labels) && !anyDuplicated(group_labels) &&
             !is.null(names(group_labels)) && !anyNA(names(group_labels)) &&
             !anyDuplicated(names(group_labels)) &&
             setequal(names(group_labels), groups)) {
    group_labels <- unname(group_labels[match(groups, names(group_labels))])
  } else {
    stop("group_ids must map every wide frequency column to a distinct group ID",
         call. = FALSE)
  }
  .ancestry_profile_stage("validation")
  coefficients <- contract$coefficients
  reference_source <- contract$reference_source
  quote_id <- function(x) as.character(DBI::dbQuoteIdentifier(con, x))
  sources <- relations
  sources[[2L]] <- reference_source
  frequency <- if (input_kind == "dosage") "dosage / 2.0" else "frequency"
  quote_str <- function(x) as.character(DBI::dbQuoteString(con, x))
  audit <- basename(tempfile("ancestry_audit_"))
  moments <- basename(tempfile("ancestry_moments_"))
  solved <- basename(tempfile("ancestry_solved_"))
  prediction <- basename(tempfile("ancestry_prediction_"))
  aligned <- tempfile("ancestry_aligned_", fileext = ".parquet")
  on.exit({
    unlink(aligned)
    for (name in c(prediction, solved, moments, audit)) {
      invisible(try(DBI::dbExecute(con, paste0("DROP TABLE IF EXISTS ",
                                            quote_id(name))), silent = TRUE))
    }
  }, add = TRUE)
  input_loci <- paste0("(SELECT DISTINCT ",
    .ancestry_chromosome_key(contract$numeric_chromosome),
    " AS chromosome, position FROM ", sources[[1L]], ")")
  duplicated <- DBI::dbGetQuery(con, paste0(
    "SELECT EXISTS (SELECT 1 FROM ", reference_source, " r JOIN ", input_loci,
    " i ON r.chromosome = i.chromosome AND r.position = i.position ",
    "GROUP BY r.chromosome, r.position HAVING count(*) > 1) AS present"))$present
  if (duplicated) {
    stop("duplicate reference locus", call. = FALSE)
  }
  .ancestry_profile_stage("duplicate")
  classification <- .ancestry_classification_query(sources, frequency,
                                                    contract$numeric_chromosome)
  DBI::dbExecute(con, paste0("CREATE TEMP TABLE ", quote_id(audit),
    " AS ", .ancestry_indexed_audit_query(classification)))
  .ancestry_profile_stage("audit")
  DBI::dbExecute(con, paste0("COPY (", .ancestry_aligned_query(classification),
    ") TO ", quote_str(aligned),
    " (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 32768)"))
  .ancestry_profile_stage("alignment")
  aligned_bytes <- file.info(aligned)$size
  aligned_source <- paste0("read_parquet(", quote_str(aligned), ")")
  one_sample <- DBI::dbGetQuery(con, paste0("SELECT count(*) AS n FROM ",
                                         quote_id(audit)))$n == 1L
  .ancestry_profile_stage("sample_index")
  indexed_source <- if (one_sample) aligned_source else paste0(
    "(SELECT a.*, i.sample_index FROM ", aligned_source, " a JOIN ",
    quote_id(audit), " i USING (sample_id))")
  numeric_columns <- paste0("r.", quote_id(c(paste0("PC", seq_along(coefficients)),
                                              groups)))
  invalid <- paste0("(", numeric_columns, " IS NULL OR NOT isfinite(",
                    numeric_columns, "))", collapse = " OR ")
  group_columns <- paste0("r.", quote_id(groups))
  out_of_range <- paste0("(", group_columns, " < 0 OR ", group_columns,
                         " > 1)", collapse = " OR ")
  invalid_site <- DBI::dbGetQuery(con, paste0(
    "SELECT coalesce(bool_or(", invalid, "), false) AS nonfinite, ",
    "coalesce(bool_or(", out_of_range, "), false) AS out_of_range FROM ",
    indexed_source, " a JOIN ", reference_source,
    " r ON a.chromosome = r.chromosome AND a.position = r.position ",
    "AND a.allele_a = r.allele_a AND a.allele_b = r.allele_b"))
  if (invalid_site$nonfinite) {
    stop("reference sites require finite frequencies and loadings", call. = FALSE)
  }
  if (invalid_site$out_of_range) {
    stop("reference sites require frequencies in [0, 1]", call. = FALSE)
  }
  .ancestry_profile_stage("finite")
  DBI::dbExecute(con, paste0("CREATE TEMP TABLE ", quote_id(moments), " AS ",
    .ancestry_moments_query(con, reference_source, groups, coefficients,
                            indexed_source, one_sample)))
  .ancestry_profile_stage("moments")
  DBI::dbExecute(con, paste0("CREATE TEMP TABLE ", quote_id(solved), " AS SELECT ",
    "m.sample_id, ", .ancestry_solver_expression(groups, length(coefficients), sum_to_one),
    " AS q FROM ", quote_id(moments), " m"))
  .ancestry_profile_stage("solver")
  single_proportions <- NULL
  if (one_sample) {
    proportions <- DBI::dbGetQuery(con, paste0("SELECT q FROM ", quote_id(solved)))
    if (nrow(proportions)) {
      single_proportions <- proportions$q[[1L]]
    }
  }
  DBI::dbExecute(con, paste0("CREATE TEMP TABLE ", quote_id(prediction), " AS ",
    .ancestry_prediction_query(con, reference_source, groups, aligned_source,
                               quote_id(solved), one_sample, single_proportions)))
  .ancestry_profile_stage("prediction")
  query <- .ancestry_wide_query(con, group_labels, quote_id(audit), quote_id(moments),
                                quote_id(solved), quote_id(prediction), sum_to_one, min_cor)
  .ancestry_profile_stage("publish")
  result <- .somalier_publish_query(con, query, table_name, overwrite)
  if (is.null(table_name)) attr(result, "aligned_bytes") <- aligned_bytes
  result
}

.ancestry_profile_stage <- function(stage) {
  callback <- getOption("duckhts.ancestry_step_hook")
  if (!is.null(callback)) callback(stage)
}

.ancestry_check_site_contract <- function(con, reference_table, relations) {
  fields <- DBI::dbListFields(con, reference_table)
  keys <- c("chromosome", "position", "allele_a", "allele_b")
  pcs <- paste0("PC", seq_len(sum(grepl("^PC[1-9][0-9]*$", fields))))
  groups <- sort(setdiff(fields, c(keys, pcs)))
  if (length(pcs) < 1L || length(pcs) > 64L ||
      !all(c(keys, pcs) %in% fields) || length(groups) < 1L ||
      length(groups) > 30L || anyDuplicated(fields)) {
    stop("wide reference requires keyed sites, consecutive PC1..PCn (n <= 64) and 1..30 groups", call. = FALSE)
  }
  chromosome_key <- .ancestry_chromosome_key(FALSE)
  null_sample <- DBI::dbGetQuery(con, paste0("SELECT EXISTS (SELECT 1 FROM ",
    relations[[1L]], " WHERE sample_id IS NULL) AS present"))$present
  if (null_sample) stop("sample_id must not be NULL", call. = FALSE)
  invalid_position <- DBI::dbGetQuery(con, paste0(
    "SELECT EXISTS (SELECT 1 FROM ", relations[[1L]],
    " WHERE position IS NULL OR NOT isfinite(position) OR position < 1 ",
    "OR position != floor(position)) AS present"))$present
  if (invalid_position) stop("input positions must be positive whole numbers",
                             call. = FALSE)
  corrections <- DBI::dbGetQuery(con, paste0("SELECT pc, coefficient FROM ", relations[[3L]],
                                             " ORDER BY pc"))
  if (nrow(corrections) != length(pcs) ||
      !.ancestry_consecutive_pcs(corrections$pc) ||
      any(!is.finite(corrections$coefficient))) {
    stop("correction requires one finite coefficient for every PC", call. = FALSE)
  }
  corrections <- corrections[order(as.numeric(corrections$pc)), , drop = FALSE]
  chromosome_type <- DBI::dbGetQuery(con, paste0("DESCRIBE SELECT chromosome FROM ",
                                             relations[[2L]]))$column_type
  numeric_chromosome <- .ancestry_integer_type(chromosome_type)
  reference_source <- if (numeric_chromosome) relations[[2L]] else paste0(
    "(SELECT ", chromosome_key, " AS chromosome, * EXCLUDE (chromosome) FROM ",
    relations[[2L]], ")")
  list(groups = groups, coefficients = corrections$coefficient,
       reference_source = reference_source, numeric_chromosome = numeric_chromosome)
}

# A reference whose chromosome column has one of these DuckDB types compares
# chromosomes as numbers.
.ancestry_integer_type <- function(column_type) {
  column_type %in% c("TINYINT", "SMALLINT", "INTEGER", "BIGINT", "HUGEINT",
                     "UTINYINT", "USMALLINT", "UINTEGER", "UBIGINT", "UHUGEINT")
}

# The chromosome join key, as SQL. Names are compared with duckhts_contig_key(),
# the one contig key of the extension. A reference whose chromosome column is
# an integer compares that key as an integer, so the reference rows need no
# conversion.
.ancestry_chromosome_key <- function(numeric_chromosome, column = "chromosome") {
  key <- paste0("duckhts_contig_key(", column, "::VARCHAR)")
  if (numeric_chromosome) paste0("try_cast(", key, " AS BIGINT)") else key
}

# The keys, as SQL, that make two input rows one locus: the join key. A name
# that gives no integer for an integer reference joins nothing; it keeps its
# own name, so two such names at one position are not duplicates of each other.
.ancestry_duplicate_keys <- function(numeric_chromosome) {
  name_key <- .ancestry_chromosome_key(FALSE)
  if (!numeric_chromosome) return(name_key)
  number_key <- .ancestry_chromosome_key(TRUE)
  paste0(number_key, ", CASE WHEN ", number_key, " IS NULL THEN ", name_key, " END")
}

.ancestry_classification_query <- function(relations, frequency,
                                           numeric_chromosome) {
  input <- relations[[1L]]
  reference <- relations[[2L]]
  paste0(
    "raw AS (SELECT sample_id, chromosome, position, allele_a, allele_b, ",
    .ancestry_chromosome_key(numeric_chromosome), " AS ref_chr, ", frequency,
    # Duplicates are counted by the key the join uses, so two spellings that
    # join to one reference locus are duplicates, not two variants.
    " AS f, count(*) OVER (PARTITION BY sample_id, ",
    .ancestry_duplicate_keys(numeric_chromosome), ", position) AS copies ",
    "FROM ", input, "), ",
    "sites AS (SELECT ", if (numeric_chromosome) "chromosome" else
      .ancestry_chromosome_key(FALSE),
    " AS chromosome, position, allele_a AS ra, allele_b AS rb ",
    "FROM ", reference, "), ",
    "classified AS (SELECT i.*, s.ra, s.rb, CASE ",
    "WHEN i.copies > 1 THEN 'duplicate' ",
    "WHEN i.allele_a IS NULL OR i.allele_b IS NULL OR length(i.allele_a) != 1 ",
    "OR length(i.allele_b) != 1 OR i.allele_a NOT IN ('A','C','G','T') ",
    "OR i.allele_b NOT IN ('A','C','G','T') OR i.allele_a = i.allele_b ",
    "THEN 'invalid_alleles' ",
    "WHEN i.allele_a || i.allele_b IN ('AT','TA','CG','GC') ",
    "OR s.ra || s.rb IN ('AT','TA','CG','GC') THEN 'ambiguous' ",
    "WHEN i.f IS NULL OR NOT isfinite(i.f) OR i.f < 0 OR i.f > 1 THEN 'missing_frequency' ",
    "WHEN s.ra IS NULL THEN 'missing_reference' ",
    "WHEN i.allele_a = s.ra AND i.allele_b = s.rb THEN 'direct' ",
    "WHEN i.allele_a = s.rb AND i.allele_b = s.ra THEN 'reversed' ",
    "WHEN translate(i.allele_a, 'ACGT', 'TGCA') = s.ra ",
    "AND translate(i.allele_b, 'ACGT', 'TGCA') = s.rb THEN 'flipped' ",
    "WHEN translate(i.allele_a, 'ACGT', 'TGCA') = s.rb ",
    "AND translate(i.allele_b, 'ACGT', 'TGCA') = s.ra THEN 'flipped_reversed' ",
    "ELSE 'allele_mismatch' END AS match_status FROM raw i LEFT JOIN sites s ",
    "ON i.ref_chr = s.chromosome AND i.position = s.position)")
}

.ancestry_aligned_query <- function(classification) {
  paste0("WITH ", classification,
    " SELECT sample_id, ref_chr AS chromosome, position, ra AS allele_a, rb AS allele_b, ",
    "CASE WHEN match_status IN ('reversed','flipped_reversed') THEN 1-f ELSE f END ",
    "AS aligned_f FROM classified WHERE match_status IN ",
    "('direct','reversed','flipped','flipped_reversed')")
}

.ancestry_audit_query <- function() {
  paste0(
    "SELECT sample_id, count(*) AS input_variants, ",
    "count(*) FILTER (WHERE match_status IN ('direct','reversed','flipped','flipped_reversed')) AS used_variants, ",
    "count(*) FILTER (WHERE match_status IN ('reversed','flipped_reversed')) AS reversed_variants, ",
    "count(*) FILTER (WHERE match_status IN ('flipped','flipped_reversed')) AS flipped_variants, ",
    "count(*) FILTER (WHERE match_status = 'duplicate') AS duplicate_variants, ",
    "count(*) FILTER (WHERE match_status = 'ambiguous') AS ambiguous_variants, ",
    "count(*) FILTER (WHERE match_status IN ('missing_reference','allele_mismatch')) AS unmatched_variants, ",
    "count(*) FILTER (WHERE match_status = 'invalid_alleles') AS invalid_variants, ",
    "count(*) FILTER (WHERE match_status = 'missing_frequency') AS missing_variants ",
    "FROM classified GROUP BY sample_id")
}

.ancestry_indexed_audit_query <- function(classification) {
  paste0("WITH ", classification, ", totals AS (", .ancestry_audit_query(),
         ") SELECT totals.*, row_number() OVER (ORDER BY sample_id)::INTEGER ",
         "AS sample_index FROM totals")
}

.ancestry_moments_query <- function(con, reference, groups, coefficients, aligned,
                                    one_sample) {
  quote_id <- function(x) as.character(DBI::dbQuoteIdentifier(con, x))
  pc <- paste0("r.", quote_id(paste0("PC", seq_along(coefficients))))
  group <- paste0("r.", quote_id(groups))
  # The solver reads X in PC-major, group-minor order.
  x_terms <- unlist(lapply(seq_along(pc), function(k) {
    paste0("sum(", pc[[k]], " * ", group, ") AS ",
           paste0("x", k, "_", seq_along(group)))
  }), use.names = FALSE)
  y_terms <- paste0("sum(", pc, " * a.aligned_f) * ",
                    vapply(coefficients, function(x) .somalier_quote_number(con, x),
                           character(1L)), " AS y", seq_along(pc))
  cor_terms <- paste0("corr(a.aligned_f, ", group, ") AS cor", seq_along(group))
  sample <- "any_value(a.sample_id) AS sample_id"
  aggregate <- if (one_sample) "HAVING count(*) > 0" else "GROUP BY a.sample_index"
  paste0("SELECT ", sample, ", ",
         paste(c(x_terms, y_terms, cor_terms), collapse = ", "),
         " FROM ", aligned, " a JOIN ", reference, " r ON a.chromosome = r.chromosome ",
         "AND a.position = r.position AND a.allele_a = r.allele_a ",
         "AND a.allele_b = r.allele_b ", aggregate)
}

.ancestry_solver_expression <- function(groups, pcs, sum_to_one) {
  x_names <- unlist(lapply(seq_len(pcs), function(k) {
    paste0("m.x", k, "_", seq_along(groups))
  }), use.names = FALSE)
  paste0("duckhts_ancestry_proportions(list_value(", paste(x_names, collapse = ", "),
         "), list_value(", paste0("m.y", seq_len(pcs), collapse = ", "), "), ",
         length(groups), ", ", if (sum_to_one) "true" else "false", ")")
}

.ancestry_prediction_query <- function(con, reference, groups, aligned,
                                       solved, one_sample, single_proportions) {
  quote_id <- function(x) as.character(DBI::dbQuoteIdentifier(con, x))
  group <- paste0("r.", quote_id(groups))
  if (one_sample && length(single_proportions)) {
    weights <- vapply(single_proportions, function(x) .somalier_quote_number(con, x),
                      character(1L))
  } else if (one_sample) {
    weights <- rep("NULL::DOUBLE", length(groups))
  } else {
    weights <- paste0("s.q[", seq_along(groups), "]")
  }
  predicted <- paste0("(", group, " * ", weights, ")", collapse = " + ")
  sample <- if (one_sample) "any_value(a.sample_id)" else "a.sample_id"
  aggregate <- if (one_sample) "HAVING count(*) > 0" else "GROUP BY a.sample_id"
  solved_join <- if (one_sample) "" else paste0("JOIN ", solved,
                                               " s ON s.sample_id = a.sample_id ")
  paste0("SELECT ", sample, " AS sample_id, corr(a.aligned_f, ", predicted,
    ") AS cor_pred FROM ", aligned, " a JOIN ", reference,
    " r USING (chromosome, position, allele_a, allele_b) ", solved_join, aggregate)
}

.ancestry_wide_query <- function(con, groups, audit, moments, solved, prediction,
                                 sum_to_one, min_cor) {
  quote_str <- function(x) as.character(DBI::dbQuoteString(con, x))
  values_cor <- paste0("m.cor", seq_along(groups), collapse = ", ")
  group_rows <- paste0("(", seq_along(groups), ", ",
                       vapply(groups, quote_str, character(1L)), ")", collapse = ", ")
  equality <- if (sum_to_one) "true" else "false"
  gate <- .somalier_quote_number(con, min_cor)
  paste0(
    "WITH each_cor AS (SELECT m.sample_id, c.group_index, c.cor_each FROM ", moments,
    " m, UNNEST(list_value(", values_cor,
    ")) WITH ORDINALITY AS c(cor_each, group_index)), ",
    "quality AS (SELECT a.sample_id, p.cor_pred, CASE ",
    "WHEN a.used_variants = 0 THEN 'no_matched_variants' ",
    "WHEN (SELECT avg(e.cor_each) FROM each_cor e WHERE e.sample_id = a.sample_id) < -0.2 ",
    "THEN 'reversed_alleles' WHEN p.cor_pred IS NULL OR NOT isfinite(p.cor_pred) ",
    "OR p.cor_pred < ", gate, " THEN 'low_correlation' ELSE 'ok' END AS status ",
    "FROM ", audit, " a LEFT JOIN ", prediction, " p USING (sample_id)) ",
    "SELECT a.sample_id, g.group_id, CASE WHEN q.status = 'ok' ",
    "THEN round(s.q[g.group_index], 7) ELSE NULL END AS proportion, ",
    "q.cor_pred, e.cor_each, q.status, a.input_variants, a.used_variants, ",
    "a.input_variants - a.used_variants AS dropped_variants, ",
    "a.reversed_variants, a.flipped_variants, a.duplicate_variants, ",
    "a.ambiguous_variants, a.unmatched_variants, a.invalid_variants, a.missing_variants, ",
    equality, " AS sum_to_one, ", gate, "::DOUBLE AS min_cor ",
    "FROM ", audit, " a JOIN quality q USING (sample_id) ",
    "CROSS JOIN (VALUES ", group_rows, ") g(group_index, group_id) ",
    "LEFT JOIN ", solved, " s ON s.sample_id = a.sample_id ",
    "LEFT JOIN each_cor e ON e.sample_id = a.sample_id AND e.group_index = g.group_index ",
    "ORDER BY a.sample_id, g.group_id"
  )
}