lanery lanery - 1 month ago 10
Python Question

How to fill numpy array of zeros with ones given indices/coordinates

Given a numpy array of zeros, say

arr = np.zeros((5, 5))


and an array of indices that represent vertices of a polygon, say

verts = np.array([[0, 2], [2, 0], [2, 4]])


1) What is the elegant way of doing

for v in verts:
arr[v[0], v[1]] = 1


such that the resulting array is

In [108]: arr
Out[108]:
array([[ 0., 0., 1., 0., 0.],
[ 0., 0., 0., 0., 0.],
[ 1., 0., 0., 0., 1.],
[ 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0.]])


2) How can I fill the array with ones such that the output array is

In [158]: arr
Out[158]:
array([[ 0., 0., 1., 0., 0.],
[ 0., 1., 1., 1., 0.],
[ 1., 1., 1., 1., 1.],
[ 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0.]])

Answer

To answer the first part of your question: arr[tuple(verts.T)] = 1

verts.T transposes your indices to a (2, n) array, where the two rows correspond to the row and column dimensions of arr. These are then unpacked into a tuple of (row_indices, col_indices), which we then use to index into arr.

We could write this a bit more verbosely as:

row_indices = verts[:, 0]
col_indices = verts[:, 1]
arr[row_indices, col_indices] = 1

For the second part, one method that will work for arbitrary polygons would be to use matplotlib.Path.contains_points, as described here:

from matplotlib.path import Path

points = np.indices(arr.shape).reshape(2, -1).T
path = Path(verts)
mask = path.contains_points(points, radius=1e-9)
mask = mask.reshape(arr.shape).astype(arr.dtype)

print(repr(mask))
# array([[ 0.,  0.,  1.,  0.,  0.],
#        [ 0.,  1.,  1.,  1.,  0.],
#        [ 1.,  1.,  1.,  1.,  1.],
#        [ 0.,  0.,  0.,  0.,  0.],
#        [ 0.,  0.,  0.,  0.,  0.]])
Comments