This question may have the same flavor as the one answered at Why is vectorized numpy code slower than for loops?
I want to implement a sigma-clipping method to reject some outliers in a series of images. Here's what I got, with first some code to generate random data cube (I normally have real images of 4k x 4k typically) with outliers.
Code for generating some random data:
import numpy as np
def make_outliers(mean=1000, err=10, size=50):
# Create 1 stack of pixel with 10% outliers.
std = np.sqrt(mean)
n_outliers = np.random.randint(0, int(size/10))
size -= n_outliers
random_err = err * np.random.normal(loc=mean, scale=std, size=size) * np.random.choice((-1, 1), size)
data = mean + random_err
outlier_int = 50 * mean
outlier_errs = outlier_int * np.random.rand(n_outliers) * np.random.choice((-1, 1), n_outliers)
data = np.concatenate((data, outlier_errs))
np.random.shuffle(data)
return data
def cube_outliers(size=(128, 128, 50)):
# Make an image series with outliers with respect to the 3rd dimension
cube = np.empty(size)
for r in range(size[0]):
for c in range(size[1]):
cube[r,c,:] = make_outliers(size=size[-1])
return cube
Sigma clipping function that returns the rejected pixels as a boolean mask:
def sigma_clip(datacube):
rejMask = np.zeros(datacube.shape, dtype=np.bool)
n = 1
n_outliers = []
while n > 0:
rejMask0 = rejMask.copy()
med = np.nanmedian(datacube, axis=-1)
sigma = np.nanstd(datacube, axis=-1)
rejMask = (np.abs(datacube - med[...,np.newaxis]) > 5*sigma[...,np.newaxis])
n = rejMask.sum()
n_outliers.append(n)
# print(n)
rejMask = rejMask0 | rejMask
datacube[rejMask] = np.nan
return rejMask, n_outliers
# Generate data cube for testing, emulating a series of 100 images.
images = cube_outliers(size=(1024, 1024, 100))
Timing with vectorized version:
%timeit rej_mask1, n_outs = sigma_clip(images.copy())
print(n_outs, sum(n_outs))
24.2 s ± 114 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
[22213, 36, 0] 22249
Timing with for-loop over all rows (in fact, list comprehension)
%%timeit
images2 = images.copy()
nout_rows = []
rej_mask2 = np.empty(images.shape)
for r in range(images.shape[0]):
rej_mask2[r, ...], nout_r = sigma_clip(images2[r,...])
nout_rows.append(nout_r)
10.6 s ± 60.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
nloops = max([len(nouts) for nouts in nout_rows])
nout_sum = sum([sum(nouts) for nouts in nout_rows])
print(nout_max, nout_sum)
np.array_equal(rej_mask1, rej_mask2)
(3, 22249)
True
So here the vectorized code is more than twice slower than when I use one for loop to run the sigma-clipping filter row by row. The example uses image sizes that are at least half of those in the real scenario in each dimension. Is that something I did in the vectorized code that's just not Numpy-optimized and there's a better way with Numpy? Or is the issue similar to the one linked above which mentions cache/L1 that are CPU-related aspect. In the latter case, I would have expected Numpy to optimize that for us, as hardware consideration is more something I would explicitely play with various compiler options and not something I'd expect Numpy to make me think of.
Without going into Cython or parallelization with Numba, is there a faster way to do this while sticking to Python/Numpy? I have ~24 GB of allocatable memory with the machine used in this timing.
[Update] Before this version, I tried with Masked Array but it was noticeably slower than directly flagging bad values with np.nan and using nanmedian() and nanstd().
[Update 2] Based on the comments, to make it clear that between the two tests there is no difference in the number of rejected pixels and overall number of while loop with respect to the 3rd dimension, i've added a counter (n_outliers) in the code above, to keep track of the number of iteration in the while-loop for each pixel stack that gets clipped. It simply helps assert that, provided the random data generation is not re-run between two tests (otherwise, you just get different data with new sets of outliers of different sizes), the code behaves identically regarding what pixels get rejected. In addition, I check that np.array_equal(rej_mask1, rej_mask2) is indeed True.
[Update 3] Using the output added in Update 2, we can indeed nail down in the 2nd test that the amount of processing is much less than in the 1st test:
nloops_per_row = np.array([len(nouts) for nouts in nout_rows])
from collections import Counter
Counter(nloops_per_row)
Counter({2: 988, 3: 36})
So, when in test 1 (vectorized version), the entire array is sent to the while loop 3 times as long as there is at least 1 outlier in the 2nd pass, in the 2nd test (for loop), only those rows that still have outliers at the end of the 2nd pass get a 3rd pass, which is, here, only 36 of them, against the 1024 rows in test 1.
[Update 4] - Solution 1
This is probably not elegant; here's a vectorized version that is much faster while sticking to Numpy. (i) It reshapes the cube in 2D, (ii) it separates the 1st pass from the other passes of the clipping so that we do not need nanmedian() and nanstd() on the 1st pass, as they are slower than median() and std(), and (iii) on the other passes, it only loads and processes the rows that have outliers instead of dealing with all the rows regardless of their content.
def sigma_clip2(datacube):
# Make a 1st pass while there's no need to consider NaN-flagged arrays
# as median() and std() are faster than nanmedian() and nanstd().
sz = datacube.shape
flatc = datacube.reshape([sz[0]*sz[1], sz[2]])
rej_mask = np.zeros(flatc.shape, dtype=np.bool)
m = np.median(flatc, axis=1)
sigma = np.std(flatc, axis=1)
mask = np.abs(flatc - m[:,np.newaxis]) > 5*sigma[:,np.newaxis]
n = mask.sum()
if n == 0:
return rej_mask
# Prepare new passes only on the rows that have outliers.
clip_rows = np.where(np.any(mask, axis=1))[0]
mask = mask[clip_rows, :]
rej_mask[clip_rows, :] = mask
while n > 0:
# Work only on the rows that have outliers. Flag them as NaN
flatc = flatc[clip_rows, :]
flatc[mask] = np.nan
m = np.nanmedian(flatc, axis=1)
sigma = np.nanstd(flatc, axis=1)
mask = np.abs(flatc - m[:,np.newaxis]) > 5*sigma[:,np.newaxis]
clip_rows0 = clip_rows.copy()
clip_rows = np.where(np.any(mask, axis=1))[0]
mask = mask[clip_rows, :]
n = mask.sum()
rej_mask[clip_rows0[clip_rows]] = rej_mask[clip_rows0[clip_rows]] | mask
rej_mask = rej_mask.reshape(sz)
return rej_mask
With the same random-generated data cube, timing for the 3 different versions give:
Slow vectorize (version 1): 24 s ± 120 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
for-loop (version 2): 12 s ± 51.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Faster Vectorize (version 3): 3.82 s ± 46.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In the latter case, it is essentially the 1st pass that is the longest, at about 3s. The 2nd pass falls down to the order of ~200ms as there are only a few dozens rows with outliers. Maybe the 1st pass could benefit from yet another optimization.
[Update 5] - Bug in Solution 1, and its bug fix
The indices in the output rejection mask rej_mask were not properly mapped when doing: clip_rows0 = clip_rows.copy(). This works if we have only have 2 passes but this is wrong if the while loop iterates a 2nd time (the random data generator was making 1 or 2 iterations so this was a misleading coincidence that the output of all solutions were equal). Changing to clip_rows0 = clip_rows0[clip_rows] gives the expected mapping. This has no effect on the above performance comparison. Solution 1 [fixed] below.
def sigma_clip2(datacube):
# Make a 1st pass while there's no need to consider NaN-flagged arrays
# as median() and std() are faster than nanmedian() and nanstd().
sz = datacube.shape
flatc = datacube.reshape([sz[0]*sz[1], sz[2]])
rej_mask = np.zeros(flatc.shape, dtype=np.bool)
m = np.median(flatc, axis=1)
sigma = np.std(flatc, axis=1)
mask = np.abs(flatc - m[:,np.newaxis]) > 5*sigma[:,np.newaxis]
n = mask.sum()
if n == 0:
return rej_mask
# Prepare new passes only on the rows that have outliers.
clip_rows = np.where(np.any(mask, axis=1))[0]
clip_rows0 = clip_rows.copy()
mask = mask[clip_rows, :]
rej_mask[clip_rows, :] = mask
while n > 0:
# Work only on the rows that have outliers. Flag them as NaN
flatc = flatc[clip_rows, :]
flatc[mask] = np.nan
m = np.nanmedian(flatc, axis=1)
sigma = np.nanstd(flatc, axis=1)
mask = np.abs(flatc - m[:,np.newaxis]) > 5*sigma[:,np.newaxis]
clip_rows = np.where(np.any(mask, axis=1))[0]
mask = mask[clip_rows, :]
n = mask.sum()
# Map back to original indices of the original array
clip_rows0 = clip_rows0[clip_rows]
rej_mask[clip_rows0] = rej_mask[clip_rows0] | mask
rej_mask = rej_mask.reshape(sz)
return rej_mask