EEG similarity search with embeddings

Each 30 s window becomes one L2-normalised 768-d vector, so a matrix product is a cosine-similarity search. This recipe embeds 20 nights of public sleep EEG, queries one window, returns its nearest windows from other people, and measures how often those neighbours carry the same scored sleep stage.

Research demonstration, not clinical validation. Sleep stages are used here only as an outside label to check what the neighbours have in common.

Data and protocol

Dataset: Sleep-EDF Expanded, sleep-cassette study, first night of subjects 0 to 19 (SC4001 to SC4191, PSG and hypnogram files). Check the downloads against PhysioNet’s SHA256SUMS.txt: a truncated EDF still opens.
Signal: the two EEG channels, Fpz-Cz and Pz-Oz, at 100 Hz as stored. Each night is cropped to 30 minutes before the first and after the last sleep epoch.
Windows: the default stride of 30 s, so windows line up with the hypnogram’s 30 s epochs. 20,650 windows were embedded; 20,626 have a stage (W, N1, N2, N3, REM; stages 3 and 4 merged into N3). The 24 windows scored as movement or unknown are left out of the search.
Search: cosine similarity over all labelled windows, with every window of the query’s own night masked, so neighbours always come from another subject. K = 10.
Chance: for a query of stage s, the fraction of other-subject windows that have stage s. This is the same-stage rate of 10 neighbours drawn at random.
Run: 10 November 2026, neuroencoder 2.1.0, hosted API, Python 3.11.

1. Embed the nights

embed_nights.py

import glob, os, time
import numpy as np, mne
from neuroencoder import embeddings

DATA, CACHE = "sleep-cassette", "cache"   # Sleep-EDF Expanded sleep-cassette files; one output file per night
NIGHTS = [f"SC4{s:02d}1" for s in range(20)]  # first night of subjects 0-19
HYPNO = {"Sleep stage W": "W", "Sleep stage 1": "N1", "Sleep stage 2": "N2", "Sleep stage 3": "N3",
         "Sleep stage 4": "N3", "Sleep stage R": "REM"}       # Movement time and "?" get no label
mne.set_log_level("ERROR"); os.makedirs(CACHE, exist_ok=True)

for night in NIGHTS:
    raw = mne.io.read_raw_edf(glob.glob(f"{DATA}/{night}E0-PSG.edf")[0], preload=True)
    n = int(raw.n_times / raw.info["sfreq"] // 30)
    y = np.full(n, "", dtype="<U3")                           # one label per 30 s epoch from recording start
    for a in mne.read_annotations(glob.glob(f"{DATA}/{night}*-Hypnogram.edf")[0]):
        y[int(a["onset"] // 30):int((a["onset"] + a["duration"]) // 30)] = HYPNO.get(a["description"], "")
    sleep = np.flatnonzero(np.isin(y, ["N1", "N2", "N3", "REM"]))
    a, b = max(sleep[0] - 60, 0), min(sleep[-1] + 61, n)      # keep 30 min before first and after last sleep epoch
    raw.crop(a * 30, b * 30, include_tmax=False).pick(["EEG Fpz-Cz", "EEG Pz-Oz"])
    t = time.time()
    r = embeddings(raw)                                       # hosted API, one window per 30 s epoch
    np.savez(f"{CACHE}/{night}.npz", emb=r.to_numpy(), y=y[a:b], t=np.array(r.t_start))
    print(f"{night}: {len(y[a:b])} windows, {time.time() - t:.1f} s, billed {r.usage.seconds_of_eeg:.0f} s", flush=True)

Needs pip install neuroencoder mne and a key (see Quickstart). The 20 nights took 1,201 s of wall time in the embedding calls (37 to 83 s per night) and billed 619,500 s, about 172 h of EEG. That is more than one trial key’s 100 h per month; with a trial key, embed 10 nights or fewer.

2. Search

similarity_search.py

import glob, json
import numpy as np

CACHE, K = "cache", 10                     # written by embed_nights.py, one file per night
STAGES = ["W", "N1", "N2", "N3", "REM"]
E, Y, S, T = [], [], [], []
for s, f in enumerate(sorted(glob.glob(f"{CACHE}/SC4*.npz"))):
    d = np.load(f); keep = d["y"] != ""
    E.append(d["emb"][keep]); Y.append(d["y"][keep]); T.append(d["t"][keep]); S.append(np.full(keep.sum(), s))
E, Y, S, T = map(np.concatenate, (E, Y, S, T))  # rows of E are L2-normalised: E @ q is cosine similarity

def search(q, subject, k=K):
    sim = E @ q
    sim[S == subject] = -np.inf             # only windows from other people
    top = np.argpartition(-sim, k)[:k]
    return top[np.argsort(-sim[top])], sim

i = np.flatnonzero((S == 0) & (Y == "REM"))[0]  # query: first REM window of night 0
top, sim = search(E[i], S[i])
print(f"query: night {S[i]}, t={T[i]:.0f} s, stage {Y[i]}")
for j in top:
    print(f"  night {S[j]:2d}  t={T[j]:6.0f} s  stage {Y[j]:3s}  cosine {sim[j]:.3f}")

same, chance = np.zeros(len(E)), np.zeros(len(E))
for i in range(len(E)):                     # every window as a query
    top, _ = search(E[i], S[i])
    same[i] = np.mean(Y[top] == Y[i])
    other = S != S[i]
    chance[i] = np.mean(Y[other] == Y[i])   # same-stage rate of a random other-subject window
res = {"windows": len(E), "nights": int(S.max() + 1), "k": K,
       "overall": [same.mean(), chance.mean()],
       "per_stage": {s: [same[Y == s].mean(), chance[Y == s].mean(), int((Y == s).sum())] for s in STAGES}}
res["stage_mean"] = [np.mean([v[0] for v in res["per_stage"].values()]), np.mean([v[1] for v in res["per_stage"].values()])]
json.dump(res, open("similarity_search.json", "w"), indent=1)
print(f"{K} nearest other-subject windows with the same stage: {res['overall'][0]:.3f} (chance {res['overall'][1]:.3f})")
for s, (a, c, n) in res["per_stage"].items():
    print(f"  {s:3s} n={n:5d}  {a:.3f}  chance {c:.3f}")

Plain NumPy, no index library. Using all 20,626 windows as queries took 101 s on a 16-core machine; a single query is one matrix-vector product. t is seconds from the start of the cropped night, and night n is the n-th file in sorted order (SC4001 is night 0).

Output

query: night 0, t=7140 s, stage REM
  night 17  t= 25230 s  stage REM  cosine 0.972
  night 17  t= 10110 s  stage W    cosine 0.971
  night 18  t= 10950 s  stage REM  cosine 0.970
  night  1  t= 28230 s  stage N2   cosine 0.968
  night  2  t= 14760 s  stage REM  cosine 0.966
  night 17  t= 13470 s  stage REM  cosine 0.966
  night 18  t= 11070 s  stage REM  cosine 0.965
  night 17  t= 21570 s  stage REM  cosine 0.963
  night 17  t= 22410 s  stage N1   cosine 0.963
  night 17  t= 20640 s  stage REM  cosine 0.962
10 nearest other-subject windows with the same stage: 0.722 (chance 0.283)
  W   n= 3449  0.808  chance 0.167
  N1  n= 1240  0.240  chance 0.059
  N2  n= 9200  0.732  chance 0.445
  N3  n= 2981  0.778  chance 0.142
  REM n= 3756  0.733  chance 0.182

Grouped bars: same-stage fraction of the 10 nearest other-subject windows versus chance. W 0.81 vs 0.17, N1 0.24 vs 0.06, N2 0.73 vs 0.44, N3 0.78 vs 0.14, REM 0.73 vs 0.18, all 0.72 vs 0.28. — Fraction of the 10 nearest other-subject windows that share the query's stage, against chance, per stage and over all 20,626 windows.

What this shows and what it does not

Across subjects, the nearest windows share the query’s stage far more often than chance for every stage: 72.2% of neighbours overall against 28.3% by chance.
N1 is the weak stage: 24.0% of neighbours are N1, against 5.9% by chance, so most neighbours of an N1 window carry another stage.
The overall figure is weighted by stage frequency; N2 alone is 45% of windows. The per-stage rows are the fairer read.
In the example, 6 of the 10 neighbours come from night 17. Neighbours can cluster on a few nights, so a search result is not a sample of the population.
Same-stage rate is one proxy. It does not show that neighbours match on anything clinically relevant, it is not a sleep stager (each window is searched on its own, with no context from neighbouring epochs), and it was measured on the Fpz-Cz and Pz-Oz channels of 20 recordings from one public dataset only.

For an interactive view of the same kind of space, see Visualization.

Get started

Workflows

Reference

EEG similarity search with embeddings

Data and protocol

1. Embed the nights

2. Search

Output

What this shows and what it does not

​Data and protocol

​1. Embed the nights

​2. Search

​Output

​What this shows and what it does not

Data and protocol

1. Embed the nights

2. Search

Output

What this shows and what it does not