Untitled

#!/usr/bin/env python3

import time
from pathlib import Path

import numpy as np
from numba import njit
from PIL import Image


@njit
def clifford_chunk(a, b, c, d, n, x0, y0):
    xs = np.empty(n, dtype=np.float32)
    ys = np.empty(n, dtype=np.float32)

    x, y = x0, y0

    for i in range(n):
        x_new = np.sin(a * y) + c * np.cos(a * x)
        y_new = np.sin(b * x) + d * np.cos(b * y)

        x, y = x_new, y_new
        xs[i] = x
        ys[i] = y

    return xs, ys, x, y


def generate_points(a, b, c, d, n, chunks=20):
    xs_parts = []
    ys_parts = []

    x, y = 0.0, 0.0
    chunk_size = n // chunks
    t0 = time.perf_counter()

    for chunk in range(chunks):
        n_chunk = chunk_size if chunk < chunks - 1 else n - chunk_size * (chunks - 1)

        xs, ys, x, y = clifford_chunk(a, b, c, d, n_chunk, x, y)

        xs_parts.append(xs)
        ys_parts.append(ys)

        done = sum(len(arr) for arr in xs_parts)
        frac = done / n
        elapsed = time.perf_counter() - t0
        eta = elapsed / frac - elapsed

        print(
            f"{done:,}/{n:,} points "
            f"({frac * 100:5.1f}%) | "
            f"elapsed {elapsed:7.1f}s | "
            f"ETA {eta:7.1f}s",
            flush=True,
        )

    return np.concatenate(xs_parts), np.concatenate(ys_parts)


def save_density_png(
    xs,
    ys,
    out_file,
    width=4962,      # A4 width at 600 dpi
    height=7014,     # A4 height at 600 dpi
    xmin=-2.2,
    xmax=2.2,
    ymin=-2.2,
    ymax=2.2,
    background=(239, 227, 211),
    gamma=0.75,
):
    print("\nRasterizing density image...")
    t0 = time.perf_counter()

    ix = ((xs - xmin) / (xmax - xmin) * (width - 1)).astype(np.int32)
    iy = ((ys - ymin) / (ymax - ymin) * (height - 1)).astype(np.int32)

    mask = (ix >= 0) & (ix < width) & (iy >= 0) & (iy < height)

    hist = np.zeros((height, width), dtype=np.uint32)

    # Flip y so image coordinates match normal plot orientation
    np.add.at(hist, (height - 1 - iy[mask], ix[mask]), 1)

    img = np.log1p(hist.astype(np.float32))

    if img.max() > 0:
        img /= img.max()

    # Gamma controls contrast.
    # Smaller gamma = darker/fuller image.
    img = img ** gamma

    bg = np.array(background, dtype=np.float32)
    fg = np.array((0, 0, 0), dtype=np.float32)

    strength = img[..., None]
    rgb = bg * (1.0 - strength) + fg * strength
    rgb = np.clip(rgb, 0, 255).astype(np.uint8)

    Image.fromarray(rgb, mode="RGB").save(out_file)

    print(f"PNG rasterization time: {time.perf_counter() - t0:.2f}s")
    print(f"Saved: {out_file}")


def main():
    # Clifford parameters
    a = -1.4
    b = 1.6
    c = 1.0
    d = 0.7

    n = 10_000_000
    chunks = 20

    output_file = Path("clifford_A4_600dpi_density.png")

    print("Compiling Numba...")
    t_compile = time.perf_counter()
    clifford_chunk(a, b, c, d, 10, 0.0, 0.0)
    print(f"Compile time: {time.perf_counter() - t_compile:.2f}s\n")

    print("Generating points...")
    t_gen = time.perf_counter()
    xs, ys = generate_points(a, b, c, d, n, chunks)
    gen_time = time.perf_counter() - t_gen

    print(f"\nGeneration time: {gen_time:.2f}s")
    print(f"Point memory: {(xs.nbytes + ys.nbytes) / 1024**2:.1f} MB")

    save_density_png(xs, ys, output_file)

    print(f"\nTotal excluding compile: {time.perf_counter() - t_gen:.2f}s")


if __name__ == "__main__":
    main()