TheBird

main.py

9.8 kB · python · 267 lines

1import hashlib2import json3import os4import random5import shutil6import time7from datetime import datetime, timezone8from typing import Any, Dict, List9import mrlypy.life10import numpy as np11from mrlypy.core.state import choice, seed as seed_state12from mrlypy.life.crop import crop_grids13from mrlypy.life.enums import Fate14from mrlypy.two import Cell2d15from config import ATTEMPTS, DATA_DIR, FLASHES, FPS, FRAMES_DIR, POSTS, HEATMAP_DIR, HEATMAP_FPS, INDEX, LIVE_DAYS, MANIFEST, MASKS_DIR, MAX_GENERATIONS, MAX_SEGMENTS, MIN_GENERATIONS, POSTER, RATE, VERSION, files16from frames import create_saga_frames, create_saga_masks, create_saga_poster, frame_path17from heatmap import create_saga_heatmap18from models import Saga, Task19from setup import setup_saga, setup_segment20from video import create_saga_videos2122FFMPEG = shutil.which("ffmpeg") or "/opt/bin/ffmpeg"2324# NAME2526def name_for(seed: int) -> str:27    return hashlib.sha256(str(seed).encode()).hexdigest()[:8]2829# SEGMENT3031def _segment_life_config(task: Task) -> mrlypy.life.Config:32    padding = (task.canvas_unit_width - task.tile.grid_unit_width) // 233    return mrlypy.life.Config(34        max_generations=MAX_GENERATIONS,35        birth_counts=task.birth_counts,36        survive_counts=task.survive_counts,37        boundary=task.boundary,38        padding=padding,39        grid_size=task.tile.grid_size,40    )4142def _alive(grids: List[Cell2d]) -> List[Cell2d]:43    for i, grid in enumerate(grids):44        if not np.any(grid.types):45            return grids[:i]46    return grids4748def _generate_segment(saga: Saga, index: int, prev_grid: Cell2d) -> Task:49    task = setup_segment(saga, index, prev_grid)50    config = _segment_life_config(task)51    result = mrlypy.life.animate(config, grid=task.tile.cell, mask=task.mask.cell)52    alive = _alive(result.grids)53    if len(alive) < len(result.grids):54        result.grids = alive55        result.fate = Fate.DEAD56    result.count = len(result.grids)57    task.result = result58    task.count = result.count59    print(f"Segment {index} animated ({result.fate.value}, count={result.count}).")60    return task6162def _length_options(count: int) -> List[int]:63    lo = min(MIN_GENERATIONS, count)64    lo_mul4 = ((lo + 3) // 4) * 465    hi_mul4 = (count // 4) * 466    if hi_mul4 < lo_mul4:67        return []68    return list(range(lo_mul4, hi_mul4 + 1, 4))6970def _pivot_length(segment: Task) -> int:71    options = _length_options(segment.count)72    if not options:73        return segment.count74    return choice(options)7576def _truncate_segment(segment: Task, length: int) -> Task:77    segment.result.grids = segment.result.grids[:length]78    segment.count = len(segment.result.grids)79    segment.result.count = segment.count80    return segment8182# SAGA8384def _pad_grids(grids: List[Cell2d], size: int) -> List[Cell2d]:85    current = grids[0].types.shape[0]86    if current >= size:87        return grids88    before = (size - current) // 289    after = size - current - before90    return [Cell2d(types=np.pad(g.types, ((before, after), (before, after)), mode="constant", constant_values=0)) for g in grids]9192def _finalize_saga(saga: Saga, attempts: int) -> Saga:93    all_grids = []94    for seg in saga.segments:95        all_grids.extend(seg.result.grids)96    largest_mask = max(seg.mask.cell.types.shape[0] for seg in saga.segments)97    saga.grids = _pad_grids(crop_grids(all_grids), largest_mask)98    saga.segment_lengths = [len(s.result.grids) for s in saga.segments]99    saga.count = len(saga.grids)100    saga.time = round(sum((s.result.time or 0.0) for s in saga.segments), 2)101    saga.fate = saga.segments[-1].result.fate if saga.segments else None102    saga.attempts = attempts103    return saga104105def generate_saga(seed: int, key: str) -> Saga:106    for attempt in range(1, ATTEMPTS + 1):107        saga = setup_saga(Saga(), seed, key)108        prev_grid = None109        for index in range(MAX_SEGMENTS):110            segment = _generate_segment(saga, index, prev_grid)111            if segment.count < MIN_GENERATIONS:112                print(f"Segment {index} ran only {segment.count} generations, retrying.")113                break114            saga.segments.append(segment)115            if segment.result.fate == Fate.LIFE:116                print(f"Saga reached LIFE on attempt {attempt}.")117                return _finalize_saga(saga, attempt)118            _truncate_segment(segment, _pivot_length(segment))119            prev_grid = segment.result.grids[-1].copy()120        print(f"Attempt {attempt} found no LIFE, retrying.")121    raise RuntimeError(f"seed {seed} found no LIFE in {ATTEMPTS} attempts")122123# STORY124125def _count(n: int, word: str) -> str:126    return f"{n} {word}" if n == 1 else f"{n} {word}s"127128def _story(saga: Saga) -> str:129    ways = " then ".join(dict.fromkeys(s.way.value for s in saga.segments))130    cells = int(saga.grids[0].types.shape[0])131    return (f"{_count(len(saga.segments), 'segment')}, {_count(saga.count, 'generation')} "132            f"on a {cells} cell grid, {saga.boundary.value} boundary, {ways}, ending {saga.fate.value}")133134# MANIFEST135136def _segment_rows(saga: Saga) -> List[Dict[str, Any]]:137    rows = []138    for index, (seg, length) in enumerate(zip(saga.segments, saga.segment_lengths)):139        rows.append({140            "index": index,141            "key": seg.key,142            "way": seg.way.value if seg.way else None,143            "path": seg.path.value if seg.path else None,144            "secondary": seg.secondary.value if seg.secondary else None,145            "fate": seg.result.fate.value if seg.result else None,146            "generations": length,147            "tile": seg.tile.to_dict() if seg.tile else None,148            "mask": seg.mask.to_dict() if seg.mask else None,149            "music": seg.music.to_dict() if seg.music else None,150        })151    return rows152153def _sizes(out_dir: str, key: str) -> Dict[str, int]:154    return {name: os.path.getsize(f"{out_dir}/{name}") for name in files(key) if not name.endswith(MANIFEST)}155156def _manifest(saga: Saga, at: str, videos: Dict[str, Any], out_dir: str) -> Dict[str, Any]:157    first = saga.segments[0]158    return {159        "v": VERSION,160        "name": saga.key,161        "seed": saga.seed,162        "at": at,163        "story": _story(saga),164        "fps": FPS,165        "flashes": FLASHES,166        "heatmap_fps": HEATMAP_FPS,167        "rate": RATE,168        "canvas": int(saga.grids[0].types.shape[0]),169        "canvas_unit_width": saga.canvas_unit_width,170        "canvas_unit_height": saga.canvas_unit_height,171        "boundary": saga.boundary.value if saga.boundary else None,172        "primary": saga.primary.value if saga.primary else None,173        "fate": saga.fate.value if saga.fate else None,174        "attempts": saga.attempts,175        "tile": first.tile.to_dict() if first.tile else None,176        "mask": first.mask.to_dict() if first.mask else None,177        "segments": _segment_rows(saga),178        "generations": saga.count,179        "poster_generation": poster_frame(saga),180        "frames": videos["frames"],181        "duration": videos["duration"],182        "size": videos["size"],183        "steps": videos["steps"],184        "sizes": _sizes(out_dir, saga.key),185        "files": files(saga.key),186    }187188def index_row(manifest: Dict[str, Any]) -> Dict[str, Any]:189    return {190        "name": manifest["name"],191        "seed": manifest["seed"],192        "at": manifest["at"],193        "duration": manifest["duration"],194        "frames": manifest["frames"],195        "size": manifest["size"],196        "segments": len(manifest["segments"]),197        "canvas": manifest["canvas"],198        "story": manifest["story"],199    }200201def expired(rows: List[Dict[str, Any]], now: datetime) -> List[Dict[str, Any]]:202    cutoff = now.timestamp() - LIVE_DAYS * 86400203    return [item for item in rows if datetime.fromisoformat(item["at"].replace("Z", "+00:00")).timestamp() < cutoff]204205def write_index(root: str, row: Dict[str, Any]) -> str:206    path = os.path.join(root, INDEX)207    rows = []208    if os.path.exists(path):209        with open(path) as handle:210            rows = [item for item in json.load(handle) if item.get("name") != row["name"]]211    for item in expired(rows, datetime.now(timezone.utc)):212        shutil.rmtree(os.path.join(root, POSTS, item["name"]), ignore_errors=True)213        rows.remove(item)214        print(f"reap {item['name']}")215    rows.insert(0, row)216    with open(path, "w") as handle:217        json.dump(rows, handle)218    return path219220# POSTER221222def poster_frame(saga: Saga) -> int:223    longest = max(range(len(saga.segment_lengths)), key=lambda i: saga.segment_lengths[i])224    return sum(saga.segment_lengths[:longest + 1])225226# MAKE227228def make(seed: int, root: str) -> Dict[str, Any]:229    marks = {}230    clock = time.time()231232    def mark(stage):233        nonlocal clock234        now = time.time()235        marks[stage] = int((now - clock) * 1000)236        clock = now237        print(f"stage {stage} {marks[stage]} ms")238239    at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")240    key = name_for(seed)241    seed_state(seed)242    saga = generate_saga(seed, key)243    mark("saga")244    work = os.path.join(root, "work", key)245    out = os.path.join(root, POSTS, key)246    shutil.rmtree(work, ignore_errors=True)247    create_saga_frames(saga, f"{work}/{FRAMES_DIR}")248    create_saga_masks(saga, f"{work}/{MASKS_DIR}")249    mark("frames")250    create_saga_heatmap(saga, f"{work}/{HEATMAP_DIR}")251    mark("heatmap")252    videos = create_saga_videos(saga, work, out, FFMPEG)253    mark("videos")254    create_saga_poster(frame_path(f"{work}/{HEATMAP_DIR}", key, poster_frame(saga)), f"{out}/{key}{POSTER}", videos["size"])255    mark("poster")256    manifest = _manifest(saga, at, videos, out)257    with open(f"{out}/{key}{MANIFEST}", "w") as handle:258        json.dump(manifest, handle)259    row = index_row(manifest)260    write_index(root, row)261    return {"manifest": manifest, "row": row, "dir": out, "ms": marks}262263if __name__ == "__main__":264    record = make(random.getrandbits(32), DATA_DIR)265    print()266    print(record["manifest"]["story"])267    print(record["dir"])