McKinney Chapter 4 - NumPy Basics: Arrays and Vectorized Computation

FINA 6333 for Spring 2025

Author

Richard Herron 

import numpy as np
%precision 4
'%.4f'

Introduction

Chapter 4 of McKinney (2022) discusses the NumPy package (an abbreviation of numerical Python), which is the foundation for numerical computing in Python, including pandas.

We will focus on:

  1. Creating arrays
  2. Slicing arrays
  3. Applying functions and methods to arrays
  4. Using conditional logic with arrays (i.e., np.where() and np.select())

Note: Indented block quotes are from McKinney (2022) unless otherwise indicated. The section numbers here differ from McKinney (2022) because we will only discuss some topics.

Here is a simple example of NumPy’s speed and syntax advantages relative to Python’s built-in data structures. First, we create a list and a NumPy array with values from 0 to 999,999.

my_list = list(range(1_000_000))
my_arr = np.arange(1_000_000)
my_list[:5]
[0, 1, 2, 3, 4]
my_arr[:5]
array([0, 1, 2, 3, 4])

We must use a for loop or a list comprehension to double each value in my_list. We will comment this code because it prints a list with one million elements!

# [2 * x for x in my_list] # list comprehension to double each value

However, we can multiply my_arr by two because math “just works” with NumPy. Jupyter will pretty print the NumPy array, showing only the first and last few elements.

my_arr * 2
array([      0,       2,       4, ..., 1999994, 1999996, 1999998],
      shape=(1000000,))

We can use the “magic” function %timeit to time these two calculations.

%timeit [x * 2 for x in my_list]
72.5 ms ± 16.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit my_arr * 2
4.44 ms ± 867 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

The NumPy version is much faster than the list version! The NumPy version is also faster to type, read, and troubleshoot, and our time is more valuable than computer time!

The NumPy ndarray: A Multidimensional Array Object

One of the key features of NumPy is its N-dimensional array object, or ndarray, which is a fast, flexible container for large datasets in Python. Arrays enable you to perform mathematical operations on whole blocks of data using similar syntax to the equivalent operations between scalar elements.

np.random.seed(42) # makes random numbers repeatable
data = np.random.randn(2, 3)
data
array([[ 0.4967, -0.1383,  0.6477],
       [ 1.523 , -0.2342, -0.2341]])

Multiplying data by 10 multiplies each element in data by 10, and adding data to itself adds each element to itself (i.e., element-wise addition). NumPy arrays must contain homogeneous data types (e.g., all floats or integers) to achieve this common-sense behavior.

data * 10
array([[ 4.9671, -1.3826,  6.4769],
       [15.2303, -2.3415, -2.3414]])
data_2 = data + data
data_2
array([[ 0.9934, -0.2765,  1.2954],
       [ 3.0461, -0.4683, -0.4683]])

NumPy arrays have attributes. Recall that IPython and Jupyter provide tab completion.

data.ndim
2
data.shape
(2, 3)
data.dtype
dtype('float64')

We slice NumPy arrays using [], the same as we slice lists and tuples.

data[0]
array([ 0.4967, -0.1383,  0.6477])

We chain []s with arrays, the same as we chain []s with lists and tuples.

data[0][0]
0.4967

However, with NumPy arrays, we can replace \(n\) chained []s with one pair of []s containing \(n\) indexes or slices, separated by commas. For example, [i][j] becomes [i, j], and [i][j][k] becomes [i, j, k].

data[0, 0] # zero row, zero column
0.4967
data[0][0] == data[0, 0]
np.True_

Creating ndarrays

The easiest way to create an array is to use the array function. This accepts any sequence-like object (including other arrays) and produces a new NumPy array containing the passed data

data1 = [6, 7.5, 8, 0, 1]
arr1 = np.array(data1)
arr1
array([6. , 7.5, 8. , 0. , 1. ])
arr1.dtype
dtype('float64')

Here, np.array() implicitly casts the integers in data1 to floats because NumPy arrays must have homogenous data types. We could explicitly cast all values to integers but would lose information.

np.array(data1, dtype=np.int64)
array([6, 7, 8, 0, 1])

We can cast a list of lists to a two-dimensional NumPy array.

data2 = [[1, 2, 3, 4], [5, 6, 7, 8]]
arr2 = np.array(data2)
arr2
array([[1, 2, 3, 4],
       [5, 6, 7, 8]])
arr2.shape
(2, 4)
arr2.dtype
dtype('int64')

There are several other ways to create NumPy arrays.

np.zeros((3, 6))
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]])
np.ones((3, 6))
array([[1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1.]])
np.ones_like(arr2)
array([[1, 1, 1, 1],
       [1, 1, 1, 1]])

The np.arange() function is similar to Python’s built-in range() but creates an array directly.

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

Table 4-1 from McKinney (2022) summarizes NumPy array creation functions.

  • array: Convert input data (list, tuple, array, or other sequence type) to an ndarray either by inferring a dtype or explicitly specifying a dtype; copies the input data by default
  • asarray: Convert input to ndarray, but do not copy if the input is already an ndarray
  • arange: Like the built-in range but returns an ndarray instead of a list
  • ones, ones_like: Produce an array of all 1s with the given shape and dtype; ones_like takes another array and produces a ones array of the - same shape and dtype
  • zeros, zeros_like: Like ones and ones_like but producing arrays of 0s instead
  • empty, empty_like: Create new arrays by allocating new memory, but do not populate with any values like ones and zeros
  • full, full_like: Produce an array of the given shape and dtype with all values set to the indicated “fill value”
  • eye, identity: Create a square N-by-N identity matrix (1s on the diagonal and 0s elsewhere)

Arithmetic with NumPy Arrays

Arrays are important because they enable you to express batch operations on data without writing any for loops. NumPy users call this vectorization. Any arithmetic operations between equal-size arrays applies the operation element-wise

arr = np.array([[1., 2., 3.], [4., 5., 6.]])
arr
array([[1., 2., 3.],
       [4., 5., 6.]])

NumPy array addition is elementwise.

arr + arr
array([[ 2.,  4.,  6.],
       [ 8., 10., 12.]])

NumPy array multiplication is elementwise.

arr * arr
array([[ 1.,  4.,  9.],
       [16., 25., 36.]])

NumPy array division is elementwise.

1 / arr
array([[1.    , 0.5   , 0.3333],
       [0.25  , 0.2   , 0.1667]])

NumPy powers are elementwise, too.

arr ** 2
array([[ 1.,  4.,  9.],
       [16., 25., 36.]])

We can also raise a single value to the elements in an array!

2 ** arr
array([[ 2.,  4.,  8.],
       [16., 32., 64.]])

Basic Indexing and Slicing

We index and slice one-dimensional arrays in the same way as lists and tuples.

arr = np.arange(10)
arr
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
arr[5]
np.int64(5)
arr[5:8]
array([5, 6, 7])
equiv_list = list(range(10))
equiv_list
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
equiv_list[5:8]
[5, 6, 7]

We must jump through some hoops to replace elements 5, 6, and 7 with the value 12 in the list equiv_list.

# # TypeError: can only assign an iterable
# equiv_list[5:8] = 12
equiv_list[5:8] = [12] * 3
equiv_list
[0, 1, 2, 3, 4, 12, 12, 12, 8, 9]

However, this operation is easy with NumPy array arr!

arr[5:8] = 12
arr
array([ 0,  1,  2,  3,  4, 12, 12, 12,  8,  9])

We call this behavior “broadcasting”.

As you can see, if you assign a scalar value to a slice, as in arr[5:8] = 12, the value is propagated (or broadcasted henceforth) to the entire selection. An important first distinction from Python’s built-in lists is that array slices are views on the original array. This means that the data is not copied, and any modifications to the view will be reflected in the source array.

arr_slice = arr[5:8]
arr_slice
array([12, 12, 12])
arr_slice[1] = 12345
arr_slice
array([   12, 12345,    12])
arr
array([    0,     1,     2,     3,     4,    12, 12345,    12,     8,
           9])

The : slices every element in arr_slice.

arr_slice[:] = 64
arr_slice
array([64, 64, 64])
arr
array([ 0,  1,  2,  3,  4, 64, 64, 64,  8,  9])

If you want a copy of a slice of an ndarray instead of a view, you will need to explicitly copy the array-for example, arr[5:8].copy().

arr_slice_2 = arr[5:8].copy()
arr_slice_2
array([64, 64, 64])
arr_slice_2[:] = 2_001
arr_slice_2
array([2001, 2001, 2001])
arr
array([ 0,  1,  2,  3,  4, 64, 64, 64,  8,  9])

Indexing with slices

We can slice across two or more dimensions and use the [i, j] notation.

arr2d = np.array([[1,2,3], [4,5,6], [7,8,9]])
arr2d
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
arr2d[:2]
array([[1, 2, 3],
       [4, 5, 6]])
arr2d[:2, 1:]
array([[2, 3],
       [5, 6]])

A colon (:) by itself selects the entire dimension and is necessary to slice higher dimensions.

arr2d[:, :1]
array([[1],
       [4],
       [7]])
arr2d[:2, 1:] = 0
arr2d
array([[1, 0, 0],
       [4, 0, 0],
       [7, 8, 9]])

Always check your output!

Boolean Indexing

We can use Booleans (i.e., True and False) to slice arrays, too. Boolean indexing in Python is like combining index() and match() in Excel.

names = np.array(['Bob', 'Joe', 'Will', 'Bob', 'Will', 'Joe', 'Joe'])
np.random.seed(42)
data = np.random.randn(7, 4)
names
array(['Bob', 'Joe', 'Will', 'Bob', 'Will', 'Joe', 'Joe'], dtype='<U4')
data
array([[ 0.4967, -0.1383,  0.6477,  1.523 ],
       [-0.2342, -0.2341,  1.5792,  0.7674],
       [-0.4695,  0.5426, -0.4634, -0.4657],
       [ 0.242 , -1.9133, -1.7249, -0.5623],
       [-1.0128,  0.3142, -0.908 , -1.4123],
       [ 1.4656, -0.2258,  0.0675, -1.4247],
       [-0.5444,  0.1109, -1.151 ,  0.3757]])

Here names provides seven names for the seven rows in data.

names == 'Bob'
array([ True, False, False,  True, False, False, False])
data[names == 'Bob']
array([[ 0.4967, -0.1383,  0.6477,  1.523 ],
       [ 0.242 , -1.9133, -1.7249, -0.5623]])

We can combine Boolean slicing with : slicing.

data[names == 'Bob', 2:]
array([[ 0.6477,  1.523 ],
       [-1.7249, -0.5623]])

We can use ~ to invert a Boolean.

cond = names == 'Bob'
data[~cond]
array([[-0.2342, -0.2341,  1.5792,  0.7674],
       [-0.4695,  0.5426, -0.4634, -0.4657],
       [-1.0128,  0.3142, -0.908 , -1.4123],
       [ 1.4656, -0.2258,  0.0675, -1.4247],
       [-0.5444,  0.1109, -1.151 ,  0.3757]])

For NumPy arrays, we must use & and | instead of and and or.

cond = (names == 'Bob') | (names == 'Will')
data[cond]
array([[ 0.4967, -0.1383,  0.6477,  1.523 ],
       [-0.4695,  0.5426, -0.4634, -0.4657],
       [ 0.242 , -1.9133, -1.7249, -0.5623],
       [-1.0128,  0.3142, -0.908 , -1.4123]])

We can also create a Boolean for each element.

data
array([[ 0.4967, -0.1383,  0.6477,  1.523 ],
       [-0.2342, -0.2341,  1.5792,  0.7674],
       [-0.4695,  0.5426, -0.4634, -0.4657],
       [ 0.242 , -1.9133, -1.7249, -0.5623],
       [-1.0128,  0.3142, -0.908 , -1.4123],
       [ 1.4656, -0.2258,  0.0675, -1.4247],
       [-0.5444,  0.1109, -1.151 ,  0.3757]])
data < 0
array([[False,  True, False, False],
       [ True,  True, False, False],
       [ True, False,  True,  True],
       [False,  True,  True,  True],
       [ True, False,  True,  True],
       [False,  True, False,  True],
       [ True, False,  True, False]])
data[data < 0] = 0
data
array([[0.4967, 0.    , 0.6477, 1.523 ],
       [0.    , 0.    , 1.5792, 0.7674],
       [0.    , 0.5426, 0.    , 0.    ],
       [0.242 , 0.    , 0.    , 0.    ],
       [0.    , 0.3142, 0.    , 0.    ],
       [1.4656, 0.    , 0.0675, 0.    ],
       [0.    , 0.1109, 0.    , 0.3757]])

Universal Functions: Fast Element-Wise Array Functions

A universal function, or ufunc, is a function that performs element-wise operations on data in ndarrays. You can think of them as fast vectorized wrappers for simple functions that take one or more scalar values and produce one or more scalar results.

arr = np.arange(10)
arr
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
np.sqrt(arr)
array([0.    , 1.    , 1.4142, 1.7321, 2.    , 2.2361, 2.4495, 2.6458,
       2.8284, 3.    ])

Like above, we can raise a single value to a NumPy array of powers.

2**arr
array([  1,   2,   4,   8,  16,  32,  64, 128, 256, 512])

np.exp(x) is \(e^x\).

np.exp(arr)
array([1.0000e+00, 2.7183e+00, 7.3891e+00, 2.0086e+01, 5.4598e+01,
       1.4841e+02, 4.0343e+02, 1.0966e+03, 2.9810e+03, 8.1031e+03])

Table 4-4 from McKinney (2022) summarizes fast, element-wise unary functions:

  • abs, fabs: Compute the absolute value element-wise for integer, floating-point, or complex values
  • sqrt: Compute the square root of each element (equivalent to arr ** 0.5)
  • square: Compute the square of each element (equivalent to arr ** 2)
  • exp: Compute the exponent \(e^x\) of each element
  • log, log10, log2, log1p: Natural logarithm (base e), log base 10, log base 2, and log(1 + x), respectively
  • sign: Compute the sign of each element: 1 (positive), 0 (zero), or –1 (negative)
  • ceil: Compute the ceiling of each element (i.e., the smallest integer greater than or equal to that number)
  • floor: Compute the floor of each element (i.e., the largest integer less than or equal to each element)
  • rint: Round elements to the nearest integer, preserving the dtype
  • modf: Return fractional and integral parts of array as a separate array
  • isnan: Return boolean array indicating whether each value is NaN (Not a Number)
  • isfinite, isinf: Return boolean array indicating whether each element is finite (non-inf, non-NaN) or infinite, respectively
  • cos, cosh, sin, sinh, tan, tanh: Regular and hyperbolic trigonometric functions
  • arccos, arccosh, arcsin, arcsinh, arctan, arctanh: Inverse trigonometric functions
  • logical_not: Compute truth value of not x element-wise (equivalent to ~arr).

These “unary” functions operate on one array and return a new array with the same shape. There are also “binary” functions that operate on two arrays and return one array.

np.random.seed(42)
x = np.random.randn(8)
y = np.random.randn(8)
x
array([ 0.4967, -0.1383,  0.6477,  1.523 , -0.2342, -0.2341,  1.5792,
        0.7674])
y
array([-0.4695,  0.5426, -0.4634, -0.4657,  0.242 , -1.9133, -1.7249,
       -0.5623])
np.maximum(x, y)
array([ 0.4967,  0.5426,  0.6477,  1.523 ,  0.242 , -0.2341,  1.5792,
        0.7674])

Table 4-5 from McKinney (2022) summarizes fast, element-wise binary functions:

  • add: Add corresponding elements in arrays
  • subtract: Subtract elements in second array from first array
  • multiply: Multiply array elements
  • divide, floor_divide: Divide or floor divide (truncating the remainder)
  • power: Raise elements in first array to powers indicated in second array
  • maximum, fmax: Element-wise maximum; fmax ignores NaN
  • minimum, fmin: Element-wise minimum; fmin ignores NaN
  • mod: Element-wise modulus (remainder of division)
  • copysign: Copy sign of values in second argument to values in first argument
  • greater, greater_equal, less, less_equal, equal, not_equal: Perform element-wise comparison, yielding boolean array (equivalent to infix operators >, >=, <, <=, ==, !=)
  • logical_and, logical_or, logical_xor: Compute element-wise truth value of logical operation (equivalent to infix operators & |, ^)

Array-Oriented Programming with Arrays

Using NumPy arrays enables you to express many kinds of data processing tasks as concise array expressions that might otherwise require writing loops. This practice of replacing explicit loops with array expressions is commonly referred to as vectorization. In general, vectorized array operations will often be one or two (or more) orders of magnitude faster than their pure Python equivalents, with the biggest impact in any kind of numerical computations. Later, in Appendix A, I explain broadcasting, a powerful method for vectorizing computations.

Expressing Conditional Logic as Array Operations

The numpy.where function is a vectorized version of the ternary expression x if condition else y.

np.where() is an if-else statement, like Excel’s if().

xarr = np.array([1.1, 1.2, 1.3, 1.4, 1.5])
yarr = np.array([2.1, 2.2, 2.3, 2.4, 2.5])
cond = np.array([True, False, True, True, False])
np.where(cond, xarr, yarr)
array([1.1, 2.2, 1.3, 1.4, 2.5])

We could use a list comprehension instead, but it takes longer to type, read, and troubleshoot.

np.array([(x if c else y) for x, y, c in zip(xarr, yarr, cond)])
array([1.1, 2.2, 1.3, 1.4, 2.5])

np.select() lets us test more than one condition and has a default value if no condition is met.

np.select(
    condlist=[cond==True, cond==False],
    choicelist=[xarr, yarr]
)
array([1.1, 2.2, 1.3, 1.4, 2.5])

Mathematical and Statistical Methods

A set of mathematical functions that compute statistics about an entire array or about the data along an axis are accessible as methods of the array class. You can use aggregations (often called reductions) like sum, mean, and std (standard deviation) either by calling the array instance method or using the top-level NumPy function.

np.random.seed(42)
arr = np.random.randn(5, 4)
arr
array([[ 0.4967, -0.1383,  0.6477,  1.523 ],
       [-0.2342, -0.2341,  1.5792,  0.7674],
       [-0.4695,  0.5426, -0.4634, -0.4657],
       [ 0.242 , -1.9133, -1.7249, -0.5623],
       [-1.0128,  0.3142, -0.908 , -1.4123]])
arr.mean()
-0.1713
arr.sum()
-3.4260

The aggregation methods above aggregated the whole array. We can use the axis argument to aggregate columns (axis=0) and rows (axis=1).

arr.mean(axis=1)
array([ 0.6323,  0.4696, -0.214 , -0.9896, -0.7547])
arr[0].mean()
0.6323
arr[1].mean()
0.4696
arr.mean(axis=0)
array([-0.1956, -0.2858, -0.1739, -0.03  ])
arr[:, 0].mean()
-0.1956
arr[:, 1].mean()
-0.2858

The .cumsum() method returns the sum of all previous elements.

arr = np.array([0, 1, 2, 3, 4, 5, 6, 7]) # same output as np.arange(8)
arr.cumsum()
array([ 0,  1,  3,  6, 10, 15, 21, 28])

We can also use the .cumsum() method along the axis of a multi-dimensional array.

arr = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
arr
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
arr.cumsum(axis=0)
array([[ 0,  1,  2],
       [ 3,  5,  7],
       [ 9, 12, 15]])
arr.cumprod(axis=1)
array([[  0,   0,   0],
       [  3,  12,  60],
       [  6,  42, 336]])

Table 4-6 from McKinney (2022) summarizes basic statistical methods:

  • sum: Sum of all the elements in the array or along an axis; zero-length arrays have sum 0
  • mean: Arithmetic mean; zero-length arrays have NaN mean
  • std, var: Standard deviation and variance, respectively, with optional degrees of freedom adjustment (default denominator \(n\))
  • min, max: Minimum and maximum
  • argmin, argmax: Indices of minimum and maximum elements, respectively
  • cumsum: Cumulative sum of elements starting from 0
  • cumprod: Cumulative product of elements starting from 1

References

McKinney, Wes. 2022. Python for Data Analysis. 3rd ed. https://wesmckinney.com/book/; O’Reilly Media, Inc.