python - 使用多个条件从 numpy 数组中有效选择元素

Question

我正在寻找最快的方法来选择满足多个条件的 numpy 数组的元素。举个例子，假设我想从数组中选择 0.2 到 0.8 之间的所有元素。我通常会做这样的事情:

the_array = np.random.random(100000)
idx = (the_array > 0.2) * (the_array < 0.8)
selected_elements = the_array[idx]

但是，这会创建两个与 the_array 大小相同的附加数组(一个用于 the_array > 0.2，另一个用于 the_array < 0.8)。如果数组很大，这可能会消耗大量内存。有什么办法可以解决这个问题吗？所有内置的 numpy 函数(例如逻辑和)似乎都在幕后做同样的事情。

Answer 1

您可以为选择实现自定义 C 调用。最基本的方法是通过 ctypes 实现。

选择.c

int select(float lower, float upper, float* in, float* out, int n)
{
  int ii;
  int outcount = 0;
  float val;
  for (ii=0;ii<n;ii++)
    {
      val = in[ii];
      if ((val>lower) && (val<upper))
        {
          out[outcount] = val;
          outcount++;
        }
    }
  return outcount;
}

编译为:

gcc -lm -shared select.c -o lib.so

在Python方面:

select.py

import ctypes as C
from numpy.ctypeslib import as_ctypes
import numpy as np

# open the library in python
lib = C.CDLL("./lib.so")

# explicitly tell ctypes the argument and return types of the function
pfloat = C.POINTER(C.c_float)
lib.select.argtypes = [C.c_float,C.c_float,pfloat,pfloat,C.c_int]
lib.select.restype = C.c_int

size = 1000000

# create numpy arrays
np_input  = np.random.random(size).astype(np.float32)
np_output = np.empty(size).astype(np.float32)

# expose the array contents to ctypes
ctypes_input = as_ctypes(np_input)
ctypes_output = as_ctypes(np_output)

# call the function and get the number of selected points
outcount = lib.select(0.2,0.8,ctypes_input,ctypes_output,size)

# select those points 
selected = np_output[:outcount]

不要指望这样的普通实现会带来疯狂的加速，但在 C 端，您可以选择添加 OpenMP 编译指示来获得快速而肮脏的并行性，这可能会给您带来显着的提升。

正如评论中提到的，numexpr可能是一种更快更简洁的方法，只需几行即可完成所有这些工作。