Restore the original cell order after the FCS roundtrip

CHANGELOG.md

@@ -239,4 +239,8 @@
   and thus can't be directly overriden (new values given by get_donor_batch_map is int) (PR #119).
 
 * Update flowsom mapping similarity so we subset to just markers to correct, and lisi to remove control samples
-  and unlabelled cells (PR #119).
+  and unlabelled cells (PR #119).
+
+* Fix bug in `gaussnorm` and the `cytonorm_*` methods where the batch corrected matrix was
+  attached to the obs of the input anndata while still being ordered per sample (PR #128).
+  Added `fcs_cell_order()` to `src/utils/anndata_to_fcs.R` to restore the original cell order.

src/methods/cytonorm_all_controls_to_goal/script.R

@@ -114,6 +114,17 @@ cat("Preparing output anndata\n")
 norm_mat <- flowCore::fsApply(norm_fset_all, exprs)
 colnames(norm_mat) <- adata$var_names
 
+# cytonorm returns the flowframes in the same order as it received them
+cells_per_sample <- flowCore::fsApply(fset_all, function(ff) nrow(exprs(ff)))
+cells_per_norm_sample <- flowCore::fsApply(norm_fset_all, function(ff) nrow(exprs(ff)))
+if (!identical(as.integer(cells_per_sample), as.integer(cells_per_norm_sample))) {
+    stop("Cytonorm returned a different number of cells per sample than it was given.")
+}
+
+# fsApply returns the cells grouped per sample, so restore the original
+# cell order before attaching the obs of the input anndata
+norm_mat <- norm_mat[order(fcs_cell_order(adata, fset_all)), ]
+
 norm_mat <- anndata::AnnData(
     obs = adata$obs[, integer(0)],
     var = adata$var[colnames(norm_mat), integer(0)],

src/methods/cytonorm_all_controls_to_mid/script.R

@@ -111,6 +111,17 @@ cat("Preparing output anndata\n")
 norm_mat <- flowCore::fsApply(norm_fset_all, exprs)
 colnames(norm_mat) <- adata$var_names
 
+# cytonorm returns the flowframes in the same order as it received them
+cells_per_sample <- flowCore::fsApply(fset_all, function(ff) nrow(exprs(ff)))
+cells_per_norm_sample <- flowCore::fsApply(norm_fset_all, function(ff) nrow(exprs(ff)))
+if (!identical(as.integer(cells_per_sample), as.integer(cells_per_norm_sample))) {
+    stop("Cytonorm returned a different number of cells per sample than it was given.")
+}
+
+# fsApply returns the cells grouped per sample, so restore the original
+# cell order before attaching the obs of the input anndata
+norm_mat <- norm_mat[order(fcs_cell_order(adata, fset_all)), ]
+
 norm_mat <- anndata::AnnData(
     obs = adata$obs[, integer(0)],
     var = adata$var[colnames(norm_mat), integer(0)],

src/methods/cytonorm_no_controls_to_goal/script.R

@@ -128,6 +128,17 @@ cat("Preparing output anndata\n")
 norm_mat <- flowCore::fsApply(norm_fset_all, exprs)
 colnames(norm_mat) <- adata$var_names
 
+# cytonorm returns the flowframes in the same order as it received them
+cells_per_sample <- flowCore::fsApply(fset_all, function(ff) nrow(exprs(ff)))
+cells_per_norm_sample <- flowCore::fsApply(norm_fset_all, function(ff) nrow(exprs(ff)))
+if (!identical(as.integer(cells_per_sample), as.integer(cells_per_norm_sample))) {
+    stop("Cytonorm returned a different number of cells per sample than it was given.")
+}
+
+# fsApply returns the cells grouped per sample, so restore the original
+# cell order before attaching the obs of the input anndata
+norm_mat <- norm_mat[order(fcs_cell_order(adata, fset_all)), ]
+
 norm_mat <- anndata::AnnData(
     obs = adata$obs[, integer(0)],
     var = adata$var[colnames(norm_mat), integer(0)],

src/methods/cytonorm_no_controls_to_mid/script.R

@@ -125,6 +125,17 @@ cat("Preparing output anndata\n")
 norm_mat <- flowCore::fsApply(norm_fset_all, exprs)
 colnames(norm_mat) <- adata$var_names
 
+# cytonorm returns the flowframes in the same order as it received them
+cells_per_sample <- flowCore::fsApply(fset_all, function(ff) nrow(exprs(ff)))
+cells_per_norm_sample <- flowCore::fsApply(norm_fset_all, function(ff) nrow(exprs(ff)))
+if (!identical(as.integer(cells_per_sample), as.integer(cells_per_norm_sample))) {
+    stop("Cytonorm returned a different number of cells per sample than it was given.")
+}
+
+# fsApply returns the cells grouped per sample, so restore the original
+# cell order before attaching the obs of the input anndata
+norm_mat <- norm_mat[order(fcs_cell_order(adata, fset_all)), ]
+
 norm_mat <- anndata::AnnData(
     obs = adata$obs[, integer(0)],
     var = adata$var[colnames(norm_mat), integer(0)],

src/methods/cytonorm_one_control_to_goal/script.R

@@ -116,6 +116,17 @@ cat("Preparing output anndata\n")
 norm_mat <- flowCore::fsApply(norm_fset_all, exprs)
 colnames(norm_mat) <- adata$var_names
 
+# cytonorm returns the flowframes in the same order as it received them
+cells_per_sample <- flowCore::fsApply(fset_all, function(ff) nrow(exprs(ff)))
+cells_per_norm_sample <- flowCore::fsApply(norm_fset_all, function(ff) nrow(exprs(ff)))
+if (!identical(as.integer(cells_per_sample), as.integer(cells_per_norm_sample))) {
+    stop("Cytonorm returned a different number of cells per sample than it was given.")
+}
+
+# fsApply returns the cells grouped per sample, so restore the original
+# cell order before attaching the obs of the input anndata
+norm_mat <- norm_mat[order(fcs_cell_order(adata, fset_all)), ]
+
 norm_mat <- anndata::AnnData(
     obs = adata$obs[, integer(0)],
     var = adata$var[colnames(norm_mat), integer(0)],

src/methods/cytonorm_one_control_to_mid/script.R

@@ -113,6 +113,17 @@ cat("Preparing output anndata\n")
 norm_mat <- flowCore::fsApply(norm_fset_all, exprs)
 colnames(norm_mat) <- adata$var_names
 
+# cytonorm returns the flowframes in the same order as it received them
+cells_per_sample <- flowCore::fsApply(fset_all, function(ff) nrow(exprs(ff)))
+cells_per_norm_sample <- flowCore::fsApply(norm_fset_all, function(ff) nrow(exprs(ff)))
+if (!identical(as.integer(cells_per_sample), as.integer(cells_per_norm_sample))) {
+    stop("Cytonorm returned a different number of cells per sample than it was given.")
+}
+
+# fsApply returns the cells grouped per sample, so restore the original
+# cell order before attaching the obs of the input anndata
+norm_mat <- norm_mat[order(fcs_cell_order(adata, fset_all)), ]
+
 norm_mat <- anndata::AnnData(
     obs = adata$obs[, integer(0)],
     var = adata$var[colnames(norm_mat), integer(0)],

src/methods/gaussnorm/script.R

@@ -44,6 +44,10 @@ for(marker in markers_to_correct){
 # Concatenate all flowFrames in a FlowSet
 integrated_matrix <- fsApply(fset,exprs)
 
+# fsApply returns the cells grouped per sample, so restore the original
+# cell order before attaching the obs of the input anndata
+integrated_matrix <- integrated_matrix[order(fcs_cell_order(adata, fset)), ]
+
 cat("Write output AnnData to file\n")
 output <- anndata::AnnData(
   obs = adata$obs[,integer(0)],

src/utils/anndata_to_fcs.R

@@ -112,3 +112,34 @@ anndata_to_fcs <- function(adata, out_dir= NULL, layer_name = 'preprocessed') {
     return(fcs_files)
   }
 }
+
+#' @title Cell order of a FlowSet relative to the AnnData object it came from
+#'
+#' @description `anndata_to_fcs()` creates one flowFrame per sample, so
+#' `flowCore::fsApply(fset, exprs)` returns the cells grouped per sample rather
+#' than in the original cell order of the AnnData object. This function returns
+#' the AnnData row indices in the order in which they appear in the concatenated
+#' matrix, so `mat[order(fcs_cell_order(adata, fset)), ]` restores the original
+#' cell order. If the cells were already grouped per sample, this is a no-op.
+#'
+#' @param adata AnnData object the FlowSet was created from.
+#' @param fset FlowSet whose sample names determine the concatenation order.
+#' @return An integer vector with one entry per cell in `adata`.
+fcs_cell_order <- function(adata, fset) {
+  samples <- as.character(adata$obs$sample)
+
+  cell_order <- unlist(lapply(
+    flowCore::sampleNames(fset),
+    function(sample) which(samples == sample)
+  ))
+
+  if (length(cell_order) != length(samples)) {
+    stop(
+      "The FlowSet does not contain the same cells as the AnnData object: ",
+      length(cell_order), " cells in the FlowSet, ",
+      length(samples), " cells in the AnnData object."
+    )
+  }
+
+  cell_order
+}