How can all dimensions left after the specified one be preserved, without explicitly listing them?

Question

^{Or equivalently, "what is the equivalent of NumPy's ellipsis indexing in Matlab"}

Say I have some high-dimensional array:

x = zeros(3, 4, 5, 6);

I want to write a function that takes an array of size (3, ...), and does some computation. In NumPy, I could write this:

def fun(x):
    return x[0]*x[1] + x[2]

However, the equivalent in MATLAB doesn't work, because indexing with one integer flattens the array to 1d

function y = fun_bad(x)
    y = x(1)*x(2) + x(3)

I can make this work for up to 3-dimensional arrays with

function y = fun_ok3d(x)
    y = x(1,:,:)*x(2,:,:) + x(3,:,:)

If I want this to work for up to 10-dimensional arrays, I can write

function y = fun_ok10d(x)
    y = x(1,:,:,:,:,:,:,:,:,:)*x(2,:,:,:,:,:,:,:,:,:) + x(3,:,:,:,:,:,:,:,:,:)

How can I avoid writing stupid numbers of colons here, and just make this work for any dimension? Is there some x(1,...) syntax that implies this?

NumPy can use the ... (Ellipsis) literal in an indexing expression to mean ": as many times as needed", which would solve this problem.

Answer 1

Approach 1: using a comma-separated list with ':'

I don't know a way to specify

: as many times as needed

while preserving shape. But you can specify

: an arbitrary number of times

where that number of times is defined at run-time. With this method you can preserve shape, provided that the number of indices coincides with the number of dimensions.

This is done using a comma-separated list generated from a cell array, and exploiting the fact that the string ':' can be used as an index instead of ::

function y = fun(x)
colons = repmat({':'}, 1, ndims(x)-1); % row cell array containing the string ':'
                                       % repeated the required number of times
y = x(1,colons{:}).*x(2,colons{:}) + x(3,colons{:});

This approach can be easily generalized to indexing along any dimension, not just the first:

function y = fun(x, dim)
% Input argument dim is the dimension along which to index
colons_pre = repmat({':'}, 1, dim-1);
colons_post = repmat({':'}, 1, ndims(x)-dim);
y = x(colons_pre{:}, 1, colons_post{:}) ...
  .*x(colons_pre{:}, 2, colons_post{:}) ...
  + x(colons_pre{:}, 3, colons_post{:});

Approach 2: splitting the array

You can split the array along the first dimension using num2cell, and then apply the operation to the resulting subarrays. Of course this uses more memory; and as noted by @Adriaan it is slower.

function y = fun(x)
xs = num2cell(x, [2:ndims(x)]); % x split along the first dimension
y = xs{1}.*xs{2} + xs{3};

Or, for indexing along any dimension:

function y = fun(x, dim)
xs = num2cell(x, [1:dim-1 dim+1:ndims(x)]); % x split along dimension dim
y = xs{1}.*xs{2} + xs{3};

Answer 2

MATLAB flattens all trailing dimensions when using a single colon, so you can use that to get from your N-D array to a 2D array, which you can reshape back into the original N dimensions after the calculation.

Along the first dimension

If you want to use the first dimension you can use a relatively simple and short piece of code:

function y = MyMultiDimensional(x)
    x_size = size(x); % Get input size
    yflat = x(1,:) .* x(2,:) + x(3,:); % Calculate "flattened" 2D function
    y = reshape(yflat, [1 x_size(2:end)]); % Reshape output back to original size
end

Along an arbitrary dimension, now featuring N-D permute.

When you want your function to act along the n-th dimension out of a total of N, you can permute that dimension to the front first:

function y = MyMultiDimensional(x,n)
    x_size = size(x); % Get input size

    Order = 1:numel(x_size);
    Order(n)=[]; % Remove n-th dimension
    Order2 = [n, Order]; % Prepend n-th dimension

    xPermuted = permute(x,Order2); % permute the n-th dimension to the front
    yTmp = xPermuted (1,:) .* xPermuted (2,:) + xPermuted (3,:); % Calculate "flattened" 2D function
    y = reshape(yTmp, x_size(Order)); % Reshape output back to original size
end

I timed the results of the two methods of Luis' and my methods:

function timeMultiDim()

x = rand(1e1,1e1,1e1,1e1,1e1,1e1,1e1,1e1);

    function y = Luis1(x)
        colons = repmat({':'}, 1, ndims(x)-1); % row cell array containing the string ':'
        % repeated the required number of times
        y = x(1,colons{:}).*x(2,colons{:}) + x(3,colons{:});
        
    end

    function y = Luis2(x)
        xs = num2cell(x, [2:ndims(x)]); % x split along the first dimension
        y = xs{1}.*xs{2} + xs{3};
    end

    function y = Adriaan(x)
        x_size = size(x); % Get input size
        yflat = x(1,:) .* x(2,:) + x(3,:); % Calculate "flattened" 2D function
        y = reshape(yflat, [1 x_size(2:end)]); % Reshape output back to original size
    end

n=1;
    function y = Adriaan2(x,n)
        x_size = size(x); % Get input size
        
        Order = 1:numel(x_size);
        Order(n)=[]; % Remove n-th dimension
        Order2 = [n, Order]; % Prepend n-th dimension
        
        xPermuted = permute(x,Order2); % permute the n-th dimension to the front
        yTmp = xPermuted (1,:) .* xPermuted (2,:) + xPermuted (3,:); % Calculate "flattened" 2D function
        y = reshape(yTmp, x_size(Order)); % Reshape output back to original size
        
    end

t1 = timeit(@() Luis1(x));
t2 = timeit(@() Luis2(x));
t3 = timeit(@() Adriaan(x));
t4 = timeit(@() Adriaan2(x,n));

format long g;
fprintf('Luis 1: %f seconds\n', t1);
fprintf('Luis 2: %f seconds\n', t2);
fprintf('Adriaan 1: %f seconds\n', t3);
fprintf('Adriaan 2: %f seconds\n', t4);

end

Luis 1: 0.698139 seconds
Luis 2: 4.082378 seconds
Adriaan 1: 0.696034 seconds
Adriaan 2: 0.691597 seconds

So, going to a cell is bad, it takes more than 5 times as long, reshape and ':' are barely apart, so that'd come down to preference.