Eli Rose 's postsabout code trinkets quotesmedia log
Numpy Indexing

[Epistemic status: quite confident, statements about a formal system, but not complete documentation.]

I'm learning Numpy. The various ways to index arrays confused me, and the docs I found inscrutable. These are my notes I wrote to myself as I was figuring it out. Note: This is just what I've needed so far; I'm sure I'm missing other features or edge cases.

Say you're given an indexing expression like X[:, [1], 1:10]. What is the shape of the result going to be? Reading from left to right, you can construct it iteratively, based on the shape of the array X you're indexing into and the type of each thing you encounter between the commas.

Let's define an array to use in examples:

>>> A = np.array([[1, 2, 3], [2, 4, 6]])
>>> A.ndim
2
>>> A.shape
(2, 3)

Here's what each type of thing you can index by means. I'll focus mostly on how to construct the shape of the result, since I am aiming for simplicity and ease of remembering, and my confusions are more often about the shape than about the result itself.

:

(..., k, ...)[..., :, ...] -> (..., k, ...)
Each : you encounter means "the shape of the result will have an entry here, and the entry is equal to the shape of A in this slot".

  • Example: A[:, :] == A, shape (2, 3) -> (2, 3)
  • Example: A[:] == A as well, because if not enough indices are specified, the rest fills in with ":"

a:b

(..., k, ...)[..., a:b, ...] -> (..., b - a, ...)
Each a:b you encounter means "the shape of the result will have an entry here, and the entry is equal to b - a".

  • This is true as long as there aren't actually more elements than you're slicing.
  • Example: A[:, 0:1] == [[1], [2]], shape (2, 3) -> (2, 1)

a

(..., k, ...)[..., a, ...] -> (..., ...)
Each integer you encounter means "the shape of the result will not have an entry here".

  • Example: A[:, 0] == [1, 2], shape (2, 3) -> (2,)
  • Example: A[0, 1] == 2, shape (2, 3) -> (,)

[a]

(..., k, ...)[..., [a], ...] -> (..., *[a].shape, ...)
Each numpy array or Python list you encounter means "the shape of the result will, have the shape of this numpy array or Python list inserted here."

  • Example: A[[0], :] == [[1], [2]], shape (2, 3) -> (1, 3)
  • Example: A[[[0]], :].shape == (1, 1, 3)
  • Example: A[[[0], [1]], :].shape == (2, 1, 3)

If you want further info, I did find the scipy-cookbook docs approachable.