logogram 0.1.0__py3-none-any.whl

logogram/spec.py

@@ -0,0 +1,518 @@
+"""The experiment spec: the single description of an experiment that the GUI and CLI both run.
+
+A spec is stored as ``experiments/<id>/spec.json``. Every methodological choice that can change
+a number (direction, baseline, position, metric, normalization, seeds, batch size, dtype) lives
+here explicitly, so a run can always be traced back to how it was produced.
+
+Specs refer to abstract sites (``resid_pre``, ``attn_out``, ``head``, ...). Backend-specific hook
+names never appear in a spec.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+NAME_MAX = 200
+
+SPEC_VERSION = 1
+
+StreamSiteKind = Literal["resid_pre", "resid_mid", "resid_post", "attn_out", "mlp_out"]
+SiteKind = Literal[
+    "resid_pre", "resid_mid", "resid_post", "attn_out", "mlp_out", "head", "sae_feature"
+]
+
+SITE_KIND_LABELS: dict[str, str] = {
+    "resid_pre": "residual stream before the layer",
+    "resid_mid": "residual stream between attention and MLP",
+    "resid_post": "residual stream after the layer",
+    "attn_out": "attention output",
+    "mlp_out": "MLP output",
+    "head": "attention head output (z)",
+    "sae_feature": "SAE feature",
+}
+
+
+class _Strict(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+# ---------------------------------------------------------------------------------------------
+# Positions
+
+
+class AllPositions(_Strict):
+    """Intervene on every token position at once."""
+
+    kind: Literal["all"] = "all"
+
+
+class LastPosition(_Strict):
+    """The final token of each prompt (the position that predicts the answer)."""
+
+    kind: Literal["last"] = "last"
+
+
+class IndexPosition(_Strict):
+    """A fixed token index. Negative values count from the end (-1 is the last token)."""
+
+    kind: Literal["index"] = "index"
+    index: int
+
+
+class LabelPosition(_Strict):
+    """A named position from the dataset, such as ``S2`` in IOI prompts."""
+
+    kind: Literal["label"] = "label"
+    label: str = Field(min_length=1)
+
+
+Position = Annotated[
+    AllPositions | LastPosition | IndexPosition | LabelPosition, Field(discriminator="kind")
+]
+
+
+def describe_position(position: AllPositions | LastPosition | IndexPosition | LabelPosition) -> str:
+    if isinstance(position, AllPositions):
+        return "all positions"
+    if isinstance(position, LastPosition):
+        return "the last token"
+    if isinstance(position, IndexPosition):
+        return f"token index {position.index}"
+    return f"position {position.label}"
+
+
+# ---------------------------------------------------------------------------------------------
+# Sites and sweep scopes
+
+
+class Site(_Strict):
+    kind: SiteKind
+    layer: int = Field(ge=0)
+    head: int | None = Field(default=None, ge=0)
+    # A feature of the spec's SAE (kind sae_feature): its index among the SAE's features.
+    feature: int | None = Field(default=None, ge=0)
+    position: Position = Field(default_factory=AllPositions)
+
+    @model_validator(mode="after")
+    def _head_matches_kind(self) -> Site:
+        if self.kind == "head" and self.head is None:
+            raise ValueError("a head site needs a head index")
+        if self.kind != "head" and self.head is not None:
+            raise ValueError(f"a {self.kind} site has no head index")
+        if self.kind == "sae_feature" and self.feature is None:
+            raise ValueError("an SAE feature site needs a feature index")
+        if self.kind != "sae_feature" and self.feature is not None:
+            raise ValueError(f"a {self.kind} site has no feature index")
+        return self
+
+
+class HeadsScope(_Strict):
+    """Every attention head in every layer, at one position: a layer x head grid."""
+
+    kind: Literal["heads"] = "heads"
+    position: Position = Field(default_factory=AllPositions)
+
+
+class LayerPositionScope(_Strict):
+    """One stream site at every layer and every position: a layer x position grid.
+
+    ``positions="each"`` uses every token index and needs prompts of equal length.
+    ``positions="labels"`` uses the dataset's named positions, which works across templates.
+    """
+
+    kind: Literal["layer_position"] = "layer_position"
+    site: StreamSiteKind = "resid_pre"
+    positions: Literal["each", "labels"] = "each"
+
+
+class LayerComponentsScope(_Strict):
+    """Attention and MLP outputs of every layer, at one position: a layer x component grid."""
+
+    kind: Literal["layer_components"] = "layer_components"
+    components: list[StreamSiteKind] = Field(default_factory=lambda: ["attn_out", "mlp_out"])
+    position: Position = Field(default_factory=AllPositions)
+
+    @model_validator(mode="after")
+    def _non_empty(self) -> LayerComponentsScope:
+        if not self.components:
+            raise ValueError("choose at least one component")
+        if len(set(self.components)) != len(self.components):
+            raise ValueError("components must not repeat")
+        return self
+
+
+class SitesScope(_Strict):
+    """An explicit list of sites, for example a single head chosen on the map."""
+
+    kind: Literal["sites"] = "sites"
+    sites: list[Site] = Field(min_length=1)
+
+
+class FeaturesScope(_Strict):
+    """Every feature of the spec's SAE, at one position (or summed over all of them).
+
+    Only attribution patching can sweep every feature: it estimates them all from one gradient
+    and keeps the ``top`` with the largest estimated effects (by magnitude) as the run's sites,
+    ready to verify by patching.
+    """
+
+    kind: Literal["features"] = "features"
+    position: Position = Field(default_factory=AllPositions)
+    top: int = Field(ge=1, le=500)
+
+
+Scope = Annotated[
+    HeadsScope | LayerPositionScope | LayerComponentsScope | SitesScope | FeaturesScope,
+    Field(discriminator="kind"),
+]
+
+
+# ---------------------------------------------------------------------------------------------
+# Experiments
+
+
+class ActivationPatching(_Strict):
+    """Copy activations from one prompt of each pair into the run on the other.
+
+    ``clean_to_corrupt`` runs the corrupt prompt and patches in clean activations
+    ("does this restore the behavior?"). ``corrupt_to_clean`` runs the clean prompt and patches
+    in corrupt activations ("does this break it?").
+    """
+
+    kind: Literal["activation_patching"] = "activation_patching"
+    direction: Literal["clean_to_corrupt", "corrupt_to_clean"]
+
+
+class ZeroBaseline(_Strict):
+    kind: Literal["zero"] = "zero"
+
+
+class MeanBaseline(_Strict):
+    """Replace the activation with its mean over a stated reference set of prompts.
+
+    With an all-positions intervention the mean is taken per position, over reference prompts
+    of the same token length. With a single position it is taken at each reference prompt's own
+    resolved position.
+    """
+
+    kind: Literal["mean"] = "mean"
+    reference: Literal["clean", "corrupt"] = "corrupt"
+
+
+class ResampleBaseline(_Strict):
+    """Replace the activation with the same activation from other prompts (donors).
+
+    Each prompt draws ``donors`` donors from the pool without replacement, never itself, using
+    ``seed``. The patched metric is averaged over donors. The same donors are used at every site.
+    """
+
+    kind: Literal["resample"] = "resample"
+    pool: Literal["clean", "corrupt"] = "corrupt"
+    donors: int = Field(default=10, ge=1, le=1000)
+    seed: int = 0
+
+
+Baseline = Annotated[ZeroBaseline | MeanBaseline | ResampleBaseline, Field(discriminator="kind")]
+
+
+class Ablation(_Strict):
+    """Run the clean prompt and replace the activation with a baseline value."""
+
+    kind: Literal["ablation"] = "ablation"
+    baseline: Baseline
+
+
+class DirectLogitAttribution(_Strict):
+    """Split the logit difference of the chosen prompts into what each component writes directly.
+
+    At the last position, the residual stream is the sum of the embeddings and every attention and
+    MLP output. With the final normalization's scale held at its value in the run, the logit
+    difference is a sum of one term per component plus a constant from biases. A component's term
+    is its direct effect: it leaves out everything the component does through later components.
+    This is a decomposition of one forward pass, not an intervention.
+    """
+
+    kind: Literal["direct_logit_attribution"] = "direct_logit_attribution"
+    prompts: Literal["clean", "corrupt"]
+
+
+class AttributionPatching(_Strict):
+    """Estimate activation patching at every site from one gradient (attribution patching).
+
+    For each site, the change patching would cause is estimated as (source activation - receiver
+    activation) · the gradient of the logit difference at the receiver run, from one forward and
+    backward pass per batch of prompts instead of one patched run per site. It is a first-order
+    estimate: it misses saturation (in attention, normalization and the softmax) and can miss or
+    even invert an effect, so verify the strongest sites with activation patching.
+    """
+
+    kind: Literal["attribution_patching"] = "attribution_patching"
+    direction: Literal["clean_to_corrupt", "corrupt_to_clean"]
+
+
+class Steering(_Strict):
+    """Add a direction to the residual stream at several strengths, and measure what it does.
+
+    At each steered site, the direction is the mean difference between the two prompts of each
+    pair, from the receiver (``apply_to``) toward the other prompt. It is computed on a training
+    split of the prompts and applied to the held-out rest, so no prompt receives its own
+    difference. A strength of 1 adds the whole mean difference. The effect is normalized like
+    patching: 1 means the steered prompt moves as far as switching to the other prompt. With
+    ``control``, a random direction of the same length, at the same strengths, runs alongside.
+    """
+
+    kind: Literal["steering"] = "steering"
+    apply_to: Literal["clean", "corrupt"]
+    coefficients: list[float] = Field(min_length=1, max_length=16)
+    train_fraction: float = Field(gt=0.0, lt=1.0)
+    seed: int
+    control: bool
+
+    @model_validator(mode="after")
+    def _distinct_finite(self) -> Steering:
+        if any(c != c or c in (float("inf"), float("-inf")) for c in self.coefficients):
+            raise ValueError("steering strengths must be finite numbers")
+        if len(set(self.coefficients)) != len(self.coefficients):
+            raise ValueError("steering strengths must not repeat")
+        return self
+
+
+class HeadReceiver(_Strict):
+    """A later attention head whose query, key or value input receives the path."""
+
+    kind: Literal["head"] = "head"
+    layer: int = Field(ge=0)
+    head: int = Field(ge=0)
+    input: Literal["q", "k", "v"]
+
+
+class LogitsReceiver(_Strict):
+    """The residual stream at the end of the model, read directly by the unembedding."""
+
+    kind: Literal["logits"] = "logits"
+
+
+PathReceiver = Annotated[HeadReceiver | LogitsReceiver, Field(discriminator="kind")]
+
+
+class PathPatching(_Strict):
+    """Patch a component's effect along chosen paths only (path patching).
+
+    For each sender (a site of the scope), the receiver prompt runs with the sender's activation
+    from the source prompt and every other attention head held at its own value, so the change
+    travels only through the residual stream (and through the MLPs, unless ``freeze_mlps``). The
+    receivers' inputs from that run are recorded, then patched into an otherwise unchanged
+    receiver run, and the metric is read there. The result is the sender's effect through those
+    receivers alone. Senders must come before at least one receiver.
+    """
+
+    kind: Literal["path_patching"] = "path_patching"
+    direction: Literal["clean_to_corrupt", "corrupt_to_clean"]
+    receivers: list[PathReceiver] = Field(min_length=1, max_length=64)
+    freeze_mlps: bool
+
+    @model_validator(mode="after")
+    def _distinct(self) -> PathPatching:
+        keys = [r.model_dump_json() for r in self.receivers]
+        if len(set(keys)) != len(keys):
+            raise ValueError("path patching receivers must not repeat")
+        return self
+
+
+def receiver_text(receiver: HeadReceiver | LogitsReceiver) -> str:
+    if isinstance(receiver, LogitsReceiver):
+        return "logits"
+    return f"L{receiver.layer} H{receiver.head} {receiver.input}"
+
+
+Experiment = Annotated[
+    ActivationPatching
+    | Ablation
+    | DirectLogitAttribution
+    | AttributionPatching
+    | Steering
+    | PathPatching,
+    Field(discriminator="kind"),
+]
+
+
+def strength_text(coefficient: float) -> str:
+    """A steering strength as written on the map, for example ×2 or ×−0.5."""
+    text = f"{coefficient:g}".replace("-", "−")
+    return f"×{text}"
+
+
+# ---------------------------------------------------------------------------------------------
+# Model, data, metric, statistics
+
+
+class ModelRef(_Strict):
+    id: str = Field(min_length=1, description="Hugging Face model id")
+    revision: str | None = Field(
+        default=None, description="Exact commit. None resolves the current main branch at run time."
+    )
+    dtype: Literal["float32", "float16", "bfloat16"] = "float32"
+    device: Literal["auto", "cpu", "cuda", "mps"] = "auto"
+    process_weights: bool = Field(
+        default=True,
+        description="Fold LayerNorm weights and center writing weights and the unembedding, as "
+        "TransformerLens does by default. Logit differences are unchanged by this.",
+    )
+
+
+class DatasetRef(_Strict):
+    path: str = Field(min_length=1, description="Project-relative path to a JSONL dataset")
+    sha256: str | None = Field(default=None, description="Expected hash of the file, if pinned")
+    limit: int | None = Field(default=None, ge=1, description="Use only the first n prompts")
+
+
+class SAERef(_Strict):
+    """A published sparse autoencoder on Hugging Face: the repository, the folder holding it
+    (empty for the top level) and the exact commit, pinned when a run starts."""
+
+    repo: str = Field(min_length=1)
+    path: str = ""
+    revision: str | None = None
+
+
+class Tokenization(_Strict):
+    prepend_bos: bool = True
+
+
+class Metric(_Strict):
+    """Logit difference at the last position: logit(answer) - logit(distractor).
+
+    The normalized effect for each prompt is (patched - receiver) divided by either the
+    dataset's mean gap (``dataset_gap``) or that prompt's own gap (``prompt_gap``). The gap is
+    source - receiver for patching, and corrupt - clean for ablation. 0 means no change and 1
+    means a change as large as swapping to the other prompt.
+    """
+
+    kind: Literal["logit_diff"] = "logit_diff"
+    normalization: Literal["dataset_gap", "prompt_gap"] = "dataset_gap"
+
+
+class Statistics(_Strict):
+    bootstrap: int = Field(default=1000, ge=100, le=100_000)
+    ci: float = Field(default=0.95, gt=0.5, lt=1.0)
+    seed: int = 0
+
+
+class Execution(_Strict):
+    batch_size: int = Field(
+        default=64,
+        ge=1,
+        le=4096,
+        description="Rows per forward pass. Part of the spec because batch shape can change "
+        "floating-point results in the last digits.",
+    )
+
+
+class PredictionSettings(_Strict):
+    """An optional, explicitly configured diagnostic; it never changes an intervention."""
+
+    method: Literal["final_norm_logit_lens"]
+    prompt_index: int = Field(ge=0)
+    which: Literal["clean", "corrupt"]
+    position: Annotated[LastPosition | IndexPosition, Field(discriminator="kind")]
+    top_k: int = Field(ge=1, le=20)
+
+
+class Spec(_Strict):
+    logogram_spec: Literal[1] = SPEC_VERSION
+    name: str = Field(min_length=1, max_length=NAME_MAX)
+    notes: str = ""
+    model: ModelRef
+    dataset: DatasetRef
+    tokenization: Tokenization = Field(default_factory=Tokenization)
+    experiment: Experiment
+    scope: Scope
+    metric: Metric = Field(default_factory=Metric)
+    statistics: Statistics = Field(default_factory=Statistics)
+    execution: Execution = Field(default_factory=Execution)
+    predictions: PredictionSettings | None = None
+    # The SAE whose features sae_feature sites and the features scope refer to.
+    sae: SAERef | None = None
+
+    @model_validator(mode="after")
+    def _features_have_an_sae(self) -> Spec:
+        uses = isinstance(self.scope, FeaturesScope) or (
+            isinstance(self.scope, SitesScope)
+            and any(s.kind == "sae_feature" for s in self.scope.sites)
+        )
+        if uses and self.sae is None:
+            raise ValueError("SAE features need the spec's sae: the SAE they belong to")
+        return self
+
+    def to_json(self) -> str:
+        return json.dumps(self.model_dump(mode="json"), indent=2) + "\n"
+
+    @classmethod
+    def from_path(cls, path: str | Path) -> Spec:
+        return cls.model_validate_json(Path(path).read_text(encoding="utf-8"))
+
+
+def describe_intervention(
+    exp: ActivationPatching
+    | Ablation
+    | DirectLogitAttribution
+    | AttributionPatching
+    | Steering
+    | PathPatching,
+) -> str:
+    """The intervention in a few words, for example 'Resample-ablate (10 corrupt donors, seed 0)'."""
+    if isinstance(exp, ActivationPatching):
+        return (
+            "Patch clean → corrupt"
+            if exp.direction == "clean_to_corrupt"
+            else "Patch corrupt → clean"
+        )
+    if isinstance(exp, DirectLogitAttribution):
+        return f"Direct logit attribution ({exp.prompts} prompts)"
+    if isinstance(exp, AttributionPatching):
+        arrow = "clean → corrupt" if exp.direction == "clean_to_corrupt" else "corrupt → clean"
+        return f"Attribution patching, estimated ({arrow})"
+    if isinstance(exp, PathPatching):
+        arrow = "clean → corrupt" if exp.direction == "clean_to_corrupt" else "corrupt → clean"
+        names = ", ".join(receiver_text(r) for r in exp.receivers[:4])
+        more = f" and {len(exp.receivers) - 4} more" if len(exp.receivers) > 4 else ""
+        frozen = "heads and MLPs" if exp.freeze_mlps else "other heads"
+        return f"Path patching {arrow} into {names}{more} ({frozen} frozen)"
+    if isinstance(exp, Steering):
+        toward = "corrupt" if exp.apply_to == "clean" else "clean"
+        strengths = ", ".join(strength_text(c) for c in exp.coefficients)
+        control = ", random control" if exp.control else ""
+        return (
+            f"Steer {exp.apply_to} prompts toward {toward} ({strengths}; "
+            f"{exp.train_fraction:.0%} of pairs train the direction, seed {exp.seed}{control})"
+        )
+    b = exp.baseline
+    if isinstance(b, ZeroBaseline):
+        return "Zero-ablate"
+    if isinstance(b, MeanBaseline):
+        return f"Mean-ablate ({b.reference} mean)"
+    return f"Resample-ablate ({b.donors} {b.pool} donors, seed {b.seed})"
+
+
+def describe_experiment(spec: Spec) -> str:
+    """A one-line plain-language description, used in history and summaries."""
+    what = describe_intervention(spec.experiment)
+    scope = spec.scope
+    if isinstance(scope, HeadsScope):
+        where = f"every head, {describe_position(scope.position)}"
+    elif isinstance(scope, LayerPositionScope):
+        unit = "every position" if scope.positions == "each" else "each labelled position"
+        where = f"{scope.site} at every layer and {unit}"
+    elif isinstance(scope, LayerComponentsScope):
+        where = f"{' and '.join(scope.components)} per layer, {describe_position(scope.position)}"
+    elif isinstance(scope, FeaturesScope):
+        where = (
+            f"every SAE feature, {describe_position(scope.position)}, keeping the top {scope.top}"
+        )
+    else:
+        where = f"{len(scope.sites)} chosen site{'s' if len(scope.sites) != 1 else ''}"
+    return f"{what} · {where}"

logogram/stats.py

@@ -0,0 +1,171 @@
+"""Per-site statistics: means, standard deviations, bootstrap confidence intervals, sign counts.
+
+Everything is vectorized over sites. Bootstrap resamples are drawn once per run from the raw
+PCG64 bit stream (stable across NumPy versions) and shared by every site, so differences between
+sites are not blurred by different resamples.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+Normalization = str  # "dataset_gap" | "prompt_gap"
+
+
+def resample_counts(n: int, n_boot: int, seed: int) -> np.ndarray:
+    """A [n_boot, n] matrix: how often each prompt appears in each bootstrap resample."""
+    raw = np.random.PCG64(seed).random_raw(n_boot * n)
+    idx = (raw % np.uint64(n)).astype(np.int64).reshape(n_boot, n)
+    counts = np.zeros((n_boot, n), dtype=np.float64)
+    rows = np.repeat(np.arange(n_boot), n)
+    np.add.at(counts, (rows, idx.ravel()), 1.0)
+    return counts
+
+
+def _quantiles(boot: np.ndarray, ci: float) -> tuple[np.ndarray, np.ndarray]:
+    alpha = (1.0 - ci) / 2.0
+    lo, hi = np.quantile(boot, [alpha, 1.0 - alpha], axis=1)
+    return lo, hi
+
+
+def _sd(values: np.ndarray) -> np.ndarray:
+    n = values.shape[1]
+    if n < 2:
+        return np.zeros(values.shape[0])
+    return values.std(axis=1, ddof=1)
+
+
+@dataclass
+class SiteStats:
+    """Statistics for S sites over n prompts. Arrays have length S unless noted."""
+
+    n: int
+    effect: np.ndarray  # [S, n] per-prompt normalized effect
+    effect_mean: np.ndarray
+    effect_sd: np.ndarray
+    effect_lo: np.ndarray
+    effect_hi: np.ndarray
+    delta: np.ndarray  # [S, n] patched - receiver logit difference
+    delta_mean: np.ndarray
+    delta_sd: np.ndarray
+    delta_lo: np.ndarray
+    delta_hi: np.ndarray
+    patched_mean: np.ndarray
+    prob_mean: np.ndarray
+    prob_delta_mean: np.ndarray
+    sign_flips: np.ndarray  # prompts whose logit difference changed sign
+    opposite_sign: np.ndarray  # prompts whose effect has the opposite sign to the mean
+    denominator: float  # mean gap (dataset_gap) or nan (prompt_gap)
+
+
+def normalize(
+    delta: np.ndarray, gap: np.ndarray, normalization: Normalization
+) -> tuple[np.ndarray, float]:
+    if normalization == "dataset_gap":
+        denom = float(gap.mean())
+        return delta / denom, denom
+    if normalization == "prompt_gap":
+        return delta / gap[None, :], float("nan")
+    raise ValueError(f"unknown normalization {normalization!r}")
+
+
+def compute_site_stats(
+    patched_ld: np.ndarray,
+    patched_prob: np.ndarray,
+    receiver_ld: np.ndarray,
+    source_ld: np.ndarray,
+    receiver_prob: np.ndarray,
+    normalization: Normalization,
+    counts: np.ndarray,
+    ci: float,
+    delta: np.ndarray | None = None,
+    gap: np.ndarray | None = None,
+) -> SiteStats:
+    """``patched_*`` are [S, n]; ``receiver_*``/``source_ld`` are [n]; ``counts`` is [B, n].
+
+    ``delta`` [S, n] and ``gap`` [n] replace patched - receiver and source - receiver for methods
+    whose per-prompt values aren't patched runs (estimates, and terms of a decomposition). Where
+    nothing ran patched, ``patched_*`` are NaN and count no sign flips.
+    """
+    patched_ld = np.asarray(patched_ld, dtype=np.float64)
+    patched_prob = np.asarray(patched_prob, dtype=np.float64)
+    receiver_ld = np.asarray(receiver_ld, dtype=np.float64)
+    source_ld = np.asarray(source_ld, dtype=np.float64)
+    n = patched_ld.shape[1]
+
+    delta = (
+        patched_ld - receiver_ld[None, :] if delta is None else np.asarray(delta, dtype=np.float64)
+    )
+    gap = source_ld - receiver_ld if gap is None else np.asarray(gap, dtype=np.float64)
+    effect, denom = normalize(delta, gap, normalization)
+
+    weights = counts.T / n  # [n, B]
+    delta_boot = delta @ weights  # [S, B]
+    if normalization == "dataset_gap":
+        gap_boot = gap @ weights  # [B]
+        effect_boot = delta_boot / gap_boot[None, :]
+    else:
+        effect_boot = effect @ weights
+    effect_lo, effect_hi = _quantiles(effect_boot, ci)
+    delta_lo, delta_hi = _quantiles(delta_boot, ci)
+    if n < 2:
+        # One prompt has no sampling distribution: report no interval rather than a false one.
+        nan = np.full(len(effect_lo), np.nan)
+        effect_lo, effect_hi, delta_lo, delta_hi = nan, nan.copy(), nan.copy(), nan.copy()
+
+    effect_mean = effect.mean(axis=1)
+    flips = ((receiver_ld[None, :] > 0) & (patched_ld < 0)) | (
+        (receiver_ld[None, :] < 0) & (patched_ld > 0)
+    )
+    mean_sign = np.sign(effect_mean)[:, None]
+    opposite = (np.sign(effect) == -mean_sign) & (mean_sign != 0)
+
+    return SiteStats(
+        n=n,
+        effect=effect,
+        effect_mean=effect_mean,
+        effect_sd=_sd(effect),
+        effect_lo=effect_lo,
+        effect_hi=effect_hi,
+        delta=delta,
+        delta_mean=delta.mean(axis=1),
+        delta_sd=_sd(delta),
+        delta_lo=delta_lo,
+        delta_hi=delta_hi,
+        patched_mean=patched_ld.mean(axis=1),
+        prob_mean=patched_prob.mean(axis=1),
+        prob_delta_mean=(patched_prob - np.asarray(receiver_prob)[None, :]).mean(axis=1),
+        sign_flips=flips.sum(axis=1),
+        opposite_sign=opposite.sum(axis=1),
+        denominator=denom,
+    )
+
+
+def rankdata(values: np.ndarray) -> np.ndarray:
+    """Ranks starting at 1, ties get the average rank (as in scipy.stats.rankdata)."""
+    values = np.asarray(values, dtype=np.float64)
+    order = np.argsort(values, kind="mergesort")
+    ranks = np.empty(len(values), dtype=np.float64)
+    sorted_vals = values[order]
+    i = 0
+    while i < len(values):
+        j = i
+        while j + 1 < len(values) and sorted_vals[j + 1] == sorted_vals[i]:
+            j += 1
+        ranks[order[i : j + 1]] = (i + j) / 2.0 + 1.0
+        i = j + 1
+    return ranks
+
+
+def spearman(a: np.ndarray, b: np.ndarray) -> float:
+    if len(a) < 2:
+        return float("nan")
+    ra, rb = rankdata(a), rankdata(b)
+    ra -= ra.mean()
+    rb -= rb.mean()
+    denom = np.sqrt((ra * ra).sum() * (rb * rb).sum())
+    if denom == 0:
+        return float("nan")
+    return float((ra * rb).sum() / denom)