NumPy Vectorization, Boolean Masks & Statistical Aggregations

NumPy powers modern numerical computing in Python. This guide covers the 2D Axis mental model, NaN-safe aggregations, Boolean masking, and statistical reductions.

1. The 2D Axis Mental Model

One of the most common points of confusion in NumPy and Pandas is the axis parameter in 2D arrays:

import numpy as np

matrix = np.array([
    [10, 20, 30],
    [40, 50, 60]
])  # Shape: (2, 3) -> 2 rows, 3 columns

Axis Setting	Direction of Collapse	Code Example	Result
None (Default)	Entire array into a single scalar	`matrix.sum()`	`210`
`axis=0`	Down columns (collapses rows; results in 1 value per column)	`matrix.sum(axis=0)`	`[50, 70, 90]` (length 3)
`axis=1`	Across rows (collapses columns; results in 1 value per row)	`matrix.sum(axis=1)`	`[60, 150]` (length 2)

2. Handling NaN Values with NaN-Safe Aggregations

In NumPy, standard aggregations propagate np.nan values. If a single element is NaN, the result is NaN:

data = np.array([10.0, 20.0, np.nan, 40.0])

print(np.mean(data))     # Output: nan (Standard aggregation fails)
print(np.nanmean(data))  # Output: 23.333 (Safely ignores NaN)
print(np.nansum(data))   # Output: 70.0
print(np.nanstd(data))   # Output: 12.472

3. Boolean Masking & Conditional Filtering

Boolean indexing allows you to filter elements in a NumPy array without writing a loop:

scores = np.array([55, 82, 90, 45, 78, 95, 60])

# 1. Create a boolean mask
passing_mask = scores >= 70
# Output: [False, True, True, False, True, True, False]

# 2. Extract matching elements
passing_scores = scores[passing_mask]
# Output: [82, 90, 78, 95]

# 3. Compound conditions (use & for AND, | for OR, ~ for NOT):
honors = scores[(scores >= 80) & (scores <= 95)]
# Output: [82, 90, 95]

# 4. Conditional replacement with np.where(condition, if_true, if_false)
grade_labels = np.where(scores >= 70, "PASS", "FAIL")
# Output: ['FAIL', 'PASS', 'PASS', 'FAIL', 'PASS', 'PASS', 'FAIL']

4. Statistical Operations: Percentiles & Cumulative Sums

sales = np.array([100, 150, 200, 250, 300])

# Cumulative Revenue over time:
cum_sales = np.cumsum(sales)
# Output: [100, 250, 450, 700, 1000]

# Percentiles / Quartiles:
q25, q50, q75 = np.percentile(sales, [25, 50, 75])
print(f"25th: {q25}, Median: {q50}, 75th: {q75}")

# Finding index positions of min and max:
arr = np.array([10, 45, 3, 99, 24])
max_idx = np.argmax(arr)  # 3 (points to 99)
min_idx = np.argmin(arr)  # 2 (points to 3)