New cf_add_meta function

NAMESPACE

@@ -4,6 +4,7 @@ S3method(c,cidlist)
 S3method(print,cidlist)
 export("%>%")
 export(abbreviate_datasets)
+export(cf_add_meta)
 export(cf_cosine_plot)
 export(cf_datasets)
 export(cf_ids)

R/meta.R

@@ -334,6 +334,113 @@ cf_meta <- function(ids, bind.rows=TRUE, integer64=FALSE, keep.all=FALSE,
   res
 }
 
+
+#' Add metadata to a dataframe containing neuron keys
+#'
+#' @description Enriches a dataframe with metadata for neurons identified by
+#'   key columns. This is useful for adding metadata to partner data or other
+#'   results that contain neuron keys.
+#'
+#' @param x A data.frame containing one or more columns with keys
+#' @param keycol Character vector of column names containing keys.
+#'   Default \code{"key"}. When multiple columns specified, metadata is joined
+#'   for each with corresponding suffixes.
+#' @param suffix Character vector of suffixes for added columns. Must match
+#'   length of keycol. Default generates \code{""} for single keycol, or
+#'   \code{".1"}, \code{".2"}, etc. for multiple. Use \code{""} for no suffix.
+#' @param cols Character vector of column names to add from metadata.
+#'   Default \code{NULL} adds all columns. Warns if any specified columns
+#'   are not found in metadata. The \code{key} column is always kept for joining.
+#' @param ... Additional arguments passed to \code{\link{cf_meta}}
+#'
+#' @return data.frame with additional metadata columns
+#' @export
+#'
+#' @examples
+#' \dontrun{
+#' # Add metadata to partner column only
+#' partners <- cf_partners(cf_ids(hemibrain='DA2_lPN'), threshold=10, details=FALSE)
+#' partners_meta <- cf_add_meta(partners, keycol="post_key")
+#'
+#' # Add metadata to both pre and post
+#' partners_both <- cf_add_meta(partners,
+#'   keycol = c("pre_key", "post_key"),
+#'   suffix = c(".pre", ".post"))
+#'
+#' # Add only type and side columns
+#' partners_minimal <- cf_add_meta(partners, keycol="post_key",
+#'   cols = c("type", "side"))
+#' }
+cf_add_meta <- function(x, keycol = "key", suffix = NULL, cols = NULL, ...) {
+  # Validate inputs
+  if (!is.data.frame(x) || nrow(x) == 0) {
+    return(x)
+  }
+
+  missing_cols <- setdiff(keycol, names(x))
+  if (length(missing_cols))
+    stop("keycol column(s) not found in x: ", paste(missing_cols, collapse = ", "),
+         "\nHint: use keys(x) to add a '", paste(missing_cols, collapse = "/"),
+         "' column, or set keycol= to identify the correct key column.")
+
+  # Set default suffixes
+  if (is.null(suffix)) {
+    suffix <- if (length(keycol) == 1) "" else paste0(".", seq_along(keycol))
+  }
+  if (length(suffix) != length(keycol))
+    stop("Length of suffix must match length of keycol")
+
+  # Extract unique keys and fetch metadata once
+  all_keys <- unique(unlist(x[keycol], use.names = FALSE))
+  all_keys <- all_keys[!is.na(all_keys) & nzchar(all_keys)]
+
+  if (length(all_keys) == 0) {
+    warning("No valid keys found in specified columns")
+    return(x)
+  }
+
+  meta <- cf_meta(all_keys, ...)
+
+  if (is.null(meta) || nrow(meta) == 0) {
+    warning("No metadata found for provided keys")
+    return(x)
+  }
+
+
+  # Remove columns that would conflict or are typically already present
+  meta_cols_to_drop <- c("id", "dataset", "tissue", "sex")
+  meta <- meta[, setdiff(names(meta), meta_cols_to_drop), drop = FALSE]
+
+  # Subset columns if requested (character vector)
+  if (!is.null(cols)) {
+    if (!is.character(cols)) {
+      stop("cols must be a character vector of column names")
+    }
+    missing <- setdiff(cols, names(meta))
+    if (length(missing) > 0) {
+      warning("Column(s) not found in metadata: ",
+              paste(missing, collapse = ", "))
+    }
+    keep_cols <- c("key", intersect(cols, names(meta)))
+    meta <- meta[, keep_cols, drop = FALSE]
+  }
+
+  # Join for each keycol
+  for (i in seq_along(keycol)) {
+    kc <- keycol[i]
+    sf <- suffix[i]
+
+    # Rename key column to match keycol, add suffix to other columns, then join
+    x <- meta %>%
+      dplyr::rename(!!kc := key) %>%
+      dplyr::rename_with(~paste0(.x, sf), .cols = -dplyr::all_of(kc)) %>%
+      dplyr::left_join(x, ., by = kc)
+  }
+
+  x
+}
+
+
 # Shared helper for fanc and banc metadata
 fancorbanc_meta <- function(table, ids=NULL, ...) {
   ol_classes=c("centrifugal", "distal medulla", "distal medulla dorsal rim area",

R/partners.R

@@ -15,6 +15,14 @@
 #' @param partners Whether to return inputs or outputs
 #' @param MoreArgs Additional arguments in the form of a hierarchical list
 #'   (expert use; see details and examples).
+#' @param details Which neurons to enrich with metadata. Options:
+#'   \itemize{
+#'     \item \code{"partner"} (default): Add metadata for partner neurons only
+#'     \item \code{"query"}: Add metadata for query neurons only
+#'     \item \code{"both"}: Add metadata for both with \code{.pre}/\code{.post} suffixes
+#'     \item \code{"neither"}: No metadata, return minimal columns for speed
+#'   }
+#'   Metadata can also be added later via \code{\link{cf_add_meta}}.
 #'
 #' @inheritParams cf_meta
 #' @return A data.frame or a named list (when \code{bind.rows=FALSE})
@@ -38,9 +46,11 @@
 #' }
 cf_partners <- function(ids, threshold=1L, partners=c("inputs", "outputs"),
                         bind.rows=TRUE, MoreArgs=list(), keep.all=FALSE,
+                        details=c("partner", "query", "both", "neither"),
                         use_superclass=getOption("coconatfly.use_superclass", FALSE),
                         harmonise_class=getOption("coconatfly.harmonise_class", FALSE)) {
   partners=match.arg(partners)
+  details=match.arg(details)
   threshold <- checkmate::assert_integerish(
     threshold, lower=0L,len = 1, null.ok = F, all.missing = F)
 
@@ -51,7 +61,7 @@ cf_partners <- function(ids, threshold=1L, partners=c("inputs", "outputs"),
   if(is.data.frame(ids)) {
     ss=split(ids$id, ids$dataset)
     res=cf_partners(ss, threshold = threshold, partners = partners,
-                    bind.rows = bind.rows, MoreArgs=MoreArgs,
+                    bind.rows = bind.rows, MoreArgs=MoreArgs, details=details,
                     use_superclass=use_superclass, harmonise_class=harmonise_class)
     return(res)
   }
@@ -90,14 +100,7 @@ cf_partners <- function(ids, threshold=1L, partners=c("inputs", "outputs"),
       do.call(PFUN, commonArgs)
     }
 
-    # Enrich with partner metadata if partnerfun returned minimal data
-    tres <- add_partner_metadata(tres, dataset = n, partners = partners)
-
     tres=coconat:::standardise_partner_summary(tres)
-    if(isTRUE(harmonise_class) && "class" %in% colnames(tres))
-      tres$class=harmonise_top_class_values(tres$class, n)
-    if("side" %in% colnames(tres))
-      tres$side=normalise_side(tres$side)
     if(nrow(tres)>0) {
       tres$dataset=n
       tres$tissue=dataset_tissue(n)
@@ -108,8 +111,37 @@ cf_partners <- function(ids, threshold=1L, partners=c("inputs", "outputs"),
       tres$sex=character()
       warning("no ", partners, " found for `", n, "` dataset.")
     }
+    # Add keys before metadata enrichment
     tres$pre_key=keys(tres, idcol="pre_id")
     tres$post_key=keys(tres, idcol='post_id')
+
+    # Enrich with metadata based on details option
+    # Always use cf_add_meta rather than relying on partnerfun metadata
+    if (details != "neither" && nrow(tres) > 0) {
+      # Keep only core connectivity columns, drop all partnerfun metadata
+      # cf_add_meta will re-add metadata from cf_meta for consistency
+      core_cols <- c("pre_id", "post_id", "weight", "dataset", "tissue", "sex",
+                     "pre_key", "post_key")
+      tres <- tres[, intersect(names(tres), core_cols), drop = FALSE]
+
+      # Determine which key columns to enrich
+      query_col <- if (partners == "outputs") "pre_key" else "post_key"
+      partner_col <- if (partners == "outputs") "post_key" else "pre_key"
+
+      keycols <- switch(details,
+        partner = partner_col,
+        query = query_col,
+        both = c("pre_key", "post_key")
+      )
+      suffixes <- switch(details,
+        partner = "",
+        query = "",
+        both = c(".pre", ".post")
+      )
+
+      tres <- cf_add_meta(tres, keycol = keycols, suffix = suffixes,
+                          harmonise_class = harmonise_class)
+    }
     res[[n]]=tres
   }
   if(isTRUE(bind.rows)) {
@@ -270,66 +302,3 @@ cf_partner_summary <- function(ids, threshold=1L, partners=c("inputs", "outputs"
       outputcol = ifelse(partners=='outputs', group.post[1], group.pre[1]),
       standardise_input = F, sparse = rval=="sparse")
 }
-
-
-# Add metadata to partner results if the partnerfun returned minimal columns
-# This allows partnerfuns to return just ids + weight and have metadata added
-# automatically via cf_meta()
-add_partner_metadata <- function(tres, dataset, partners) {
-  if (is.null(tres) || !is.data.frame(tres) || nrow(tres) == 0) {
-    return(tres)
-  }
-
-  # Check if we need to fetch partner metadata:
-  # - 3 columns (pre_id, post_id, weight) always needs enrichment
-  # - <=5 columns without 'type' likely needs enrichment
-  needs_enrichment <- ncol(tres) == 3 ||
-    (ncol(tres) <= 5 && !"type" %in% names(tres))
-
-  if (!needs_enrichment) {
-    return(tres)
-  }
-
-  # Find the partner column based on query direction
-  # For outputs: partners are post_id; for inputs: partners are pre_id
-  partner_col <- if (partners == "outputs") "post_id" else "pre_id"
-
-  # Check if expected column exists, otherwise try to find it
-
-  if (!partner_col %in% colnames(tres)) {
-    # Fallback: look for a single *_id column or 'partner' column
-    id_cols <- grep("_id$", colnames(tres), value = TRUE)
-    if (length(id_cols) == 1) {
-      partner_col <- id_cols
-    } else {
-      partner_col <- grep("^partner$", colnames(tres), value = TRUE)
-    }
-  }
-
-  if (length(partner_col) != 1 || !partner_col %in% colnames(tres)) {
-    warning("Unable to find partner column for dataset: ", dataset,
-            "\nPartnerfuns should return pre_id/post_id columns or include metadata")
-    return(tres)
-  }
-
-  # Fetch metadata for unique partner IDs
-  pids <- unique(tres[[partner_col]])
-  # Convert to character for keys() and ensure dataset uses abbreviation
-  pids_char <- coconat::id2char(pids)
-  metadf <- cf_meta(keys(data.frame(id = pids_char, dataset = dataset)))
-
-  if (is.null(metadf) || nrow(metadf) == 0) {
-    return(tres)
-  }
-
-  # Remove columns that will be added by cf_partners later
-  metadf <- metadf[setdiff(colnames(metadf), c("dataset", "key"))]
-
-  # Rename id column to match partner column for joining
-  colnames(metadf)[1] <- partner_col
-
-  # Ensure partner column is character for joining
-  tres[[partner_col]] <- coconat::id2char(tres[[partner_col]])
-
-  dplyr::left_join(tres, metadf, by = partner_col)
-}