NumPy Cheatsheet

Things about NumPy that I am prone to forgetting.

Axis Ordering

• 2D NumPy arrays are ordered as shown.

Indexing

• To get a row of a two-dimensional array as a one-dimensional array, you can use a single index. For example:

  import numpy as np
  
  bar = np.array(
      [
          [5, 10, 15, 20],
          [25, 30, 35, 40],
          [45, 50, 55, 60],
      ]
  )
  
  print(bar[0])  # prints array([ 5, 10, 15, 20])
  

• To get a row of a two-dimensional array as a two-dimensional array, you can use None or slicing. For example:

  import numpy as np
  
  bar = np.array(
      [
          [5, 10, 15, 20],
          [25, 30, 35, 40],
          [45, 50, 55, 60],
      ]
  )
  
  print(bar[0, None])  # prints array([[5, 10, 15, 20]])
  
  print(bar[:1])  # also prints array([[5, 10, 15, 20]])
  

• You can index a list of specific elements in a two-dimensional array by passing in a list of the points’ row and column indices. For example:

  import numpy as np
  
  bar = np.array(
      [
          [5, 10, 15, 20],
          [25, 30, 35, 40],
          [45, 50, 55, 60],
      ]
  )
  
  print(bar[[0, 1, 2], [0, 1, 2]])
  # array([5, 30, 55])
  

• When every dimension of an array is indexed with an array, and each of those arrays is the same shape, the output of the indexing operation will be the same shape as the index arrays. For example:

  import numpy as np
  
  foo = np.arange(3 * 2 * 4).reshape((3, 2, 4))
  print(foo)
  # [[[ 0  1  2  3]
  #   [ 4  5  6  7]]
  #
  #  [[ 8  9 10 11]
  #   [12 13 14 15]]
  #
  #  [[16 17 18 19]
  #   [20 21 22 23]]]
  
  print(foo[[[0, 2], [2, 0], [1, 1]], [[0, 0], [0, 0], [1, 1]], [[0, 1], [0, 2], [0, 3]]])
  # [[ 0 17]
  #  [16  2]
  #  [12 15]]
  

• If index arrays are not the same shape, but can be broadcasted, the output of the indexing operation will be the same shape as the broadcasted array.

Broadcasting

• Suppose we want to add two arrays, A and B.

  • Moving backwards from the last dimension of each array, we check if their dimensions are compatible.
  • Dimensions are compatible if they are equal or either of them is one.
  • If all of A’s dimensions are compatible with B’s dimensions, or vice versa, they are compatible arrays.

Reshaping Arrays

• You can insert new axes in front of and behind existing axes using None or np.newaxis (which is just an alias for None). For example:

  import numpy as np
  
  A = np.array(
      [5, 10, 15, 20],
  )
  
  print(A[:, np.newaxis].shape)  # prints (4, 1)
  print(A[np.newaxis, :].shape)  # prints (1, 4)
  

• You can reshape arrays in-place by modifying the shape attribute instead of using reshape. For example:

  import numpy as np
  
  A = np.array(
      [
          [1, 2, 3, 4],
          [5, 6, 7, 8],
      ],
  )
  
  A.reshape(4, 2)  # creates a copy of `A` and reshapes the copy
  A.shape = (4, 2)  # directly modifies `A`
  

Logical Indexing

• To combine logical statements, use bitwise operators like & and |.
• When combining logical statements, each statement should be wrapped with parentheses, so that they are evaluated before the any bitwise operators.

NaN

• np.nan is a special floating point constant. It cannot be inserted into arrays of non-floating point values.
• Use np.isnan() to find np.nan values in an array. np.nan is designed in such a way that np.nan == np.nan evaluates to False, but np.nan != np.nan evaluates to True.

Math Functions

• NumPy has a number of math functions that behave similarly.

  import numpy as np
  
  bar = np.array(
      [
          [5.0, 2.0, 9.0],
          [1.0, 0.0, 2.0],
          [1.0, 7.0, 8.0],
      ]
  )
  
  # prints 35.0, the sum of all elements in `bar`
  print(np.sum(bar))
  
  # prints array([7., 9., 19.]), the column sums of the array
  print(np.sum(bar, axis=0))
  
  # prints array([16., 3., 16.]), the row sums of the array
  print(np.sum(bar, axis=1))
  
  # Note that in the above cases, `sum` returns a 1D array. To retain the
  # dimensions of `bar`, set `keepdims=True`. For example, the following prints
  # array[[16.], [3.], [16.]].
  print(np.sum(bar, axis=1, keepdims=True))
  

• If any values in an array are nan and you try to apply sum to the array, it will return nan, since by default it is not expecting nan values in the array. np.nansum will instead treat nan values as 0.