umair-akbar-1*XbZqTBXZmnQtVZ llbGJTw - NumPy Illustrated: The Visual Guide to NumPy
Image credit: Author
Brush up your NumPy or learn it from scratch

NumPy is a fundamental library that most of the widely used Python data processing libraries are built upon (pandasOpenCV), inspired by (PyTorch), or can efficiently share data with (TensorFlowKeras, 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:

  1. Vectors, the 1D Arrays
  2. Matrices, the 2D Arrays
  3. 3D and above

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:

umair-akbar-1*ND8LvMjQOX19G Yg0ANPxw - NumPy Illustrated: The Visual Guide to NumPy

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
umair-akbar-1*D I8hK4WXC8wtpR5tvR0fw - NumPy Illustrated: The Visual Guide to NumPy

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

Vector initialization

One way to create a NumPy array is to convert a Python list. The type will be auto-deduced from the list element types:

umair-akbar-1*oQRt7v zDO7DNRLRUlCm5w - NumPy Illustrated: The Visual Guide to NumPy

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 np.zeros or np.empty:

umair-akbar-1*  CdjlwqtzCj98i2oCr5Tw - NumPy Illustrated: The Visual Guide to NumPy

It is often necessary to create an empty array which matches the existing one by shape and elements type:

umair-akbar-1*WXny4Z0frysfsVllO84Jyw - NumPy Illustrated: The Visual Guide to NumPy

Actually, all the functions that create an array filled with a constant value have a _like counterpart:

umair-akbar-1*cyN FxUVbkdDyrULhfTIGw - NumPy Illustrated: The Visual Guide to NumPy

There are as many as two functions for array initialization with a monotonic sequence in NumPy:

umair-akbar-1*NNCWBPDCPfLtNYDc ul5SA - NumPy Illustrated: The Visual Guide to NumPy

If you need a similar-looking array of floats, like [0., 1., 2.], you can change the type of the arange output: 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.

But arange is not especially good at handling floats:

umair-akbar-1*ps2PYTthKUdNbzynKjW4oQ - NumPy Illustrated: The Visual Guide to NumPy

This 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:

umair-akbar-1*oFoQ5Vw4s9mx7RyUi1wYOA - NumPy Illustrated: The Visual Guide to NumPy

Vector indexing

Once you have your data in the array, NumPy is brilliant at providing easy ways of giving it back:

umair-akbar-1*4xpufyWZWcIbabsOHVlc4g - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*Dv9HzG2dSf3QuhjfpSIhqA - NumPy Illustrated: The Visual Guide to NumPy

Another super-useful way of getting data from NumPy arrays is boolean indexing, which allows using all kinds of logical operators:

umair-akbar-1*nFGcXav xxD7TXGiRYMpHg - NumPy Illustrated: The Visual Guide to NumPy
any and all act just like their python peers, but don’t short-circuit

Careful though; Python “ternary” comparisons like 3<=a<=5 don’t work here.

As seen above, boolean indexing is also writable. Two common use cases of it spun off as dedicated functions: the excessively overloaded np.where function (see both meanings below) and np.clip.

umair-akbar-1*UH2blpvIaFhkSYvkROcnhw - NumPy Illustrated: The Visual Guide to NumPy

Vector operations

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:

umair-akbar-1*RNfQubSwH 6 GnWHVjn9CQ - NumPy Illustrated: The Visual Guide to NumPy
As usual in Python, a//b means a div b (quotient from division), x**n means xⁿ

The same way ints are promoted to floats when adding or subtracting, scalars are promoted (aka broadcasted) to arrays:

umair-akbar-1*VyadDu7CyuF5 7rKVoFSrw - NumPy Illustrated: The Visual Guide to NumPy

Most of the math functions have NumPy counterparts that can handle vectors:

umair-akbar-1*R1Py e56Tss5yAQog49RTQ - NumPy Illustrated: The Visual Guide to NumPy

Scalar product has an operator of its own:

umair-akbar-1*GCRJripd ZE30Pefwq1RKQ - NumPy Illustrated: The Visual Guide to NumPy

You don’t need loops for trigonometry either:

umair-akbar-1*8Hd09fz8kZzWQvYMxA eGw - NumPy Illustrated: The Visual Guide to NumPy

Arrays can be rounded as a whole:

umair-akbar-1*GkjXtrFriVBXTQove8vEPw - NumPy Illustrated: The Visual Guide to NumPy
floor rounds to -∞, ceil to +∞ and around — to the nearest integer (.5 to even)

The name 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:

umair-akbar-1*Rl3CrJLU1O8Ocr9QASAGVQ - NumPy Illustrated: The Visual Guide to NumPy
Each of these functions has a nan-resistant variant: eg nansumnanmax, etc

Sorting functions have fewer capabilities than their Python counterparts:

umair-akbar-1*mmS5 eCb8onu7MYjjdYpFw - NumPy Illustrated: The Visual Guide to NumPy

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.

umair-akbar-1*bQXR BLBC91f5WvH4TaFcw - NumPy Illustrated: The Visual Guide to NumPy
The square brackets in the definition of index() mean that either j or both i and j can be omitted
  • One way of finding an element is np.where(a==x)[0][0], 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[0] for i, v in np.ndenumerate(a) if v==x), -1) with Numba⁷ (otherwise it’s way slower in the worst case than where).
  • Once the array is sorted though, the situation gets better:v = np.searchsorted(a, x); return v if a[v]==x else -1 is 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.

Comparing floats

The function np.allclose(a, b) compares arrays of floats with a given tolerance

umair-akbar-1*LZWByfm3vyHXhNreHGy3hQ - NumPy Illustrated: The Visual Guide to NumPy
There is no silver bullet!
  • np.allclose assumes 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 atol argument value by 1e9: np.allclose(1e-9, 2e-9, atol=1e-17) == False.
  • math.isclose makes no assumptions about the numbers to be compared but relies on a user to give a reasonable abs_tol value instead (taking the default np.allclose atol value of 1e-8 is good enough for numbers with a typical scale of 1): math.isclose(0.1+0.2–0.3, abs_tol=1e-8)==True.

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:

umair-akbar-1*aLMuXA81pDXaw0J0QdKvRQ - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*tf4GsJUwjGTNHhPDpj5NAw - NumPy Illustrated: The Visual Guide to NumPy

Two-dimensional indexing syntax is more convenient than that of nested lists:

umair-akbar-1*brbsl7QFZGWfmvgFHMwt9Q - NumPy Illustrated: The Visual Guide to NumPy

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 andaxis=1 means row-wise.

umair-akbar-1*jmXqsVUNaBaUsBAkHgqb3A - NumPy Illustrated: The Visual Guide to NumPy

Matrix arithmetic

In addition to ordinary operators (like +,-,*,/,// and **) which work element-wise, there’s a @ operator that calculates a matrix product:

umair-akbar-1*z0vpti4m5rnRdG7tjFBAAA - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*qbi kCjxBE9Wd6Y0cKHfaQ - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*AxPMBAD67SbeXV4wEdcamQ - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*PYk9Ukt2OpnWXrlYqjXEBw - NumPy Illustrated: The Visual Guide to NumPy

Two operations that are capable of making a 2D column vector out of a 1D array are reshaping and indexing with newaxis:

umair-akbar-1*ettWjwGl17oyxeSvAHtbCQ - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*DjVBs4arPuSLIfFIBO5Dfw - NumPy Illustrated: The Visual Guide to NumPy

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.

Matrix manipulations

There are two main functions for joining the arrays:

umair-akbar-1*SJbmp5 qJixVztVMJQn9sw - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*ZUFxtl3Z2cdDZhcLUNpE0Q - NumPy Illustrated: The Visual Guide to NumPy

The inverse of stacking is splitting:

umair-akbar-1*ADtgaXz1wxze7c6K6kK9CQ - NumPy Illustrated: The Visual Guide to NumPy

Matrix replication can be done in two ways: tile acts like copy-pasting and repeat like collated printing:

umair-akbar-1*kCClP8kiegtYZNyEdsCPbw - NumPy Illustrated: The Visual Guide to NumPy

Specific columns and rows can be deleted like that:

umair-akbar-1*OSCZ2ADLw YN S6reXOiMQ - NumPy Illustrated: The Visual Guide to NumPy

The inverse operation is insert:

umair-akbar-1*IEVGMjetDiL27pSOAO7FuA - NumPy Illustrated: The Visual Guide to NumPy

The 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:

umair-akbar-1*ji2JnZUnVcG6IHFSl0FWFg - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*B2rBWsT iRbX 9kbjnnePw - NumPy Illustrated: The Visual Guide to NumPy


The broadcasting rules make it simpler to work with meshgrids. Suppose, you need the following matrix (but of a very large size):

umair-akbar-1*WdYuQI6m7C6oTAdYNPLoyQ - NumPy Illustrated: The Visual Guide to NumPy

Two obvious approaches are slow, as they use Python loops. The MATLAB way of dealing with such problems is to create a meshgrid:

umair-akbar-1*1pM5EfQuHuAAH 5VL r9aw - NumPy Illustrated: The Visual Guide to NumPy

The 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:

umair-akbar-1*Af4Z0dCrwxyZRzImmtoeBA - NumPy Illustrated: The Visual Guide to NumPy

Without the indexing=’ij’ argument, 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:

umair-akbar-1*PVIGDRv9TAqVwJV4mm69Aw - NumPy Illustrated: The Visual Guide to NumPy
Works with sparse meshgrids, too

Matrix statistics

Just like sum, all the other stats functions (min/maxargmin/argmaxmean/median/percentilestd/var) accept the axis parameter and act accordingly:

umair-akbar-1*rTNRworY2lLrtnEW92tQbQ - NumPy Illustrated: The Visual Guide to NumPy
np.amin is just an alias of np.min to avoid shadowing the Python min when you write ‘from numpy import *'

The argmin and 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:

umair-akbar-1*bbkABjr18h q5cPQmgkehg - NumPy Illustrated: The Visual Guide to NumPy

The quantifiers all and any are also aware of the axis argument:

umair-akbar-1*y95ZnO1wB9U MFAYf4r4qw - NumPy Illustrated: The Visual Guide to NumPy

Matrix sorting

As helpful as the axis argument is for the functions listed above, it is as unhelpful for the 2D sorting:

umair-akbar-1*Ilp0Mn0vS5ly0NhxUs3bhw - NumPy Illustrated: The Visual Guide to NumPy

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:

1. a[a[:,0].argsort()] sorts the array by the first column:

umair-akbar-1*way2tJEnsXOjtfo1PkTv8w - NumPy Illustrated: The Visual Guide to NumPy

Here 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')]

umair-akbar-1*r98wxRa5nD4uK9KMcV7pEA - NumPy Illustrated: The Visual Guide to NumPy

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.flipud(a[2,5].T))] sorts by column 2 first and then (where the values in column 2 are equal) by column 5.
– a[np.lexsort(np.flipud(a.T))] sorts by all columns in left-to-right order.

umair-akbar-1*Wlv5bukk35rMsNB84X0Y8g - NumPy Illustrated: The Visual Guide to NumPy

Here 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 flipud, not fliplr, that flips the 1D arrays).

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:

umair-akbar-1*Dt1eP hrgy8wu1xoaoa2 w - NumPy Illustrated: The Visual Guide to NumPy

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):

umair-akbar-1*CLQKyU9p2QqARvCImz4zuA - NumPy Illustrated: The Visual Guide to NumPy

This way it is convenient to reference a specific pixel: a[i,j] gives an RGB tuple of the (i,j) pixel.

So the actual command to create a certain geometrical shape depends on the conventions of the domain you’re working on:

umair-akbar-1*axngp0hJbvpQldg5kqr78A - NumPy Illustrated: The Visual Guide to NumPy

Obviously, NumPy functions like hstackvstack, or dstack are not aware of those conventions. The index order hardcoded in them is (y,x,z), the RGB images order:

umair-akbar-1*c XjRbJHN26dMw1dcRS6qQ - NumPy Illustrated: The Visual Guide to NumPy
Stacking RGB images (only two colors here)

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 axis argument:

umair-akbar-1*bUfuXiTZOas6ZSD lOSbXA - NumPy Illustrated: The Visual Guide to NumPy
Stacking generic 3D arrays

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:

umair-akbar-1*5yGMnpHwEo8XeJET U5CMA - NumPy Illustrated: The Visual Guide to NumPy

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:

umair-akbar-1*EJiDVb9xb8VL9XmFDevSGg - NumPy Illustrated: The Visual Guide to NumPy

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.

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):

umair-akbar-1*uGKflexdelNxiireoI8YJw - NumPy Illustrated: The Visual Guide to NumPy

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.

If you want to test your NumPy skills, there’s a tricky crowd-sourced set of 100 NumPy exercises¹⁰ on GitHub.

Let me know (on reddithackernews, etc) if I have missed your favorite NumPy feature, and I’ll try to fit it in!


  1. Scott Sievert, NumPy GPU acceleration
  2. Jay Alammar, A Visual Intro to NumPy and Data Representation
  3. Big-O Cheat Sheet site
  4. Python Time Complexity wiki page
  5. NumPy Issue #14989, Reverse param in ordering functions
  6. NumPy Issue #2269, First nonzero element
  7. Numba library homepage
  8. The Floating-Point Guide, Comparison
  9. NumPy Issue #10161, numpy.isclose vs math.isclose
  10. 100 NumPy exercises on GitHub

About the Author


Umair Akbar | Cloud Engineer

Umair Akbar is a Senior Information Security Engineer with over 5 years of experience leading the development and daily management of InfoSec systems.

View All Articles