使用带有 numba njit 功能的字典
Posted
技术标签:
【中文标题】使用带有 numba njit 功能的字典【英文标题】:Using Dictionaries with numba njit function 【发布时间】:2019-07-31 09:30:11 【问题描述】:当输入和返回是字典时,如何使用numba加速函数?
我熟悉将 numba 用于接受数字并返回数组的函数,如下所示:
@numba.jit('float64[:](int32,int32)',nopython=True)
def f(a, b):
# returns array 1d array
现在我有一个接受和返回字典的函数。我如何在这里申请 numba?
def collocation(aeolus_data,val_data):
...
return sample_aeolus, sample_valdata
【问题讨论】:
嗨@Annerl,欢迎来到SO。为了帮助我们回答您的问题,您能否详细说明您正在尝试访问 numba 的哪些优化和模式?您是否尝试过faqs 中的故障排除技巧?您是否能够将您的功能分解为更小的部分并分别对它们应用@numba
注释?
【参考方案1】:
现在 Numba 版本 43.0
中添加了对字典的支持。虽然它非常有限(不支持列表和设置为键/值)。但是,您可以阅读更新的文档 here for more info。
这是一个例子
import numpy as np
from numba import njit
from numba import types
from numba.typed import Dict
# First create a dictionary using Dict.empty()
# Specify the data types for both key and value pairs
# Dict with key as strings and values of type float array
dict_param1 = Dict.empty(
key_type=types.unicode_type,
value_type=types.float64[:],
)
# Dict with keys as string and values of type float
dict_param2 = Dict.empty(
key_type=types.unicode_type,
value_type=types.float64,
)
# Type-expressions are currently not supported inside jit functions.
float_array = types.float64[:]
@njit
def add_values(d_param1, d_param2):
# Make a result dictionary to store results
# Dict with keys as string and values of type float array
result_dict = Dict.empty(
key_type=types.unicode_type,
value_type=float_array,
)
for key in d_param1.keys():
result_dict[key] = d_param1[key] + d_param2[key]
return result_dict
dict_param1["hello"] = np.asarray([1.5, 2.5, 3.5], dtype='f8')
dict_param1["world"] = np.asarray([10.5, 20.5, 30.5], dtype='f8')
dict_param2["hello"] = 1.5
dict_param2["world"] = 10
final_dict = add_values(dict_param1, dict_param2)
print(final_dict)
# Output : hello: [3. 4. 5.], world: [20.5 30.5 40.5]
Link to Google colab notebook.
参考资料: - https://github.com/numba/numba/issues/3644 - https://numba.pydata.org/numba-doc/dev/reference/pysupported.html#dict
【讨论】:
以上是关于使用带有 numba njit 功能的字典的主要内容,如果未能解决你的问题,请参考以下文章
python 带有排名功能的python字典类,来自python cookbook 2 Ch5.14