vectorize a loop with iterative conditions to generate a mask

Question

I am trying to vectorize the following problem:

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)

Ideally I would obtain an mask array with shape (49,100,100) corresponding to the dt.shape and time_lag.size-1. I am looking for a true vectorization of the for loop (i.e. not with np.vectorize). Is this possible with np.where?

Answer 1

Not paying attention to what the code does, but this gives the equivalent result, as if you ask do mask = np.array(mask) at the end of your code.

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)