R: Apply FUN to kxk subsections of array

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.

回答1:

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) ) )

回答2:

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.