AI-assisted Python Coding

Day 3: NumPy and Matplotlib

Klaus Reygers

Studierendentage Physik, Uni Heidelberg, 23.-27.03.2026

Why NumPy in physics?

NumPy gives us:

Fast operations on large numeric arrays
Compact, readable code close to math notation
Reliable building blocks for SciPy, pandas, matplotlib

Key object: numpy.ndarray (usually called an “array”).

NumPy vs Python lists

Python lists store generic Python objects
NumPy arrays store homogeneous numeric data in contiguous memory
Most heavy computations run in optimized compiled code (C/Fortran under the hood)
Fewer Python-level loops → less interpreter overhead

Quick speed demo: list vs NumPy

import time
import numpy as np

n = 1_000_000

# Python list approach
t0 = time.perf_counter()
xs = list(range(n))
t1 = time.perf_counter()

# NumPy approach
t2 = time.perf_counter()
a = np.arange(n)
b = a * a + 1
t3 = time.perf_counter()

print(f"list time:  {t1 - t0:.4f} s")
print(f"numpy time: {t3 - t2:.4f} s")

list time:  0.0090 s
numpy time: 0.0087 s

Creating arrays: `np.array`

a = np.array([1.0, 2.0, 3.0])
b = np.array([[1, 2, 3], [4, 5, 6]])

print(a)
print(a.shape, a.dtype)
print(b)
print(b.shape, b.dtype)

[1. 2. 3.]
(3,) float64
[[1 2 3]
 [4 5 6]]
(2, 3) int64

shape tells dimensions, dtype tells numeric type.

Creating arrays: `np.zeros` and `np.linspace`

z = np.zeros(5)
grid = np.linspace(0.0, 10.0, 6)
print("zeros:", z)
print("linspace:", grid)

zeros: [0. 0. 0. 0. 0.]
linspace: [ 0.  2.  4.  6.  8. 10.]

np.zeros(n): n zeros
np.linspace(start, stop, num): evenly spaced grid (including endpoints)

Physics mini-example with `linspace`

Compute free-fall distance \(s(t)=\tfrac{1}{2}gt^2\) on a time grid.

g = 9.81
t = np.linspace(0.0, 2.0, 5)
s = 0.5 * g * t**2

print("s:", s)

s: [ 0.       1.22625  4.905   11.03625 19.62   ]

Basic indexing and slicing (1D)

x = np.array([10, 20, 30, 40, 50])

print("x[0]   =", x[0])
print("x[-1]  =", x[-1])
print("x[1:4] =", x[1:4])
print("x[:3]  =", x[:3])
print("x[::2] =", x[::2])

x[0]   = 10
x[-1]  = 50
x[1:4] = [20 30 40]
x[:3]  = [10 20 30]
x[::2] = [10 30 50]

2D arrays: indexing and slicing

A = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])

print("A shape:", A.shape)
print("A[1, 2] =", A[1, 2])
print("row 1   =", A[1, :])
print("col 0   =", A[:, 0])
print("block   =\n", A[:2, 1:3])

A shape: (3, 3)
A[1, 2] = 6
row 1   = [4 5 6]
col 0   = [1 4 7]
block   =
 [[2 3]
 [5 6]]

Visual guide to 2D slicing

Concatenation: combining arrays

a = np.array([1, 2, 3])
b = np.array([4, 5])
print("1D concatenate:", np.concatenate([a, b]))

M1 = np.array([[1, 2], [3, 4]])
M2 = np.array([[5, 6]])
print("axis=0:\n", np.concatenate([M1, M2], axis=0))
print("axis=1:\n", np.concatenate([M1, M2.T], axis=1))

1D concatenate: [1 2 3 4 5]
axis=0:
 [[1 2]
 [3 4]
 [5 6]]
axis=1:
 [[1 2 5]
 [3 4 6]]

Elementwise math: operators and functions

x = np.array([1.0, 2.0, 3.0])
y = np.array([10.0, 20.0, 30.0])

print("x + y      ->", x + y)
print("np.add     ->", np.add(x, y))
print("x * y      ->", x * y)
print("np.multiply->", np.multiply(x, y))
print("x**2       ->", x**2)
print("np.sqrt(y) ->", np.sqrt(y))
print("method sum ->", x.sum())

x + y      -> [11. 22. 33.]
np.add     -> [11. 22. 33.]
x * y      -> [10. 40. 90.]
np.multiply-> [10. 40. 90.]
x**2       -> [1. 4. 9.]
np.sqrt(y) -> [3.16227766 4.47213595 5.47722558]
method sum -> 6.0

Math functions are elementwise by default.

Summing over a given axis

M = np.array([[1, 2, 3],
              [4, 5, 6]])

print("sum all      ->", np.sum(M))
print("sum axis=0   ->", np.sum(M, axis=0))
print("sum axis=1   ->", np.sum(M, axis=1))

sum all      -> 21
sum axis=0   -> [5 7 9]
sum axis=1   -> [ 6 15]

Masking (boolean indexing)

values = np.array([-2.0, -0.5, 0.0, 1.2, 3.4])
mask = values > 0

print("mask:", mask)
print("positive values:", values[mask])
print("same shortcut:", values[values > 0])

mask: [False False False  True  True]
positive values: [1.2 3.4]
same shortcut: [1.2 3.4]

Exercise: clean measurements with a mask

You recorded count rates (counts/min). Accept only values in a plausible window.

rates = np.array([1200, 1180, 1050, 980, 930, 760, 870, 820])
low, high = 850, 1100

Tasks:

Build a boolean mask for accepted values in the interval [low, high].
Extract accepted values into accepted.
Compute the mean only of accepted.
Print accepted values below 900.

Solution: clean measurements with a mask

rates = np.array([1200, 1180, 1050, 980, 930, 760, 870, 820])
low, high = 850, 1100

window_mask = (rates >= low) & (rates <= high)
accepted = rates[window_mask]

low_mask = accepted < 900

print("window mask:", window_mask)
print("accepted:", accepted)
print("mean(accepted):", np.mean(accepted))
print("accepted below 900:", accepted[low_mask])

window mask: [False False  True  True  True False  True False]
accepted: [1050  980  930  870]
mean(accepted): 957.5
accepted below 900: [870]

Numerical Pitfalls: floating-point comparisons

import numpy as np

a = 0.1 + 0.2
b = 0.3

print("a == b:", a == b)
print("isclose:", np.isclose(a, b))

a == b: False
isclose: True

Internally (floating point representation):

0.1 ≈ 0.10000000000000000555...
0.2 ≈ 0.20000000000000001110...
0.5 = 0.10000000000000000000  (exact)

Floats are stored as binary numbers. Values like 0.5 (= 1/2, a power of 2 denominator) are exactly representable, while 0.1 or 0.2 are not exact in binary.

Use np.isclose / np.allclose instead of direct equality for floating-point data.

Special values inNumPy: np.nan, np.inf

import numpy as np

values = np.array([1.0, np.nan, np.inf, 4.0])
mask = np.isfinite(values)

print("finite mask:", mask)
print("clean values:", values[mask])

finite mask: [ True False False  True]
clean values: [1. 4.]

Before math operations, check for np.nan or np.inf:

check np.isfinite(data).all()
or filter with data[np.isfinite(data)]

Broadcasting: same operation, different shapes

Broadcasting means NumPy can combine arrays with different shapes when dimensions are compatible.

X = np.array([[1, 2, 3],
              [4, 5, 6]])      # shape (2, 3)
v = np.array([10, 20, 30])     # shape (3,)

Y = X + v
print("X shape:", X.shape)
print("v shape:", v.shape)
print("Y = X + v:\n", Y)

X shape: (2, 3)
v shape: (3,)
Y = X + v:
 [[11 22 33]
 [14 25 36]]

Here, v is added to each row of X.

Broadcasting with a 2D column vector

X = np.array([[1, 2, 3],
              [4, 5, 6]])
c = np.array([[100], [200]])    # shape (2, 1)

print("X + c:\n", X + c)

X + c:
 [[101 102 103]
 [204 205 206]]

Now each row gets a different offset.

Exercise: sensor calibration with broadcasting

You measured raw temperatures from 3 sensors over 4 time points.

Use broadcasting to apply a per-sensor calibration and a per-time drift correction:

raw = np.array([
  [20.10, 20.40, 20.80, 21.00],
  [19.70, 20.00, 20.30, 20.60],
  [20.50, 20.80, 21.10, 21.50]
])

sensor_offset = np.array([[0.15], [-0.05], [0.10]])  # shape (3, 1)
time_drift = np.array([0.00, -0.02, -0.04, -0.06])   # shape (4,)

Tasks:

Compute corrected = raw + sensor_offset + time_drift.
Print the shapes of all arrays involved.
Compute the mean corrected temperature per sensor (axis=1).

Solution: sensor calibration with broadcasting

raw = np.array([
  [20.10, 20.40, 20.80, 21.00],
  [19.70, 20.00, 20.30, 20.60],
  [20.50, 20.80, 21.10, 21.50]
])

sensor_offset = np.array([[0.15], [-0.05], [0.10]])
time_drift = np.array([0.00, -0.02, -0.04, -0.06])

corrected = raw + sensor_offset + time_drift

print("raw shape:", raw.shape)
print("sensor_offset shape:", sensor_offset.shape)
print("time_drift shape:", time_drift.shape)
print("corrected shape:", corrected.shape)
print("mean per sensor:", np.mean(corrected, axis=1))

raw shape: (3, 4)
sensor_offset shape: (3, 1)
time_drift shape: (4,)
corrected shape: (3, 4)
mean per sensor: [20.695 20.07  21.045]

sensor_offset is broadcast across columns and time_drift across rows, so NumPy can align dimensions automatically.

Outer product of two arrays via broadcasting

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

c = a[:, np.newaxis] * b[np.newaxis, :]

a[:, np.newaxis]
[[ 1]
 [ 2]
 [ 3]]

b[np.newaxis, :]
[[10 20 30]]

outer product:
[[ 10  20  30]
 [ 20  40  60]
 [ 30  60  90]]

None can be used instead of np.newaxis, for example:

a[:, None] * b[None, :]

Exercise: pairwise distances via broadcasting

Given 2D particle positions:

pos = np.array([
  [0.0, 0.0],
  [1.0, 0.0],
  [0.0, 2.0],
  [2.0, 1.0]
])

Tasks:

Build diff with shape (N, N, 2) containing all pairwise vectors pos[i] - pos[j].
From diff, compute the Euclidean distance matrix D with shape (N, N).
Verify that D is symmetric and has zeros on the diagonal.
Find the index of the nearest neighbor for each particle (excluding itself).

Solution: pairwise distances via broadcasting

pos = np.array([
  [0.0, 0.0],
  [1.0, 0.0],
  [0.0, 2.0],
  [2.0, 1.0]
])

diff = pos[:, None, :] - pos[None, :, :]   # (N, N, 2) -> None inserts a new axis
D = np.linalg.norm(diff, axis=2)           # (N, N)

print("D =\n", D)
print("symmetric:", np.allclose(D, D.T))
print("diag zero:", np.allclose(np.diag(D), 0.0))

# Exclude self-distance by setting diagonal to +inf
D_no_self = D.copy()
np.fill_diagonal(D_no_self, np.inf)
nearest = np.argmin(D_no_self, axis=1)
print("nearest neighbor index:", nearest)

D =
 [[0.         1.         2.         2.23606798]
 [1.         0.         2.23606798 1.41421356]
 [2.         2.23606798 0.         2.23606798]
 [2.23606798 1.41421356 2.23606798 0.        ]]
symmetric: True
diag zero: True
nearest neighbor index: [1 0 0 1]

Broadcasting explanation: what `None` does

pos = np.array([
  [0.0, 0.0],
  [1.0, 0.0],
  [0.0, 2.0],
  [2.0, 1.0]
])

print("pos.shape       ->", pos.shape)          # (4, 2)
print("pos[:, None, :].shape ->", pos[:, None, :].shape)   # (4, 1, 2)
print("pos[None, :, :].shape ->", pos[None, :, :].shape)   # (1, 4, 2)

pos.shape       -> (4, 2)
pos[:, None, :].shape -> (4, 1, 2)
pos[None, :, :].shape -> (1, 4, 2)

Shape alignment in broadcasting:

(4, 1, 2)
(1, 4, 2)
---------
(4, 4, 2)

Broadcasting explanation: values of `pos[:, None, :]` and `pos[None, :, :]`

ASCII diagram (great for slides / explanation)

pos[:, None, :]        pos[None, :, :]

(4, 1, 2)              (1, 4, 2)

[ p0 ]                 [ p0  p1  p2  p3 ]
[ p1 ]
[ p2 ]
[ p3 ]

Broadcasting expands them to:
(4, 4, 2)

[ p0-p0  p0-p1  p0-p2  p0-p3 ]
[ p1-p0  p1-p1  p1-p2  p1-p3 ]
[ p2-p0  p2-p1  p2-p2  p2-p3 ]
[ p3-p0  p3-p1  p3-p2  p3-p3 ]

Broadcasting intuition: `None` means “add a loop axis”

Think of None as: “I want another loop here”.

pos[:, None, :] means:

“For each point i, compare it against all j.”

So mentally translate

for i in range(N):
  for j in range(N):
    diff[i, j] = pos[i] - pos[j]

None lets NumPy create this pairwise structure without writing explicit Python loops.

Physics setup: two cables in static equilibrium

We model a mass \(m\) hanging from two cables with unknown forces \(F_1\) and \(F_2\).

At the knot, horizontal and vertical force balance gives a \(2\times2\) linear system.

Solving linear equations (physics example)

Static equilibrium with two cables holding a mass:

\[ \begin{aligned} F_1\cos(30^\circ) - F_2\cos(45^\circ) &= 0,\\ F_1\sin(30^\circ) + F_2\sin(45^\circ) &= mg. \end{aligned} \]

Unknowns are forces \(F_1, F_2\).

m = 10.0
g = 9.81

A = np.array([
  [np.cos(np.deg2rad(30)), -np.cos(np.deg2rad(45))],
  [np.sin(np.deg2rad(30)),  np.sin(np.deg2rad(45))]
])
b = np.array([0.0, m * g])

F1, F2 = np.linalg.solve(A, b)
print(f"F1 = {F1:.2f} N, F2 = {F2:.2f} N")

F1 = 71.81 N, F2 = 87.95 N

Determinant and inverse

detA = np.linalg.det(A)
A_inv = np.linalg.inv(A)

print("det(A) =", detA)
print("A^{-1} =\n", A_inv)
print("check A @ A^{-1} =\n", A @ A_inv)

det(A) = 0.9659258262890682
A^{-1} =
 [[ 0.73205081  0.73205081]
 [-0.51763809  0.89657547]]
check A @ A^{-1} =
 [[1.00000000e+00 1.71557540e-17]
 [1.46197479e-17 1.00000000e+00]]

det(A) = 0 means singular matrix (no unique solution)
In practice, prefer np.linalg.solve(A, b) over explicitly computing the inverse

Exercise: system of linear equations

Solve the following \(3 \times 3\) system for \(x_1, x_2, x_3\):

\[ \begin{aligned} 2x_1 - x_2 + 3x_3 &= 9\\ 4x_1 + 2x_2 - x_3 &= 5\\ -x_1 + 3x_2 + 2x_3 &= 7 \end{aligned} \]

Tasks:

Assemble the coefficient matrix A and right-hand side b as NumPy arrays.
Check that np.linalg.det(A) != 0.
Solve with np.linalg.solve.
Verify: check that A @ x ≈ b using np.allclose.

Solution: system of linear equations

A = np.array([
  [ 2.0, -1.0,  3.0],
  [ 4.0,  2.0, -1.0],
  [-1.0,  3.0,  2.0]
])
b = np.array([9.0, 5.0, 7.0])

print(f"det(A) = {np.linalg.det(A):.4f}")

x = np.linalg.solve(A, b)
print(f"x1 = {x[0]:.4f}, x2 = {x[1]:.4f}, x3 = {x[2]:.4f}")
print("A @ x ≈ b:", np.allclose(A @ x, b))

det(A) = 63.0000
x1 = 1.3175, x2 = 1.1111, x3 = 2.4921
A @ x ≈ b: True

NumPy Summary

NumPy arrays are fast and expressive for scientific computing
Core skills:
- create arrays (array, zeros, linspace)
- index/slice 1D and 2D arrays
- combine arrays (concatenate)
- use elementwise math, masking, axis sums
- use broadcasting for shape-aware computations
- solve linear systems with np.linalg.solve

Next: plotting NumPy results with Matplotlib.

Matplotlib: first plot of a 1D function

matplotlib is the standard plotting library in Python.

import matplotlib.pyplot as plt

x = np.linspace(0, 2 * np.pi, 400)
y = np.sin(x)

plt.figure(figsize=(6, 3.5))
plt.plot(x, y)
plt.title("y = sin(x)")
plt.xlabel("x")
plt.ylabel("y")
plt.show()

1D function: example plot

Line styles and colors

x = np.linspace(0, 2 * np.pi, 400)

plt.figure(figsize=(6.5, 3.8))
plt.plot(x, np.sin(x), color="tab:blue", linestyle="-", linewidth=2, label="sin(x)")
plt.plot(x, np.cos(x), color="darkorange", linestyle="--", linewidth=2, label="cos(x)")
plt.legend()
plt.xlabel("x")
plt.ylabel("value")
plt.title("Styles and colors")
plt.show()

Line styles and colors: example plot

Axis labels, title, and plot ranges

x = np.linspace(0, 10, 300)
y = np.exp(-0.3 * x) * np.sin(2 * x)

plt.figure(figsize=(6.5, 3.8))
plt.plot(x, y, color="teal")
plt.title("Damped oscillation")
plt.xlabel("time t (s)")
plt.ylabel("amplitude")
plt.xlim(0, 8)
plt.ylim(-1.0, 1.0)
plt.grid(True, alpha=0.3)
plt.show()

Plot ranges: example plot

Log scale

Useful when values span many orders of magnitude.

x = np.linspace(0, 5, 200)
y = np.exp(x)

plt.figure(figsize=(6.2, 3.8))
plt.plot(x, y)
plt.yscale("log")
plt.xlabel("x")
plt.ylabel("exp(x)")
plt.title("Same data with logarithmic y-axis")
plt.grid(True, which="both", alpha=0.3)
plt.show()

Log scale: example plot

Read measurement data with NumPy

Now we use NumPy file I/O (np.loadtxt) instead of manual parsing.

file_path = "data/data_rate_vs_thickness_pb.csv"

data = np.loadtxt(file_path, delimiter=",", comments="#")
thickness_mm = data[:, 0]
count_rate = data[:, 1]

thickness_mm, count_rate

(array([ 0.,  3.,  6.,  9., 12., 15., 18., 21., 24.]),
 array([52644., 43823., 35922., 29968., 24771., 20633., 16839., 13816.,
        11576.]))

Plot data points: count rate vs absorber thickness

plt.figure(figsize=(6.6, 4.0))
plt.plot(thickness_mm, count_rate, "o", color="darkorange", label="measured")
plt.xlabel("Pb thickness (mm)")
plt.ylabel("count rate (counts/min)")
plt.title("Gamma attenuation data")
plt.grid(True, alpha=0.25)
plt.legend()
plt.show()

Count rate vs absorber thickness: example plot

Error bars

Example data with measured values \(y\) and uncertainties \(\sigma_y\).

x = np.array([1, 2, 3, 4, 5])
y = np.array([1.7, 2.3, 3.5, 3.3, 4.3])
sigma_y = np.array([0.5, 0.3, 0.4, 0.4, 0.6])

plt.figure(figsize=(6.4, 4.0))
plt.errorbar(
  x, y,
  yerr=sigma_y,
  fmt="o",
  markersize=4,
  color="darkorange",
  ecolor="0.25",
  elinewidth=1.2,
  capsize=0,
  label="data with uncertainty"
)
plt.xlabel("x")
plt.ylabel("y")
plt.title("Data with y error bars")
plt.grid(True, alpha=0.25)
plt.legend()
plt.show()

Error bars: example plot

Histogram plotting

Example: distribution of many simulated count values around one mean.

rng = np.random.default_rng(42)
sim_counts = rng.poisson(lam=np.mean(count_rate), size=1000)

plt.figure(figsize=(6.8, 3.9))
plt.hist(sim_counts, bins=25, color="steelblue", edgecolor="white")
plt.xlabel("counts/min")
plt.ylabel("frequency")
plt.title("Histogram of simulated counting measurements")
plt.show()

Histogram: example plot

Planck spectrum of the Sun

Use Planck’s law

\[ B(\lambda, T) = \frac{2hc^2}{\lambda^5}\,\frac{1}{\exp\!\left(\frac{hc}{\lambda k_B T}\right)-1} \]

with \(h\), \(c\), and \(k_B\) as usual constants.

For \(T = 5778\,\mathrm{K}\):

Build \(\lambda\) from \(200\) to \(3000\) nm (convert to meters) and compute \(B(\lambda,T)\).
Convert to per-nanometer units and plot \(B\) vs wavelength (nm) with labels and grid.
Print the peak wavelength in nm.

Solution: Planck spectrum of the Sun

h = 6.62607015e-34      # Planck constant (J s)
c = 2.99792458e8       # speed of light (m/s)
k_B = 1.380649e-23     # Boltzmann constant (J/K)
T = 5778.0             # solar surface temperature (K)

lam_nm = np.linspace(200.0, 3000.0, 1000)
lam = lam_nm * 1e-9

B = (2 * h * c**2) / (lam**5) / (np.exp((h * c) / (lam * k_B * T)) - 1)
B_nm = B * 1e-9  # convert from per meter to per nanometer

imax = np.argmax(B)
lam_max_nm = lam_nm[imax]
print(f"Peak wavelength: {lam_max_nm:.1f} nm")

Peak wavelength: 502.7 nm

plt.figure(figsize=(7, 4))
plt.plot(lam_nm, B_nm, color="tab:blue", linewidth=2)
plt.xlabel("wavelength (nm)")
plt.ylabel(r"$B_\lambda$ (W m$^{-2}$ sr$^{-1}$ nm$^{-1}$)")
plt.title("Planck spectrum at T = 5778 K")
plt.grid(True, alpha=0.3)
plt.show()

Contour plot (2D scalar field)

x = np.linspace(-3, 3, 200)
y = np.linspace(-3, 3, 180)
X, Y = np.meshgrid(x, y)
Z = np.sin(X) * np.cos(Y)

plt.figure(figsize=(6.2, 4.6))
cs = plt.contour(X, Y, Z, levels=12, cmap="viridis")
plt.clabel(cs, inline=True, fontsize=8)
plt.xlabel("x")
plt.ylabel("y")
plt.title("Contour lines of z = sin(x) cos(y)")
plt.show()

Contour plot: example plot

Heatmap (same field as color map)

plt.figure(figsize=(6.2, 4.6))
im = plt.imshow(
  Z,
  origin="lower",
  extent=[x.min(), x.max(), y.min(), y.max()],
  aspect="auto",
  cmap="magma"
)
plt.colorbar(im, label="z value")
plt.xlabel("x")
plt.ylabel("y")
plt.title("Heatmap of z = sin(x) cos(y)")
plt.show()

Heatmap: example plot

3D surface plot

fig = plt.figure(figsize=(7.0, 5.0))
ax = fig.add_subplot(111, projection="3d")

surf = ax.plot_surface(X, Y, Z, cmap="viridis", linewidth=0, antialiased=True)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.set_title("3D surface: z = sin(x) cos(y)")
fig.colorbar(surf, shrink=0.65, label="z value")
plt.show()

3D surface plot: example plot

2 x 2 panel plot

Use plt.subplots(2, 2) to create a grid of four axes in one figure.

t = np.linspace(0, 2 * np.pi, 300)

fig, axes = plt.subplots(2, 2, figsize=(8, 6))

axes[0, 0].plot(t, np.sin(t), color="tab:blue")
axes[0, 0].set_title("sin(t)")

axes[0, 1].plot(t, np.cos(t), color="darkorange")
axes[0, 1].set_title("cos(t)")

axes[1, 0].plot(t, np.sin(2 * t), color="teal")
axes[1, 0].set_title("sin(2t)")

axes[1, 1].plot(t, np.exp(-0.3 * t) * np.sin(3 * t), color="crimson")
axes[1, 1].set_title("damped oscillation")

for ax in axes.flat:
  ax.set_xlabel("t")
  ax.set_ylabel("y")
  ax.grid(True, alpha=0.25)

fig.tight_layout()
plt.show()

2 x 2 panel plot: example plot

Matplotlib summary

1D function plots: plot, labels, ranges, styles, colors, log scales
Data visualization: points + error bars with errorbar
Distributions: hist
2D fields: contour, imshow (heatmap)
3D views: plot_surface
Multi-panel figures: plt.subplots(2, 2)

Next step: combine NumPy + Matplotlib for complete mini data analyses.

Mandelbrot set (Apfelmännchen)

For each complex number \(c = x + iy\), define the iteration

\[ z_{n+1} = z_n^2 + c, \qquad z_0 = 0. \]

If the sequence stays bounded, \(c\) belongs to the Mandelbrot set. In practice, we stop after max_iter steps and mark points that did not escape beyond \(|z| > 2\).

Tasks:

Build a grid in the complex plane, e.g. \(x \in [-2.2, 1.0]\), \(y \in [-1.4, 1.4]\).
Iterate \(z \leftarrow z^2 + c\) for all grid points at once (vectorized NumPy).
Store for each point the first iteration where \(|z| > 2\) (escape time).
Plot the result with imshow and a perceptually clear colormap.

Solution: Mandelbrot set (Apfelmännchen)

import numpy as np
import matplotlib.pyplot as plt

# Grid in the complex plane
nx, ny = 900, 700
x = np.linspace(-2.2, 1.0, nx)
y = np.linspace(-1.4, 1.4, ny)
X, Y = np.meshgrid(x, y)
c = X + 1j * Y

# Mandelbrot iteration
max_iter = 120
z = np.zeros_like(c)
escape_iter = np.full(c.shape, max_iter, dtype=int)
active = np.ones(c.shape, dtype=bool)

for n in range(max_iter):
  z[active] = z[active] * z[active] + c[active]
  escaped_now = active & (np.abs(z) > 2.0)
  escape_iter[escaped_now] = n
  active[escaped_now] = False

print("escape_iter shape:", escape_iter.shape)
print("max stored iteration:", escape_iter.max())

escape_iter shape: (700, 900)
max stored iteration: 120

Solution: Mandelbrot plot

plt.figure(figsize=(4.4, 5.6))
im = plt.imshow(
  escape_iter,
  extent=[x.min(), x.max(), y.min(), y.max()],
  origin="lower",
  cmap="magma",
  aspect="equal"
)
plt.colorbar(im, orientation="vertical", pad=0.08, shrink=0.5, label="escape iteration")
plt.xlabel("Re(c)")
plt.ylabel("Im(c)")
plt.title("Mandelbrot set (Apfelmännchen)")
plt.show()

AI-assisted Python Coding

Why NumPy in physics?

NumPy vs Python lists

Quick speed demo: list vs NumPy

Creating arrays: np.array

Creating arrays: np.zeros and np.linspace

Physics mini-example with linspace

Basic indexing and slicing (1D)

2D arrays: indexing and slicing

Visual guide to 2D slicing

Concatenation: combining arrays

Elementwise math: operators and functions

Summing over a given axis

Masking (boolean indexing)

Exercise: clean measurements with a mask

Solution: clean measurements with a mask

Numerical Pitfalls: floating-point comparisons

Special values inNumPy: np.nan, np.inf

Broadcasting: same operation, different shapes

Broadcasting with a 2D column vector

Exercise: sensor calibration with broadcasting

Solution: sensor calibration with broadcasting

Outer product of two arrays via broadcasting

Exercise: pairwise distances via broadcasting

Solution: pairwise distances via broadcasting

Broadcasting explanation: what None does

Broadcasting explanation: values of pos[:, None, :] and pos[None, :, :]

Broadcasting intuition: None means “add a loop axis”

Physics setup: two cables in static equilibrium

Solving linear equations (physics example)

Determinant and inverse

Exercise: system of linear equations

Solution: system of linear equations

NumPy Summary

Matplotlib: first plot of a 1D function

1D function: example plot

Line styles and colors

Line styles and colors: example plot

Axis labels, title, and plot ranges

Plot ranges: example plot

Log scale

Log scale: example plot

Read measurement data with NumPy

Plot data points: count rate vs absorber thickness

Count rate vs absorber thickness: example plot

Error bars

Error bars: example plot

Histogram plotting

Histogram: example plot

Planck spectrum of the Sun

Solution: Planck spectrum of the Sun

Contour plot (2D scalar field)

Contour plot: example plot

Heatmap (same field as color map)

Heatmap: example plot

3D surface plot

3D surface plot: example plot

2 x 2 panel plot

2 x 2 panel plot: example plot

Matplotlib summary

Mandelbrot set (Apfelmännchen)

Solution: Mandelbrot set (Apfelmännchen)

Solution: Mandelbrot plot

Creating arrays: `np.array`

Creating arrays: `np.zeros` and `np.linspace`

Physics mini-example with `linspace`

Broadcasting explanation: what `None` does

Broadcasting explanation: values of `pos[:, None, :]` and `pos[None, :, :]`

Broadcasting intuition: `None` means “add a loop axis”