2D Minkowski animation showing Lorentz transformation

minkowski.py

'''Minkowski 2D space-time animation showing Lorentz boost.

Requires Python 3.7, numpy and matplotlib, and an ffmpeg binary in your PATH.'''

import math

import numpy as np
import matplotlib
import matplotlib.animation as anim
from matplotlib import pyplot as plot

# Constants

vel_max = 0.5
V = visual_limit = 3

# We can simplify the Lorentz transformations with β=v/c
#      x' = γ(x - vt)      |   x'  = γ(x - βct)
#      t' = γ(t - vx/c²)   |   ct' = γ(ct - βx)
# This makes for a wonderfully orthogonal Lorentz matrix Λ(v)
# with lightspeed eigenvectors [±1, 1].

γ = lambda β: (1 - β**2) ** (-1/2)
Λ = lambda β: np.array([[1, -β], [-β, 1]]) * γ(β)

L = limit = V * int(γ(vel_max)) * 2
y = np.array((-L, L))

figure = plot.figure()
ax = plot.axes(xlim=(-V, V), ylim=(-V, V))

ax.set(xlabel='space / x', ylabel='time / ct')

# axis lines - extend four times as far
ax.plot(y*4, y*0, color='black', lw=0.8)
ax.plot(y*0, y*4, color='black', lw=0.8)

# lorentz transformable lines
lines = []

for grid_line in np.linspace(-L, L, L*2+1):
    static = [grid_line] * 2
    lines += ax.plot(static, y, color='silver', lw=0.25)
    lines += ax.plot(y, static, color='silver', lw=0.25)

# worldlines
for v in (-vel_max, 0, vel_max):
    lines += ax.plot(y*v, y, color='blue', lw=1)

for c in (-1, 1):
    lines += ax.plot(y*c, y, color='red', lw=1)

# save data at v=0 so lines can be modified each time
for eq in lines:
    eq.orig_data = eq.get_data()

FRAMES = 300
def animate(i):
    β = vel_max * math.sin(i*math.tau / FRAMES)
    ax.set(title=f'Minkowski graph with Lorentz boost v={β:+.02f}c')
    return [l.set_data(*np.matmul(Λ(β), l.orig_data)) for l in lines]

if __name__ == '__main__':
    an = anim.FuncAnimation(
        figure, animate, frames=FRAMES, interval=1000//60)
    
    an.save('minkowski.mp4', anim.FFMpegWriter(fps=60))
    plot.show()