使用 numpy.where 创建掩码

Question

我有这个用于创建掩码（布尔数组）的函数，我希望更快。

def get_validity_1(ts, times):
    validity = numpy.zeros(len(ts))
    indexes = []
    for start, end in times:
        index_start = numpy.argmax(ts >= start)
        index_end = numpy.argmax(ts >= end)
        indexes.append([index_start, index_end])
    for start, end in indexes:
        validity[start:end] = 1
    return validity
res_1 = get_validity_1(numpy.linspace(0, 1, 100000000), numpy.array([[0.01, 0.1], [0.5, 0.8]]))

这个问题的问题是如何使用 numpy.where 条件来制作它。我试过这个：

def get_validity_2(ts, times):
    return numpy.where(numpy.logical_or([t1<ts.all()<t2 for t1, t2 in times]))

但是python会提高：

ValueError: invalid number of arguments

这里有一些输入断言：

• ts[n-1]
• 次[n][0]
• 次[n-1][1]

这是一个脚本作为输入：

import time, numpy

def get_validity_1(ts, times):
    validity = numpy.zeros(len(ts))
    indexes = []
    for start, end in times:
        index_start = numpy.argmax(ts >= start)
        index_end = numpy.argmax(ts >= end)
        indexes.append([index_start, index_end])
    for start, end in indexes:
        validity[start:end] = 1
    return validity
    
def get_validity_2(ts, times):
    return numpy.where(numpy.logical_or([t1<ts.all()<t2 for t1, t2 in times]))

if __name__ == "__main__":
    n = 100000000
    ts = numpy.linspace(0, 1, n)
    
    times = numpy.array([[0.01, 0.1], [0.5, 0.8]])
    
    t0 = time.time()
    res_1 = get_validity_1(ts, times)
    t_1 = time.time() - t0
    
    t0 = time.time()
    res_2 = get_validity_2(ts, times)
    t_2 = time.time() - t0
    
    print("t_1: " + str(t_1))
    print("t_2: " + str(t_2))
    
    assert res_1 == res_2
    assert t_1 > t_2

有谁知道如何完成函数“get_validity_2”并传递断言？ 或者只是一个包的功能来解决这个问题？

Answer 1

np.logical_or(*[np.logical_and(t1<ts, ts<t2) for t1, t2 in times])

如果您想要一个像您想要实现的 1-liner。但是，这仍然是低效的，因为您要在 O(N) 中比较大型数组。

由于 ts 已排序，这里有一个更快的方法来使用二进制搜索在 O(log(N)) 中找到开始/结束索引：

def get_validity_3(ts, times):
    validity = numpy.zeros(len(ts))
    for start, end in times:
        index_start = np.searchsorted(ts, start)
        index_end = np.searchsorted(ts, end)
        validity[index_start:index_end] = 1
    return validity

整体代码：

import time, numpy

def get_validity_1(ts, times):
    validity = numpy.zeros(len(ts))
    indexes = []
    for start, end in times:
        index_start = numpy.argmax(ts >= start)
        index_end = numpy.argmax(ts >= end)
        indexes.append([index_start, index_end])
    for start, end in indexes:
        validity[start:end] = 1
    return validity
    
def get_validity_2(ts, times):
    return np.logical_or(*[np.logical_and(t1<ts, ts<t2) for t1, t2 in times])
    
def get_validity_3(ts, times):
    validity = numpy.zeros(len(ts))
    for start, end in times:
        index_start = np.searchsorted(ts, start)
        index_end = np.searchsorted(ts, end)
        validity[index_start:index_end] = 1
    return validity
    

if __name__ == "__main__":
    n = 100000000
    ts = numpy.linspace(0, 1, n)
    
    times = numpy.array([[0.01, 0.1], [0.5, 0.8]])
    
    t0 = time.time()
    res_1 = get_validity_1(ts, times)
    t_1 = time.time() - t0
    
    t0 = time.time()
    res_2 = get_validity_2(ts, times)
    t_2 = time.time() - t0
    
    t0 = time.time()
    res_3 = get_validity_3(ts, times)
    t_3 = time.time() - t0
    
    print("t_1: " + str(t_1))
    print("t_2: " + str(t_2))
    print("t_3: " + str(t_3))
    
    assert (res_1 == res_2).all()
    assert (res_1 == res_3).all()

输出：

t_1: 0.4412200450897217
t_2: 0.3446168899536133
t_3: 0.14597129821777344