AI/Biz

Python Primer·Chapter 10

NumPy: Thinking in Arrays

What an array actually is in memory, why shape matters more than value, the difference between a view and a copy, and the broadcasting rules that make mismatched arrays work exactly as you intended.

Almost every difficult bug you will hit in numerical Python comes from one of three misunderstandings: what an array is in memory, whether you are holding a view or a copy, and which axis an operation is collapsing.

None of them are hard. All of them are invisible until they cost you an afternoon. This chapter is longer than the others because getting these right makes everything above NumPy easier.

What an Array Actually Is

A Python list of a million integers is a million separate objects scattered across memory, plus an array of a million pointers to find them. Each object carries a type tag, a reference count, and its value.

A NumPy array is one contiguous block of memory holding raw values, plus a small header describing how to interpret it.

flowchart LR
    subgraph L["Python list — [1, 2, 3]"]
        direction TB
        LP["pointer array"] --> O1["int object<br/>type · refcount · 1"]
        LP --> O2["int object<br/>type · refcount · 2"]
        LP --> O3["int object<br/>type · refcount · 3"]
    end
    subgraph N["NumPy array — np.array([1, 2, 3])"]
        direction TB
        H["header<br/>dtype: int64<br/>shape: (3,)<br/>strides: (8,)"]
        B["buffer<br/>01 00 00 00 00 00 00 00<br/>02 00 00 00 00 00 00 00<br/>03 00 00 00 00 00 00 00"]
        H -.describes.-> B
    end
    style N fill:#eef1fc,stroke:#2141c8
    style L fill:#f0efe9,stroke:#5c6670

That difference produces everything else. The buffer is compact, so it fits in cache. The values are uniform, so a compiled loop can process them without checking types. And because the header is separate from the buffer, two arrays can describe the same memory differently, which is where views come from and where a whole class of bugs lives.

An array has four attributes worth knowing by name:

import numpy as np

a = np.arange(12).reshape(3, 4)

a.shape      # (3, 4)      — length along each dimension
a.ndim       # 2           — number of dimensions
a.dtype      # dtype('int64') — how to read each element
a.strides    # (32, 8)     — bytes to step to move one along each axis
a.nbytes     # 96          — total buffer size

Read strides carefully, because it explains a lot. To move one row down, step 32 bytes. To move one column right, step 8 bytes. The buffer is a flat run of twelve integers, and shape plus strides is the only thing making it two-dimensional.

dtype: The Decision You Make Once

Every array has a single element type, fixed at creation. This is what buys you the speed, and it has consequences.

np.array([1, 2, 3])              # int64   on most platforms
np.array([1.0, 2, 3])            # float64 — one float promotes everything
np.array([1, 2, 3], dtype=np.int8)   # 1 byte each instead of 8
np.array([1, 2, 3], dtype=np.float32) # half the memory of float64

Three things to keep in mind.

Integers overflow silently. Python integers grow without limit. NumPy integers do not.

x = np.array([127], dtype=np.int8)
x + 1        # array([-128]) — wrapped around, no warning

This is not a bug, it is what fixed-width integers do. It bites when you shrink a dtype to save memory without checking your value range.

Floats do not compare exactly. This is true everywhere but you meet it constantly here.

0.1 + 0.2 == 0.3            # False
np.isclose(0.1 + 0.2, 0.3)  # True — use this
np.allclose(a, b)           # for whole arrays

Precision is a real trade. float32 halves your memory and often doubles your throughput, at about seven decimal digits of precision instead of sixteen. For most machine learning that is fine and it is why GPUs favour it. For accumulating a sum over ten million values, it is not, and you will watch your total quietly drift.

dtype Bytes Use it for
int8 / uint8 1 Image pixels, small categorical codes
int32 4 Counts, indices, IDs
int64 8 The default. Anything you have not thought about
float32 4 Neural network weights, large feature matrices
float64 8 The default. Statistics, accumulation, anything financial
bool 1 Masks

Views and Copies: The Bug You Will Write

Because shape and strides are separate from the buffer, NumPy can hand you a new array object that points at memory you already have. That is a view, and it is fast because nothing is copied. It is also dangerous, because writing through one name changes the other.

a = np.arange(10)
b = a[2:5]        # a VIEW — no data copied
b[0] = 999
a                 # array([  0,   1, 999,   3, ...])  ← a changed

Here is the rule, and it is worth memorising:

flowchart TD
    S["You index an array"] --> Q{"How?"}
    Q -->|"basic slicing<br/>a[2:5], a[::2], a[1]"| V["VIEW<br/>shares memory<br/>writes propagate"]
    Q -->|"boolean mask<br/>a[a > 5]"| C["COPY<br/>independent memory"]
    Q -->|"fancy indexing<br/>a[[0, 3, 7]]"| C
    V --> W["Use .copy() if you<br/>need independence"]
    style V fill:#f7e8e8,stroke:#c02020
    style C fill:#eef1fc,stroke:#2141c8

Slicing gives a view because a slice can always be expressed as a new offset and stride into the same buffer. A boolean mask or a list of indices cannot, because the elements you asked for are not evenly spaced, so NumPy has no choice but to gather them into fresh memory.

Check when you are unsure:

b = a[2:5]
b.base is a       # True  → b is a view onto a
c = a[a > 5]
c.base is None    # True  → c owns its own data

The practical habit: when a function receives an array and intends to modify it, either document that clearly or call .copy() first. Silent mutation of a caller’s data is the most common source of “the numbers changed and nobody touched them”.

Axes: The Thing Everyone Gets Backwards

axis=0 does not mean rows. It means the axis that disappears.

That reframing fixes the confusion permanently. A reduction collapses the axis you name, and the result has that dimension removed.

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

a.sum(axis=0)    # array([5, 7, 9])   shape (3,) — the 2 is gone
a.sum(axis=1)    # array([6, 15])     shape (2,) — the 3 is gone
a.sum()          # 21                 scalar — everything is gone
flowchart TB
    A["a — shape (2, 3)<br/>[[1, 2, 3],<br/> [4, 5, 6]]"]
    A -->|"sum(axis=0)<br/>collapse the rows"| B["shape (3,)<br/>[5, 7, 9]<br/><i>column totals</i>"]
    A -->|"sum(axis=1)<br/>collapse the columns"| C["shape (2,)<br/>[6, 15]<br/><i>row totals</i>"]
    A -->|"sum()"| D["scalar<br/>21"]
    style B fill:#eef1fc,stroke:#2141c8
    style C fill:#eef1fc,stroke:#2141c8

Say it out loud when you write it: “sum along axis zero” means “sum away axis zero”. Column totals come from axis=0 because the rows are what get consumed.

keepdims=True retains the collapsed axis with length one, which is exactly what you want when the result has to broadcast back against the original:

col_means = a.mean(axis=0, keepdims=True)   # shape (1, 3) not (3,)
centred = a - col_means                     # broadcasts cleanly

Broadcasting

Broadcasting is how NumPy operates on arrays of different shapes without you writing a loop or physically duplicating data. It is the single most useful thing in the library and it follows two rules.

Compare shapes right to left. For each pair of dimensions:

  1. They match, or
  2. One of them is 1, in which case it is stretched to match the other.

If neither holds, it is an error. Missing dimensions on the left are treated as 1.

flowchart TB
    subgraph OK1["Works — trailing dimensions align"]
        direction LR
        A1["(3, 4)"] --- B1["(4,)"] --> R1["→ (3, 4)"]
    end
    subgraph OK2["Works — the 1 stretches"]
        direction LR
        A2["(3, 1)"] --- B2["(1, 4)"] --> R2["→ (3, 4)"]
    end
    subgraph OK3["Works — scalar against anything"]
        direction LR
        A3["(2, 3, 4)"] --- B3["()"] --> R3["→ (2, 3, 4)"]
    end
    subgraph BAD["Fails — 3 and 4 are neither equal nor 1"]
        direction LR
        A4["(3, 4)"] --- B4["(3,)"] --> R4["✗ ValueError"]
    end
    style OK1 fill:#eef1fc,stroke:#2141c8
    style OK2 fill:#eef1fc,stroke:#2141c8
    style OK3 fill:#eef1fc,stroke:#2141c8
    style BAD fill:#f7e8e8,stroke:#c02020

That last failure is the one people hit. An array of shape (3, 4) and one of shape (3,) look compatible, because both have a three in them. But alignment is from the right, so the (3,) is compared against the 4. Fix it by making the intent explicit:

a = np.ones((3, 4))
b = np.array([10, 20, 30])       # shape (3,)

a + b                            # ValueError
a + b[:, np.newaxis]             # works — b becomes (3, 1)
a + b.reshape(-1, 1)             # identical, different spelling

Nothing is copied. Broadcasting is implemented with a stride of zero, which means “stay put as you walk along this axis”. A (1, 4) array broadcast against (1000000, 4) does not create a million rows. It reads the same four values a million times, from cache.

The classic use is computing all pairwise differences without a loop:

x = np.array([1.0, 5.0, 9.0])

# (3, 1) against (1, 3) gives every combination
diffs = x[:, None] - x[None, :]
# array([[ 0., -4., -8.],
#        [ 4.,  0., -4.],
#        [ 8.,  4.,  0.]])

That is a distance matrix in one line, and it runs entirely in the compiled layer.

Boolean Masks

Selecting by condition is the most common thing you will do, and it composes well.

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

a > 5              # array([False, True, False, True, False])
a[a > 5]           # array([8, 9])          — the values
np.where(a > 5)    # (array([1, 3]),)       — the positions
(a > 5).sum()      # 2  — True counts as 1, so this counts matches
(a > 5).mean()     # 0.4 — and this gives the proportion

Combine conditions with &, | and ~, not and, or and not. The Python keywords ask for a single truth value and an array cannot provide one. Parentheses are mandatory because & binds tighter than comparison:

a[(a > 2) & (a < 9)]      # correct
a[a > 2 & a < 9]          # wrong, and the error will not be obvious

np.where also works as a vectorised if-else, which replaces a great many loops:

np.where(a > 5, 'high', 'low')
# array(['low', 'high', 'low', 'high', 'low'], dtype='<U4')

# Clip outliers without touching the rest
np.where(a > 8, 8, a)

Missing Values

NumPy represents missing floating-point data as np.nan. It has one property worth committing to memory now: it is not equal to itself.

np.nan == np.nan     # False
np.isnan(np.nan)     # True — the only correct test

a = np.array([1.0, np.nan, 3.0])
a.sum()              # nan — one missing value poisons the result
np.nansum(a)         # 4.0 — the nan-aware version
a[~np.isnan(a)]      # array([1., 3.]) — drop them

Most reductions have a nan-prefixed twin: nanmean, nanstd, nanmax, nanpercentile. Note that nan only exists for floats. An integer array cannot hold a missing value, which is why loading data with gaps silently converts your integer column to float.

Writing Fast Code

Three habits, in order of how much they matter.

Vectorise the loop. This is worth more than everything else combined.

# Slow — the interpreter runs a million times
out = np.empty(len(data))
for i in range(len(data)):
    out[i] = data[i] ** 2 + 3 * data[i]

# Fast — one crossing into compiled code
out = data ** 2 + 3 * data

Preallocate rather than grow. np.append in a loop reallocates and copies the entire array every time, which turns linear work into quadratic.

# Quadratic. Avoid.
result = np.array([])
for chunk in chunks:
    result = np.append(result, chunk)

# Linear.
result = np.concatenate(chunks)

Operate in place when arrays are large. a = a + 1 allocates a whole new array; a += 1 does not.

a += 1            # in place, no allocation
np.sqrt(a, out=a) # most ufuncs accept an out= target

For a sense of scale, here is the same computation four ways on a million elements:

Approach Relative time
Python for loop with list append ~200×
List comprehension ~120×
np.vectorize (a convenience wrapper, still a Python loop) ~90×
Vectorised NumPy expression 1×

np.vectorize deserves a specific warning. Its name promises speed and it delivers none. It exists to make a scalar function accept arrays, and it loops in Python to do it.

A Worked Example

Standardising a feature matrix, which you will do constantly, uses most of this chapter at once:

import numpy as np

rng = np.random.default_rng(0)          # seed for reproducibility
X = rng.normal(loc=[10, 200, 3], scale=[2, 50, 0.5], size=(1000, 3))

# Per-column mean and standard deviation.
# axis=0 collapses the rows, leaving one value per feature.
mu = X.mean(axis=0)                     # shape (3,)
sigma = X.std(axis=0)                   # shape (3,)

# Broadcasting: (1000, 3) against (3,) aligns from the right. Correct.
Z = (X - mu) / sigma

Z.mean(axis=0)   # ~[0, 0, 0]
Z.std(axis=0)    # ~[1, 1, 1]

Note np.random.default_rng(0) rather than np.random.seed(0). The generator API is the current one, it is faster, and it avoids a global piece of state that any library you import can quietly change underneath you.

Mistakes Worth Knowing In Advance

  • Mutating a view and surprising your caller. If a function might modify its input, copy first or say so in the docstring.
  • Confusing (3,) with (3, 1). They broadcast completely differently. Print .shape when something behaves oddly. It is almost always shape.
  • Using == on floats. Use np.isclose or np.allclose.
  • Reaching for np.vectorize for speed. It is a convenience, not an optimisation.
  • Integer division surprises. np.array([5]) / 2 gives 2.5 as a float; // floors. Be deliberate.
  • Growing arrays in a loop. Collect into a Python list, then np.array or np.concatenate once at the end.
  • Forgetting that a reduction removes an axis. When the next operation fails to broadcast, keepdims=True is usually the fix.

What Carries Forward

The three ideas that matter most from this chapter:

Shape is the thing to reason about. Most errors are shape errors wearing a disguise. When something is wrong, print shapes before you print values.

Know whether you hold a view or a copy. Basic slicing gives views, boolean and fancy indexing give copies.

Broadcasting aligns from the right. Once that is automatic, you will stop writing loops without noticing.

Next: pandas, which is built directly on everything here. A DataFrame column is a NumPy array with a label attached, and that label, the index, is where most of pandas’ behaviour comes from.