heatmap.py

1from mrlypy.core.colors import gradient2from mrlypy.two import Cell2d3import numpy as np4import os5from mrlypy.paint.colors import get_color6from mrlypy.paint.enums import Ink7from .helpers import logger8from .models import Life910BLACK = get_color(Ink.BLACK)11WHITE = get_color(Ink.WHITE)1213def create_heatmap(result: Life, scale=10, output_dir="data/life/heatmap", format="png", prefix="life") -> str:14    logger.debug(f"Creating heatmap for {len(result.grids)} grids")15    os.makedirs(output_dir, exist_ok=True)16    final_heatmap = np.zeros_like(result.grids[0].types, dtype=np.int32)17    for grid in result.grids:18        final_heatmap += grid.types19    max_count = int(np.max(final_heatmap))20    logger.debug(f"Max heatmap count: {max_count}")21    gradient_colors = gradient([WHITE, BLACK], max_count)22    gradient_colors.insert(0, BLACK)23    gradient_rgba = [c.to_rgba() for c in gradient_colors]24    gradient_array = np.array(gradient_rgba, dtype=np.uint8)25    cumulative_heatmap = np.zeros_like(result.grids[0].types, dtype=np.int32)26    for i, grid in enumerate(result.grids):27        cumulative_heatmap += grid.types28        colored_heatmap = gradient_array[cumulative_heatmap]29        cell = Cell2d(colors=colored_heatmap)30        image = cell.to_image(scale).convert("RGB")31        fp = f"{output_dir}/{prefix}_{i+1:03d}.{format}"32        image.save(fp, format=format)33        logger.debug(f"Saved: {fp}")34    logger.info(f"Saved {len(result.grids)} heatmaps to {output_dir}")35    return output_dir