用迭代条件向量化循环以生成掩码

Question

我正在尝试向量化以下问题：

time_lag = np.linspace(0, 10, 50)
time = np.arange(100)
dt = np.abs(time[:,None]-time[None, :]) ## calculate matrix of differences
mask = []
for num in range(len(time_lag)-1):
  m0 = (time_lag[num] < dt) & (dt < time_lag[num+1])
  ## do something with mask 
  mask.append(m0)

mask = np.array(mask)

理想情况下，我会获得一个掩码数组，其形状 (49,100,100) 对应于 dt.shape 和 time_lag.size-1。我正在寻找 for 循环的真正矢量化（即不使用 np.vectorize）。 np.where 这可能吗？

Answer 1

没有注意代码的作用，但这给出了等效的结果，就好像您在代码末尾询问 do mask = np.array(mask) 一样。

time_lag = np.linspace(0, 10, 50)
time = np.arange(100)
dt = np.abs(time[:,None]-time[None, :]) ## calculate matrix of differences
# vectorized
mask_p = (time_lag[:-1, None, None] < dt[None]) & \
       (dt[None] < time_lag[1:, None, None])
# iterative
mask = []
for num in range(len(time_lag)-1):
  m0 = (time_lag[num] < dt) & (dt < time_lag[num+1])
  ## do something with mask 
  mask.append(m0)
mask = np.asarray(mask)
assert np.all(mask == mask_p)

用迭代条件向量化循环以生成掩码

vectorize a loop with iterative conditions to generate a mask

python

numpy

array-broadcasting