Slicing a 3D Numpy array using boolean and index arrays

Question

Given a 3D numpy matrix D , I need to extract a 3D slice given a boolean vector b and index vector k.

D is a 3D Numpy array with dimensions (P, n, n) and data type float.
b is a 1D Numpy array with dimensions (P,) and data type bool.
k is a 1D Numpy array with dimensions (P,) and data type int.

b is populated based on a prior comparison check on D.

Each element p of k corresponds to the desired start index (for the slicing operation) for the p^th 2D matrix in D, where:

0≤p<P and 0≤k[p]<n-1

It is important to keep in mind that the elements of k are not uniform.

I can do element wise assignment and extraction i.e,

F[b,k[b],k[b]] = D[b,k[b],k[b]]

where F had the same dimensions as D.

But I cannot generate a 3D matrix G, which should be created by extracting Rx2x2 slices from D. R≤P and is the number of elements in b that satisfy the boolean condition.

Ignoring b, I wish to do something like the following:

G = D[:,k:(k+1), k:(k+1)]

where the p^th 2D slice of D can be visualized below:

+------------------+----------------------+
| (k[p], k[p])     | (k[p], k[p] + 1)     |
+------------------+----------------------+
| (k[p] + 1, k[p]) | (k[p] + 1, k[p] + 1) |
+------------------+----------------------+

Executing the line of code preceding the tabular representation understandably returns a Type Error.

TypeError: only integer scalar arrays can be converted to a scalar index

How do I generate G, without using a for loop to iterate through the first axis (axis 0)? Is a vectorized solution possible?

A small python script is attached below as a minimal example and in order to recreate the problem i am facing.

import numpy as np


P = 4096
n = 4

D = np.random.default_rng().normal(scale=1, size=(P,n,n))
k = np.zeros((P,), dtype=int)


# Is the top-left element in each pth n*n slice of D <= 0.5 ?
b = np.less_equal(D[:,0,0], 0.5) 


# Create a 3D array with dimensions similar to D
F = np.zeros((P,n,n),)

# Replace for top left element at axis 0 indices (p) for which b is True
F[b,k[b],k[b]] = D[b,k[b],k[b]]



# Attempt to generate G. Type Error:  only integer scalar arrays can be converted to a scalar index
G = D[b,k[b]:k[b]+1, k[b]:k[b]+1]



# Init G as a P*2*2 array.
G =np.zeros((P,2,2),)

#Attempt to populate G with 2x2 slices from D if b[p] is True and given base indices in k for pth slice.
# TypeError: only integer scalar arrays can be converted to a scalar index
G[b,k[b]:k[b]+1, k[b]:k[b]+1] = D[b,k[b]:k[b]+1, k[b]:k[b]+1]

Tls Chris · Answer 1 · 2023-02-24T14:03:53.327

I'm not sure I understand the question completely but it may be better to do the indexing in two steps. Filter on b then extract the required slice(s).

import numpy as np

arr = np.arange( 10*4*4 ).reshape( 10, 4, 4 )

b = np.array( [ False, True, False, True, True, True, False, True, False, True ] )

arr[b]

Result

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

       [[ 48,  49,  50,  51],
        [ 52,  53,  54,  55],
        [ 56,  57,  58,  59],
        [ 60,  61,  62,  63]],

       [[ 64,  65,  66,  67],
        [ 68,  69,  70,  71],
        [ 72,  73,  74,  75],
        [ 76,  77,  78,  79]],

       [[ 80,  81,  82,  83],
        [ 84,  85,  86,  87],
        [ 88,  89,  90,  91],
        [ 92,  93,  94,  95]],

       [[112, 113, 114, 115],
        [116, 117, 118, 119],
        [120, 121, 122, 123],
        [124, 125, 126, 127]],

       [[144, 145, 146, 147],
        [148, 149, 150, 151],
        [152, 153, 154, 155],
        [156, 157, 158, 159]]])

arr[b][:, :2, :2 ] # With k = 2

Result:

array([[[ 16,  17],
        [ 20,  21]],

       [[ 48,  49],
        [ 52,  53]],

       [[ 64,  65],
        [ 68,  69]],

       [[ 80,  81],
        [ 84,  85]],

       [[112, 113],
        [116, 117]],

       [[144, 145],
        [148, 149]]])

If k is an int array

k = np.array( [1, 2, 0, 1, 0, 1, 2, 2, 0, 1])

result = []
for s, kix  in zip(arr[b], k[b]):
    result.append( s[ kix:kix+2, kix:kix+2 ] )
result= np.array( result )

result
# array([[[ 26,  27],
#         [ 30,  31]],

#        [[ 53,  54],
#         [ 57,  58]],

#        [[ 64,  65],
#         [ 68,  69]],

#        [[ 85,  86],
#         [ 89,  90]],

#        [[122, 123],
#         [126, 127]],

#        [[149, 150],
#         [153, 154]]])

Slicing a 3D Numpy array using boolean and index arrays

1 Answers1