python - 为什么这个 python 类不能与 numba jitclass 一起使用？

Question

我在 numpy 的帮助下编写了以下代码，我想提高 numba 的性能。我不确定为什么它不起作用，因为我已经按照 numba 系统设置了所有变量。我正在尝试加快此代码的速度，因为我将来会处理大型数据集。

import numpy as np
import math
from numba import jitclass 
from numba import float64,int64

spec =[
       ('spacing',float64),
       ('n_iterations',int64),
       ('np_emptyhouses',float64[:,:]),
       ('np_agenthouses',float64[:,:]),
       ('similarity_threshhold',float64),
       ('n_changes',int64)
       ]

@jitclass(spec)
class geo_schelling_update:

    def __init__(self,n_iterations,spacing,np_agenthouses,np_emptyhouses,similarity_threshhold):
        self.spacing=spacing
        self.n_iterations=n_iterations
        self.np_emptyhouses=np_emptyhouses
        self.np_agenthouses=np_agenthouses
        self.similarity_threshhold=similarity_threshhold

    def distance_vectorize(self,pointA1, pointA2,agent):
        x_square=np.square(pointA1-agent[0])
        y_square=np.square(pointA2-agent[1])
        dist=np.sqrt(np.array(x_square,dtype=np.float32)+np.array(y_square,dtype=np.float32))
        return np.round(dist,4)

    def is_unsatisfied_vectorize(self,x,y):
        race = np.extract(np.logical_and(np.equal(self.np_agenthouses[:,0],x),np.equal(self.np_agenthouses[:,1],y)),self.np_agenthouses[:,2])[0]
        euclid_distance1=round(math.hypot(self.spacing,self.spacing),4)
        euclid_distance2=self.spacing
        total_agents=np.extract(np.logical_or(np.equal(np.round(np.hypot((self.np_agenthouses[:,0]-(x)),(self.np_agenthouses[:,1]-(y))),4),euclid_distance1),np.equal(np.round(np.hypot((self.np_agenthouses[:,0]-(x)),(self.np_agenthouses[:,1]-(y))),4),euclid_distance2)),self.np_agenthouses[:,2])
        if total_agents.size ==0:
            return False
        else:
            return total_agents[total_agents==race].size/total_agents.size<self.similarity_threshhold    

    def move_to_empty(self,x,y):
        race = np.extract(np.logical_and(np.equal(self.np_agenthouses[:,0],x),np.equal(self.np_agenthouses[:,1],y)),self.np_agenthouses[:,2])[0]
        x_new,y_new=self.np_emptyhouses[np.random.choice(self.np_emptyhouses.shape[0],1),:][0]
        self.np_agenthouses=self.np_agenthouses[~(np.logical_and(self.np_agenthouses[:,0]==x, self.np_agenthouses[:,1]==y))]
        self.np_agenthouses=np.vstack([self.np_agenthouses,[x_new,y_new,race]])
        self.np_emptyhouses=self.np_emptyhouses[~(np.logical_and(self.np_emptyhouses[:,0]==x_new, self.np_emptyhouses[:,1]==y_new))]
        self.np_emptyhouses=np.vstack([self.np_emptyhouses,[x,y]])

    def update_helper(self,agent):
        if self.is_unsatisfied_vectorize(agent[0],agent[1]):
            self.move_to_empty(agent[0],agent[1])
            return 1
        else:
            return 0

    def update(self):
        for i in np.arange(self.n_iterations):
            np_oldagenthouses=self.np_agenthouses.copy()
            n_changes=0
            for row in np_oldagenthouses:
                n=self.update_helper(row)
                n_changes+=n
            print(n_changes)
            print(i)
            if n_changes == 0:
                break



np_agenthouses=np.array([[-71.8,    41.4,   2.0],
                        [-71.6, 41.4,   2.0],
                        [-71.6, 41.6,   2.0],
                        [-71.4, 41.6,   1.0],
                        [-71.6, 41.8,   1.0],
                        [-71.4, 41.8,   2.0],
                        [-71.6, 42.0,   2.0],
                        [-71.4, 42.0,   1.0],
                        [-71.4, 41.4,   2.0],
                        [-71.2, 41.4,   1.0]])

np_emptyhouses=np.array([[-71.8,  41.3],[-71.8,  41.4],[-71.5,  41.5],
                [-71.5,  41.6],[-71.7,  41.8],[-71.7,  41.9],
                [-71.5,  41.9],[-71.2,  41.4],[-71.6,  41.7]])

spacing=0.1
similarity_threshhold=0.65
n_iterations=100
schelling= geo_schelling_update(n_iterations,
                         spacing,
                         np_agenthouses,
                         np_emptyhouses,similarity_threshhold)
schelling.update() 

这是我得到的错误:

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<function round_ at 0x000001909ED270D0>) with argument(s) of type(s): (array(float64, 1d, C), Literal[int](4))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: resolving callee type: Function(<function round_ at 0x000001909ED270D0>)
[2] During: typing of call at C:/Users/ksharma/Documents/geoschelling/test2.py (42)


File "test2.py", line 42:
    def is_unsatisfied_vectorize(self,x,y):
        <source elided>
        euclid_distance2=self.spacing
        total_agents=np.extract(np.logical_or(np.equal(np.round(np.hypot((self.np_agenthouses[:,0]-(x)),(self.np_agenthouses[:,1]-(y))),4),euclid_distance1),np.equal(np.round(np.hypot((self.np_agenthouses[:,0]-(x)),(self.np_agenthouses[:,1]-(y))),4),euclid_distance2)),self.np_agenthouses[:,2])
        ^

[1] During: resolving callee type: BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'is_unsatisfied_vectorize') for instance.jitclass.geo_schelling_update#190b49eec18<spacing:float64,n_iterations:int64,np_emptyhouses:array(float64, 2d, A),np_agenthouses:array(float64, 2d, A),similarity_threshhold:float64,n_changes:int64>)
[2] During: typing of call at C:/Users/ksharma/Documents/geoschelling/test2.py (57)

[3] During: resolving callee type: BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'is_unsatisfied_vectorize') for instance.jitclass.geo_schelling_update#190b49eec18<spacing:float64,n_iterations:int64,np_emptyhouses:array(float64, 2d, A),np_agenthouses:array(float64, 2d, A),similarity_threshhold:float64,n_changes:int64>)
[4] During: typing of call at C:/Users/ksharma/Documents/geoschelling/test2.py (57)


File "test2.py", line 57:
    def update_helper(self,agent):
        if self.is_unsatisfied_vectorize(agent[0],agent[1]):
        ^

[1] During: resolving callee type: BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'update_helper') for instance.jitclass.geo_schelling_update#190b49eec18<spacing:float64,n_iterations:int64,np_emptyhouses:array(float64, 2d, A),np_agenthouses:array(float64, 2d, A),similarity_threshhold:float64,n_changes:int64>)
[2] During: typing of call at C:/Users/ksharma/Documents/geoschelling/test2.py (68)

[3] During: resolving callee type: BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'update_helper') for instance.jitclass.geo_schelling_update#190b49eec18<spacing:float64,n_iterations:int64,np_emptyhouses:array(float64, 2d, A),np_agenthouses:array(float64, 2d, A),similarity_threshhold:float64,n_changes:int64>)
[4] During: typing of call at C:/Users/ksharma/Documents/geoschelling/test2.py (68)


File "test2.py", line 68:
    def update(self):
        <source elided>
            for row in np_oldagenthouses:
                n=self.update_helper(row)
                ^

[1] During: resolving callee type: BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'update') for instance.jitclass.geo_schelling_update#190b49eec18<spacing:float64,n_iterations:int64,np_emptyhouses:array(float64, 2d, A),np_agenthouses:array(float64, 2d, A),similarity_threshhold:float64,n_changes:int64>)
[2] During: typing of call at <string> (3)

[3] During: resolving callee type: BoundFunction((<class 'numba.types.misc.ClassInstanceType'>, 'update') for instance.jitclass.geo_schelling_update#190b49eec18<spacing:float64,n_iterations:int64,np_emptyhouses:array(float64, 2d, A),np_agenthouses:array(float64, 2d, A),similarity_threshhold:float64,n_changes:int64>)
[4] During: typing of call at <string> (3)

我也在 IDE 中运行这段代码。如果上述代码不适用于 numba，那么使该代码工作以获得相同结果的最佳方法是什么。

Answer 1

问题出在 np.round 上。从文档中并不完全清楚，但您可以通过查看 source 来了解，如果您在数组输入上使用该函数，则需要提供所有 3 个参数。所以以下不起作用:

nb.jit(nopython=True)
def func(x):
    return np.round(x)

但以下内容按预期工作:

nb.jit(nopython=True)
def func(x):
    out = np.empty_like(x)
    np.round(x, 0, out)
    return out

参见 docs for np.around完整的描述。我将在 numba 问题跟踪器上提出一个问题，因为从查看文档时这并不明显。