Python NumPy

NumPy is the library behind almost all numerical work in Python: data science, machine learning, image processing and scientific computing. Its core is the array, a grid of numbers of one type that you can calculate with as a whole, without writing loops. NumPy is installed in our editor; the first import numpy sets it up in your browser, which takes a moment.

Creating arrays

import numpy as np

a = np.array([3, 1, 4, 1, 5])
print(a, a.size, a.shape)
print(np.arange(0, 10, 2))
print(np.zeros(3), np.ones(2, dtype=int))
print(np.linspace(0, 1, 5))

[3 1 4 1 5] 5 (5,)
[0 2 4 6 8]
[0. 0. 0.] [1 1]
[0.   0.25 0.5  0.75 1.  ]

np is the usual short name. size is the number of elements and shape the length along each dimension; every element has the same type, shown by a.dtype. arange works like range but makes an array; linspace gives evenly spaced values between two ends, including both. Arrays print without commas.

Maths on whole arrays

Operators work on every element at once, which is both shorter and much faster than a loop:

import numpy as np

prices = np.array([120, 80, 250, 40])
qty = np.array([2, 5, 1, 10])
print(prices * qty)
print((prices * qty).sum())
print(prices * 1.18)
print(prices > 100)
print(prices[prices > 100])

[240 400 250 400]
1290
[141.6  94.4 295.   47.2]
[ True False  True False]
[120 250]

Comparing an array gives an array of booleans, a mask, and indexing with a mask keeps only the elements where it is True. This is how you filter data in NumPy.

Two-dimensional arrays

import numpy as np

marks = np.array([[72, 85, 90],
                  [64, 70, 58],
                  [88, 92, 95]])
print(marks.shape)
print(marks[0], marks[:, 1])
print(marks.mean(axis=0).round(1))
print(marks.sum(axis=1))
print(np.arange(6).reshape(2, 3))

(3, 3)
[72 85 90] [85 70 92]
[74.7 82.3 81. ]
[247 192 275]
[[0 1 2]
 [3 4 5]]

Each row is a student and each column a subject. marks[:, 1] takes column 1 of every row. axis=0 works down the columns (per subject) and axis=1 across the rows (per student). reshape arranges the same numbers in a new shape.

Random numbers and statistics

np.random.default_rng(seed) makes a random generator: rng.integers(1, 7, size=10) rolls ten dice at once, and rng.normal(170, 8, 1000) makes a thousand heights. mean, median, std, min, max and percentile summarise them.

Practice exercises

1. Make an array of the numbers 1 to 10 and print their squares and their sum.
2. Convert an array of temperatures in Celsius to Fahrenheit in one expression.
3. From an array of marks, print only the marks of 60 or more, and how many there are.
4. Make a 3 by 4 array with arange and reshape, and print the sum of each column.