Obtain matrix of indices in octave / matlab

Question

Given some multidimensional matrix A in Octave / Matlab,

What's the easiest way to get a matrix of the same size as A where all elements are replaced by their index along the k'th dimension

ie for the matrix

A =

ans(:,:,1) =

   0.095287   0.191905
   0.226278   0.749100

ans(:,:,2) =

   0.076826   0.131639
   0.862747   0.699016

I want a function f such that f(A,1) =

ans(:,:,1) =

   1   1
   2   2

ans(:,:,2) =

   1   1
   2   2

f(A,2) =

ans(:,:,1) =

   1   2
   1   2

ans(:,:,2) =

   1   2
   1   2

and

f(A, 3) =

ans(:,:,1) =

   1   1
   1   1

ans(:,:,2) =

   2   2
   2   2

Also, given a sparse matrix B

What's the easiest way to get another sparse matrix of the same size where the nonzero elements are replaced by their index along the k'th dimension? (so same problem as above, but for only the nonzero elements)

Ideally I'm looking for a way which is well-vectorized for octave (meaning it doesn't explicitly loop over anything)

CLARIFICATION: For the sparse matrix one, I'm looking for a solution which does not involve creating a full size(B) matrix at any point

Answer 1

ndgrid() does what you want, although not in the format you are looking for. If you know the dims of the input A beforehand, you can use the following line to create the N-dimentional mesh grid:

% for matrix a where ndims(a) == 3
[x, y, z] = ndgrid (1:size(a,1), 1:size(a,2), 1:size(a,3));
% x is like f(a, 1)
% y is like f(a, 2)
% z is like f(a, 3)

You may be able to write a custom wrapper around ndgrid() to convert it to the function format you are looking for.

Answer 2

In case anyone's curious, since I didn't know about ndgrid, here's the answer I came up with:

function [y] = indices(a,k)
    s = size(a);
    n = s(k);
    D = length(s);
    x = permute(a,[k,1:(k-1),(k+1):D]);

    y = reshape(x,n,[]);
    y = diag(1:n) * ones(size(y));
    y = reshape(y,size(x));

    y = permute(y,[(2:k),1,(k+1):D]);
endfunction

function [y] = spindices(a,k)
    s = size(a);
    n = s(k);
    D = length(s);
    x = permute(a,[k,1:(k-1),(k+1):D]);

    y = reshape(x,n,[]);
    y = spdiag(1:n) * spones(y);
    y = reshape(y,size(x));

    y = permute(y,[(2:k),1,(k+1):D]);
endfunction