Particle Life on GPU CuPy port for long-form renders

#!/usr/bin/env python3
"""GPU port of render_seed_480_panoramic.py using CuPy — optimised.

Streams raw uint8 frames directly into ffmpeg's stdin (no intermediate
raw file on disk) so very long renders are bounded by GPU memory, not
storage. At 540,000 frames × 720×480 the un-streamed raw bytes would
have been 186 GB — would not fit on /tmp.

Two top-level parameter sets at the bottom: one for the
panoramic-pan-piece (4K-wide aspect, short duration) and one for the
disc-fill target (TV native aspect, multi-hour duration).  Flip the
`MODE` constant to select.
"""
import os
import subprocess
import sys
from pathlib import Path
import numpy as np
import cupy as cp


try:
    cp.cuda.set_allocator(cp.cuda.MemoryAsyncPool().malloc)
except Exception:
    pass


# --- which render to do --------------------------------------------------
# 'panoramic'  : 7680×2160 × 4000 frames  (200 sec, pan-piece, ~2 GB)
# 'dvd_mp4'    : 720×480 × 540,000 frames (7.5 hours, data DVD-R 4.7 GB)
# 'dvd_video'  : 720×480 × 216,000 frames (2 hours, real DVD-Video, any player)
# 'bd_4k_6h'   : 3840×2160 × 432,000 frames (6 hours, fills BD-R 25 GB)
MODE = 'dvd_video'


_PRESETS = {
    'panoramic': dict(width=7680, height=2160, N=115000, iters=4000,
                      crf=23, aspect=None),
    # DVD-R MP4 fill (data disc, NOT DVD-Video).  ~1.4 Mbps avg h264.
    'dvd_mp4':   dict(width=720,  height=480,  N=2500,   iters=540_000,
                      bitrate='1400k', maxrate='2500k',
                      bufsize='5000k', aspect='16:9'),
    # Real DVD-Video — bootable in any DVD player + Xbox 360.
    # 720×480 NTSC, ~30 fps, MPEG-2 + silent AC-3 audio, DVD-PS muxer.
    # Uses ffmpeg's `-target ntsc-dvd` which fixes codec/bitrate/GOP/mux
    # to spec.  Output is a .mpg (MPEG-2 Program Stream); dvdauthor
    # turns that into the VIDEO_TS/ folder you burn to disc.
    # DVD-R single layer = 4.7 GB marketed = 4.38 GiB binary.  Target
    # final file at ~4.2 GiB so dvdauthor's IFO/BUP overhead still fits.
    # `-target ntsc-dvd` defaults to 6 Mbps video + 448 kbps audio
    # (5.8 GB total over 2 hr — too big).  Audio is SILENT so 192 kbps
    # is plenty; the savings go to the video budget.
    'dvd_video': dict(width=720, height=480, N=2500,
                      iters=216_000, fps=30,
                      dvd_target='ntsc-dvd',
                      video_bitrate='4400k',
                      audio_bitrate='192k',
                      aspect='16:9',
                      out_ext='.mpg', add_silent_audio=True),
    # BD-R fill (~25 GB) at 3840×2160, 6 hr → ~9.3 Mbps avg, HEVC.
    'bd_4k_6h':  dict(width=3840, height=2160, N=115000, iters=432_000,
                      bitrate='9000k', maxrate='15000k',
                      bufsize='30000k', aspect=None, codec='libx265'),
}

PARAMS = _PRESETS[MODE]


# --- splat kernel (one CUDA launch per frame) ----------------------------
_SPLAT_KERNEL = cp.RawKernel(r"""
extern "C" __global__
void splat(const float* __restrict__ pos,
           float* __restrict__ buf,
           const float* __restrict__ kern,
           const int N, const int H, const int W) {
    int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= N) return;
    float px = pos[p * 2 + 1];
    float py = pos[p * 2];
    int ix = ((int)px % W + W) % W;
    int iy = ((int)py % H + H) % H;
    #pragma unroll
    for (int dy = -2; dy <= 2; ++dy) {
        int ky = ((iy + dy) % H + H) % H;
        #pragma unroll
        for (int dx = -2; dx <= 2; ++dx) {
            int kx = ((ix + dx) % W + W) % W;
            float w = kern[(dy + 2) * 5 + (dx + 2)];
            atomicAdd(&buf[ky * W + kx], w);
        }
    }
}
""", 'splat')


_KERN_HOST = np.array([
    [0.05, 0.20, 0.30, 0.20, 0.05],
    [0.20, 0.60, 0.85, 0.60, 0.20],
    [0.30, 0.85, 1.00, 0.85, 0.30],
    [0.20, 0.60, 0.85, 0.60, 0.20],
    [0.05, 0.20, 0.30, 0.20, 0.05],
], dtype=np.float32)
KERN_GPU = cp.asarray(_KERN_HOST.ravel())

_DY_OFFSETS = cp.asarray(np.array([-1, -1, -1, 0, 0, 0, 1, 1, 1],
                                  dtype=np.int32))
_DX_OFFSETS = cp.asarray(np.array([-1, 0, 1, -1, 0, 1, -1, 0, 1],
                                  dtype=np.int32))


def render_points(buf, pos, S):
    buf *= 0.80
    H, W = int(S[0]), int(S[1])
    threads = 256
    blocks = (pos.shape[0] + threads - 1) // threads
    _SPLAT_KERNEL((blocks,), (threads,),
                  (pos.astype(cp.float32), buf, KERN_GPU,
                   pos.shape[0], H, W))
    buf_max = float(buf.max())
    bright = cp.clip(cp.log1p(buf) / cp.log1p(buf_max + 1e-9) * 255,
                     0, 255).astype(cp.uint8)
    return cp.asnumpy(bright)


@cp.fuse()
def _force_value(rn, Aij_pairs, beta):
    rep = rn / beta - 1.0
    att = Aij_pairs * (1.0 - cp.abs(2.0 * rn - 1.0 - beta) / (1.0 - beta))
    return cp.where(rn < beta, rep, cp.where(rn < 1.0, att, 0.0))


def particle_life_stream(width, height, iters, ffmpeg_stdin,
                          N=2000, K=5, seed=480):
    """Run the simulation; write each uint8 grayscale frame directly
    to the supplied open file/pipe (so encoding can happen in parallel
    with simulation and no large raw file is materialised)."""
    rng = np.random.default_rng(seed)
    S_host = np.array([height, width], dtype=np.float32)
    S = cp.asarray(S_host)

    pos = cp.asarray(rng.random((N, 2)).astype(np.float32)) * S
    vel = cp.zeros((N, 2), dtype=cp.float32)
    typ = cp.asarray(rng.integers(0, K, N).astype(np.int32))
    A = cp.asarray(rng.uniform(-1, 1, (K, K)).astype(np.float32))

    rmax = cp.float32(88.0)
    rmax_val = 88.0
    beta = cp.float32(0.3)
    friction = cp.float32(0.85)
    fs = cp.float32(0.8)
    dt = cp.float32(0.6)

    buf = cp.zeros((height, width), dtype=cp.float32)
    cell_size = rmax_val
    grid_h = int(np.ceil(height / cell_size))
    grid_w = int(np.ceil(width / cell_size))
    n_cells = grid_h * grid_w

    print(f"Rendering {iters} frames at {width}×{height}, N={N} on GPU...",
          flush=True)
    mem_pool = cp.get_default_memory_pool()

    import time
    t_start = time.time()
    last_report = t_start

    for t in range(iters):
        cell_y = (pos[:, 0] / cell_size).astype(cp.int32) % grid_h
        cell_x = (pos[:, 1] / cell_size).astype(cp.int32) % grid_w
        cell_idx = cell_y * grid_w + cell_x

        sort_order = cp.argsort(cell_idx)
        cell_idx_sorted = cell_idx[sort_order]

        all_cells = cp.arange(n_cells, dtype=cp.int32)
        cell_start_all = cp.searchsorted(cell_idx_sorted, all_cells,
                                          side='left')
        cell_end_all = cp.searchsorted(cell_idx_sorted, all_cells,
                                        side='right')
        cell_start = cp.where(cell_end_all > cell_start_all,
                              cell_start_all, -1)
        cell_end = cell_end_all

        ny = (cell_y[:, None] + _DY_OFFSETS[None, :]) % grid_h
        nx = (cell_x[:, None] + _DX_OFFSETS[None, :]) % grid_w
        neighbour_ids = (ny * grid_w + nx).astype(cp.int32)
        flat_ne = neighbour_ids.reshape(-1)
        starts = cell_start[flat_ne]
        ends = cell_end[flat_ne]
        counts = cp.where(starts >= 0, ends - starts, 0).astype(cp.int32)

        total = int(counts.sum())

        per_particle = counts.reshape(N, 9).sum(axis=1)
        idx_i = cp.repeat(cp.arange(N, dtype=cp.int32), per_particle)

        cum_counts = cp.concatenate(
            [cp.zeros(1, dtype=cp.int64),
             cp.cumsum(counts.astype(cp.int64))])
        positions = cp.arange(total, dtype=cp.int64)
        pair_idx = cp.searchsorted(cum_counts[1:], positions,
                                    side='right').astype(cp.int32)
        local_offset = (positions - cum_counts[pair_idx]).astype(cp.int32)
        sort_idx = starts[pair_idx] + local_offset
        idx_j = sort_order[sort_idx].astype(cp.int32)

        keep = idx_i != idx_j
        idx_i = idx_i[keep]
        idx_j = idx_j[keep]

        d = pos[idx_j] - pos[idx_i]
        cp.subtract(d[:, 0], S[0] * cp.round(d[:, 0] / S[0]), out=d[:, 0])
        cp.subtract(d[:, 1], S[1] * cp.round(d[:, 1] / S[1]), out=d[:, 1])
        dist = cp.sqrt((d * d).sum(axis=1))

        within = dist <= rmax
        idx_i = idx_i[within]
        idx_j = idx_j[within]
        d = d[within]
        dist = dist[within]

        rn = dist / rmax
        dirv = d / (dist[:, None] + cp.float32(1e-9))

        Aij_pairs = A[typ[idx_i], typ[idx_j]]
        F_val = _force_value(rn, Aij_pairs, beta)
        force_contrib = fs * F_val[:, None] * dirv

        acc = cp.zeros((N, 2), dtype=cp.float32)
        cp.add.at(acc, idx_i, force_contrib)

        vel = vel * friction + acc * dt
        pos = (pos + vel * dt) % S

        frame_host = render_points(buf, pos, S)
        ffmpeg_stdin.write(frame_host.tobytes())

        if (t & 0x1F) == 0:
            mem_pool.free_all_blocks()

        # Progress report every ~10 sec wall clock
        now = time.time()
        if now - last_report > 10:
            elapsed = now - t_start
            fps = (t + 1) / elapsed
            eta = (iters - t - 1) / fps
            print(f"  frame {t + 1}/{iters}  "
                  f"({(t + 1) / iters * 100:5.1f}%)  "
                  f"fps={fps:5.1f}  "
                  f"eta={eta / 60:5.1f} min", flush=True)
            last_report = now

    print(f"Completed {iters} frames in {time.time() - t_start:.0f} sec")


if __name__ == "__main__":
    seed = 480
    P = PARAMS
    width, height = P['width'], P['height']
    iters, N = P['iters'], P['N']
    outdir = (f"/home/fox/Downloads/py/output/even_more/particle_life/"
              f"{seed}_{MODE}")
    Path(outdir).mkdir(parents=True, exist_ok=True)
    fps = P.get('fps', 20)
    out_ext = P.get('out_ext', '.mp4')
    mp4_path = Path(outdir) / f"animation{out_ext}"

    print(f"=== Particle Life Rendering: MODE={MODE} ===")
    print(f"Seed: {seed}   Resolution: {width}×{height}   "
          f"Duration: {iters / fps:.1f} sec ({iters} frames @ {fps} fps)")
    print(f"N: {N}   GPU: "
          f"{cp.cuda.runtime.getDeviceProperties(0)['name'].decode()}")
    print()

    # ffmpeg command construction.  Two distinct flavours: real DVD-Video
    # (-target ntsc-dvd, MPEG-2 PS, silent AC-3 audio track) vs the
    # h264/x265 .mp4 path.  Both stream video frames from stdin so no
    # giant raw file is materialised.
    if P.get('dvd_target'):
        # DVD-Video pipeline — outputs a .mpg you feed to dvdauthor.
        cmd = [
            'ffmpeg', '-y',
            '-f', 'rawvideo', '-pix_fmt', 'gray',
            '-s', f'{width}x{height}',
            '-framerate', str(fps),
            '-i', '-',
        ]
        if P.get('add_silent_audio'):
            cmd += ['-f', 'lavfi',
                    '-i', 'anullsrc=channel_layout=stereo:sample_rate=48000']
        cmd += ['-target', P['dvd_target']]
        # Override video bitrate that -target ntsc-dvd would default to.
        # The DVD-Video max combined bitrate is 10.08 Mbps; we deliberately
        # stay well under so we hit a specific final file size.
        if P.get('video_bitrate'):
            cmd += ['-b:v', P['video_bitrate']]
        if P.get('audio_bitrate'):
            cmd += ['-b:a', P['audio_bitrate']]
        if P.get('aspect'):
            cmd += ['-aspect', P['aspect']]
        cmd += ['-shortest', str(mp4_path)]
    else:
        cmd = [
            'ffmpeg', '-y',
            '-f', 'rawvideo', '-pix_fmt', 'gray',
            '-s', f'{width}x{height}',
            '-framerate', str(fps),
            '-i', '-',
            '-c:v', P.get('codec', 'libx264'),
            '-pix_fmt', 'yuv420p',
            '-preset', 'medium',
        ]
        if P.get('bitrate'):
            cmd += ['-b:v', P['bitrate']]
            if P.get('maxrate'):
                cmd += ['-maxrate', P['maxrate'], '-bufsize', P['bufsize']]
        elif P.get('crf') is not None:
            cmd += ['-crf', str(P['crf'])]
        if P.get('aspect'):
            cmd += ['-aspect', P['aspect']]
        cmd += [str(mp4_path)]

    print(f"ffmpeg cmd: {' '.join(cmd)}")
    print()
    ff = subprocess.Popen(cmd, stdin=subprocess.PIPE,
                          stderr=subprocess.DEVNULL)
    try:
        particle_life_stream(width, height, iters, ff.stdin,
                             N=N, seed=seed)
    finally:
        ff.stdin.close()
        ff.wait()

    sz = mp4_path.stat().st_size if mp4_path.exists() else 0
    print(f"\nMP4 saved: {mp4_path} ({sz / 1024**3:.2f} GB, "
          f"{sz / 1024**2:.1f} MB)")
    print(f"\n=== Complete ===")