Blame for experiments/analysis/write_alignment_causal.py

#!/usr/bin/env python3
"""Causal test: do write-aligned SAE features steer more effectively than non-aligned features?

The paper reports 17.5% of bilinear decoder atoms align with actual GDN write vectors
(combined alignment > 0.3). A reviewer asks whether this alignment is causally relevant
or merely geometric coincidence.

Protocol:
  1. Load the write alignment data (memory_alignment_L9_H4.json) to split features into
     write-aligned (combined > 0.3) and non-aligned (combined < 0.1) groups.
  2. From each group, select the top 10 by boundary-differential activation (same
     criterion as the boost experiment).
  3. Boost each group at 5x and 10x, single head (H4 only).
  4. Compare newline shift: if write-aligned features steer harder, alignment identifies
     causally relevant features.

Conditions per prompt (7 total):
  1. baseline: no intervention
  2. write_aligned_5x: boost top-10 write-aligned features at 5x
  3. write_aligned_10x: boost top-10 write-aligned features at 10x
  4. non_aligned_5x: boost top-10 non-aligned features at 5x
  5. non_aligned_10x: boost top-10 non-aligned features at 10x
  6. random_5x: boost 10 random alive features at 5x
  7. random_10x: boost 10 random alive features at 10x

Uses the same 20 instruction prompts, model (Qwen3.5-4B-Base), layer 9, head 4,
generation parameters (400 tokens, temperature 0.7) as the existing experiments.

Usage:
    modal run --detach write_alignment_causal.py
"""

from __future__ import annotations

import json
import os
import subprocess
import time
from pathlib import Path

import modal

# Modal infrastructure


def _current_code_sha() -> str:
    try:
        sha = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            capture_output=True, text=True,
            cwd=Path(__file__).resolve().parent,
        ).stdout.strip()
    except (FileNotFoundError, OSError):
        sha = ""
    return sha or os.environ.get("MATRIX_SAE_CODE_SHA", "unknown")


CURRENT_CODE_SHA = _current_code_sha()
CAUSAL_CONV1D_WHEEL = (
    "https://github.com/Dao-AILab/causal-conv1d/releases/download/"
    "v1.6.1.post4/"
    "causal_conv1d-1.6.1%2Bcu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl"
)
MAMBA_SSM_WHEEL = (
    "https://github.com/state-spaces/mamba/releases/download/"
    "v2.3.1/"
    "mamba_ssm-2.3.1%2Bcu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl"
)

image = (
    modal.Image.from_registry("nvidia/cuda:12.6.3-devel-ubuntu22.04", add_python="3.12")
    .apt_install("git", "build-essential")
    .env({
        "CUDA_HOME": "/usr/local/cuda",
        "TORCH_CUDA_ARCH_LIST": "8.0;8.6;8.9;9.0",
        "MAX_JOBS": "4",
        "CC": "gcc",
        "CXX": "g++",
        "CUDAHOSTCXX": "g++",
    })
    .run_commands("python -m pip install --upgrade pip setuptools wheel")
    .run_commands(
        "python -m pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu126"
    )
    .pip_install(
        "transformers>=5.0", "datasets", "numpy", "tqdm",
        "matplotlib", "wandb", "accelerate", "sentencepiece", "scipy", "scikit-learn",
        "einops", "ninja", "flash-linear-attention",
    )
    .run_commands(
        f"python -m pip install --no-deps '{CAUSAL_CONV1D_WHEEL}'",
        f"python -m pip install --no-deps '{MAMBA_SSM_WHEEL}'",
        "python -c \""
        "import causal_conv1d; "
        "from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined; "
        "from fla.modules import FusedRMSNormGated; "
        "from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule; "
        "print('ALL_FASTPATH_IMPORTS_OK')\"",
    )
    .add_local_file("extract_states.py", "/root/extract_states.py", copy=True)
    .add_local_file("sae.py", "/root/sae.py", copy=True)
    .add_local_file("split_utils.py", "/root/split_utils.py", copy=True)
    .add_local_file("generation_intervention.py", "/root/generation_intervention.py", copy=True)
    .add_local_file("memory_alignment.py", "/root/memory_alignment.py", copy=True)
    .env({"MATRIX_SAE_CODE_SHA": CURRENT_CODE_SHA})
)

app = modal.App("matrix-sae-write-alignment-causal")

vol = modal.Volume.from_name(
    os.environ.get("MATRIX_SAE_MODAL_DATA_VOLUME", "").strip() or "matrix-sae-data",
    create_if_missing=True,
)
model_vol = modal.Volume.from_name(
    os.environ.get("MATRIX_SAE_MODAL_MODEL_VOLUME", "").strip() or "hf-model-cache",
    create_if_missing=True,
)

DATA = "/data"
MODELS = "/models"

# Prompts (identical to existing experiments)

QUALITATIVE_PROMPTS = [
    "Write a paragraph about the history of bridges.",
    "Describe the fall of the Roman Empire in a few sentences.",
    "Summarize the key events of the French Revolution.",
    "Explain how photosynthesis works.",
    "Describe what happens inside a star during nuclear fusion.",
    "Explain why the sky is blue.",
    "Describe the process of making bread from scratch.",
    "Explain how to make a simple tomato sauce.",
    "Write a paragraph about the history of chocolate.",
    "Describe what a visitor would see walking through the streets of Tokyo.",
    "Write a paragraph about the geography of Iceland.",
    "Explain how a computer processor executes instructions.",
    "Describe how the internet routes data between computers.",
    "Write a paragraph about the invention of the printing press.",
    "Describe the water cycle from ocean to rainfall.",
    "Explain how birds migrate thousands of miles each year.",
    "Write a paragraph about the ecosystem of a coral reef.",
    "Explain why we have seasons on Earth.",
    "Describe how human memory works.",
    "Write a paragraph about the construction of the Great Wall of China.",
]


# SAE checkpoint resolution (from existing experiments)


def _normalize_corpus_source(corpus_source: str) -> str:
    aliases = {
        "openwebtext": "openwebtext", "owt": "openwebtext",
        "skylion007/openwebtext": "openwebtext",
    }
    return aliases.get(corpus_source.strip().lower(), corpus_source)


def _states_dir(corpus_source: str) -> Path:
    slug = _normalize_corpus_source(corpus_source)
    if slug == "openwebtext":
        return Path(f"{DATA}/states")
    return Path(f"{DATA}/states_{slug}")


def _experiment_tag(model_name: str, seq_len: int, n_samples: int, corpus_source: str = "openwebtext") -> str:
    model_slug = model_name.split("/")[-1].lower().replace(".", "_")
    tag = f"{model_slug}_sl{seq_len}_ns{n_samples}"
    slug = _normalize_corpus_source(corpus_source)
    if slug != "openwebtext":
        tag = f"{tag}_{slug}"
    return tag


def _resolve_sae_checkpoint(
    *, layer: int, head: int, n_features_target: int,
    corpus_source: str = "openwebtext",
    preferred_types: tuple[str, ...] = ("bilinear", "bilinear_tied", "rank1"),
) -> tuple[Path, dict, str]:
    """Find the best SAE checkpoint for a given layer/head."""
    states_dir = _states_dir(corpus_source)
    meta_path = states_dir / "metadata.json"
    if not meta_path.exists():
        raise FileNotFoundError(f"No metadata at {meta_path}")
    meta = json.loads(meta_path.read_text())
    exp_tag = _experiment_tag(meta["model"], meta["seq_len"], meta["n_samples"], corpus_source)
    ckpt_root = Path(f"{DATA}/checkpoints") / exp_tag

    best_ckpt = None
    best_cfg = None
    best_tag = None

    if not ckpt_root.exists():
        raise FileNotFoundError(f"No checkpoint root at {ckpt_root}")

    for d in sorted(ckpt_root.iterdir()):
        cp, bp = d / "config.json", d / "best.pt"
        if not cp.exists() or not bp.exists():
            continue
        cfg = json.loads(cp.read_text())
        if cfg.get("layer") != layer or cfg.get("head", 0) != head:
            continue
        if cfg.get("n_features") != n_features_target:
            continue
        sae_type = cfg.get("sae_type", "")
        if sae_type not in preferred_types:
            continue
        candidate_key = (
            preferred_types.index(sae_type) if sae_type in preferred_types else 99,
            0 if cfg.get("seed") == 42 else 1,
            cfg.get("seed", 999),
        )
        current_key = None
        if best_cfg is not None:
            best_sae_type = best_cfg.get("sae_type", "")
            current_key = (
                preferred_types.index(best_sae_type) if best_sae_type in preferred_types else 99,
                0 if best_cfg.get("seed") == 42 else 1,
                best_cfg.get("seed", 999),
            )
        if current_key is None or candidate_key < current_key:
            best_ckpt = bp
            best_cfg = cfg
            best_tag = d.name

    if best_ckpt is None:
        raise FileNotFoundError(f"No checkpoint for L{layer}H{head} nf={n_features_target}")
    return best_ckpt, best_cfg, best_tag


# Main experiment


@app.function(
    volumes={DATA: vol, MODELS: model_vol},
    gpu="A10G", image=image, timeout=14400, memory=32768,
)
def run_write_alignment_causal(
    layer: int = 9,
    head: int = 4,
    n_tokens: int = 400,
    n_boundary_features: int = 10,
    n_features_target: int = 2048,
    model_name: str = "Qwen/Qwen3.5-4B-Base",
    temperature: float = 0.7,
    random_seed: int = 42,
    period_token_id: int = 13,
    alignment_threshold_high: float = 0.3,
    alignment_threshold_low: float = 0.1,
    boost_scales: tuple[float, ...] = (5.0, 10.0),
    sae_types: tuple[str, ...] = ("bilinear", "bilinear_tied", "rank1"),
) -> dict:
    """Test whether write-aligned features steer more effectively than non-aligned features.

    Single-head experiment (H4 only, where alignment data exists).

    Feature selection pipeline:
      1. Load alignment data to partition features into write-aligned and non-aligned.
      2. Load extracted states to compute boundary-differential activation per feature.
      3. From each alignment group, select top n_boundary_features by |boundary_diff|.
      4. Boost each group at each dose; compare newline shift.
    """
    import sys
    import numpy as np
    import torch

    os.environ["HF_HOME"] = f"{MODELS}/hf_cache"
    sys.path.insert(0, "/root")

    from sae import build_sae_from_config, infer_sae_type
    from extract_states import load_model_and_tokenizer
    from generation_intervention import (
        select_boundary_features_fast,
        generate_with_intervention,
        compute_generation_stats,
    )

    vol.reload()
    t0 = time.time()
    rng = np.random.RandomState(random_seed)

    # ---------------------------------------------------------------
    # ---------------------------------------------------------------
    # The memory_alignment stage saves to /data/memory_alignment/L{layer}_H{head}/alignment_results.json
    candidates = [
        Path(f"{DATA}/memory_alignment/L{layer}_H{head}/alignment_results.json"),
        Path(f"{DATA}/analysis/memory_alignment_L{layer}_H{head}.json"),
        Path(f"{DATA}/memory_alignment_L{layer}_H{head}.json"),
    ]
    alignment_path = None
    for candidate in candidates:
        if candidate.exists():
            alignment_path = candidate
            print(f"Found alignment data at: {alignment_path}")
            break
    if alignment_path is None:
        raise FileNotFoundError(
            f"No alignment data found. Checked: {[str(c) for c in candidates]}"
        )

    alignment_data = json.loads(alignment_path.read_text())
    alignment_results = alignment_data["results"]

    # Compute combined alignment = mean(|k_cos|, |v_cos|)
    alignment_map: dict[int, float] = {}
    for r in alignment_results:
        if r.get("alive", False):
            combined = (r.get("mean_abs_k_cos", 0) + r.get("mean_abs_v_cos", 0)) / 2
            alignment_map[r["feature"]] = combined

    write_aligned_set = {f for f, c in alignment_map.items() if c > alignment_threshold_high}
    non_aligned_set = {f for f, c in alignment_map.items() if c < alignment_threshold_low}
    all_alive_set = set(alignment_map.keys())

    print(f"Alignment data loaded: {len(alignment_map)} alive features")
    print(f"Write-aligned (combined > {alignment_threshold_high}): {len(write_aligned_set)}")
    print(f"Non-aligned (combined < {alignment_threshold_low}): {len(non_aligned_set)}")

    # ---------------------------------------------------------------
    # ---------------------------------------------------------------
    corpus_source = "openwebtext"
    states_dir = _states_dir(corpus_source)
    meta_path = states_dir / "metadata.json"
    if not meta_path.exists():
        raise FileNotFoundError(f"No metadata at {meta_path}. Run extraction first.")
    meta = json.loads(meta_path.read_text())
    key_dim, val_dim = meta["key_head_dim"], meta["value_head_dim"]

    ckpt_path, cfg, resolved_tag = _resolve_sae_checkpoint(
        layer=layer, head=head, n_features_target=n_features_target,
        corpus_source=corpus_source, preferred_types=sae_types,
    )
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    sae_model = build_sae_from_config(
        cfg, state_dict=ckpt["model_state_dict"],
        default_d_k=key_dim, default_d_v=val_dim,
    )
    sae_model.load_state_dict(ckpt["model_state_dict"])
    sae_model = sae_model.cuda().eval()
    sae_type = infer_sae_type(cfg, ckpt["model_state_dict"])
    print(f"Loaded SAE: {resolved_tag} (type={sae_type})")

    # ---------------------------------------------------------------
    # Load corpus, build boundary mask, compute boundary-differential
    # ---------------------------------------------------------------
    corpus_path = states_dir / "corpus.npy"
    corpus_arr = np.load(str(corpus_path), mmap_mode="r")
    n_corpus_seqs = corpus_arr.shape[0]
    seq_len = corpus_arr.shape[1]

    corpus_full = np.array(corpus_arr)
    boundary_mask = np.zeros((n_corpus_seqs, seq_len), dtype=bool)
    boundary_mask[:, :-1] = (corpus_full[:, 1:] == period_token_id)

    head_states_path = states_dir / f"layer_{layer}" / f"head_{head}.npy"
    if not head_states_path.exists():
        raise FileNotFoundError(f"No extracted states at {head_states_path}")

    head_states = np.load(str(head_states_path), mmap_mode="r")
    n_states = head_states.shape[0]
    n_use_seqs = min(n_states // seq_len, n_corpus_seqs)
    flat_mask = boundary_mask[:n_use_seqs].reshape(-1)
    states_flat = head_states[:n_use_seqs * seq_len]

    boundary_indices = np.where(flat_mask)[0]
    nonboundary_indices = np.where(~flat_mask)[0]
    n_sample = min(len(boundary_indices), len(nonboundary_indices), 10000)

    b_idx = rng.choice(boundary_indices, size=n_sample, replace=False)
    nb_idx = rng.choice(nonboundary_indices, size=n_sample, replace=False)

    print(f"Corpus: {n_corpus_seqs} seqs x {seq_len} tokens")
    print(f"Boundary positions: {int(boundary_mask.sum())}")
    print(f"Sampled: {n_sample} boundary + {n_sample} non-boundary states")

    # Get boundary-differential ranking for ALL features
    all_ranked = select_boundary_features_fast(
        sae_model, sae_type,
        np.array(states_flat[b_idx]), np.array(states_flat[nb_idx]),
        n_features=n_features_target,  # rank all features
    )

    # Build lookup: feature_idx -> rank info
    ranked_lookup = {f["feature_idx"]: f for f in all_ranked}

    # ---------------------------------------------------------------
    # Select features for each group (boundary-differential within group)
    # ---------------------------------------------------------------

    def _select_top_boundary_diff(feature_set: set[int], n: int, sign: str = "positive") -> list[dict]:
        """From a set of feature indices, select top n by boundary-differential."""
        candidates = []
        for fidx in feature_set:
            info = ranked_lookup.get(fidx)
            if info is None:
                continue
            if sign == "positive" and info["mean_diff"] <= 0:
                continue
            if sign == "negative" and info["mean_diff"] >= 0:
                continue
            candidates.append(info)
        candidates.sort(key=lambda x: abs(x["mean_diff"]), reverse=True)
        return candidates[:n]

    write_aligned_features = _select_top_boundary_diff(write_aligned_set, n_boundary_features, sign="positive")
    non_aligned_features = _select_top_boundary_diff(non_aligned_set, n_boundary_features, sign="positive")

    # Random alive features (excluding both groups)
    excluded = set(f["feature_idx"] for f in write_aligned_features + non_aligned_features)
    random_pool = [idx for idx in all_alive_set if idx not in excluded]
    random_selection = rng.choice(random_pool, size=n_boundary_features, replace=False).tolist()
    random_features = [ranked_lookup.get(fi, {"feature_idx": fi, "mean_boundary_nonzero": 0.01}) for fi in random_selection]

    print("\n--- Feature Selection ---")
    print(f"Write-aligned group ({len(write_aligned_features)} features):")
    for f in write_aligned_features:
        alignment = alignment_map.get(f["feature_idx"], 0)
        print(f"  F{f['feature_idx']}: boundary_diff={f['mean_diff']:+.4f}, alignment={alignment:.3f}")

    print(f"\nNon-aligned group ({len(non_aligned_features)} features):")
    for f in non_aligned_features:
        alignment = alignment_map.get(f["feature_idx"], 0)
        print(f"  F{f['feature_idx']}: boundary_diff={f['mean_diff']:+.4f}, alignment={alignment:.3f}")

    print(f"\nRandom control ({len(random_features)} features):")
    for f in random_features:
        alignment = alignment_map.get(f["feature_idx"], 0)
        print(f"  F{f['feature_idx']}: alignment={alignment:.3f}")

    if len(write_aligned_features) < n_boundary_features:
        print(f"\nWARNING: Only {len(write_aligned_features)} write-aligned boundary-differential "
              f"features found (need {n_boundary_features})")
    if len(non_aligned_features) < n_boundary_features:
        print(f"\nWARNING: Only {len(non_aligned_features)} non-aligned boundary-differential "
              f"features found (need {n_boundary_features})")

    # ---------------------------------------------------------------
    # Compute calibrated push values (same calibration as boost experiment)
    # ---------------------------------------------------------------

    def _make_push_dict(features: list[dict], strength: float) -> dict[int, float]:
        """Build additive push dict: feature_idx -> push_value * strength."""
        pushes = {}
        for f in features:
            fi = f["feature_idx"]
            # Use mean nonzero boundary activation as base push
            base = max(abs(float(f.get("mean_boundary_nonzero", 0.0))), 0.01)
            pushes[fi] = base * strength
        return pushes

    def _make_random_push_dict(features: list[dict], strength: float) -> dict[int, float]:
        """Build push dict for random features using their own mean nonzero."""
        pushes = {}
        for f in features:
            fi = f["feature_idx"]
            base = max(abs(float(f.get("mean_boundary_nonzero", f.get("mean_nonboundary_nonzero", 0.01)))), 0.01)
            pushes[fi] = base * strength
        return pushes

    rand_indices = [f["feature_idx"] for f in random_features]

    conditions: dict[str, dict[int, dict[int, float]]] = {"baseline": {}}

    for scale in boost_scales:
        scale_tag = f"{int(scale)}x" if scale == int(scale) else f"{scale:g}x"

        # Write-aligned boost
        conditions[f"write_aligned_{scale_tag}"] = {
            head: _make_push_dict(write_aligned_features, scale),
        }
        # Non-aligned boost
        conditions[f"non_aligned_{scale_tag}"] = {
            head: _make_push_dict(non_aligned_features, scale),
        }
        # Random boost
        conditions[f"random_{scale_tag}"] = {
            head: _make_random_push_dict(random_features, scale),
        }

    cond_names = list(conditions.keys())
    print(f"\nConditions: {cond_names}")
    print(f"Boost scales: {boost_scales}")

    # Report calibration values
    for label, features in [("write_aligned", write_aligned_features),
                            ("non_aligned", non_aligned_features),
                            ("random", random_features)]:
        base_pushes = [max(abs(float(f.get("mean_boundary_nonzero",
                       f.get("mean_nonboundary_nonzero", 0.01)))), 0.01) for f in features]
        if base_pushes:
            print(f"  {label} base push: min={min(base_pushes):.4f}, max={max(base_pushes):.4f}, "
                  f"mean={np.mean(base_pushes):.4f}")

    # ---------------------------------------------------------------
    # ---------------------------------------------------------------
    print(f"\nLoading model: {model_name}")
    model, tokenizer, _ = load_model_and_tokenizer(model_name, device="cuda")

    raw_prompts = list(QUALITATIVE_PROMPTS)
    if hasattr(tokenizer, "apply_chat_template"):
        formatted_prompts = []
        for p in raw_prompts:
            messages = [{"role": "user", "content": p}]
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True,
            )
            text = text.replace("<think>\n\n</think>\n\n", "")
            formatted_prompts.append(text)
        print(f"Formatted {len(formatted_prompts)} prompts with chat template")
    else:
        formatted_prompts = raw_prompts

    # ---------------------------------------------------------------
    # ---------------------------------------------------------------
    device = next(model.parameters()).device
    sae_per_head = {head: sae_model}
    all_results: list[dict] = []
    total_runs = len(formatted_prompts) * len(conditions)
    run_idx = 0
    t_gen = time.time()

    for i, (raw_prompt, fmt_prompt) in enumerate(zip(raw_prompts, formatted_prompts)):
        prompt_ids = tokenizer(fmt_prompt, return_tensors="pt")["input_ids"].to(device)
        entry = {
            "prompt_idx": i,
            "prompt_text": raw_prompt,
            "prompt_tokens": int(prompt_ids.shape[1]),
        }

        for cond_name, updates_per_head in conditions.items():
            gen_ids, meta = generate_with_intervention(
                model=model,
                tokenizer=tokenizer,
                sae_type=sae_type,
                layer_idx=layer,
                prompt_ids=prompt_ids,
                n_tokens=n_tokens,
                temperature=temperature,
                sae_per_head=sae_per_head if updates_per_head else {},
                feature_updates_per_head=updates_per_head,
                additive=True,
                additive_min_zero=False,
            )
            gen_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
            gen_stats = compute_generation_stats(gen_text)

            entry[cond_name] = {
                "text": gen_text,
                "stats": gen_stats,
                "n_generated": meta["n_generated"],
                "mean_intervention_norm": meta["mean_intervention_norm"],
            }
            run_idx += 1

        elapsed = time.time() - t_gen
        rate = run_idx / elapsed if elapsed > 0 else 0
        remaining = (total_runs - run_idx) / rate if rate > 0 else 0

        base_nl = entry["baseline"]["stats"]["n_newlines"]
        wa_nl = entry.get("write_aligned_10x", {}).get("stats", {}).get("n_newlines", 0)
        na_nl = entry.get("non_aligned_10x", {}).get("stats", {}).get("n_newlines", 0)
        print(
            f"  [{i+1}/{len(formatted_prompts)}] {elapsed:.0f}s elapsed, "
            f"{remaining:.0f}s remaining | "
            f"newlines: baseline={base_nl:.0f} write_aligned_10x={wa_nl:.0f} non_aligned_10x={na_nl:.0f}"
        )
        all_results.append(entry)

    total_time = time.time() - t0

    # ---------------------------------------------------------------
    # ---------------------------------------------------------------
    from scipy import stats as sp_stats

    stat_names = list(compute_generation_stats("test text.").keys())
    condition_means: dict[str, dict[str, float]] = {}
    per_cond_arrays: dict[str, dict[str, list[float]]] = {}

    for cond_name in cond_names:
        means = {}
        arrays = {}
        for stat in stat_names:
            vals = [r[cond_name]["stats"][stat] for r in all_results]
            means[stat] = float(np.mean(vals))
            arrays[stat] = vals
        condition_means[cond_name] = means
        per_cond_arrays[cond_name] = arrays

    # ---------------------------------------------------------------
    # Statistical tests
    # ---------------------------------------------------------------
    n_bootstrap = 10000

    def bootstrap_ci(vals_a, vals_b, n_boot=n_bootstrap, alpha=0.05):
        a = np.array(vals_a)
        b = np.array(vals_b)
        diffs = a - b
        observed = float(np.mean(diffs))
        boot_means = np.array([
            np.mean(rng.choice(diffs, size=len(diffs), replace=True))
            for _ in range(n_boot)
        ])
        ci_lo = float(np.percentile(boot_means, 100 * alpha / 2))
        ci_hi = float(np.percentile(boot_means, 100 * (1 - alpha / 2)))
        if observed >= 0:
            p = 2 * float(np.mean(boot_means < 0))
        else:
            p = 2 * float(np.mean(boot_means > 0))
        return observed, ci_lo, ci_hi, min(p, 1.0)

    # Key comparisons
    comparisons = []
    for scale in boost_scales:
        tag = f"{int(scale)}x" if scale == int(scale) else f"{scale:g}x"
        comparisons.append((f"write_aligned_{tag}_vs_baseline", f"write_aligned_{tag}", "baseline"))
        comparisons.append((f"non_aligned_{tag}_vs_baseline", f"non_aligned_{tag}", "baseline"))
        comparisons.append((f"random_{tag}_vs_baseline", f"random_{tag}", "baseline"))
        # Direct comparison: write-aligned vs non-aligned
        comparisons.append((f"write_aligned_vs_non_aligned_{tag}", f"write_aligned_{tag}", f"non_aligned_{tag}"))

    tests: dict[str, dict[str, dict]] = {}
    for comparison, cond_a, cond_b in comparisons:
        comp = {}
        for stat in stat_names:
            vals_a = per_cond_arrays[cond_a][stat]
            vals_b = per_cond_arrays[cond_b][stat]

            t_stat, t_pval = sp_stats.ttest_rel(vals_a, vals_b)
            diff = np.array(vals_a) - np.array(vals_b)
            mean_diff = float(np.mean(diff))
            std_diff = float(np.std(diff, ddof=1))
            cohens_d = mean_diff / std_diff if std_diff > 1e-12 else 0.0

            obs, ci_lo, ci_hi, boot_p = bootstrap_ci(vals_a, vals_b)

            direction = "higher" if mean_diff > 0 else "lower"
            comp[stat] = {
                "mean_a": float(np.mean(vals_a)),
                "mean_b": float(np.mean(vals_b)),
                "mean_diff": mean_diff,
                "direction": direction,
                "t_stat": float(t_stat),
                "t_pvalue": float(t_pval),
                "cohens_d": cohens_d,
                "bootstrap_mean_diff": obs,
                "bootstrap_ci_95": [ci_lo, ci_hi],
                "bootstrap_pvalue": boot_p,
            }
        tests[comparison] = comp

    # ---------------------------------------------------------------
    # ---------------------------------------------------------------
    analysis_dir = Path(f"{DATA}/analysis")
    analysis_dir.mkdir(parents=True, exist_ok=True)

    result = {
        "experiment": "write_alignment_causal",
        "layer": layer,
        "head": head,
        "n_prompts": len(formatted_prompts),
        "n_tokens": n_tokens,
        "n_boundary_features": n_boundary_features,
        "temperature": temperature,
        "sae_type": sae_type,
        "model_name": model_name,
        "alignment_threshold_high": alignment_threshold_high,
        "alignment_threshold_low": alignment_threshold_low,
        "boost_scales": list(boost_scales),
        "total_time_s": total_time,
        "n_write_aligned_available": len(write_aligned_set),
        "n_non_aligned_available": len(non_aligned_set),
        "write_aligned_features": [
            {"feature_idx": f["feature_idx"],
             "mean_diff": f["mean_diff"],
             "alignment": alignment_map.get(f["feature_idx"], 0)}
            for f in write_aligned_features
        ],
        "non_aligned_features": [
            {"feature_idx": f["feature_idx"],
             "mean_diff": f["mean_diff"],
             "alignment": alignment_map.get(f["feature_idx"], 0)}
            for f in non_aligned_features
        ],
        "random_features": [
            {"feature_idx": fi,
             "alignment": alignment_map.get(fi, 0)}
            for fi in rand_indices
        ],
        "conditions": condition_means,
        "condition_names": cond_names,
        "tests": tests,
        "per_prompt": all_results,
    }

    filename = f"write_alignment_causal_L{layer}_H{head}.json"
    out_path = analysis_dir / filename
    out_path.write_text(json.dumps(result, indent=2, default=str))
    vol.commit()
    print(f"\nSaved to {out_path}")
    result["output_path"] = str(out_path)

    # ---------------------------------------------------------------
    # ---------------------------------------------------------------
    print(f"\n{'='*80}")
    print(f"WRITE ALIGNMENT CAUSAL TEST: L{layer}H{head}")
    print(f"{'='*80}")

    for cond_name in cond_names:
        m = condition_means[cond_name]
        print(f"  {cond_name:25s}: newlines={m['n_newlines']:.1f}  paragraphs={m['n_paragraphs']:.1f}  "
              f"words={m['n_words']:.0f}  mean_word_len={m['mean_word_length']:.2f}")

    print("\n--- Key Statistical Tests ---")
    key_stats = ["n_newlines", "n_paragraphs"]
    for comp_name, comp_data in tests.items():
        print(f"\n  {comp_name}:")
        for stat in key_stats:
            s = comp_data[stat]
            sig = "*" if s["t_pvalue"] < 0.05 else ""
            print(f"    {stat:22s}: diff={s['mean_diff']:+.2f}  d={s['cohens_d']:+.3f}  "
                  f"p={s['t_pvalue']:.4f}{sig}  "
                  f"boot_ci=[{s['bootstrap_ci_95'][0]:+.2f}, {s['bootstrap_ci_95'][1]:+.2f}]")

    # The key question
    print(f"\n{'='*80}")
    print("KEY QUESTION: Do write-aligned features steer harder?")
    print(f"{'='*80}")
    for scale in boost_scales:
        tag = f"{int(scale)}x" if scale == int(scale) else f"{scale:g}x"
        wa_nl = condition_means[f"write_aligned_{tag}"]["n_newlines"]
        na_nl = condition_means[f"non_aligned_{tag}"]["n_newlines"]
        rand_nl = condition_means[f"random_{tag}"]["n_newlines"]
        base_nl = condition_means["baseline"]["n_newlines"]
        wa_shift = wa_nl - base_nl
        na_shift = na_nl - base_nl
        rand_shift = rand_nl - base_nl

        direct = tests[f"write_aligned_vs_non_aligned_{tag}"]
        nl_direct = direct["n_newlines"]
        print(f"\n  At {tag} dose:")
        print(f"    Write-aligned newlines: {wa_nl:.1f} (shift {wa_shift:+.1f})")
        print(f"    Non-aligned newlines:   {na_nl:.1f} (shift {na_shift:+.1f})")
        print(f"    Random newlines:        {rand_nl:.1f} (shift {rand_shift:+.1f})")
        print(f"    Write-aligned vs non-aligned: diff={nl_direct['mean_diff']:+.2f}, "
              f"p={nl_direct['t_pvalue']:.4f}, d={nl_direct['cohens_d']:+.3f}")

    print(f"\nTotal time: {total_time:.0f}s")
    return result


# Local entrypoint


@app.local_entrypoint()
def main():
    result = run_write_alignment_causal.remote()

    local_out = os.path.join(os.path.dirname(__file__), "results", "data")
    os.makedirs(local_out, exist_ok=True)
    local_name = os.path.basename(result.get("output_path", ""))
    if not local_name:
        local_name = "write_alignment_causal_L9_H4.json"
    local_path = os.path.join(local_out, local_name)
    with open(local_path, "w") as f:
        json.dump(result, f, indent=2, default=str)
    print(f"\nResults saved locally to {local_path}")

    tests = result.get("tests", {})
    conds = result.get("conditions", {})

    print("\n" + "=" * 70)
    print("VERDICT")
    print("=" * 70)

    for tag in ["10x"]:
        direct_key = f"write_aligned_vs_non_aligned_{tag}"
        if direct_key in tests:
            nl = tests[direct_key].get("n_newlines", {})
            p = nl.get("t_pvalue", 1.0)
            d = nl.get("cohens_d", 0)
            wa_nl = conds.get(f"write_aligned_{tag}", {}).get("n_newlines", 0)
            na_nl = conds.get(f"non_aligned_{tag}", {}).get("n_newlines", 0)

            if p < 0.05:
                if wa_nl < na_nl:
                    print(f"  Write-aligned features steer HARDER at {tag} (p={p:.4f}, d={d:+.3f})")
                    print("  Alignment metric identifies causally relevant features.")
                else:
                    print(f"  Non-aligned features steer harder at {tag} (p={p:.4f}, d={d:+.3f})")
                    print("  Alignment metric does NOT predict causal relevance.")
            else:
                print(f"  No significant difference at {tag} (p={p:.4f}, d={d:+.3f})")
                print(f"  Inconclusive with {result.get('n_prompts', 20)} prompts.")