diff --git a/src/multi_swarm/dashboard/aquarium.py b/src/multi_swarm/dashboard/aquarium.py new file mode 100644 index 0000000..a94ba13 --- /dev/null +++ b/src/multi_swarm/dashboard/aquarium.py @@ -0,0 +1,281 @@ +"""Aquarium 2D visualization helpers. + +Builds a list of fish records from a merged DataFrame (evaluations + genomes) +and renders a self-contained HTML/JS canvas animation, embeddable in Streamlit +via ``streamlit.components.v1.html``. +""" + +from __future__ import annotations + +import json +import math +from typing import Any + +import pandas as pd # type: ignore[import-untyped] + +# Color palette per cognitive style. Default fallback for unknown styles is grey. +STYLE_COLORS: dict[str, str] = { + "physicist": "#4cc9f0", + "biologist": "#52b788", + "historian": "#e76f51", + "meteorologist": "#ffd166", + "ecologist": "#a78bfa", + "engineer": "#fb6f92", +} +DEFAULT_COLOR: str = "#94a3b8" + + +def build_fish_dataset(merged: pd.DataFrame, max_fish: int = 30) -> list[dict[str, Any]]: + """Build a list of fish records from a merged evaluations+genomes DataFrame. + + Expects columns: ``genome_id``, ``fitness``, ``cognitive_style``, ``n_trades``, + ``dsr``. Rows with NaN ``fitness`` are dropped. Result is sorted by fitness + descending and capped at ``max_fish`` entries. + """ + if merged.empty: + return [] + + cols_needed = ["genome_id", "fitness", "cognitive_style", "n_trades", "dsr"] + available = [c for c in cols_needed if c in merged.columns] + df = merged[available].copy() + + if "fitness" not in df.columns: + return [] + + df = df.dropna(subset=["fitness"]) + df = df.sort_values("fitness", ascending=False).head(max_fish) + + fish: list[dict[str, Any]] = [] + for _, row in df.iterrows(): + gid = str(row.get("genome_id", "")) + fitness_val = float(row.get("fitness", 0.0)) + if math.isnan(fitness_val): + fitness_val = 0.0 + style_raw = row.get("cognitive_style", None) + style = str(style_raw) if style_raw is not None and not _is_nan(style_raw) else "unknown" + n_trades_raw = row.get("n_trades", 0) + n_trades = int(n_trades_raw) if not _is_nan(n_trades_raw) else 0 + dsr_raw = row.get("dsr", 0.0) + dsr = float(dsr_raw) if not _is_nan(dsr_raw) else 0.0 + fish.append( + { + "id": gid, + "fitness": fitness_val, + "cognitive_style": style, + "n_trades": n_trades, + "dsr": dsr, + } + ) + return fish + + +def _is_nan(v: Any) -> bool: + try: + return bool(pd.isna(v)) + except (TypeError, ValueError): + return False + + +def build_aquarium_html( + fish: list[dict[str, Any]], + canvas_w: int = 1000, + canvas_h: int = 600, + show_labels: bool = False, +) -> str: + """Build the self-contained HTML/JS string for the aquarium canvas.""" + fish_json = json.dumps(fish) + palette_json = json.dumps(STYLE_COLORS) + default_color = DEFAULT_COLOR + show_labels_js = "true" if show_labels else "false" + + # All braces inside