无法使用 np.where() 在复杂的 2D Python 数组中找到元素的索引

Question

I am trying to use numpy.where() to find out the index of an element, but it returns an empty array.

import numpy as np
grid= np.mgrid[-2:2:5*1j, -2:2:11*1j]
X , Y = grid[0], grid[1]
complex_grid = X+1j*Y
xid, yid= np.where(complex_grid == -1.0 - 0.8j)
print(xid, yid)

It should return index (1,3), but it returns an empty array and its data type. What am I doing wrong?

EDIT :- My main aim is to find the index corresponding to a given coordinate (x,y) from a grid. I made a complex grid just because I can fuse the two 2D matrices I get from mgrid.

Answer 1

Floating point values are seldom exact, so it's generally considered a bad idea to directly compare them.

However you can use something like abs(difference) < epsilon to check for a given closeness to a value, for example:

>>> xid, yid= np.where(np.abs(complex_grid -(-1.0 - 0.8j)) < 1e-10)
>>> print(xid, yid)
[1] [3]

or even better: Use numpy.isclose which already does that but allows relative and absolute tolerances and nan handling (if you need these) as well:

>>> xid, yid= np.where(np.isclose(complex_grid, -1 - 0.8j)))
>>> print(xid, yid)
[1] [3]

Answer 2

isclose()是更好的选择，因为比较浮点数总是很冒险：

xid, yid = np.where(np.isclose(complex_grid, np.ones(complex_grid.shape) * (-1.0 - 0.8j)))