| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| from driftwm.sim.flow import sample_flow |
| from driftwm.utils import ensure_dir |
|
|
|
|
| PAPER_FLOW_FAMILIES = [ |
| "noflow", |
| "uniform", |
| "vortex_center", |
| "double_gyre", |
| "source_sink", |
| "source_sink_pair", |
| "gradient", |
| "shear", |
| "turbulent_patch", |
| "random_fourier", |
| ] |
|
|
|
|
| def color_map(values: np.ndarray, vmax: float = 0.38) -> np.ndarray: |
| x = np.clip(values / vmax, 0.0, 1.0) |
| stops = np.array( |
| [ |
| [18, 35, 61], |
| [31, 91, 107], |
| [64, 145, 108], |
| [164, 190, 105], |
| [244, 220, 102], |
| ], |
| dtype=np.float32, |
| ) |
| scaled = x * (len(stops) - 1) |
| lo = np.floor(scaled).astype(np.int32) |
| hi = np.clip(lo + 1, 0, len(stops) - 1) |
| frac = scaled[..., None] - lo[..., None] |
| return ((1.0 - frac) * stops[lo] + frac * stops[hi]).astype(np.uint8) |
|
|
|
|
| def draw_arrow(draw: ImageDraw.ImageDraw, start: tuple[float, float], vec: tuple[float, float], scale: float) -> None: |
| vx, vy = vec |
| norm = float(np.hypot(vx, vy)) |
| if norm < 1.0e-4: |
| return |
| sx, sy = start |
| ex = sx + scale * vx |
| ey = sy - scale * vy |
| draw.line((sx, sy, ex, ey), fill=(255, 255, 255), width=2) |
| ux = (ex - sx) / max(float(np.hypot(ex - sx, ey - sy)), 1.0e-6) |
| uy = (ey - sy) / max(float(np.hypot(ex - sx, ey - sy)), 1.0e-6) |
| px, py = -uy, ux |
| size = 6.0 |
| p1 = (ex, ey) |
| p2 = (ex - size * ux + 0.55 * size * px, ey - size * uy + 0.55 * size * py) |
| p3 = (ex - size * ux - 0.55 * size * px, ey - size * uy - 0.55 * size * py) |
| draw.polygon((p1, p2, p3), fill=(255, 255, 255)) |
|
|
|
|
| def make_panel( |
| family: str, |
| flow_seed: int, |
| panel_size: int = 320, |
| *, |
| show_header: bool = True, |
| show_seed: bool = True, |
| ) -> Image.Image: |
| rng = np.random.default_rng(flow_seed) |
| flow = sample_flow(family, rng, flow_id=flow_seed) |
| n_bg = 96 |
| xs = np.linspace(0.0, 10.0, n_bg, dtype=np.float32) |
| ys = np.linspace(0.0, 10.0, n_bg, dtype=np.float32) |
| bg_grid = np.stack(np.meshgrid(xs, ys), axis=-1).astype(np.float32) |
| bg_velocity = flow.velocity(bg_grid, t=0.0) |
| speed = np.linalg.norm(bg_velocity, axis=-1) |
| bg = Image.fromarray(color_map(np.flipud(speed)), mode="RGB").resize((panel_size, panel_size), Image.Resampling.BICUBIC) |
| header = 44 if show_header else 0 |
| panel = Image.new("RGB", (panel_size, panel_size + header), (245, 245, 240)) |
| panel.paste(bg, (0, header)) |
| draw = ImageDraw.Draw(panel) |
| font = ImageFont.load_default() |
| draw.rectangle((0, 0, panel_size - 1, panel_size + header - 1), outline=(22, 28, 33), width=2) |
| if show_header: |
| draw.text((10, 10), f"{family}", fill=(20, 24, 28), font=font) |
| if show_seed: |
| draw.text((10, 25), f"seed={flow_seed}", fill=(55, 61, 68), font=font) |
|
|
| n_arrow = 13 |
| ax = np.linspace(0.8, 9.2, n_arrow, dtype=np.float32) |
| ay = np.linspace(0.8, 9.2, n_arrow, dtype=np.float32) |
| arrow_grid = np.stack(np.meshgrid(ax, ay), axis=-1).astype(np.float32) |
| arrow_velocity = flow.velocity(arrow_grid, t=0.0) |
| for pos, vel in zip(arrow_grid.reshape(-1, 2), arrow_velocity.reshape(-1, 2), strict=True): |
| px = 0.5 + pos[0] / 10.0 * (panel_size - 1) |
| py = header + 0.5 + (1.0 - pos[1] / 10.0) * (panel_size - 1) |
| draw_arrow(draw, (float(px), float(py)), (float(vel[0]), float(vel[1])), scale=72.0) |
| return panel |
|
|
|
|
| def _family_seeds(seed: int) -> list[int]: |
| rng = np.random.default_rng(seed) |
| return [int(rng.integers(1, 2**31 - 1)) for _ in PAPER_FLOW_FAMILIES] |
|
|
|
|
| def make_flow_atlas(seed: int, out: str | Path, panel_size: int = 320) -> None: |
| seeds = _family_seeds(seed) |
| panels = [make_panel(family, flow_seed, panel_size=panel_size) for family, flow_seed in zip(PAPER_FLOW_FAMILIES, seeds, strict=True)] |
| width, height = panels[0].size |
| margin = 22 |
| header = 58 |
| cols = 5 |
| rows = int(np.ceil(len(panels) / cols)) |
| atlas = Image.new("RGB", (cols * width + (cols + 1) * margin, rows * height + (rows + 1) * margin + header), (250, 250, 246)) |
| draw = ImageDraw.Draw(atlas) |
| font = ImageFont.load_default() |
| draw.text((margin, 18), f"FlowMo-WM paper flow atlas, deterministic atlas seed={seed}", fill=(15, 19, 24), font=font) |
| for idx, panel in enumerate(panels): |
| row, col = divmod(idx, cols) |
| x = margin + col * (width + margin) |
| y = header + margin + row * (height + margin) |
| atlas.paste(panel, (x, y)) |
| out = Path(out) |
| ensure_dir(out.parent) |
| atlas.save(out) |
| print(f"wrote {out}") |
| for family, flow_seed in zip(PAPER_FLOW_FAMILIES, seeds, strict=True): |
| print(f"{family}: {flow_seed}") |
|
|
|
|
| def export_flow_panels(seed: int, out_dir: str | Path, panel_size: int = 640) -> None: |
| out_dir = Path(out_dir) |
| clean_dir = out_dir / "clean" |
| labeled_dir = out_dir / "labeled" |
| ensure_dir(clean_dir) |
| ensure_dir(labeled_dir) |
| rows = ["family\tseed\tclean_png\tlabeled_png"] |
| for family, flow_seed in zip(PAPER_FLOW_FAMILIES, _family_seeds(seed), strict=True): |
| clean = make_panel(family, flow_seed, panel_size=panel_size, show_header=False) |
| labeled = make_panel(family, flow_seed, panel_size=panel_size, show_header=True, show_seed=False) |
| clean_path = clean_dir / f"{family}.png" |
| labeled_path = labeled_dir / f"{family}.png" |
| clean.save(clean_path) |
| labeled.save(labeled_path) |
| rows.append(f"{family}\t{flow_seed}\t{clean_path}\t{labeled_path}") |
| print(f"wrote {clean_path}") |
| print(f"wrote {labeled_path}") |
| manifest = out_dir / "manifest.tsv" |
| manifest.write_text("\n".join(rows) + "\n") |
| print(f"wrote {manifest}") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--seed", type=int, default=20260525) |
| parser.add_argument("--out", default="experiments/reports/figures/flow_family_atlas.png") |
| parser.add_argument("--panel-dir", default=None) |
| parser.add_argument("--panel-size", type=int, default=640) |
| args = parser.parse_args() |
| make_flow_atlas(args.seed, args.out) |
| if args.panel_dir is not None: |
| export_flow_panels(args.seed, args.panel_dir, panel_size=args.panel_size) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
