将一些东西从迭代的numpy数组改为向量化 [英] Changing something from iterating over a numpy array to vectorization

问题描述

我正在尝试通过矢量化来加速下面的代码：

I am trying to speed up the piece of code below by vectorization:

[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

这就是我的代码目前的样子：

This is how my code looks at the moment:

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]

麻烦我有的是他们最后没有给我相同的结果我出错了吗？

The trouble I am having is they aren't giving me the same result at the end any suggestions where I am going wrong?

推荐答案

你可以使用 np.where 显着提高您的性能，以获取您的条件h的索引appen：

You can significantly improve your performance using np.where to get the indices where your conditions happen:

ind = np.where( flow_direction_np==32 )

你会看到 ind 是一个包含两个元素的元组，第一个是第一个轴的索引和 flow_direction_np 数组的第二个轴的第二个。

you will see that ind is a tuple with two elements, the first is the indices of the first axis and the second of the second axis of your flow_direction_np array.

您可以使用此索引来应用轮班： i-1 ， j-1 依此类推...

You can work out with this indices to apply the shifts: i-1, j-1 and so on...

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

然后你使用花式索引来更新数组：

Then you use fancy indexing to update the arrays:

elevation_gain[ ind_32 ] += sediment_transport_np[ ind ]

---

编辑：应用此概念对你的情况会给出这样的东西：

---

applying this concept to your case would give something like this:

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 ]

这篇关于将一些东西从迭代的numpy数组改为向量化的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！