Matlab: Arrayfun with Matrices

Question

I am trying to rewrite the following code using arrayfun

A = ones(3,3,3)
for i = 1:3
    B(i) = trace(A(:,:,i));
end

I hopefully tried

f = @(x) trace(x)
B = arrayfun(f, A);

But this just (as you would expect) traces over each individual A(i,j,k) not the A(:,:,i) as I would like. I then tried declaring A{i}=ones(3,3) as a cell and passing to arrayfun but this didn't work either.

How can I vectorise functions on matrices in Matlab?

Answer 1

bsxfun based vectorized solution that ab(uses) how trace is defined - sum of diagonal elements -

%// Get size of A
[m,n,r] = size(A) 

%// Get indices of the diagonal elements for each 3D "slice" as columns of idx
idx = bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n) %//'

%// Thus, for your 3 x 3 x 3 case, idx would be -
%//idx =
%//      1    10    19
%//      5    14    23
%//      9    18    27
%// and these are the linear indices to the diagonal elements to each `3D` slide.

%//Index into A with idx and sum along columns to get each element of desired output
B = sum(A(idx),1)

If you would like to save on workspace cluttering with not-so-necessary extra variables, avoid idx with -

B = sum(A(bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n)),1)

---

For working with GPUs

If you have to work with GPUs, you can declare them as gpuArrays with gpuArray(A), then subsequent work involving A would be done on GPU and you would get the output as a gpuArray, which you can get back as a CPU variable with gather(..).

Thus, a complete code would look like this -

[m,n,r] = size(A); %// Get size
gpu_A = gpuArray(A); %// copy data from CPU to GPU

%// Perform calculations on GPU
gpu_B = sum(gpu_A(bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n)),1); %//'

B = gather(gpu_B); %// get back output onto CPU

Quick tests: With GTX 750 Ti (that I have access to), this appears to give me a 3x speedup over your loopy code.

Answer 2

If you really want to use arrayfun you could try a trick like this:

arrayfun(@(i)trace(A(:,:,i)), 1:size(A,3))

but note that arrayfun is NOT VECTORIZATION!!, it's just a wrapper for a loop and often is slower than a loop due to added overhead.

While equally not vectorization, to do it your second way you should have just changed to cellfun. i.e. if A{i} = ones(3,3) then

cellfun(@(x)trace(x), A)