experiments/analyze_stability.py

from __future__ import annotations

import argparse
import csv
import json
from collections import defaultdict
from pathlib import Path

import numpy as np
import scipy.sparse as sp

from experiments.common import ARTIFACT_DIR


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Estimate prompt-wide candidate selection stability from saved SAE activations.'
    )
    parser.add_argument('--activation-dir', type=Path, default=ARTIFACT_DIR / 'activations')
    parser.add_argument('--output', type=Path, default=ARTIFACT_DIR / 'selection_stability.csv')
    parser.add_argument('--resamples', type=int, default=128)
    parser.add_argument('--top-features', type=int, default=20)
    parser.add_argument('--seed', type=int, default=42)
    return parser.parse_args()


def _vector_mean(matrix: sp.csr_matrix, indices: np.ndarray) -> np.ndarray:
    if indices.size == 0:
        return np.zeros(matrix.shape[1], dtype=np.float32)
    return np.asarray(matrix[indices].mean(axis=0)).ravel().astype(np.float32, copy=False)


def _vector_rate(binary: sp.csr_matrix, indices: np.ndarray) -> np.ndarray:
    if indices.size == 0:
        return np.zeros(binary.shape[1], dtype=np.float32)
    return np.asarray(binary[indices].mean(axis=0)).ravel().astype(np.float32, copy=False)


def balanced_candidate_score(
    target_mean: np.ndarray,
    other_mean: np.ndarray,
    target_rate: np.ndarray,
) -> np.ndarray:
    """Live-compatible selectivity × coverage × log-magnitude candidate score."""
    target = np.asarray(target_mean, dtype=np.float64)
    other = np.asarray(other_mean, dtype=np.float64)
    rate = np.asarray(target_rate, dtype=np.float64)
    difference = target - other
    denom = np.abs(target) + np.abs(other) + 1e-12
    selectivity = np.where(difference > 0.0, difference / denom, 0.0)
    return selectivity * np.clip(rate, 0.0, 1.0) * np.log1p(np.maximum(target, 0.0))


def _rank_top(scores: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
    positive = np.flatnonzero(scores > 0.0)
    if positive.size == 0:
        return np.empty(0, dtype=int), np.empty(0, dtype=float)
    k = min(int(k), int(positive.size))
    candidate_scores = scores[positive]
    local = np.argpartition(candidate_scores, -k)[-k:]
    ids = positive[local]
    ordered = np.argsort(scores[ids])[::-1]
    ids = ids[ordered]
    return ids.astype(int), scores[ids].astype(float)


def _sample_indices(
    concept_indices: dict[str, np.ndarray],
    rng: np.random.Generator,
) -> dict[str, np.ndarray]:
    return {
        concept: rng.choice(indices, size=indices.size, replace=True)
        for concept, indices in concept_indices.items()
    }


def _concept_statistics(
    matrix: sp.csr_matrix,
    binary: sp.csr_matrix,
    concept_indices: dict[str, np.ndarray],
) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
    means = {
        concept: _vector_mean(matrix, indices)
        for concept, indices in concept_indices.items()
    }
    rates = {
        concept: _vector_rate(binary, indices)
        for concept, indices in concept_indices.items()
    }
    return means, rates


def _score_for_concept(
    concept: str,
    means: dict[str, np.ndarray],
    rates: dict[str, np.ndarray],
) -> np.ndarray:
    other_concepts = [name for name in means if name != concept]
    other_mean = np.mean(np.stack([means[name] for name in other_concepts]), axis=0)
    return balanced_candidate_score(means[concept], other_mean, rates[concept])


def main() -> None:
    args = parse_args()
    if args.resamples < 1:
        raise ValueError('--resamples must be at least 1.')
    if args.top_features < 1:
        raise ValueError('--top-features must be at least 1.')

    metadata = json.loads((args.activation_dir / 'metadata.json').read_text(encoding='utf-8'))
    pooling = metadata.get('feature_pooling', '')
    if 'prompt-wide' not in pooling:
        raise RuntimeError(
            'Selection stability requires v0.14 prompt-wide activation artifacts. '
            'Rerun experiments.collect_activations before this analysis.'
        )

    rows = metadata['rows']
    layers = [int(layer) for layer in metadata['layers']]
    concepts = sorted({row['concept'] for row in rows})
    concept_indices = {
        concept: np.array(
            [idx for idx, row in enumerate(rows) if row['concept'] == concept],
            dtype=int,
        )
        for concept in concepts
    }
    rng = np.random.default_rng(args.seed)
    output_rows: list[dict] = []

    for layer in layers:
        matrix = sp.load_npz(args.activation_dir / f'features_layer{layer}.npz').tocsr()
        binary = matrix.copy()
        binary.data = np.ones_like(binary.data, dtype=np.float32)

        full_means, full_rates = _concept_statistics(matrix, binary, concept_indices)
        full_scores = {
            concept: _score_for_concept(concept, full_means, full_rates)
            for concept in concepts
        }
        full_orders: dict[str, np.ndarray] = {
            concept: np.argsort(scores)[::-1]
            for concept, scores in full_scores.items()
        }
        support: dict[tuple[str, int], int] = defaultdict(int)
        ranks: dict[tuple[str, int], list[int]] = defaultdict(list)

        for _ in range(args.resamples):
            sampled = _sample_indices(concept_indices, rng)
            means, rates = _concept_statistics(matrix, binary, sampled)
            for concept in concepts:
                scores = _score_for_concept(concept, means, rates)
                ids, _ = _rank_top(scores, args.top_features)
                for rank, feature_id in enumerate(ids.tolist(), start=1):
                    key = (concept, int(feature_id))
                    support[key] += 1
                    ranks[key].append(rank)

        for concept in concepts:
            full_rank_map = np.empty(matrix.shape[1], dtype=np.int32)
            full_rank_map[full_orders[concept]] = np.arange(1, matrix.shape[1] + 1, dtype=np.int32)
            seen = {
                feature_id
                for (seen_concept, feature_id), count in support.items()
                if seen_concept == concept and count > 0
            }
            full_ids, _ = _rank_top(full_scores[concept], max(args.top_features, 50))
            seen.update(int(feature_id) for feature_id in full_ids.tolist())

            for feature_id in seen:
                key = (concept, feature_id)
                feature_ranks = ranks.get(key, [])
                output_rows.append(
                    {
                        'layer': layer,
                        'concept': concept,
                        'feature_id': feature_id,
                        'full_score': float(full_scores[concept][feature_id]),
                        'full_rank': int(full_rank_map[feature_id]),
                        'resample_support': float(support.get(key, 0) / args.resamples),
                        'median_resample_rank': (
                            float(np.median(feature_ranks)) if feature_ranks else float('nan')
                        ),
                        'mean_resample_rank': (
                            float(np.mean(feature_ranks)) if feature_ranks else float('nan')
                        ),
                        'resamples': int(args.resamples),
                        'top_features_per_resample': int(args.top_features),
                    }
                )
        print(f'Stability analysis complete for layer {layer}', flush=True)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        'layer',
        'concept',
        'feature_id',
        'full_score',
        'full_rank',
        'resample_support',
        'median_resample_rank',
        'mean_resample_rank',
        'resamples',
        'top_features_per_resample',
    ]
    with args.output.open('w', newline='', encoding='utf-8') as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(output_rows)
    print(f'Wrote {len(output_rows)} selection-stability rows to {args.output}')


if __name__ == '__main__':
    main()