2.6 Aggregation and Array Methods

: 15 minutes

An aggregation, or reduction, collapses many numbers into fewer. Where a ufunc preserves shape, a reduction deliberately destroys an axis—and the whole skill of this section is controlling which axis it destroys.

Some of the most commonly used statistical methods are listed below.

Mathematical and Statistical Methods (McKinney 2017)

Method	Description
sum	Sum of the elements, over the whole array or along an axis; an empty sum is 0
mean	Arithmetic mean; returns nan (with a warning) on a zero-length axis
std, var	Standard deviation and variance
min, max	Minimum and maximum
argmin, argmax	Index of the minimum or maximum, not the value itself
any, all	Whether any / every element is truthy—most useful on boolean arrays
cumsum	Running totals: element i of the result is the sum of the first i+1 inputs
cumprod	Running products, formed the same way

McKinney, Wes. 2017. Python for Data Analysis. 2nd ed. O’Reilly Media. https://www.oreilly.com/library/view/python-for-data/9781491957653/.

Note that cumsum and cumprod are the odd ones out: they are not reductions. They return an array of the same shape as the input, since each position keeps its own running total.

The axis argument

With no argument, a reduction flattens the entire array and returns a scalar:

Far more often we want a summary along an axis. If A is a matrix with observations in rows and features in columns, A.mean(axis=0) (equivalently numpy.mean(A, axis=0)) averages down the rows and gives one number per column:

ImportantThe axis you name is the axis that disappears

This is the single sentence to memorise for this chapter, and it resolves the axis=0 versus axis=1 question every time without you having to remember anything about rows or columns.

A has shape (5, 4). Naming axis=0 deletes the 5 and leaves (4,). Naming axis=1 deletes the 4 and leaves (5,). Check by looking at the shape, not by squinting at the numbers.

Counter-intuitively, axis=0 gives you the column means. That is not a quirk: you are summing over the row axis, so the row axis is what vanishes.

For higher-order arrays the result has one axis fewer; its shape is the original shape with the named entry removed.

You may also collapse several axes at once by passing a tuple, which is occasionally exactly what you want:

keepdims

A reduction removes an axis, which is a problem the moment you want to combine the summary back with the array it came from—the shapes no longer line up. Passing keepdims=True retains the reduced axis with length 1:

A length-1 axis is precisely what broadcasting knows how to stretch, so the keepdims version slots straight back in:

Without keepdims, A / A.sum(axis=1) would right-align (3,) against (3, 4), compare 3 with 4, and raise a ValueError—which, as errors go, is the lucky outcome. On a square array it would have broadcast silently and given you the wrong answer.

Rule of thumb: if you reduce along an axis and then use the result together with the original array, you want keepdims=True.

argmin and argmax

These return a position, not a value—a distinction that matters because the position is usually what you actually want (“which patient had the highest reading?”). With no axis, the index refers to the flattened array:

To recover the value, index with the result: A.max() and A.ravel()[A.argmax()] agree. To turn a flat index back into a coordinate pair, use np.unravel_index(A.argmax(), A.shape).

any and all

On a boolean array—typically the output of a comparison—any and all reduce exactly like sum does, and take the same axis argument:

(A < 0).sum() counts how many, because True sums as 1; (A < 0).any() asks merely whether there are any. Both are more idiomatic and far faster than a loop.

Other Functions

numpy.where

The general form np.where(cond, arr1, arr2) takes a boolean array cond and, position by position, chooses the corresponding value from arr1 where cond is True and from arr2 where it is False. Either or both of arr1 and arr2 may be a scalar, and all three broadcast against each other. Think of it as a vectorized ternary expression: it selects, it does not skip (see the warning in 2.5).

Suppose you have drawn from the standard normal distribution and regard values beyond \pm 2 as too extreme to keep.

There are two such values here:

Now watch a bug that is easy to write and hard to see. The intent is “pull the extremes back to the boundary”:

Both offending values became +2—including the one that was -2.06, which has just had its sign flipped. The condition tested the magnitude but the replacement forgot it. What we wanted was to move each extreme to the nearer boundary:

This is the archetype of the bugs this module is about: it runs, it produces plausible numbers, and it is wrong. In practice you would not write either version—np.clip(A, -2, 2), below, says exactly what you mean and cannot get the sign wrong.

numpy.sort

A NumPy array arr can be sorted along an axis using numpy.sort(arr), which returns a new sorted array, or arr.sort(), which sorts in place and returns None.

Cautionnp.sort(A) and A.sort() are not interchangeable

B = np.sort(A) leaves A untouched. B = A.sort() destroys the ordering of A and leaves B bound to None—after which the next line fails with a puzzling AttributeError: 'NoneType' object has no attribute ....

The same asymmetry holds throughout NumPy and pandas: a bare method name that returns None is a hint that it worked in place.

The following sorts along the last axis—that is, within each row. This is the default (axis=-1) when no axis is given.

A specific axis can also be supplied. Here axis=0 sorts within each column.

Note that sorting a 2-D array along an axis sorts each row (or column) independently; it does not reorder whole rows. To reorder rows by the values in one column—which is nearly always what you want with real data—sort the indices with argsort and use them to index:

numpy.clip

The numpy.clip() function limits the values in an array to a specified range. It takes an array and a minimum and maximum value as arguments. Any elements in the array that are less than the specified minimum value are replaced by the minimum value, and any elements greater than the specified maximum value are replaced by the maximum value.

numpy.stack

A list of NumPy arrays of the same shape can be stacked along an axis using the numpy.stack function. The function takes a list of arrays to stack and (optionally) an axis of choice.

Let us first consider stacking two 1D arrays or vectors vertically along axis=0, to form a (fat-short) 2D array.

Let us now stack them horizontally along axis=1, to form a (skinny-tall) 2D array.

Next, we stack two matrices along different axes.

Tipstack creates an axis; concatenate extends one

This is the distinction to hold on to, and it is entirely predictable from the shapes.

• np.stack requires all inputs to have identical shapes and adds a new axis. Two (3, 5) arrays give (2, 3, 5), and the order rises from 2 to 3.
• np.concatenate joins along an existing axis and adds no new one. Two (3, 5) arrays concatenated on axis=0 give (6, 5); the order stays at 2. The inputs must agree on every axis except the one being joined.

If the order of your result surprised you, you reached for the wrong one of the two.

Exercises

Note Which axis vanishes?

A has shape (5, 4, 2). Without running anything, give the shape of A.sum(axis=1) and assign that tuple to answer.

answer = (5, 2)

Summing over axis=1 deletes the 4. With keepdims=True you would instead get (5, 1, 2).

Note Rows that sum to one

X is a 4\times 5 table of non-negative counts. Convert each row into proportions, so that every row of the result sums to 1.

The naive X / X.sum(axis=1) raises a ValueError here—work out why from the shapes before you fix it.

answer = X / X.sum(axis=1, keepdims=True)

X.reshape(-1, 1) on the sums works too, but keepdims=True states the intent. Whenever you reduce and then recombine with the original array, reach for it.

Note Which observation is largest overall?

X holds 6 patients in rows and 4 measurements in columns. Find the index of the patient whose measurements sum to the largest total—a single integer between 0 and 5.

Compose two steps: reduce to a per-patient total, then locate the largest.

answer = X.sum(axis=1).argmax()

X.sum(axis=1) deletes the measurement axis and leaves shape (6,); argmax then returns the position of the largest entry. Had you written X.argmax() you would have got a flat index into all 24 entries, which answers a different question entirely.

Note Reading an Aggregation

sales has shape (12, 50, 3), with axes (month, store, product). You evaluate

sales.sum(axis=1)

What is the shape of the result, and what does a single entry of it mean?

Shape (12, 3); entry (m, p) is the total units of product p sold across all 50 stores in month m.Shape (50, 3); entry (s, p) is the total units of product p sold by store s across all months.Shape (12, 50); entry (m, s) is the total units sold by store s in month m.Shape (12, 1, 3); the store axis is retained with length 1.

The rule is: the axis you name is the axis that disappears. Naming axis=1 sums over the stores, so the store axis is removed from the shape and (12, 50, 3) becomes (12, 3).

The remaining axes keep their meaning, so entry (m, p) is the month-m, product-p total, aggregated over all stores.

The other options describe real operations, just different ones:

• (50, 3) would come from axis=0, summing over months.
• (12, 50) would come from axis=2, summing over products.
• (12, 1, 3) is what sales.sum(axis=1, keepdims=True) gives. That length-1 axis is not decoration—it is exactly what lets you write sales / sales.sum(axis=1, keepdims=True) to express each store’s share of its month’s total, which the plain (12, 3) result cannot broadcast into.

The practical habit: after any aggregation, read the surviving shape back as a sentence about the data. “12 months by 3 products” tells you at once that the stores were summed away.

Note argmax Returns What?

Let

A = np.array([[3., 9., 1.],
              [7., 2., 8.]])

What does A.argmax() return?

9.0, the largest value in A.(0, 1), the row and column of the largest value.1, the position of the largest value in the flattened array.array([1, 0, 1]), the winning row index for each column.

argmax returns a position, not a value—that is max’s job. And with no axis argument, NumPy first flattens the array in row-major order, giving [3., 9., 1., 7., 2., 8.]. The largest entry, 9, sits at flat index 1.

The other options each correspond to a different call:

• A.max() returns the value 9.0.
• np.unravel_index(A.argmax(), A.shape) converts the flat index back into the coordinate pair (0, 1).
• A.argmax(axis=0) returns array([1, 0, 1])—for each column, which row holds its maximum.

Two further points worth knowing. Ties are broken by taking the first occurrence. And if the array contains nan, argmax will report the nan’s position, since nan compares as neither smaller nor larger than anything; use np.nanargmax when missing values are possible.