Home   Research   Notes

Numpy Tutorial

This is an introductory tutorial for Numpy, which assumes you already know the basics of Python. I packed a fair amount of material into it, so it's a bit of an information firehose. If you aren't in a rush, you may learn better by reading one section a day and then experimenting with what you learned rather than blasting through in one sitting, despite how short the sections are. Also, I've scattered relevant links to the Numpy documentation throughout in case you want more detail or an alternate framing.

To begin with, we need to import Numpy. np is the standard abbreviation, so we'll use that everywhere.

import numpy as np

Working with Arrays

First, here's a few common ways to make an array. You can turn an existing iterable (like a list) into an array, generate an array of regularly spaced numbers, create an array filled with zeros, ones, or some other number, or even generate the contents randomly (standard normal or 0-1 uniform).

  1. np.array(some_iterable)
  2. np.linspace(start, stop, count), np.logspace(start, stop, count)
  3. np.zeros(shape), np.ones(shape), np.full(shape, value)
  4. np.random.randn(shape), np.random.rand(shape)

These names are mostly self-explanatory, but if you aren't sure of something, remember you can run e.g. help(np.full) to read the documentation. Here is an example which creates an array of six numbers evenly spaced from zero to one:

x = np.linspace(0, 1, 6)
print(x)
[0.  0.2 0.4 0.6 0.8 1. ]

The nice thing about Numpy arrays is that you can use operations and functions on every element in the array using a clean syntax. This is also much faster than using a for loop.

y = 2.5 + 3 * np.sin(2.5 * x)
print(y)
[2.5        3.93827662 5.02441295 5.49248496 5.22789228 4.29541643]

Numpy has a ton of functions for these operations. If you can't find something you're looking for, you can find even more functions in scipy.special (see here). Some examples are

Numpy also provides a number of summary statistics that iterate over the array, such as np.sum, np.mean, np.std, np.median, np.percentile, and many others.

print(np.mean(y), np.std(y))
4.413080540472895 1.0084575809987286

Of course, you can also just index an array to get a particular element, the same way you would for a Python list or tuple.

print(y[2])
5.024412954423689

Data Types and Comparisons

Numpy supports more data types than just the double-precision floats we've defaulted to so far, such as integers, complex numbers, strings, and booleans. You can specify which datatype you want when you create an array, or use np.astype to change the type afterwards. Complex numbers (Numpy, like C, calls them cdouble) are actually a built in part of python, so you can specify imaginary numbers by adding j to the end of a number, e.g. 3.5j and 1 + .5j.

cx = np.linspace(-3, 3, 5, dtype="cdouble")
cy = np.sqrt(4j + cx)
print("cx:", cx)
print("cy:", cy)
cx: [-3. +0.j -1.5+0.j  0. +0.j  1.5+0.j  3. +0.j]
cy: [1.        +2.j         1.17728541+1.6988234j  1.41421356+1.41421356j
 1.6988234 +1.17728541j 2.        +1.j        ]

Boolean arrays are particularly useful, as you can use them to filter your existing arrays. The main way to do this is by comparing an array with something, and then using the resulting boolean array to index any array of the same size. For example, here is every integer between 1 and 100 whose cos is greater than 0.9:

x = np.arange(100)
x_filtered = x[np.sin(x) > 0.9]
print(x_filtered)
[ 2  8 14 20 27 33 39 46 52 58 64 71 77 83 96]

This is extremely valuable for data processing, especially when combined with the operators && (and), || (or), ~ (not). Unlike the Python keywords and, or, and not (which raise an error), these work by comparing each element one-by-one. Some examples:

N-Dimensional Arrays

So far, I've shown you one-dimensional arrays, meaning they only have one index and are basically equivalent to vectors in math. Numpy can do up to a 32-dimensional arrays. This is what the shape parameter of the array-creation functions is for -- a shape of (10, 3) means the array is two-dimensional with 10 rows and 3 columns. Some useful functions for dealing with shapes are here. The ones I use the most are:

For example of a few of these together:

x = np.arange(12)
print("Setup:", x.ndim, x.shape, "\n", x)
x = x.reshape(3, 4)
print("Reshaped:", x.ndim, x.shape, "\n", x)
x = x.T
print("Transposed:", x.ndim, x.shape, "\n", x)
Setup: 1 (12,) 
 [ 0  1  2  3  4  5  6  7  8  9 10 11]
Reshaped: 2 (3, 4) 
 [[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
Transposed: 2 (4, 3) 
 [[ 0  4  8]
 [ 1  5  9]
 [ 2  6 10]
 [ 3  7 11]]

A helpful feature Numpy provides is the axis keyword on many of its functions that operate across a 1-d array. This lets you select which axis to operate on, rather than flattening the array and operating on the whole thing. For example, here's the average across each row:

np.mean(x, axis=1)
array([4., 5., 6., 7.])

Broadcasting and Slicing

To do an operation involving two Numpy arrays, they need to have compatible shapes. The simplest way to do this is for them to have the same shape, but Numpy has some rules for more complicated situations it calls broadcasting. They are:

  1. If the number of dimensions differs, implicitly add length-one dimensions to start of the lower-dimension array's shape until they match. E.g. (5,) -> (1, 1, 5) if 3 dimensions are needed.
  2. If the length of each dimension matches, operate on them element-wise (like a for loop does). If instead one array's length in that dimension is 1, reuse it for every element of the other array.

This probably sounds more complicated than it actually is, and these are incredibly useful once you get the hang of them. Here's an example of broadcasting:

a = np.array([[1, 2, 3], [4, 5, 6]])
b = np.array([-5, 10, 3])
print(a + b)
[[-4 12  6]
 [-1 15  9]]

This works because a has shape (2, 3) and b has shape (3,). When added, b is broadcast to (1, 3) following rule 1, so that b is added to each row of a (reusing it, following rule 2). On the other hand, if we tried a.T * b, it would fail; rule 1 would again broadcast b to (1, 3), but this isn't compatible with a.T, which has shape (3, 2). The second dimension lengths aren't equal nor is one of them 1, so Numpy will raise an exception.

You might imagine this means we're going to have to use reshape a ton just to get everything to line up properly. While it does get some use, a lot of it can be avoided by using Numpy's powerful slicing system. Similarly to Python slicing, this lets you select portions of an iterable using the indexing syntax arr[start:stop:step]. For example, here is every number between 0 and 99, sliced to select every 12th element beginning with the 20th element but before the 80th:

x = np.arange(100)
print(x[20:80:12])
[20 32 44 56 68]

You don't have to specify any of these, since start defaults to 0 (the start of the array), stop defaults to -1 (the end), and step defaults to 1. So x[::] and x[:] select everything (a no-op). More interestingly, the step can be negative to work backwards from start to stop (meaning stop must be < start to select anything):

print(x[50:20:-8])
[50 42 34 26]

All of that is standard Python slicing. Numpy arrays support slicing in several dimensions at once by separating the indices with commas, as in a[1, 2]; you can use : to skip an axis if you need to index later axes instead:

print(a[::-1])
print(a[:, ::-1])
[[4 5 6]
 [1 2 3]]
[[3 2 1]
 [6 5 4]]

Throwing one last wrench into things, you can also add additional length-1 dimensions as-needed with np.newaxis. This lets you adjust the dimensions quickly just before broadcasting. For example, here's how to make a multiplication table from a 1-d array (this is known as an outer product):

x = np.arange(1, 11)
print(x[:, np.newaxis] * x)
[[  1   2   3   4   5   6   7   8   9  10]
 [  2   4   6   8  10  12  14  16  18  20]
 [  3   6   9  12  15  18  21  24  27  30]
 [  4   8  12  16  20  24  28  32  36  40]
 [  5  10  15  20  25  30  35  40  45  50]
 [  6  12  18  24  30  36  42  48  54  60]
 [  7  14  21  28  35  42  49  56  63  70]
 [  8  16  24  32  40  48  56  64  72  80]
 [  9  18  27  36  45  54  63  72  81  90]
 [ 10  20  30  40  50  60  70  80  90 100]]

I could have written the second term as x[np.newaxis, :], but that was unneeded due to broadcasting rule 1 (above).

Combining slicing, broadcasting, filtering, and Numpy functions allow you to do some very complex operations on large amounts of data in just a few lines. In addition to being more terse than for loops, it is also faster since Numpy's internal iteration is handled in compiled C code.

Linear Algebra

There's a lot of interesting things you can do with multidimensional arrays, but linear algebra is central enough that I'll discuss it here (but here's the Numpy linalg documentation for good measure).

Starting things off, the symbol for matrix multiplication is "@" (they were scraping the bottom of the barrel on symbols):

M = np.random.rand(3, 4)
x = np.linspace(3, 10, 4)
print(M)
print("x:", x)
print("M@x:", M @ x)
[[0.24267934 0.00402404 0.09694515 0.5378315 ]
 [0.05639292 0.29902839 0.01299276 0.1725355 ]
 [0.69550555 0.07249993 0.17264741 0.72448763]]
x: [ 3.          5.33333333  7.66666667 10.        ]
M@x: [ 6.87106068  3.58896303 11.04168935]

Notice that x was treated as a column vector despite being 1-d and printed as a row of numbers. This is convenient, but can easily lead to confusion and programming errors, so be careful. If you transpose x, it is still a column vector since the first and last dimension are the same. Worse, if you think about how matrix multiplication works (e.g. AB ≠ BA in most cases), you'll see that this operation cannot follow the broadcasting rules described in the previous section. To avoid confusion, you might consider making all of your vectors 2-d using arr.reshape or [:, np.newaxis], so that it's completely clear which are row vectors and which are column vectors, and transposing works the way you'd expect.

As long as you can keep that straight, Numpy is very good for doing linear algebra. Here are a handful of particularly useful functions:

  1. np.dot(A, B) computes the dot product of two matrices or vectors.
  2. np.cross(A, B) computes the cross product of two vectors, so long as they are length 3 (or 2, which causes Numpy to assume the third component is zero).
  3. np.linalg.eye(n) creates an n × n identity matrix
  4. np.linalg.det(A) computes the determinant of a matrix.
  5. np.norm(x) computes the norm of a vector or matrix. There's quite a few options for the type of norm, but if you use the defaults (as I have here) you'll get the length of the vector.
  6. np.linalg.eig(A) computes the eigenvalues and eigenvectors of a square matrix.
  7. np.linalg.solve(A, x) computes the solution to the linear system A−1x. You can technically use np.linalg.inv(A) @ b but solve is more numerically stable so you should use it when you have a choice.
  8. np.linalg.lstsq(X, y) solves the least squares problem Xβ = y.

Of course, there is far more than this for you to use, but these are some of the core functions. If you can't find something you're looking for, it's probably in Scipy Linalg, which has a lot of functions for special types of matrices.