Создание массива индекса в numpy - исключение двойного цикла для цикла

Question

У меня есть код физической симуляции, написанный на python и использующий numpy/scipy. Профилирование кода показывает, что 38% времени процессора расходуется в одном двойном вложенном цикле - это кажется чрезмерным, поэтому я пытался его сократить.

Целью цикла является создание массива индексов, показывающих, какие элементы 1D-массива имеют элементы двумерного массива.

indices[i,j] = where(1D_array == 2D_array[i,j]) 

В качестве примера, если 1D_array = [7.2, 2.5, 3.9] и

2D_array = [[7.2, 2.5] 
      [3.9, 7.2]] 

Мы должны

indices = [[0, 1] 
      [2, 0]] 

я в настоящее время это реализуется как

for i in range(ni): 
    for j in range(nj): 
     out[i, j] = (1D_array - 2D_array[i, j]).argmin() 

argmin необходимо, поскольку я имею дело с числами с плавающей запятой, и поэтому равенство не обязательно точное. Я знаю, что каждое число в массиве 1D уникально и каждый элемент в 2D-массиве имеет совпадение, поэтому этот подход дает правильный результат.

Есть ли способ устранить двойной цикл?

Примечание:

мне нужен массив индексов, чтобы выполнить следующую операцию:

f = complex_function(1D_array) 
output = f[indices] 

Это быстрее, чем альтернатива, так как 2D массив имеет размер NxN по сравнению с 1xN для 1D-массив, а 2D-массив имеет много повторяющихся значений. Если кто-то может предложить другой способ получить тот же результат, не пройдя через массив индексов, это также может быть решением

Answer 1

Словарь метод, что некоторые другие предполагают, может работать, но для этого требуется, чтобы вы знали заранее, что каждый элемент вашего целевого массива (массив 2d) имеет точное соответствие в вашем массиве поиска (ваш 1-й массив). Даже когда это должно быть верно в принципе, вам все равно придется иметь дело с проблемами точности с плавающей запятой, например, попробуйте это .1 * 3 == .3.

Другой подход заключается в использовании функцииnumpy. searchsorted принимает отсортированный массив поиска 1d, и любой массив трассировок затем находит самые близкие элементы в массиве поиска для каждого элемента в целевом массиве. Я адаптировал этот answer для вашей ситуации, взгляните на него, чтобы описать, как работает функция find_closest.

import numpy as np 

def find_closest(A, target): 
    order = A.argsort() 
    A = A[order] 

    idx = A.searchsorted(target) 
    idx = np.clip(idx, 1, len(A)-1) 
    left = A[idx-1] 
    right = A[idx] 
    idx -= target - left < right - target 
    return order[idx] 

array1d = np.array([7.2, 2.5, 3.9]) 
array2d = np.array([[7.2, 2.5], 
        [3.9, 7.2]]) 

indices = find_closest(array1d, array2d) 
print(indices) 
# [[0 1] 
# [2 0]] 

Answer 2

В чистом Python вы можете сделать это с помощью словаря в O(N) времени, только время штрафа идет чтобы цикл Python Подключайтесь:

>>> arr1 = np.array([7.2, 2.5, 3.9]) 
>>> arr2 = np.array([[7.2, 2.5], [3.9, 7.2]]) 
>>> indices = dict(np.hstack((arr1[:, None], np.arange(3)[:, None]))) 
>>> np.fromiter((indices[item] for item in arr2.ravel()), dtype=arr2.dtype).reshape(arr2.shape) 
array([[ 0., 1.], 
     [ 2., 0.]]) 

Answer 3

чтобы избавиться от двух Python for петель, вы можете сделать все сравнения равенства «в один присест» путем добавления новых осей к массивам (что делает их broadcastable с каждым Другие).

Имейте в виду, что это создает новый массив, содержащий значения len(arr1)*len(arr2). Если это очень большое число, этот подход может быть неосуществимым в зависимости от ограничений вашей памяти. В противном случае, он должен быть достаточно быстро:

>>> (arr1[:,np.newaxis] == arr2[:,np.newaxis]).argmax(axis=1) 
array([[0, 1], 
     [2, 0]], dtype=int32) 

Если вам необходимо получить индекс значения соответствия ближайший в arr1 используйте:

np.abs(arr1[:,np.newaxis] - arr2[:,np.newaxis]).argmin(axis=1)