NumPy

NumPy#

Recitation Goals#

  • Represent scientific data with NumPy arrays.

  • Use vectorized calculations instead of handling one measurement at a time.

  • Practice axis, Boolean masks, reductions, and broadcasting in a physical context.

  • Build a simple particle model of an ideal gas and compare its trends with the gas laws.

In the next recitation we will use Matplotlib to plot the simulation results.

Importing NumPy#

NumPy is usually imported with the short name np.

import numpy as np

Warm-up: Chemical Measurements as Arrays#

Suppose we measured absorbance for several known solution concentrations. Instead of storing each measurement in a separate variable, we store the whole series in one array.

concentration_mM = np.array([0.00, 0.20, 0.40, 0.60, 0.80, 1.00])
absorbance = np.array([0.01, 0.15, 0.30, 0.44, 0.61, 0.74])

print(concentration_mM)
print(absorbance)
print(type(absorbance), absorbance.shape)

[0.  0.2 0.4 0.6 0.8 1. ]
[0.01 0.15 0.3  0.44 0.61 0.74]
<class 'numpy.ndarray'> (6,)

Elementwise Operations#

An operation between an array and a number is applied to every element. For example, converting from mM to M:

concentration_M = concentration_mM / 1000
print(concentration_M)

[0.     0.0002 0.0004 0.0006 0.0008 0.001 ]

Boolean Masks#

A Boolean mask is an array of True/False values. We can use it to select only measurements that satisfy a condition.

valid = absorbance > 0.05
print(valid)
print(concentration_mM[valid])
print(absorbance[valid])

[False  True  True  True  True  True]
[0.2 0.4 0.6 0.8 1. ]
[0.15 0.3  0.44 0.61 0.74]

Reductions: Mean, Standard Deviation, and Extremes#

Functions such as np.mean, np.std, np.min, and np.max reduce an array to a smaller result, often a single number. These operations are common in measurement analysis.

print('mean absorbance:', np.mean(absorbance))
print('std absorbance:', np.std(absorbance))
print('max absorbance:', np.max(absorbance))

mean absorbance: 0.375
std absorbance: 0.25237207980810134
max absorbance: 0.74

From Measurements to a Model: An Ideal Gas in a 2D Box#

We will build a simple model: point particles move in a square box and collide elastically with the walls. The purpose is to practice NumPy with a physical system whose ingredients are familiar from chemistry.

We use dimensionless units:

  • Each particle has mass m = 1 unless stated otherwise.

  • Boltzmann’s constant is k_B = 1.

  • The box is two-dimensional.

We will construct the simulation from several short functions. Each function represents one physical operation.

Random Numbers and Thermal Velocities#

A gas at temperature \(T\) does not have one fixed velocity. In each spatial direction, the particle velocities are normally distributed around zero. With \(k_B=1\), the standard deviation is

\[ \sigma = \sqrt{\frac{T}{m}}. \]

np.random.default_rng(seed) creates a random-number generator. The seed makes the results reproducible. rng.normal(...) then samples the velocity components.

The \(x\) and \(y\) components are normally distributed; the resulting particle speeds follow the Maxwell-Boltzmann speed distribution.

def generate_velocities(N, temperature, mass, rng):
    sigma = np.sqrt(temperature / mass)
    velocities = rng.normal(0, sigma, size=(N, 2))
    return velocities

Test the function with a small number of particles. Each row represents one particle, and the two columns contain \(v_x\) and \(v_y\).

rng = np.random.default_rng(7)
velocities = generate_velocities(N=5, temperature=2.0, mass=1.0, rng=rng)

print(velocities)
print('shape:', velocities.shape)

[[ 1.73969956e-03  4.22489991e-01]
 [-3.87689473e-01 -1.25948706e+00]
 [-6.43001591e-01 -1.40240001e+00]
 [ 8.50558985e-02  1.89535058e+00]
 [-6.96085134e-01 -8.77484018e-01]]
shape: (5, 2)

Representing Particle Positions#

Positions also use an array of shape (N, 2): one row per particle, with one column for \(x\) and one for \(y\).

row

meaning

column 0

column 1

0

particle 0

\(x_0\)

\(y_0\)

1

particle 1

\(x_1\)

\(y_1\)

2

particle 2

\(x_2\)

\(y_2\)

positions[:, 0] selects all \(x\) coordinates, while positions[:, 1] selects all \(y\) coordinates.

positions = np.array([
    [2.0, 8.0],
    [4.5, 1.0],
    [9.0, 5.5],
])

print('x coordinates:', positions[:, 0])
print('y coordinates:', positions[:, 1])

x coordinates: [2.  4.5 9. ]
y coordinates: [8.  1.  5.5]

Short Exercise: Selecting Particles#

Create Boolean masks that identify:

  1. Particles in the left half of a box of length 10.

  2. Particles in the upper half of the box.

box_length = 10.0

# Complete these expressions.
in_left_half = positions[:, 0] < 5.0
in_upper_half = positions[:, 1] > 5.0

print('left half:', in_left_half)
print('upper half:', in_upper_half)

left half: [ True  True False]
upper half: [ True False  True]

Moving the Particles#

For a short time interval dt, motion at constant velocity is

\[ \vec{r}_{\mathrm{new}} = \vec{r}_{\mathrm{old}} + \vec{v}\,dt. \]

The position and velocity arrays have the same shape, so NumPy performs this calculation for every coordinate of every particle.

def move_particles(positions, velocities, dt):
    new_positions = positions + velocities * dt
    return new_positions

Use a simple example whose result can be predicted before running the code.

test_positions = np.array([
    [1.0, 2.0],
    [5.0, 4.0],
])

test_velocities = np.array([
    [2.0, 0.0],
    [-1.0, 3.0],
])

moved_positions = move_particles(test_positions, test_velocities, dt=0.5)
print(moved_positions)

[[2.  2. ]
 [4.5 5.5]]

Detecting Wall Collisions#

A particle crossed a wall if one of its coordinates is outside the interval from 0 to box_length. We test each wall separately and store the four Boolean masks in a dictionary.

def find_wall_collisions(positions, box_length):
    collisions = {
        'left': positions[:, 0] < 0,
        'right': positions[:, 0] > box_length,
        'bottom': positions[:, 1] < 0,
        'top': positions[:, 1] > box_length,
    }
    return collisions

positions_after_move = np.array([
    [-0.2, 4.0],
    [10.3, 7.0],
    [5.0, 10.1],
    [6.0, 3.0],
])

collisions = find_wall_collisions(positions_after_move, box_length=10.0)
for wall in collisions:
    print(wall, collisions[wall])

left [ True False False False]
right [False  True False False]
bottom [False False False False]
top [False False  True False]

Momentum Transferred to the Walls#

For an elastic collision with a vertical wall, \(v_x\) changes to \(-v_x\). The magnitude of the momentum transferred to that wall is

\[ |\Delta p_x| = 2m|v_x|. \]

For a horizontal wall, the same expression uses \(v_y\). Boolean masks select only the particles that hit each wall.

def wall_momentum_transfer(velocities, collisions, mass):
    hit_vertical_wall = collisions['left'] | collisions['right']
    hit_horizontal_wall = collisions['bottom'] | collisions['top']

    vertical_momentum = 2 * mass * np.sum(
        np.abs(velocities[hit_vertical_wall, 0])
    )
    horizontal_momentum = 2 * mass * np.sum(
        np.abs(velocities[hit_horizontal_wall, 1])
    )

    return vertical_momentum + horizontal_momentum

Short Exercise: Read the Masks#

Before running the next cell, identify which velocity components contribute to the transferred momentum. A collision with a vertical wall uses column 0; a collision with a horizontal wall uses column 1.

test_velocities = np.array([
    [-2.0, 1.0],
    [3.0, -1.0],
    [0.5, 4.0],
    [1.0, 1.0],
])

momentum = wall_momentum_transfer(test_velocities, collisions, mass=1.0)
print(momentum)

18.0

Reflecting Particles from the Walls#

After detecting a collision, we must correct both the position and the relevant velocity component. The function works on copies, so the input arrays are not modified.

We handle the \(x\) and \(y\) directions explicitly. This is longer than a compact NumPy expression, but each line corresponds directly to one wall.

def reflect_from_walls(positions, velocities, box_length, collisions):
    left = collisions['left']
    right = collisions['right']
    bottom = collisions['bottom']
    top = collisions['top']

    hit_x_wall = left | right
    hit_y_wall = bottom | top

    new_positions = positions.copy()
    new_velocities = velocities.copy()

    # Vertical walls affect x and vx.
    new_positions[left, 0] = -positions[left, 0]
    new_positions[right, 0] = 2 * box_length - positions[right, 0]
    new_velocities[hit_x_wall, 0] = -velocities[hit_x_wall, 0]

    # Horizontal walls affect y and vy.
    new_positions[bottom, 1] = -positions[bottom, 1]
    new_positions[top, 1] = 2 * box_length - positions[top, 1]
    new_velocities[hit_y_wall, 1] = -velocities[hit_y_wall, 1]

    return new_positions, new_velocities

reflected_positions, reflected_velocities = reflect_from_walls(
    positions_after_move,
    test_velocities,
    box_length=10.0,
    collisions=collisions,
)

print(reflected_positions)
print(reflected_velocities)

[[0.2 4. ]
 [9.7 7. ]
 [5.  9.9]
 [6.  3. ]]
[[ 2.   1. ]
 [-3.  -1. ]
 [ 0.5 -4. ]
 [ 1.   1. ]]

Combining the Operations into One Step#

The detailed work is now contained in short functions that we have discussed. simulation_step only states the order of operations:

  1. Move.

  2. Detect collisions.

  3. Measure momentum transfer.

  4. Reflect.

def simulation_step(positions, velocities, box_length, dt, mass):
    moved_positions = move_particles(positions, velocities, dt)
    collisions = find_wall_collisions(moved_positions, box_length)
    momentum = wall_momentum_transfer(velocities, collisions, mass)
    new_positions, new_velocities = reflect_from_walls(
        moved_positions, velocities, box_length, collisions
    )
    return new_positions, new_velocities, momentum

Creating the Initial State#

The initial positions are uniformly distributed throughout the box. The initial velocities are generated at the requested temperature.

The same random-number generator is used for both operations so that one seed completely determines the initial state.

def initial_state(N, box_length, temperature, mass, seed):
    rng = np.random.default_rng(seed)
    positions = rng.uniform(0, box_length, size=(N, 2))
    velocities = generate_velocities(N, temperature, mass, rng)
    return positions, velocities

positions, velocities = initial_state(
    N=5,
    box_length=10.0,
    temperature=2.0,
    mass=1.0,
    seed=7,
)

print(positions)
print(velocities)

[[6.25095467 8.97213801]
 [7.7568569  2.2520719 ]
 [3.00166285 8.73553445]
 [0.05265305 8.21228418]
 [7.97069429 4.67934953]]
[[ 0.69274127  0.50471445]
 [ 0.14907826 -1.31588053]
 [-0.04136832  0.98330721]
 [-1.90100644 -0.64716642]
 [-2.68873498 -1.82368176]]

Why Use a Class for the Particle System?#

The helper functions above describe individual operations: generating velocities, moving particles, finding collisions, measuring momentum transfer, and reflecting from walls. Each function receives arrays and returns a result. This makes the NumPy operations easy to inspect and test separately.

A complete simulation, however, also has state that persists from one step to the next:

  • the current particle positions,

  • the current particle velocities,

  • the box length and particle mass,

  • the time step,

  • the elapsed simulation time,

  • the total momentum transferred to the walls.

These quantities belong together. A ParticleSystem object stores them as attributes, so we do not need to pass the same values between many function calls. Its methods describe operations on that state:

  • step() advances the system once,

  • run(steps) repeats that operation,

  • calculate_pressure() uses the accumulated measurement.

The class coordinates the physical model; the shorter helper functions still contain the detailed NumPy calculations.

The ParticleSystem Class#

The constructor creates the initial state and stores the fixed simulation parameters. It also initializes two measured quantities to zero:

  • elapsed_time records how long the system has evolved.

  • total_wall_momentum accumulates the momentum transferred during all wall collisions.

During step(), positions and velocities change, while the box length, particle mass, and dt remain fixed. This distinction between changing state and fixed parameters is one reason a class fits this problem naturally.

class ParticleSystem:
    def __init__(
        self,
        N,
        box_length,
        temperature,
        mass,
        dt,
        seed,
    ):
        self.box_length = box_length
        self.mass = mass
        self.dt = dt

        self.positions, self.velocities = initial_state(
            N, box_length, temperature, mass, seed
        )

        self.elapsed_time = 0.0
        self.total_wall_momentum = 0.0

    def step(self):
        self.positions, self.velocities, momentum = simulation_step(
            self.positions,
            self.velocities,
            self.box_length,
            self.dt,
            self.mass,
        )

        self.total_wall_momentum += momentum
        self.elapsed_time += self.dt

    def run(self, steps):
        for _ in range(steps):
            self.step()

    def calculate_pressure(self):
        perimeter = 4 * self.box_length
        pressure = self.total_wall_momentum / (
            self.elapsed_time * perimeter
        )
        return pressure

Reading the Class as a Physical Model#

The class methods correspond to a sequence of physical ideas:

  1. step() advances all particles by one time interval, processes wall collisions, and adds their momentum transfer to the running total.

  2. run(steps) repeats the same physical update. It does not need to know the details of collisions because those details are already handled by step() and the helper functions.

  3. calculate_pressure() divides total momentum transfer by elapsed time and wall length.

For a square box, the total wall length is its perimeter, \(4L\). Therefore our two-dimensional pressure estimate is

\[ P = \frac{\text{total momentum transferred}}{\text{elapsed time} \times 4L}. \]

Pressure is meaningful only after time has elapsed. Calling calculate_pressure() immediately after constructing the object would divide by zero, so the intended order is: create the system, run it, then calculate pressure.

Basic Run#

The main program is now short because the object keeps track of its own state. The names of the commands also describe the scientific workflow directly.

gas = ParticleSystem(
    N=300,
    box_length=10.0,
    temperature=1.5,
    mass=1.0,
    dt=0.01,
    seed=4,
)

gas.run(steps=4000)
pressure = gas.calculate_pressure()

print('elapsed time:', gas.elapsed_time)
print('box area:', gas.box_length**2)
print('total wall momentum:', gas.total_wall_momentum)
print('measured pressure:', pressure)

elapsed time: 40.00000000000061
box area: 100.0
total wall momentum: 7854.335090773291
measured pressure: 4.908959431733232