NumPy: fast arrays of numbers

Why not just Python lists?

A Python list can hold anything: numbers, text, other lists. That flexibility makes it slow for maths, because Python checks the type of every item in a loop. A NumPy array holds one type of value (for example, 64-bit decimal numbers) packed tightly in memory, and its maths runs in fast compiled code. On a million numbers, a NumPy sum is often 10 to 100 times faster than a Python loop.

Creating arrays

import numpy as np

prices = np.array([120, 250, 99, 480])
print(prices)          # [120 250  99 480]
print(prices.dtype)    # int64 (on most systems)
print(prices.shape)    # (4,)

marks = np.array([[78, 65, 90],
                  [88, 72, 95]])
print(marks.shape)     # (2, 3): 2 rows, 3 columns

Vectorisation: maths on the whole array

Vectorisation means writing an operation once for the whole array instead of looping. prices * 1.18 adds 18% GST to every price. Two arrays of the same shape can be added element by element: the first with the first, the second with the second.

with_gst = prices * 1.18
print(with_gst)        # [141.6  295.   116.82 566.4 ]
print(prices + np.array([10, 10, 10, 10]))   # [130 260 109 490]

Broadcasting

Broadcasting is NumPy's rule for combining arrays of different shapes. A single number is "stretched" to match the array, which is why prices * 1.18 works. A row of 3 numbers can be added to every row of a (2, 3) table. If the shapes cannot be matched, NumPy raises an error such as "operands could not be broadcast together".

bonus = np.array([5, 0, 2])      # extra marks per subject
print(marks + bonus)
# [[83 65 92]
#  [93 72 97]]

Boolean masks: filtering without loops

print(prices > 100)              # [ True  True False  True]
print(prices[prices > 100])      # [120 250 480]

Axis: along rows or down columns

For a 2-D array, axis=0 means "down the rows" (one answer per column) and axis=1 means "across the columns" (one answer per row). marks.mean(axis=0) gives each subject's average; marks.mean(axis=1) gives each student's average.

print(marks.mean(axis=0))   # [83.  68.5 92.5]   per subject
print(marks.mean(axis=1))   # [77.67 85.  ] roughly, per student

Common mistakes

• Writing Python loops over arrays out of habit. Fix: look for a vectorised expression first.
• Using and / or with arrays: (a > 1 and a < 5) raises "truth value of an array is ambiguous". Fix: use & and | with brackets: (a > 1) & (a < 5).
• Mixing up axis=0 and axis=1. Fix: check the .shape of the result.
• Expecting arrays to grow like lists. np.append makes a full copy every time. Fix: build a Python list first, then convert once.