scripts/run_embedding_analysis.py

"""Exploratory embedding analysis for serialized AX-CPT representations.

This script reads deterministic text representations produced by
scripts/rebuild_release_outputs.py and computes local hashed token n-gram
embeddings. These are actual vector embeddings of the serialized text, but they
are not neural model embeddings and they are not latent model hidden states.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import re
import zlib
from collections import defaultdict
from pathlib import Path

import numpy as np


MODEL_NAME = "local_hashing_token_ngram_v1"
EMBEDDING_DIM = 256
NGRAM_RANGE = (1, 2)
TOKEN_RE = re.compile(r"[a-z0-9_]+")


def read_jsonl(path: Path) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def write_csv(path: Path, rows: list[dict[str, object]], fieldnames: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def tokenize(text: str) -> list[str]:
    return TOKEN_RE.findall(text.lower())


def token_ngrams(tokens: list[str], ngram_range: tuple[int, int]) -> list[str]:
    min_n, max_n = ngram_range
    out: list[str] = []
    for n in range(min_n, max_n + 1):
        if n <= 0 or len(tokens) < n:
            continue
        for i in range(0, len(tokens) - n + 1):
            out.append(" ".join(tokens[i : i + n]))
    return out


def stable_bucket(feature: str, dim: int) -> tuple[int, float]:
    data = feature.encode("utf-8")
    bucket = zlib.crc32(data) % dim
    sign = 1.0 if (zlib.crc32(b"sign:" + data) % 2 == 0) else -1.0
    return bucket, sign


def embed_text(text: str, dim: int = EMBEDDING_DIM) -> np.ndarray:
    vector = np.zeros(dim, dtype=np.float32)
    features = token_ngrams(tokenize(text), NGRAM_RANGE)
    for feature in features:
        bucket, sign = stable_bucket(feature, dim)
        vector[bucket] += sign

    norm = float(np.linalg.norm(vector))
    if norm > 0:
        vector /= norm
    return vector


def embed_rows(rows: list[dict[str, object]], dim: int = EMBEDDING_DIM) -> np.ndarray:
    embeddings = np.zeros((len(rows), dim), dtype=np.float32)
    for idx, row in enumerate(rows):
        embeddings[idx] = embed_text(str(row["serialized_text"]), dim=dim)
    return embeddings


def save_embedding_bundle(path: Path, rows: list[dict[str, object]], embeddings: np.ndarray) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        path,
        representation_id=np.array([str(row["representation_id"]) for row in rows]),
        dataset=np.array([str(row["dataset"]) for row in rows]),
        condition=np.array([str(row["condition"]) for row in rows]),
        embedding=embeddings.astype(np.float32),
    )


def metadata_rows(rows: list[dict[str, object]]) -> list[dict[str, object]]:
    out: list[dict[str, object]] = []
    for idx, row in enumerate(rows):
        out.append(
            {
                "row_idx": idx,
                "representation_id": row["representation_id"],
                "representation_level": row["representation_level"],
                "dataset": row["dataset"],
                "condition": row["condition"],
                "n_trials": row.get("n_trials", ""),
                "window_size": row.get("window_size", ""),
                "window_start_trial_idx": row.get("window_start_trial_idx", ""),
                "window_end_trial_idx": row.get("window_end_trial_idx", ""),
            }
        )
    return out


def condition_vector_rows(rows: list[dict[str, object]], embeddings: np.ndarray) -> list[dict[str, object]]:
    out: list[dict[str, object]] = []
    for idx, row in enumerate(rows):
        out.append(
            {
                "row_idx": idx,
                "representation_id": row["representation_id"],
                "dataset": row["dataset"],
                "condition": row["condition"],
                "embedding_model": MODEL_NAME,
                "embedding_dim": embeddings.shape[1],
                "embedding_vector_json": json.dumps([round(float(x), 8) for x in embeddings[idx]], separators=(",", ":")),
            }
        )
    return out


def cosine_pair_summary(
    rows: list[dict[str, object]],
    embeddings: np.ndarray,
    label: str,
) -> list[dict[str, object]]:
    groups: dict[tuple[str, str], list[int]] = defaultdict(list)
    for idx, row in enumerate(rows):
        groups[(str(row["dataset"]), str(row["condition"]))].append(idx)

    out: list[dict[str, object]] = []
    keys = sorted(groups)
    for i, key_a in enumerate(keys):
        idx_a = groups[key_a]
        emb_a = embeddings[idx_a]
        for key_b in keys[i:]:
            idx_b = groups[key_b]
            emb_b = embeddings[idx_b]
            sims = emb_a @ emb_b.T

            if key_a == key_b:
                if len(idx_a) < 2:
                    values = np.array([], dtype=np.float32)
                else:
                    mask = np.triu(np.ones(sims.shape, dtype=bool), k=1)
                    values = sims[mask]
            else:
                values = sims.reshape(-1)

            if values.size == 0:
                mean_cos = min_cos = max_cos = std_cos = mean_dist = None
            else:
                mean_cos = float(values.mean())
                min_cos = float(values.min())
                max_cos = float(values.max())
                std_cos = float(values.std(ddof=0))
                mean_dist = float((1.0 - values).mean())

            out.append(
                {
                    "representation_level": label,
                    "dataset_a": key_a[0],
                    "condition_a": key_a[1],
                    "dataset_b": key_b[0],
                    "condition_b": key_b[1],
                    "pair_type": "within_condition" if key_a == key_b else "between_condition",
                    "n_pairs": int(values.size),
                    "mean_cosine_similarity": round(mean_cos, 6) if mean_cos is not None else "",
                    "mean_cosine_distance": round(mean_dist, 6) if mean_dist is not None else "",
                    "min_cosine_similarity": round(min_cos, 6) if min_cos is not None else "",
                    "max_cosine_similarity": round(max_cos, 6) if max_cos is not None else "",
                    "std_cosine_similarity": round(std_cos, 6) if std_cos is not None else "",
                }
            )
    return out


def pca_2d(embeddings: np.ndarray) -> tuple[np.ndarray, list[float]]:
    if embeddings.shape[0] == 0:
        return np.zeros((0, 2), dtype=np.float32), [0.0, 0.0]

    centered = embeddings.astype(np.float64) - embeddings.astype(np.float64).mean(axis=0, keepdims=True)
    _, singular_values, vt = np.linalg.svd(centered, full_matrices=False)
    components = vt[:2].copy()

    for component_idx in range(components.shape[0]):
        pivot = int(np.argmax(np.abs(components[component_idx])))
        if components[component_idx, pivot] < 0:
            components[component_idx] *= -1

    coords = centered @ components.T
    total_variance = float((singular_values**2).sum())
    explained = []
    for idx in range(2):
        if idx < len(singular_values) and total_variance > 0:
            explained.append(float((singular_values[idx] ** 2) / total_variance))
        else:
            explained.append(0.0)
    return coords.astype(np.float32), explained


def projection_rows(
    rows: list[dict[str, object]],
    coords: np.ndarray,
    explained: list[float],
) -> list[dict[str, object]]:
    out: list[dict[str, object]] = []
    for idx, row in enumerate(rows):
        out.append(
            {
                "row_idx": idx,
                "representation_id": row["representation_id"],
                "representation_level": row["representation_level"],
                "dataset": row["dataset"],
                "condition": row["condition"],
                "window_size": row.get("window_size", ""),
                "window_start_trial_idx": row.get("window_start_trial_idx", ""),
                "window_end_trial_idx": row.get("window_end_trial_idx", ""),
                "pc1": round(float(coords[idx, 0]), 8),
                "pc2": round(float(coords[idx, 1]), 8),
                "pc1_explained_variance_ratio": round(explained[0], 8),
                "pc2_explained_variance_ratio": round(explained[1], 8),
            }
        )
    return out


def write_report(
    path: Path,
    condition_rows: list[dict[str, object]],
    sliding_rows: list[dict[str, object]],
    condition_explained: list[float],
    sliding_explained: list[float],
) -> None:
    report = f"""# Exploratory Embedding Analysis

This is a compact exploratory analysis of serialized AX-CPT representations. It should not be treated as evidence about latent model states or mechanistic representations.

## Inputs

- `outputs/condition_level_representations.jsonl`: {len(condition_rows)} rows.
- `outputs/sliding_window_representations.jsonl`: {len(sliding_rows)} rows.

Trial-level representations are not embedded in this first pass.

## Embedding Model

- Model/library: `{MODEL_NAME}` implemented locally in `scripts/run_embedding_analysis.py`.
- Dependency: `numpy` for vector math, cosine similarity, and PCA.
- Text processing: lowercase alphanumeric tokenization with regex `{TOKEN_RE.pattern}`.
- Features: token unigrams and bigrams.
- Vectorization: deterministic signed feature hashing with CRC32 into {EMBEDDING_DIM} dimensions.
- Normalization: L2 normalization per row.

These are actual text-derived embedding vectors for the serialized representations. They are not neural embeddings, latent model embeddings, hidden states, logits, probabilities, reaction times, costs, or latency measurements.

## Similarity

Cosine similarity and cosine distance are computed on L2-normalized hashed text embeddings. Summary files report within-condition and between-condition comparisons. Similarities reflect overlap in the serialized representation text and should be interpreted cautiously.

## Projection

2D projections use deterministic PCA via `numpy.linalg.svd` on centered embedding matrices. Component signs are fixed by forcing the largest absolute component loading to be positive.

- Condition-level PCA explained variance ratio: PC1={condition_explained[0]:.6f}, PC2={condition_explained[1]:.6f}
- Sliding-window PCA explained variance ratio: PC1={sliding_explained[0]:.6f}, PC2={sliding_explained[1]:.6f}

## Outputs

- `condition_embeddings.npz`
- `condition_embedding_vectors.csv`
- `condition_embedding_metadata.csv`
- `condition_embedding_similarity_pairs.csv`
- `condition_embedding_projection_2d.csv`
- `sliding_window_embeddings.npz`
- `sliding_window_embedding_metadata.csv`
- `sliding_window_embedding_similarity_summary.csv`
- `sliding_window_embedding_projection_2d.csv`
- `embedding_model_config.json`
"""
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(report, encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description="Run exploratory embedding analysis for AX-CPT representations.")
    parser.add_argument("--input-dir", type=Path, default=Path("outputs"))
    parser.add_argument("--output-dir", type=Path, default=Path("outputs/embedding_analysis"))
    parser.add_argument("--dim", type=int, default=EMBEDDING_DIM)
    args = parser.parse_args()

    condition_rows = read_jsonl(args.input_dir / "condition_level_representations.jsonl")
    sliding_rows = read_jsonl(args.input_dir / "sliding_window_representations.jsonl")
    args.output_dir.mkdir(parents=True, exist_ok=True)

    condition_embeddings = embed_rows(condition_rows, dim=args.dim)
    sliding_embeddings = embed_rows(sliding_rows, dim=args.dim)

    save_embedding_bundle(args.output_dir / "condition_embeddings.npz", condition_rows, condition_embeddings)
    save_embedding_bundle(args.output_dir / "sliding_window_embeddings.npz", sliding_rows, sliding_embeddings)

    metadata_fields = [
        "row_idx",
        "representation_id",
        "representation_level",
        "dataset",
        "condition",
        "n_trials",
        "window_size",
        "window_start_trial_idx",
        "window_end_trial_idx",
    ]
    write_csv(args.output_dir / "condition_embedding_metadata.csv", metadata_rows(condition_rows), metadata_fields)
    write_csv(args.output_dir / "sliding_window_embedding_metadata.csv", metadata_rows(sliding_rows), metadata_fields)

    write_csv(
        args.output_dir / "condition_embedding_vectors.csv",
        condition_vector_rows(condition_rows, condition_embeddings),
        [
            "row_idx",
            "representation_id",
            "dataset",
            "condition",
            "embedding_model",
            "embedding_dim",
            "embedding_vector_json",
        ],
    )

    similarity_fields = [
        "representation_level",
        "dataset_a",
        "condition_a",
        "dataset_b",
        "condition_b",
        "pair_type",
        "n_pairs",
        "mean_cosine_similarity",
        "mean_cosine_distance",
        "min_cosine_similarity",
        "max_cosine_similarity",
        "std_cosine_similarity",
    ]
    write_csv(
        args.output_dir / "condition_embedding_similarity_pairs.csv",
        cosine_pair_summary(condition_rows, condition_embeddings, "condition"),
        similarity_fields,
    )
    write_csv(
        args.output_dir / "sliding_window_embedding_similarity_summary.csv",
        cosine_pair_summary(sliding_rows, sliding_embeddings, "sliding_window"),
        similarity_fields,
    )

    condition_coords, condition_explained = pca_2d(condition_embeddings)
    sliding_coords, sliding_explained = pca_2d(sliding_embeddings)
    projection_fields = [
        "row_idx",
        "representation_id",
        "representation_level",
        "dataset",
        "condition",
        "window_size",
        "window_start_trial_idx",
        "window_end_trial_idx",
        "pc1",
        "pc2",
        "pc1_explained_variance_ratio",
        "pc2_explained_variance_ratio",
    ]
    write_csv(
        args.output_dir / "condition_embedding_projection_2d.csv",
        projection_rows(condition_rows, condition_coords, condition_explained),
        projection_fields,
    )
    write_csv(
        args.output_dir / "sliding_window_embedding_projection_2d.csv",
        projection_rows(sliding_rows, sliding_coords, sliding_explained),
        projection_fields,
    )

    config = {
        "analysis_label": "exploratory_embedding_analysis",
        "embedding_model": MODEL_NAME,
        "embedding_dim": args.dim,
        "library": "numpy",
        "tokenizer_regex": TOKEN_RE.pattern,
        "ngram_range": list(NGRAM_RANGE),
        "hash_function": "zlib.crc32 signed feature hashing",
        "normalization": "l2",
        "projection": "PCA via numpy.linalg.svd with deterministic component sign convention",
        "inputs": {
            "condition_level": str(args.input_dir / "condition_level_representations.jsonl"),
            "sliding_window": str(args.input_dir / "sliding_window_representations.jsonl"),
        },
        "not_included": [
            "neural embeddings",
            "latent model hidden states",
            "logits",
            "probabilities",
            "reaction times",
            "API costs",
            "latency measurements",
        ],
    }
    (args.output_dir / "embedding_model_config.json").write_text(
        json.dumps(config, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    write_report(
        args.output_dir / "exploratory_embedding_report.md",
        condition_rows=condition_rows,
        sliding_rows=sliding_rows,
        condition_explained=condition_explained,
        sliding_explained=sliding_explained,
    )

    print(f"Embedded {len(condition_rows)} condition-level rows.")
    print(f"Embedded {len(sliding_rows)} sliding-window rows.")
    print(f"Wrote exploratory embedding outputs to {args.output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())