NumPy Illustrated: The Visual Guide to NumPy
NumPy is a fundamental library that most of the widely used Python data processing libraries are built upon (pandas, OpenCV), inspired by (PyTorch), or can efficiently share data with (TensorFlow, Keras, etc). Understanding how NumPy works gives a boost to your skills in those libraries as well. It is also possible to run NumPy code with no or minimal changes on GPU¹.
The central concept of NumPy is an n-dimensional array. The beauty of it is that most operations look just the same, no matter how many dimensions an array has. But 1D and 2D cases are a bit special. The article consists of three parts:
I took a great article, “A Visual Intro to NumPy” by Jay Alammar², as a starting point, significantly expanded its coverage, and amended a pair of nuances.
Numpy Array vs. Python List
At first glance, NumPy arrays are similar to Python lists. They both serve as containers with fast item getting and setting and somewhat slower inserts and removals of elements.
The hands-down simplest example when NumPy arrays beat lists is arithmetic:
Other than that, NumPy arrays are:
- more compact, especially when there’s more than one dimension
- faster than lists when the operation can be vectorized
- slower than lists when you append elements to the end
- usually homogeneous: can only work fast with elements of one type
Here O(N) means that the time necessary to complete the operation is proportional to the size of the array (see Big-O Cheat Sheet³ site), and O*(1) (the so-called “amortized” O(1)) means that the time does not generally depend on the size of the array (see Python Time Complexity⁴ wiki page)
1. Vectors, the 1D Arrays
One way to create a NumPy array is to convert a Python list. The type will be auto-deduced from the list element types:
Be sure to feed in a homogeneous list, otherwise you’ll end up with
dtype=’object’, which annihilates the speed and only leaves the syntactic sugar contained in NumPy.
NumPy arrays cannot grow the way a Python list does: No space is reserved at the end of the array to facilitate quick appends. So it is a common practice to either grow a Python list and convert it to a NumPy array when it is ready or to preallocate the necessary space with
It is often necessary to create an empty array which matches the existing one by shape and elements type:
Actually, all the functions that create an array filled with a constant value have a
There are as many as two functions for array initialization with a monotonic sequence in NumPy:
If you need a similar-looking array of floats, like
[0., 1., 2.], you can change the type of the
arange(3).astype(float), but there’s a better way. The
arange function is type-sensitive: If you feed ints as arguments, it will generate ints, and if you feed floats (e.g.,
arange(3.)) it will generate floats.
arange is not especially good at handling floats:
0.1 looks like a finite decimal number to us but not to the computer: In binary, it is an infinite fraction and has to be rounded somewhere thus an error. That’s why feeding a step with fractional part to
arange is generally a bad idea: You might run into an off-by-one error. You can make an end of the interval fall into a non-integer number of steps (solution1) but that reduces readability and maintainability. This is where
linspace might come in handy. It is immune to rounding errors and always generates the number of elements you ask for. There’s a common gotcha with
linspace, though. It counts points, not intervals, thus the last argument is always plus one to what you would normally think of. So it is 11, not 10 in the example above.
For testing purposes it is often necessary to generate random arrays:
Once you have your data in the array, NumPy is brilliant at providing easy ways of giving it back:
All of the indexing methods presented above except fancy indexing are actually so-called “views”: They don’t store the data and reflect the changes in the original array if it happens to get changed after being indexed.
All of those methods including fancy indexing are mutable: They allow modification of the original array contents through assignment, as shown above. This feature breaks the habit of copying arrays by slicing them:
Also, such assignments must not change the size of the array, so tricks like
won’t work in NumPy — use
np.append, etc. instead (described in the “2D” section below).
Another super-useful way of getting data from NumPy arrays is boolean indexing, which allows using all kinds of logical operators:
Careful though; Python “ternary” comparisons like
3<=a<=5 don’t work here.
Arithmetic is one of the places where NumPy speed shines most. Vector operators are shifted to the c++ level and allow us to avoid the costs of slow Python loops. NumPy allows the manipulation of whole arrays just like ordinary numbers:
The same way ints are promoted to floats when adding or subtracting, scalars are promoted (aka broadcasted) to arrays:
Most of the math functions have NumPy counterparts that can handle vectors:
Scalar product has an operator of its own:
You don’t need loops for trigonometry either:
Arrays can be rounded as a whole:
np.around is just an alias to
np.round introduced to avoid shadowing Python
round when you write
from numpy import * (as opposed to a more common
import numpy as np). You can use
a.round() as well.
NumPy is also capable of doing the basic stats:
Sorting functions have fewer capabilities than their Python counterparts:
In the 1D case, the absence of the
reversed keyword can be easily compensated for by reversing the result. In 2D it is somewhat trickier (feature request⁵).
Searching for an element in a vector
As opposed to Python lists, a NumPy array does not have an
index method. The corresponding feature request⁶ has been hanging there for quite a while now.
- One way of finding an element is
np.where(a==x), which is neither elegant nor fast as it needs to look through all elements of the array even if the item to find is in the very beginning.
- A faster way to do it is via accelerating
next((i for i, v in np.ndenumerate(a) if v==x), -1)with Numba⁷ (otherwise it’s way slower in the worst case than
- Once the array is sorted though, the situation gets better:
v = np.searchsorted(a, x); return v if a[v]==x else -1is really fast with O(log N) complexity, but it requires O(N log N) time to sort first.
Actually, it is not a problem to speed up searching by implementing it in C. The problem is float comparisons. This is a task that simply does not work out of the box for arbitrary data.
np.allclose(a, b) compares arrays of floats with a given tolerance
np.allcloseassumes all the compared numbers to be of a typical scale of 1. For example, if you work with nanoseconds, you need to divide the default
atolargument value by 1e9:
np.allclose(1e-9, 2e-9, atol=1e-17) == False.
math.isclosemakes no assumptions about the numbers to be compared but relies on a user to give a reasonable
abs_tolvalue instead (taking the default
atolvalue of 1e-8 is good enough for numbers with a typical scale of 1):
Aside from that,
np.allclose has some minor issues in a formula for absolute and relative tolerances, for example, for certain a,b
allclose(a, b) != allclose(b, a). Those issues were resolved in the (scalar) function
math.isclose (which was introduced later). To learn more on that, take a look at the excellent floating-point guide⁸ and the corresponding NumPy issue⁹ on GitHub.
2. Matrices, the 2D Arrays
There used to be a dedicated
matrix class in NumPy, but it is deprecated now, so I’ll use the words matrix and 2D array interchangeably.
Matrix initialization syntax is similar to that of vectors:
Double parentheses are necessary here because the second positional parameter is reserved for the (optional)
dtype (which also accepts integers).
Random matrix generation is also similar to that of vectors:
Two-dimensional indexing syntax is more convenient than that of nested lists:
The “view” sign means that no copying is actually done when slicing an array. When the array is modified, the changes are reflected in the slice as well.
The axis argument
In many operations (e.g.,
sum) you need to tell NumPy if you want to operate across rows or columns. To have a universal notation that works for an arbitrary number of dimensions, NumPy introduces a notion of axis: The value of the
axis argument is, as a matter of fact, the number of the index in question: The first index is
axis=0, the second one is
axis=1, and so on. So in 2D
axis=0 is column-wise and
axis=1 means row-wise.
Generally speaking, the first dimension
i (axis=0) is responsible for indexing the rows, so
sum(axis=0) should be read like “for any given column sum over all of its rows” rather than just “column-wise”. The 2D case is somewhat counter-intuitive: you need to specify the dimension to be eliminated, instead of the remaining one you would normally think about. In higher dimensional cases this is more natural, though: it’d be a burden to enumerate all the remaining dimensions if you only need to sum over a single one.
The notion of the ‘axis’ argument (called ‘dimension’ in MATLAB) is actually very simple: it is just number of the index you want the operation to carry over. It is used in many functions in NumPy so it pays off to understand it, but if for some reason it doesn’t work out for you, in this particular case you can use the Einstein summation — that’s a way to do the sums without ‘axes’:
np.einsum(‘ij->j’, a) is summing column-wise (same as
sum(axis=0) in the image above) and
np.einsum(‘ij->i’, a) is row-wise (same as
In addition to ordinary operators (like +,-,*,/,// and **) which work element-wise, there’s a @ operator that calculates a matrix product:
As a generalization of broadcasting from scalar that we’ve seen already in the first part, NumPy allows mixed operations between a vector and a matrix, and even between two vectors:
Note that in the last example it is a symmetric per-element multiplication. To calculate the outer product using an asymmetric linear algebra matrix multiplication the order of the operands should be reversed:
Row vectors and column vectors
As seen from the example above, in the 2D context, the row and column vectors are treated differently. This contrasts with the usual NumPy practice of having one type of 1D arrays wherever possible (e.g.,
a[:,j] — the j-th column of a 2D array
a— is a 1D array). By default 1D arrays are treated as row vectors in 2D operations, so when multiplying a matrix by a row vector, you can use either shape (n,) or (1, n) — the result will be the same. If you need a column vector, there are a couple of ways to cook it from a 1D array, but surprisingly
transpose is not one of them:
Two operations that are capable of making a 2D column vector out of a 1D array are reshaping and indexing with
Here the -1 argument tells
reshape to calculate one of the dimension sizes automatically and
None in the square brackets serves as a shortcut for
np.newaxis, which adds an empty axis at the designated place.
So, there’s a total of three types of vectors in NumPy: 1D arrays, 2D row vectors, and 2D column vectors. Here’s a diagram of explicit conversions between those:
By the rules of broadcasting, 1D arrays are implicitly interpreted as 2D row vectors, so it is generally not necessary to convert between those two — thus the corresponding area is shaded.
Strictly speaking, any array, all but one dimensions of which are single-sized, is a vector (eg.
a.shape==[1,1,1,5,1,1]), so there’s an infinite number of vector types in numpy, but only these three are commonly used. You can use
np.reshape to convert a ‘normal’ 1D vector to this form and
np.squeeze to get it back. Both functions act as views.
There are two main functions for joining the arrays:
Those two work fine with stacking matrices only or vectors only, but when it comes to mixed stacking of 1D arrays and matrices, only the
vstack works as expected: The
hstack generates a dimensions-mismatch error because as described above, the 1D array is interpreted as a row vector, not a column vector. The workaround is either to convert it to a row vector or to use a specialized
column_stack function which does it automatically:
The inverse of stacking is splitting:
Specific columns and rows can be
deleted like that:
The inverse operation is
append function, just like
hstack, is unable to automatically transpose 1D arrays, so once again, either the vector needs to be reshaped or a dimension added, or
column_stack needs to be used instead:
Actually, if all you need to do is add constant values to the border(s) of the array, the (slightly overcomplicated)
pad function should suffice:
The broadcasting rules make it simpler to work with meshgrids. Suppose, you need the following matrix (but of a very large size):
Two obvious approaches are slow, as they use Python loops. The MATLAB way of dealing with such problems is to create a meshgrid:
meshgrid function accepts an arbitrary set of indices,
mgrid — just slices and
indices can only generate the complete index ranges.
fromfunction calls the provided function just once, with the I and J argument as described above.
But actually, there is a better way to do it in NumPy. There’s no need to spend memory on the whole I and J matrices (even though
meshgrid is smart enough to only store references to the original vectors if possible). It is sufficient to store only vectors of the correct shape, and the broadcasting rules take care of the rest:
meshgrid will change the order of the arguments:
J, I= np.meshgrid(j, i) — it is an ‘xy’ mode, useful for visualizing 3D plots (see the example from the docs).
Aside from initializing functions over a two- or three-dimensional grid, the meshes can be useful for indexing arrays:
argmax functions in 2D and above have an annoyance of returning the flattened index (of the first instance of the min and max value). To convert it to two coordinates, an
unravel_index function is required:
As helpful as the
axis argument is for the functions listed above, it is as unhelpful for the 2D sorting:
It is just not what you would usually want from sorting a matrix or a spreadsheet:
axis is in no way a replacement for the
key argument. But luckily, NumPy has several helper functions which allow sorting by a column — or by several columns, if required:
a[a[:,0].argsort()] sorts the array by the first column:
argsort returns an array of indices of the original array after sorting.
This trick can be repeated, but care must be taken so that the next sort does not mess up the results of the previous one:
a = a[a[:,2].argsort()]
a = a[a[:,1].argsort(kind='stable')]
a = a[a[:,0].argsort(kind='stable')]
2. There’s a helper function
lexsort which sorts in the way described above by all available columns, but it always performs row-wise, and the order of rows to be sorted is inverted (i.e., from bottom to top) so its usage is a bit contrived, e.g.
a[np.lexsort(np.rot90(a))] sorts by all columns in left-to-right order.
a[np.lexsort(np.rot90(a[:,[2,5]]))] sorts by column 2 first and then (where the values in column 2 are equal) by column 5. An alternative form is
flipud flips the matrix in the up-down direction (to be precise, in the
axis=0 direction, same as
a[::-1,...], where three dots mean “all other dimensions’”— so it’s all of a sudden
fliplr, that flips the 1D arrays) and
rot90 rotates the matrix in the positive direction (i.e., counter-clockwise) by 90 degrees. Both return views.
3. There also is an
order argument to
sort, but it is neither fast nor easy to use if you start with an ordinary (unstructured) array.
4. It might be a better option to do it in pandas since this particular operation is way more readable and less error-prone there:
pd.DataFrame(a).sort_values(by=[2,5]).to_numpy() sorts by column 2, then by column 5.
pd.DataFrame(a).sort_values().to_numpy() sorts by all columns in the left-to-right order.
3. 3D and Above
When you create a 3D array by reshaping a 1D vector or converting a nested Python list, the meaning of the indices is (z,y,x). The first index is the number of the plane, then the coordinates go in that plane:
This index order is convenient, for example, for keeping a bunch of gray-scale images:
a[i] is a shortcut for referencing the i-th image.
But this index order is not universal. When working with RGB images, the (y,x,z) order is usually used: First are two pixel coordinates, and the last one is the color coordinate (RGB in Matplotlib, BGR in OpenCV):
This way it is convenient to reference a specific pixel:
a[i,j] gives an
(r,g,b) tuple of the
So the actual command to create a certain geometrical shape depends on the conventions of the domain you’re working on:
Obviously, NumPy functions like
dstack are not aware of those conventions. The index order hardcoded in them is (y,x,z), the RGB images order:
If your data is laid out differently, it is more convenient to stack images using the
concatenate command, feeding it the explicit index number in an
If thinking in terms of axes numbers is not convenient to you, you can convert the array to the form that is hardcoded into
hstack and co:
This conversion is cheap: No actual copying takes place; it just mixes the order of indices on the fly.
Another operation that mixes the order of indices is array transposition. Examining it might make you feel more familiar with 3D arrays. Depending on the order of axes you’ve decided on, the actual command to transpose all planes of your array will be different: For generic arrays, it is swapping indices 1 and 2, for RGB images it is 0 and 1:
Interesting though that the default
axes argument for
transpose (as well as the only
a.T operation mode) reverses the index order, which coincides with neither of the index order conventions described above.
Broadcasting works with higher dimensions as well, see my note “Broadcasting in NumPy” for details.
And finally, here’s a function that can save you a lot of Python loops when dealing with the multidimensional arrays and can make your code cleaner —
einsum (Einstein summation):
It sums the arrays along the repeated indices. In this particular example
np.tensordot(a, b, axis=1) would suffice in both cases, but in the more complex cases
einsum might work faster and is generally easier to both write and read — once you understand the logic behind it.
- Scott Sievert, NumPy GPU acceleration
- Jay Alammar, A Visual Intro to NumPy and Data Representation
- Big-O Cheat Sheet site
- Python Time Complexity wiki page
- NumPy Issue #14989, Reverse param in ordering functions
- NumPy Issue #2269, First nonzero element
- Numba library homepage
- The Floating-Point Guide, Comparison
- NumPy Issue #10161, numpy.isclose vs math.isclose
- 100 NumPy exercises on GitHub