I have a numpy array with shape (M, N, N) – it’s effectively a bunch (M) of (N,N) covariance matrices. I want to be able to extract submatrices out of this with shape (M, P, P). But I’m trying to access non-contiguous indices. Here’s an example:

import numpy as np

# Display all the columns
np.set_printoptions(threshold=False, edgeitems=50, linewidth=200)

# Create a 6 x 6 matrix.
x = np.arange(36).reshape(6,6)

# Now make multiple copies to practice with.
y = np.array([x, x])

print(f"{y.shape=}n")

print(f"{y=}n")

# We want to extract the submatrices containting the first 2 indices
# and the last 2 indices. There are an "unknown" number of intermediate
# indices - in this example 2. Thus I'm using negative indices to get the
# last two indicies.

# Extraction using advanced indexing
s = np.array([[0, 1] + [-2, -1]])

subset = y[:, s.T, s]
print(f"{subset=}n")

# Now try it with numpy slices. This approach doesn't work
first_slice = np.s_[0:2]
second_slice = np.s_[-2:]
combined_slice = np.r_[first_slice, second_slice]

subset = y[:, combined_slice, combined_slice]
print(subset)

That produces the following output:

y.shape=(2, 6, 6)

y=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]],

       [[ 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]]])

subset=array([[[ 0,  1,  4,  5],
        [ 6,  7, 10, 11],
        [24, 25, 28, 29],
        [30, 31, 34, 35]],

       [[ 0,  1,  4,  5],
        [ 6,  7, 10, 11],
        [24, 25, 28, 29],
        [30, 31, 34, 35]]])

[[0 7]
 [0 7]]

Using array based indexing works, but it seems like there should be a way to do this with slice objects too. Any tips would be appreciated. Thanks!