From d87474e574788754c02fd40500c0d557d7151239 Mon Sep 17 00:00:00 2001
From: Elias Projahn <elias@johrpan.de>
Date: Mon, 11 Oct 2021 10:10:28 +0200
Subject: [PATCH] Update correlation method

This adds some performance optimizations and removes the penalty for
missing values.
---
 correlation.R | 56 ++++++++++++++++++++++++++-------------------------
 1 file changed, 29 insertions(+), 27 deletions(-)

diff --git a/correlation.R b/correlation.R
index 14a71f5..18ad4da 100644
--- a/correlation.R
+++ b/correlation.R
@@ -1,6 +1,4 @@
 library(data.table)
-library(progress)
-library(rlog)
 
 #' Compute the mean correlation coefficient comparing gene distances with a set
 #' of reference genes.
@@ -8,7 +6,7 @@ library(rlog)
 #' The result will be a data.table with the following columns:
 #'
 #'  - `gene` Gene ID of the processed gene.
-#'  - `r_mean` Mean correlation coefficient.
+#'  - `correlation` Mean correlation coefficient.
 #'
 #' @param distances Distance data to use.
 #' @param species_ids Species, whose data should be included.
@@ -17,55 +15,59 @@ library(rlog)
 process_correlation <- function(distances, species_ids, gene_ids,
                                 reference_gene_ids) {
     results <- data.table(gene = gene_ids)
-    gene_count <- length(gene_ids)
     reference_count <- length(reference_gene_ids)
 
-    log_info(sprintf(
-        "Correlating %i genes from %i species with %i reference genes",
-        gene_count,
-        length(species_ids),
-        reference_count
-    ))
-
-    progress <- progress_bar$new(
-        total = gene_count,
-        format = "Correlating genes [:bar] :percent (ETA :eta)"
-    )
-
     # Prefilter distances by species.
     distances <- distances[species %chin% species_ids]
 
-    for (i in 1:gene_count) {
-        progress$tick()
+    # Add an index for quickly accessing data per gene.
+    setkey(distances, gene)
 
-        gene_id <- gene_ids[i]
-        gene_distances <- distances[gene == gene_id]
+    # Prepare the reference genes' data.
+    reference_distances <- distances[gene %chin% reference_gene_ids]
 
-        if (nrow(gene_distances) < 10) {
-            next
+    #' Perform the correlation for one gene.
+    compute <- function(gene_id) {
+        gene_distances <- distances[gene_id]
+        gene_species_count <- nrow(gene_distances)
+
+        # Return a score of 0.0 if there is just one or no value at all.
+        if (gene_species_count <= 1) {
+            return(0.0)
         }
 
         #' Buffer for the sum of correlation coefficients.
-        r_sum <- 0
+        correlation_sum <- 0
 
         # Correlate with all reference genes but not with the gene itself.
         for (reference_gene_id in
             reference_gene_ids[reference_gene_ids != gene_id]) {
             data <- merge(
                 gene_distances,
-                distances[gene == reference_gene_id],
+                reference_distances[reference_gene_id],
                 by = "species"
             )
 
+            # Skip this reference gene, if there are not enough value pairs.
+            # This will lessen the final score, because it effectively
+            # represents a correlation coefficient of 0.0.
+            if (nrow(data) <= 1) {
+                next
+            }
+
             # Order data by the reference gene's distance to get a monotonic
             # relation.
             setorder(data, distance.y)
 
-            r_sum <- r_sum + abs(cor(data[, distance.x], data[, distance.y]))
+            correlation_sum <- correlation_sum + abs(cor(
+                data[, distance.x], data[, distance.y],
+                method = "spearman"
+            ))
         }
 
-        results[gene == gene_id, r_mean := r_sum / reference_count]
+        # Compute the score as the mean correlation coefficient.
+        score <- correlation_sum / reference_count
     }
 
-    results
+    results[, correlation := compute(gene), by = 1:nrow(results)]
 }
\ No newline at end of file