我正在尝试通过矢量化来加速下面的代码:
[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]
我遇到的麻烦是他们最后没有给我相同的结果,有什么建议我哪里出错了?