Skip to content

NumPy Essentials

The foundation of numerical Python. NumPy provides fast, memory-efficient N-dimensional arrays and vectorized operations — pandas, scikit-learn, and PyTorch all build on it.

Table of Contents

Installation / Setup

pip install numpy            # or: uv add numpy

Works with any Python 3.10+; conda also ships it. Import convention is non-negotiable:

import numpy as np

Core Syntax

Creating arrays

Function Result
np.array([1, 2, 3]) From a list
np.zeros((2, 3)) / np.ones((2, 3)) Filled 0.0 / 1.0
np.full((2, 2), 7) Filled with a value
np.arange(5) [0 1 2 3 4]
np.linspace(0, 1, 5) 5 evenly spaced points
np.eye(3) Identity matrix
np.random.default_rng(42).random((2, 3)) Random uniform
a = np.array([[1, 2], [3, 4]])
a.shape        # (2, 2)
a.ndim         # 2
a.size         # 4
a.dtype        # dtype('int64')
a.nbytes       # 32 — memory in bytes

dtypes

dtype Size Notes
int8 … int64 1–8 B Signed ints
uint8 … uint64 1–8 B Unsigned (e.g. image pixels)
float16 / float32 / float64 2 / 4 / 8 B float64 is the default
bool 1 B Masks
complex64 / complex128 8 / 16 B FFT / signal work
str_ / object — Text or Python objects (slow)
datetime64[ns] 8 B Dates/timestamps
a.astype(np.float32)
np.array([1, 2], dtype=np.int8)

Indexing & slicing

a = np.arange(12).reshape(3, 4)
a[1]                # row 1
a[-1]               # last row
a[1, 2]             # single element (row, col)
a[:, 1]             # column 1 (all rows)
a[1:, :2]           # rows 1.., cols 0..2
a[[0, 2]]           # fancy indexing: rows 0 and 2
a[a > 5]            # boolean mask -> 1-D values

Slicing returns a view (shares memory); fancy/boolean indexing returns a copy.

view = a[0]         # view: modifies a
copy = a[[0]].copy()

Reshaping & transposing

a.reshape(2, 6)      # same data, new shape
a.reshape(-1)        # flatten to 1-D (ravel)
np.arange(4)[:, None]       # add axis: (4, 1) column vector
np.newaxis                  # same as None
a.T                  # transpose (view)
a.transpose(1, 0)    # explicit axes order
a.squeeze()          # drop length-1 axes
a.ravel()            # flatten (view when possible)
a.flatten()          # always a copy

Common Operations

Broadcasting

NumPy lines up shapes from the right; dimensions must be equal or one of them must be 1.

a = np.ones((3, 1))        # (3, 1)
b = np.arange(4)           # (4,)
a + b                      # -> (3, 4): a stretches right, b stretches down

X = np.random.rand(5, 3)   # (5, 3) data
mean = X.mean(axis=0)      # (3,)
X - mean                   # (5, 3) — row-wise centering with no loop

Aggregations

Operation Functions
Sum / mean / std / var a.sum(), a.mean(), a.std(), a.var()
Min / max a.min(), a.max(), a.argmin(), a.argmax()
Median / percentile np.median(a), np.percentile(a, 95)
Cumulative np.cumsum(a), np.cumprod(a)
Logical a.any(), a.all()

axis controls the direction: X.mean(axis=0) = down rows (per column), X.mean(axis=1) = across columns (per row).

np.random.seed(0)  # legacy but still common in tutorials
X = np.random.rand(100, 5)
col_means = X.mean(axis=0)
row_totals = X.sum(axis=1)

Mathematical operations

a + 1, a * 2, a ** 2, a // 2, a % 3      # elementwise with scalars
a + b, a * b                              # elementwise between arrays
np.sqrt(a), np.exp(a), np.log(a), np.abs(a)
np.sign(a), np.round(a, 2), np.clip(a, 0, 1)
np.sin(a), np.cos(a)
np.where(cond, x, y)                      # elementwise ternary
np.power(a, 2), np.maximum(a, b)          # ufunc forms

Random numbers (modern Generator API)

rng = np.random.default_rng(42)           # seeded, reproducible
rng.random((3, 2))                        # uniform [0, 1)
rng.normal(loc=0.0, scale=1.0, size=5)    # Gaussian
rng.integers(0, 10, size=5)               # discrete ints [0, 10)
rng.choice(np.arange(10), size=3, replace=False)  # sample without replacement
rng.permutation(np.arange(10))            # shuffled copy
rng.shuffle(X)                            # in-place shuffle of rows

np.random.default_rng(seed) is preferred over legacy np.random.seed(seed) + module-level draws.

Linear algebra

np.linalg.inv(A)          # inverse
np.linalg.pinv(A)         # pseudo-inverse (least squares)
np.linalg.det(A)
np.linalg.eig(A)          # eigenvalues / vectors
np.linalg.svd(A)          # SVD: U, s, Vt
np.linalg.qr(A)
np.linalg.solve(A, b)     # solve A x = b
np.linalg.lstsq(X, y, rcond=None)   # least-squares fit
np.linalg.norm(v)         # vector/matrix norm
X @ y                     # matmul; np.matmul(X, y)
X.dot(y)                  # equivalent

Sorting

np.sort(a)                # returns sorted copy (default last axis)
np.sort(a, axis=0)        # sort along a chosen axis
np.argsort(a)             # indices that would sort
np.partition(a, 3)        # 3 smallest first (faster than full sort)
np.argpartition(a, 3)
np.searchsorted(sorted_a, 5)   # insertion position for binary search

Masking & filtering

mask = X > 0.5                    # boolean array
X[mask] = 0.0                     # set values where mask is True
X[mask].mean()                    # mean of values matching
np.where(X > 0.5, 1, 0)           # binarize
np.isin(X, [0, 1, 2])             # membership mask
(X > 0.5).sum(axis=0)             # count matches per column
X[X[:, 0] > 0.5]                  # keep rows where first col > 0.5

Missing values

x = np.array([1.0, np.nan, 3.0])
np.isnan(x)                       # [False  True False]
x[~np.isnan(x)]                   # drop NaN
np.nanmean(x), np.nansum(x), np.nanstd(x)
np.nanmin(x), np.nanmax(x), np.nanargmax(x)
np.isfinite(x)                    # not NaN and not inf
x[np.isfinite(x)]                 # drop NaN and inf together

NumPy's nan differs from pandas' pd.NA — pandas handles missing data more richly (see the pandas sheet). For masked operations NumPy also offers np.ma.MaskedArray.

Performance

# Vectorize: never loop elementwise in Python
# SLOW:  y = [f(xi) for xi in x]      FAST: y = f(x)  (if f is vectorized)

np.add(a, b, out=a)          # reuse buffer, avoid allocation
a += b                       # in-place variant
np.ascontiguousarray(a)      # ensure C-contiguous layout for fast C loops
np.einsum("ij,jk->ik", A, B) # fuse multiply-accumulate; beats A @ B chains
  • Work in float32 when precision allows — half the memory and often faster on GPUs.
  • Prefer broadcasting over np.tile/np.repeat (no extra memory).
  • For custom elementwise logic, use Numba (@njit) — Python loops stay slow.
  • Profile with %timeit in notebooks before optimizing.

Practical Examples

1. Standardize a feature matrix

X = np.random.default_rng(1).random((100, 4))
mean = X.mean(axis=0)
std = X.std(axis=0)
X_std = (X - mean) / std            # broadcasting: (100, 4) - (4,)
X_std.mean(axis=0)                  # ~0
X_std.std(axis=0)                   # ~1

2. Remove outliers by percentile

low, high = np.percentile(X_std, [1, 99])
clean = X_std[(X_std >= low).all(axis=1) & (X_std <= high).all(axis=1)]

3. Linear regression via least squares

X = np.random.rand(200, 3)
y = X @ np.array([1.0, -2.0, 0.5]) + 0.1 * np.random.randn(200)

Xb = np.c_[np.ones(len(X)), X]              # add bias column
coef, *_ = np.linalg.lstsq(Xb, y, rcond=None)
pred = Xb @ coef
r2 = 1 - ((y - pred) ** 2).sum() / ((y - y.mean()) ** 2).sum()

Common Patterns

# One-hot labels -> matrix
labels = np.array([0, 1, 2, 1])
one_hot = np.eye(3)[labels]

# Argmax = predicted class (works with any row-major logits)
y_pred = logits.argmax(axis=1)

# Pairwise distances via broadcasting
diff = X[:, None, :] - X[None, :, :]        # (n, n, d)
dists = np.sqrt((diff ** 2).sum(axis=-1))

# Clip + normalize a gradient/weights
w = np.clip(w, -1, 1); w /= np.linalg.norm(w)

# Shuffle rows in unison (same order in X and y)
rng = np.random.default_rng(0)
idx = rng.permutation(len(X)); X, y = X[idx], y[idx]

Best Practices

  • Vectorize everything — a NumPy loop in Python is a bug waiting to happen at scale.
  • Mind axis semantics: axis=0 aggregates down the rows (per column).
  • Understand view vs copy — slices are views; mutate copies deliberately with .copy().
  • Seed your RNG (default_rng(seed)) for reproducible experiments.
  • Downcast ints and floats (int32/float32) for large data to cut memory.
  • Use np.errstate to control divide/overflow warnings: with np.errstate(divide="ignore"):.
  • Prefer np.linalg.lstsq/np.linalg.solve over hand-computed inverses (inv(X.T @ X) is numerically fragile).

Common Mistakes

Mistake Consequence Correct approach
a = b then mutate a Both arrays change (same object) a = b.copy()
In-place op on a view Unexpected edits to the parent Work on .copy() when needed
np.mean(x, 1) ambiguous Obscure axis errors Use axis= keyword
Python loop over elements 10–100× slower Vectorize or Numba
X[y] on a 2-D mask Shape confusion Ensure boolean mask is 1-D per axis
Mixing int and nan in one array NaN silently cast away (int cannot hold NaN) Use float64 for data with NaN

Troubleshooting

Symptom Cause Fix
operands could not be broadcast together Mismatched shapes Align shapes: add axes [:, None] or reshape
object dtype after np.array([...]) Ragged/inconsistent rows Pad rows or use list-of-arrays
RuntimeWarning: invalid value encountered 0/0, log of negative np.errstate; check data ranges
Silent wrong dtype (int truncation) astype before division Divide in float first
MemoryError on large data Full copy in memory float32, np.memmap, chunked processing
Slicing edits the original View semantics Call .copy() on the slice

Official Documentation

Last Updated

2026-09-06 — Initial version.