NumPy arrays and vectorization

NumPy stores homogeneous values in compact multidimensional arrays and applies operations in compiled loops. It is a strong fit for numerical workloads where shape and dtype are part of the data contract.

Install and inspect an array

uv add numpy

import numpy as np

temperatures = np.array(
    [
        [18.5, 20.0, 21.5],
        [17.0, 19.5, 22.0],
    ],
    dtype=np.float64,
)

print(temperatures.shape)  # (2, 3)
print(temperatures.ndim)   # 2
print(temperatures.dtype)  # float64
print(temperatures.size)   # 6

Here, axis 0 has two observations and axis 1 has three measurements. Name axis meaning in variable names or documentation; (2, 3) alone does not explain the domain.

Construct arrays deliberately

zeros = np.zeros((3, 4), dtype=np.float32)
identifiers = np.arange(100, 105, dtype=np.int64)
grid = np.linspace(0.0, 1.0, num=5)

Avoid np.empty unless every element will be overwritten; it exposes uninitialized memory values.

Vectorization replaces Python loops

celsius = np.array([0.0, 10.0, 20.0, 30.0])
fahrenheit = celsius * 9 / 5 + 32

Filtering is also vectorized:

readings = np.array([18.2, np.nan, 22.4, 31.0])
valid = readings[~np.isnan(readings)]
comfortable = valid[(valid >= 20.0) & (valid <= 25.0)]

Use &, |, and ~ for element-wise Boolean logic and parenthesize each comparison. Python's and and or expect one truth value and are ambiguous for multi-element arrays.

Vectorization is not automatically faster if it creates several huge temporary arrays. Measure memory and runtime for the real workload.

Shapes and reshaping

values = np.arange(12)
matrix = values.reshape(3, 4)
rows = matrix.reshape(-1, 2)

assert matrix.shape == (3, 4)
assert rows.shape == (6, 2)

-1 asks NumPy to infer one dimension. Reshape must preserve the number of elements.

Add and remove singleton axes intentionally:

column = np.array([1.0, 2.0, 3.0])[:, np.newaxis]
assert column.shape == (3, 1)

flat = np.squeeze(column, axis=1)
assert flat.shape == (3,)

Specify the axis to squeeze; an unrestricted squeeze can silently remove more dimensions after an input shape changes.

Broadcasting

Broadcasting aligns trailing dimensions. Dimensions are compatible when they are equal or one of them is 1.

measurements = np.array(
    [
        [10.0, 20.0, 30.0],
        [12.0, 18.0, 33.0],
    ]
)  # (2, 3)
offsets = np.array([0.5, -1.0, 2.0])  # (3,)

calibrated = measurements + offsets  # (2, 3)

For a per-row offset, reshape:

row_offsets = np.array([1.0, 10.0])[:, None]  # (2, 1)
adjusted = measurements + row_offsets         # (2, 3)

An operation that succeeds can still broadcast along the wrong semantic axis. Assert expected shapes at module boundaries.

Dtypes affect correctness

small = np.array([120, 10], dtype=np.int8)
overflowed = small + np.int8(20)

Fixed-width integers can overflow. Select dtypes from valid ranges, not only current sample values.

Floating-point arithmetic is approximate:

assert np.isclose(0.1 + 0.2, 0.3)

Use np.isclose or np.allclose, choosing tolerances based on the domain. Do not use binary floating point for exact currency; use integer minor units or decimal.Decimal outside NumPy.

Mixing types may promote the array:

mixed = np.array([1, 2.5])
assert mixed.dtype == np.float64

An object dtype usually loses vectorized performance and may indicate heterogeneous or malformed input.

Views versus copies

Basic slicing usually returns a view:

source = np.array([1, 2, 3, 4])
view = source[1:3]
view[0] = 99
assert source.tolist() == [1, 99, 3, 4]

Copy when independent ownership is required:

independent = source[1:3].copy()

Advanced indexing, such as integer arrays or Boolean masks, generally returns copies. When mutation matters, verify ownership with np.shares_memory.

Reductions and axes

scores = np.array(
    [
        [80.0, 90.0, 100.0],
        [70.0, 85.0, 95.0],
    ]
)

per_student = scores.mean(axis=1)  # shape (2,)
per_exam = scores.mean(axis=0)     # shape (3,)
overall = scores.mean()            # scalar

For missing floats:

average = np.nanmean(readings)

Treat missing-data policy as a domain decision. Silently ignoring NaN may conceal failed sensors.

Reproducible random numbers

rng = np.random.default_rng(seed=2026)
sample = rng.normal(loc=0.0, scale=1.0, size=(100, 4))

Pass a Generator into functions rather than relying on global random state. Tests can inject a fixed seed while production can create an unseeded generator.

Common pitfalls

• Comparing arrays with == inside a plain assert; use np.array_equal or testing helpers.
• Using and/or instead of element-wise Boolean operators.
• Ignoring dtype overflow or unintended object arrays.
• Confusing matrix multiplication (@) with element-wise multiplication (*).
• Mutating a view and unexpectedly changing the source.
• Broadcasting successfully along the wrong axis.
• Vectorizing tiny one-off operations where plain Python is clearer.

Practice

1. Normalize each column of a 2D array using broadcasting.
2. Reject input unless it has shape (n, 3) and a floating dtype.
3. Demonstrate one slicing view and one advanced-indexing copy.
4. Compare a Python loop with a vectorized operation using realistic data size.
5. Write tests with numpy.testing.assert_allclose and domain-appropriate tolerances.