""" Result models returned by the analysis pipeline. """ from __future__ import annotations import hashlib import json from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, List, Optional import numpy as np from tracescope.models.trace import TraceSession @dataclass class ClusterResult: """Clustering output. Attributes: n_clusters: Number of clusters found. labels: Cluster label per entry (length N). clusters: Per-cluster info dicts with keys: cluster_id, indices, centroid. """ n_clusters: int labels: List[int] clusters: List[dict] @dataclass class AxisInfo: """PCA axis information for the 3D projection. Attributes: axes: (3, D) array – principal directions in high-dim space. lengths: Variance explained per axis (length 3). labels: LLM-generated semantic labels (length 3). min_point_idx: Index of the point at the negative extreme per axis. max_point_idx: Index of the point at the positive extreme per axis. min_cluster_idx: Cluster at negative extreme per axis. max_cluster_idx: Cluster at positive extreme per axis. """ axes: np.ndarray lengths: List[float] labels: List[str] = field(default_factory=lambda: ["Axis 1", "Axis 2", "Axis 3"]) min_point_idx: List[int] = field(default_factory=list) max_point_idx: List[int] = field(default_factory=list) min_cluster_idx: List[int] = field(default_factory=list) max_cluster_idx: List[int] = field(default_factory=list) @dataclass class AnalysisResult: """Full result of the analysis pipeline. Attributes: session: The analyzed TraceSession. embedding_model: Name of the embedding model used. embeddings: (N, D) high-dim embedding matrix. clusters: Clustering result. projected_3d: (N, 3) projected coordinates. axis_info: PCA axis information with semantic labels. cluster_labels: LLM-generated description per cluster. flow_model_trained: Whether the flow model was successfully trained. segments: Optional segmentation result. """ session: TraceSession embedding_model: str embeddings: np.ndarray clusters: ClusterResult projected_3d: np.ndarray axis_info: AxisInfo cluster_labels: List[str] = field(default_factory=list) flow_model_trained: bool = False segments: Optional[List[dict]] = None # Flow field data (from MDN model) velocity_grid: Optional[np.ndarray] = None # (40,40,40,3) velocity field confidence_grid: Optional[np.ndarray] = None # (40,40,40) MDN confidence [0,1] axis_min: Optional[np.ndarray] = None # bounding box min [3] axis_max: Optional[np.ndarray] = None # bounding box max [3] mdn_simulate: Optional[Any] = None # MDN simulate() callable fitted_reducer: Optional[Any] = None # fitted UMAP/tSNE for .transform() # Cluster geometry in 3D projected space cluster_centroids_3d: Optional[np.ndarray] = None # (K,3) cluster centers max_cluster_distance: float = 1.0 # max pairwise dist between centroids cache_path: Optional[str] = None # base cache path for attractor caching flow_mode: str = "mdn" # flow model type used ("mdn" or "rbf") seed: int = 42 # global seed used by every randomized stage deterministic: bool = True # whether RNGs were seeded this run @property def n_entries(self) -> int: return len(self.session) @property def score_channels(self) -> List[str]: """Return sorted list of all score channel names available in the data.""" return self.session.score_channels def get_entry_scores(self, channel: str) -> List[Optional[float]]: """Return score values for a channel across all entries (None if missing).""" return [ e.scores.get(channel) for e in self.session.entries ] def get_path_scores(self, channel: str) -> dict: """Return {path_id: score} for a channel from path_scores.""" return { pid: scores[channel] for pid, scores in self.session.path_scores.items() if channel in scores } def find_attractors(self, score_channel: Optional[str] = None) -> list: """Find flow attractors in the velocity field. Detects regions where flow converges and particles accumulate, distinguishing real attractors from boundary effects and dead zones. Computed on demand from the velocity grid (fast, ~0.1s). Args: score_channel: Optional score channel name to compute mean score within each attractor's basin (e.g. "solved" to see if the attractor is positive or negative). Returns: List of attractor dicts with keys: position, strength, divergence, basin_mask, basin_size, basin_fraction, mean_score """ if self.velocity_grid is None: return [] from tracescope.visualization.flow_field import FlowFieldSystem flow = FlowFieldSystem( self.velocity_grid, self.axis_min, self.axis_max, particle_grid=2, # minimal — we only need the grid, not particles confidence_grid=self.confidence_grid, ) # Build score grid if requested score_grid = None if score_channel and self.score_channels: if score_channel in self.score_channels: from tracescope.visualization.gl_renderer import FlowRenderer # Get normalized scores per entry raw = self.get_entry_scores(score_channel) path_scores = self.get_path_scores(score_channel) vals = [] for i, entry in enumerate(self.session.entries): v = raw[i] if v is None and entry.path_id is not None: v = path_scores.get(entry.path_id) vals.append(float(v) if v is not None else float('nan')) arr = np.array(vals, dtype=np.float32) valid = ~np.isnan(arr) if valid.any(): lo, hi = float(np.nanmin(arr)), float(np.nanmax(arr)) if hi - lo > 1e-8: arr[valid] = (arr[valid] - lo) / (hi - lo) else: arr[valid] = 0.5 arr[~valid] = 0.5 flow.build_score_grid(self.projected_3d, arr) score_grid = flow._score_grid return flow.find_attractors(score_grid=score_grid) def get_cluster_texts(self, cluster_id: int) -> List[str]: """Return texts belonging to a specific cluster.""" return [ self.session.entries[i].text for i, lbl in enumerate(self.clusters.labels) if lbl == cluster_id ] def get_entry_at_3d(self, idx: int) -> dict: """Return entry info with its 3D coordinates.""" entry = self.session.entries[idx] coords = self.projected_3d[idx] return { "text": entry.text, "role": entry.role, "x": float(coords[0]), "y": float(coords[1]), "z": float(coords[2]), "cluster": self.clusters.labels[idx], "step_index": entry.step_index, } def fingerprint(self) -> str: """Compute a fingerprint from (texts + embedding_model + flow_mode + score channel names [+ seed]). Score channels are included so that changing the channel set (e.g. dropping a score or adding a new one) invalidates the cache — otherwise a cached result would be loaded with a stale score layout. Backward-compat: when run with the historical defaults (deterministic=True AND seed=42) the seed is *omitted* from the fingerprint, producing the exact same hash as pre-seed versions of TraceScope. This keeps existing cached results valid bit-for-bit on disk, so users who never touch the seed don't need to recompute anything. Non-default seeds and non-deterministic runs append a seed suffix, giving each alternate run its own distinct cached result. """ texts = sorted(e.text for e in self.session.entries) score_channels = sorted(self.session.score_channels) is_default_seed = bool(self.deterministic) and int(self.seed) == 42 if is_default_seed: blob = ( json.dumps(texts, sort_keys=True) + "|" + self.embedding_model + "|" + self.flow_mode + "|" + json.dumps(score_channels, sort_keys=True) ) else: seed_key = ( f"{int(self.seed)}" if self.deterministic else "nondet" ) blob = ( json.dumps(texts, sort_keys=True) + "|" + self.embedding_model + "|" + self.flow_mode + "|" + json.dumps(score_channels, sort_keys=True) + "|" + seed_key ) return hashlib.sha256(blob.encode("utf-8")).hexdigest() def save_result(self, path: str): """Save the full AnalysisResult to disk. Creates two files: - {path}.npz — numpy arrays (embeddings, projected_3d, velocity_grid, etc.) - {path}.json — metadata (session, cluster info, axis labels, etc.) Args: path: Base path (without extension). E.g. "my_result" creates my_result.npz and my_result.json. """ base = Path(path) # ── Numpy arrays ─────────────────────────────── arrays = { "embeddings": self.embeddings, "projected_3d": self.projected_3d, } if self.velocity_grid is not None: arrays["velocity_grid"] = self.velocity_grid if self.confidence_grid is not None: arrays["confidence_grid"] = self.confidence_grid if self.axis_min is not None: arrays["axis_min"] = self.axis_min if self.axis_max is not None: arrays["axis_max"] = self.axis_max if self.axis_info.axes is not None: arrays["axis_info_axes"] = self.axis_info.axes if self.cluster_centroids_3d is not None: arrays["cluster_centroids_3d"] = self.cluster_centroids_3d base.parent.mkdir(parents=True, exist_ok=True) np.savez_compressed(str(base) + ".npz", **arrays) # ── JSON sidecar ─────────────────────────────── meta = { "fingerprint": self.fingerprint(), "embedding_model": self.embedding_model, "session": self.session.to_dict(), "clusters": { "n_clusters": self.clusters.n_clusters, "labels": self.clusters.labels, "clusters": self.clusters.clusters, }, "axis_info": { "lengths": self.axis_info.lengths, "labels": self.axis_info.labels, "min_point_idx": self.axis_info.min_point_idx, "max_point_idx": self.axis_info.max_point_idx, "min_cluster_idx": self.axis_info.min_cluster_idx, "max_cluster_idx": self.axis_info.max_cluster_idx, }, "cluster_labels": self.cluster_labels, "flow_model_trained": self.flow_model_trained, "max_cluster_distance": self.max_cluster_distance, "flow_mode": self.flow_mode, "seed": int(self.seed), "deterministic": bool(self.deterministic), } with open(str(base) + ".json", "w", encoding="utf-8") as f: json.dump(meta, f, indent=2, ensure_ascii=False) @classmethod def load_result(cls, path: str) -> "AnalysisResult": """Load an AnalysisResult from disk. Args: path: Base path (without extension) matching a previous save_result() call. Returns: Reconstructed AnalysisResult (without mdn_simulate or fitted_reducer). """ base = Path(path) # ── JSON sidecar ─────────────────────────────── with open(str(base) + ".json", "r", encoding="utf-8") as f: meta = json.load(f) # ── Numpy arrays ─────────────────────────────── data = np.load(str(base) + ".npz", allow_pickle=False) session = TraceSession.from_dict(meta["session"]) # Restore embeddings on entries embeddings = data["embeddings"] for i, entry in enumerate(session.entries): entry.embedding = embeddings[i] entry.model_name = meta["embedding_model"] clusters_meta = meta["clusters"] clusters = ClusterResult( n_clusters=clusters_meta["n_clusters"], labels=clusters_meta["labels"], clusters=clusters_meta["clusters"], ) ai = meta["axis_info"] axis_info = AxisInfo( axes=data["axis_info_axes"] if "axis_info_axes" in data else np.eye(3), lengths=ai["lengths"], labels=ai["labels"], min_point_idx=ai.get("min_point_idx", []), max_point_idx=ai.get("max_point_idx", []), min_cluster_idx=ai.get("min_cluster_idx", []), max_cluster_idx=ai.get("max_cluster_idx", []), ) import warnings warnings.warn( "AnalysisResult.load_result() does not restore the fitted UMAP/t-SNE " "reducer. TraceQuery projections (query_flow_at, explain_path, etc.) " "will use k-NN interpolation instead of the original transform. " "For highest accuracy, use pipeline.analyze() with cache_path.", UserWarning, stacklevel=2, ) return cls( session=session, embedding_model=meta["embedding_model"], embeddings=embeddings, clusters=clusters, projected_3d=data["projected_3d"], axis_info=axis_info, cluster_labels=meta.get("cluster_labels", []), flow_model_trained=meta.get("flow_model_trained", False), velocity_grid=data["velocity_grid"] if "velocity_grid" in data else None, confidence_grid=data["confidence_grid"] if "confidence_grid" in data else None, axis_min=data["axis_min"] if "axis_min" in data else None, axis_max=data["axis_max"] if "axis_max" in data else None, cluster_centroids_3d=data["cluster_centroids_3d"] if "cluster_centroids_3d" in data else None, max_cluster_distance=meta.get("max_cluster_distance", 1.0), flow_mode=meta.get("flow_mode", "mdn"), seed=int(meta.get("seed", 42)), deterministic=bool(meta.get("deterministic", True)), )