Fastest way to mask a 4D Boolean numpy array

Viewed 202

Most efficient way to mask a 4D Boolean array with a 2D Boolean mask.

I tried two methods:

A. Reshape mask to 4D and mask

B. Reshape the inverted mask to 4D and product two matrixes

import numpy as np
import time

I = 150
J = 2000
K = I
S = 25

matrix_to_mask = np.random.choice(a=[True, False], size=(I, J, K, S))
mask_2d = np.random.choice(a=[True, False], size=(I, S))

t = time.time()
matrix_to_mask[np.tile(mask_2d[:, np.newaxis, np.newaxis, :], (1, J, K, 1))] = False
print("Mask: " + str(time.time() - t)) # 2.77 sec

t = time.time()
a = matrix_to_mask * np.tile(np.invert(mask_2d)[:, np.newaxis, np.newaxis, :], (1, J, K, 1))
print("Product: " + str(time.time() - t)) # 1 sec

Is there any other method which will speed up the mask?

Thanks

1 Answers

I would use np.broadcast_to for this. Broadcasting produces a 'virtual' copy of your lower dimension array. E.g. the array is repeated along the dimensions that you specify, but no new memory is allocated, it will only point to the original values. So this is more efficient because it uses less memory and it is faster because it doesn't spend time allocating.

t = time.time()
mask_4d = np.broadcast_to(mask_2d, matrix_to_mask.shape)

a = matrix_to_mask * mask_4d
# matrix_to_mask[mask_4d] = False
print("Broadcast: " + str(time.time() - t))

So:

import numpy as np
import time

I = 150
J = 1000
K = I
S = 25

matrix_to_mask = np.random.choice(a=[True, False], size=(I, J, K, S))
mask_2d = np.random.choice(a=[True, False], size=(I, S))

t = time.time()
matrix_to_mask[np.tile(mask_2d[:, np.newaxis, np.newaxis, :], (1, J, K, 1))] = False
print("Mask: " + str(time.time() - t)) # 2.77 sec

t = time.time()
a = matrix_to_mask * np.tile(mask_2d[:, np.newaxis, np.newaxis, :], (1, J, K, 1))
print("Product: " + str(time.time() - t)) # 1 sec

t = time.time()
mask_4d = np.broadcast_to(mask_2d, matrix_to_mask.shape)

a = matrix_to_mask * mask_4d
# matrix_to_mask[mask_4d] = False
print("Broadcast: " + str(time.time() - t))

Mask: 2.0803346633911133

Product: 0.8886129856109619

Broadcast: 0.4950692653656006

That said, I checked the values of the obtained mask and the one you generated with np.tile(mask_2d[:, np.newaxis, np.newaxis, :], (1, J, K, 1)) and they do not correspond, so something is off. Perhaps check that your method actually produces the mask that you intend, because normally np.broadcast_to does exactly what you are looking for, so I would think that that is correct...

Related