speed improvement to generateNoiseImage

[hvalev]

Hi,

the following pull request modifies the generateNoiseImage function, where the apply function is replaced with a new implementation which leverages vectorization for computing the dot product. Vectorization allows the matrices to be computed in parallel, whereas apply does it sequentially and thus much slower. The impact is high since the generateNoiseImage function is used for both computing CIs as well as zmaps. Below is a comparison of the outputs of both implementation variants as well as a benchmark script comparing the computation speed. The following RData file has been used for both scripts outlined below and has been extracted from a genuine dataset. This pull request should be a drop-in replacement, since it does not use any additional packages.

The script below evaluates the output of the original apply function and that of the newer implementation:

# Sample data
set.seed(42)

load("test.RData")

# Original apply function
original_noise <- apply(p$patches * array(params[p$patchIdx], dim(p$patches)), 1:2, mean)

# Improved code for performance
patch_indices <- p$patchIdx
patch_params <- array(params[patch_indices], dim(p$patches))
reshaped_matrix <- array(p$patches * patch_params, dim(p$patches))
noise <- array(rowMeans(reshaped_matrix), dim(p$patches)[1:2])

# Print array dimensions of each alternative
print(dim(original_noise))
print(dim(improved_noise))

# Check if the arrays are identical
are_identical <- identical(original_noise, improved_noise)

if (are_identical) {
  cat("The computed arrays are identical.\n")
} else {
  cat("The computed arrays are NOT identical.\n")
}

if (length(original_noise) != length(improved_noise)) {
  stop("Arrays must have the same length.")
}

# Calculate absolute differences element-wise
abs_diff <- abs(original_noise - improved_noise)

# Sum of the deviations across each compared pair
total_deviation <- sum(abs_diff)

# Maximum deviation
max_deviation <- max(abs_diff)

# Minimum deviation
min_deviation <- min(abs_diff)

# Print the results
cat("Sum of deviations:", total_deviation, "\n")
cat("Maximum deviation:", max_deviation, "\n")
cat("Minimum deviation:", min_deviation, "\n")

# Compare element-wise
comparison_result <- original_noise == improved_noise

# Check if all elements are the same
all_same <- all(comparison_result)

# Print the aggregated result
if (all_same) {
  print("All elements are the same.")
} else {
  print("Not all elements are the same.")
}

# Count the number of identical (TRUE) and different (FALSE) values
value_counts <- table(comparison_result)

# Print the counts
print(value_counts)

the output from running the script on the linked file is:

[1] 512 512
[1] 512 512
The computed arrays are NOT identical.
Sum of deviations: 1.250356e-16 
Maximum deviation: 1.387779e-17 
Minimum deviation: 0 
[1] "Not all elements are the same."
comparison_result
 FALSE   TRUE 
   955 261189 

Differences occur only for 0.003% (or 955 elements) of the noise matrix, with the largest deviation being $\approx1.38\times10^{-17}$. The results differ on a marginal part of the test data, and the aggregate sum of errors is also negligible ($\approx1.25\times10^{-16}$). The reason for this deviation is presumably due to differences in how the last bits are rounded.

Below is a benchmark evaluating the old and new implementation on performing a computation on RData file linked above:

# Load the necessary data
load("test.RData")

# Define the original apply function
original_apply <- function() {
  original_noise <- apply(p$patches * array(params[p$patchIdx], dim(p$patches)), 1:2, mean)
}

# Define the improved code
improved_code <- function() {
  # Improved code for calculating 3D dot product
  patch_indices <- p$patchIdx
  patch_params <- array(params[patch_indices], dim(p$patches))
  reshaped_matrix <- array(p$patches * patch_params, dim(p$patches))
  noise <- array(rowMeans(reshaped_matrix), dim(p$patches)[1:2])
}

# Benchmark the original apply function
original_time <- system.time(original_apply())

# Benchmark the improved code
improved_time <- system.time(improved_code())

# Print the benchmark results
cat("Original apply function time:", original_time[["elapsed"]], "seconds\n")
cat("Improved code time:", improved_time[["elapsed"]], "seconds\n")

Below are the results:

Original apply function time: 2.23 seconds
Improved code time: 0.23 seconds

The newer implementation yields around 9x speed increase. Realistically, since this method is (almost) always called in parallel from multiple threads on multiple different matrices, the speed increase is around 6x due to CPU bottlenecks. In terms of memory usage, both implementations are roughly equal.

[rdotsch]

This reshaped matrix does not change dimensionality: dim(p$patches) = 512 x 512 x n_layers, dim(patch_params) = 512 x 512 x n_layers, so array(p$patches * patch_params, dim(p$patches)) will have that same dimensionality. I think what this should do, in order for the rowMeans to be correct is to change reshaped_matrix to have dimensionality n_layers x 512 x 512 (assuming rowMeans takes the average in the first dimension). The reason is that you want to take the average of every pixels across every layer (made up of patches).

[rdotsch]

When I reran your example script with the same Rdata file, I actually got sum of absolute deviations of:

Sum of deviations: 10980.83 

I checked and indeed we need to reorder the dimensions for this to work:

reshaped_matrix <- aperm(p$patches * patch_params, c(3, 1, 2))
noise <- array(colMeans(reshaped_matrix), dim(p$patches)[1:2])

When I use this, we get:

Sum of deviations: 3.58144e-12 

Wdyt @hvalev?

[hvalev]

Hi Ron,

you're absolutely right. For development I had used some hard-coded values for the dimensions. When I was switching to using variables for this PR I think somehow something got cached. The function below should yield the results as described in the PR. Just replace it in both scripts.

patch_indices <- p$patchIdx
pd <- dim(p$patches)
patch_params <- array(params[patch_indices], dim = pd)
reshaped_matrix <- array(p$patches * patch_params, dim = c(pd[1]*pd[2], pd[3]))
improved_noise <- array(rowMeans(reshaped_matrix), dim = pd[1:2])

[rdotsch]

Thank you for your PR @hvalev; as you can see I think we need to make some changes to make this work correctly. Please take a look at my comments and see if you agree.