NumPy: Fast Numbers

Ava has a year of daily sales, three hundred and sixty five numbers, and she wants the average, the best day, and prices with tax, for every day. A loop works, but with millions of numbers it is slow. NumPy is the foundation of data science in Python. It stores numbers in compact arrays, and calculates on all of them at once, at the speed of compiled code.

• A year of daily sales
• Loops get slow with big data
• NumPy: whole-array maths, at C speed

Creating arrays

Install numpy with pip, and import it as np, by convention. An array holds numbers of one type, in a grid. Arange makes a sequence, linspace spreads points evenly, and zeros and ones make filled arrays. The shape tells the dimensions, and dtype tells the type.

import numpy as np

a = np.array([4, 9, 14, 40])
print(a, a.dtype, a.shape)
print(np.arange(0, 10, 2))
print(np.linspace(0, 1, 5))
print(np.zeros((2, 3)))

[ 4  9 14 40] int64 (4,)
[0 2 4 6 8]
[0.   0.25 0.5  0.75 1.  ]
[[0. 0. 0.]
 [0. 0. 0.]]

Vectorised maths

Here is the magic. Arithmetic on an array applies to every element, with no loop. Multiply all prices by one point two, and every price gets tax. Compare an array to a number, and you get true and false for every element. Use that mask, to pick the matching items. This is called vectorising, and it is the heart of NumPy.

import numpy as np

prices = np.array([9.5, 7.0, 24.99, 12.0])
print(prices * 1.2)
print(prices + prices)
print(prices > 10)
print(prices[prices > 10])
print(prices.sum(), prices.mean(), prices.max())

[11.4    8.4   29.988 14.4  ]
[19.   14.   49.98 24.  ]
[False False  True  True]
[24.99 12.  ]
53.489999999999995 13.372499999999999 24.99

Two dimensions and statistics

Arrays can have several dimensions, like a table. Here, rows are weeks, and columns are days. The axis argument says in which direction to calculate. Axis zero goes down the rows, and axis one goes across. Reshape changes the layout, and the indexing uses a comma, row and column. Real data analysis is built on these few ideas.

import numpy as np

sales = np.array([[5, 8, 6], [9, 4, 7]])
print(sales.shape)
print(sales.sum(axis=0))
print(sales.sum(axis=1))
print(sales[1, 2], sales[:, 0])
print(np.arange(6).reshape(2, 3))

(2, 3)
[14 12 13]
[19 20]
7 [5 9]
[[0 1 2]
 [3 4 5]]

Speed: loop versus vector

Now the proof. Squaring one million numbers, with a Python loop, takes a good fraction of a second. NumPy does it in a few milliseconds. We measure both, and print how many times faster the vector version is. For big data, the difference is the difference between a coffee break, and an instant answer.

import numpy as np, time

n = 1_000_000
x = np.arange(n)
t = time.perf_counter(); a = [i * i for i in range(n)]; loop = time.perf_counter() - t
t = time.perf_counter(); b = x * x; vec = time.perf_counter() - t
print("same result:", a[-1] == b[-1])
print("numpy is faster:", loop / vec > 5)

same result: True
numpy is faster: True

Prices with a discount

Your turn. Make an array of four prices. Apply a ten percent discount to all of them, in one line, and print the new prices and the total. Ready? Here is one solution. Multiply the array by point nine, and call sum.

import numpy as np

prices = np.array([10, 20, 30, 40])
new = prices * 0.9
print(new)
print(new.sum())

[ 9. 18. 27. 36.]
90.0

Summary