ускорение класса с помощью массива с использованием Cython

Question

У меня есть следующие коды:

class _Particles:
    def __init__(self, num_particle, dim, fun, lower_bound, upper_bound):
        self.lower_bound = lower_bound   # np.array of shape (dim,)
        self.upper_bound = upper_bound   # np.array of shape (dim,)
        self.num_particle = num_particle   # a scalar
        self.dim = dim   # dimension, a scalar
        self.fun = fun   # a function

        self.pos = np.empty((num_particle,dim))
        self.val = np.empty(num_particle)
        self.randomize()


    def randomize(self):
        self.pos = np.random.rand(self.num_particle, self.dim)*(self.upper_bound\
                -self.lower_bound)+self.lower_bound
        self.val = self.fun(np.transpose(self.pos))
        self.best_idx = np.argmin(self.val)
        self.best_val = self.val[self.best_idx]
        self.best_pos = self.pos[self.best_idx]


    def move(self, displacement, idx='all', check_bound=True):
        if idx is 'all':
            self.pos += displacement
        elif isinstance(idx,(tuple,list,np.ndarray)):
            self.pos[idx] += displacement
        else:
            raise TypeError('Check the type of idx!',type(idx))

        self.pos = np.maximum(self.pos, self.lower_bound[np.newaxis,:])
        self.pos = np.minimum(self.pos, self.upper_bound[np.newaxis,:])
        self.val = self.fun(np.transpose(self.pos))
        self.best_idx = np.argmin(self.val)
        self.best_val = self.val[self.best_idx]
        self.best_pos = self.pos[self.best_idx]

Я хочу посмотреть, смогу ли я ускорить приведенный выше код, и я думаю об использовании cython, но я не уверен, возможно ли это, поскольку он в основном использует массив numpy, а большинство выполнений выполняется с векторизацией. Я пытаюсь что-то вроде:

# the .pyx file that will be compiled
cdef class _Particles(object):
    cdef int num_particle
    cdef int dim
    cdef fun
    cdef np.ndarray lower_bound
    cdef np.ndarray upper_bound
    cdef np.ndarray pos
    cdef np.ndarray val
    cdef int best_idx
    cdef double best_val
    cdef np.ndarray[np.float64_t, ndim=1] best_pos

    def __init__(self, int num_particle, int dim, fun,
                 np.ndarray lower_bound, np.ndarray upper_bound):
        self.num_particle = num_particle
        self.dim = dim
        self.fun = fun
        self.lower_bound = lower_bound
        self.upper_bound = upper_bound

        self.pos = np.empty((num_particle,dim))
        self.val = np.empty(num_particle)
        self.randomize()

    def randomize(self):
        self.pos = npr.rand(self.num_particle,self.dim)*(self.upper_bound\
                -self.lower_bound)+self.lower_bound

        self.val = self.fun(np.transpose(self.pos))
        self.best_idx = np.argmin(self.val)
        self.best_val = self.val[self.best_idx]
        self.best_pos = self.pos[self.best_idx]

Это быстрее, но только немного, что ожидаемо, так как это все еще в основном код на python. Так есть ли способы ускорить приведенный выше код с помощью Cython (или указать мне на некоторые другие полностью методы)? В частности, как ускорить такие коды, как self.fun(self.pos), np.argmin(self.val)?

Спасибо.

Answer 1

На самом деле, я боюсь, что в приведенном выше коде оптимизировать особо нечего. Чтобы сделать argmin быстрее, я бы посоветовал вам получить (или иным образом скомпилировать) NumPy с многопоточной поддержкой (или вы можете заново реализовать несколько многопоточных аргументов argmin).

Что касается Cython, то вы получаете реальную сделку, когда начинаете использовать типы C, но я бы не увидел в этом большого улучшения в опубликованном вами коде. Это в основном клеевой код, здесь не нужно ничего перемалывать.

Я бы ожидал, что в функции произойдет сокращение числа fun, и это, вероятно, единственное место, где фактическая ручная оптимизация может иметь значение, если ее не так просто векторизовать (читай: for или другой ручной цикл). Затем я бы начал с numba, что намного проще для вашего кода, если оно работает. Если этого не произойдет, вероятно, стоит начать изучать Cython.