NumPy Aggregation, Sorting and Linear Algebra in Python (Part 4)

NumPy aggregation reduces an array to answers such as totals, means, maxima and positions, either for the whole array or along one axis with axis=0 (per column) or axis=1 (per row). The functions are simple. The defaults are where results go quietly wrong.

One example: the same six numbers give a standard deviation of 1.708 in NumPy and 1.871 in pandas, because the two libraries use different divisors. This final part of the series is for readers who know arrays, indexing and broadcasting. It covers axis reductions, NaN handling, positions, sorting and the linear algebra calls modelling work relies on.

The outputs below come from NumPy 2.2 and pandas 2.3.

How does NumPy aggregation work along an axis?

Pass axis to name the dimension that collapses. With no axis, the reduction covers every element:

import numpy as np

a = np.array([[1, 2, 3],
              [4, 5, 6]])

a.sum()          # 21   - everything
a.sum(axis=0)    # [5 7 9]     - one value per column
a.sum(axis=1)    # [ 6 15]     - one value per row

The mental model that works: the named axis disappears. For example, an array of shape (2, 3) summed with axis=0 loses the first dimension and returns shape (3,).

The same applies to every reduction:

a.mean(axis=0)      # [2.5 3.5 4.5]
a.max(axis=1)       # [3 6]
a.std()             # 1.707825127659933 - population std by default
a.std(ddof=1)       # 1.8708286933869707 - sample std

That ddof default is a genuine trap. NumPy's std uses ddof=0 (population), whereas pandas' std uses ddof=1 (sample). Consequently, the same six numbers give 1.708 in NumPy and 1.871 in pandas, which is a confusing afternoon if you do not know why.

Keep dimensions when you need the result to broadcast back:

row_means = a.mean(axis=1, keepdims=True)   # shape (2, 1) not (2,)
a - row_means                                # [[-1. 0. 1.] [-1. 0. 1.]]

Without keepdims, this is exactly the broadcasting bug from part three.

In a notebook, a single result such as a.sum() displays as np.int64(21). That is how NumPy 2 represents scalars; print() shows 21.

How do you handle NaN in NumPy aggregation?

Use the nan-aware functions such as np.nanmean and np.nansum, because NaN propagates and one missing value contaminates the whole result:

data = np.array([1.0, 2.0, np.nan, 4.0])

data.mean()        # nan
np.nanmean(data)   # 2.3333333333333335

The nan-aware family covers the common cases:

np.nansum(data)      # 7.0
np.nanmean(data)     # 2.3333333333333335
np.nanstd(data)      # 1.247219128924647
np.nanmax(data)      # 4.0
np.nanmedian(data)   # 2.0

Detecting and removing:

np.isnan(data)              # [False False  True False]
data[~np.isnan(data)]       # [1. 2. 4.]
np.isnan(data).sum()        # 1  - count of missing

Note that np.nan == np.nan is False, so equality never finds NaN. Always use np.isnan instead. Also, per the nanmean documentation, an all-NaN slice still returns NaN and raises a RuntimeWarning.

How do you find where the maximum value is in NumPy?

Use argmax, which returns the index rather than the value. argmin and argmax return indices rather than values, which is what you want when the array is aligned with something else:

temps = np.array([18, 25, 31, 12, 28])
cities = np.array(['Pune', 'Delhi', 'Nagpur', 'Shimla', 'Bhopal'])

hottest = temps.argmax()      # 2
cities[hottest]               # 'Nagpur'

For 2D arrays, convert the flat index into coordinates:

grid = np.array([[1, 8], [3, 5]])
flat = grid.argmax()                        # 1
np.unravel_index(flat, grid.shape)          # (0, 1) - shown as (np.int64(0), np.int64(1))

Other useful position functions:

np.where(temps > 25)          # (array([2, 4]),)  - indices meeting a condition
np.argsort(temps)             # [3 0 1 4 2] - indices that would sort
temps[np.argsort(temps)]      # [12 18 25 28 31]

argsort is how you sort one array by another. For instance, cities[np.argsort(temps)] gives ['Shimla' 'Pune' 'Delhi' 'Bhopal' 'Nagpur'], coldest first.

How do you sort a NumPy array?

Use np.sort for a sorted copy or the .sort() method to sort in place. Both sort along the last axis unless you pass axis:

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

np.sort(a)             # [[1 2 3] [7 8 9]] - each row (last axis by default)
np.sort(a, axis=0)     # [[3 1 2] [9 7 8]] - each column
a.sort()               # in place, modifies a

np.sort returns a new array, whereas the .sort() method modifies in place and returns None. Assigning the result of .sort() to a variable is therefore a common and puzzling bug.

Which NumPy linear algebra functions matter?

In modelling work, the everyday set is @ for matrix multiplication, np.linalg.solve for linear systems, eigh for symmetric matrices and norm for lengths:

A = np.array([[1., 2.],
              [3., 4.]])
b = np.array([5., 6.])

A @ A                     # [[ 7. 10.] [15. 22.]] - matrix multiplication
A.T                       # transpose
np.linalg.det(A)          # -2.0000000000000004
np.linalg.matrix_rank(A)  # 2

The determinant shows a tiny floating-point error rather than exactly -2.0, which is normal. Compare floats with np.allclose, never with ==.

To solve Ax = b, use solve rather than inverting:

x = np.linalg.solve(A, b)     # [-4.   4.5]
np.allclose(A @ x, b)         # True
# not: np.linalg.inv(A) @ b   - slower and less numerically stable

This matters more than it looks. Explicit inversion loses precision on ill-conditioned matrices, while np.linalg.solve uses the LAPACK gesv routine, a factorisation that avoids forming the inverse. The advice holds in every numerical library, not just NumPy.

For eigenvalues, choose the function by matrix type:

values, vectors = np.linalg.eig(A)      # general square matrix

S = np.array([[2., 1.],
              [1., 2.]])
values, vectors = np.linalg.eigh(S)     # symmetric: values [1. 3.]

If you are implementing PCA by hand, a covariance matrix is symmetric, so use eigh. It is faster and returns eigenvalues in ascending order, whereas eig guarantees no order.

Finally, norms:

np.linalg.norm(b)            # 7.810249675906654 - Euclidean length
np.linalg.norm(A, axis=1)    # [2.23606798 5.        ] - row-wise norms

Worked example: z-scores and outliers

This small example combines most of the series, computing per-feature z-scores and finding outliers:

rng = np.random.default_rng(0)
X = rng.normal(loc=50, scale=10, size=(200, 4))
X[5, 2] = np.nan                       # inject a gap

means = np.nanmean(X, axis=0, keepdims=True)
stds = np.nanstd(X, axis=0, keepdims=True)

z = (X - means) / stds                 # broadcasting
outlier_rows = np.where(np.nanmax(np.abs(z), axis=1) > 3)[0]

print(f"{len(outlier_rows)} rows contain an extreme value")
# 4 rows contain an extreme value

It uses aggregation with an axis, NaN handling, keepdims for broadcasting, and where for positions: the whole series in eight lines.

NumPy aggregation quick reference

Common NumPy aggregation mistakes

• Swapping axis=0 and axis=1. Check the result's shape: the named axis should be gone.
• Comparing NumPy and pandas standard deviations. Pass ddof explicitly so both use the same divisor.
• Finding NaN with ==. Use np.isnan, since NaN is not equal to itself.
• Writing sorted_a = a.sort(). The method returns None; use np.sort(a) for a new array.
• Inverting to solve. Use np.linalg.solve, and eigh for symmetric matrices.

Related reading

Start the series at NumPy Part 1, or revisit Part 3 on broadcasting. For visualising the results, see the Plotly tutorial and the matplotlib styling guide. The NumPy cheat sheet summarises every function in one place.