Adam Idris
Adam Idris

Reputation: 75

How do search through a numpy array to find a corresponding value of two coordinates?

I have an array which has 3 columns: X, Y, and velocity. How do I write a function such that I can search through the array by inputting the X,Y coordinates and the function returns the corresponding velocity value for that index? Let's say my array is the following:

srcxy_deltv = [[2500 0 3.4123]
              [23000 0 3.4213]
              [23500 0 3.4303]
              ...
              [33675 25500 0.49377]
              [33725 25500 0.49878]
              [33775 25500 0.50381]]

The first column is the X, second coordinate is the Y and third is the velocity. I want to be able to feed both X and Y values into the function to return the velocity. For example:

srcxy(2500,0)

That should return the value 3.4123

Upvotes: 0

Views: 117

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150735

You can do a lookup like this:

srcxy_deltv = np.array(srcxy_deltv)

def srcxy(x,y, data=srcxy_deltv):
    return srcxy_deltv[(srcxy_deltv[:,:2] == [2500,0]).all(1),-1][0]

srcxy(2500,0)
# 3.4123

If you are open for other packages, pandas can be a good choice:

df = pd.DataFrame(srcxy_deltv).set_index([0,1])

df.loc[(2500,0)].iloc[0]
# 3.4123

df.loc[(33725, 25500)].iloc[0]
# 0.49878

Upvotes: 1

Related Questions