Basic Python through physics¶
Python syntax is easiest to learn when it is needed to answer a physical question. We will begin with relativistic kinematics and then build a simulation of diffusion from a two-dimensional random walk.
Along the way we will encounter numbers, strings, lists, tuples, dictionaries, conditionals, loops, functions, classes, numerical integration, and plotting.
Variables, numerical types, and a relativistic particle¶
A Python variable is a name bound to an object. Python determines the object's type at runtime, so a name can later be rebound to an object of another type.
For a particle moving at $v=\beta c$, the Lorentz factor is
\begin{equation} \gamma=\frac{1}{\sqrt{1-\beta^2}}. \end{equation}
import math
import numpy as np
import matplotlib.pyplot as plt
c = 299_792_458.0 # float: speed of light in m/s
particle_name = "electron" # str
number_of_particles = 1 # int
relativistic = True # bool
wave_amplitude = 1.0 + 0.5j # complex
for value in (c, particle_name, number_of_particles, relativistic, wave_amplitude):
print(f"{value!r:>20} has type {type(value).__name__}")
299792458.0 has type float
'electron' has type str
1 has type int
True has type bool
(1+0.5j) has type complex
f"{value!r:>20}"
│ │ │
│ │ └── right-align in a field 20 characters wide
│ └────── use repr(value), rather than str(value)
└─────────── value to format
value = "electron"
print(f"{value}") # uses str() to convert
print(f"{value!r}") # uses repr() to convert
electron 'electron'
Python reserves some names as keywords, including for, if, def, class, and lambda. The authoritative list is available from the language itself:
import keyword
print(keyword.kwlist)
['False', 'None', 'True', 'and', 'as', 'assert', 'async', 'await', 'break', 'class', 'continue', 'def', 'del', 'elif', 'else', 'except', 'finally', 'for', 'from', 'global', 'if', 'import', 'in', 'is', 'lambda', 'nonlocal', 'not', 'or', 'pass', 'raise', 'return', 'try', 'while', 'with', 'yield']
Lists, loops, and formatted output¶
A list is an ordered, mutable collection. Here it stores several velocities. A for loop visits each entry, and an if statement identifies the most relativistic case.
The expression inside an f-string, such as {beta:5.2f}, is evaluated and formatted when the string is created.
betas = [0.01, 0.10, 0.50, 0.90, 0.99, 0.999]
for beta in betas:
gamma = 1.0/math.sqrt(1.0-beta**2)
message = f"beta = {beta:6.3f}, gamma = {gamma:9.4f}"
if gamma > 10:
message += " strongly relativistic"
print(message)
beta = 0.010, gamma = 1.0001 beta = 0.100, gamma = 1.0050 beta = 0.500, gamma = 1.1547 beta = 0.900, gamma = 2.2942 beta = 0.990, gamma = 7.0888 beta = 0.999, gamma = 22.3663 strongly relativistic
Useful list operations include:
values[index]: select one entry;values[start:stop:step]: take a slice (stopis excluded);append(value): add one entry;value in values: test membership;- negative indices such as
values[-1]: count backward from the end.
print("first beta:", betas[0])
print("last beta:", betas[-1])
print("every second beta:", betas[::2])
betas.append(0.9999)
print("0.9 is present:", 0.9 in betas)
print("updated list:", betas)
first beta: 0.01 last beta: 0.999 every second beta: [0.01, 0.5, 0.99] 0.9 is present: True updated list: [0.01, 0.1, 0.5, 0.9, 0.99, 0.999, 0.9999]
Comprehensions¶
A list comprehension is a compact way to construct a list. It should be used when the expression remains easy to read.
gammas = [1.0/math.sqrt(1.0-beta**2) for beta in betas]
print(gammas)
# The equivalent explicit loop:
gammas_loop = []
for beta in betas:
gammas_loop.append(1.0/math.sqrt(1.0-beta**2))
print("The results agree:", gammas == gammas_loop)
[1.0000500037503126, 1.005037815259212, 1.1547005383792517, 2.294157338705618, 7.088812050083354, 22.36627204212937, 70.71244595191452] The results agree: True
A note about older formatting styles¶
Older scientific programs often use % formatting or str.format. Students should be able to recognize them, but f-strings are normally clearer for new code.
beta = 0.9
gamma = 1.0/math.sqrt(1.0-beta**2)
print(f"f-string: beta={beta:.2f}, gamma={gamma:.5f}")
print("percent style: beta=%.2f, gamma=%.5f" % (beta, gamma))
print("str.format: beta={:.2f}, gamma={:.5f}".format(beta, gamma))
f-string: beta=0.90, gamma=2.29416 percent style: beta=0.90, gamma=2.29416 str.format: beta=0.90, gamma=2.29416
Dictionaries: giving physical parameters names¶
A dictionary maps hashable keys to values. Named parameters are usually clearer than remembering positions in a list.
electron = {
"name": "electron",
"mass": 9.109_383_7139e-31, # kg
"charge": -1.602_176_634e-19, # C
"spin": 0.5,
}
print(electron["name"])
print(f"mass = {electron['mass']:.6e} kg")
electron["antiparticle"] = "positron"
for key, value in electron.items():
print(f"{key:>12}: {value}")
electron
mass = 9.109384e-31 kg
name: electron
mass: 9.1093837139e-31
charge: -1.602176634e-19
spin: 0.5
antiparticle: positron
Dictionary keys can be numbers, strings, or other hashable objects. A tuple of immutable objects can be a key, which is useful for a sparse mapping:
matrix_elements = {
(0, 0): 1.0,
(1, 1): 2.0,
(0, 2): -0.25,
}
print("H[0,2] =", matrix_elements[(0, 2)])
print("H[2,2] =", matrix_elements.get((2, 2), 0.0))
H[0,2] = -0.25 H[2,2] = 0.0
Functions: package a physical calculation¶
A function definition begins with def. A docstring explains what the function calculates, and return sends a result back to the caller.
def lorentz_factor(beta):
'''Return the Lorentz factor for a speed beta*c.'''
if abs(beta) >= 1:
raise ValueError("A massive particle must have |beta| < 1")
return 1.0/math.sqrt(1.0-beta**2)
def relativistic_energy(momentum, mass, speed_of_light=c):
'''Return sqrt((pc)^2 + (mc^2)^2) in joules.'''
return math.sqrt((momentum*speed_of_light)**2 + (mass*speed_of_light**2)**2)
print("gamma(0.99) =", lorentz_factor(0.99))
energy = relativistic_energy(1.0e-22, electron["mass"])
print(f"energy = {energy/electron['charge']*-1:.6e} eV")
gamma(0.99) = 7.088812050083354 energy = 5.441803e+05 eV
The preceding call uses positional arguments. Keyword arguments document their meaning and can be given in a different order:
energy = relativistic_energy(mass=electron["mass"],momentum=1.0e-22,)
print(energy)
8.718730009052618e-14
Assignment, aliasing, and mutability¶
Assignment binds a name to an object. It does not automatically copy that object. Lists and dictionaries are mutable; numbers, strings, and tuples are immutable.
initial_state = [1.0, 0.0]
state_alias = initial_state
state_copy = initial_state.copy()
state_alias[0] = 2.0
print("initial_state:", initial_state)
print("state_alias: ", state_alias)
print("state_copy: ", state_copy)
initial_state: [2.0, 0.0] state_alias: [2.0, 0.0] state_copy: [1.0, 0.0]
Python passes references to objects into functions. A function can mutate a supplied list, but rebinding a local parameter does not rebind the caller's name.
def change_state(state, time):
state[0] = 10.0 # Mutates the caller's list.
time = 5.0 # Rebinds only the local name.
state = [1.0, 0.0]
t = 0.0
change_state(state, t)
print("state =", state)
print("t =", t)
state = [10.0, 0.0] t = 0.0
A physical problem: diffusion from a random walk¶
Consider a particle on a square lattice. During each time step it moves a distance $\ell$ in one of four directions with equal probability.
After $N$ independent steps, symmetry predicts
\begin{equation} \langle x\rangle=\langle y\rangle=0, \qquad \langle r^2\rangle=N\ell^2. \end{equation}
We will first simulate one trajectory using lists, a loop, and conditionals.
def random_walk_2d(nsteps, step_length=1.0, rng=np.random.default_rng()):
"""Return x and y coordinates of one unbiased square-lattice walk."""
# Displacements corresponding to right, left, up, and down.
moves = step_length*np.array([
[ 1.0, 0.0],
[-1.0, 0.0],
[ 0.0, 1.0],
[ 0.0, -1.0],
])
# Generate all random directions at once.
directions = rng.integers(0, 4, size=nsteps) # nstep integers between 0<=direction[i]<4
# Convert the direction indices into displacement vectors. Vectorized call.
displacements = moves[directions] # displacements[i]=[1.0,0.0]*sl or [-1.0,0.0]*sl or ...
# Add all preceding displacements to obtain the trajectory.
positions = np.empty((nsteps + 1, 2)) # position[i,xy] = \sum_{k=0}^{k=i} displacement[k,xy]
positions[0] = 0.0
positions[1:] = np.cumsum(displacements, axis=0) # sum over first index
# returning reached positions x[i] and y[i] for all steps
return positions[:, 0], positions[:, 1]
The function returns two objects. Python packs them into a tuple; the assignment below unpacks that tuple into x and y.
The random-number generator is seeded so that the lecture produces the same trajectory every time. Changing the seed produces a different valid realization.
rng = np.random.default_rng(509)
x, y = random_walk_2d(1000, step_length=1.0, rng=rng)
plt.figure(figsize=(6, 6))
plt.plot(x, y, linewidth=0.8)
plt.scatter(x[0], y[0], color="green", label="start", zorder=3)
plt.scatter(x[-1], y[-1], color="red", label="finish", zorder=3)
plt.xlabel(r"$x/\ell$")
plt.ylabel(r"$y/\ell$")
plt.title("One two-dimensional random walk")
plt.axis("equal")
plt.legend()
plt.show()
One random trajectory does not look like a smooth diffusion law. The prediction concerns an ensemble average, so we must repeat the experiment for many particles.
A dictionary provides a readable collection of simulation parameters.
parameters = {
"nsteps": 500,
"nwalkers": 2000,
"step_length": 1.0,
"seed": 509,
}
print(parameters)
{'nsteps': 500, 'nwalkers': 2000, 'step_length': 1.0, 'seed': 509}
def mean_square_displacement(nsteps, nwalkers, step_length=1.0, seed=None):
'''Return the ensemble-averaged r^2 after every step.'''
rng = np.random.default_rng(seed)
mean_r2 = np.zeros(nsteps + 1)
for _ in range(nwalkers):
x, y = random_walk_2d(nsteps, step_length, rng)
mean_r2 += x**2 + y**2
return mean_r2/nwalkers
mean_r2 = mean_square_displacement(**parameters)
steps = np.arange(parameters["nsteps"] + 1)
theory = steps*parameters["step_length"]**2
plt.figure(figsize=(7, 5))
plt.plot(steps, mean_r2, label="simulation")
plt.plot(steps, theory, "--", label=r"$N\ell^2$")
plt.xlabel("number of steps, $N$")
plt.ylabel(r"$\langle r^2\rangle$")
plt.title("Diffusion emerges from an ensemble of random walks")
plt.legend()
plt.show()
If $\langle r^2\rangle\propto N^\alpha$, the slope of a log-log graph is $\alpha$. Normal diffusion should give approximately one.
How to conventiently pull out part of the array?
print('mask for steps>400=', steps>400)
print('steps for that mask=', steps[steps>400])
print('array for that mask=', mean_r2[steps>400])
mask for steps>400= [False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False False True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True True] steps for that mask= [401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500] array for that mask= [404.368 404.605 406.064 407.1 407.368 407.974 411.126 412.372 414.102 414.271 415.466 417.12 418.206 420.156 421.124 421.732 422.758 423.639 425.402 426.764 427.536 428.614 429.288 430.855 433.054 434.183 435.456 437.005 439.126 439.745 439.592 441.387 442.21 442.762 443.536 444.601 444.59 445.294 445.722 447.552 447.504 449.032 449.556 449.894 451.04 452.451 453.738 455.572 457.332 458.977 460.44 460.8 462.83 463.64 465.464 466.895 467.806 468.693 469.738 469.48 470.372 472.479 473.684 475.258 476.666 478.706 479.398 480.337 480.884 480.745 481.804 482.795 485.668 486.011 488.02 488.637 489.33 490.696 492.5 493.197 493.646 494.077 495.672 497.662 498.956 500.145 502.086 503.209 504.05 505.41 506.028 507.619 508.234 509.201 510.68 511.437 511.88 512.93 513.996 514.438]
fit_region = steps >= 20
slope, intercept = np.polyfit(
np.log(steps[fit_region]),
np.log(mean_r2[fit_region]),
1,
)
print(f"measured log-log slope = {slope:.4f}")
plt.figure(figsize=(7, 5))
plt.loglog(steps[1:], mean_r2[1:], label="simulation")
plt.loglog(steps[1:], theory[1:], "--", label="slope 1")
plt.xlabel("number of steps, $N$")
plt.ylabel(r"$\langle r^2\rangle$")
plt.legend()
plt.show()
measured log-log slope = 1.0026
Lambda expressions and numerical integration¶
A lambda expression creates a short anonymous function. It is useful when a function is needed only once as an argument to another function.
For example, SciPy can verify the Gaussian integral
\begin{equation} \int_{-\infty}^{\infty} e^{-x^2}\,dx=\sqrt{\pi}. \end{equation}
from scipy import integrate
value, estimated_error = integrate.quad(
lambda x: np.exp(-x**2),
-np.inf,
np.inf,
)
print(f"numerical value = {value:.15f}")
print(f"exact value = {np.sqrt(np.pi):.15f}")
print(f"estimated error = {estimated_error:.3e}")
numerical value = 1.772453850905516 exact value = 1.772453850905516 estimated error = 1.420e-08
An optional object-oriented version¶
A class can combine the state of one walker with the operations that update it. Ordinary instance methods receive the object as their first argument, conventionally named self.
For a short calculation, the function-based version is often simpler. Classes become useful when many related operations and pieces of state must remain together.
class RandomWalker:
"""A two-dimensional square-lattice random walker."""
# Shared class attribute: unit moves right, left, up, and down.
_unit_moves = np.array([
[ 1.0, 0.0],
[-1.0, 0.0],
[ 0.0, 1.0],
[ 0.0, -1.0],
])
def __init__(self, step_length=1.0, seed=None):
"""Initialize a walker at the origin."""
self.step_length = step_length
self.rng = np.random.default_rng(seed)
self.x = 0.0
self.y = 0.0
def step(self):
"""Advance by one random lattice step."""
direction = self.rng.integers(0, 4)
displacement = self.step_length*self._unit_moves[direction]
self.x += displacement[0]
self.y += displacement[1]
return self.x, self.y
def run(self, nsteps):
"""Generate a trajectory and leave the walker at its final point."""
# Generate all directions and corresponding displacements at once.
directions = self.rng.integers(0, 4, size=nsteps)
displacements = self.step_length*self._unit_moves[directions]
initial_position = np.array([self.x, self.y])
trajectory = np.empty((nsteps + 1, 2))
trajectory[0] = initial_position
# trajectory[i] contains the initial point plus displacements 0,...,i-1.
trajectory[1:] = initial_position + np.cumsum(displacements, axis=0)
# Store the final point so that the next call continues this walk.
self.x = float(trajectory[-1, 0])
self.y = float(trajectory[-1, 1])
return trajectory
def reset(self, x=0.0, y=0.0):
"""Move the walker to a specified position."""
self.x = float(x)
self.y = float(y)
@property
def position(self):
"""Return the current position as a NumPy array."""
return np.array([self.x, self.y])
def __repr__(self):
return (
f"RandomWalker(x={self.x}, y={self.y}, "
f"step_length={self.step_length})"
)
walker = RandomWalker(step_length=0.5, seed=12)
trajectory = walker.run(20)
print(walker)
print(trajectory[-5:])
RandomWalker(x=1.5, y=-0.5, step_length=0.5) [[ 1.5 0.5] [ 2. 0.5] [ 2. 0. ] [ 1.5 0. ] [ 1.5 -0.5]]
From a notebook to a module¶
Reusable functions and classes should eventually be moved into a .py module. For example, a file named random_walk.py could contain random_walk_2d, mean_square_displacement, and RandomWalker. A notebook in the same directory could then use
from random_walk import random_walk_2d, mean_square_displacement
Python imports a module only once per kernel session. After editing an already imported module during interactive work, reload it explicitly:
import importlib
import random_walk
importlib.reload(random_walk)
In ordinary work, create modules with an editor or IDE. Jupyter's %%writefile magic can demonstrate file creation, but it overwrites an existing file with the same name.
Homework: drift and diffusion in a biased random walk¶
Modify the square-lattice walk so that the probabilities of moving right, left, up, and down are
\begin{equation} (p_R,p_L,p_U,p_D)=(0.40,0.10,0.25,0.25). \end{equation}
For a step length $\ell$, predict and measure:
- the mean position $\langle x_N\rangle$ and $\langle y_N\rangle$;
- the variances $\mathrm{Var}(x_N)$ and $\mathrm{Var}(y_N)$;
- how the mean displacement and the fluctuations scale with $N$.
For one step,
\begin{align} \langle\Delta x\rangle &= \ell(p_R-p_L),\\ \mathrm{Var}(\Delta x) &= \ell^2(p_R+p_L) -\langle\Delta x\rangle^2. \end{align}
Independent steps imply that the mean and variance after $N$ steps are $N$ times their one-step values.
Homework solution¶
rng.choice?
Signature: rng.choice(a, size=None, replace=True, p=None, axis=0, shuffle=True) Docstring: choice(a, size=None, replace=True, p=None, axis=0, shuffle=True) Generates a random sample from a given array Parameters ---------- a : {array_like, int} If an ndarray, a random sample is generated from its elements. If an int, the random sample is generated from np.arange(a). size : {int, tuple[int]}, optional Output shape. If the given shape is, e.g., ``(m, n, k)``, then ``m * n * k`` samples are drawn from the 1-d `a`. If `a` has more than one dimension, the `size` shape will be inserted into the `axis` dimension, so the output ``ndim`` will be ``a.ndim - 1 + len(size)``. Default is None, in which case a single value is returned. replace : bool, optional Whether the sample is with or without replacement. Default is True, meaning that a value of ``a`` can be selected multiple times. p : 1-D array_like, optional The probabilities associated with each entry in a. If not given, the sample assumes a uniform distribution over all entries in ``a``. axis : int, optional The axis along which the selection is performed. The default, 0, selects by row. shuffle : bool, optional Whether the sample is shuffled when sampling without replacement. Default is True, False provides a speedup. Returns ------- samples : single item or ndarray The generated random samples Raises ------ ValueError If a is an int and less than zero, if p is not 1-dimensional, if a is array-like with a size 0, if p is not a vector of probabilities, if a and p have different lengths, or if replace=False and the sample size is greater than the population size. See Also -------- integers, shuffle, permutation Notes ----- Setting user-specified probabilities through ``p`` uses a more general but less efficient sampler than the default. The general sampler produces a different sample than the optimized sampler even if each element of ``p`` is 1 / len(a). ``p`` must sum to 1 when cast to ``float64``. To ensure this, you may wish to normalize using ``p = p / np.sum(p, dtype=float)``. When passing ``a`` as an integer type and ``size`` is not specified, the return type is a native Python ``int``. Examples -------- Generate a uniform random sample from np.arange(5) of size 3: >>> rng = np.random.default_rng() >>> rng.choice(5, 3) array([0, 3, 4]) # random >>> #This is equivalent to rng.integers(0,5,3) Generate a non-uniform random sample from np.arange(5) of size 3: >>> rng.choice(5, 3, p=[0.1, 0, 0.3, 0.6, 0]) array([3, 3, 0]) # random Generate a uniform random sample from np.arange(5) of size 3 without replacement: >>> rng.choice(5, 3, replace=False) array([3,1,0]) # random >>> #This is equivalent to rng.permutation(np.arange(5))[:3] Generate a uniform random sample from a 2-D array along the first axis (the default), without replacement: >>> rng.choice([[0, 1, 2], [3, 4, 5], [6, 7, 8]], 2, replace=False) array([[3, 4, 5], # random [0, 1, 2]]) Generate a non-uniform random sample from np.arange(5) of size 3 without replacement: >>> rng.choice(5, 3, replace=False, p=[0.1, 0, 0.3, 0.6, 0]) array([2, 3, 0]) # random Any of the above can be repeated with an arbitrary array-like instead of just integers. For instance: >>> aa_milne_arr = ['pooh', 'rabbit', 'piglet', 'Christopher'] >>> rng.choice(aa_milne_arr, 5, p=[0.5, 0.1, 0.1, 0.3]) array(['pooh', 'pooh', 'pooh', 'Christopher', 'piglet'], # random dtype='<U11') Type: method
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
axes[0].plot(steps, mean_x, label=r"simulation $\langle x\rangle$")
axes[0].plot(steps, theory_mean_x, "--", label="theory")
axes[0].plot(steps, mean_y, label=r"simulation $\langle y\rangle$")
axes[0].plot(steps, theory_mean_y, "--", label="theory")
axes[0].set_xlabel("number of steps")
axes[0].set_ylabel("mean position")
axes[0].set_title("Drift")
axes[0].legend()
axes[1].plot(steps, var_x, label=r"simulation $\mathrm{Var}(x)$")
axes[1].plot(steps, theory_var_x, "--", label="theory")
axes[1].plot(steps, var_y, label=r"simulation $\mathrm{Var}(y)$")
axes[1].plot(steps, theory_var_y, "--", label="theory")
axes[1].set_xlabel("number of steps")
axes[1].set_ylabel("variance")
axes[1].set_title("Diffusive fluctuations")
axes[1].legend()
plt.tight_layout()
plt.show()
print(f"At N={nsteps}:")
print(f" <x>: simulation={mean_x[-1]:.3f}, theory={theory_mean_x[-1]:.3f}")
print(f" <y>: simulation={mean_y[-1]:.3f}, theory={theory_mean_y[-1]:.3f}")
print(f" Var(x): simulation={var_x[-1]:.3f}, theory={theory_var_x[-1]:.3f}")
print(f" Var(y): simulation={var_y[-1]:.3f}, theory={theory_var_y[-1]:.3f}")
At N=300: <x>: simulation=90.085, theory=90.000 <y>: simulation=-0.183, theory=0.000 Var(x): simulation=127.272, theory=123.000 Var(y): simulation=146.706, theory=150.000
The biased walk contains both kinds of motion:
- drift: the mean position grows linearly because $p_R\ne p_L$;
- diffusion: the variance also grows linearly because the individual steps remain random.
This distinction between systematic motion and fluctuations appears throughout statistical mechanics and transport theory.