import os
os.environ.update(OMP_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1", MKL_NUM_THREADS="1")
import io
import time
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import requests
from scipy import ndimage, optimize, stats

import bwm

pd.set_option("display.width", 220)
pd.set_option("display.max_columns", 40)
pd.set_option("display.max_rows", 300)
warnings.filterwarnings("ignore", category=RuntimeWarning)

REPO = "https://raw.githubusercontent.com/int-brain-lab/paper-brain-wide-map/118fc36cb3602934466ad2c6087c2b3b441f9f1f"
region_info = pd.read_csv(io.StringIO(requests.get(f"{REPO}/brainwidemap/meta/region_info.csv", timeout=60).text))
inclusion_info = pd.read_csv(io.StringIO(requests.get(f"{REPO}/brainwidemap/meta/inclusion_info.csv", timeout=60).text), index_col=0)

started = time.time()
assets = bwm.list_assets()
processed = assets[assets.kind == "processed"].sort_values("path").reset_index(drop=True)
scans = bwm.run_parallel(bwm.scan_session, list(processed.asset_id))
scan_seconds = time.time() - started
print(f"dandiset {bwm.DANDISET_ID} version {bwm.DANDISET_VERSION}: {len(assets)} assets; "
      f"scanned {len(scans)} processed files in {scan_seconds:.0f} s")

dandiset 000409 version 0.260309.1324: 2048 assets; scanned 459 processed files in 131 s

# One row per session, and the trials of all sessions in the paper's terms.
rows, trial_tables = [], []
for s in scans:
    i, u = s["info"], s["units"]
    t = bwm.task_table(s["trials"])
    mask, n_kept = bwm.analysis_mask(t)
    t["analysed"] = mask
    t["eid"] = i["session_id"]
    t["subject"] = i["subject_id"]
    t["trial"] = np.arange(len(t))
    t["from_end"] = len(t) - t.trial
    trial_tables.append(t)
    rows.append({"eid": i["session_id"], "subject": i["subject_id"], "lab": i["lab"], "sex": i["sex"],
                 "dob": pd.Timestamp(i["date_of_birth"]), "start": pd.Timestamp(i["session_start"]),
                 "weight_g": float(i["weight"].split()[0]) * 1000 if i["weight"] else np.nan,
                 "protocol": i["protocol"], "revision": i["revision"], "n_trials": len(t),
                 "n_kept": n_kept, "n_analysed": int(mask.sum()), "n_units": len(u),
                 "n_good": int((u.ibl_quality_score == 1).sum()), "n_probes": u.probe_name.nunique(),
                 "modules": ",".join(i["processing"]),
                 "passive": ",".join(s["passive"]) if s["passive"] else "none"})
sessions = pd.DataFrame(rows)
trials = pd.concat(trial_tables, ignore_index=True)
sessions["start"] = pd.to_datetime(sessions.start, utc=True)
sessions["dob"] = pd.to_datetime(sessions.dob, utc=True)
sessions["age_weeks"] = (sessions.start - sessions.dob).dt.days / 7
print(sessions.drop(columns=["modules", "passive"]).head().to_string())

                                    eid      subject       lab sex                       dob                            start  weight_g                             protocol    revision  n_trials  n_kept  n_analysed  n_units  n_good  n_probes  age_weeks
0  3e7ae7c0-fe8b-487c-9354-036236fa1010  CSH_ZAD_001  zadorlab   F 2019-04-16 04:00:00+00:00 2020-01-16 20:53:21.500926+00:00     22.37  _iblrig_tasks_ephysChoiceWorld6.2.5  2025-05-06       854     854         406      519      50         1  39.285714
1  03cf52f6-fba6-4743-a42e-dd1ac3072343  CSH_ZAD_011  zadorlab   F 2019-09-17 04:00:00+00:00        2020-03-22 16:56:40+00:00     20.27  _iblrig_tasks_ephysChoiceWorld6.4.0  2025-05-06       415     415         305     1341     119         2  26.714286
2  5b44c40f-80f4-44fb-abfb-c7f19e27a6ca  CSH_ZAD_011  zadorlab   F 2019-09-17 04:00:00+00:00        2020-03-21 18:17:23+00:00     20.27  _iblrig_tasks_ephysChoiceWorld6.4.0  2025-05-06       591     591         437     1927     134         2  26.571429
3  da188f2c-553c-4e04-879b-c9ea2d1b9a93  CSH_ZAD_011  zadorlab   F 2019-09-17 04:00:00+00:00        2020-03-23 16:58:00+00:00     20.27  _iblrig_tasks_ephysChoiceWorld6.4.0  2025-05-06       927     927         676     2019     205         2  26.857143
4  49e0ab27-827a-4c91-bcaa-97eea27a1b8d  CSH_ZAD_019  zadorlab   F 2019-11-12 05:00:00+00:00        2020-08-18 11:41:02+00:00     23.00  _iblrig_tasks_ephysChoiceWorld6.4.1  2025-05-06       884     884         576     1689     141         2  40.000000

print("Assets by kind:", assets.kind.value_counts().to_dict())
per_session = assets.groupby("eid").kind.apply(lambda k: ",".join(sorted(k)))
print("Sessions with a processed and a raw file:", int((per_session.str.contains("processed") & per_session.str.contains("raw")).sum()),
      "of", per_session.size)
print("Sessions:", len(sessions), "; subjects:", sessions.subject.nunique(), "; labs:", sessions.lab.nunique())
print("Sessions per lab:", sessions.lab.value_counts().to_dict())
print("Sessions per subject:", sessions.subject.value_counts().value_counts().sort_index().to_dict(), "(sessions: subjects)")
print("Spike-sorting revision in the files:", sessions.revision.value_counts().to_dict())
print("Task protocols:", sessions.protocol.str.replace("_iblrig_tasks_", "").value_counts().to_dict())
print("Processing modules:", sessions.modules.value_counts().to_dict())
print("Passive protocol phases:", sessions.passive.value_counts().to_dict())
print("Groups under stimulus/ (presentation, templates) in the processed files:", pd.Series([s["stimulus"] for s in scans]).value_counts().to_dict())
epochs = pd.concat([s["epochs"] for s in scans])
print("Epoch types:", epochs.protocol_type.value_counts().to_dict())

Assets by kind: {'raw': 459, 'processed': 459, 'VideoLeftCamera': 437, 'VideoRightCamera': 433, 'VideoBodyCamera': 260}
Sessions with a processed and a raw file: 459 of 459
Sessions: 459 ; subjects: 139 ; labs: 12
Sessions per lab: {'hausserlab': 62, 'danlab': 46, 'mainenlab': 44, 'cortexlab': 42, 'churchlandlab_ucla': 41, 'angelakilab': 41, 'zadorlab': 37, 'wittenlab': 37, 'churchlandlab': 36, 'steinmetzlab': 31, 'mrsicflogellab': 25, 'hoferlab': 17}
Sessions per subject: {1: 32, 2: 25, 3: 29, 4: 20, 5: 16, 6: 4, 7: 6, 8: 3, 9: 3, 13: 1} (sessions: subjects)
Spike-sorting revision in the files: {'2025-05-06': 459}
Task protocols: {'ephysChoiceWorld6.4.2': 182, 'ephysChoiceWorld': 64, 'ephysChoiceWorld6.2.5': 58, 'ephysChoiceWorld6.4.0': 50, 'ephysChoiceWorld6.5.3': 29, 'ephysChoiceWorld6.6.2': 24, 'ephysChoiceWorld6.4.1': 20, 'ephysChoiceWorld6.1.3': 15, 'ephysChoiceWorld6.4.3': 6, 'ephysChoiceWorld6.6.4': 3, 'ephysChoiceWorld6.6.1': 3, 'ephysChoiceWorld6.5.0': 3, 'ephysChoiceWorld6.2.3': 2}
Processing modules: {'lick_times,motion_energy,passive_protocol,pose_estimation,pupil,wheel': 378, 'lick_times,passive_protocol,wheel': 32, 'lick_times,motion_energy,pose_estimation,pupil,wheel': 24, 'passive_protocol,wheel': 13, 'lick_times,motion_energy,passive_protocol,pose_estimation,wheel': 5, 'lick_times,wheel': 4, 'wheel': 3}
Passive protocol phases: {'spontaneousActivity,taskReplay': 428, 'none': 31}
Groups under stimulus/ (presentation, templates) in the processed files: {'empty': 459}
Epoch types: {'task': 428, 'passive': 428, 'none': 31}

subj = sessions.groupby("subject").agg(sex=("sex", "first"), n_sex=("sex", "nunique"), dob=("dob", "nunique"),
                                       w_min=("weight_g", "min"), w_max=("weight_g", "max"), n=("eid", "size"))
print("Sex per mouse:", subj.sex.value_counts().to_dict(), "; mice with more than one sex or birth date:",
      int((subj.n_sex > 1).sum()), int((subj.dob > 1).sum()))
print("Weight in the files (g): range", sessions.weight_g.min(), "to", sessions.weight_g.max(),
      "; sessions without weight:", int(sessions.weight_g.isna().sum()))
print("Mice whose weight differs between their sessions:", int((subj.w_max > subj.w_min).sum()), "of", int((subj.n > 1).sum()),
      "with more than one session")
print("Age at recording (weeks): range", round(sessions.age_weeks.min(), 1), "to", round(sessions.age_weeks.max(), 1))

Sex per mouse: {'M': 94, 'F': 45} ; mice with more than one sex or birth date: 0 0
Weight in the files (g): range 16.1 to 36.2 ; sessions without weight: 0
Mice whose weight differs between their sessions: 0 of 107 with more than one session
Age at recording (weeks): range 13.1 to 65.9

print("Trial columns:", list(scans[0]["trials"].columns))
print("Trials:", len(trials), "; per session", sessions.n_trials.min(), "to", sessions.n_trials.max())
print("Contrasts (%):", (trials.contrast * 100).value_counts().sort_index().to_dict())
c = trials[trials.p_left != 0.5].contrast.value_counts()
print("Ratio of contrasts 1, 0.25, 0.125, 0.0625, 0 in biased blocks:", (c.loc[[1, 0.25, 0.125, 0.0625, 0]] / c.loc[0]).round(2).tolist())
print("p_left values:", trials.p_left.value_counts().to_dict())
first = trials.groupby("eid").apply(lambda g: int((g.block == 0).sum()), include_groups=False)
print("Length of the first (unbiased) block:", first.value_counts().to_dict())
biased = trials[trials.p_left != 0.5]
match = biased.groupby("eid").apply(lambda g: np.mean((g.side < 0) == (g.p_left == 0.8)), include_groups=False)
sessions["block_match"] = sessions.eid.map(match)
sessions["block_ok"] = sessions.block_match >= 0.7
BLOCK_OK = set(sessions.eid[sessions.block_ok])
print("Fraction of biased-block trials whose stimulus is on the block's likely side, per session:")
print("   sessions at or above 0.7:", len(BLOCK_OK), "range", round(match[match >= 0.7].min(), 3), "to", round(match[match >= 0.7].max(), 3))
print("   sessions below 0.7:", int((~sessions.block_ok).sum()), "range", round(match[match < 0.7].min(), 3), "to", round(match[match < 0.7].max(), 3))
bad_s = sessions[~sessions.block_ok]
print("   their protocols:", bad_s.protocol.value_counts().to_dict(), "; labs:", bad_s.lab.value_counts().to_dict(),
      "; dates", bad_s.start.min().date(), "to", bad_s.start.max().date())
print("   all sessions with that protocol name:", int((sessions.protocol == "_iblrig_tasks_ephysChoiceWorld").sum()))
blen = trials[trials.block > 0].groupby(["eid", "block"]).size()
last = trials.groupby("eid").block.max()
for name, ok in (("block_ok", True), ("other", False)):
    b = blen[[(e in BLOCK_OK) == ok and bl < last[e] for e, bl in blen.index]]
    print(f"Biased block lengths in {name} sessions, last block left out: min {b.min()}, max {b.max()}, mean {b.mean():.1f}, blocks {len(b)}")
print("Stimulus side x choice x rewarded:")
print(pd.crosstab([trials.side, trials.choice], trials.correct).to_string())
print("Reward volume on rewarded trials (uL):", trials[trials.correct].reward_ul.round(2).value_counts().head(6).to_dict())
q = pd.concat([s["trials"].quiescence_period for s in scans])
print("Quiescence period (s): range", round(q.min(), 3), "to", round(q.max(), 3))
print("Missing values:", {k: int(trials[k].isna().sum()) for k in ["stim_on", "go_cue", "first_move", "feedback_time"]})
print("Trials without a response:", int((trials.choice == 0).sum()))
print("Stimulus onset minus go cue (ms), median and 1st/99th percentiles:",
      np.nanpercentile((trials.stim_on - trials.go_cue) * 1000, [50, 1, 99]).round(1).tolist())
print("First movement before stimulus onset:", round(float(np.mean(trials.rt < 0)) * 100, 2), "% of trials; before the go cue:",
      round(float(np.mean(trials.first_move < trials.go_cue)) * 100, 2), "%")
print("Trials analysed (complete events, first movement 0.08 to 2 s):", int(trials.analysed.sum()), "of", len(trials))
trunc = sessions[sessions.n_kept < sessions.n_trials]
print("Sessions truncated to pass the 90% criterion on easy (contrast >= 0.5) trials:", len(trunc))
for r in trunc.itertuples():
    tt = trials[trials.eid == r.eid].reset_index(drop=True)
    print(f"   {r.eid[:8]}: {r.n_trials} trials, kept {r.n_kept} (floor 400: {r.n_kept == 400}); easy-trial performance "
          f"{bwm.perf_easy(tt):.3f} over all trials, {bwm.perf_easy(tt.iloc[:r.n_kept]):.3f} over those kept; analysed {r.n_analysed}")
print("Sessions that do not pass even at 400 trials, and are excluded as the code excludes them:", int((sessions.n_analysed == 0).sum()))

Trial columns: ['auditory_cue_time', 'block_index', 'block_type', 'choice_registration_time', 'feedback_time', 'gabor_stimulus_contrast', 'gabor_stimulus_offset_time', 'gabor_stimulus_onset_time', 'gabor_stimulus_side', 'is_mouse_rewarded', 'mouse_wheel_choice', 'probability_left', 'quiescence_period', 'reward_volume_uL', 'start_time', 'stop_time', 'wheel_movement_onset_time']
Trials: 296090 ; per session 401 to 1525
Contrasts (%): {0.0: 34713, 6.25: 64999, 12.5: 64904, 25.0: 65283, 100.0: 66191}
Ratio of contrasts 1, 0.25, 0.125, 0.0625, 0 in biased blocks: [1.89, 1.86, 1.85, 1.85, 1.0]
p_left values: {0.8: 130533, 0.2: 124247, 0.5: 41310}

Length of the first (unbiased) block:

 {90: 459}
Fraction of biased-block trials whose stimulus is on the block's likely side, per session:
   sessions at or above 0.7: 395 range 0.759 to 0.828
   sessions below 0.7: 64 range 0.475 to 0.57
   their protocols: {'_iblrig_tasks_ephysChoiceWorld': 64} ; labs: {'hausserlab': 48, 'churchlandlab_ucla': 13, 'mainenlab': 2, 'cortexlab': 1} ; dates 2022-01-31 to 2023-10-19
   all sessions with that protocol name: 64

Biased block lengths in block_ok sessions, last block left out: min 20, max 99, mean 50.2, blocks 4158
Biased block lengths in other sessions, last block left out: min 1, max 85, mean 1.1, blocks 31041
Stimulus side x choice x rewarded:
correct      False   True 
side choice               
-1   -1          0  123603
      0        626       0
      1      27097       0
 1   -1      25213       0
      0        526       0
      1          0  119025
Reward volume on rewarded trials (uL): {1.5: 242628}
Quiescence period (s): range 0.4 to 0.7
Missing values: {'stim_on': 94, 'go_cue': 954, 'first_move': 3009, 'feedback_time': 199}
Trials without a response: 1152
Stimulus onset minus go cue (ms), median and 1st/99th percentiles: [-1.0, -3.3, 21.9]
First movement before stimulus onset: 11.93 % of trials; before the go cue: 11.97 %
Trials analysed (complete events, first movement 0.08 to 2 s): 196962 of 296090
Sessions truncated to pass the 90% criterion on easy (contrast >= 0.5) trials: 2
   d71e565d: 600 trials, kept 457 (floor 400: False); easy-trial performance 0.872 over all trials, 0.902 over those kept; analysed 296
   283ecb4c: 487 trials, kept 460 (floor 400: False); easy-trial performance 0.882 over all trials, 0.908 over those kept; analysed 319
Sessions that do not pass even at 400 trials, and are excluded as the code excludes them: 0

units = pd.concat([s["units"].assign(eid=s["info"]["session_id"]) for s in scans], ignore_index=True)
passes = ((units.median_spike_amplitude_uV > 50).astype(int) + (units.noise_cutoff < 20).astype(int)
          + (units.sliding_rp_violation == 1).astype(int)) / 3
print("Units:", len(units), "; quality score values:", units.ibl_quality_score.round(2).value_counts().sort_index().to_dict())
print("Units whose score equals the fraction of the three metrics they pass, noise cutoff below 20 (Methods):",
      int(np.isclose(passes, units.ibl_quality_score, atol=0.01).sum()), "of", len(units))
passes5 = ((units.median_spike_amplitude_uV > 50).astype(int) + (units.noise_cutoff < 5).astype(int)
           + (units.sliding_rp_violation == 1).astype(int)) / 3
print("The same with a noise cutoff below 5:", int(np.isclose(passes5, units.ibl_quality_score, atol=0.01).sum()))
print("Noise cutoff: units with a value", int(units.noise_cutoff.notna().sum()), "; range", round(units.noise_cutoff.min(), 2), "to",
      round(units.noise_cutoff.max(), 1), "; well-isolated units with a cutoff of 5 or more:",
      int(((units.ibl_quality_score == 1) & (units.noise_cutoff >= 5)).sum()))
print("Units whose spike_count equals the length of their spike train:", int((units.spike_count == units.n_spikes_index).sum()))
print("Units whose probe matches the probe of their peak channel:", int((units.probe_name == units.region_probe).sum()))
probe_models = pd.Series([v.split(" probe.")[0].replace("A ", "", 1) for s in scans for v in s["probes"].values()])
print("Probe devices by model:", probe_models.value_counts().to_dict())
print("Units per probe: mean", round(len(units) / sessions.n_probes.sum(), 1))
descs = pd.Series([s["info"]["units_description"] for s in scans])
print("Distinct units-table descriptions:", descs.nunique(), "; first sentence:", descs.iloc[0].split(". ")[0])
print("Descriptions that mention a 'label' column:", int(descs.str.contains("'label' column").sum()), "; columns of the units table include 'label':",
      "label" in scans[0]["units"].columns)

Units: 621733 ; quality score values: {0.0: 112878, 0.33: 255535, 0.67: 177612, 1.0: 75708}
Units whose score equals the fraction of the three metrics they pass, noise cutoff below 20 (Methods): 540399 of 621733
The same with a noise cutoff below 5: 621733
Noise cutoff: units with a value 618252 ; range -3.54 to 239469.7 ; well-isolated units with a cutoff of 5 or more: 0
Units whose spike_count equals the length of their spike train: 621733
Units whose probe matches the probe of their peak channel: 621733
Probe devices by model: {'PRB_1_4_0480_1': 403, 'PRB_1_4_0480_1_C': 183, 'NP1010': 109, 'NP2010': 4}
Units per probe: mean 889.5
Distinct units-table descriptions: 1 ; first sentence: Spike-sorted units from Neuropixels 1.0 probes
Descriptions that mention a 'label' column: 459 ; columns of the units table include 'label': False

# Recordings deposited twice: identical trial timings or identical unit spike counts.
fp_trials = trials.groupby("eid").apply(lambda g: (len(g), round(float(g.stim_on.iloc[0]), 4), round(float(g.feedback_time.iloc[-1]), 4)), include_groups=False)
dup_t = fp_trials[fp_trials.duplicated(keep=False)]
fp_units = units.groupby(["eid", "probe_name"]).spike_count.apply(lambda c: tuple(np.sort(c.values)[-20:]))
dup_u = fp_units[fp_units.duplicated(keep=False)]
print("Sessions sharing trial count, first stimulus time and last feedback time with another:", len(dup_t))
print("Probe insertions sharing their 20 largest unit spike counts with another:", len(dup_u))
same_day = sessions.groupby(["subject", sessions.start.dt.date]).size()
print("Mice with more than one session on one day:", int((same_day > 1).sum()))

Sessions sharing trial count, first stimulus time and last feedback time with another: 0
Probe insertions sharing their 20 largest unit spike counts with another: 0
Mice with more than one session on one day: 7

# Clocks of the wheel and lick streams against the trials table.
near, lick_ratio = [], []
for s in scans:
    t = bwm.task_table(s["trials"])
    fm = t.first_move.dropna().values
    wm = np.sort(s["wheel_moves"])
    if len(wm) and len(fm):
        j = np.clip(np.searchsorted(wm, fm), 1, len(wm) - 1)
        near.append(np.minimum(np.abs(wm[j] - fm), np.abs(wm[j - 1] - fm)))
    lk = np.sort(s["licks"])
    fb = t.feedback_time[t.correct].dropna().values
    if len(lk) > 100:
        after = (np.searchsorted(lk, fb + 1) - np.searchsorted(lk, fb)).sum()
        before = (np.searchsorted(lk, fb) - np.searchsorted(lk, fb - 1)).sum()
        lick_ratio.append(after / max(before, 1))
near = np.concatenate(near)
print("First movement time to the nearest wheel-movement onset (ms): median", round(np.median(near) * 1000, 2),
      "; within 50 ms:", round(np.mean(near < 0.05) * 100, 1), "% of", len(near), "trials")
print("Licks in the 1 s after reward over the 1 s before, per session: median", round(np.median(lick_ratio), 1),
      "; sessions above 2:", int(np.sum(np.array(lick_ratio) > 2)), "of", len(lick_ratio))

First movement time to the nearest wheel-movement onset (ms): median 0.33 ; within 50 ms: 96.8 % of 293081 trials
Licks in the 1 s after reward over the 1 s before, per session: median 7.1 ; sessions above 2: 407 of 442

print("Sessions:", len(sessions), "(paper 459); mice:", sessions.subject.nunique(), "(139); labs:", sessions.lab.nunique(), "(12)")

Sessions: 459 (paper 459); mice: 139 (139); labs: 12 (12)

n_probes = int(sessions.n_probes.sum())
print("Insertions:", n_probes, "(paper 699); units:", len(units), "(621,733); per probe:", round(len(units) / n_probes, 1), "(889)")

Insertions: 699 (paper 699); units: 621733 (621,733); per probe: 889.5 (889)

n_good = int((units.ibl_quality_score == 1).sum())
print("Well-isolated units:", n_good, "(paper 75,708); per probe:", round(n_good / n_probes, 1), "(108)")
print("Well-isolated by recomputing the three metrics: noise cutoff below 20 (Methods)", int((passes == 1).sum()),
      "; below 5", int((passes5 == 1).sum()))

Well-isolated units: 75708 (paper 75,708); per probe: 108.3 (108)
Well-isolated by recomputing the three metrics: noise cutoff below 20 (Methods) 95430 ; below 5 75708

canon, steps = bwm.canonical_units(scans)
paper_steps = inclusion_info.iloc[[0, 3, 4, 5, 6, 7]].reset_index(drop=True)
steps["paper_sessions"] = paper_steps.n_sessions.values
steps["paper_units"] = paper_steps.n_units.values
steps["paper_regions"] = paper_steps.n_regions.values
print(steps.to_string(index=False))
canon_regions = sorted(canon.region.unique())
paper_canonical = set(region_info[region_info.canonical].Beryl)
print("\nCanonical regions:", len(canon_regions), "; regions marked canonical in region_info.csv:", len(paper_canonical),
      "; of ours among them:", len(set(canon_regions) & paper_canonical))
print("Marked canonical in region_info.csv but with fewer than 20 units:", sorted(paper_canonical - set(canon_regions)))

            step  sessions  units  regions  paper_sessions  paper_units  paper_regions
       all units       459 621733      281             459       621733            281
   well-isolated       459  75708      268             459        75708            268
     grey matter       459  65336      266             459        65336            266
>= 5 per session       455  63357      245             455        63357            245
   >= 2 sessions       454  62990      210             454        62990            210
>= 20 per region       454  62857      201             454        62857            201

Canonical regions: 201 ; regions marked canonical in region_info.csv: 210 ; of ours among them: 201
Marked canonical in region_info.csv but with fewer than 20 units: ['AD', 'APr', 'CA2', 'ECU', 'FC', 'IA', 'IGL', 'PBG', 'PH']

grey = units[~units.region.isin(["void", "root"])]
print("Beryl regions with at least one unit, void and root left out:", grey.region.nunique(), "(paper 279)")
print("With at least one well-isolated unit:", grey[grey.ibl_quality_score == 1].region.nunique())

Beryl regions with at least one unit, void and root left out: 279 (paper 279)
With at least one well-isolated unit: 266

print(subj.sex.value_counts().to_dict(), "(paper 94 male, 45 female)")

{'M': 94, 'F': 45} (paper 94 male, 45 female)

s_sorted = sessions.sort_values("start")
readings = {"a: first session": s_sorted.groupby("subject").first(), "b: last session": s_sorted.groupby("subject").last(),
            "c: every session": s_sorted}
expected = {"age_weeks": [13, 178, 44.96, 27.0], "weight_g": [16.1, 35.7, 23.9, 23.84]}
for name, d in readings.items():
    for var, exp in expected.items():
        v = d[var].dropna()
        got = [v.min(), v.max(), v.mean(), v.median()]
        ok = [abs(g - e) <= 0.1 * e for g, e in zip(got, exp)]
        print(f"{name:17s} {var:9s} min {got[0]:6.1f} max {got[1]:6.1f} mean {got[2]:6.2f} median {got[3]:6.2f}  "
              f"within 10%: {ok}  (n = {len(v)})")
print("Paper's code (meta/mice_info.py) takes age_weeks and reference_weight from the IBL database, not from the recording day.")

a: first session  age_weeks min   13.1 max   63.1 mean  28.09 median  26.00  within 10%: [np.True_, np.False_, np.False_, np.True_]  (n = 139)
a: first session  weight_g  min   16.1 max   36.2 mean  24.19 median  24.00  within 10%: [np.True_, np.True_, np.True_, np.True_]  (n = 139)
b: last session   age_weeks min   13.3 max   65.9 mean  28.86 median  26.29  within 10%: [np.True_, np.False_, np.False_, np.True_]  (n = 139)
b: last session   weight_g  min   16.1 max   36.2 mean  24.19 median  24.00  within 10%: [np.True_, np.True_, np.True_, np.True_]  (n = 139)
c: every session  age_weeks min   13.1 max   65.9 mean  28.53 median  26.00  within 10%: [np.True_, np.False_, np.False_, np.True_]  (n = 459)
c: every session  weight_g  min   16.1 max   36.2 mean  24.32 median  24.10  within 10%: [np.True_, np.True_, np.True_, np.True_]  (n = 459)
Paper's code (meta/mice_info.py) takes age_weeks and reference_weight from the IBL database, not from the recording day.

group = probe_models.map(lambda m: "3A (NP1010)" if m == "NP1010" else "3B2 (PRB_1_4_0480_1)" if m.startswith("PRB_1_4_0480_1")
                         else "2.x (NP2010)" if m == "NP2010" else m)
print(group.value_counts().to_dict(), "(paper 3A 109, 3B2 586, 2.4 4)")

{'3B2 (PRB_1_4_0480_1)': 586, '3A (NP1010)': 109, '2.x (NP2010)': 4} (paper 3A 109, 3B2 586, 2.4 4)

n = sessions.n_trials
print(f"mean {n.mean():.1f} (645), median {n.median():.0f} (602), range {n.min()} to {n.max()} (401 to 1,525)")

mean 645.1 (645), median 601 (602), range 401 to 1525 (401 to 1,525)

def summary(v):
    v = np.asarray(v, float)
    return f"mean {v.mean():.2f}%, s.d. {v.std(ddof=1):.2f}, s.e.m. {v.std(ddof=1) / np.sqrt(len(v)):.2f} (n = {len(v)})"

perf = {
    "a: mice, all trials": trials.groupby("subject").correct.mean() * 100,
    "b: mice, analysed trials": trials[trials.analysed].groupby("subject").correct.mean() * 100,
    "c: sessions, all trials": trials.groupby("eid").correct.mean() * 100,
}
for k, v in perf.items():
    print(f"{k:26s}", summary(v), " within 81.4 +/- 0.4:", abs(v.mean() - 81.4) <= 0.4)
pooled = trials.correct.mean() * 100
print(f"{'d: all trials pooled':26s} {pooled:.2f}%  within 81.4 +/- 0.4: {abs(pooled - 81.4) <= 0.4}")
canon_t = trials[trials.eid.isin(set(canon.eid))]
v454 = {"a: mice, all trials": canon_t.groupby("subject").correct.mean() * 100,
        "c: sessions, all trials": canon_t.groupby("eid").correct.mean() * 100}
print(f"Variant on the {canon_t.eid.nunique()} sessions of the canonical set ({canon_t.subject.nunique()} mice):")
for k, v in v454.items():
    print(f"   {k:26s}", summary(v), " within 81.4 +/- 0.4:", abs(v.mean() - 81.4) <= 0.4)

a: mice, all trials        mean 80.76%, s.d. 5.12, s.e.m. 0.43 (n = 139)  within 81.4 +/- 0.4: False
b: mice, analysed trials   mean 84.62%, s.d. 5.40, s.e.m. 0.46 (n = 139)  within 81.4 +/- 0.4: False
c: sessions, all trials    mean 81.58%, s.d. 5.29, s.e.m. 0.25 (n = 459)  within 81.4 +/- 0.4: True

d: all trials pooled       81.94%  within 81.4 +/- 0.4: False

Variant on the 454 sessions of the canonical set (139 mice):
   a: mice, all trials        mean 80.77%, s.d. 5.12, s.e.m. 0.43 (n = 139)  within 81.4 +/- 0.4: False
   c: sessions, all trials    mean 81.59%, s.d. 5.28, s.e.m. 0.25 (n = 454)  within 81.4 +/- 0.4: True

canon_t = trials[trials.eid.isin(set(canon.eid))]
for name, d in (("a: all trials", trials), ("b: analysed trials", trials[trials.analysed]),
                ("variant, a on the 454 canonical sessions", canon_t), ("variant, b on the 454", canon_t[canon_t.analysed])):
    v = d[d.contrast == 0].groupby("subject").correct.mean() * 100
    t_, p_ = stats.ttest_1samp(v, 50)
    print(f"{name:41s} {summary(v)}; t{len(v) - 1} = {t_:.2f}, P = {p_:.2g}; "
          f"mean within 0.4: {abs(v.mean() - 58.7) <= 0.4}, t within 2.02: {abs(t_ - 20.18) <= 2.02}")
print("Paper: t138 = 20.18 gives P =", f"{stats.t.sf(20.18, 138) * 2:.2g}", "two-sided")

a: all trials                             mean 58.26%, s.d. 4.95, s.e.m. 0.42 (n = 139); t138 = 19.68, P = 6.9e-42; mean within 0.4: False, t within 2.02: True

b: analysed trials                        mean 59.29%, s.d. 6.76, s.e.m. 0.57 (n = 139); t138 = 16.21, P = 9.3e-34; mean within 0.4: False, t within 2.02: False
variant, a on the 454 canonical sessions  mean 58.27%, s.d. 4.96, s.e.m. 0.42 (n = 139); t138 = 19.68, P = 6.9e-42; mean within 0.4: False, t within 2.02: True
variant, b on the 454                     mean 59.31%, s.d. 6.76, s.e.m. 0.57 (n = 139); t138 = 16.23, P = 8.4e-34; mean within 0.4: False, t within 2.02: False
Paper: t138 = 20.18 gives P = 5.3e-43 two-sided

early = trials.rt < 0.08
has = trials.rt.notna()
print(f"a: over all trials {early.mean() * 100:.2f}% of {len(trials)}")
print(f"b: over trials with both times {early[has].mean() * 100:.2f}% of {int(has.sum())}")
print(f"c: mean over sessions {trials[has].groupby('eid').rt.apply(lambda r: np.mean(r < 0.08)).mean() * 100:.2f}%")
print(f"Of which before stimulus onset: {np.mean(trials.rt[has] < 0) * 100:.2f}% of trials with both times")

a: over all trials 22.75% of 296090
b: over trials with both times 22.99% of 292997
c: mean over sessions 22.49%

Of which before stimulus onset: 12.06% of trials with both times

resp = trials[trials.choice != 0]
by_c = resp.groupby(["subject", "contrast"]).agg(correct=("correct", "mean"), rt=("rt", "median")).reset_index()
pc = by_c.pivot(index="subject", columns="contrast", values="correct") * 100
rt = by_c.pivot(index="subject", columns="contrast", values="rt") * 1000
print("Median over mice, percent correct by contrast:", pc.median().round(1).to_dict())
print("Median over mice, median first movement time (ms) by contrast:", rt.median().round(0).to_dict())
w1 = stats.wilcoxon(pc[1.0], pc[0.0])
w2 = stats.wilcoxon(rt[1.0], rt[0.0])
print(f"100% against 0%: correct higher in {(pc[1.0] > pc[0.0]).sum()} of {len(pc)} mice, p = {w1.pvalue:.2g}; "
      f"first movement shorter in {(rt[1.0] < rt[0.0]).sum()} of {len(rt)} mice, p = {w2.pvalue:.2g}")
print("Median percent correct rises at every contrast step:", bool(np.all(np.diff(pc.median().values) > 0)))
print("Median first movement time falls at every contrast step:", bool(np.all(np.diff(rt.median().values) < 0)))

Median over mice, percent correct by contrast: {0.0: 59.2, 0.0625: 67.7, 0.125: 81.9, 0.25: 91.5, 1.0: 96.8}
Median over mice, median first movement time (ms) by contrast: {0.0: 196.0, 0.0625: 191.0, 0.125: 168.0, 0.25: 152.0, 1.0: 137.0}
100% against 0%: correct higher in 139 of 139 mice, p = 1.5e-24; first movement shorter in 133 of 139 mice, p = 5.9e-24
Median percent correct rises at every contrast step: True

Median first movement time falls at every contrast step: True

nz = trials[(trials.contrast > 0) & (trials.choice != 0)]
end = nz[nz.from_end <= 50].groupby("subject").agg(c=("correct", "mean"), rt=("rt", "median"))
mid = nz[(nz.from_end > 50) & (nz.from_end <= 300)].groupby("subject").agg(c=("correct", "mean"), rt=("rt", "median"))
w1 = stats.wilcoxon(end.c, mid.c)
w2 = stats.wilcoxon(end.rt, mid.rt)
print(f"Percent correct, last 50 trials: median over mice {end.c.median() * 100:.1f}; trials 51 to 300 from the end {mid.c.median() * 100:.1f}; "
      f"lower at the end in {(end.c < mid.c).sum()} of {len(end)} mice, p = {w1.pvalue:.2g}")
print(f"Median first movement time (ms), last 50: {end.rt.median() * 1000:.0f}; trials 51 to 300: {mid.rt.median() * 1000:.0f}; "
      f"longer at the end in {(end.rt > mid.rt).sum()} of {len(end)} mice, p = {w2.pvalue:.2g}")
curve = nz[nz.from_end <= 400].assign(bin=lambda d: (d.from_end - 1) // 50).groupby(["subject", "bin"]).agg(c=("correct", "mean"), rt=("rt", "median"))
print("By 50-trial bins from the end (bin 0 is the last 50), median over mice:")
print(curve.groupby("bin").median().assign(c=lambda d: (d.c * 100).round(1), rt=lambda d: (d.rt * 1000).round(0)).T.to_string())

Percent correct, last 50 trials: median over mice 80.0; trials 51 to 300 from the end 85.0; lower at the end in 105 of 139 mice, p = 7.9e-13
Median first movement time (ms), last 50: 279; trials 51 to 300: 165; longer at the end in 131 of 139 mice, p = 2.6e-22
By 50-trial bins from the end (bin 0 is the last 50), median over mice:
bin      0      1      2      3      4      5      6      7
c     80.0   84.5   85.8   86.4   85.4   86.7   86.7   85.7
rt   279.0  197.0  176.0  165.0  158.0  152.0  145.0  142.0

t2 = trials.copy()
t2["pos"] = t2.groupby(["eid", "block"]).cumcount() + 1
sw = t2[(t2.block >= 2) & (t2.pos <= 40) & (t2.contrast == 0) & (t2.choice != 0) & t2.eid.isin(BLOCK_OK)]
mouse_curve = sw.groupby(["subject", "pos"]).correct.mean().unstack()
mean_curve = mouse_curve.mean()
n = mean_curve.index.values.astype(float)
f = lambda x, a, b, tau: a - b * np.exp(-(x - 1) / tau)
(a, b, tau), _ = optimize.curve_fit(f, n, mean_curve.values, p0=(0.7, 0.4, 5), maxfev=20000)
print("Mean over mice of the fraction correct on 0% trials, by position after the switch:")
print(mean_curve.round(3).to_string())
print(f"Fit: a = {a:.3f}, b = {b:.3f}, tau = {tau:.2f} trials; 90% of the change at {1 + tau * np.log(10):.1f} trials")
print("Mice with 0% trials at position 1:", int(mouse_curve[1].notna().sum()), "; 0% trials in the curve:", len(sw),
      "; sessions with valid blocks:", len(BLOCK_OK))
fits = []
for k, row in mouse_curve.iterrows():
    try:
        p, _ = optimize.curve_fit(f, n[row.notna()], row.dropna().values, p0=(0.7, 0.4, 5), bounds=([0, -1, 0.1], [1, 1, 100]))
        fits.append(p[2])
    except RuntimeError:
        fits.append(np.nan)
fits = np.array(fits)
print(f"Per-mouse tau: median {np.nanmedian(fits):.2f}, interquartile {np.nanpercentile(fits, 25):.2f} to {np.nanpercentile(fits, 75):.2f}; "
      f"failed fits {int(np.isnan(fits).sum())} of {len(fits)}")
rng_boot = np.random.default_rng(60)
boot = []
mc = mouse_curve.values
for _ in range(1000):
    m = np.nanmean(mc[rng_boot.integers(0, len(mc), len(mc))], 0)
    try:
        boot.append(optimize.curve_fit(f, n, m, p0=(0.7, 0.4, 5), maxfev=20000)[0][2])
    except RuntimeError:
        boot.append(np.nan)
boot = np.array(boot)
lo, hi = np.nanpercentile(boot, [2.5, 97.5])
print(f"Bootstrap over mice (1000 resamples, {int(np.isnan(boot).sum())} failed fits): tau 95% interval {lo:.2f} to {hi:.2f}; "
      f"1 + tau ln 10 interval {1 + lo * np.log(10):.1f} to {1 + hi * np.log(10):.1f} trials; "
      f"resamples with 1 + tau ln 10 in 4.5 to 11: {np.nanmean((1 + boot * np.log(10) >= 4.5) & (1 + boot * np.log(10) <= 11)):.3f}")
fig, ax = plt.subplots(figsize=(5, 3))
ax.plot(n, mean_curve.values, "o", ms=3)
ax.plot(n, f(n, a, b, tau))
ax.set_xlabel("trial after block switch (0% contrast)")
ax.set_ylabel("fraction correct")
print("Plotted range:", round(mean_curve.min(), 3), "to", round(mean_curve.max(), 3))

Mean over mice of the fraction correct on 0% trials, by position after the switch:
pos
1     0.338
2     0.399
3     0.426
4     0.530
5     0.509
6     0.573
7     0.609
8     0.491
9     0.534
10    0.603
11    0.613
12    0.611
13    0.646
14    0.620
15    0.604
16    0.726
17    0.623
18    0.670
19    0.606
20    0.641
21    0.646
22    0.553
23    0.564
24    0.627
25    0.664
26    0.364
27    0.512
28    0.749
29    0.586
30    0.742
31    0.497
32    0.657
33    0.656
34    0.601
35    0.719
36    0.707
37    0.716
38    0.543
39    0.593
40    0.659
Fit: a = 0.625, b = 0.290, tau = 3.86 trials; 90% of the change at 9.9 trials
Mice with 0% trials at position 1: 90 ; 0% trials in the curve: 16447 ; sessions with valid blocks: 395

Per-mouse tau: median 3.88, interquartile 1.86 to 11.30; failed fits 0 of 125

Bootstrap over mice (1000 resamples, 0 failed fits): tau 95% interval 2.79 to 5.30; 1 + tau ln 10 interval 7.4 to 13.2 trials; resamples with 1 + tau ln 10 in 4.5 to 11: 0.759
Plotted range: 0.338 to 0.749

rb = resp[(resp.p_left != 0.5) & resp.eid.isin(BLOCK_OK)]
pr = rb.assign(right=rb.choice > 0).groupby(["subject", "p_left", "signed_contrast"]).right.mean().unstack(["p_left"])
diff = (pr[0.2] - pr[0.8]).unstack("signed_contrast")
rows = []
for c in diff.columns:
    d = diff[c].dropna()
    rows.append({"signed_contrast": c, "mice": len(d), "median_diff": round(d.median(), 3), "higher_in": int((d > 0).sum()),
                 "p": stats.wilcoxon(d).pvalue})
bias = pd.DataFrame(rows)
print(bias.to_string(index=False))
print("Contrasts with a higher right-block fraction (median > 0 and p < 0.05):", int(((bias.median_diff > 0) & (bias.p < 0.05)).sum()), "of 9",
      "; mice:", rb.subject.nunique(), "; sessions with valid blocks:", rb.eid.nunique())
allb = resp[resp.p_left != 0.5].assign(right=lambda d: d.choice > 0).groupby(["subject", "p_left", "signed_contrast"]).right.mean().unstack("p_left")
d_all = (allb[0.2] - allb[0.8]).unstack("signed_contrast")
print("Variant with all sessions, median difference by contrast:", d_all.median().round(3).to_dict())

 signed_contrast  mice  median_diff  higher_in            p
         -1.0000   125        0.013         76 9.931702e-09
         -0.2500   125        0.064        104 9.945649e-16
         -0.1250   125        0.171        115 3.557317e-21
         -0.0625   125        0.271        120 1.197511e-21
          0.0000   125        0.333        123 3.344639e-22
          0.0625   125        0.282        123 4.154633e-22
          0.1250   125        0.170        119 1.096578e-20
          0.2500   125        0.073         97 1.635191e-11
          1.0000   125        0.028         78 1.011955e-09
Contrasts with a higher right-block fraction (median > 0 and p < 0.05): 9 of 9 ; mice: 125 ; sessions with valid blocks: 395
Variant with all sessions, median difference by contrast: {-1.0: 0.012, -0.25: 0.058, -0.125: 0.151, -0.0625: 0.241, 0.0: 0.314, 0.0625: 0.269, 0.125: 0.137, 0.25: 0.061, 1.0: 0.017}

EXAMPLES = ["VISp", "GRN", "IRN"]
LATENCY = ["LGd", "VISp", "LP", "VISpm", "VISam"]
n_sess = canon.groupby("region").eid.nunique()
pool = sorted(n_sess[n_sess <= 12].index.difference(EXAMPLES + LATENCY))
RANDOM = sorted(np.random.default_rng(409).choice(pool, 7, replace=False).tolist())
DECODING = EXAMPLES + RANDOM
TRAJ = LATENCY + ["IRN", "GRN"]
TESTED = sorted(set(DECODING) | set(TRAJ))
print("Random regions:", RANDOM, "from", len(pool), "candidates")
sel = canon[canon.region.isin(TESTED)]
print(sel.groupby("region").agg(sessions=("eid", "nunique"), units=("unit", "size")).T.to_string())
jobs = []
for aid, g in sel.groupby("asset_id"):
    jobs.append((aid, g.unit.tolist(), g[g.region.isin(TRAJ)].unit.tolist(), g[g.region == "IRN"].unit.tolist()))
t0 = time.time()
extracted = bwm.run_parallel(bwm.extract_session, jobs)
extract_seconds = time.time() - t0
by_asset = {e["asset_id"]: e for e in extracted}
t_eid = dict(zip(processed.asset_id, sessions.eid))
unit_region = {(r.asset_id, r.unit): r.region for r in sel.itertuples()}
print(f"Streamed {len(extracted)} sessions, {sum(e['n_spikes'] for e in extracted) / 1e6:.0f} million spikes of {len(sel)} units, "
      f"in {extract_seconds:.0f} s ({extract_seconds / len(extracted):.2f} s per session with {bwm.N_WORKERS} processes)")
print("Units whose spike times were not sorted:", sum(e["unsorted_units"] for e in extracted))

Random regions: ['ACAv', 'CL', 'HATA', 'LDT', 'PA', 'PVT', 'TR'] from 140 candidates
region    ACAv  CL  GRN  HATA  IRN  LDT  LGd    LP   PA  PVT  TR  VISam  VISp  VISpm
sessions    11   2   14     5   27    2   22    58    7    5   3     18    35     10
units      187  23  555    58  739   51  551  1780  106   77  35    363   844    240

Streamed 168 sessions, 222 million spikes of 5609 units, in 71 s (0.42 s per session with 8 processes)
Units whose spike times were not sorted: 0

# Every canonical unit of every canonical session: trajectory sums for traj-stim-regions (one set of
# permutations per session, shared by its regions, with the Methods' and the code's classes) and the
# clock check.
t0 = time.time()
canon_jobs = [(aid, g.unit.values, g.region.values, 40000 + k) for k, (aid, g) in enumerate(canon.groupby("asset_id", sort=True))]
acc_traj = {k: ({}, {}) for k in ("methods_ok", "methods_all", "code_ok", "code_all")}   # summed squares, cells
clock_all = []
for r in bwm.run_parallel_iter(bwm.region_traj_job, canon_jobs):
    ok = t_eid[r["asset_id"]] in BLOCK_OK
    clock_all.append({"asset_id": r["asset_id"], "n_analysed": r["n_analysed"], "acc": r["clock"]["acc"],
                      "shifted": r["clock"]["shifted"], "block_ok": ok})
    for reg in r["sq"]:
        for key, part in (("methods_ok", "sq"), ("methods_all", "sq"), ("code_ok", "sq_code"), ("code_all", "sq_code")):
            if key.endswith("_ok") and not ok:
                continue
            sq_, n_ = acc_traj[key]
            sq_[reg] = sq_.get(reg, 0) + r[part][reg].astype(np.float64)
            n_[reg] = n_.get(reg, 0) + r["n"][reg]
all_seconds = time.time() - t0
print(f"Streamed and binned {len(canon)} canonical units of {len(canon_jobs)} sessions in {all_seconds:.0f} s; "
      f"sessions with no analysed trial: {sum(c['n_analysed'] == 0 for c in clock_all)}")

Streamed and binned 62857 canonical units of 454 sessions in 707 s; sessions with no analysed trial: 0

# Spikes against the trials' clock in every canonical session, with all its canonical units.
acc_all = np.array([c["acc"] for c in clock_all])
jump_all = np.nanmean(acc_all[:, bwm.CLOCK_OFFSETS >= 0], 1) - np.nanmean(acc_all[:, bwm.CLOCK_OFFSETS <= -0.4], 1)
print("Balanced accuracy of feedback decoding by window start relative to feedback (s), median over sessions:")
print(dict(zip(bwm.CLOCK_OFFSETS.tolist(), np.nanmedian(acc_all, 0).round(3).tolist())))
print(f"Jump from the windows before to those from feedback on: median {np.nanmedian(jump_all):.3f}; above 0.05 in "
      f"{int(np.sum(jump_all > 0.05))} of {len(jump_all)} sessions; labels shifted by one trial: median "
      f"{np.nanmedian([c['shifted'] for c in clock_all]):.3f}")
bad_all = [processed.set_index("asset_id").loc[c["asset_id"], "path"].split("/")[0] + " " + t_eid[c["asset_id"]][:8]
           for c, j in zip(clock_all, jump_all) if not j > 0.05]
print("Sessions where the jump is not established:", len(bad_all), bad_all)

Balanced accuracy of feedback decoding by window start relative to feedback (s), median over sessions:
{-1.0: 0.52, -0.8: 0.526, -0.6: 0.532, -0.4: 0.54, -0.2: 0.577, -0.0: 0.902, 0.2: 0.94, 0.4: 0.921}
Jump from the windows before to those from feedback on: median 0.374; above 0.05 in 448 of 454 sessions; labels shifted by one trial: median 0.558
Sessions where the jump is not established: 6 ['sub-ibl-witten-29 58b1e920', 'sub-ibl-witten-29 af55d16f', 'sub-MFD-07 c46b8def', 'sub-PL024 30af8629', 'sub-NYU-12 b182b754', 'sub-UCLA034 8a3a0197']

# Region-session inputs.
def region_counts(region, var_keys=("stim", "choice", "feedback")):
    out = []
    for aid, g in sel[sel.region == region].groupby("asset_id"):
        e = by_asset[aid]
        cols = np.isin(e["units"], g.unit.values)
        out.append((aid, e["trials"], {k: e["counts"][k][:, cols] for k in var_keys}))
    return out

import sklearn
from sklearn.linear_model import LogisticRegression
aid0, t_0, c_0 = region_counts("VISp")[0]
rows0, y0, _ = bwm.decoding_targets(t_0, "stim")
X0 = c_0["stim"][rows0].astype(float)
print("scikit-learn", sklearn.__version__, "; decoder parameters:", {k: v for k, v in bwm._logreg(1.0).get_params().items() if k in ("penalty", "l1_ratio", "solver")})
for C_ in (0.01, 0.1, 1.0):
    l1 = bwm._logreg(C_).fit(X0, y0)
    l2 = LogisticRegression(C=C_, l1_ratio=0.0, solver="liblinear", tol=1e-3, max_iter=20000, class_weight="balanced", random_state=0).fit(X0, y0)
    print(f"   one VISp session, stimulus, C = {C_}: nonzero coefficients {int((l1.coef_ != 0).sum())} of {l1.coef_.size} "
          f"(the same fit with l1_ratio 0, L2: {int((l2.coef_ != 0).sum())})")

t0 = time.time()
dec_jobs = []
for region in DECODING:
    for k, (aid, t, c) in enumerate(region_counts(region)):
        dec_jobs.append(((region, aid), t, c, 1000 * (DECODING.index(region) + 1) + k))
dec = bwm.run_parallel(bwm.decode_job, dec_jobs, workers=4)
decode_seconds = time.time() - t0
rows = []
for d in dec:
    for var in ("stim", "choice", "feedback"):
        rows.append({"region": d["key"][0], "asset_id": d["key"][1], "var": var, **{k: d[var][k] for k in ("n_trials", "n_minor", "n_null", "score", "p", "corrected")}})
dec_df = pd.DataFrame(rows)
print(f"Logistic fits: {sum(d['fits'] for d in dec)}; not converged within 20000 iterations: {sum(d['not_converged'] for d in dec)}")
print(f"{len(dec_jobs)} region-sessions decoded in {decode_seconds:.0f} s; null sessions per decoding: "
      f"{dec_df.n_null.min()} to {dec_df.n_null.max()}; learning rates of the action-kernel model: "
      f"{pd.Series([d['alpha'] for d in dec]).value_counts().sort_index().to_dict()}")
region_res = []
for (region, var), g in dec_df.groupby(["region", "var"]):
    region_res.append({"region": region, "var": var, "sessions": len(g), "p_fisher": bwm.fisher_p(g.p), "effect": g.corrected.median(),
                       "median_score": g.score.median(), "sessions_p05": int((g.p < 0.05).sum())})
region_res = pd.DataFrame(region_res)
region_res["significant"] = False
for var, g in region_res.groupby("var"):
    region_res.loc[g.index, "significant"] = bwm.bh(g.p_fisher.values, 0.01)
ri = region_info.set_index("Beryl")
paper_col = {"stim": "stim_dec", "choice": "choice_dec", "feedback": "fback_dec"}
region_res["paper_effect"] = [ri.loc[r, paper_col[v]] for r, v in zip(region_res.region, region_res["var"])]
region_res["paper_significant"] = [bool(ri.loc[r, paper_col[v] + "_sig"]) for r, v in zip(region_res.region, region_res["var"])]
print(region_res.round(4).sort_values(["var", "region"]).to_string(index=False))
for q in (0.01, 0.05):
    print(f"Variant, BH at q = {q}:", {v: int(bwm.bh(g.p_fisher.values, q).sum()) for v, g in region_res.groupby("var")},
          "; Fisher p < 0.01 uncorrected:", {v: int((g.p_fisher < 0.01).sum()) for v, g in region_res.groupby("var")})
subj_of = dict(zip(processed.asset_id, sessions.subject))
per = dec_df[dec_df.region.isin(EXAMPLES)].assign(subject=lambda d: d.asset_id.map(subj_of))
print("Per session, example regions (score and p), by mouse:")
print(per.pivot_table(index=["region", "subject", "asset_id"], columns="var", values=["score", "p"]).round(3).droplevel(2).to_string())

scikit-learn 1.9.1 ; decoder parameters: {'l1_ratio': 1.0, 'penalty': 'deprecated', 'solver': 'liblinear'}
   one VISp session, stimulus, C = 0.01: nonzero coefficients 0 of 13 (the same fit with l1_ratio 0, L2: 12)
   one VISp session, stimulus, C = 0.1: nonzero coefficients 2 of 13 (the same fit with l1_ratio 0, L2: 12)
   one VISp session, stimulus, C = 1.0: nonzero coefficients 6 of 13 (the same fit with l1_ratio 0, L2: 12)

Logistic fits: 226440; not converged within 20000 iterations: 6
111 region-sessions decoded in 334 s; null sessions per decoding: 100 to 100; learning rates of the action-kernel model: {0.02: 6, 0.05: 1, 0.1: 3, 0.15: 5, 0.2: 18, 0.3: 33, 0.4: 22, 0.5: 11, 0.7: 12}
region      var  sessions  p_fisher  effect  median_score  sessions_p05  significant  paper_effect  paper_significant
  ACAv   choice        11    0.0000  0.0760        0.5817             7         True          0.07               True
    CL   choice         2    0.0010  0.2467        0.7598             2         True          0.25               True
   GRN   choice        14    0.0000  0.2774        0.8214            14         True          0.28               True
  HATA   choice         5    0.0123  0.0234        0.5326             1        False          0.02              False
   IRN   choice        27    0.0000  0.1800        0.7182            22         True          0.17               True
   LDT   choice         2    0.0154  0.1300        0.6422             1        False          0.15               True
    PA   choice         7    0.4726 -0.0165        0.4960             1        False         -0.02              False
   PVT   choice         5    0.0413  0.0161        0.5221             1        False          0.03              False
    TR   choice         3    0.4482  0.0072        0.5161             0        False          0.02              False
  VISp   choice        35    0.0000  0.0423        0.5612            15         True          0.06               True
  ACAv feedback        11    0.0000  0.0528        0.5681             6         True          0.09               True
    CL feedback         2    0.0010  0.2672        0.7666             2         True          0.25               True
   GRN feedback        14    0.0000  0.4061        0.9130            14         True          0.42               True
  HATA feedback         5    0.0255  0.0025        0.5048             2        False          0.04               True
   IRN feedback        27    0.0000  0.4066        0.9052            25         True          0.42               True
   LDT feedback         2    0.0010  0.3045        0.8071             2         True          0.32               True
    PA feedback         7    0.1455  0.0000        0.5000             2        False         -0.01              False
   PVT feedback         5    0.0008  0.0888        0.5960             3         True          0.04               True
    TR feedback         3    0.0227  0.0493        0.5473             1        False          0.03              False
  VISp feedback        35    0.0000  0.1133        0.6149            20         True          0.11               True
  ACAv     stim        11    0.5733  0.0058        0.5146             0        False          0.01              False
    CL     stim         2    0.0010  0.1655        0.6710             2         True          0.17               True
   GRN     stim        14    0.0000  0.0559        0.5950             7         True          0.08               True
  HATA     stim         5    0.1914  0.0000        0.5051             1        False          0.01              False
   IRN     stim        27    0.0007  0.0290        0.5460             5         True          0.01               True
   LDT     stim         2    0.0912  0.0141        0.5294             1        False          0.04              False
    PA     stim         7    0.4257  0.0049        0.5122             0        False          0.00              False
   PVT     stim         5    0.9696 -0.0188        0.4812             0        False         -0.01              False
    TR     stim         3    0.9720 -0.0338        0.4836             0        False         -0.01              False
  VISp     stim        35    0.0000  0.0387        0.5462            15         True          0.02               True
Variant, BH at q = 0.01: {'choice': 5, 'feedback': 7, 'stim': 4} ; Fisher p < 0.01 uncorrected: {'choice': 5, 'feedback': 7, 'stim': 4}
Variant, BH at q = 0.05: {'choice': 7, 'feedback': 9, 'stim': 4} ; Fisher p < 0.01 uncorrected: {'choice': 5, 'feedback': 7, 'stim': 4}
Per session, example regions (score and p), by mouse:
                          p                  score                
var                  choice feedback   stim choice feedback   stim
region subject                                                    
GRN    CSHL045        0.010    0.010  0.208  0.686    0.698  0.579
       CSHL060        0.010    0.010  0.347  0.656    0.809  0.526
       CSHL060        0.010    0.010  0.347  0.782    0.918  0.543
       KS045          0.010    0.010  0.287  0.924    0.738  0.542
       NR_0028        0.010    0.010  0.010  0.940    0.968  0.746
       NYU-37         0.010    0.010  0.010  0.885    0.983  0.715
       SWC_054        0.010    0.010  0.010  0.826    0.974  0.627
       SWC_054        0.010    0.010  0.059  0.813    0.841  0.591
       SWC_054        0.010    0.010  0.010  0.862    0.924  0.652
       UCLA033        0.010    0.010  0.010  0.817    0.984  0.642
       ZFM-01592      0.010    0.010  0.010  0.841    0.951  0.691
       ibl_witten_20  0.010    0.010  0.337  0.605    0.870  0.530
       ibl_witten_20  0.010    0.010  0.020  0.850    0.895  0.575
       ibl_witten_27  0.010    0.010  0.149  0.707    0.908  0.599
IRN    CSHL045        0.030    0.010  0.109  0.582    0.905  0.572
       CSHL060        0.010    0.010  0.050  0.660    0.804  0.593
       CSHL060        0.020    0.010  0.743  0.641    0.928  0.507
       CSH_ZAD_022    0.010    0.010  0.069  0.800    0.722  0.605
       CSH_ZAD_025    0.129    0.010  0.406  0.545    0.748  0.520
       CSH_ZAD_025    0.129    0.277  0.644  0.581    0.548  0.518
       CSH_ZAD_029    0.010    0.010  0.040  0.804    0.963  0.574
       DY_020         0.010    0.010  0.832  0.625    0.979  0.481
       DY_020         0.010    0.010  0.079  0.850    1.000  0.594
       DY_020         0.010    0.010  0.347  0.822    0.975  0.540
       KS043          0.010    0.010  0.584  0.695    0.689  0.520
       MFD_08         0.327    0.515  0.634  0.530    0.499  0.494
       PL030          0.267    0.010  0.416  0.581    0.741  0.538
       SWC_052        0.010    0.010  0.238  0.750    0.835  0.589
       SWC_052        0.010    0.010  0.941  0.654    0.760  0.480
       SWC_053        0.010    0.010  0.109  0.740    0.966  0.573
       SWC_054        0.010    0.010  0.010  0.849    0.984  0.639
       SWC_058        0.010    0.010  0.624  0.800    0.986  0.494
       SWC_058        0.010    0.010  0.129  0.650    0.724  0.546
       UCLA033        0.010    0.010  0.059  0.718    0.993  0.575
       ZFM-01576      0.010    0.010  0.604  0.757    0.974  0.510
       ZFM-01592      0.010    0.010  0.010  0.883    0.956  0.611
       ibl_witten_20  0.010    0.010  0.149  0.733    0.896  0.573
       ibl_witten_20  0.010    0.010  0.089  0.772    0.937  0.566
       ibl_witten_20  0.010    0.010  0.010  0.895    0.990  0.702
       ibl_witten_26  0.099    0.010  0.772  0.561    0.905  0.496
       ibl_witten_27  0.010    0.010  0.703  0.670    0.848  0.511
VISp   CSHL052        0.337    0.010  0.891  0.533    0.635  0.480
       CSH_ZAD_011    0.059    0.931  0.010  0.558    0.455  0.634
       CSH_ZAD_022    0.069    0.010  0.010  0.609    0.762  0.704
       CSH_ZAD_025    0.287    0.010  0.901  0.601    0.766  0.490
       CSH_ZAD_025    0.010    0.040  0.010  0.782    0.585  0.810
       CSH_ZAD_026    0.010    0.010  0.010  0.675    0.691  0.773
       CSH_ZAD_029    0.010    0.010  0.772  0.582    0.649  0.495
       DY_009         0.010    0.010  0.743  0.772    0.626  0.484
       DY_009         0.178    0.069  0.871  0.537    0.583  0.481
       DY_010         0.713    0.020  0.040  0.499    0.625  0.579
       KS014          0.406    0.010  0.030  0.513    0.793  0.546
       KS042          0.010    0.010  0.218  0.605    0.629  0.536
       MFD_05         0.683    0.109  0.980  0.489    0.575  0.447
       MFD_08         0.673    0.317  0.980  0.495    0.519  0.450
       NR_0020        0.010    0.010  0.010  0.676    0.662  0.709
       NR_0021        0.010    0.277  0.010  0.746    0.542  0.809
       NR_0024        0.238    0.772  0.168  0.554    0.473  0.566
       NR_0028        0.069    0.010  0.099  0.584    0.599  0.551
       NR_0028        0.010    0.010  0.030  0.677    0.702  0.567
       NR_0031        0.436    0.020  0.851  0.532    0.640  0.485
       NR_0031        0.386    0.109  0.792  0.534    0.576  0.492
       NYU-12         0.188    0.109  0.139  0.540    0.559  0.547
       PL050          0.297    0.297  0.446  0.528    0.516  0.514
       PL050          0.010    0.010  0.010  0.738    0.754  0.698
       PL050          0.990    0.307  0.921  0.471    0.525  0.480
       PL050          0.010    0.416  0.218  0.596    0.505  0.529
       SWC_052        0.653    0.149  0.257  0.493    0.550  0.531
       UCLA015        0.010    0.020  0.010  0.561    0.615  0.645
       UCLA033        0.010    0.010  0.050  0.683    0.863  0.585
       UCLA034        0.257    0.010  0.723  0.544    0.687  0.501
       UCLA037        0.376    0.109  0.693  0.509    0.547  0.485
       UCLA044        0.050    0.010  0.010  0.569    0.659  0.614
       ZFM-01577      0.030    0.010  0.010  0.576    0.788  0.605
       ibl_witten_20  0.505    0.822  0.010  0.528    0.477  0.606
       ibl_witten_29  0.010    0.950  0.356  0.591    0.436  0.514

nsig = region_res.groupby("var").significant.sum()
print("Tested regions significant (BH q = 0.01 over the 10 regions):", nsig.to_dict())
visp = region_res[(region_res.region == "VISp") & (region_res["var"] == "stim")].iloc[0]
print(f"VISp stimulus: {visp.sessions} sessions, Fisher p = {visp.p_fisher:.2g}, effect {visp.effect:.3f} (paper {visp.paper_effect}), "
      f"sessions with p < 0.05: {visp.sessions_p05}")
s = region_res[region_res["var"] == "stim"]
print("Agreement with the paper's significance:", int((s.significant == s.paper_significant).sum()), "of", len(s))
print("Spearman correlation of effects with the paper's:", stats.spearmanr(s.effect, s.paper_effect))
print("Paper over 201 regions: stimulus 40, choice 172, feedback 198 significant")

Tested regions significant (BH q = 0.01 over the 10 regions): {'choice': 5, 'feedback': 7, 'stim': 4}
VISp stimulus: 35 sessions, Fisher p = 1.3e-08, effect 0.039 (paper 0.02), sessions with p < 0.05: 15
Agreement with the paper's significance: 10 of 10
Spearman correlation of effects with the paper's: SignificanceResult(statistic=np.float64(0.9108770332470622), pvalue=np.float64(0.000247580323683572))
Paper over 201 regions: stimulus 40, choice 172, feedback 198 significant

grn = region_res[(region_res.region == "GRN") & (region_res["var"] == "choice")].iloc[0]
print(f"GRN choice: {grn.sessions} sessions, Fisher p = {grn.p_fisher:.2g}, effect {grn.effect:.3f} (paper {grn.paper_effect})")
c = region_res[region_res["var"] == "choice"]
print("Significant for choice:", int(c.significant.sum()), "; for stimulus:", int(nsig["stim"]))
print("Agreement with the paper's significance:", int((c.significant == c.paper_significant).sum()), "of", len(c))
print("Spearman correlation of effects with the paper's:", stats.spearmanr(c.effect, c.paper_effect))
print(c[["region", "sessions", "p_fisher", "effect", "paper_effect", "significant", "paper_significant"]].round(4).to_string(index=False))

GRN choice: 14 sessions, Fisher p = 6e-15, effect 0.277 (paper 0.28)
Significant for choice: 5 ; for stimulus: 4
Agreement with the paper's significance: 9 of 10
Spearman correlation of effects with the paper's: SignificanceResult(statistic=np.float64(0.9787279253249042), pvalue=np.float64(8.73148286763855e-07))
region  sessions  p_fisher  effect  paper_effect  significant  paper_significant
  ACAv        11    0.0000  0.0760          0.07         True               True
    CL         2    0.0010  0.2467          0.25         True               True
   GRN        14    0.0000  0.2774          0.28         True               True
  HATA         5    0.0123  0.0234          0.02        False              False
   IRN        27    0.0000  0.1800          0.17         True               True
   LDT         2    0.0154  0.1300          0.15        False               True
    PA         7    0.4726 -0.0165         -0.02        False              False
   PVT         5    0.0413  0.0161          0.03        False              False
    TR         3    0.4482  0.0072          0.02        False              False
  VISp        35    0.0000  0.0423          0.06         True               True

fb = region_res[region_res["var"] == "feedback"]
print("Significant for feedback:", int(fb.significant.sum()), "of", len(fb), "; IRN significant:", bool(fb[fb.region == "IRN"].significant.iloc[0]))
print(fb[["region", "sessions", "p_fisher", "effect", "median_score", "paper_effect", "significant", "paper_significant"]].round(4).to_string(index=False))
f_sess = dec_df[dec_df["var"] == "feedback"]
print(f"Region-sessions with balanced accuracy above 0.9: {int((f_sess.score > 0.9).sum())} of {len(f_sess)}; median accuracy {f_sess.score.median():.3f}")
print("Spearman correlation of effects with the paper's:", stats.spearmanr(fb.effect, fb.paper_effect))

Significant for feedback: 7 of 10 ; IRN significant: True
region  sessions  p_fisher  effect  median_score  paper_effect  significant  paper_significant
  ACAv        11    0.0000  0.0528        0.5681          0.09         True               True
    CL         2    0.0010  0.2672        0.7666          0.25         True               True
   GRN        14    0.0000  0.4061        0.9130          0.42         True               True
  HATA         5    0.0255  0.0025        0.5048          0.04        False               True
   IRN        27    0.0000  0.4066        0.9052          0.42         True               True
   LDT         2    0.0010  0.3045        0.8071          0.32         True               True
    PA         7    0.1455  0.0000        0.5000         -0.01        False              False
   PVT         5    0.0008  0.0888        0.5960          0.04         True               True
    TR         3    0.0227  0.0493        0.5473          0.03        False              False
  VISp        35    0.0000  0.1133        0.6149          0.11         True               True
Region-sessions with balanced accuracy above 0.9: 23 of 111; median accuracy 0.662
Spearman correlation of effects with the paper's: SignificanceResult(statistic=np.float64(0.9573348703109935), pvalue=np.float64(1.3767647149143998e-05))

# Variant for the three decoding claims: the null with the full nested search of C, against the
# null that reuses the real labels' C, on three VISp, GRN and IRN sessions with 20 nulls each.
nn_jobs = []
for k, region in enumerate(EXAMPLES):
    aid, t, c = region_counts(region)[0]
    nn_jobs.append(((region, aid), t, c, 5000 + k, 20))
nn = bwm.run_parallel(bwm.nested_null_job, nn_jobs, workers=3)
for r in nn:
    for var in ("stim", "choice", "feedback"):
        a, b = r[var]["nested"], r[var]["reused"]
        print(f"{r['key'][0]:5s} {var:8s} nested null: median {np.median(a):.3f}, 95th pct {np.percentile(a, 95):.3f}; "
              f"reused C: median {np.median(b):.3f}, 95th pct {np.percentile(b, 95):.3f}")

VISp  stim     nested null: median 0.518, 95th pct 0.575; reused C: median 0.508, 95th pct 0.531
VISp  choice   nested null: median 0.499, 95th pct 0.574; reused C: median 0.524, 95th pct 0.567
VISp  feedback nested null: median 0.507, 95th pct 0.584; reused C: median 0.512, 95th pct 0.565
GRN   stim     nested null: median 0.554, 95th pct 0.622; reused C: median 0.559, 95th pct 0.598
GRN   choice   nested null: median 0.527, 95th pct 0.577; reused C: median 0.524, 95th pct 0.577
GRN   feedback nested null: median 0.506, 95th pct 0.564; reused C: median 0.509, 95th pct 0.561
IRN   stim     nested null: median 0.529, 95th pct 0.562; reused C: median 0.523, 95th pct 0.560
IRN   choice   nested null: median 0.529, 95th pct 0.555; reused C: median 0.532, 95th pct 0.584
IRN   feedback nested null: median 0.498, 95th pct 0.559; reused C: median 0.517, 95th pct 0.569

# Variant for dec-stim-restricted: VISp stimulus with the full nested search for every null, on the
# first 30 of the same pseudo-sessions in each of its sessions, against the reused C on those 30.
N_NESTED = 30
visp_jobs, reused_p = [], []
for d, job in zip(dec, dec_jobs):
    if d["key"][0] != "VISp":
        continue
    visp_jobs.append((d["key"], job[1], job[2], job[3], d["stim"]["score"], N_NESTED))
    nl = d["stim"]["null"][:N_NESTED]
    reused_p.append((1 + np.sum(nl >= d["stim"]["score"])) / (1 + len(nl)))
t0 = time.time()
vn = bwm.run_parallel(bwm.nested_p_job, visp_jobs, workers=4)
nested_p = np.array([r["p"] for r in vn])
reused_p = np.array(reused_p)
print(f"VISp stimulus, {len(vn)} sessions, {N_NESTED} nulls each, in {time.time() - t0:.0f} s")
print(f"   per-session p, nested null: median {np.median(nested_p):.3f}, p < 0.05 in {int((nested_p < 0.05).sum())}; "
      f"reused C: median {np.median(reused_p):.3f}, p < 0.05 in {int((reused_p < 0.05).sum())}")
print(f"   95th percentile of the null, median over sessions: nested {np.median([np.percentile(r['null'], 95) for r in vn]):.3f}, "
      f"reused {np.median([np.percentile(d['stim']['null'][:N_NESTED], 95) for d in dec if d['key'][0] == 'VISp']):.3f}")
fp_nested, fp_reused = bwm.fisher_p(nested_p), bwm.fisher_p(reused_p)
s_p = region_res[region_res["var"] == "stim"].set_index("region").p_fisher.copy()
s_p["VISp"] = fp_nested
print(f"   Fisher p, nested null {fp_nested:.2g}, reused C on the same {N_NESTED} nulls {fp_reused:.2g}; "
      f"VISp significant with BH at q = 0.01 among the 10 stimulus p-values when its p is the nested one: "
      f"{bool(bwm.bh(s_p.values, 0.01)[list(s_p.index).index('VISp')])}")
print(pd.DataFrame({"session": [r["key"][1][:8] for r in vn], "score": [r_["stim"]["score"] for r_ in dec if r_["key"][0] == "VISp"],
                    "p_reused": reused_p, "p_nested": nested_p}).round(3).to_string(index=False))

VISp stimulus, 35 sessions, 30 nulls each, in 154 s
   per-session p, nested null: median 0.129, p < 0.05 in 13; reused C: median 0.161, p < 0.05 in 13
   95th percentile of the null, median over sessions: nested 0.565, reused 0.558
   Fisher p, nested null 1.8e-05, reused C on the same 30 nulls 2.8e-05; VISp significant with BH at q = 0.01 among the 10 stimulus p-values when its p is the nested one: True
 session  score  p_reused  p_nested
0032161b  0.485     0.968     0.710
0e06cf91  0.514     0.323     0.323
0f52c917  0.514     0.452     0.548
1287de45  0.566     0.161     0.129
2bf1e46b  0.480     0.871     0.774
36d20ff1  0.773     0.032     0.032
398fa808  0.546     0.032     0.129
49b2ec1c  0.809     0.032     0.032
4a9e1f83  0.484     0.774     0.677
58360097  0.531     0.290     0.323
58834f11  0.536     0.161     0.161
5a368a20  0.698     0.032     0.032
6023d1a1  0.480     0.903     0.903
6f479578  0.634     0.032     0.032
766adddb  0.492     0.903     0.742
7dbe1f56  0.551     0.065     0.129
8e42dd5f  0.495     0.839     0.710
9c6bc203  0.481     0.839     0.871
a2929247  0.605     0.032     0.032
a71b50e7  0.450     1.000     0.968
b61d8516  0.490     0.871     0.968
bfbdfa2d  0.810     0.032     0.032
c04a3528  0.485     0.742     0.774
c0a278a1  0.614     0.032     0.032
c28a733b  0.547     0.161     0.065
c6bfe1cf  0.645     0.032     0.032
c919fa56  0.709     0.032     0.032
cc1b0a9f  0.606     0.032     0.032
d9da8c43  0.704     0.032     0.032
e9c2b6f3  0.567     0.032     0.032
f2f596d6  0.501     0.710     0.613
faca19fd  0.585     0.097     0.032
fbebc211  0.447     1.000     0.935
fc66281f  0.529     0.194     0.290
fd25a403  0.579     0.097     0.065

t0 = time.time()
sc_jobs = []
for region in TESTED:
    for k, (aid, t, c) in enumerate(region_counts(region, ("stim", "choice"))):
        if t_eid[aid] not in BLOCK_OK:
            continue
        sc_jobs.append(((region, aid), t, c, 20000 + 100 * TESTED.index(region) + k))
sc = bwm.run_parallel(bwm.single_cell_job, sc_jobs, workers=4)
left_out = sel[~sel.eid.isin(BLOCK_OK)]
print(f"Left out for invalid blocks: {left_out.groupby(['region', 'eid']).ngroups} region-sessions, {len(left_out)} units")
rows = []
for r in sc:
    for i in range(len(r["stim_real"])):
        rows.append({"region": r["key"][0], "asset_id": r["key"][1], **{k: bool(r[k][i]) for k in r if k != "key"}})
cells = pd.DataFrame(rows)
print(f"{len(cells)} units in {len(sc_jobs)} region-sessions, in {time.time() - t0:.0f} s")
frac = cells.groupby("region")[["stim_real", "choice_real", "stim_null", "choice_null", "stim_shuf", "choice_shuf"]].mean()
frac["units"] = cells.groupby("region").size()
frac["paper_stim"] = [ri.loc[r, "stim_man"] for r in frac.index]
frac["paper_choice"] = [ri.loc[r, "choice_man"] for r in frac.index]
print(frac.round(4).to_string())

Left out for invalid blocks: 19 region-sessions, 381 units
5228 units in 200 region-sessions, in 29 s
        stim_real  choice_real  stim_null  choice_null  stim_shuf  choice_shuf  units  paper_stim  paper_choice
region                                                                                                         
ACAv       0.0000       0.0214     0.0000       0.0053     0.0000       0.0000    187        0.00          0.01
CL         0.0588       0.3529     0.0000       0.0000     0.0000       0.0000     17        0.04          0.30
GRN        0.0162       0.2198     0.0036       0.0054     0.0000       0.0018    555        0.01          0.21
HATA       0.0000       0.0000     0.0000       0.0000     0.0000       0.0000     58        0.00          0.00
IRN        0.0056       0.1379     0.0014       0.0000     0.0014       0.0000    718        0.00          0.13
LDT        0.0000       0.0588     0.0000       0.0000     0.0000       0.0000     51        0.00          0.02
LGd        0.0375       0.0094     0.0038       0.0019     0.0000       0.0000    533        0.03          0.01
LP         0.0130       0.0173     0.0025       0.0019     0.0019       0.0006   1621        0.01          0.01
PA         0.0000       0.0000     0.0094       0.0000     0.0000       0.0000    106        0.00          0.00
PVT        0.0000       0.0260     0.0000       0.0000     0.0000       0.0000     77        0.00          0.00
TR         0.0000       0.0000     0.0000       0.0000     0.0000       0.0000     35        0.00          0.00
VISam      0.0207       0.0385     0.0030       0.0059     0.0000       0.0000    338        0.02          0.02
VISp       0.0368       0.0141     0.0028       0.0042     0.0000       0.0014    707        0.03          0.02
VISpm      0.0311       0.0133     0.0044       0.0000     0.0000       0.0000    225        0.02          0.00

def ci(k, n):
    lo = stats.beta.ppf(0.025, k, n - k + 1) if k > 0 else 0.0
    hi = stats.beta.ppf(0.975, k + 1, n - k)
    return lo, hi

for var in ("stim", "choice"):
    k, n = int(cells[f"{var}_null"].sum()), len(cells)
    lo, hi = ci(k, n)
    print(f"{var}: null run selects {k} of {n} units ({k / n:.2e}), 95% interval {lo:.2e} to {hi:.2e}; "
          f"paper's rate 5e-5, band 1.7e-5 to 1.5e-4; interval excludes the band: {lo > 1.5e-4 or hi < 1.7e-5}")
    print(f"   null units passing the simple test alone: {cells[f'{var}_null_mw'].mean():.2e}; the combined test alone: {cells[f'{var}_null_cc'].mean():.3f}")
    k2 = int(cells[f"{var}_shuf"].sum())
    lo2, hi2 = ci(k2, n)
    print(f"   labels shuffled over the tested trials (one permutation per region-session): {k2} of {n} units ({k2 / n:.2e}), "
          f"95% interval {lo2:.2e} to {hi2:.2e}; simple test alone {cells[f'{var}_shuf_mw'].mean():.2e}, combined test alone "
          f"{cells[f'{var}_shuf_cc'].mean():.3f}; product of the two single rates {cells[f'{var}_shuf_mw'].mean() * cells[f'{var}_shuf_cc'].mean():.1e}")
print("\nStimulus: largest regional fraction", round(frac.stim_real.max(), 4), "in", frac.stim_real.idxmax(),
      "; pooled over the 14 regions", round(cells.stim_real.mean(), 4))
agree = (frac.stim_real - frac.paper_stim).abs() <= 0.02
print("Regions within 0.02 of the paper's fraction:", int(agree.sum()), "of", len(frac))
print("Paper, pooled over the canonical regions (from region_info.csv):",
      round(float((region_info[region_info.canonical].stim_man * region_info[region_info.canonical]["# good neurons"]).sum()
                  / region_info[region_info.canonical]["# good neurons"].sum()), 4))

stim: null run selects 14 of 5228 units (2.68e-03), 95% interval 1.46e-03 to 4.49e-03; paper's rate 5e-5, band 1.7e-5 to 1.5e-4; interval excludes the band: True
   null units passing the simple test alone: 7.08e-03; the combined test alone: 0.069
   labels shuffled over the tested trials (one permutation per region-session): 4 of 5228 units (7.65e-04), 95% interval 2.09e-04 to 1.96e-03; simple test alone 9.56e-04, combined test alone 0.050; product of the two single rates 4.7e-05
choice: null run selects 13 of 5228 units (2.49e-03), 95% interval 1.32e-03 to 4.25e-03; paper's rate 5e-5, band 1.7e-5 to 1.5e-4; interval excludes the band: True
   null units passing the simple test alone: 5.16e-03; the combined test alone: 0.061
   labels shuffled over the tested trials (one permutation per region-session): 3 of 5228 units (5.74e-04), 95% interval 1.18e-04 to 1.68e-03; simple test alone 5.74e-04, combined test alone 0.058; product of the two single rates 3.3e-05

Stimulus: largest regional fraction 0.0588 in CL ; pooled over the 14 regions 0.0182
Regions within 0.02 of the paper's fraction: 14 of 14
Paper, pooled over the canonical regions (from region_info.csv): 0.0041

only_c = int((cells.choice_real & ~cells.stim_real).sum())
only_s = int((cells.stim_real & ~cells.choice_real).sum())
bt = stats.binomtest(only_c, only_c + only_s, 0.5, alternative="greater") if only_c + only_s else None
print(f"Pooled: choice {cells.choice_real.mean():.4f}, stimulus {cells.stim_real.mean():.4f}; selective for choice only {only_c}, "
      f"stimulus only {only_s}, binomial p = {bt.pvalue if bt else float('nan'):.2g}")
print("Spearman correlation of regional choice fractions with the paper's:", stats.spearmanr(frac.choice_real, frac.paper_choice))
print(frac[["units", "choice_real", "paper_choice"]].sort_values("choice_real", ascending=False).round(4).to_string())

Pooled: choice 0.0564, stimulus 0.0182; selective for choice only 270, stimulus only 70, binomial p = 3.8e-29
Spearman correlation of regional choice fractions with the paper's: SignificanceResult(statistic=np.float64(0.8408608989062024), pvalue=np.float64(0.00016491202306419137))
        units  choice_real  paper_choice
region                                  
CL         17       0.3529          0.30
GRN       555       0.2198          0.21
IRN       718       0.1379          0.13
LDT        51       0.0588          0.02
VISam     338       0.0385          0.02
PVT        77       0.0260          0.00
ACAv      187       0.0214          0.01
LP       1621       0.0173          0.01
VISp      707       0.0141          0.02
VISpm     225       0.0133          0.00
LGd       533       0.0094          0.01
HATA       58       0.0000          0.00
PA        106       0.0000          0.00
TR         35       0.0000          0.00

t0 = time.time()
traj_res, traj_all, traj_long = {}, {}, {}
for k, region in enumerate(TRAJ):
    sess, sess_all = [], []
    for aid, g in sel[sel.region == region].groupby("asset_id"):
        e = by_asset[aid]
        cols = np.isin(e["traj_units"], g.unit.values)
        t = e["trials"]
        a = t.analysed.values
        item = (e["traj"][a][:, cols], t.side.values[a], t.p_left.values[a], t.choice.values[a])
        sess_all.append(item)
        if t_eid[aid] in BLOCK_OK:
            sess.append(item)
    short = [(b[:, :, :bwm.TRAJ_TIMES.size],) + tuple(rest) for b, *rest in sess]
    traj_res[region] = bwm.trajectory(short, bwm.N_PERM_TRAJ, np.random.default_rng(3000 + k))
    traj_all[region] = bwm.trajectory([(b[:, :, :bwm.TRAJ_TIMES.size],) + tuple(rest) for b, *rest in sess_all], bwm.N_PERM_TRAJ,
                                      np.random.default_rng(3100 + k))
    traj_long[region] = bwm.trajectory(sess, bwm.N_PERM_TRAJ, np.random.default_rng(3200 + k), bwm.TRAJ_TIMES_LONG)
def as_table(res):
    df = pd.DataFrame({r: {k: v for k, v in d.items() if k != "d"} for r, d in res.items()}).T
    for c in ("latency", "latency_70", "latency_70_bin_start", "latency_canonical"):
        df[c] = (df[c].astype(float) * 1000).round(1)
    return df
traj_df = as_table(traj_res)
print("Sessions left out for invalid blocks:", {r: int(sel[(sel.region == r) & ~sel.eid.isin(BLOCK_OK)].eid.nunique()) for r in TRAJ})
print(f"Trajectories of {len(TRAJ)} regions with {bwm.N_PERM_TRAJ} permutations, three variants, in {time.time() - t0:.0f} s")
print("latency_70: first time above 70% of the range (Methods); latency: also significant point-wise (code); ms on the code's axis;")
print("latency_canonical: the code's canonical 0 to 100 ms window (first 48 points, full-window baseline) with significance")
print(traj_df.to_string())
print("\nVariant with the 64 sessions of invalid blocks included (permutations within their recorded blocks):")
print(as_table(traj_all)[["n_cells", "amp", "p", "latency_70", "latency"]].to_string())
print("\nVariant with the window 0 to 200 ms (96 points, axis linspace(0, 0.2, 96)):")
print(as_table(traj_long)[["n_cells", "amp", "p", "latency_70", "latency"]].to_string())
print("\nDistance (spikes/s, minimum subtracted) every 10th point of the 0 to 200 ms window, by time (ms):")
print(pd.DataFrame({r: d["d"][::10] for r, d in traj_long.items()}, index=(bwm.TRAJ_TIMES_LONG[::10] * 1000).round(1)).round(2).T.to_string())
fig, ax = plt.subplots(figsize=(6, 3.5))
for r, d in traj_res.items():
    ax.plot(bwm.TRAJ_TIMES * 1000, d["d"], label=r)
ax.set_xlabel("time from stimulus onset (ms)")
ax.set_ylabel("distance (spikes/s)")
ax.legend(fontsize=7)
print("Plotted distances range from 0 to", round(max(d["d"].max() for d in traj_res.values()), 2))

Sessions left out for invalid blocks: {'LGd': 2, 'VISp': 6, 'LP': 6, 'VISpm': 1, 'VISam': 1, 'IRN': 2, 'GRN': 0}
Trajectories of 7 regions with 1000 permutations, three variants, in 40 s
latency_70: first time above 70% of the range (Methods); latency: also significant point-wise (code); ms on the code's axis;
latency_canonical: the code's canonical 0 to 100 ms window (first 48 points, full-window baseline) with significance
       n_cells        amp         p  latency  latency_70  latency_70_bin_start  latency_canonical
LGd      533.0   3.561933  0.000999     31.7        31.7                  31.0               31.9
VISp     707.0   3.645451  0.000999     46.5        46.5                  45.5               46.8
LP      1621.0   1.777594  0.000999     48.6        48.6                  47.5               48.9
VISpm    225.0   3.564427  0.000999     61.3        61.3                  60.0               61.7
VISam    338.0   2.906685  0.000999     59.2        59.2                  58.0               59.6
IRN      718.0  10.076835  0.000999    118.3       118.3                 116.5                NaN
GRN      555.0  14.319785  0.000999    109.9       109.9                 108.0              100.0

Variant with the 64 sessions of invalid blocks included (permutations within their recorded blocks):
       n_cells        amp         p  latency_70  latency
LGd      551.0   3.514707  0.000999        31.7     31.7
VISp     844.0   3.216251  0.000999        46.5     46.5
LP      1780.0   1.673447  0.000999        48.6     48.6
VISpm    240.0   3.427658  0.000999        61.3     61.3
VISam    363.0   2.753979  0.000999        59.2     59.2
IRN      739.0   9.896538  0.000999       118.3    118.3
GRN      555.0  14.319785  0.000999       109.9    109.9

Variant with the window 0 to 200 ms (96 points, axis linspace(0, 0.2, 96)):
       n_cells        amp         p  latency_70  latency
LGd      533.0   3.561933  0.000999        31.6     31.6
VISp     707.0   3.645451  0.000999        46.3     46.3
LP      1621.0   1.777594  0.000999        48.4     48.4
VISpm    225.0   3.564428  0.000999        61.1     61.1
VISam    338.0   2.906685  0.000999        58.9     58.9
IRN      718.0  10.350676  0.000999       117.9    117.9
GRN      555.0  14.534202  0.000999       109.5    109.5

Distance (spikes/s, minimum subtracted) every 10th point of the 0 to 200 ms window, by time (ms):
       0.0    21.1   42.1   63.2   84.2   105.3  126.3  147.4  168.4  189.5
LGd     0.03   0.46   3.53   2.16   1.65   1.26   1.19   1.18   1.55   1.75
VISp    0.12   0.29   2.11   3.35   2.53   1.97   1.63   1.68   1.92   1.95
LP      0.10   0.00   0.99   1.65   1.22   1.11   1.02   1.46   1.48   1.40
VISpm   0.18   0.48   0.38   2.91   3.46   2.69   2.24   1.57   1.25   1.59
VISam   0.32   0.13   0.40   2.32   2.58   2.09   1.16   1.79   1.93   1.64
IRN     0.53   0.99   0.65   0.78   2.44   5.34   8.49   9.68   9.95  10.22
GRN     0.00   0.32   0.05   1.88   5.21   9.56  12.26  14.06  14.22  13.21
Plotted distances range from 0 to 14.32

exp = {"LGd": 34, "VISp": 42, "LP": 42, "VISpm": 57, "VISam": 78}
for col in ("latency_70", "latency", "latency_canonical"):
    L = traj_df[col].astype(float)
    order = (L["LGd"] < min(L["VISp"], L["LP"]) and max(L["VISp"], L["LP"]) < L["VISpm"] < L["VISam"]
             and abs(L["VISp"] - L["LP"]) <= 10)
    within = {r: bool(abs(L[r] - v) <= 0.1 * v) for r, v in exp.items()}
    print(f"{col:11s}", {r: L[r] for r in exp}, "order holds:", order, "within 10%:", within)
L = as_table(traj_all)["latency_70"].astype(float)
print("Variant with the 64 sessions included, latency_70:", {r: L[r] for r in exp}, "order holds:",
      bool(L["LGd"] < min(L["VISp"], L["LP"]) and max(L["VISp"], L["LP"]) < L["VISpm"] < L["VISam"] and abs(L["VISp"] - L["LP"]) <= 10),
      "within 10%:", {r: bool(abs(L[r] - v) <= 0.1 * v) for r, v in exp.items()})
print("Significance of the distance (permutation p):", traj_df.loc[list(exp), "p"].to_dict())
paper_tab = pd.read_pickle(io.BytesIO(requests.get(f"{REPO}/brainwidemap/meta/Results/stim.pkl", timeout=60).content))
print("The paper's results table meta/Results/stim.pkl, euclidean_latency (ms):",
      (paper_tab.set_index("region").loc[list(exp) + ["IRN", "GRN"], "euclidean_latency"] * 1000).round(1).to_dict())

latency_70  {'LGd': np.float64(31.7), 'VISp': np.float64(46.5), 'LP': np.float64(48.6), 'VISpm': np.float64(61.3), 'VISam': np.float64(59.2)} order holds: False within 10%: {'LGd': True, 'VISp': False, 'LP': False, 'VISpm': True, 'VISam': False}
latency     {'LGd': np.float64(31.7), 'VISp': np.float64(46.5), 'LP': np.float64(48.6), 'VISpm': np.float64(61.3), 'VISam': np.float64(59.2)} order holds: False within 10%: {'LGd': True, 'VISp': False, 'LP': False, 'VISpm': True, 'VISam': False}
latency_canonical {'LGd': np.float64(31.9), 'VISp': np.float64(46.8), 'LP': np.float64(48.9), 'VISpm': np.float64(61.7), 'VISam': np.float64(59.6)} order holds: False within 10%: {'LGd': True, 'VISp': False, 'LP': False, 'VISpm': True, 'VISam': False}
Variant with the 64 sessions included, latency_70: {'LGd': np.float64(31.7), 'VISp': np.float64(46.5), 'LP': np.float64(48.6), 'VISpm': np.float64(61.3), 'VISam': np.float64(59.2)} order holds: False within 10%: {'LGd': True, 'VISp': False, 'LP': False, 'VISpm': True, 'VISam': False}
Significance of the distance (permutation p): {'LGd': 0.000999000999000999, 'VISp': 0.000999000999000999, 'LP': 0.000999000999000999, 'VISpm': 0.000999000999000999, 'VISam': 0.000999000999000999}

The paper's results table meta/Results/stim.pkl, euclidean_latency (ms): {'LGd': 34.0, 'VISp': 44.7, 'LP': 46.8, 'VISpm': 59.6, 'VISam': 66.0, 'IRN': 89.4, 'GRN': 89.4}

for r in ("IRN", "GRN"):
    d = traj_df.loc[r]
    print(f"{r}: latency (Methods) {d.latency_70} ms, with significance (code) {d.latency} ms, distance p = {d.p:.3g}, "
          f"{int(d.n_cells)} cells; in 90 to 132 ms: {90 <= d.latency_70 <= 132}")
long_df = as_table(traj_long)
for r in ("IRN", "GRN"):
    print(f"{r}, variant 0 to 200 ms: latency (Methods) {long_df.loc[r, 'latency_70']} ms, with significance {long_df.loc[r, 'latency']} ms; "
          f"canonical 0 to 100 ms window with significance {traj_df.loc[r, 'latency_canonical']} ms")

IRN: latency (Methods) 118.3 ms, with significance (code) 118.3 ms, distance p = 0.000999, 718 cells; in 90 to 132 ms: True
GRN: latency (Methods) 109.9 ms, with significance (code) 109.9 ms, distance p = 0.000999, 555 cells; in 90 to 132 ms: True
IRN, variant 0 to 200 ms: latency (Methods) 117.9 ms, with significance 117.9 ms; canonical 0 to 100 ms window with significance nan ms
GRN, variant 0 to 200 ms: latency (Methods) 109.5 ms, with significance 109.5 ms; canonical 0 to 100 ms window with significance 100.0 ms

def region_table(sq, ncell):
    rows = []
    for reg in canon_regions:
        n_c = ncell.get(reg, 0)
        st = bwm.traj_stats(sq[reg], n_c) if n_c >= 20 else {}
        rows.append({"region": reg, "n_cells": n_c, "amp": st.get("amp", np.nan), "p": st.get("p", np.nan),
                     "latency_70": st.get("latency_70", np.nan) * 1000})
    df = pd.DataFrame(rows).set_index("region")
    tested = df.p.notna()
    for name, fn, q in (("bh_01", bwm.bh, 0.01), ("bh_05", bwm.bh, 0.05), ("by_05", bwm.by, 0.05)):
        df[name] = False
        df.loc[tested, name] = fn(df.p[tested].values, q)
    df["raw_01"] = df.p < 0.01
    df["raw_05"] = df.p < 0.05
    df["paper"] = [bool(ri.loc[r, "stim_euc_sig"]) for r in df.index]
    return df

tables = {k: region_table(*v) for k, v in acc_traj.items()}
traj_main = tables["methods_ok"]
tested = traj_main.p.notna()
print(f"Canonical regions {len(traj_main)}; tested (20 or more units in sessions with valid blocks) {int(tested.sum())}; "
      f"not tested {sorted(traj_main.index[~tested])}")
print(f"Smallest attainable p: {1 / (bwm.N_PERM_TRAJ + 1):.4f}; regions at it: {int((traj_main.p <= 1 / (bwm.N_PERM_TRAJ + 1) + 1e-12).sum())}")
print(f"Paper: 104 (region_info.csv stim_euc_sig marks {int(traj_main.paper.sum())} of these {len(traj_main)} regions)")
names = {"methods_ok": "Methods' classes, valid-block sessions (main)", "methods_all": "Methods' classes, all 454 sessions",
         "code_ok": "code's classes, valid-block sessions", "code_all": "code's classes, all 454 sessions"}
labels = {"bh_01": "(a) BH q = 0.01", "raw_01": "(b) P < 0.01", "by_05": "code, BY at 0.05", "bh_05": "BH q = 0.05", "raw_05": "P < 0.05"}
rows = []
for k, df in tables.items():
    for c, lab in labels.items():
        s_ = df[c]
        rows.append({"permutations and sessions": names[k], "threshold": lab, "significant": int(s_.sum()),
                     "within 104 +/- 10.4": abs(int(s_.sum()) - 104) <= 10.4, "agree with paper": int((s_ == df.paper).sum()),
                     "both": int((s_ & df.paper).sum()), "here only": int((s_ & ~df.paper).sum()), "paper only": int((~s_ & df.paper).sum())})
print(pd.DataFrame(rows).to_string(index=False))
chk = pd.DataFrame({"amp_all_regions_pass": traj_main.loc[TRAJ, "amp"], "amp_trajectory_section": traj_df.loc[TRAJ, "amp"].astype(float)})
print("Amplitude of the seven trajectory regions in this pass and in the section above (same cells and trials):")
print(chk.round(4).T.to_string())
print("Median p over the tested regions: Methods' classes", round(float(traj_main.p.median()), 4), "; code's classes, all sessions",
      round(float(tables["code_all"].p.median()), 4))
print("Regions significant in the paper's column and not under (a), with their p here:",
      traj_main[traj_main.paper & ~traj_main.bh_01].p.round(4).to_dict())

Canonical regions 201; tested (20 or more units in sessions with valid blocks) 195; not tested ['CL', 'LING', 'MH', 'P5', 'PC5', 'SOC']
Smallest attainable p: 0.0010; regions at it: 42
Paper: 104 (region_info.csv stim_euc_sig marks 104 of these 201 regions)
                    permutations and sessions        threshold  significant  within 104 +/- 10.4  agree with paper  both  here only  paper only
Methods' classes, valid-block sessions (main)  (a) BH q = 0.01           46                False               133    41          5          63
Methods' classes, valid-block sessions (main)     (b) P < 0.01           62                False               137    51         11          53
Methods' classes, valid-block sessions (main) code, BY at 0.05           46                False               133    41          5          63
Methods' classes, valid-block sessions (main)      BH q = 0.05           76                False               139    59         17          45
Methods' classes, valid-block sessions (main)         P < 0.05           91                False               136    65         26          39
           Methods' classes, all 454 sessions  (a) BH q = 0.01           51                False               134    44          7          60
           Methods' classes, all 454 sessions     (b) P < 0.01           70                False               141    57         13          47
           Methods' classes, all 454 sessions code, BY at 0.05           51                False               134    44          7          60
           Methods' classes, all 454 sessions      BH q = 0.05           84                False               143    65         19          39
           Methods' classes, all 454 sessions         P < 0.05           96                 True               147    73         23          31
         code's classes, valid-block sessions  (a) BH q = 0.01           47                False               132    41          6          63
         code's classes, valid-block sessions     (b) P < 0.01           67                False               138    54         13          50
         code's classes, valid-block sessions code, BY at 0.05           47                False               132    41          6          63
         code's classes, valid-block sessions      BH q = 0.05           77                False               138    59         18          45
         code's classes, valid-block sessions         P < 0.05           95                 True               140    69         26          35
             code's classes, all 454 sessions  (a) BH q = 0.01           53                False               130    43         10          61
             code's classes, all 454 sessions     (b) P < 0.01           76                False               145    62         14          42
             code's classes, all 454 sessions code, BY at 0.05           53                False               130    43         10          61
             code's classes, all 454 sessions      BH q = 0.05           81                False               142    63         18          41
             code's classes, all 454 sessions         P < 0.05           98                 True               151    76         22          28
Amplitude of the seven trajectory regions in this pass and in the section above (same cells and trials):
                           LGd    VISp      LP   VISpm   VISam      IRN      GRN
amp_all_regions_pass    3.5619  3.6455  1.7776  3.5644  2.9067  10.0768  14.3198
amp_trajectory_section  3.5619  3.6455  1.7776  3.5644  2.9067  10.0768  14.3198
Median p over the tested regions: Methods' classes 0.0569 ; code's classes, all sessions 0.0549
Regions significant in the paper's column and not under (a), with their p here: {'AIv': 0.7872, 'ANcr2': 0.0739, 'AON': 0.0569, 'AUDd': 0.007, 'AUDv': 0.3077, 'AV': 0.039, 'CEA': 0.5634, 'CENT3': 0.003, 'CL': nan, 'CLA': 0.8322, 'CUL4 5': 0.1748, 'ENTl': 0.2208, 'Eth': 0.1119, 'GPi': 0.035, 'IC': 0.009, 'IP': 0.014, 'LAV': 0.033, 'LD': 0.5295, 'LDT': 0.018, 'LRN': 0.4306, 'LSv': 0.5425, 'MDRN': 0.006, 'MG': 0.5395, 'MOp': 0.0509, 'MS': 0.0699, 'MV': 0.025, 'NB': 0.019, 'NPC': 0.004, 'OP': 0.3017, 'ORBm': 0.4755, 'PAR': 0.3197, 'PGRN': 0.018, 'POST': 0.012, 'PPT': 0.2877, 'PRE': 0.0679, 'PRM': 0.1299, 'PRP': 0.018, 'PSV': 0.0879, 'PeF': 0.0589, 'ProS': 0.0529, 'RN': 0.006, 'RPF': 0.0939, 'RSPagl': 0.013, 'RSPv': 0.0619, 'SCs': 0.004, 'SI': 0.0519, 'SIM': 0.0759, 'SMT': 0.1808, 'SNc': 0.005, 'SOC': nan, 'SPF': 0.0529, 'SPIV': 0.029, 'SPVI': 0.011, 'SSp-ul': 0.1479, 'SUV': 0.1818, 'TEa': 0.006, 'V': 0.1459, 'VAL': 0.2418, 'VCO': 0.02, 'VISa': 0.2537, 'VISal': 0.1229, 'VM': 0.003, 'VPL': 0.0839}

rows = []
lick_by_asset = {s["asset_id"]: np.sort(s["licks"]) for s in scans}
for aid, g in sel[sel.region == "IRN"].groupby("asset_id"):
    e = by_asset[aid]
    t = e["trials"]
    pop = e["psth"].mean(0)          # 2 x bins: correct, incorrect
    f_c, p_c, freqs = bwm.spectrum_peak(pop[0], bwm.PSTH_BIN)
    _, p_i, _ = bwm.spectrum_peak(pop[1], bwm.PSTH_BIN)
    band = np.abs(freqs - f_c) <= 1
    in_band = (freqs >= 4) & (freqs <= 20)
    # Variants: each unit's PSTH z-scored over both conditions; power as a fraction of the 4 to 20 Hz power
    z = e["psth"]
    z = (z - z.mean((1, 2), keepdims=True)) / np.where(z.std((1, 2), keepdims=True) > 0, z.std((1, 2), keepdims=True), 1)
    zf_c, zp_c, _ = bwm.spectrum_peak(z.mean(0)[0], bwm.PSTH_BIN)
    _, zp_i, _ = bwm.spectrum_peak(z.mean(0)[1], bwm.PSTH_BIN)
    zband = np.abs(freqs - zf_c) <= 1
    lk = lick_by_asset[aid]
    fb = t.feedback_time[t.analysed & t.correct].values
    edges = fb[:, None] + np.arange(0, bwm.PSTH_POST, bwm.PSTH_BIN)[None, :]
    lick_psth = bwm.window_counts(lk, edges, bwm.PSTH_BIN).mean(0) / bwm.PSTH_BIN if len(lk) else np.full(edges.shape[1], np.nan)
    f_l = bwm.spectrum_peak(lick_psth, bwm.PSTH_BIN)[0] if len(lk) else np.nan
    # Variants for (ii) that do not use the correct-trial PSTH to pick the band: the band within 1 Hz of
    # the session's lick peak, and the paper's 9 to 11 Hz; each on the raw PSTHs and on the PSTHs with
    # each condition's own slow evoked response removed (the PSTH minus its Gaussian smoothing,
    # sigma 60 ms, which passes 3% of the amplitude at 7 Hz). The last compares each condition at its own peak.
    hp = [x - ndimage.gaussian_filter1d(x, 0.06 / bwm.PSTH_BIN, mode="nearest") for x in pop]
    hf_c, hp_c, _ = bwm.spectrum_peak(hp[0], bwm.PSTH_BIN)
    hf_i, hp_i, _ = bwm.spectrum_peak(hp[1], bwm.PSTH_BIN)
    lband = np.abs(freqs - f_l) <= 1
    pband = (freqs >= 9) & (freqs <= 11)
    nolick = np.nan if not len(lk) else 1.0
    rows.append({"session": aid[:8], "units": len(g), "peak_correct_hz": f_c, "power_correct": p_c[band].sum(),
                 "power_incorrect": p_i[band].sum(), "licks": len(lk), "lick_peak_hz": f_l,
                 "error_trials": int((t.analysed & ~t.correct).sum()),
                 "rel_correct": p_c[band].sum() / p_c[in_band].sum(), "rel_incorrect": p_i[band].sum() / p_i[in_band].sum(),
                 "z_peak_hz": zf_c, "z_power_correct": zp_c[zband].sum(), "z_power_incorrect": zp_i[zband].sum(),
                 "lick_c": nolick * p_c[lband].sum(), "lick_i": nolick * p_i[lband].sum(),
                 "paper_c": p_c[pband].sum(), "paper_i": p_i[pband].sum(),
                 "hp_lick_c": nolick * hp_c[lband].sum(), "hp_lick_i": nolick * hp_i[lband].sum(),
                 "hp_paper_c": hp_c[pband].sum(), "hp_paper_i": hp_i[pband].sum(),
                 "hp_peak_c_hz": hf_c, "hp_peak_i_hz": hf_i,
                 "hp_own_c": hp_c[np.abs(freqs - hf_c) <= 1].sum(), "hp_own_i": hp_i[np.abs(freqs - hf_i) <= 1].sum()})
osc = pd.DataFrame(rows)
print(osc.drop(columns=[k for k in osc.columns if k.startswith(("lick_c", "lick_i", "paper_", "hp_"))]).round(2).to_string(index=False))
w = stats.wilcoxon(osc.power_correct, osc.power_incorrect)
print(f"Median peak frequency, correct trials: {osc.peak_correct_hz.median():.2f} Hz; power at the peak higher on correct trials in "
      f"{int((osc.power_correct > osc.power_incorrect).sum())} of {len(osc)} sessions, Wilcoxon p = {w.pvalue:.2g}")
print(f"Sessions whose correct-trial peak is at the 4 Hz edge of the band (below 4.1 Hz): {int((osc.peak_correct_hz < 4.1).sum())} of {len(osc)}; "
      f"peak frequencies range {osc.peak_correct_hz.min():.2f} to {osc.peak_correct_hz.max():.2f} Hz")
for name, a_, b_ in (("power as a fraction of the 4 to 20 Hz power", "rel_correct", "rel_incorrect"),
                     ("PSTHs z-scored per unit", "z_power_correct", "z_power_incorrect")):
    w_ = stats.wilcoxon(osc[a_], osc[b_])
    print(f"Variant, {name}: higher on correct trials in {int((osc[a_] > osc[b_]).sum())} of {len(osc)} sessions, Wilcoxon p = {w_.pvalue:.2g}; "
          f"median ratio correct / incorrect {np.median(osc[a_] / osc[b_]):.2f}")
print(f"Variant, PSTHs z-scored per unit: median peak frequency {osc.z_peak_hz.median():.2f} Hz")
print(f"Median lick PSTH peak: {osc.lick_peak_hz.median():.2f} Hz over {int(osc.lick_peak_hz.notna().sum())} sessions with licks")
print("Variants for (ii) with the band chosen without the correct-trial PSTH, and with each condition's slow evoked response removed:")
for name, a_, b_ in (("raw PSTH, within 1 Hz of the session's lick peak", "lick_c", "lick_i"),
                     ("raw PSTH, 9 to 11 Hz", "paper_c", "paper_i"),
                     ("evoked response removed, within 1 Hz of the lick peak", "hp_lick_c", "hp_lick_i"),
                     ("evoked response removed, 9 to 11 Hz", "hp_paper_c", "hp_paper_i"),
                     ("evoked response removed, each condition at its own peak", "hp_own_c", "hp_own_i")):
    d_ = osc[[a_, b_]].dropna()
    w_ = stats.wilcoxon(d_[a_], d_[b_])
    print(f"   {name}: higher on correct trials in {int((d_[a_] > d_[b_]).sum())} of {len(d_)} sessions, Wilcoxon p = {w_.pvalue:.2g}; "
          f"median ratio correct / incorrect {np.median(d_[a_] / d_[b_]):.2f}")
print(f"   Evoked response removed: median peak frequency {osc.hp_peak_c_hz.median():.2f} Hz on correct trials "
      f"(range {osc.hp_peak_c_hz.min():.2f} to {osc.hp_peak_c_hz.max():.2f}), {osc.hp_peak_i_hz.median():.2f} Hz on incorrect trials; "
      f"sessions whose correct peak is within 1 Hz of the lick peak: {int((np.abs(osc.hp_peak_c_hz - osc.lick_peak_hz) <= 1).sum())} of {int(osc.lick_peak_hz.notna().sum())}")

 session  units  peak_correct_hz  power_correct  power_incorrect  licks  lick_peak_hz  error_trials  rel_correct  rel_incorrect  z_peak_hz  z_power_correct  z_power_incorrect
223c5ed1     23             8.15      289991.69          6774.49  83375          8.06            86         0.80           0.06       8.15           967.90             254.15
2826fe96      9             6.88       69258.77          2466.56  37370          6.93            57         0.69           0.05      13.53           237.38              82.00
2a95ae6f     45             7.23      973756.52         18877.06  39281          7.32            65         0.74           0.31       7.08          1053.27              57.24
31cf56c6     41             4.00        2245.06          3705.58  41906          4.00            75         0.50           0.43       4.00            90.55              64.34
378f1572     12             4.00       31557.91         67952.09  28208          7.23            59         1.02           0.55       7.81           432.78            1344.04
4c20f70a     29             5.96       12227.82          4132.85  10422          4.49            61         0.58           0.10       5.91            46.85              45.57
4ffe4a8c      6            10.06        1674.74         51775.43  45102          7.62            79         0.37           0.51       9.47           139.27            1562.80
533bfe1c     11             7.23      167263.41         16588.22  68715          7.28            70         0.79           0.11       7.23          1982.79             194.81
54a0a5cd     38             6.88      526880.51           566.72  72262          6.88           104         0.81           0.04       6.84           714.68              21.59
5a66e6f2      7             8.35      344248.80         67786.53  72661          8.15            38         0.81           0.13       8.30          2345.09             462.57
5c0e9a3d     29             7.96        6446.27         11146.54  67697          7.96            61         0.52           0.44       8.11           123.59             104.46
5f5fbdaa     11             9.47       16480.33          8320.07    133          4.25            22         0.37           0.05       9.38            69.89             196.53
6853c903     47             6.54        7083.24          1547.69  81102          6.54           117         0.75           0.20       6.59           148.00              12.08
6b18ebcf     28             4.25        2085.77          2929.05  54766          9.03            67         0.69           0.38       4.00            79.32             149.22
729d63fd     11             5.91     1330593.57          3571.52  37125          6.05            57         0.84           0.17       6.01          5529.95             273.13
7b5ef954     25            15.62       12979.33         23044.80  12259         13.43            72         0.40           0.12      19.97           106.58             112.29
8b5345b8     31             7.28       88595.10          7430.71  71193          7.08            78         0.81           0.19       7.28           214.62              68.69
8b73f33f     42             4.00        4122.61          3852.30  72601          7.23            87         0.70           0.17       8.50            94.16             144.91
95259d76     39             6.10       18212.44          1056.90  48835          6.20            99         0.82           0.10       5.91           105.56              14.15
972c258e      8             6.15      702626.74          3143.89  47114          6.15            61         0.82           0.07       5.91           432.41              73.66
999a7b5e     73             5.32        3153.89           938.15 100473          8.74            93         0.34           0.16       5.81            26.72              73.48
a79329d2     29             7.96        4826.22           747.17 106741          8.35            78         0.55           0.08       8.30           323.40              91.44
b816bcd5     35             4.00        2285.24          2229.81 157145          7.42            83         0.33           0.24       7.42            69.45              70.17
cc1b0a9f     68             7.08       51194.20          1259.31   8438          5.37            98         0.66           0.22       7.08           200.43              50.95
d124f0db     15             8.79       21063.27         10307.43      0           NaN            68         0.33           0.10       9.13           290.81              72.29
d9da8c43     18             7.42        5996.07         22631.90   8384          7.18            48         0.31           0.14       6.98            44.08             210.10
e7597b80      9             8.11       15079.40          1212.19  86474          8.11           174         0.68           0.11       7.86           250.58             124.18
Median peak frequency, correct trials: 7.08 Hz; power at the peak higher on correct trials in 20 of 27 sessions, Wilcoxon p = 0.0042
Sessions whose correct-trial peak is at the 4 Hz edge of the band (below 4.1 Hz): 4 of 27; peak frequencies range 4.00 to 15.62 Hz
Variant, power as a fraction of the 4 to 20 Hz power: higher on correct trials in 26 of 27 sessions, Wilcoxon p = 1e-07; median ratio correct / incorrect 4.11
Variant, PSTHs z-scored per unit: higher on correct trials in 18 of 27 sessions, Wilcoxon p = 0.028; median ratio correct / incorrect 2.89
Variant, PSTHs z-scored per unit: median peak frequency 7.28 Hz
Median lick PSTH peak: 7.23 Hz over 26 sessions with licks
Variants for (ii) with the band chosen without the correct-trial PSTH, and with each condition's slow evoked response removed:
   raw PSTH, within 1 Hz of the session's lick peak: higher on correct trials in 15 of 26 sessions, Wilcoxon p = 0.038; median ratio correct / incorrect 4.13
   raw PSTH, 9 to 11 Hz: higher on correct trials in 12 of 27 sessions, Wilcoxon p = 0.7; median ratio correct / incorrect 0.91
   evoked response removed, within 1 Hz of the lick peak: higher on correct trials in 15 of 26 sessions, Wilcoxon p = 0.031; median ratio correct / incorrect 4.50
   evoked response removed, 9 to 11 Hz: higher on correct trials in 12 of 27 sessions, Wilcoxon p = 0.66; median ratio correct / incorrect 0.91
   evoked response removed, each condition at its own peak: higher on correct trials in 16 of 27 sessions, Wilcoxon p = 0.077; median ratio correct / incorrect 1.55
   Evoked response removed: median peak frequency 7.42 Hz on correct trials (range 4.35 to 15.62), 7.91 Hz on incorrect trials; sessions whose correct peak is within 1 Hz of the lick peak: 17 of 26

per_session = extract_seconds / len(extracted)
per_decode = decode_seconds * 4 / len(dec_jobs) / 3            # CPU seconds per variable per region-session
n_region_sessions = canon.groupby(["region", "eid"]).ngroups
print(f"Streaming the spikes of all {canon.eid.nunique()} canonical sessions at {per_session:.2f} s each: about {canon.eid.nunique() * per_session / 60:.0f} min")
print(f"Decoding one variable for all {n_region_sessions} region-sessions as here (1 run, 100 nulls with reused C): about "
      f"{n_region_sessions * per_decode / 4 / 60:.0f} min on 4 processes; with the paper's 10 runs and 200 nested-search nulls the "
      f"null fits alone grow by more than 10 x 2 x 35 times")
print(f"Streaming all {len(canon)} canonical units of {len(canon_jobs)} sessions with their trajectories and clock check took {all_seconds / 60:.1f} min")
print(f"Total run time of this notebook so far: {(time.time() - started) / 60:.1f} min of the 30 min limit")

Streaming the spikes of all 454 canonical sessions at 0.42 s each: about 3 min
Decoding one variable for all 2458 region-sessions as here (1 run, 100 nulls with reused C): about 41 min on 4 processes; with the paper's 10 runs and 200 nested-search nulls the null fits alone grow by more than 10 x 2 x 35 times
Streaming all 62857 canonical units of 454 sessions with their trajectories and clock check took 11.8 min
Total run time of this notebook so far: 25.2 min of the 30 min limit

Claim	Tested
dataset-sessions-mice-labs: 459 sessions, 139 mice, 12 labs	yes
dataset-probes-units: 699 insertions, 621,733 units	yes
well-isolated-neurons: 75,708 well-isolated neurons	yes
canonical-regions: 201 regions, 62,857 neurons	yes
brain-areas-279: 279 brain areas	yes
mice-sex: 94 male, 45 female	yes
mice-age-weight: 13 to 178 weeks, 16.1 to 35.7 g	yes
probe-versions: 109 3A, 586 3B2, 4 2.4	yes
trials-per-session: mean 645, median 602, 401 to 1,525	yes
rf-mapping: receptive fields (Supplementary Fig. 5)	blocked, data missing
behav-performance: 81.4% correct	yes
behav-zero-contrast: 58.7% at 0% contrast, t138 = 20.18	yes
behav-early-movements: 22.8% of first movements under 80 ms	yes
fig1d-contrast: better and faster at high contrast	yes
fig1ef-end-of-session: worse and slower at session end	yes
fig1g-block-switch: 5 to 10 trials to adjust	yes
fig1h-block-bias: block biases choices at all contrasts	yes
dec-stim-restricted: stimulus decodable in fewer regions, VISp	yes, 10 regions
dec-choice-widespread: choice in more regions than stimulus, GRN	yes, 10 regions
dec-feedback-nearly-all: feedback in nearly all regions	yes, 10 regions
sc-stim-fraction: few stimulus-selective neurons	yes, 14 regions
sc-choice-fraction: more choice- than stimulus-selective neurons	yes, 14 regions
traj-latency-visual: LGd < VISp = LP < VISpm < VISam	yes
traj-latency-late: IRN, GRN at 100 to 120 ms	yes, 2 of 5 regions
traj-stim-regions: 104 regions with stimulus trajectories	yes, all canonical regions
fback-oscillation-licking: about 10 Hz with licking	yes, IRN
wheel-speed-decoding: 163 of 201 regions	blocked, not attempted
encoding-feedback-largest: correct-feedback kernel largest	blocked, not attempted
video-correlates: half of neurons correlate with video	blocked, not attempted
granger-interactions	blocked, not attempted

Claim	Outcome
dataset-sessions-mice-labs	reproduced
dataset-probes-units	reproduced
well-isolated-neurons	reproduced
canonical-regions	reproduced
brain-areas-279	reproduced
mice-sex	reproduced
mice-age-weight	partially reproduced (weight yes, age no)
probe-versions	reproduced
trials-per-session	reproduced
rf-mapping	blocked, data missing
behav-performance	reproduced (mean over sessions)
behav-zero-contrast	not reproduced (mean 0.44 points below)
behav-early-movements	reproduced
fig1d-contrast	reproduced
fig1ef-end-of-session	reproduced
fig1g-block-switch	reproduced (reading b)
fig1h-block-bias	reproduced
dec-stim-restricted	reproduced
dec-choice-widespread	reproduced
dec-feedback-nearly-all	partially reproduced
sc-stim-fraction	blocked, methods underspecified
sc-choice-fraction	blocked, methods underspecified
traj-latency-visual	not reproduced
traj-latency-late	reproduced
traj-stim-regions	not reproduced (46 regions)
fback-oscillation-licking	partially reproduced
wheel-speed-decoding	blocked, not attempted
encoding-feedback-largest	blocked, not attempted
video-correlates	blocked, not attempted
granger-interactions	blocked, not attempted

Reproducing the IBL brain-wide map (2025) from dandiset 000409¶

Data check¶

dataset-sessions-mice-labs¶

dataset-probes-units¶

well-isolated-neurons¶

canonical-regions¶

brain-areas-279¶

mice-sex¶

mice-age-weight¶

probe-versions¶

trials-per-session¶

rf-mapping (blocked)¶

behav-performance¶

behav-zero-contrast¶

behav-early-movements¶

fig1d-contrast¶

fig1ef-end-of-session¶

fig1g-block-switch¶

fig1h-block-bias¶

Neural data: the regions tested and their spikes¶

dec-stim-restricted¶

dec-choice-widespread¶

dec-feedback-nearly-all¶

Single-cell selectivity (sc-stim-fraction, sc-choice-fraction)¶

sc-stim-fraction¶

sc-choice-fraction¶

Population trajectories (traj-latency-visual, traj-latency-late)¶

traj-latency-visual¶

traj-latency-late¶

traj-stim-regions¶

fback-oscillation-licking¶

Claims over all regions (blocked, not attempted)¶

Summary¶