使用 Numpy 对除第一个维度之外的所有维度进行平面索引

Question

有没有什么方法可以通过 NumPy 对剩余的维度使用平面索引？我正在尝试将以下 MATLAB 函数转换为 Python

function [indices, weights] = locate(values, gridpoints)
    indices = ones(size(values));
    weights = zeros([2, size(values)]);

    for ix = 1:numel(values)
        if values(ix) <= gridpoints(1)
            indices(ix) = 1;
            weights(:, ix) = [1; 0];
        elseif values(ix) >= gridpoints(end)
            indices(ix) = length(gridpoints) - 1;
            weights(:, ix) = [0; 1];
        else
            indices(ix) = find(gridpoints <= values(ix), 1, 'last');    
            weights(:, ix) = ...
                [gridpoints(indices(ix) + 1) - values(ix); ...
                 values(ix) - gridpoints(indices(ix))] ...
                / (gridpoints(indices(ix) + 1) - gridpoints(indices(ix)));
        end
    end
end

但我无法理解与 MATLAB 的 weights(:, ix) 等效的 NumPy 是什么——也就是说，仅在剩余维度中进行线性索引。

我希望语法可以直接翻译，但是假设 values 是一个 3×4 的数组，那么 weights 就变成了 2×3×4大批。在 MATLAB 中，weights(:, ix) 是一个 2×1 数组，而在 Python 中，weights[:, ix] 是一个 2×3 数组。

我认为我已经处理了下面函数中的所有其他内容。

import numpy as np


def locate(values, gridpoints):
    indices = np.zeros(np.shape(values), dtype=int)
    weights = np.zeros((2,) + np.shape(values))

    for ix in range(values.size):
        if values.flat[ix] <= gridpoints[0]:
            indices.flat[ix] = 0
            # weights[:, ix] = [1, 0]
        elif values.flat[ix] >= gridpoints[-1]:
            indices.flat[ix] = gridpoints.size - 2
            # weights[:, ix] = [0, 1]
        else:
            indices.flat[ix] = (
                np.argwhere(gridpoints <= values.flat[ix]).flatten()[-1]
            )
            # weights[:, ix] = (
            #         np.array([gridpoints[indices.flat[ix] + 1] - values.flat[ix],
            #                   values.flat[ix] - gridpoints[indices.flat[ix]]])
            #         / (gridpoints[indices.flat[ix] + 1] - gridpoints[indices.flat[ix]])
            # )

    return indices, weights

你有什么建议吗？也许我只是在想这个问题。我也尝试尽可能简单地编写代码，因为我打算稍后使用 Numba 来加速它。

Answer 1

根据 hpaulj's , there doesn't seem to be a direct NumPy equivalent. In lack thereof, the best I can think of is to reshape the weights array as in the code below and the suggestion from NumPy for Matlab Users.

import numpy as np


def locate(values, gridpoints):
    indices = np.zeros(values.shape, dtype=int)
    weights = np.zeros((2, values.size))  # Temporarily make weights 2-by-N

    for ix in range(values.size):
        if values.flat[ix] <= gridpoints[0]:
            indices.flat[ix] = 0
            weights[:, ix] = [1, 0]
        elif values.flat[ix] >= gridpoints[-1]:
            indices.flat[ix] = gridpoints.size - 2
            weights[:, ix] = [0, 1]
        else:
            indices.flat[ix] = (
                np.argwhere(gridpoints <= values.flat[ix]).flatten()[-1]
            )
            weights[:, ix] = (
                    np.array([gridpoints[indices.flat[ix] + 1] - values.flat[ix],
                              values.flat[ix] - gridpoints[indices.flat[ix]]])
                    / (gridpoints[indices.flat[ix] + 1] - gridpoints[indices.flat[ix]])
            )
    
    # Give weights correct dimensions
    weights.shape = (2,) + values.shape
    
    return indices, weights

使用 Numpy 对除第一个维度之外的所有维度进行平面索引

Flat indexing of all but first dimension with Numpy

python

indexing

matlab

numpy

matrix-indexing