仅使用NumPy优化step函数的代码

我试图在下面的代码中只使用NumPy函数(或者列表理解)来优化函数“pw”。

from time import time
import numpy as np

def pw(x, udata):
    """
    Creates the step function
                 | 1,  if d0 <= x < d1
                 | 2,  if d1 <= x < d2
    pw(x,data) = ...
                 | N, if d(N-1) <= x < dN
                 | 0, otherwise
    where di is the ith element in data.
    INPUT:      x   --  interval which the step function is defined over
              data  --  an ordered set of data (without repetitions)
    OUTPUT: pw_func --  an array of size x.shape[0]
    """
    vals = np.arange(1,udata.shape[0]+1).reshape(udata.shape[0],1)
    pw_func = np.sum(np.where(np.greater_equal(x,udata)*np.less(x,np.roll(udata,-1)),vals,0),axis=0)
    return pw_func


N = 50000
x = np.linspace(0,10,N)
data = [1,3,4,5,5,7]
udata = np.unique(data)

ti = time()
pw(x,udata)
tf = time()
print(tf - ti)

import cProfile
cProfile.run('pw(x,udata)')

cProfile。run告诉我,大部分开销来自np。在哪里(大约1毫秒),但如果可能的话,我想创建更快的代码。似乎行操作和列操作的执行会产生一些不同,除非我弄错了,但我认为我已经解释了这一点。我知道有时列表理解会更快,但我想不出比我用它做什么更快的方法。

Searchsorted似乎可以产生更好的性能,但我的电脑上仍然保留着1毫秒的时间:

(modified)
def pw(xx, uu):
    """
    Creates the step function
                 | 1,  if d0 <= x < d1
                 | 2,  if d1 <= x < d2
    pw(x,data) = ...
                 | N, if d(N-1) <= x < dN
                 | 0, otherwise
    where di is the ith element in data.
    INPUT:      x   --  interval which the step function is defined over
              data  --  an ordered set of data (without repetitions)
    OUTPUT: pw_func --  an array of size x.shape[0]
    """
    inds = np.searchsorted(uu, xx, side='right')
    vals = np.arange(1,uu.shape[0]+1)
    pw_func = vals[inds[inds != uu.shape[0]]]
    num_mins = np.sum(xx < np.min(uu))
    num_maxs = np.sum(xx > np.max(uu))

    pw_func = np.concatenate((np.zeros(num_mins), pw_func, np.zeros(xx.shape[0]-pw_func.shape[0]-num_mins)))
    return pw_func

这 answer 使用分段似乎非常接近,但这是在标量x0和x1上。如何在阵列上执行该操作?这样会更有效率吗?

可以理解的是,x可能相当大,但我正在尝试进行压力测试。

我仍在学习,所以一些可以帮助我的提示或技巧会很好。

编辑

第二个函数中似乎有一个错误,因为第二个函数生成的数组与第一个不匹配(我相信它可以工作):

N1 = pw1(x,udata.reshape(udata.shape[0],1)).shape[0]
N2 = np.sum(pw1(x,udata.reshape(udata.shape[0],1)) == pw2(x,udata))
print(N1 - N2)

产量

15000

数据点是 不 相同的。所以我似乎不知道如何使用“searchsorted”。

编辑2

实际上我已经修好了:

pw_func = vals[inds[inds != uu.shape[0]]]

改为

pw_func = vals[inds[inds[(inds != uu.shape[0])*(inds != 0)]-1]]

所以至少得到的数组是匹配的。但问题仍然在于是否有更有效的方法来实现这一点。

编辑3

谢谢田丽指出错误。这个应该有用

pw_func = vals[inds[(inds != uu.shape[0])*(inds != 0)]-1]

也许更具可读性的表达方式是

non_endpts = (inds != uu.shape[0])*(inds != 0) # only consider the points in between the min/max data values
shift_inds = inds[non_endpts]-1       # searchsorted side='right' includes the left end point and not right end point so a shift is needed
pw_func = vals[shift_inds]

我想我在那些括号里迷路了!我想这就是可读性的重要性。