feature_neighbors_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin feature neighbor --check` — save JSON + diagram to the current working directory."""
 
from __future__ import annotations
 
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
_JSON_NAME = "feature-neighbors-check.json"
_PNG_NAME = "feature-neighbors-check.png"
 
 
def _flatten_feature_neighbors_result(result: dict[str, Any]) -> dict[str, Any]:
    if result.get("error"):
        return dict(result)
 
    data = result.get("content") or result
    if not isinstance(data, dict):
        return {"error": "Invalid feature-neighbors result"}
 
    neighbors = data.get("neighbors") or []
    return {**data, "neighbors": neighbors}
 
 
def write_feature_neighbors_check(
    result: dict[str, Any],
    *,
    tool_name: str | None,
    cwd: str | Path,
) -> tuple[Path, Path]:
    cwd = Path(cwd)
    json_path = cwd / _JSON_NAME
    png_path = cwd / _PNG_NAME
 
    flat = _flatten_feature_neighbors_result(result)
    payload = {
        "saved_at": datetime.now(timezone.utc).isoformat(),
        "tool": tool_name,
        **flat,
    }
    json_path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
    _plot_feature_neighbors_check(flat, png_path)
    return json_path, png_path
 
 
def _neighbor_label(n: dict[str, Any]) -> str:
    ref = n.get("feature_ref")
    if ref:
        return str(ref)[:28]
    idx = n.get("feature_idx", "?")
    label = str(n.get("label") or "")
    if label:
        return f"{idx}:{label[:20]}"
    return f"F{idx}"
 
 
def _plot_feature_neighbors_check(result: dict[str, Any], png_path: Path) -> None:
    import matplotlib
 
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
 
    if result.get("error"):
        fig, ax = plt.subplots(figsize=(6, 2), facecolor="#0f1117")
        ax.axis("off")
        ax.text(0.5, 0.5, f"Error: {result['error']}", ha="center", va="center", wrap=True, color="#e5e7eb")
        fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
        plt.close(fig)
        return
 
    ref = result.get("feature_ref") or result.get("feature_idx", "?")
    layer = result.get("layer", "?")
    label = str(result.get("label") or result.get("causal_label") or "")
    neighbors = sorted(
        result.get("neighbors") or [],
        key=lambda n: float(n.get("similarity") or 0),
        reverse=True,
    )
 
    title = f"Feature neighbors — {ref}  ·  layer {layer}"
    if label:
        title += f"  ·  {label[:32]}"
 
    fig_h = max(4, min(12, len(neighbors) * 0.35 + 1.8))
    fig, ax = plt.subplots(figsize=(9, fig_h), facecolor="#0f1117")
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=0.98)
 
    ax.set_facecolor("#0f1117")
    ax.set_title("Decoder cosine similarity (nearest features)", color="#e5e7eb", fontsize=10, pad=10)
 
    if not neighbors:
        ax.axis("off")
        ax.text(0.5, 0.5, "No neighbors found.", ha="center", va="center", color="#9ca3af")
    else:
        labels = [_neighbor_label(n) for n in neighbors]
        vals = [float(n.get("similarity") or 0) for n in neighbors]
        y_pos = list(range(len(neighbors)))
        colors = ["#38bdf8" if v >= 0.9 else "#6366f1" if v >= 0.8 else "#78716c" for v in vals]
        ax.barh(y_pos, vals, color=colors, height=0.72, alpha=0.9)
        ax.set_yticks(y_pos)
        ax.set_yticklabels(labels, fontsize=8, color="#d1d5db")
        ax.invert_yaxis()
        ax.set_xlim(0, 1.02)
        ax.set_xlabel("Cosine similarity", color="#9ca3af", fontsize=9)
        ax.tick_params(axis="x", colors="#6b7280", labelsize=8)
        ax.grid(axis="x", alpha=0.2, color="#4b5563")
        for spine in ax.spines.values():
            spine.set_color("#374151")
        for i, v in enumerate(vals):
            ax.text(min(v + 0.01, 0.98), i, f"{v:.4f}", va="center", fontsize=7, color="#9ca3af")
 
    fig.tight_layout()
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)