'What is the best way in numpy to use a vector of positions as indexes in a an nd array?
I have a numpy array of agents positions:
positions = np.array([[row_0, col_0],
[row_1, col_1],
[row_2, col_2]])
I have a 2d grid, grid = np.zeros((N, N))
. Now, let's say I want to mark all the cells where the agents are with True
. What I have been doing is this:
grid[positions[:, 0], positions[:, 1]] = True
However, this is wordy and not scalable. Let's say now the positions are in 3d space. Now we have to modify the code like so:
grid[positions[:, 0], positions[:, 1], positions[:, 2]] = True
There is another way to solve the scalability problem, which is to do this:
grid[tuple(positions.T)] = True
However, it a involves a function call. Additionally, let's say there is not one, but two grids, such as:
grid = np.zeros((2, N, N))
and I wanted to mark only the first grid with True, then the method above does not work unless I do:
grid[0, (*tuple(positions.T))] = True
which does not look readable, and I wonder if there is a better way to do this. Any help is appreciated!
EDIT: I have been looking into this, and it finally hit me that adding a comma ,
in the parenthesis will allow tuple unpacking without explicitly calling the tuple constructor, like so:
grid[(*positions.T,)] = True
For the case of modifying the grid at index n of many grids combined at axis 0, the command becomes:
grid[(n, *positions.T,)] = True
Sources
This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.
Source: Stack Overflow
Solution | Source |
---|