NumPy: array thinking for computational physics¶

In the previous notebook we already used NumPy arrays to simulate random walks. Here we study the ideas that make those calculations work: shape, axes, data types, indexing, views and copies, broadcasting, vectorization, and numerical linear algebra.

The goal is not to memorize a catalog of functions. It is to learn how to organize an entire physical calculation as operations on arrays.

In [1]:
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt

A brief array recap¶

An ndarray contains elements of one data type arranged along one or more axes. The most useful descriptive attributes are:

  • shape: length along each axis;
  • ndim: number of axes;
  • size: total number of elements;
  • dtype: stored numerical type;
  • itemsize: bytes used by one element;
  • nbytes: bytes used by all array elements.
In [2]:
state = np.array([
    [1.0, 0.0],
    [0.0, 1.0],
    [1.0, 1.0],
])

print(state)
print("shape   =", state.shape)
print("ndim    =", state.ndim)
print("size    =", state.size)
print("dtype   =", state.dtype)
print("itemsize=", state.itemsize, "bytes")
print("nbytes  =", state.nbytes, "bytes")

[[1. 0.]
 [0. 1.]
 [1. 1.]]
shape   = (3, 2)
ndim    = 2
size    = 6
dtype   = float64
itemsize= 8 bytes
nbytes  = 48 bytes

Several common constructors are worth recognizing. We use integer arange for indices and linspace when a floating-point interval and its endpoints matter.

In [3]:
indices = np.arange(6)
x_grid = np.linspace(0.0, 1.0, 6)
zeros = np.zeros((2, 3))
identity = np.eye(3)

print("indices:", indices)
print("x grid:", x_grid)
print("zeros:\n", zeros)
print("identity:\n", identity)

indices: [0 1 2 3 4 5]
x grid: [0.  0.2 0.4 0.6 0.8 1. ]
zeros:
 [[0. 0. 0.]
 [0. 0. 0.]]
identity:
 [[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]

Floating-point arange can accumulate step-size error and may not produce an expected endpoint. Prefer linspace when the number of samples or both endpoints are important.

In [4]:
print("arange:  ", np.arange(0.0, 1.0, 0.1))
print("linspace:", np.linspace(0.0, 1.0, 11))

arange:   [0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9]
linspace: [0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1. ]

Shapes and axes through the random walk¶

Generate many complete trajectories at once. The position array will have shape

(nwalkers, nsteps + 1, 2)
      │          │       │
   walkers      time    x,y

Each axis has a physical meaning. Keeping those meanings explicit is central to reliable array programming.

In [5]:
def random_walk_ensemble(nsteps, nwalkers, rng, step_length=1.0):
    '''Return positions with shape (nwalkers, nsteps + 1, 2).'''
    moves = step_length*np.array([
        [ 1.0,  0.0],
        [-1.0,  0.0],
        [ 0.0,  1.0],
        [ 0.0, -1.0],
    ])

    directions = rng.integers(0, 4, size=(nwalkers, nsteps))
    displacements = moves[directions]

    positions = np.empty((nwalkers, nsteps + 1, 2))
    positions[:, 0, :] = 0.0
    positions[:, 1:, :] = np.cumsum(displacements, axis=1)
    return positions
In [6]:
rng = np.random.default_rng(509)
positions = random_walk_ensemble(
    nsteps=400,
    nwalkers=3000,
    rng=rng,
)

print("positions.shape =", positions.shape)
print("one walker:      ", positions[0].shape)
print("one time slice:  ", positions[:, 100, :].shape)
print("all x values:    ", positions[:, :, 0].shape)

positions.shape = (3000, 401, 2)
one walker:       (401, 2)
one time slice:   (3000, 2)
all x values:     (3000, 401)

Reductions and the meaning of axis¶

A reduction removes one or more axes. For example, averaging over axis=0 removes the walker axis but retains time and coordinate:

In [9]:
mean_position = positions.mean(axis=0)
mean_over_time = positions.mean(axis=1)

print("positions.shape=", positions.shape)
print("mean over walkers:", mean_position.shape)
print("mean over time:   ", mean_over_time.shape)

r2 = np.sum(positions**2, axis=2)  # Sum x^2+y^2; remove coordinate axis for each random walker and each step
mean_r2 = r2.mean(axis=0)          # Average over walkers; retain time.

print("r2.shape      =", r2.shape)
print("mean_r2.shape =", mean_r2.shape)

positions.shape= (3000, 401, 2)
mean over walkers: (401, 2)
mean over time:    (3000, 2)
r2.shape      = (3000, 401)
mean_r2.shape = (401,)
In [10]:
steps = np.arange(positions.shape[1])
plt.figure(figsize=(7, 5))
plt.plot(steps, mean_r2, label="simulation")
plt.plot(steps, steps, "--", label=r"$N\ell^2$")
plt.xlabel("number of steps")
plt.ylabel(r"$\langle r^2\rangle$")
plt.title("A reduction over the walker axis")
plt.legend()
plt.show()
No description has been provided for this image

Useful reductions include sum, mean, var, std, min, max, argmin, and argmax. Always ask which physical dimension should disappear and which should remain.

The keyword keepdims=True retains reduced axes with length one, which can make later broadcasting clearer.

In [12]:
center_of_mass = positions.mean(axis=0, keepdims=True)
fluctuations = positions-center_of_mass

print("center_of_mass.shape =", center_of_mass.shape)
print("fluctuations.shape   =", fluctuations.shape)

center_of_mass.shape = (1, 401, 2)
fluctuations.shape   = (3000, 401, 2)

Indexing: selections, views, and copies¶

Basic slicing normally returns a view that shares memory with the original array. Integer-array and Boolean indexing return copies.

This distinction matters because modifying a view can modify the original data.

In [13]:
a = np.arange(10.0)

view = a[2:7]
integer_copy = a[[2, 4, 6]]
boolean_copy = a[a > 5]

print("shares with slice:  ", np.shares_memory(a, view))
print("shares with integer:", np.shares_memory(a, integer_copy))
print("shares with mask:   ", np.shares_memory(a, boolean_copy))

view[0] = 100.0
print("a after changing the view:", a)

shares with slice:   True
shares with integer: False
shares with mask:    False
a after changing the view: [  0.   1. 100.   3.   4.   5.   6.   7.   8.   9.]

Use .copy() when an independent array is required:

In [14]:
a = np.arange(10.0)
independent = a[2:7].copy()
independent[0] = 100.0

print("original:   ", a)
print("independent:", independent)

original:    [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]
independent: [100.   3.   4.   5.   6.]

Boolean and integer-array indexing¶

Boolean masks select elements satisfying a condition. Integer-array indexing selects specified entries and can reorder or repeat them.

In [15]:
final_positions = positions[:, -1, :]
final_radius = np.sqrt(np.sum(final_positions**2, axis=1))

far_mask = final_radius > 30
far_walkers = final_positions[far_mask]
largest_indices = np.argsort(final_radius)[-5:]

print("number with r > 30:", np.count_nonzero(far_mask))
print("selected shape:", far_walkers.shape)
print("five largest radii:", final_radius[largest_indices])

number with r > 30: 328
selected shape: (328, 2)
five largest radii: [51.4781507  54.40588203 56.35601121 56.92099788 58.03447251]

Data types and casting¶

Every array has one dtype. NumPy chooses a common type when an array is created, and assignment converts incoming values to the existing type.

In [16]:
integers = np.array([1, 2, 3], dtype=np.int64)
reals = np.array([1, 2, 3], dtype=np.float64)
complex_values = np.array([1, 2+1j], dtype=np.complex128)

print(integers.dtype, reals.dtype, complex_values.dtype)

integers[0] = 3.9  # Converted to integer; the fractional part is lost.
print(integers)

converted = integers.astype(np.float64)
print(converted, converted.dtype)

int64 float64 complex128
[3 2 3]
[3. 2. 3.] float64

Avoid silently storing physical data in an inappropriate type. For large calculations, smaller types such as float32 save memory, but they also provide less precision and may change numerical behavior.

In [17]:
n = 10_000_000
print("float64:", np.empty(n, dtype=np.float64).nbytes/1e6, "MB")
print("float32:", np.empty(n, dtype=np.float32).nbytes/1e6, "MB")

float64: 80.0 MB
float32: 40.0 MB

Broadcasting¶

Broadcasting combines arrays with compatible shapes without explicitly copying a smaller array. Dimensions are compared from right to left; two dimensions are compatible when they are equal or one of them is one.

For example, construct several plane waves

\begin{equation} \psi_k(x)=e^{ikx} \end{equation}

on one spatial grid.

In [18]:
x = np.linspace(0.0, 2*np.pi, 400)  # shape (400,)
k = np.array([1.0, 2.0, 3.0])       # shape (3,)

phase = k[:, None]*x[None, :]
psi = np.exp(1j*phase)

print("k[:,None].shape =", k[:, None].shape)
print("x[None,:].shape =", x[None, :].shape)
print("phase.shape     =", phase.shape)
print("psi.shape       =", psi.shape)

k[:,None].shape = (3, 1)
x[None,:].shape = (1, 400)
phase.shape     = (3, 400)
psi.shape       = (3, 400)
In [19]:
plt.figure(figsize=(8, 4))
for index, wave_number in enumerate(k):
    plt.plot(x, psi[index].real, label=fr"$k={wave_number:g}$")
plt.xlabel(r"$x$")
plt.ylabel(r"$\mathrm{Re}\,\psi_k(x)$")
plt.legend()
plt.show()
No description has been provided for this image

Broadcasting also explains the eigenvector check used later. Multiplying a matrix of eigenvectors with a one-dimensional array of eigenvalues scales each eigenvector column.

Reshaping and combining arrays¶

reshape and ravel return views when possible, but they may copy when the requested layout is incompatible with the original memory arrangement. Use np.shares_memory when the distinction matters.

In [20]:
A = np.arange(12).reshape(3, 4)
flat_view = A.ravel()
flat_copy = A.flatten()

print(A)
print("ravel shares memory: ", np.shares_memory(A, flat_view))
print("flatten shares memory:", np.shares_memory(A, flat_copy))

flat_view[0] = 100
print("A after changing ravel result:\n", A)

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
ravel shares memory:  True
flatten shares memory: False
A after changing ravel result:
 [[100   1   2   3]
 [  4   5   6   7]
 [  8   9  10  11]]
In [21]:
left = np.array([[1, 2], [3, 4]])
right = np.array([[5, 6], [7, 8]])

print("vertical stack:\n", np.vstack((left, right)))
print("horizontal stack:\n", np.hstack((left, right)))
print("new leading axis:\n", np.stack((left, right), axis=0))

vertical stack:
 [[1 2]
 [3 4]
 [5 6]
 [7 8]]
horizontal stack:
 [[1 2 5 6]
 [3 4 7 8]]
new leading axis:
 [[[1 2]
  [3 4]]

 [[5 6]
  [7 8]]]

Vectorization and performance¶

NumPy array expressions perform their elementwise loops in compiled code. A Python loop is often slower because the interpreter handles each individual element.

In [22]:
values = np.linspace(-5.0, 5.0, 100_000)

def square_with_loop(values):
    result = np.empty_like(values)
    for index, value in enumerate(values):
        result[index] = value**2
    return result

np.allclose(square_with_loop(values), values**2)
Out[22]:
True
In [23]:
%timeit -n 10 -r 3 square_with_loop(values)
%timeit -n 100 -r 3 values**2

9.6 ms ± 1.72 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
13.8 μs ± 493 ns per loop (mean ± std. dev. of 3 runs, 100 loops each)

np.vectorize is a convenience wrapper, not a compiler. It normally executes a Python function once per element.

For a step function, use direct array operations:

In [24]:
theta = np.where(values >= 0, 1, 0)
positive = values >= 0

print(theta[:5], theta[-5:])
print("number of positive entries:", np.count_nonzero(positive))

[0 0 0 0 0] [1 1 1 1 1]
number of positive entries: 50000

Application: Stockholm daily temperatures¶

The local data file contains columns for year, month, day, minimum temperature, mean temperature, an adjusted temperature, and a quality flag. We use the daily mean temperature.

In [25]:
data_file = Path("stockholm_td_adj.dat.txt")
temperature_data = np.loadtxt(data_file)

year = temperature_data[:, 0].astype(int)
month = temperature_data[:, 1].astype(int)
day = temperature_data[:, 2].astype(int)
temperature = temperature_data[:, 4]

print("data shape:", temperature_data.shape)
print("year range:", year.min(), "to", year.max())
print("temperature range:", temperature.min(), "to", temperature.max())

data shape: (77431, 7)
year range: 1800 to 2011
temperature range: -25.8 to 27.5

For visualization, construct an approximate decimal year. Subtracting one from month and day places January 1 at the integer year.

In [26]:
decimal_year = year+(month-1)/12+(day-1)/365

plt.figure(figsize=(14, 4))
plt.plot(decimal_year, temperature, linewidth=0.4)
plt.xlabel("year")
plt.ylabel(r"daily mean temperature [$^\circ$C]")
plt.title("Stockholm daily temperature")
plt.show()
No description has been provided for this image

Boolean masks and reductions on measured data¶

Select one year with a mask and calculate its statistics.

In [27]:
mask_1973 = year == 1973
temperature_1973 = temperature[mask_1973]

print("days:", temperature_1973.size)
print("mean:", temperature_1973.mean())
print("minimum:", temperature_1973.min())
print("maximum:", temperature_1973.max())

days: 365
mean: 6.646027397260275
minimum: -13.4
maximum: 24.6

Calculate the mean temperature for each calendar month across the complete record:

In [28]:
monthly_mean = np.array([
    temperature[month == selected_month].mean()
    for selected_month in range(1, 13)
])

plt.bar(np.arange(1, 13), monthly_mean)
plt.xlabel("month")
plt.ylabel(r"mean temperature [$^\circ$C]")
plt.title("Seasonal cycle in Stockholm")
plt.show()
No description has been provided for this image

np.histogram can count how many selected dates fall in each year. Pass only the year array—not the entire date table—to the histogram.

In [29]:
positive_years = year[temperature > 0]
year_edges = np.arange(year.min(), year.max()+2)
positive_days, edges = np.histogram(positive_years, bins=year_edges)
histogram_years = edges[:-1].astype(int)

plt.figure(figsize=(12, 4))
plt.plot(histogram_years, positive_days)
plt.xlabel("year")
plt.ylabel("days with mean temperature above 0°C")
plt.show()
No description has been provided for this image

Text and binary NumPy files¶

Useful NumPy file functions include:

np.loadtxt("data.txt")
np.savetxt("data.txt", array)
np.save("data.npy", array)
array = np.load("data.npy")
np.savez_compressed("several_arrays.npz", x=x, y=y)

Text files are portable and human-readable. NumPy's binary formats preserve shape and dtype and are usually faster and smaller. When exchanging labeled tables with other software, formats such as CSV, HDF5, or NetCDF may be more appropriate.

Linear algebra through coupled oscillators¶

Consider identical masses connected by identical springs with fixed endpoints. In suitable units, the stiffness matrix is

\begin{equation} K=\begin{pmatrix} 2&-1&0&\cdots\\ -1&2&-1&\cdots\\ 0&-1&2&\cdots\\ \vdots&\vdots&\vdots&\ddots \end{pmatrix}. \end{equation}

The normal modes satisfy

\begin{equation} K v_n=\omega_n^2 v_n. \end{equation}

In [30]:
nmasses = 6
K = 2*np.eye(nmasses)
K += -np.eye(nmasses, k=1)-np.eye(nmasses, k=-1)

print(K)
print("symmetric:", np.allclose(K, K.T))

[[ 2. -1.  0.  0.  0.  0.]
 [-1.  2. -1.  0.  0.  0.]
 [ 0. -1.  2. -1.  0.  0.]
 [ 0.  0. -1.  2. -1.  0.]
 [ 0.  0.  0. -1.  2. -1.]
 [ 0.  0.  0.  0. -1.  2.]]
symmetric: True

For a real symmetric or complex Hermitian matrix, use np.linalg.eigh. It exploits the matrix structure and returns real eigenvalues in ascending order.

In [31]:
omega_squared, modes = np.linalg.eigh(K)
omega = np.sqrt(omega_squared)

print("frequencies:", omega)
print("modes.shape:", modes.shape)

# Eigenvectors are stored in the columns of modes.
print("eigenvalue equation satisfied:",
      np.allclose(K @ modes, modes*omega_squared))
print("orthonormal eigenvectors:",
      np.allclose(modes.T @ modes, np.eye(nmasses)))

frequencies: [0.44504187 0.86776748 1.2469796  1.56366296 1.80193774 1.94985582]
modes.shape: (6, 6)
eigenvalue equation satisfied: True
orthonormal eigenvectors: True
In [32]:
mass_index = np.arange(1, nmasses+1)
fig, axes = plt.subplots(2, 3, figsize=(11, 6), sharex=True)

for mode_index, axis in enumerate(axes.flat):
    axis.plot(mass_index, modes[:, mode_index], "o-")
    axis.axhline(0, color="black", linewidth=0.6)
    axis.set_title(fr"$\omega={omega[mode_index]:.3f}$")
    axis.set_xlabel("mass index")

axes[0, 0].set_ylabel("displacement")
axes[1, 0].set_ylabel("displacement")
plt.tight_layout()
plt.show()
No description has been provided for this image

Elementwise multiplication versus matrix multiplication¶

For ordinary arrays, * multiplies element by element and @ performs matrix multiplication.

In [33]:
print("K*K, elementwise:\n", K*K)
print("K@K, matrix product:\n", K@K)

K*K, elementwise:
 [[4. 1. 0. 0. 0. 0.]
 [1. 4. 1. 0. 0. 0.]
 [0. 1. 4. 1. 0. 0.]
 [0. 0. 1. 4. 1. 0.]
 [0. 0. 0. 1. 4. 1.]
 [0. 0. 0. 0. 1. 4.]]
K@K, matrix product:
 [[ 5. -4.  1.  0.  0.  0.]
 [-4.  6. -4.  1.  0.  0.]
 [ 1. -4.  6. -4.  1.  0.]
 [ 0.  1. -4.  6. -4.  1.]
 [ 0.  0.  1. -4.  6. -4.]
 [ 0.  0.  0.  1. -4.  5.]]

Solve linear systems without constructing the inverse¶

For a static force $f$, the equilibrium displacement satisfies $Kx=f$. Use solve(K, f) rather than explicitly forming inv(K) @ f.

In [34]:
force = np.zeros(nmasses)
force[nmasses//2] = 1.0

displacement = np.linalg.solve(K, force)
residual = K@displacement-force

print("displacement:", displacement)
print("residual norm:", np.linalg.norm(residual))

displacement: [0.42857143 0.85714286 1.28571429 1.71428571 1.14285714 0.57142857]
residual norm: 4.440892098500626e-16

For complex arrays, .T transposes without conjugating. The Hermitian conjugate is A.conj().T.

In [35]:
H = np.array([
    [1.0, 1.0j],
    [-1.0j, 2.0],
])

print("transpose:\n", H.T)
print("Hermitian conjugate:\n", H.conj().T)
print("Hermitian:", np.allclose(H, H.conj().T))
print("eigenvalues:", np.linalg.eigvalsh(H))

transpose:
 [[ 1.+0.j -0.-1.j]
 [ 0.+1.j  2.+0.j]]
Hermitian conjugate:
 [[ 1.-0.j -0.+1.j]
 [ 0.-1.j  2.-0.j]]
Hermitian: True
eigenvalues: [0.38196601 2.61803399]

Homework: temperature trends and freezing days¶

Using the Stockholm arrays already loaded:

  1. Calculate the annual mean temperature for every year.
  2. Fit a straight line to the annual means and report the slope in degrees Celsius per century.
  3. Count the number of days per year whose mean temperature is below $0^\circ$C.
  4. Compare the mean temperature during the first 30 complete years with the last 30 complete years.
  5. Plot the annual means, fitted trend, and freezing-day counts.

Use masks, reductions, array construction, and np.polyfit. Do not loop over individual days.

Homework solution¶

Main ideas¶

  • An array's axes should correspond to clearly identified physical dimensions.
  • Reductions remove selected axes; axis is part of the physical meaning of a calculation.
  • Basic slices usually return views, while Boolean and integer-array indexing return copies.
  • Broadcasting combines compatible shapes without manually replicating data.
  • Vectorized NumPy expressions execute elementwise loops in compiled code; np.vectorize does not provide that acceleration.
  • Ordinary arrays and @ replace the old np.matrix class.
  • Use structure-aware linear algebra such as eigh for Hermitian matrices and solve for linear systems.