R: Apply FUN to kxk subsections of array

Question

The language is R.

I have an nxm matrix, and I'd like to partition it into 3x3 sections and calculate the mean (or any function) within each. (If there's a leftover bit that isn't 3x3 then use just what's left).

I'm sure there's an apply-ish way to do this -- it's on the tip of my tongue -- but my brain is currently failing me. I suppose it's a bit like a moving window question except I want non-overlapping windows (so it's easier).

Can anyone think of an inbuilt function that does this? Or a vectorised way?

Here's my loopy version:

winSize <- 3
mat <- matrix(runif(6*11),nrow=6,ncol=11)
nr <- nrow(mat)
nc <- ncol(mat)
outMat <- matrix(NA,nrow=ceiling(nr/winSize),
                    ncol=ceiling(nc/winSize))
FUN <- mean
for ( i in seq(1,nr,by=winSize) ) {
    for ( j in seq(1,nc,by=winSize) ) {
        # work out mean in 3x3 window, fancy footwork
        #  with pmin just to make sure we don't go out of bounds
        outMat[ ceiling(i/winSize), ceiling(j/winSize) ] <-
               FUN(mat[ pmin(i-1 + 1:winSize,nr), pmin(j-1 + 1:winSize,nc)])
    }
}

cheers.

Answer 1

Just want to summarise the different methods for this.

First, @VincentZoonekynd's solution. This is very general -- it allows me to apply any function to my matrix. However it is a little slow because I am applying these to matrices of order ~5000x1000x3 and want back out a (5000/kernelSize) x (1000/kernelSize) x 3 image.

First, generate a matrix to test on (I made it smaller so as not to kill my computer whilst testing various methods):

sz <- c(1000,300,3)
img <- array(runif(prod(sz)),dim=sz)
kernelSize <- 3
outSz <- c(ceiling(sz[1:2]/kernelSize),3)
FUN <- mean

Method 0: loop

############
# METHOD 0 #
############
# Loopy. base standard.
t0 <- system.time({
out0 <- array(NA,dim=outSz)
for ( i in seq(1,sz[1],by=kernelSize) ) {
    for ( j in seq(1,sz[2],by=kernelSize) ) {
        for ( c in 1:sz[3] ) {
        # work out mean in 3x3 window, fancy footwork
        #  with pmin just to make sure we don't go out of bounds
        out0[ ceiling(i/kernelSize), ceiling(j/kernelSize),c ] <-
               FUN(img[ pmin(i-1 + 1:kernelSize,sz[1]), 
                        pmin(j-1 + 1:kernelSize,sz[2]),
                        c]) 
        }
    }
}})

Method 1: tapply

############
# METHOD 1 #
############
# @Vincent Zoonekynd.
# I can apply *any* function I want. how awesome!
# NOTE: I just realised that there is a slice.index(img,i)
#       is the same as his a(img,i) function.
t1 <- system.time({
out1 <- tapply(
         img,
         list( floor((slice.index(img,1)-1)/kernelSize), 
               floor((slice.index(img,2)-1)/kernelSize),
               slice.index(img,3) ),
         FUN )
})

cat('METHOD 0:',t0['elapsed'],'\n')
cat('METHOD 1:',t1['elapsed'],'\n')
cat(all(out0==out1),'\n')

This gives:

METHOD 0: 13.549 
METHOD 1: 19.415 
TRUE

Which are a bit slow, given that I would like to apply this to bigger img matrices.

What surprised me (at first) was that METHOD 0 (loops) was faster than METHOD 1 (tapply).

However, then I remembered that tapply has a reputation for being not much faster than an explicit loop (why is that? I remember reading it somewhere...the function code looks like it might do the for loop anyway, as opposed to calling external code).

I also have this general feeling that vapply and sapply are the fast versions of apply (again, not sure if this is definitively true but I've certainly found so).

Method 2: vapply

So, I tried to rewrite my loopy version using vapply. (There's probably a better way to do handle the 3rd dimension, but oh well...). This basically generates a big list of coordinates into img. The coordinates give the (i,j) of the corner of each kernelSize*kernelSize square.

Then vapply loops through them all and calculates the mean.

##########
# METHOD 2 
##########
# use 'vapply'
t2 <- system.time({
is <- seq(1,sz[1],by=kernelSize)
js <- seq(1,sz[2],by=kernelSize)
# generate a (nrow*nsize) x 2 array with
# all (i,j) combinations for corners of
# kernelSize*kernelSize squares.
# Do it column-major so we can reshape after.
coords <- cbind( rep.int(is,length(js)), rep(js,each=length(is)) ) 
out2 <- array(NA,dim=outSz)
for ( c in 1:sz[3] ) { 
    out2[,,c] <- array(
    vapply( 1:nrow(coords), function(i) {
          FUN(img[coords[i,1]:pmin(sz[1],coords[i,1]+kernelSize-1),
                   coords[i,2]:pmin(sz[2],coords[i,2]+kernelSize-1),
                   c])
            }, 0 ),
                dim=outSz[1:2] ) 
}})
cat('METHOD 2:',t2['elapsed'],'\n')
cat(all(out0==out2),'\n')

This gives:

METHOD 2: 12.627 
TRUE

So, it is a bit faster than a loop to use vapply (I do feel like I'm not getting as much out of vapply as I could be though...like I'm not using it in the right way).

Method 3: filter

This still isn't quite fast enough, so I then incorporated the information that I only want a mean in each window, and this is basically a convolution of [ 1/3 1/3 1/3 ] with the matrix in each dimension.

This loses the general applicability of applying an arbitrary FUN but gets big speedups in return.

Basically, I make a kernel [1/3, 1/3, 1/3] and convolve it with img twice, once in the x direction and once in the y. Then I only extract every 3rd value (since I wanted non-overlapping windows).

This seems a bit wasteful to me in that I calculate the mean for every 3x3 window in my original matrix, instead of just non-overlapping windows, but I don't know how to tell R not to calculate those values that I'm going to throw away anyway.

However you have to take a bit of care at the borders -- say there's only a 2x2 patch left over, then the mean is over 4 instead of 9 values. My current code doesn't handle this, but I don't mind if it's just the border that's out, because I'm only doing the downsampling for display purposes.

(It would be nice to fix this one last thing though...)

##########
# METHOD 3 
##########
# Convolve using `filter`,
# since the mean in a window is just a 
# convolution.
t3 <- system.time({
is <- pmin(seq(1,sz[1],by=kernelSize) + floor(kernelSize/2),sz[1]-1)
js <- pmin(seq(1,sz[2],by=kernelSize) + floor(kernelSize/2),sz[2]-1)
out3 <- array(NA,dim=outSz)
for ( c in 1:3 ) {
    out3[,,c] <- (t(filter(
                    t(filter(img[,,c],rep(1,kernelSize))),
                    rep(1,kernelSize))))[is,js]
}
out3 <- out3/(kernelSize*kernelSize)
})
cat('METHOD 3:',t3['elapsed'],'\n')
cat(sum(out0!=out3),'\n')

This returns:

METHOD 3: 1.593 
300

So this method is by far the quickest, and the error just along the last column of out3 (once per channel), since (I guess) there are border conditions.

Answer 2

You can use row and col to extract the row and column numbers, and then compute the coordinates of each block.

tapply( 
  mat, 
  list( floor((row(mat)-1)/winSize), floor((col(mat)-1)/winSize) ), 
  mean 
)

Edit: This can be generalized to higher-dimensional arrays, by replacing row and col with the following function.

a <- function( m, k ) {
  stopifnot( "array" %in% class(m) || "matrix" %in% class(m) )
  stopifnot( k == floor(k) )
  stopifnot( k > 0 )
  n <- length(dim(m))
  stopifnot( k <= n )
  i <- rep(
    1:dim(m)[k],
    each  = prod(dim(m)[ 1:n < k ]),
    times = prod(dim(m)[ 1:n > k ])
  )  
  array(i, dim=dim(m))
}

# A few tests
m <- array(NA, dim=c(2,3))
all( row(m) == a(m,1) )
all( col(m) == a(m,2) )
# In dimension 3, it can be done manually:
m <- array(NA, dim=c(2,3,5))
all( a(m,1) == array( rep(1:dim(m)[1], times=prod(dim(m)[2:3])), dim=dim(m) ) )
all( a(m,2) == array( rep(1:dim(m)[2], each=dim(m)[1], times=dim(m)[3]), dim=dim(m) ) )
all( a(m,3) == array( rep(1:dim(m)[3], each=prod(dim(m)[-3])), dim=dim(m) ) )