Numpy：在一个数组中创建元素数组，这些元素在另一个数组元素周围的范围内

Question

这是我的问题：

鉴于：
数组 A 具有这种形状：12000,3 和 数组 B 这个形状：150，

A 的第一列包含时间值，B 也包含时间值，采样率不同，因此它们不完全匹配。

问题： 创建一个形状为 150,3 的数组 C，其中包含数组 A 的行，其中第一列位于数组 B 中一个时间点周围的时间窗口内。时间窗口由“之前”时间和“之后”时间定义'

解决方案：

它适用于使用列表推导的一维列表，例如： C = [e for e in A if e > (B - before) 和 e

但是用数组尝试这个是行不通的。

我从尝试使用逻辑索引的整数数组开始，但这已经失败了。要么我得到全部 False，要么得到一条错误消息。

A = np.array([1,2,3,4,5,6])
B = np.array([1,3,5])

C = A[A in B]
C = A[A in B.any]
C = A[A == B]

np.select 也是如此。

理想情况下应该是这样的：

C = A[A > (B.any-before) and A < (B.any+after)]

非常感谢您的帮助！

Answer 1

我现在正在使用它。我必须遍历较短的数组，但它仍然比列表理解要快得多：

C_all = []
for i in range(len(B):
   #A and B must be np.arrays 
   bool_array = ((A > B - before) * (A < B + after), axis=1).astype('bool')
   C = B[bool_array]
   C_all.append(C)

Answer 2

这里的诀窍是广播A数组，让A和B的不同可以一起广播，然后减少原始数组，让我们重新获得B的维度。另外值得注意的是，将两个1相乘给出 1，加 1 和零也给出 1。此外，任何正整数的布尔值都是 True。

A = A.reshape(A.size,1)
B = np.ravel(B)
bool_array = np.sum((A > B - before) * (A < B + after), axis=1).astype('bool')
C = A[bool_array] 
A = np.ravel(A)

Answer 3

这看起来很像我最近遇到的一个问题（在执行分层重采样时，如果有人关心的话）。我最终使用的解决方案是这样的：

编辑：根据数组的大小，使用广播可能会更快（并且更 Pythonic），所以我已经包含了这个和一个小测试来比较两者。

import scipy
import timeit
import itertools

def AinBiterative(A,B):
    M = B.shape[0]
    N = A.shape[0]

    which = scipy.arange(M)*(N-1)/M
    # which will contain which elements from A are selected by B, this is a first guess
    shift = (B >= A[:-1][which,0]).astype(int) - (B <= A[1:][which,0]).astype(int)
    while not (shift == 0).all():
        which+= shift
        shift = (B >= A[:-1][which,0]).astype(int) - (B <= A[1:][which,0]).astype(int)

    return A[which,:]

def AinBbroadcast(A,B):
    comp = A[:,0][scipy.newaxis,:] < B[:,scipy.newaxis]
    return A[comp.sum(axis=-1)-1,:]

def wrapper(func, *args, **kwargs):
    def wrapped():
        return func(*args, **kwargs)
    return wrapped

Ms = scipy.power(10, range(1,5))
Ns = scipy.power(10, range(1,6))

for M, N in itertools.product(Ms, Ns):
    A = scipy.int_(N*scipy.randn(N,3))
    A[:,0] = scipy.arange(N, dtype=int)
    B = scipy.rand(M)*(N-1)
    B = scipy.sort(B)

    wrapped = wrapper(AinBiterative, A, B)
    iterative_time = timeit.timeit(wrapped, number=10)

    wrapped = wrapper(AinBbroadcast, A, B)
    broadcast_time = timeit.timeit(wrapped, number=10)

    print '{:<6d}, {:<6d}: {:1.2e} vs {:1.2e}'.format(M, N, iterative_time, broadcast_time)

我假设您正在寻找的所有 B 实际上都在 A 跨越的范围内（因此是 N-1）并且 A 和 B 是单调的。 A 不必是整数，时间也不必等距。

Answer 4

您可以使用numpy's broadcasting 来完成此操作。

>>> A = np.random.randn(12000, 3)
>>> B = np.random.randn(150)
>>> Ad = A[:, 0]

这里的Ad 是包含来自A 的时间的向量。 Numpy 有非常强大的广播功能，

Ad[:, None] - B[None, :]

将返回一个12000 x 150 数组，将每个A 与每个B 进行比较。如果您有足够的内存来保存该数组，这将非常有用且快速。 12000 x 150 很小，所以很好，但要注意尺寸是否增长很多。

因此，使用上面的广播思想，您可以将您的规则广播到整个产品空间：

>>> before = 0.001
>>> after = 0.001
>>> rule1 = (Ad[:, None] > B[None, :] - before) # 12000 x 150
>>> rule2 = (Ad[:, None] < B[None, :] + after)  # 12000 x 150
>>> mask = rule1 & rule2                        # 12000 x 150
>>> valid = mask.any(axis=1)

请注意，rule1 和 rule2 是密集比较（在 numpy 中速度非常快），它们在 Ad 的每个元素和 B 的每个元素之间比较上述规则。

此处有效检查A 的每个条目是否在B 的至少 1 个窗口内。请注意，您可以在不创建临时 rule1 和 rule2 的情况下进行 1 行掩码，只是想让代码更具可读性。

>>> C = A[valid]
>>> C.shape, A.shape
(943, 3), (12000, 3)

我们从 12000 个元素中过滤了 943 个元素！

---

其中的一个函数：

def window_filter(A, B, before=0.001, after=0.001):
    """Returns indexes from A that are within a window of B"""
    rule1 = (A[:, None] > B[None, :] - before)
    rule2 = (A[:, None] < B[None, :] + after)
    mask = rule1 & rule2
    valid = mask.any(axis=1)
    return valid

>>> C = A[window_filter(A[:, 0], B)]
>>> C.shape, A.shape
(943, 3), (12000, 3)

而且速度也很快（4ms）：

>>> %timeit window_filter(A[:, 0], B)
100 loops, best of 3: 4.53 ms per loop

---

此外，如果您想进一步指定更多约束，例如elements from A that are inside **only 1** window of B，您可以将valid 替换为

valid = (mask.sum(axis=1) == 1)