Flat indexing of all but first dimension with Numpy

Question

Is there some way to use flat indexing for the remaining dimensions with NumPy? I'm trying to translate the following MATLAB function to 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

but I can't wrap my head around what the NumPy equivalent of MATLAB's weights(:, ix) would be---that is, linear indexing only in the remaining dimensions.

I was hoping that the syntax could be directly translated, but suppose that values is a 3-by-4 array, then weights becomes a 2-by-3-by-4 array. In MATLAB, weights(:, ix) is then a 2-by-1 array, whereas in Python weights[:, ix] is a 2-by-3 array.

I think that I have handled everything else in the function below.

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

Do you have any suggestions? Perhaps I'm just thinking about the problem all wrong. I have also tried to write the code as simply as possible as I intend to use Numba to speed it up later.

`ix` is just a scalar, so the equivalent for `x(:,ix)` in MATLAB would be `x[:,ix]` in Python (take care of off-by-one errors due to 1-based vs 0-based indexing). There's no remaining dimensions or something here. In MATLAB, this takes column `ix` from a 2D matrix, which is trivial in NumPy as well. — Adriaan, Dec 17 '21 at 10:48
@Adriaan Not quite. Suppose in MATLAB (Python) `values = zeros(3, 4)` (`values = np.zeros((3, 4))`. Then `weights` is a 2-by-3-by-4 array and `weights(:, 1)` a 2-by-1 array (`weights[:, 1]` a 2-by-3 array). — Fredrik P, Dec 17 '21 at 11:13
In that case you can either do `weights[:, :, ix]` expliticly or `weights[..., ix]`, although I'm not familiar enough with Python whether front-ellipses work. — Adriaan, Dec 17 '21 at 11:22
@Adriaan I think that both of those would give me 2-by-3-by-1 arrays in the example from my last comment. But thank you for trying! — Fredrik P, Dec 17 '21 at 11:44
Could you please [edit] the question to make a [mcve]? I.e. take everything out except this indexing issue and add some sample data (`rand(3,4)` for example), showing what MATLAB does and what you currently obtain in Python? All the loops and `if` statements seem irrelevant to this indexing issue, making it harder to grasp. — Adriaan, Dec 17 '21 at 12:07
There isn't a direct `numpy` equivalent. It's MATLAB that's playing funny games here. — hpaulj, Dec 17 '21 at 13:59
To understand what MATLAB is doing I'd have to run the code line by line in an Octave session, paying attention to shapes at each. — hpaulj, Dec 17 '21 at 14:06
@hpaulj These “funny games” is just linear indexing, which is terribly useful and I often miss when working with NumPy. — Cris Luengo, Dec 17 '21 at 14:15

Fredrik P · Accepted Answer · 2021-12-17T14:31:16.807

As per hpaulj's comment, 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

This is indeed the only solution AFAIK. Python doesn’t do linear indexing, which is what you are doing in MATLAB in the 2nd and 3rd dimensions. — Cris Luengo, Dec 17 '21 at 14:13

Flat indexing of all but first dimension with Numpy

1 Answers1