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"])