CODEWITHSUNDEEP
rloopsvectorization
user1480926
user1480926

Reputation: 634

R - vectorizing a which operation

Hi I have a function in R that I'm trying to optimize for performance. I need to vectorize a for loop. My problem is the slightly convoluted data structure and the way I need to perform lookups using the 'which' command.

Lets say we are dealing with 5 elements (1,2,3,4,5), the 10x2 matrix pairs is a combination of all unique pairs the 5 elements (i.e. (1,2), (1,3),(1,4) ....(4,5)). all_prods is a 10x1 matrix that I need to look up using the pairs while iterating through all the 5 elements.

So for 1, I need to index rows 1, 2, 3, 4 (pairs 1,2 1,3 1,4 and 1,5) from all_prods and so on for 1, 2, 3, 4, 5.

I have only recently switched to R from matlab so would really appreciate any help.

foo <- function(AA , BB , CC ){
    pa <- AA*CC;
    pairs <-  t(combn(seq_len(length(AA)),2));

    all_prods <- pa[pairs[,1]] * pa[pairs[,2]];

    result <- matrix(0,1,length(AA));

    # WANT TO VECTORIZE THIS BLOCK
    for(st in seq(from=1,to=length(AA))){
       result[st] <- sum(all_prods[c(which(pairs[,1]==st), which(pairs[,2]==st))])*BB[st];
    }
   return(result);
}
AA <- seq(from=1,to=5); BB<-seq(from=11,to=15); CC<-seq(from=21,to=25);
results <- foo(AA,BB,CC);

#final results is [7715 164208 256542 348096 431250]

I want to convert the for loop into a vectorised version. Instead of looping through every element st, I'd like to do it in one command that gives me a results vector (rather than building it up element by element)

Upvotes: 3

Views: 282

Answers (1)

MvG
MvG

Reputation: 60988

You could write your function like this:

foo <- function(AA, BB, CC) {
  pa <- AA*CC
  x <- outer(pa, pa)
  diag(x) <- 0
  res <- colSums(x)*BB
  return(res)
}

The key idea is to not break the symmetry. Your use of ordered pairs corresponds to the upper right triangle of my matrix x. Although this seems like just half as many values to compute, the syntactic and computational overhead becomes quite large. You are distinguishing situations where st is the first element in the pair from those where it is the second. Later on this leads to quite some trouble to get rid of that distinction. Having the full symmetric matrix, you don't have to worry about order, and things vectorize smoothly.

Upvotes: 8

Related Questions