我有以下嵌套的for循环(为简单起见,使用随机数):
import numpy as np
lat_idx = np.random.randint(121, size = 4800)
lon_idx = np.random.randint(201, size = (4800,4800))
sum_cell = np.zeros((121,201))
data = np.random.rand(4800,4800)
for j in range(4800):
for i in range(4800):
if lat_idx[i] < 0 or lon_idx[i, j] < 0:
continue
sum_cell[lat_idx[i], lon_idx[i, j]] += data[i, j]
#print(sum_cell)
有没有一种将其编写为矩阵运算或带有一些“ numpy动作”的方法?此刻确实很慢。我的问题是lon_idx
都依赖于i
和j
。
答案 0 :(得分:1)
这是向量化方法:
import numpy as np
# Make input data
np.random.seed(0)
data = np.random.rand(4800, 4800)
# Add some negative values in indices
lat_idx = np.random.randint(-20, 121, size=4800)
lon_idx = np.random.randint(-50, 201, size=(4800, 4800))
# Output array
sum_cell = np.zeros((121, 201))
# Make mask for positive indices
lat_idx2 = lat_idx[:, np.newaxis]
m = (lat_idx2 >= 0) & (lon_idx >= 0)
# Get positive indices
lat_pos, lon_pos = np.broadcast_to(lat_idx2, m.shape)[m], lon_idx[m]
# Add values
np.add.at(sum_cell, (lat_pos, lon_pos), data[m])
# Check result with previous method
sum_cell2 = np.zeros((121, 201))
for j in range(4800):
for i in range(4800):
if lat_idx[i] < 0 or lon_idx[i, j] < 0:
continue
sum_cell2[lat_idx[i], lon_idx[i, j]] += data[i, j]
print(np.allclose(sum_cell, sum_cell2))
# True