Collision prediction and resolution in Python
Computing when particles will next hit a wall or each other, then updating their velocities after the collision.
- File
collision_prediction_and_resolution.py, 158 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/collision-prediction-and-resolution/. The authors are listed in prov/F02YVH.jsonl.
"""Hard disks in the unit box: predict the next contact, resolve it elastically.
Each disk has a position (rx, ry), a velocity (vx, vy), a radius s > 0 and a
mass m > 0. The disks move in straight lines in the box [0, 1] x [0, 1]. The
predictions need disks that do not overlap and that lie inside the box.
"""
import heapq
import math
from dataclasses import dataclass
INF = math.inf
@dataclass
class Particle:
rx: float
ry: float
vx: float
vy: float
s: float = 0.01
m: float = 1.0
count: int = 0
def move(self, dt: float) -> None:
"""Advance the disk in a straight line for dt units of time."""
self.rx += self.vx * dt
self.ry += self.vy * dt
def kinetic_energy(self) -> float:
return 0.5 * self.m * (self.vx * self.vx + self.vy * self.vy)
def time_to_hit(p: Particle, q: Particle) -> float:
"""Return the time until p and q first touch, or INF if they never do.
Precondition: p and q do not overlap. Cost: O(1).
"""
if p is q: # @line L1
return INF
dx, dy = q.rx - p.rx, q.ry - p.ry # @line L2
dvx, dvy = q.vx - p.vx, q.vy - p.vy
b = dvx * dx + dvy * dy # @line L3
if b >= 0: # @line L4
return INF
a = dvx * dvx + dvy * dvy # @line L5
sigma = p.s + q.s
d = b * b - a * (dx * dx + dy * dy - sigma * sigma) # @line L6
if d < 0: # @line L7
return INF
return -(b + math.sqrt(d)) / a # @line L8
def time_to_hit_vertical_wall(p: Particle) -> float:
"""Time until p touches the wall x = 0 or x = 1, or INF. Needs s <= rx <= 1 - s."""
if p.vx > 0: # @line L9
return (1.0 - p.s - p.rx) / p.vx
if p.vx < 0: # @line L10
return (p.s - p.rx) / p.vx
return INF # @line L11
def time_to_hit_horizontal_wall(p: Particle) -> float:
"""Time until p touches the wall y = 0 or y = 1, or INF. Needs s <= ry <= 1 - s."""
if p.vy > 0: # @line L12
return (1.0 - p.s - p.ry) / p.vy
if p.vy < 0:
return (p.s - p.ry) / p.vy
return INF
def bounce_off(p: Particle, q: Particle) -> None:
"""Resolve an elastic collision of p and q, which touch. Cost: O(1)."""
dx, dy = q.rx - p.rx, q.ry - p.ry # @line L13
dvx, dvy = q.vx - p.vx, q.vy - p.vy
sigma = p.s + q.s
j = 2.0 * p.m * q.m * (dvx * dx + dvy * dy) / (sigma * (p.m + q.m)) # @line L14
jx, jy = j * dx / sigma, j * dy / sigma # @line L15
p.vx += jx / p.m # @line L16
p.vy += jy / p.m
q.vx -= jx / q.m # @line L17
q.vy -= jy / q.m
p.count += 1 # @line L18
q.count += 1
def bounce_off_vertical_wall(p: Particle) -> None:
p.vx = -p.vx # @line L19
p.count += 1
def bounce_off_horizontal_wall(p: Particle) -> None:
p.vy = -p.vy # @line L20
p.count += 1
# An event is (time, seq, a, b, count_a, count_b). The index a or b is -1
# for a wall: (i, -1) is a vertical wall, (-1, i) a horizontal wall.
Event = tuple[float, int, int, int, int, int]
def predict(
ps: list[Particle], i: int, t: float, limit: float, pq: list[Event], seq: list[int]
) -> None:
"""Push every future event of particle i that happens by time limit."""
if i < 0: # @line L21
return
p = ps[i]
def push(dt: float, a: int, b: int) -> None:
if dt < INF and t + dt <= limit:
ca = ps[a].count if a >= 0 else -1
cb = ps[b].count if b >= 0 else -1
seq[0] += 1
heapq.heappush(pq, (t + dt, seq[0], a, b, ca, cb))
for k in range(len(ps)): # @line L22
push(time_to_hit(p, ps[k]), i, k) # @line L23
push(time_to_hit_vertical_wall(p), i, -1) # @line L24
push(time_to_hit_horizontal_wall(p), -1, i)
def simulate(ps: list[Particle], limit: float) -> list[tuple[float, int, int]]:
"""Run the event-driven simulation of ps up to time limit.
Return the list of resolved events (time, a, b), with -1 for a wall as in
Event. On return, every particle is at its position at time limit.
Precondition: no two disks overlap and every disk is inside the box.
Cost: O(n^2 log n) to start, then O(n log n) per resolved event.
"""
pq: list[Event] = [] # @line L25
seq = [0]
t = 0.0
for i in range(len(ps)): # @line L26
predict(ps, i, t, limit, pq, seq)
done: list[tuple[float, int, int]] = []
while pq: # @line L27
te, _, a, b, ca, cb = heapq.heappop(pq) # @line L28
if (a >= 0 and ps[a].count != ca) or (b >= 0 and ps[b].count != cb): # @line L29
continue
for p in ps: # @line L30
p.move(te - t)
t = te # @line L31
if a >= 0 and b >= 0: # @line L32
bounce_off(ps[a], ps[b])
elif b < 0:
bounce_off_vertical_wall(ps[a])
else:
bounce_off_horizontal_wall(ps[b])
done.append((te, a, b))
predict(ps, a, t, limit, pq, seq) # @line L33
predict(ps, b, t, limit, pq, seq)
for p in ps: # @line L34
p.move(limit - t)
return done
test_collision_prediction_and_resolution.py 203 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/collision-prediction-and-resolution/. The authors are listed in prov/F02YVH.jsonl.
import math
import random
from collision_prediction_and_resolution import (
INF,
Particle,
bounce_off,
bounce_off_horizontal_wall,
bounce_off_vertical_wall,
simulate,
time_to_hit,
time_to_hit_horizontal_wall,
time_to_hit_vertical_wall,
)
TOL = 1e-9
def close(x: float, y: float, tol: float = TOL) -> bool:
return abs(x - y) <= tol * max(1.0, abs(x), abs(y))
def momentum(ps: list[Particle]) -> tuple[float, float]:
return sum(p.m * p.vx for p in ps), sum(p.m * p.vy for p in ps)
def energy(ps: list[Particle]) -> float:
return sum(p.kinetic_energy() for p in ps)
def example() -> tuple[Particle, Particle]:
return Particle(0.2, 0.4, 0.3, 0.0, 0.05, 1.0), Particle(0.78, 0.46, -0.2, 0.0, 0.05, 3.0)
def test_worked_example() -> None:
p, q = example()
assert close(time_to_hit(p, q), 1.0)
assert close(time_to_hit(q, p), 1.0)
assert close(time_to_hit_vertical_wall(p), 2.5)
assert close(time_to_hit_vertical_wall(q), 3.65)
assert time_to_hit_horizontal_wall(p) == INF
p.move(1.0)
q.move(1.0)
assert close(q.rx - p.rx, 0.08) and close(q.ry - p.ry, 0.06)
bounce_off(p, q)
assert close(p.vx, -0.18) and close(p.vy, -0.36)
assert close(q.vx, -0.04) and close(q.vy, 0.12)
assert p.count == 1 and q.count == 1
assert close(energy([p, q]), 0.105)
def test_worked_trace() -> None:
p, q = example()
done = simulate([p, q], 3.0)
assert [(a, b) for _, a, b in done] == [(0, 1), (-1, 0)]
assert close(done[0][0], 1.0) and close(done[1][0], 1.0 + 0.35 / 0.36)
assert close(p.rx, 0.14) and close(p.ry, 0.42)
assert close(q.rx, 0.50) and close(q.ry, 0.70)
assert p.count == 2 and q.count == 1
def test_head_on_equal_masses_swap() -> None:
p = Particle(0.3, 0.5, 0.7, 0.0, 0.05, 2.0)
q = Particle(0.7, 0.5, -0.1, 0.0, 0.05, 2.0)
dt = time_to_hit(p, q)
assert close(dt, 0.3 / 0.8)
p.move(dt)
q.move(dt)
bounce_off(p, q)
assert close(p.vx, -0.1) and close(q.vx, 0.7)
assert close(p.vy, 0.0) and close(q.vy, 0.0)
def test_no_collision_cases() -> None:
p = Particle(0.3, 0.5, 0.0, 0.0)
assert time_to_hit(p, p) == INF
# moving apart
assert time_to_hit(Particle(0.3, 0.5, -0.1, 0.0), Particle(0.6, 0.5, 0.1, 0.0)) == INF
# same velocity: dv = 0
assert time_to_hit(Particle(0.3, 0.5, 0.2, 0.1), Particle(0.6, 0.5, 0.2, 0.1)) == INF
# passing by: approach but miss (d < 0)
assert time_to_hit(Particle(0.2, 0.2, 0.1, 0.0, 0.01), Particle(0.8, 0.5, -0.1, 0.0, 0.01)) == INF
# tangent: graze exactly once
t = time_to_hit(Particle(0.2, 0.5, 0.1, 0.0, 0.01), Particle(0.8, 0.52, -0.1, 0.0, 0.01))
assert close(t, 3.0)
# at rest: no wall hit
assert time_to_hit_vertical_wall(p) == INF
assert time_to_hit_horizontal_wall(p) == INF
def test_walls() -> None:
p = Particle(0.5, 0.5, -0.25, 0.5, 0.1)
assert close(time_to_hit_vertical_wall(p), 1.6)
assert close(time_to_hit_horizontal_wall(p), 0.8)
bounce_off_vertical_wall(p)
bounce_off_horizontal_wall(p)
assert p.vx == 0.25 and p.vy == -0.5 and p.count == 2
def oracle_time(p: Particle, q: Particle) -> float:
"""First root of f(t) = |dr + t dv|^2 - sigma^2 by bisection, without the formula."""
dx, dy = q.rx - p.rx, q.ry - p.ry
dvx, dvy = q.vx - p.vx, q.vy - p.vy
sigma = p.s + q.s
def f(t: float) -> float:
return (dx + t * dvx) ** 2 + (dy + t * dvy) ** 2 - sigma * sigma
a = dvx * dvx + dvy * dvy
if a == 0:
return INF
tmin = -(dvx * dx + dvy * dy) / a # the time of the closest approach
if tmin <= 0 or f(tmin) > 0:
return INF
lo, hi = 0.0, tmin # f(lo) > 0 >= f(hi) and f decreases on [lo, hi]
for _ in range(200):
mid = (lo + hi) / 2
if f(mid) > 0:
lo = mid
else:
hi = mid
return hi
def random_pair(rng: random.Random) -> tuple[Particle, Particle]:
while True:
p = Particle(rng.random(), rng.random(), rng.uniform(-1, 1), rng.uniform(-1, 1),
rng.uniform(0.01, 0.2), rng.uniform(0.5, 5))
q = Particle(rng.random(), rng.random(), rng.uniform(-1, 1), rng.uniform(-1, 1),
rng.uniform(0.01, 0.2), rng.uniform(0.5, 5))
if math.hypot(q.rx - p.rx, q.ry - p.ry) > p.s + q.s:
return p, q
def test_time_to_hit_against_bisection() -> None:
rng = random.Random(61)
hits = 0
for _ in range(4000):
p, q = random_pair(rng)
got, want = time_to_hit(p, q), oracle_time(p, q)
if want == INF:
assert got == INF
else:
hits += 1
assert close(got, want, 1e-7)
assert got >= 0
assert hits >= 500
def test_resolution_conserves_momentum_and_energy() -> None:
rng = random.Random(62)
for _ in range(1000):
p, q = random_pair(rng)
dt = time_to_hit(p, q)
if dt == INF:
continue
p.move(dt)
q.move(dt)
mx, my = momentum([p, q])
e = energy([p, q])
dx, dy = q.rx - p.rx, q.ry - p.ry
dvn = (q.vx - p.vx) * dx + (q.vy - p.vy) * dy
dvt = (q.vx - p.vx) * -dy + (q.vy - p.vy) * dx
bounce_off(p, q)
mx2, my2 = momentum([p, q])
assert close(mx, mx2) and close(my, my2)
assert close(e, energy([p, q]))
assert close((q.vx - p.vx) * dx + (q.vy - p.vy) * dy, -dvn)
assert close((q.vx - p.vx) * -dy + (q.vy - p.vy) * dx, dvt)
assert time_to_hit(p, q) == INF
def test_simulation_keeps_disks_apart() -> None:
rng = random.Random(63)
for _ in range(20):
ps: list[Particle] = []
while len(ps) < 6:
s = rng.uniform(0.02, 0.06)
c = Particle(rng.uniform(s, 1 - s), rng.uniform(s, 1 - s),
rng.uniform(-1, 1), rng.uniform(-1, 1), s, rng.uniform(0.5, 2))
if all(math.hypot(c.rx - p.rx, c.ry - p.ry) > c.s + p.s for p in ps):
ps.append(c)
e = energy(ps)
done = simulate(ps, 5.0)
assert all(done[k][0] <= done[k + 1][0] for k in range(len(done) - 1))
assert close(e, energy(ps), 1e-7)
for i, p in enumerate(ps):
assert -1e-9 <= p.rx - p.s and p.rx + p.s <= 1 + 1e-9
assert -1e-9 <= p.ry - p.s and p.ry + p.s <= 1 + 1e-9
for q in ps[i + 1:]:
assert math.hypot(q.rx - p.rx, q.ry - p.ry) >= p.s + q.s - 1e-7
def test_empty_and_single() -> None:
assert simulate([], 1.0) == []
p = Particle(0.5, 0.5, 0.5, 0.0, 0.1)
done = simulate([p], 1.0)
assert [(a, b) for _, a, b in done] == [(0, -1)]
assert close(done[0][0], 0.8)
assert close(p.rx, 0.8) and p.vx == -0.5
What the next level needs
- the tests did not run in this build
Read the full page about Collision prediction and resolution