-2

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?

Sebastiano1991
  • 867
  • 1
  • 10
  • 26

1 Answers1

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)
Bob
  • 13,867
  • 1
  • 5
  • 27
  • Thank you for your answer - this is excatly what I was asking. However, when I compare your mask array with the one from the minimal example they are not equal in all elements - even when i replace the "<=" with "<". Am I overlooking something trivial? – Sebastiano1991 Mar 04 '22 at 14:40