過濾（減少）NumPy 數組

Question

假設我有一個 NumPy 數組arr ，我想根據（可廣播的）function 的真值進行逐元素過濾（減少），例如，我只想獲得低於某個閾值k的值：

def cond(x):
    return x < k

有幾種方法，例如：

1. 使用生成器： np.fromiter((x for x in arr if cond(x)), dtype=arr.dtype) （這是使用列表理解的 memory 高效版本： np.array([x for x in arr if cond(x)])因為np.fromiter()會直接產生一個 NumPy 數組，而不需要分配一個中間 Python list ）
2. 使用 boolean 掩碼： arr[cond(arr)]
3. 使用 integer 索引： arr[np.nonzero(cond(arr))] （或等效地使用np.where()因為它默認為np.nonzero()只有一個條件）
4. 使用顯式循環：
   • 單遍和最終復制/調整大小
   • 兩遍：一是確定結果的大小，一是實際執行計算

（使用Cython或Numba加速的最后兩種方法）

哪個最快？ memory效率怎么樣？

---

（編輯：根據@ShadowRanger 評論直接使用np.nonzero()而不是np.where() ）

Answer 1

概括

使用基於循環的方法進行單遍和復制，並通過 Numba 加速，在速度、memory 效率和靈活性方面提供了最佳的整體權衡。 如果條件 function 的執行速度足夠快，則兩次通過 ( filter2_nb() ) 可能會更快，而無論如何它們的 memory 效率更高。 此外，對於足夠大的輸入，調整大小而不是復制（ filter_resize_xnb() ）會導致更快的執行。

如果提前知道數據類型（和條件函數）並且可以編譯，Cython 加速似乎更快。 類似的條件硬編碼很可能會導致與 Numba 加速相當的加速。

當談到僅基於 NumPy 的方法時，boolean 掩碼或 integer 索引具有相當的速度，並且哪個更快出來很大程度上取決於過濾因子，即通過過濾條件的值的部分。

np.fromiter()方法要慢得多（它會在圖中超出圖表），但不會產生大型臨時對象。

請注意，以下測試旨在提供對不同方法的一些見解，應謹慎使用。 最相關的假設是條件是可廣播的，並且它最終會非常快速地計算。

---

定義

1. 使用生成器：

def filter_fromiter(arr, cond):
    return np.fromiter((x for x in arr if cond(x)), dtype=arr.dtype)

2. 使用 boolean 屏蔽：

def filter_mask(arr, cond):
    return arr[cond(arr)]

3. 使用 integer 索引：

def filter_idx(arr, cond):
    return arr[np.nonzero(cond(arr))]

4a。 使用顯式循環，單遍和最終復制/調整大小

• Cython 加速復制（預編譯條件）

%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True


import numpy as np


cdef long NUM = 1048576
cdef long MAX_VAL = 1048576
cdef long K = 1048576 // 2


cdef int cond_cy(long x, long k=K):
    return x < k


cdef size_t _filter_cy(long[:] arr, long[:] result, size_t size):
    cdef size_t j = 0
    for i in range(size):
        if cond_cy(arr[i]):
            result[j] = arr[i]
            j += 1
    return j


def filter_cy(arr):
    result = np.empty_like(arr)
    new_size = _filter_cy(arr, result, arr.size)
    return result[:new_size].copy()

• Cython 加速調整大小（預編譯條件）

def filter_resize_cy(arr):
    result = np.empty_like(arr)
    new_size = _filter_cy(arr, result, arr.size)
    result.resize(new_size)
    return result

• Numba 加速復制

import numba as nb


@nb.njit
def cond_nb(x, k=K):
    return x < k


@nb.njit
def filter_nb(arr, cond_nb):
    result = np.empty_like(arr)
    j = 0
    for i in range(arr.size):
        if cond_nb(arr[i]):
            result[j] = arr[i]
            j += 1
    return result[:j].copy()

• Numba 加速調整大小

@nb.njit
def _filter_out_nb(arr, out, cond_nb):
    j = 0
    for i in range(arr.size):
        if cond_nb(arr[i]):
            out[j] = arr[i]
            j += 1
    return j


def filter_resize_xnb(arr, cond_nb):
    result = np.empty_like(arr)
    j = _filter_out_nb(arr, result, cond_nb)
    result.resize(j, refcheck=False)  # unsupported in NoPython mode
    return result

• Numba 使用生成器和np.fromiter()加速

@nb.njit
def filter_gen_nb(arr, cond_nb):
    for i in range(arr.size):
        if cond_nb(arr[i]):
            yield arr[i]


def filter_gen_xnb(arr, cond_nb):
    return np.fromiter(filter_gen_nb(arr, cond_nb), dtype=arr.dtype)

4b。 使用帶有兩遍的顯式循環：一是確定結果的大小，一是實際執行計算

• Cython 加速（預編譯條件）

%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True


cdef size_t _filtered_size_cy(long[:] arr, size_t size):
    cdef size_t j = 0
    for i in range(size):
        if cond_cy(arr[i]):
            j += 1
    return j


def filter2_cy(arr):
    cdef size_t new_size = _filtered_size_cy(arr, arr.size)
    result = np.empty(new_size, dtype=arr.dtype)
    new_size = _filter_cy(arr, result, arr.size)
    return result

• Numba 加速

@nb.njit
def filter2_nb(arr, cond_nb):
    j = 0
    for i in range(arr.size):
        if cond_nb(arr[i]):
            j += 1
    result = np.empty(j, dtype=arr.dtype)
    j = 0
    for i in range(arr.size):
        if cond_nb(arr[i]):
            result[j] = arr[i]
            j += 1
    return result

---

時序基准

（基於生成器的filter_fromiter()方法比其他方法慢得多 - 大約 2 個數量級。從列表理解中可以預期類似（並且可能稍差）的性能。這對於任何顯式循環都是正確的非加速代碼。）

時間將取決於輸入數組大小和過濾項目的百分比。

作為輸入大小的 function

第一張圖將時序描述為輸入大小的 function（對於 ~50% 的過濾因子——即 50% 的元素出現在結果中）：

通常，使用一種加速形式的顯式循環會導致最快的執行，但會根據輸入大小略有變化。

在 NumPy 中，integer 索引方法基本上與 boolean 掩碼相當。

使用np.fromiter() （無預分配）的好處可以通過編寫 Numba 加速生成器來獲得，這將比其他方法慢（在一個數量級內），但比純 Python 快得多循環。

作為function的灌裝

第二張圖將時序描述為通過過濾器的項目的 function（對於大約 100 萬個元素的固定輸入大小）：

第一個觀察結果是，所有方法在接近 50% 填充時最慢，而填充更少或更多時，它們更快，並且最快達到無填充（過濾出值的最高百分比，通過值的最低百分比，如中所示圖表的 x 軸）。

同樣，具有某種加速度平均值的顯式循環會導致最快的執行。

在 NumPy 中，integer 索引和 boolean 掩蔽方法再次基本相同。

（ 此處提供完整代碼）

---

Memory 注意事項

基於生成器的filter_fromiter()方法只需要最少的臨時存儲，與輸入的大小無關。 在內存方面，這是最有效的方法。 使用 Numba 加速生成器可以有效地加速這種方法。

類似 memory 的效率是 Cython / Numba 兩遍方法，因為 output 的大小是在第一遍期間確定的。 這里需要注意的是，為了使這些方法快速，計算條件必須快速。

在 memory 端，Cython 和 Numba 的單通道解決方案需要輸入大小的臨時數組。 因此，與兩遍或基於生成器的一遍相比，這些並不是非常節省內存。

然而，與掩蔽相比，它們具有相似的漸近臨時 memory 足跡，但常數項通常大於掩蔽。

The boolean masking solution requires a temporary array of the size of the input but of type bool , which in NumPy is 1 byte, so this is ~8 times smaller than the default size of a NumPy array on a typical 64-bit system.

integer 索引解決方案與第一步中的 boolean 掩碼切片具有相同的要求（在np.nonzero()調用中），在第二步中將其轉換為一系列int s（通常是 64 位系統上的int64 ） （ np.nonzero()的 output ）。 因此，第二步具有可變的 memory 要求，具體取決於過濾元素的數量。

---

評論

• boolean 掩碼和 integer 索引都需要某種形式的調節，能夠產生 boolean 掩碼（或者，索引列表）； 在上述實現中，條件是可廣播的
• 在指定不同的過濾條件時，生成器和 Numba 加速方法也是最靈活的
• Numba 加速方法要求條件與 Numba 兼容才能在 NoPython 模式下訪問 Numba 加速
• Cython 解決方案需要為它指定快速的數據類型，或者為多種類型的調度付出額外的努力，以及額外的努力（此處未探討）以獲得與其他方法相同的靈活性
• 對於 Numba 和 Cython，過濾條件可以是硬編碼的，從而導致邊際但明顯的速度差異
• 單通道解決方案需要額外的代碼來處理未使用的（但最初分配的）memory。
• 由於高級索引，NumPy 方法不返回輸入視圖，而是返回副本：

arr = np.arange(100)
k = 50
print('`arr[arr > k]` is a copy: ', arr[arr > k].base is None)
# `arr[arr > k]` is a copy:  True
print('`arr[np.where(arr > k)]` is a copy: ', arr[np.where(arr > k)].base is None)
# `arr[np.where(arr > k)]` is a copy:  True
print('`arr[:k]` is a copy: ', arr[:k].base is None)
# `arr[:k]` is a copy:  False

---

（編輯：基於@ShadowRanger、@PaulPanzer、@max9111 和@DavidW 評論的各種改進。）