Изменение что-то из итерацию над Numpy массива векторизации

Question

Я пытаюсь ускорить кусок кода ниже по векторизации:

[rows,cols] = flow_direction_np.shape 
elevation_gain = np.zeros((rows,cols), np.float) 

for [i, j], flow in np.ndenumerate(flow_direction_np): 
    try: 
     if flow == 32: 
      elevation_gain[i - 1, j - 1] = elevation_gain[i - 1, j - 1] + sediment_transport_np[i, j] 
     elif flow == 64: 
      elevation_gain[i - 1, j] = elevation_gain[i - 1, j] + sediment_transport_np[i, j] 
     elif flow == 128: 
      elevation_gain[i - 1, j + 1] = elevation_gain[i - 1, j + 1] + sediment_transport_np[i, j] 
     elif flow == 16: 
      elevation_gain[i, j - 1] = elevation_gain[i, j - 1] + sediment_transport_np[i, j] 
     elif flow == 1: 
      elevation_gain[i, j + 1] = elevation_gain[i, j + 1] + sediment_transport_np[i, j] 
     elif flow == 2: 
      elevation_gain[i + 1, j + 1] = elevation_gain[i + 1, j + 1] + sediment_transport_np[i, j] 
     elif flow == 4: 
      elevation_gain[i + 1, j] = elevation_gain[i + 1, j] + sediment_transport_np[i, j] 
     elif flow == 8: 
      elevation_gain[i + 1, j - 1] = elevation_gain[i + 1, j - 1] + sediment_transport_np[i, j] 
    except IndexError: 
      elevation_gain[i, j] = 0 

Это как мой код выглядит на данный момент:

elevation_gain = np.zeros_like(sediment_transport_np) 
nrows, ncols = flow_direction_np.shape 
lookup = {32: (-1, -1), 
      16: (0, -1), 
      8: (+1, -1), 
      4: (+1, 0), 
      64: (-1, 0), 
      128:(-1, +1), 
      1: (0, +1), 
      2: (+1, +1)} 

# Initialize an array for the "shifted" mask 
shifted = np.zeros((nrows+2, ncols+2), dtype=bool) 

# Pad elevation gain with zeros 
tmp = np.zeros((nrows+2, ncols+2), elevation_gain.dtype) 
tmp[1:-1, 1:-1] = elevation_gain 
elevation_gain = tmp 

for value, (row, col) in lookup.iteritems(): 
    mask = flow_direction_np == value 

    # Reset the "shifted" mask 
    shifted.fill(False) 
    shifted[1:-1, 1:-1] = mask 

    # Shift the mask by the right amount for the given value 
    shifted = np.roll(shifted, row, 0) 
    shifted = np.roll(shifted, col, 1) 

    # Set the values in elevation change to the offset value in sed_trans 
    elevation_gain[shifted] = elevation_gain[shifted] + sediment_transport_np[mask] 

Проблема, с которой я столкнулась, это то, что они не дают мне тот же результат в конце любых предложений, в которых я ошибаюсь?

Answer 1

Вы можете значительно улучшить производительность с помощью np.where, чтобы получить индексы, где ваши условия происходят:

ind = np.where(flow_direction_np==32) 

вы увидите, что ind является кортеж с двумя элементами, первый индексы первой оси и вторая из второй оси вашего массива flow_direction_np.

Вы можете работать с этими индексами, чтобы применить изменения: i-1, j-1 и так далее ...

ind_32 = (ind[0]-1, ind[1]-1) 

Затем вы используете фантазии индексации обновить массивы:

elevation_gain[ ind_32 ] += sediment_transport_np[ ind ] 

---

EDIT: применение этой концепции на вашем корпусе даст примерно следующее:

lookup = {32: (-1, -1), 
      16: (0, -1), 
      8: (+1, -1), 
      4: (+1, 0), 
      64: (-1, 0), 
     128: (-1, +1), 
      1: (0, +1), 
      2: (+1, +1)} 

for num, shift in lookup.iteritems(): 
    ind = np.where(flow_direction_np==num) 
    ind_num = ind[0] + shift[0], ind[1] + shift[1] 
    elevation_gain[ ind_num] += sediment_transport_np[ ind ] 

Answer 2

Причина, по которой вы получаете разные результаты, связана с тем, как питон обрабатывает отрицательную индексацию.

---

Для других людей, читающих этот вопрос (и ответ) являются прослеживание здесь: Iterating through a numpy array and then indexing a value in another array

Во-первых, я извиняюсь, что «Векторизованных» код плотной, как она есть. В my earlier answer есть объяснение, поэтому я не буду повторять его здесь.

---

Ваш исходный код (в исходном вопросе) на самом деле тонко отличается от версии, которую вы разместили здесь.

В принципе, прежде чем был

for [i, j], flow in np.ndenumerate(flow_direction_np): 
    try: 
     if flow == 32: 
      ... 
     elif ... 
      ... 

и вы получаете сообщение об ошибке, когда индекс i+1 или j+1 был больше, чем размер сетки.

Просто делать:

for [i, j], flow in np.ndenumerate(flow_direction_np): 
    try: 
     if flow == 32: 
      ... 
     elif ... 
      ... 
    except IndexError: 
     elevation_change[i, j] = 0 

на самом деле неправильно, потому что она определяет различные граничные условия на разных сторонах сетки.

Во втором случае, когда j-1 или i-1 является отрицательным, то значение из противоположной стороны сетки будет возвращено.Однако, когда j+1 или i+1 больше, чем размер сетки, возвращается 0. (Таким образом, «различные граничные условия».)

В векторизованной версии коды, 0 возвращаются как когда индексы отрицательны, когда они за пределы сетки.

---

Как простой пример, обратите внимание на то, что происходит со следующим:

In [1]: x = [1, 2, 3] 

In [2]: x[0] 
Out[2]: 1 

In [3]: x[1] 
Out[3]: 2 

In [4]: x[2] 
Out[4]: 3 

In [5]: x[3] 
--------------------------------------------------------------------------- 
IndexError        Traceback (most recent call last) 
<ipython-input-5-ed224ad0520d> in <module>() 
----> 1 x[3] 

IndexError: list index out of range 

In [6]: x[-1] 
Out[6]: 3 

In [7]: x[-2] 
Out[7]: 2 

In [8]: x[-3] 
Out[8]: 1 

In [9]: x[-4] 
--------------------------------------------------------------------------- 
IndexError        Traceback (most recent call last) 
<ipython-input-9-f9c639f21256> in <module>() 
----> 1 x[-4] 

IndexError: list index out of range 

In [10]: 

Обратите внимание, что отрицательные показатели до размера последовательности являются действительными и вернуть «противоположный конец» последовательности. Итак, x[3] вызывает ошибку, а x[-1] просто возвращает другой конец.

Надеюсь, это немного яснее.