Mandelbrot tile rendering

Partition the complex plane into tiles, compute each tile’s escape-iteration counts on a worker, and stitch them on the client. Shared viewport parameters are uploaded once with send_object.

[ ]:
# Connection settings -- edit these to point at your running scheduler.
SCHEDULER_ADDRESS = "ws://127.0.0.1:2345"  # supports tcp:// or ws://; only ws:// works from JupyterLite (browser)
OBJECT_STORAGE_ADDRESS = None  # leave None to use whatever the scheduler advertises

# Image / compute parameters. Defaults render a 4096x4096 image at 2000 iter into a
# deep-zoom pocket where most pixels need the full iteration budget; on 16 workers
# expect roughly a minute of wall-clock time.
TILES_PER_SIDE = 16               # final image is TILES_PER_SIDE * TILE_PIXELS on each side
TILE_PIXELS = 256
MAX_ITER = 2000
VIEWPORT = (-0.745, -0.735, 0.105, 0.115)  # xmin, xmax, ymin, ymax -- seahorse valley pocket
[ ]:
import time

from scaler import Client


def mandelbrot_tile(params: dict, tile_row: int, tile_col: int) -> bytes:
    """Worker-side: escape-iteration kernel for one tile, returned as packed bytes."""
    import numpy as np

    xmin, xmax, ymin, ymax = params["viewport"]
    tiles = params["tiles_per_side"]
    pixels = params["tile_pixels"]
    max_iter = params["max_iter"]

    tile_width = (xmax - xmin) / tiles
    tile_height = (ymax - ymin) / tiles
    x0 = xmin + tile_col * tile_width
    y0 = ymin + tile_row * tile_height

    xs = np.linspace(x0, x0 + tile_width, pixels, endpoint=False, dtype=np.float64)
    ys = np.linspace(y0, y0 + tile_height, pixels, endpoint=False, dtype=np.float64)
    c = xs[np.newaxis, :] + 1j * ys[:, np.newaxis]
    z = np.zeros_like(c)
    out = np.full(c.shape, max_iter, dtype=np.uint16)
    active = np.ones(c.shape, dtype=bool)
    for i in range(max_iter):
        z[active] = z[active] * z[active] + c[active]
        escaped = active & (z.real * z.real + z.imag * z.imag > 4.0)
        out[escaped] = i
        active &= ~escaped
        if not active.any():
            break
    return out.tobytes()


params = {
    "viewport": VIEWPORT,
    "tiles_per_side": TILES_PER_SIDE,
    "tile_pixels": TILE_PIXELS,
    "max_iter": MAX_ITER,
}

with Client(address=SCHEDULER_ADDRESS, object_storage_address=OBJECT_STORAGE_ADDRESS) as client:
    params_ref = client.send_object(params, name="mandelbrot-params")
    started = time.perf_counter()
    futures = {
        (row, col): client.submit(mandelbrot_tile, params_ref, row, col)
        for row in range(TILES_PER_SIDE)
        for col in range(TILES_PER_SIDE)
    }
    tile_bytes = {key: fut.result() for key, fut in futures.items()}
    elapsed = time.perf_counter() - started

print(f"rendered {TILES_PER_SIDE * TILES_PER_SIDE} tiles in {elapsed:.2f}s")
[ ]:
# Stitch and display -- client-side only and intentionally cheap.
import numpy as np
import matplotlib.pyplot as plt

side = TILES_PER_SIDE * TILE_PIXELS
image = np.empty((side, side), dtype=np.uint16)
for (row, col), blob in tile_bytes.items():
    tile = np.frombuffer(blob, dtype=np.uint16).reshape(TILE_PIXELS, TILE_PIXELS)
    image[row * TILE_PIXELS : (row + 1) * TILE_PIXELS, col * TILE_PIXELS : (col + 1) * TILE_PIXELS] = tile

fig, ax = plt.subplots(figsize=(6, 6))
ax.imshow(np.log1p(image), cmap="twilight_shifted", extent=VIEWPORT, origin="lower")
ax.set_title(f"Mandelbrot set ({side}x{side}, {MAX_ITER} iter)")
ax.set_xlabel("Re(c)")
ax.set_ylabel("Im(c)")
plt.show()