A Random Walk in ℤ⁵ - Vizualization with NumPy and Pillow

rwz5.py

from PIL import Image
import numpy as np
from typing import Annotated, TypeAlias, Tuple
import numpy.typing as npt

# %%

# five dimensional integer vector
Z5 = Annotated[npt.NDArray[np.int_], (5,)]
RGB: TypeAlias = Annotated[npt.NDArray[np.uint8], (3,)]
RGBTensor: TypeAlias = Annotated[npt.NDArray[np.uint8], (..., ..., 3)]

# %%

def random_walk_z5(
    width: int = 64,
    height: int = 64,
    strides: Tuple[int, int, int, int, int] = (1, 1, 3, 3, 3),
    n_iterations: int = 100_000,
    seed: int | None = None,
) -> Image.Image:
    """
    Visualizes a random walk in ℤ⁵, the five dimensional space of integers,
    by using two coordinates for position and the other three for color.
    """
    # validate arguments
    for n in (height, width, n_iterations) + strides:
        assert isinstance(n, int)
        assert n > 0
    if seed is not None:
        assert isinstance(seed, int)

    # initialize the random number generator with the given seed.
    rng: np.random.Generator = np.random.default_rng(seed)

    # we'll stay inside this cube by wrapping
    bounds: Z5 = np.array([height, width, 256, 256, 256], dtype=int)

    # mutable vector state, starting at a random point
    vector: Z5 = np.zeros(5, int)
    vector[:] = rng.integers(0, bounds, dtype=int)

    # image is initially black
    image: RGBTensor = np.zeros((height, width, 3), dtype=np.uint8)
    
    for iteration in range(n_iterations):
        # take a step in a random direction along a random axis
        axis = rng.integers(0, 5)
        stride = strides[axis]
        step = rng.choice((-stride, stride))
        vector[axis] += step
        vector %= bounds

        # interpret vector as (y, x, red, green, blue) and plot one pixel
        y: int = vector[0]
        x: int = vector[1]
        color: RGB = vector[2:].astype(np.uint8)
        image[y, x] = color
        
    return Image.fromarray(image)

# %%

# see a quick example scaled up x8
random_walk_z5().resize((512, 512), Image.Resampling.NEAREST)

# %%

# save to disk
image = random_walk_z5()
image.save("rwz5_64x64.png", format="PNG")

# %%

# a perfect cube of breathtaking symmetry
random_walk_z5(256, 256, strides=(1, 1, 1, 1, 1), seed=0, n_iterations=2**24)

# %%

# 37 is the most random number: https://www.youtube.com/watch?v=d6iQrh2TK98
# 100 million iterations all but guarentees few black pixels remain, but
# takes the better part of an hour to render.
random_walk_z5(512, 512, n_iterations=int(1e8), seed=37)

Example output image:
https://x.com/oranlooney/status/1979549474386321809/photo/1