CODEWITHSUNDEEP
pythonarraysnumpy
student
student

Reputation: 1646

Get element of an array that corresponds to the element of another array in numpy

Suppose I have two arrays in numpy, t and v and suppose that t is strictly monotonically increasing. In my example t represents an array of time points and v the corresponding speed of a body. Now I want for example get the velocity at time t = 3. How can I do this?

Upvotes: 2

Views: 370

Answers (2)

unutbu
unutbu

Reputation: 880707

For linear interpolation using only NumPy, you could use np.interp. For example,

import numpy as np
t = np.linspace(0, 5, 100)
v = np.sin(t)

To find the linearly interpolated value of v when t=3:

In [266]: np.interp(3, t, v)
Out[266]: 0.14107783526460238

Note that if you wish to interpolate v at many values of t, you can pass an iterable as the first argument to np.interp:

In [292]: np.interp(np.linspace(t.min(), t.max(), 10), t, v)
Out[292]: 
array([ 0.        ,  0.52741539,  0.8961922 ,  0.99540796,  0.79522006,
        0.35584199, -0.19056796, -0.67965796, -0.96431712, -0.95892427])

This is much more efficient than calling np.interp repeatedly for one value at a time.

To get an element of the array, v, which corresponds to t=3 you could use np.searchsorted:

In [272]: v[np.searchsorted(t, 3)]
Out[272]: 0.11106003812412972

Note, however, that np.searchsorted returns the index where 3 would be inserted into t to maintain its sortedness. So v[np.searchsorted(t, 3)] and v[np.searchsorted(t, 3)+1] sandwich the velocity at t=3.

Also beware that np.searchsorted may return an index which is 1 greater than the biggest valid index for t (and v). This happens if 3 > t.max():

For example, if t were [1,2,3]: 

In [277]: np.searchsorted([1,2,3], 5)
Out[277]: 3

So to protect against the possible IndexError, use np.clip to ensure the index is between 0 and len(v)-1:

idx = np.clip(np.searchsorted(t, 3), 0, len(v)-1)
v[idx]

Like np.interp, np.searchsorted can accept an iterable (here, for the second argument):

In [306]: v[np.clip(np.searchsorted(t, [3,4,5,6]), 0, len(v)-1)]
Out[306]: array([ 0.11106004, -0.7825875 , -0.95892427, -0.95892427])

Upvotes: 3

Hugh Bothwell
Hugh Bothwell

Reputation: 56694

If you have scipy there is a function that does precisely this:

from scipy.interpolate import interp1d

velocity = interp1d(t, v, kind="cubic")

print(velocity(3.0))

See the documentation at http://docs.scipy.org/doc/scipy-0.14.0/reference/tutorial/interpolate.html

Upvotes: 3

Related Questions