File size: 3,505 Bytes
37e13f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """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 # noqa: E402
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")
|