如何正确地将 numpy 逻辑函数传递给 Cython？

Question

我应该将哪些声明与逻辑函数/索引操作结合起来，以便 Cython 完成繁重的工作？

我有两个大小相等的 numpy 数组形式的大型栅格。第一个数组包含植被指数值，第二个数组包含字段 ID。目标是按字段平均植被指数值。这两个数组都有我想忽略的讨厌的 nodata 值（-9999）。

目前该函数的执行时间超过 60 秒，通常我不会太介意，但我可能会处理数百张图像。即使是 30 秒的改进也会很重要。所以我一直在探索 Cython 来帮助加快速度。我一直使用Cython numpy tutorial 作为指导。

Example data

test_cy.pyx 代码：

import numpy as np
cimport numpy as np
cimport cython
@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False)  # turn off negative index wrapping for entire function 

cpdef test():
  cdef np.ndarray[np.int16_t, ndim=2] ndvi_array = np.load("Z:cython_test/data/ndvi.npy")

  cdef np.ndarray[np.int16_t, ndim=2] field_array = np.load("Z:cython_test/data/field_array.npy")

  cdef np.ndarray[np.int16_t, ndim=1] unique_field = np.unique(field_array)
  unique_field = unique_field[unique_field != -9999]

  cdef int field_id
  cdef np.ndarray[np.int16_t, ndim=1] f_ndvi_values
  cdef double f_avg

  for field_id in unique_field :
      f_ndvi_values = ndvi_array[np.logical_and(field_array == field_id, ndvi_array != -9999)]
      f_avg = np.mean(f_ndvi_values)

Setup.py 代码：

try:
    from setuptools import setup
    from setuptools import Extension
except ImportError:
    from distutils.core import setup
    from distutils.extension import Extension

from Cython.Build import cythonize
import numpy

setup(ext_modules = cythonize('test_cy.pyx'),
      include_dirs=[numpy.get_include()])

经过一番研究和运行：

cython -a test_cy.pyx

看来索引操作ndvi_array[np.logical_and(field_array == field_id, ndvi_array != -9999)]是瓶颈，还在依赖Python。我怀疑我在这里遗漏了一些重要的声明。包括ndim 没有任何效果。

我对 numpy 也很陌生，所以我可能遗漏了一些明显的东西。

Answer 1

对我来说，您的问题看起来相当可矢量化，因此 Cython 可能不是最好的方法。 （当存在不可避免的细粒度循环时，Cython 会发光。）由于您的 dtype 是 int16，因此可能的标签范围有限，因此使用 np.bincount 应该是相当有效的。尝试类似的事情（假设您所有的有效值都 >= 0，如果不是这种情况，您必须转移 - 或者（更便宜）视图转换为 uint16（因为我们没有对标签进行任何算术运算那应该是安全的）-在使用bincount之前）：

mask = (ndvi_array != -9999) & (field_array != -9999)
nd = ndvi_array[mask]
fi = field_array[mask]
counts = np.bincount(fi, minlength=2**15)
sums = np.bincount(fi, nd, minlength=2**15)
valid = counts != 0
avgs = sums[valid] / counts[valid]