"""Persist pipeline outputs for reports and the Streamlit demo.""" from __future__ import annotations import json import math from pathlib import Path from typing import Any import matplotlib.pyplot as plt import numpy as np import pandas as pd from .adapters.base import DatasetSchema from .config import DEMO_CACHE, FIGURES, GROUP_SIZE, RESULTS, TABLES, ensure_dirs from .plot_style import ( ACCENT, CENTROID_COLOR, CLUSTER_FILLS, CLUSTER_MARKERS, INK, INK_FAINT, INK_MUTED, RULE, SURFACE, apply_style, use_mono_ticks, ) # Apply project rcParams once at module import. Idempotent. apply_style() def _json_safe(value: Any) -> Any: if value is None or value is pd.NA: return None if isinstance(value, (np.integer,)): return int(value) if isinstance(value, (np.floating, float)): value = float(value) if not math.isfinite(value): return None return value if isinstance(value, (np.bool_)): return bool(value) if isinstance(value, np.ndarray): return _json_safe(value.tolist()) if isinstance(value, pd.Series): return {key: _json_safe(val) for key, val in value.to_dict().items()} if isinstance(value, dict): return {str(key): _json_safe(val) for key, val in value.items()} if isinstance(value, (list, tuple)): return [_json_safe(item) for item in value] try: if pd.isna(value): return None except (TypeError, ValueError): pass return value def _write_json(path: Path, payload: dict[str, Any]) -> None: path.write_text(json.dumps(_json_safe(payload), indent=2, allow_nan=False), encoding="utf-8") def groups_to_frame( ids: np.ndarray, labels: np.ndarray, groups: list[list[int]], features: pd.DataFrame, schema: DatasetSchema | None = None, ) -> pd.DataFrame: rows = [] feature_lookup = features.reset_index(drop=True) if schema is not None: display_cols = [col for col in schema.display_cols if col in feature_lookup.columns] else: display_cols = [ col for col in [ "total_clicks", "weighted_score", "imd_band_ord", "collaborative_clicks", "collaboration_click_ratio", "final_result", ] if col in feature_lookup.columns ] for group_id, members in enumerate(groups, start=1): for idx in members: row = { "id_student": ids[idx], "group_id": group_id, "cluster": int(labels[idx]), } for col in display_cols: row[col] = feature_lookup.loc[idx, col] rows.append(row) return pd.DataFrame(rows) def write_pipeline_diagram(path: Path) -> None: labels = [ "Learner dataset", "Adapter", "Features", "Preprocess", "Reducers", "Clusterers", "Stability", "Selector", "Groups", "Evaluation", ] fig, ax = plt.subplots(figsize=(12, 2.4)) ax.axis("off") xs = np.linspace(0.05, 0.95, len(labels)) for x, label in zip(xs, labels): ax.text( x, 0.55, label, ha="center", va="center", bbox=dict( boxstyle="round,pad=0.35", facecolor=SURFACE, edgecolor=INK_MUTED, linewidth=0.8, ), fontsize=9, color=INK, ) for x1, x2 in zip(xs[:-1], xs[1:]): ax.annotate( "", xy=(x2 - 0.035, 0.55), xytext=(x1 + 0.035, 0.55), arrowprops=dict(arrowstyle="->", color=INK_MUTED, linewidth=0.6), ) fig.tight_layout() fig.savefig(path, dpi=200, bbox_inches="tight") plt.close(fig) def _fig7_validity_vs_stability( metrics_df: pd.DataFrame, stability_df: pd.DataFrame, winner: pd.Series, ) -> None: """Validity (silhouette) vs Stability (bootstrap-ARI) scatter, primary configs.""" joined = metrics_df.drop(columns=["labels"], errors="ignore").merge( stability_df[["config_id", "bootstrap_ari_mean"]], on="config_id", how="left", ) primary = joined[joined["clusterer"].ne("hdbscan")] if primary.empty: return fig, ax = plt.subplots(figsize=(8, 5.5)) # All points in ink; differentiate clusterers by marker shape only clusterer_markers = {"kmeans": "o", "gmm": "s", "agglo": "^"} for clusterer, sub in primary.groupby("clusterer"): ax.scatter( sub["silhouette"], sub["bootstrap_ari_mean"], s=64, marker=clusterer_markers.get(clusterer, "o"), facecolor=INK, edgecolor="white", linewidth=0.8, label=clusterer, alpha=0.85, ) # Stability threshold reference line ax.axhline(0.40, linestyle="--", color=INK_MUTED, linewidth=0.8) ax.text( 0.005, 0.40, "stability floor (ARI = 0.40)", transform=ax.get_yaxis_transform(), ha="left", va="bottom", fontsize=8, color=INK_FAINT, style="italic", ) # Highlight the winner — terracotta diamond, leader line, label winner_id = winner.get("config_id") win_row = primary[primary["config_id"] == winner_id] if not win_row.empty: wx = float(win_row["silhouette"].iloc[0]) wy = float(win_row["bootstrap_ari_mean"].iloc[0]) ax.scatter( wx, wy, marker="D", s=120, facecolor=ACCENT, edgecolor="white", linewidth=1.2, zorder=10, ) ax.annotate( f"{winner_id} (winner)", xy=(wx, wy), xytext=(wx + 0.04, wy - 0.06), fontsize=10, color=INK, weight=600, arrowprops=dict(arrowstyle="-", color=INK_MUTED, linewidth=0.6), ) # Tiny config_id labels on non-winner points for _, row in primary.iterrows(): if row["config_id"] == winner_id: continue ax.annotate( row["config_id"], xy=(row["silhouette"], row["bootstrap_ari_mean"]), xytext=(5, 4), textcoords="offset points", fontsize=8, color=INK_FAINT, ) ax.set_xlabel("silhouette (higher is better)") ax.set_ylabel("bootstrap-ARI mean (higher is better)") ax.grid(axis="x", visible=False) ax.legend(loc="lower right", title=None) use_mono_ticks(ax) fig.tight_layout() fig.savefig(FIGURES / "fig7_scatter.png", dpi=200) plt.close(fig) def _fig8_embedding( labels: np.ndarray, X_vis: np.ndarray, winner: pd.Series, ) -> None: """UMAP 2D embedding — the hero figure. Cluster differentiation by SHAPE (per design.md §5.3 amendment), not colour. All centroids in terracotta accent — signals "structural finding," not "this cluster is the bad one." """ if X_vis.shape[1] < 2: return labels_arr = np.asarray(labels) fig, ax = plt.subplots(figsize=(8, 8)) for cluster_id in sorted(set(int(l) for l in labels_arr)): mask = labels_arr == cluster_id if not mask.any(): continue ax.scatter( X_vis[mask, 0], X_vis[mask, 1], s=18, alpha=0.7, marker=CLUSTER_MARKERS.get(cluster_id, "o"), facecolor=CLUSTER_FILLS.get(cluster_id, INK_MUTED), edgecolor="white", linewidth=0.5, ) # Centroid (skip noise cluster -1) if cluster_id != -1: cx, cy = X_vis[mask, 0].mean(), X_vis[mask, 1].mean() ax.scatter( cx, cy, marker="D", s=90, facecolor=CENTROID_COLOR, edgecolor="white", linewidth=1.2, zorder=10, ) # In-cluster size annotation near each centroid; whitespace, no box ax.annotate( f"cluster {cluster_id} n={int(mask.sum())}", xy=(cx, cy), xytext=(10, 10), textcoords="offset points", fontsize=10, color=INK, weight=600, ) ax.set_xlabel("umap-1") ax.set_ylabel("umap-2") ax.set_aspect("equal", adjustable="datalim") ax.grid(axis="x", visible=False) use_mono_ticks(ax) # Caption sourced from data — never hard-code n ax.text( 0.99, 0.01, f"n = {len(labels_arr)} · winner = {winner.get('config_id', '—')}", transform=ax.transAxes, ha="right", va="bottom", fontsize=9, color=INK_FAINT, family="monospace", ) fig.tight_layout() fig.savefig(FIGURES / "fig8_embedding.png", dpi=200) plt.close(fig) def _fig10_group_metrics(group_metrics: pd.DataFrame) -> None: """Group metrics by strategy — 2×3 small-multiples grid. One metric per panel. Mode B is the only terracotta bar; everything else rides the ink ramp. Reading order: top-row = compactness/separation, bottom-row = social/equity metrics. """ if group_metrics.empty: return metric_order = [ "intra_group_distance", "inter_group_variance", "complementarity", "engagement_balance", "demographic_fairness", "cluster_coverage", ] strategies = ["random", "stratified", "mode_a", "mode_b"] strategy_labels = ["random", "stratified", "mode A", "mode B"] bar_colors = [INK_FAINT, INK_MUTED, INK, ACCENT] fig, axes = plt.subplots(2, 3, figsize=(13, 7.5)) for ax, metric in zip(axes.flat, metric_order): if metric not in group_metrics.columns: ax.set_visible(False) continue vals = [] present_labels = [] present_colors = [] for s, lbl, c in zip(strategies, strategy_labels, bar_colors): row = group_metrics.loc[group_metrics["strategy"] == s, metric] if row.empty: continue vals.append(float(row.iloc[0])) present_labels.append(lbl) present_colors.append(c) ax.bar( range(len(vals)), vals, color=present_colors, edgecolor="white", linewidth=0.5, ) ax.set_xticks(range(len(vals))) ax.set_xticklabels(present_labels, fontsize=10) ax.set_title(metric.replace("_", " "), fontsize=11, loc="left", color=INK) ax.tick_params(labelsize=9) use_mono_ticks(ax) # Drop x-grid; small-multiples already have visual rhythm from the grid ax.grid(axis="x", visible=False) fig.suptitle("") fig.tight_layout() fig.savefig(FIGURES / "fig10_group_metrics.png", dpi=200) plt.close(fig) def _fig2_feature_families(columns: list[str]) -> None: """Horizontal bar chart showing feature count by family.""" families = { "Demographic": [ "age_band_ord", "imd_band_ord", "highest_education_ord", "disability_bin", "gender_M", "num_prev_attempts", "studied_credits", "registration_day", ], "Engagement": [ "total_clicks", "active_days", "first_click_day", "last_click_day", "mean_clicks_per_active_day", "engagement_span", "click_std", "click_cv", "weekend_ratio", ], "Collaboration": [ "collaborative_clicks", "forum_clicks", "live_collab_clicks", "collaborative_active_days", "collaboration_click_ratio", ], "VLE activity": [ "clicks_oucontent", "clicks_homepage", "clicks_subpage", "clicks_url", "clicks_resource", "clicks_dataplus", "clicks_glossary", ], "Performance": [ "mean_tma_score", "weighted_score", "n_assessments_submitted", "mean_submission_lateness", "score_trajectory_slope", "no_submissions", ], } col_set = set(columns) family_names = [] family_counts = [] for name, members in families.items(): count = sum(1 for m in members if m in col_set) if count > 0: family_names.append(name) family_counts.append(count) # Any unclassified columns classified = set() for members in families.values(): classified.update(members) unclassified = col_set - classified if unclassified: family_names.append("Other") family_counts.append(len(unclassified)) fig, ax = plt.subplots(figsize=(8, 4)) bars = ax.barh( range(len(family_names)), family_counts, color=[ACCENT if n == "Collaboration" else INK_MUTED for n in family_names], edgecolor="white", linewidth=0.5, ) ax.set_yticks(range(len(family_names))) ax.set_yticklabels(family_names) ax.set_xlabel("number of features") ax.invert_yaxis() ax.grid(axis="y", visible=False) # Count labels on bars for bar, count in zip(bars, family_counts): ax.text( bar.get_width() + 0.15, bar.get_y() + bar.get_height() / 2, str(count), va="center", fontsize=10, color=INK, ) use_mono_ticks(ax) fig.tight_layout() fig.savefig(FIGURES / "fig2_feature_families.png", dpi=200) plt.close(fig) def _fig3_config_heatmap( metrics_df: pd.DataFrame, stability_df: pd.DataFrame, winner: pd.Series, ) -> None: """12-config matrix heatmap — composite metrics, winner row ringed.""" joined = metrics_df.drop(columns=["labels"], errors="ignore").merge( stability_df[["config_id", "bootstrap_ari_mean"]], on="config_id", how="left", ) display_metrics = [ "silhouette", "davies_bouldin", "calinski_harabasz", "bootstrap_ari_mean", ] available = [m for m in display_metrics if m in joined.columns] if not available: return matrix = joined.set_index("config_id")[available].copy() # Normalise each column to [0, 1] for visual comparability for col in available: mn, mx = matrix[col].min(), matrix[col].max() if mx > mn: # DBI is lower-is-better: invert if col == "davies_bouldin": matrix[col] = 1.0 - (matrix[col] - mn) / (mx - mn) else: matrix[col] = (matrix[col] - mn) / (mx - mn) else: matrix[col] = 0.5 fig, ax = plt.subplots(figsize=(9, 6)) im = ax.imshow(matrix.values, aspect="auto", cmap="Greys", vmin=0, vmax=1) ax.set_xticks(range(len(available))) ax.set_xticklabels([m.replace("_", "\n") for m in available], fontsize=9) ax.set_yticks(range(len(matrix))) ax.set_yticklabels(matrix.index, fontsize=10) # Annotate cells with raw values from joined raw = joined.set_index("config_id")[available] for i, config_id in enumerate(matrix.index): for j, col in enumerate(available): val = raw.loc[config_id, col] fmt = f"{val:.3f}" if abs(val) < 10 else f"{val:.1f}" text_color = "white" if matrix.values[i, j] > 0.6 else INK ax.text(j, i, fmt, ha="center", va="center", fontsize=8, color=text_color) # Ring the winner row winner_id = winner.get("config_id") if winner_id in matrix.index: row_idx = list(matrix.index).index(winner_id) ax.add_patch(plt.Rectangle( (-0.5, row_idx - 0.5), len(available), 1, fill=False, edgecolor=ACCENT, linewidth=2.0, zorder=10, )) ax.set_title("Normalised validity metrics (higher = better)", loc="left", fontsize=12) use_mono_ticks(ax) fig.colorbar(im, ax=ax, shrink=0.6, label="normalised score") fig.tight_layout() fig.savefig(FIGURES / "fig3_config_heatmap.png", dpi=200) plt.close(fig) def _fig9_bootstrap_boxplot( stability_df: pd.DataFrame, winner: pd.Series | None = None, ) -> None: """Bootstrap-ARI distribution boxplot, sorted by mean ARI descending.""" if stability_df.empty or "bootstrap_ari_dist" not in stability_df.columns: return # Sort by mean ARI descending ordered = stability_df.sort_values("bootstrap_ari_mean", ascending=False) config_ids = ordered["config_id"].tolist() distributions = [] for _, row in ordered.iterrows(): dist = row["bootstrap_ari_dist"] if isinstance(dist, (list, np.ndarray)): distributions.append(np.asarray(dist)) else: distributions.append(np.array([row.get("bootstrap_ari_mean", 0)])) fig, ax = plt.subplots(figsize=(10, 5.5)) bp = ax.boxplot( distributions, vert=True, patch_artist=True, labels=config_ids, widths=0.6, ) # Style boxes — highlight the actual winner, not just highest ARI winner_id = winner.get("config_id") if winner is not None else None for i, (box, median) in enumerate(zip(bp["boxes"], bp["medians"])): if config_ids[i] == winner_id: box.set(facecolor=ACCENT, alpha=0.4) median.set(color=ACCENT, linewidth=1.5) else: box.set(facecolor=SURFACE, alpha=0.8) median.set(color=INK, linewidth=1.0) box.set(edgecolor=INK_MUTED, linewidth=0.8) for whisker in bp["whiskers"]: whisker.set(color=INK_MUTED, linewidth=0.6) for cap in bp["caps"]: cap.set(color=INK_MUTED, linewidth=0.6) for flier in bp["fliers"]: flier.set(marker="o", markersize=3, markerfacecolor=INK_FAINT, markeredgecolor="none", alpha=0.6) # Stability floor ax.axhline(0.40, linestyle="--", color=INK_MUTED, linewidth=0.8) ax.text( len(distributions) + 0.3, 0.40, "stability floor", va="bottom", fontsize=8, color=INK_FAINT, style="italic", ) ax.set_ylabel("pairwise Adjusted Rand Index") ax.set_xlabel("configuration (sorted by mean ARI)") ax.tick_params(axis="x", labelrotation=45) ax.grid(axis="x", visible=False) use_mono_ticks(ax) fig.tight_layout() fig.savefig(FIGURES / "fig9_bootstrap_boxplot.png", dpi=200) plt.close(fig) def write_figures( metrics_df: pd.DataFrame, stability_df: pd.DataFrame, labels: np.ndarray, X_vis: np.ndarray, group_metrics: pd.DataFrame, winner: pd.Series, columns: list[str] | None = None, ) -> None: """Write all matplotlib figures. Pass 1 (fig7/8/10) + Pass 2 (fig2/3/9). All use Variant B tokens via plot_style. """ FIGURES.mkdir(parents=True, exist_ok=True) _fig7_validity_vs_stability(metrics_df, stability_df, winner) _fig8_embedding(labels, X_vis, winner) _fig10_group_metrics(group_metrics) # Pass 2 figures if columns is not None: _fig2_feature_families(columns) _fig3_config_heatmap(metrics_df, stability_df, winner) _fig9_bootstrap_boxplot(stability_df, winner=winner) def write_report( path: Path, winner: pd.Series, ranked: pd.DataFrame, group_metrics: pd.DataFrame, constraints: dict[str, Any], ) -> None: lines = [ "# Pipeline Run Report", "", f"Selected config: {winner['config_id']} ({winner['reducer']} + {winner['clusterer']})", "", "## Top Ranked Configurations", "", "```text", ranked.head(12).to_string(index=False), "```", "", "## Group Metrics", "", "```text", group_metrics.to_string(index=False), "```", "", "## Constraint Summary", "", "```json", json.dumps(_json_safe(constraints), indent=2, allow_nan=False), "```", "", ] path.write_text("\n".join(lines), encoding="utf-8") def write( ids: np.ndarray, features: pd.DataFrame, X_scaled: np.ndarray, reductions: dict[str, np.ndarray], labels_by_config: dict[str, np.ndarray], metrics_df: pd.DataFrame, stability_df: pd.DataFrame, winner: pd.Series, ranked: pd.DataFrame, winner_labels: np.ndarray, groups_a: list[list[int]], groups_b: list[list[int]], group_metrics: pd.DataFrame, random_baseline_metrics: pd.DataFrame | None, group_significance: pd.DataFrame | None, cluster_summary: pd.DataFrame | None, constraints: dict[str, Any], run_metadata: dict[str, Any] | None = None, cache_dir: Path = DEMO_CACHE, columns: list[str] | None = None, schema: DatasetSchema | None = None, ) -> None: ensure_dirs() cache_dir.mkdir(parents=True, exist_ok=True) write_global_artifacts = cache_dir.resolve() == DEMO_CACHE.resolve() features.to_parquet(cache_dir / "features.parquet", index=False) np.save(cache_dir / "X_scaled.npy", X_scaled) for name, values in reductions.items(): np.save(cache_dir / f"reduced_{name}.npy", values) if "umap_2d" in reductions: np.save(cache_dir / "reduced_umap_2d.npy", reductions["umap_2d"]) elif "pca" in reductions: np.save(cache_dir / "reduced_umap_2d.npy", reductions["pca"][:, :2]) cluster_labels = pd.DataFrame({"id_student": ids}) for cid, labels in labels_by_config.items(): cluster_labels[cid] = labels cluster_labels.to_parquet(cache_dir / "cluster_labels.parquet", index=False) metrics_clean = metrics_df.drop(columns=["labels"], errors="ignore") metrics_clean.to_parquet(cache_dir / "config_metrics.parquet", index=False) if write_global_artifacts: metrics_clean.to_csv(TABLES / "config_metrics.csv", index=False) stability_df.to_parquet(cache_dir / "stability.parquet", index=False) if write_global_artifacts: stability_df.drop(columns=["bootstrap_ari_dist"], errors="ignore").to_csv( TABLES / "stability.csv", index=False, ) ranked.to_parquet(cache_dir / "ranked_configs.parquet", index=False) _write_json(cache_dir / "winner.json", winner.to_dict()) groups_a_df = groups_to_frame(ids, winner_labels, groups_a, features, schema) groups_b_df = groups_to_frame(ids, winner_labels, groups_b, features, schema) groups_a_df.to_parquet(cache_dir / "groups_mode_a.parquet", index=False) groups_b_df.to_parquet(cache_dir / "groups_mode_b.parquet", index=False) if write_global_artifacts: groups_a_df.to_csv(TABLES / "groups_mode_a.csv", index=False) groups_b_df.to_csv(TABLES / "groups_mode_b.csv", index=False) group_metrics.to_parquet(cache_dir / "group_metrics.parquet", index=False) if write_global_artifacts: group_metrics.to_csv(TABLES / "group_metrics.csv", index=False) if random_baseline_metrics is not None: random_baseline_metrics.to_parquet(cache_dir / "random_baseline_metrics.parquet", index=False) if write_global_artifacts: random_baseline_metrics.to_csv(TABLES / "random_baseline_metrics.csv", index=False) if group_significance is not None: group_significance.to_parquet(cache_dir / "group_significance.parquet", index=False) if write_global_artifacts: group_significance.to_csv(TABLES / "group_significance.csv", index=False) if cluster_summary is not None: cluster_summary.to_parquet(cache_dir / "cluster_summary.parquet", index=False) if write_global_artifacts: cluster_summary.to_csv(TABLES / "cluster_summary.csv", index=False) if schema is not None: _write_json(cache_dir / "schema.json", schema.to_dict()) dataset_name = schema.dataset_name if schema is not None else winner.get("dataset_name", "") adapter_name = schema.adapter_name if schema is not None else winner.get("adapter_name", "") meta = { "n_learners": int(len(ids)), "n_features": int(len(columns) if columns is not None else len([col for col in features.columns if col != "id_student"])), "n_groups": int(len(groups_b)), "group_size": GROUP_SIZE, "dataset_name": dataset_name, "adapter_name": adapter_name, "presentation": f"{winner.get('presentation_module', '')}_{winner.get('presentation_code', '')}".strip("_"), "winner_config": winner.get("config_id"), "winner_reducer": winner.get("reducer"), "winner_clusterer": winner.get("clusterer"), } if schema is not None: meta.update( { "source_id_col": schema.source_id_col, "id_col": schema.id_col, "fairness_cols": schema.fairness_cols, "engagement_col": schema.engagement_col, "performance_col": schema.performance_col, "outcome_col": schema.outcome_col, "stratification_col": schema.stratification_col, "display_cols": schema.display_cols, } ) if run_metadata: meta.update(run_metadata) _write_json(cache_dir / "meta.json", meta) _write_json(cache_dir / "constraints.json", constraints) # Persist feature columns for regen_figures.py if columns is not None: _write_json(cache_dir / "columns.json", columns) write_pipeline_diagram(cache_dir / "pipeline_diagram.png") if write_global_artifacts: write_figures( metrics_clean, stability_df, winner_labels, reductions.get("umap_2d", reductions.get("pca", X_scaled[:, :2])), group_metrics, winner, columns=columns, ) # Graphviz diagrams (fig1/4/5/6) try: from .diagrams import render_all render_all(FIGURES) except ImportError: pass # graphviz not installed — skip flowcharts write_report(RESULTS / "pipeline_report.md", winner, ranked, group_metrics, constraints)