| """Shared flatmap plotting helpers.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import matplotlib |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| matplotlib.use("Agg") |
| from matplotlib import colormaps |
|
|
|
|
| def load_font(size: int, bold: bool = False): |
| candidates = [ |
| "DejaVuSans-Bold.ttf" if bold else "DejaVuSans.ttf", |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", |
| ] |
| for candidate in candidates: |
| try: |
| return ImageFont.truetype(candidate, size) |
| except OSError: |
| continue |
| return ImageFont.load_default() |
|
|
|
|
| def text_size(draw: ImageDraw.ImageDraw, text: str, font) -> tuple[int, int]: |
| box = draw.textbbox((0, 0), text, font=font) |
| return box[2] - box[0], box[3] - box[1] |
|
|
|
|
| def split_rows(mask: np.ndarray) -> tuple[int, int]: |
| row_counts = mask.sum(axis=1) |
| height = mask.shape[0] |
| lo = int(height * 0.35) |
| hi = int(height * 0.65) |
| split = lo + int(np.argmin(row_counts[lo:hi])) |
| return split, int(row_counts[split]) |
|
|
|
|
| def half_mask(mask: np.ndarray, split: int, half: str) -> np.ndarray: |
| out = np.zeros_like(mask, dtype=bool) |
| if half == "left": |
| out[: split + 1] = mask[: split + 1] |
| elif half == "right": |
| out[split + 1 :] = mask[split + 1 :] |
| else: |
| raise ValueError(f"Unknown hemisphere {half!r}") |
| return out |
|
|
|
|
| def crop_box_for_half(mask: np.ndarray, split: int, half: str, margin: int) -> tuple[slice, slice]: |
| hemi = half_mask(mask, split, half) |
| return crop_box_for_mask(hemi, margin) |
|
|
|
|
| def crop_box_for_mask(mask: np.ndarray, margin: int) -> tuple[slice, slice]: |
| rows, cols = np.where(mask) |
| if rows.size == 0: |
| raise ValueError("No pixels found for mask.") |
| r0 = max(int(rows.min()) - margin, 0) |
| r1 = min(int(rows.max()) + margin + 1, mask.shape[0]) |
| c0 = max(int(cols.min()) - margin, 0) |
| c1 = min(int(cols.max()) + margin + 1, mask.shape[1]) |
| return slice(r0, r1), slice(c0, c1) |
|
|
|
|
| def colorize_scalar(values: np.ndarray, mask: np.ndarray, limit: float, cmap_name: str = "coolwarm") -> Image.Image: |
| cmap = colormaps[cmap_name] |
| normed = np.clip((values / max(limit, 1e-8) + 1.0) / 2.0, 0.0, 1.0) |
| rgba = (cmap(normed) * 255).astype(np.uint8) |
| rgba[~mask, 3] = 0 |
| rgba[mask & ~np.isfinite(values), 3] = 0 |
| return Image.fromarray(rgba, mode="RGBA") |
|
|
|
|
| def render_colorbar(width: int, height: int, cmap_name: str = "coolwarm") -> Image.Image: |
| cmap = colormaps[cmap_name] |
| gradient = np.linspace(-1.0, 1.0, width, dtype=np.float32) |
| rgba = (cmap((gradient + 1.0) / 2.0) * 255).astype(np.uint8) |
| bar = np.repeat(rgba[None, :, :], height, axis=0) |
| return Image.fromarray(bar, mode="RGBA") |
|
|
|
|
| def colorize_positive(values: np.ndarray, mask: np.ndarray, limit: float, cmap_name: str = "magma") -> Image.Image: |
| cmap = colormaps[cmap_name] |
| normed = np.clip(values / max(limit, 1e-8), 0.0, 1.0) |
| rgba = (cmap(normed) * 255).astype(np.uint8) |
| rgba[~mask, 3] = 0 |
| rgba[mask & ~np.isfinite(values), 3] = 0 |
| return Image.fromarray(rgba, mode="RGBA") |
|
|
|
|
| def render_positive_colorbar(width: int, height: int, cmap_name: str = "magma") -> Image.Image: |
| cmap = colormaps[cmap_name] |
| gradient = np.linspace(0.0, 1.0, width, dtype=np.float32) |
| rgba = (cmap(gradient) * 255).astype(np.uint8) |
| bar = np.repeat(rgba[None, :, :], height, axis=0) |
| return Image.fromarray(bar, mode="RGBA") |
|
|
