Fill blocks at random places on each 2D slice of a 3D array

Question

I have 3D numpy array, for example, like this:

array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23],
        [24, 25, 26, 27],
        [28, 29, 30, 31]]])

Is there a way to index it in such a way that I select, for example, top right corner of 2x2 elements in the first plane, and a center 2x2 elements subarray from the second plane? So that I could then zero out the elements 2,3,6,7,21,22,25,26:

array([[[ 0,  1,  0,  0],
        [ 4,  5,  0,  0],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[16, 17, 18, 19],
        [20,  0,  0, 23],
        [24,  0,  0, 27],
        [28, 29, 30, 31]]])

I have a batch of images, and I need to zero out a small window of fixed size, but at different (random) locations for each image in the batch. The first dimension is number of images.

Something like this: a[:, x: x+2, y: y+2] = 0

where x and y are vectors which have different values for each first dimension of a.

Answer 1

Approach #1 : Here'e one approach that's mostly based on linear-indexing -

def random_block_fill_lidx(a, N, fillval=0):
    # a is input array
    # N is blocksize

    # Store shape info
    m,n,r = a.shape

    # Get all possible starting linear indices for each 2D slice
    possible_start_lidx = (np.arange(n-N+1)[:,None]*r + range(r-N+1)).ravel()

    # Get random start indices from all possible ones for all 2D slices
    start_lidx = np.random.choice(possible_start_lidx, m)

    # Get linear indices for the block of (N,N)
    offset_arr = (a.shape[-1]*np.arange(N)[:,None] + range(N)).ravel()

    # Add in those random start indices with the offset array
    idx = start_lidx[:,None] + offset_arr

    # On a 2D view of the input array, use advance-indexing to set fillval.
    a.reshape(m,-1)[np.arange(m)[:,None], idx] = fillval
    return a

Approach #2 : Here's another and possibly more efficient one (for large 2D slices) using advanced-indexing -

def random_block_fill_adv(a, N, fillval=0):
    # a is input array
    # N is blocksize

    # Store shape info
    m,n,r = a.shape

    # Generate random start indices for second and third axes keeping proper
    # distance from the boundaries for the block to be accomodated within.
    idx0 = np.random.randint(0,n-N+1,m)
    idx1 = np.random.randint(0,r-N+1,m)

    # Setup indices for advanced-indexing.

    # First axis indices would be simply the range array to select one per elem.
    # We need to extend this to 3D so that the latter dim indices could be aligned.
    dim0 = np.arange(m)[:,None,None]

    # Second axis indices would idx0 with broadcasted additon of blocksized 
    # range array to cover all block indices along this axis. Repeat for third.
    dim1 = idx0[:,None,None] + np.arange(N)[:,None]
    dim2 = idx1[:,None,None] + range(N)
    a[dim0, dim1, dim2] = fillval
    return a

Approach #3 : With the old-trusty loop -

def random_block_fill_loopy(a, N, fillval=0):
    # a is input array
    # N is blocksize

    # Store shape info
    m,n,r = a.shape

    # Generate random start indices for second and third axes keeping proper
    # distance from the boundaries for the block to be accomodated within.
    idx0 = np.random.randint(0,n-N+1,m)
    idx1 = np.random.randint(0,r-N+1,m)

    # Iterate through first and use slicing to assign fillval.
    for i in range(m):
        a[i, idx0[i]:idx0[i]+N, idx1[i]:idx1[i]+N] = fillval        
    return a

Sample run -

In [357]: a = np.arange(2*4*7).reshape(2,4,7)

In [358]: a
Out[358]: 
array([[[ 0,  1,  2,  3,  4,  5,  6],
        [ 7,  8,  9, 10, 11, 12, 13],
        [14, 15, 16, 17, 18, 19, 20],
        [21, 22, 23, 24, 25, 26, 27]],

       [[28, 29, 30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39, 40, 41],
        [42, 43, 44, 45, 46, 47, 48],
        [49, 50, 51, 52, 53, 54, 55]]])

In [359]: random_block_fill_adv(a, N=3, fillval=0)
Out[359]: 
array([[[ 0,  0,  0,  0,  4,  5,  6],
        [ 7,  0,  0,  0, 11, 12, 13],
        [14,  0,  0,  0, 18, 19, 20],
        [21, 22, 23, 24, 25, 26, 27]],

       [[28, 29, 30, 31, 32, 33, 34],
        [35, 36, 37, 38,  0,  0,  0],
        [42, 43, 44, 45,  0,  0,  0],
        [49, 50, 51, 52,  0,  0,  0]]])

Fun stuff : Being in-place filling, if we keep running random_block_fill_adv(a, N=3, fillval=0), we will eventually end up with all zeros a. Thus, also verifying the code.

---

Runtime test

In [579]: a = np.random.randint(0,9,(10000,4,4))

In [580]: %timeit random_block_fill_lidx(a, N=2, fillval=0)
     ...: %timeit random_block_fill_adv(a, N=2, fillval=0)
     ...: %timeit random_block_fill_loopy(a, N=2, fillval=0)
     ...: 
1000 loops, best of 3: 545 µs per loop
1000 loops, best of 3: 891 µs per loop
100 loops, best of 3: 10.6 ms per loop

In [581]: a = np.random.randint(0,9,(1000,40,40))

In [582]: %timeit random_block_fill_lidx(a, N=10, fillval=0)
     ...: %timeit random_block_fill_adv(a, N=10, fillval=0)
     ...: %timeit random_block_fill_loopy(a, N=10, fillval=0)
     ...: 
1000 loops, best of 3: 739 µs per loop
1000 loops, best of 3: 671 µs per loop
1000 loops, best of 3: 1.27 ms per loop

So, which one to choose depends on the first axis length and blocksize.