5.1.1 NumPy Basics

What NumPy is for

NumPy stores numbers in arrays and does maths on whole arrays at once, without Python loops. That makes it fast, and almost every other library builds on it: pandas columns, the inputs and outputs of scikit-learn, and images are all NumPy arrays underneath. PyTorch's tensors, in the next module, work almost the same way, so time spent here pays off twice.

You do not need to know much of NumPy for IOAI. This page lists what you do need, shows each idea in a line or two, and points to good places to learn it properly.

What you need for IOAI

• Making arrays, and reading their shape (how many rows and columns) and dtype (the kind of number).
• Maths and comparisons on whole arrays: a + 5, a * b, a > 80.
• Picking values: indexing, slicing, and selecting with a True/False mask.
• Summaries along an axis: sum, mean, max, and argmax (the position of the largest value).
• Changing shape with reshape.
• Broadcasting: how NumPy does maths between arrays of different shapes.
• Matrix multiplication with @.
• Random numbers with a fixed seed, so results can be repeated.

Where to learn it

Start with the official absolute basics for beginners. It covers the whole list above. If you prefer something shorter, Stanford's CS231n NumPy tutorial fits it into one page; if you want more explanation, read chapter 2 of the Python Data Science Handbook. There is also a one-hour video from freeCodeCamp.

For practice, HackerRank's NumPy problems are short and checked instantly, and CodeChef has a NumPy course with problems after each lesson.

The basics in code

Each result below is what a notebook shows for the last line of the cell.

Arrays and maths on whole arrays

import numpy as np

scores = np.array([72, 85, 90, 64])
scores + 5

array([77, 90, 95, 69])

The + 5 applies to every value; no loop needed. The same goes for comparisons, and a comparison gives a mask you can select with:

scores > 80

array([False,  True,  True, False])

scores[scores > 80]

array([85, 90])

Two dimensions, shape and axis

A table of numbers is a 2-D array: one row per sample, one column per feature.

X = np.array([[5.1, 3.5],
              [4.9, 3.0],
              [6.2, 2.9]])
X.shape

(3, 2)

X[1] is the second row, and X[:, 0] is the first column (: means "all rows"):

X[:, 0]

array([5.1, 4.9, 6.2])

Many functions take an axis. axis=0 works down the columns, giving one answer per column; axis=1 works across the rows, giving one answer per row:

X.mean(axis=0)

array([5.4       , 3.13333333])

A common use of axis=1: a model gives a probability for each of three classes, and argmax finds the most likely class for each sample.

probs = np.array([[0.1, 0.7, 0.2],
                  [0.8, 0.1, 0.1]])
probs.argmax(axis=1)

array([1, 0])

Broadcasting

Subtracting the column means from every row works even though the shapes differ, (3, 2) and (2,). NumPy repeats the smaller array across the rows:

(X - X.mean(axis=0)).round(2)

array([[-0.3 ,  0.37],
       [-0.5 , -0.13],
       [ 0.8 , -0.23]])

This is exactly what centring features before a model does. The broadcasting guide has the full rules.

Matrix multiplication

@ multiplies matrices. With one weight per feature, X @ w gives one number per sample, which is how a linear model makes its predictions (see Linear Regression):

w = np.array([0.5, -1.0])
X @ w

array([-0.95, -0.55,  0.2 ])

Reshaping and random numbers

np.arange(6).reshape(2, 3)

array([[0, 1, 2],
       [3, 4, 5]])

rng = np.random.default_rng(0)
rng.integers(0, 10, size=5)

array([8, 6, 5, 2, 3])

With the same seed (0) you get the same numbers every time, which lets you and others repeat an experiment exactly.

From pandas

A DataFrame turns into an array with to_numpy(), as in df.to_numpy(). scikit-learn does this for you, but you will need it when you move to PyTorch.

Common mistakes

When something goes wrong, print .shape first. Most NumPy bugs are shape bugs.

What comes next

• PyTorch Basics: tensors, which are NumPy arrays that can also run on a GPU and work out gradients.