Vectorised Operations & Aggregation

In the last two lessons you created arrays and picked values out of them. Now you will calculate with them. NumPy lets you write maths for a whole array at once, which is called a vectorised operation, and then squeeze the result into a summary such as a total or an average, which is called an aggregation.

These two ideas are the heart of all fast data work in Python. pandas uses them under the hood, so learning them here makes pandas much easier.

Element-wise arithmetic

When two arrays have the same shape, +, -, *, /, //, % and ** work on matching elements. Comparisons such as > and == work the same way and return a boolean array.

import numpy as np

quantity = np.array([2, 1, 5, 3])
price = np.array([499.0, 1299.0, 99.0, 250.0])
discount = np.array([0.10, 0.0, 0.05, 0.20])     # 10%, 0%, 5%, 20%

amount = quantity * price * (1 - discount)
print(amount)
print("Bill total: ₹", amount.sum())
print(amount > 500)
print("Orders above ₹500:", (amount > 500).sum())

[ 898.2  1299.    470.25  600.  ]
Bill total: ₹ 3267.45
[ True  True False  True]
Orders above ₹500: 3

Each line runs a loop in fast compiled code without you writing for. The expression quantity * price * (1 - discount) reads like the business rule itself, which also makes vectorised code easier to check.

Universal functions (ufuncs)

NumPy's maths functions, such as np.sqrt(), are universal functions or ufuncs: they take an array and apply the function to every element. Some useful ones:

import numpy as np

change = np.array([-2.5, 1.25, 0.0, -0.75, 3.1])     # daily share price change, ₹
print(np.abs(change))
print(np.round(change, 1))

marks = np.array([34, 104, 77, -3, 88])              # data entry errors included
print(np.clip(marks, 0, 100))

bonus = np.array([5, 0, 10, 0, 2])
print(np.maximum(marks + bonus, 40))                 # nobody below 40 after grace marks

grade = np.where(marks >= 60, "First", "Other")
print(grade)

[2.5  1.25 0.   0.75 3.1 ]
[-2.5  1.2  0.  -0.8  3.1]
[ 34 100  77   0  88]
[ 40 104  87  40  90]
['Other' 'First' 'First' 'Other' 'First']

np.where(condition, a, b) is the vectorised version of if ... else: where the condition is True it takes the value from a, otherwise from b. You will use it constantly in pandas too.

Why vectorised code is faster

A Python for loop has to check the type of every value and call a function for every step. NumPy checks the type once and then runs a tight loop in C. The difference grows with the data.

import time
import numpy as np

prices = np.arange(1_000_000, dtype=float)

start = time.perf_counter()
with_tax_loop = [p * 1.18 for p in prices]
loop_time = time.perf_counter() - start

start = time.perf_counter()
with_tax = prices * 1.18
vector_time = time.perf_counter() - start

print("Same result:", np.allclose(with_tax_loop, with_tax))
print("Vectorised was faster:", vector_time < loop_time)

Same result: True
Vectorised was faster: True

The exact times depend on your device, which is why the example prints only the comparison. Run it and print loop_time and vector_time yourself: the vectorised version is usually tens of times faster. np.allclose() checks that two sets of decimals are equal apart from tiny rounding differences.

Aggregation and the axis argument

An aggregation reduces many values to one: sum(), mean(), min(), max(), std() (standard deviation, the typical distance from the mean), argmin() and argmax() (positions of the smallest and largest value). For 2-D arrays, the axis argument decides the direction:

• axis=0 works down the rows and gives one result per column.
• axis=1 works across the columns and gives one result per row.

import numpy as np

# rows = 3 stores (Rohtak, Jaipur, Pune), columns = Jan, Feb, Mar, Apr sales in ₹ thousand
sales = np.array([
    [120, 135, 150, 110],
    [200, 180, 210, 230],
    [160, 170, 155, 190],
])

print("Grand total:", sales.sum())
print("Per month:  ", sales.sum(axis=0))
print("Per store:  ", sales.sum(axis=1))
print("Store means:", sales.mean(axis=1).round(1))
print("Best month for each store:", sales.argmax(axis=1))
print("Running total, Rohtak:", sales[0].cumsum())

Grand total: 2010
Per month:   [480 485 515 530]
Per store:   [515 820 675]
Store means: [128.8 205.  168.8]
Best month for each store: [2 3 3]
Running total, Rohtak: [120 255 405 515]

A good way to remember: the axis you name is the one that disappears. The array has shape (3, 4); summing over axis=0 removes the 3 rows and leaves 4 values, one per month.

Missing values and sorting

Real data has gaps. NumPy marks a missing decimal value with np.nan (not a number), and any normal aggregation that touches a nan returns nan. The nan-functions skip them instead. This example reads real columns from weather.csv, which has a few empty cells.

import numpy as np

data = np.genfromtxt("weather.csv", delimiter=",", skip_header=1, usecols=(2, 4))
temp_max = data[:, 0]
rain = data[:, 1]

print(data.shape)
print("Missing rain values:", np.isnan(rain).sum())
print("Plain sum:", rain.sum())
print("nansum:   ", np.nansum(rain).round(1))
print("Median max temp:", np.nanmedian(temp_max))

top5 = np.sort(temp_max[~np.isnan(temp_max)])[-5:]
print("Five highest readings:", top5)

(1825, 2)
Missing rain values: 12
Plain sum: nan
nansum:    7524.4
Median max temp: 31.7
Five highest readings: [42.5 42.9 43.  99.9 99.9]

np.genfromtxt() turned the empty cells into nan. The plain sum() is ruined by them, while np.nansum() gives a real total. Sorting the highest temperatures shows something odd: two readings of 99.9 °C. No Indian city has ever been that hot, so these are sensor errors. You will learn how to catch such values in Finding & Handling Outliers.

Common mistakes

• Mixing up the axis. axis=0 gives one value per column, axis=1 one value per row. Check the shape of the result.
• Ignoring nan. One missing value turns sum() or mean() into nan. Use np.nansum(), np.nanmean() and friends, or clean the data first.
• Using Python's built-in sum() on a 2-D array. It adds the rows together instead of giving one number. Use array.sum().
• Writing loops out of habit. If you are looping over array elements to do maths, there is almost always a vectorised way.

Key points

• Arithmetic and comparisons on arrays work element by element, with no loop needed.
• Ufuncs such as np.sqrt(), np.round(), np.clip() and np.maximum() apply to every element.
• np.where(condition, a, b) is a vectorised if-else.
• Aggregations like sum(), mean() and std() take an axis: 0 for per-column, 1 for per-row results.
• nan spreads through normal aggregations; use the nan-functions such as np.nanmean().

import numpy as np

marks = np.array([[45, 60, 72], [80, 85, 90], [33, 41, 38], [66, 70, 58]])
avg = marks.mean(axis=1).round(1)
print("Averages:", avg)
print("Top per test:", marks.max(axis=0))
print(np.where(avg >= 40, "Pass", "Fail"))

Averages: [59.  85.  37.3 64.7]
Top per test: [80 85 90]
['Pass' 'Pass' 'Fail' 'Pass']