Skip to main content
Back to module
Scientific Python · NumPy

Arrays and vectorisation

Reading1h 30mLesson 1 of 3 in this module

Replace slow Python loops with NumPy array math, and use broadcasting to combine arrays of different shapes.

Why this matters

Scientific and engineering data comes in bulk: 10,000 thermocouple readings, a 2,048-point spectrum, a strain gauge sampled at 1 kHz for an hour. Plain Python lists and for loops can handle that, but they are slow and wordy. NumPy gives you one object, the ndarray, that holds numbers of one type in a single block of memory and runs whole-array operations in compiled code. The result is code that is shorter, closer to the maths on paper, and often 10 to 100 times faster.

Almost every scientific Python library (SciPy, pandas, matplotlib, scikit-learn) is built on NumPy arrays, so this is the foundation for the rest of the course.

Core idea: an array is a grid of one type

import numpy as np

t = np.array([0.0, 0.5, 1.0, 1.5, 2.0])   # from a list
print(t.shape, t.dtype)                     # (5,) float64

x = np.linspace(0, 10, 101)   # 101 evenly spaced points from 0 to 10 inclusive
k = np.arange(0, 1, 0.25)     # [0.  0.25 0.5  0.75] (end excluded)
grid = np.zeros((3, 4))       # 3 rows, 4 columns of 0.0

Key attributes:

  • shape: the size along each axis, such as (3, 4).
  • ndim: the number of axes.
  • dtype: the element type, usually float64 for measurements or int64 for counts.

Vectorisation: operate on the whole array

Arithmetic on arrays works element by element, with no loop in your code:

celsius = np.array([20.0, 25.5, 30.1, 18.9])

# Loop version
kelvin_loop = []
for c in celsius:
    kelvin_loop.append(c + 273.15)

# Vectorised version
kelvin = celsius + 273.15

NumPy's "universal functions" (ufuncs) such as np.sqrt, np.exp, np.sin and np.log also work on every element. Reductions collapse an array to a summary: x.sum(), x.mean(), x.std(), x.min(), x.max(), np.argmax(x).

Worked example: RMS voltage

The root-mean-square of a signal is sqrt(mean(v^2)). For a sine wave of amplitude 10 V, it should be 10 / sqrt(2) ≈ 7.071 V.

t = np.linspace(0, 1, 10_001)               # one second
v = 10 * np.sin(2 * np.pi * 50 * t)         # 50 Hz, 10 V peak
v_rms = np.sqrt(np.mean(v ** 2))
print(round(v_rms, 3))                       # 7.071

Three lines, no loop, and each line reads like the formula.

Boolean masks: select by condition

Comparisons produce arrays of True/False. Use them to filter or count:

readings = np.array([4.98, 5.02, 7.40, 5.01, 4.99, -1.00])
valid = (readings > 0) & (readings < 6)     # use & and |, not "and"/"or"
print(readings[valid])                       # [4.98 5.02 5.01 4.99]
print(valid.sum())                           # 4 valid readings

Brackets matter: (readings > 0) & (readings < 6). Without them Python's operator precedence evaluates 0 & readings first.

Broadcasting: combine different shapes

When two arrays have different shapes, NumPy tries to broadcast them: it lines the shapes up from the right, and any axis of size 1 is stretched to match the other array. No data is copied.

data = np.array([[1.0, 2.0, 3.0],
                 [4.0, 5.0, 6.0]])           # shape (2, 3)
col_means = data.mean(axis=0)                # shape (3,)
centred = data - col_means                   # (2, 3) - (3,) -> (2, 3)

axis=0 means "collapse the rows", giving one value per column. To subtract row means you need the shape (2, 1), which keepdims=True or [:, None] gives you:

row_means = data.mean(axis=1, keepdims=True) # shape (2, 1)
centred_rows = data - row_means

Worked example: all pairwise distances

Given n points in 2D (shape (n, 2)), the distance matrix D[i, j] can be built with no loops:

pts = np.array([[0.0, 0.0], [3.0, 4.0], [6.0, 8.0]])
diff = pts[:, None, :] - pts[None, :, :]     # (3,1,2) - (1,3,2) -> (3,3,2)
D = np.sqrt((diff ** 2).sum(axis=-1))        # (3,3)
print(D[0, 1], D[0, 2])                      # 5.0 10.0

Inserting a new axis with None (also called np.newaxis) is the most common broadcasting trick.

Common mistakes

  • Using lists where you meant arrays. [1, 2] * 2 is [1, 2, 1, 2]; np.array([1, 2]) * 2 is [2, 4].
  • Integer arrays. np.array([1, 2, 3]) is int64. Assigning 0.5 into it silently truncates to 0. Create float arrays (np.array([1.0, 2, 3]) or dtype=float) for measurements.
  • Views, not copies. Slicing (a[2:5]) returns a view; changing it changes the original. Call .copy() when you need an independent array.
  • Shape surprises. A (3,) array and a (3, 1) array broadcast to (3, 3). Print .shape whenever a result looks too big.
  • Writing a loop and wrapping it in np.vectorize. It is convenient but runs a Python loop underneath; it is not real vectorisation.

Quick recap

  • An ndarray holds one dtype in a fixed shape; arithmetic and ufuncs act on every element.
  • Reductions (sum, mean, std, max) take an axis argument.
  • Boolean masks filter and count without loops.
  • Broadcasting aligns shapes from the right and stretches size-1 axes; [:, None] adds an axis.

Try it

Open the Do tab and replace three loop-style functions with vectorised NumPy versions: a unit conversion, an RMS calculation and a pairwise distance matrix.

Also in this lesson: