在没有明确列出的情况下，如何保留指定维度之后的所有维度？

Question

^{或者等价的，“Matlab中NumPy的省略号索引是什么等价物”}

假设我有一些高维数组：

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

我想编写一个函数，它接受一个大小为(3, ...) 的数组，并进行一些计算。在 NumPy 中，我可以这样写：

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

但是，MATLAB 中的等效项不起作用，因为使用一个整数进行索引会将数组展平为 1d

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

我可以用

使这项工作适用于最多 3 维数组

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

如果我希望它适用于最多 10 维的数组，我可以编写

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

我怎样才能避免在这里写愚蠢的冒号数字，而让它适用于任何维度？是否有一些 x(1,...) 语法暗示了这一点？

NumPy 可以在索引表达式中使用... (Ellipsis) 文字来表示“: 任意多次”，这将解决这个问题。

Answer 1

方法 1：使用逗号分隔的列表和 ':'

我不知道如何指定

: 根据需要多次

同时保持形状。但是你可以指定

:任意次

该次数是在运行时定义的。使用这种方法，您可以保留形状，前提是索引数与维数一致。

这是使用从元胞数组生成的comma-separated list 完成的，并利用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{:});

这种方法可以很容易地推广到任何维度的索引，而不仅仅是第一个：

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{:});

方法二：拆分数组

您可以使用num2cell 沿第一个维度拆分数组，然后将操作应用于生成的子数组。当然这会占用更多内存；和 noted by @Adriaan 一样慢。

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

或者，对于任何维度的索引：

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 在使用单个冒号时会展平所有尾随维度，因此您可以使用它从您的 N-D 数组获取二维数组，您可以将其 reshape 恢复为原始 N 个维度计算后。

沿第一个维度

如果你想使用第一个维度你可以使用一段相对简单而简短的代码：

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

沿任意维度，现在具有 N-D 置换功能。

当您希望您的函数在总共 N 个维度中的第 n 个维度上运行时，您可以先permute 该维度到前面：

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

我对Luis和我的两种方法的结果进行了计时：

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

所以，去一个牢房是不好的，它需要超过 5 倍的时间，reshape 和 ':' 几乎没有分开，所以这归结为偏好。