Различия в выводах numba

Question

Я выполнил поиск основных ближайших соседей в моей учебной работе. Дело в том, что основная реализация numpy работает хорошо, но просто добавляет декоратор «@jit» (компиляция в Numba), выходы различаются (он дублирует некоторые соседи в конце по неизвестной причине ...)

Вот основной алгоритм:

import numpy as np 
from numba import jit 

@jit(nopython=True) 
def knn(p, points, k): 
    '''Find the k nearest neighbors (brute force) of the point p 
    in the list points (each row is a point)''' 

    n = p.size # Lenght of the points 
    M = points.shape[0] # Number of points 
    neighbors = np.zeros((k,n)) 
    distances = 1e6*np.ones(k) 

    for i in xrange(M): 
     d = 0 
     pt = points[i, :] # Point to compare 
     for r in xrange(n): # For each coordinate 
      aux = p[r] - pt[r] 
      d += aux * aux 
     if d < distances[k-1]: # We find a new neighbor 
      pos = k-1 
      while pos>0 and d<distances[pos-1]: # Find the position 
       pos -= 1 
      pt = points[i, :] 
      # Insert neighbor and distance: 
      neighbors[pos+1:, :] = neighbors[pos:-1, :] 
      neighbors[pos, :] = pt 
      distances[pos+1:] = distances[pos:-1] 
      distances[pos] = d 

    return neighbors, distances 

Для тестирования:

p = np.random.rand(10) 
points = np.random.rand(250, 10) 
k = 5 
neighbors = knn(p, points, k) 

БЕЗ @jit декоратор, один получает правильный ответ:

In [1]: distances 
Out[1]: array([ 0.3933974 , 0.44754336, 0.54548715, 0.55619749, 0.5657846 ]) 

Но компиляция Numba дает странные результаты:

Out[2]: distances 
Out[2]: array([ 0.3933974 , 0.44754336, 0.54548715, 0.54548715, 0.54548715]) 

Кто-нибудь может помочь? Я не понимаю, почему это происходит ...

Спасибо.

Answer 1

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

neighbors[pos+1:, :] = neighbors[pos:-1, :].copy() 
... 
distances[pos+1:] = distances[pos:-1].copy()