字典与python相同的Cython

Posted 2021-04-13

我已经尝试了一些可用的解决方案，但是无法使用下面的Cython代码加快速度。

cython代码花费的时间与其python等价的时间相同。

calculate_sum.pyx

# key_id: it is a string
# values_dict: it is a dictionary with key as str and values as numpy.ndarray consisting of floats
# e.g.: print(values_dict['abc']) will give out numpy.ndarray([0.01, 1.01, 2.05]). values_dict has many such entries.

cpdef dict calculate_sum(str key_id, dict values_dict):
     cdef dict result_dict = {}
     cdef str check_id

     for check_id, values in values_dict.items():
           if check_id != key_id:
                result_dict[check_id] = sum(values)

     return result_dict

setup.py

from distutils.core import setup
from Cython.Build import cythonize

setup(ext_modules = cythonize('calculate_sum.pyx'))

main.py

import calculate_sum

# ...
# piece of code which computes values_dict
# ...

key_id = 'abc'
sum_value = calculate_sum.calculate_sum(key_id, values_dict)

感谢任何建议/帮助。

答案

总的来说，我认为Cython不会为这段代码提供帮助：它主要是在处理Python对象，并且这样做的速度往往很小。一些建议：

• 对于numpy数组，x.sum()比sum(x)快得多（对于我的PC上的快速测试用例，速度要快60倍）。如果您知道值将是Numpy数组，则直接执行此操作。
• 可以重写代码以使用字典理解。这可能会提高一点速度（主要是因为它可能能够预分配正确的大小）。
• 可能值得省略循环内的if语句，并计算每个键的总和（然后再删除key_id）。现在时间。

组合：

def calculate_sum(key_id, values_dict):
    result_dict = { key: values.sum() for key, values in values_dict.items() }
    del result_dict[key_id]
    return result_dict

请注意，由于我认为此处没有意义，因此我删除了所有Cython内容。我对基准测试没有太大兴趣，但是我怀疑对sum的更改将带来主要的不同。

---

您可能要看的另一种方法是使用Pandas（此处的代码未经测试...）。假设您使用values创建了values = pd.DataFrame.from_dict(values_dict) DataFrame，则：

def calculate_sum(key_id, values):
    return values.sum(axis=1).drop(key_id)

这将在所有Numpy数组的长度相同的情况下最好地工作（尽管我认为这不是绝对要求）。您的目标是普遍使用熊猫，而不是经常往返于dict。同样，Cython在这里也无济于事。