NumPy dimensions and axes

Use this notebook

Use Notebook 2: EEG arrays for this chapter: open it in the book or download the notebook. It applies indexing, Boolean masks, aggregation, and reshaping to an EEG recording.

Start with meaning, then shape

An epoch is a short segment of recorded EEG cut out around a known event, such as a stimulus onset. Keeping many event-aligned segments together lets us compare trials and calculate an average response. We represent them as:

trials × channels × time

import numpy as np

rng = np.random.default_rng(42)          # fixed seed makes this example repeatable
epochs = rng.normal(size=(80, 32, 500))  # trials × channels × time

print(epochs.shape)  # array dimensions: (80, 32, 500)
print(epochs.ndim)   # number of dimensions: 3

Axis	Meaning	Size
0	trials	80
1	channels	32
2	time samples	500

Axis explorer

Start with (80, 32, 500), meaning trials × channels × time.

Average over axis 0: trials axis 1: channels axis 2: time samples

Check with Python: change axis and compare the printed shape with the explorer.

import numpy as np epochs = np.zeros((80, 32, 500)) axis = 0 evoked = epochs.mean(axis=axis) print("input:", epochs.shape) print("output:", evoked.shape) Run Python

Output will appear here.

Indexing removes selected dimensions

epochs[0].shape          # first trial; trial axis removed
epochs[:, 0, :].shape    # channel 0 across all trials and times
epochs[:, :, 100].shape  # time sample 100 across trials and channels
epochs[0, 0, 100]        # one scalar: one trial, channel, and time point

Check with Python: change the indexing expression on the final line.

import numpy as np epochs = np.zeros((80, 32, 500)) selection = epochs[:, 0, :] print(selection.shape) Run Python

Output will appear here.

Aggregation removes an axis

Common mistake

The number passed to axis identifies the dimension being removed, not the dimension you want to keep. Write the meaning above every dimension before calculating a mean.

evoked = epochs.mean(axis=0)       # average across trials
print(evoked.shape)                # (32, 500): trials axis removed

We averaged trials, so the trial dimension disappeared. This is both a Python operation and a scientific decision about which observations to combine.

Exercise 9 (Predict the shapes)

For epochs.shape == (80, 32, 500), predict:

1. epochs.mean(axis=1).shape

2. epochs.mean(axis=2).shape

3. epochs.mean(axis=(0, 1)).shape

4. epochs.mean(axis=0, keepdims=True).shape

Hint

Indexing with an integer removes that selected axis. Calling .mean(axis=n) removes axis n; all other axes remain in the same order.

Solution to Exercise 9 (Predict the shapes)

1. (80, 500): channels removed.

2. (80, 32): time removed.

3. (500,): trials and channels removed.

4. (1, 32, 500): trial axis retained with size one.

Broadcasting

Broadcasting is NumPy’s way of applying an operation to arrays with compatible shapes. When a dimension has size one, NumPy repeats that value across the matching larger dimension instead of requiring you to copy the data yourself.

Subtract a baseline for every trial and channel:

baseline = epochs[:, :, :100].mean(axis=2, keepdims=True)  # one baseline per trial/channel
corrected = epochs - baseline                              # broadcast across time

print(baseline.shape)   # (80, 32, 1): keepdims preserves the time axis
print(corrected.shape)  # (80, 32, 500): same shape as the input

NumPy stretches the final size-one dimension across time.

Check with Python: remove keepdims=True and inspect why subtraction then fails.

import numpy as np epochs = np.ones((4, 3, 10)) baseline = epochs[:, :, :2].mean(axis=2, keepdims=True) corrected = epochs - baseline print("epochs:", epochs.shape) print("baseline:", baseline.shape) print("corrected:", corrected.shape) Run Python

Output will appear here.

Shape-first debugging

Before a transformation, write:

print("epochs:", epochs.shape)       # input shape
print("baseline:", baseline.shape)   # shape used for broadcasting

If you cannot describe what every dimension means, pause before continuing.

Transfer

The same reasoning applies to NLP tensors such as batch × tokens × embedding dimensions and data-science matrices such as samples × features.