import hashlib
import io
import time
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import requests
import statsmodels.formula.api as smf
from scipy import stats
from statsmodels.stats.multicomp import pairwise_tukeyhsd

import seiler

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

started = time.time()
sessions = seiler.list_sessions()
print(f"{len(sessions)} assets in dandiset {seiler.DANDISET_ID} version {seiler.DANDISET_VERSION}")

# Each subject's experiment and arm, from its files other than the 'Unknown' DLS files.
named = sessions[(sessions.arm != "Unknown") & sessions.experiment.notna()]
subj_arm = named.groupby("subject").group.agg(lambda g: g.value_counts().index[0])
sessions["subject_group"] = sessions.subject.map(subj_arm)
sessions["subject_experiment"] = sessions.subject_group.str.extract(
    r"^(FP|Opto-DLS-Excitatory|Opto-DMS-Excitatory|Opto-DMS-Inhibitory)")[0]
sessions["subject_arm"] = sessions.subject_group.str.split("-").str[-1]

is_fp = sessions.subject_experiment == "FP"
is_opto = sessions.subject_experiment.str.startswith("Opto", na=False)
need = is_fp | (is_opto & sessions.stage.isin(["shock", "omission"])) | (
    sessions.subject_experiment.isin(["Opto-DMS-Excitatory", "Opto-DLS-Excitatory"]) & sessions.stage.isin(["FR1", "RI30"]))
# one file of every optogenetic subject, for its notes
need |= is_opto & ~sessions.subject.duplicated() & sessions.subject_experiment.notna()
to_load = sessions[need]
t = time.time()
events = seiler.read_all_events(to_load.asset_id.tolist())
print(f"read the events of {len(events)} files in {time.time() - t:.0f} s")

4139 assets in dandiset 000971 version 0.261013.1707

read the events of 2202 files in 1006 s

subjects = sessions[sessions.subject_experiment.notna()].drop_duplicates("subject")
print("Subjects by experiment and arm (arm from the subject's files other than the 'Unknown' DLS files):")
print(pd.crosstab(subjects.subject_group, subjects.sex).to_string())
print("\nWestern blot assets:", int((sessions.group == "western_blot").sum()))
print("Files by experiment and stage:")
print(pd.crosstab(sessions.subject_experiment.fillna("western_blot"), sessions.stage).to_string())
two = sessions[sessions.experiment.notna() & (sessions.arm != "Unknown")].groupby("subject").experiment.nunique()
print("\nSubjects with files in two experiments:", two[two > 1].index.tolist())
for s in two[two > 1].index:
    print(" ", s, sessions[sessions.subject == s].group.value_counts().to_dict())
print("Paper: fiber photometry 36 RI60 mice (9 PR, 9 DPR, 18 PS) and RR20 9 (Fig 1D) or 7 (Methods);"
      " DMS excitatory 39 mice, DMS inhibitory 26 (Methods) or 47 (Fig 6F), DLS excitatory 36.")

Subjects by experiment and arm (arm from the subject's files other than the 'Unknown' DLS files):
sex                                 F  M
subject_group                           
FP-DPR                              4  5
FP-PR                               1  8
FP-PS                              17  1
FP-RR20                             5  1
Opto-DLS-Excitatory-ChR2            5  6
Opto-DLS-Excitatory-ChR2Scrambled   6  6
Opto-DLS-Excitatory-EYFP            6  6
Opto-DMS-Excitatory-ChR2            7  7
Opto-DMS-Excitatory-ChR2Scrambled   6  6
Opto-DMS-Excitatory-EYFP            3  4
Opto-DMS-Inhibitory-EYFP            9  7
Opto-DMS-Inhibitory-NpHR            8  5
Opto-DMS-Inhibitory-NpHRScrambled   6  7

Western blot assets: 16
Files by experiment and stage:
stage                FR1  FR1habit  RI30  RI60  RR20  RR5  omission  other  probetest  shock  unknown
subject_experiment                                                                                   
FP                   231        62    93   640    84    1        23      0         75     97       18
Opto-DLS-Excitatory  218         0    56   538     0    0        39      0          1     83        0
Opto-DMS-Excitatory  172         0    52   449     0    0        35      0         26     68        0
Opto-DMS-Inhibitory  316        31    64   536     0    0        37      0          2     76        0
western_blot           0         0     0     0     0    0         0     16          0      0        0

Subjects with files in two experiments: ['290.407']
  290.407 {'Opto-DMS-Excitatory-ChR2Scrambled': 24, 'Opto-DLS-Excitatory-ChR2': 2, 'Opto-DLS-Excitatory-Unknown': 2}
Paper: fiber photometry 36 RI60 mice (9 PR, 9 DPR, 18 PS) and RR20 9 (Fig 1D) or 7 (Methods); DMS excitatory 39 mice, DMS inhibitory 26 (Methods) or 47 (Fig 6F), DLS excitatory 36.

def event_hash(ev):
    h = hashlib.md5()
    for k in sorted(ev):
        h.update(k.encode())
        h.update(np.round(np.asarray(ev[k]), 3).tobytes())
    return h.hexdigest()

loaded = sessions[sessions.asset_id.isin(list(events))].copy()
loaded["hash"] = [event_hash(events[a]["events"]) if events[a]["events"] else "" for a in loaded.asset_id]
dups = loaded[(loaded.hash != "") & loaded.duplicated("hash", keep=False)].sort_values("hash")
pairs = dups.groupby("hash").agg(n=("path", "size"), subjects=("subject", "nunique"),
                                 groups=("group", lambda g: ",".join(sorted(set(g)))), stage=("stage", "first"),
                                 n_events=("asset_id", lambda a: sum(len(v) for v in events[a.iloc[0]]["events"].values())))
print("Loaded files whose event series are identical to another loaded file:", len(dups), "in", len(pairs), "sets")
print(pairs.groupby(["subjects", "groups", "stage"]).size().to_string())
exp_files = sessions[sessions.experiment.notna()]
same_start = exp_files[exp_files.duplicated(["subject", "start"], keep=False)]
print("\nFiles in the 'Unknown' DLS excitatory group:", int((sessions.arm == "Unknown").sum()))
print("All files sharing subject and start time with another file:", len(same_start),
      "; their groups:", same_start.group.value_counts().to_dict())
odd = pairs[pairs.subjects > 1]
for h in odd.index:
    d = dups[dups.hash == h]
    r = seiler.rewards(events[d.asset_id.iloc[0]]["events"])
    print("\nIdentical across subjects:", d.subject.tolist(), d.stage.iloc[0], "; reward intervals (s) from",
          np.round(np.diff(r).min(), 2), "to", np.round(np.diff(r).max(), 2), "; series:", list(events[d.asset_id.iloc[0]]["events"]))

Loaded files whose event series are identical to another loaded file: 77 in 38 sets
subjects  groups                                                         stage   
1         Opto-DLS-Excitatory-ChR2,Opto-DLS-Excitatory-Unknown           FR1          8
                                                                         RI30         3
                                                                         omission     3
          Opto-DLS-Excitatory-ChR2Scrambled,Opto-DLS-Excitatory-Unknown  FR1         15
                                                                         RI30         4
                                                                         omission     4
3         Opto-DMS-Inhibitory-NpHR,Opto-DMS-Inhibitory-NpHRScrambled     omission     1

Files in the 'Unknown' DLS excitatory group: 137
All files sharing subject and start time with another file: 270 ; their groups: {'Opto-DLS-Excitatory-Unknown': 135, 'Opto-DLS-Excitatory-ChR2Scrambled': 82, 'Opto-DLS-Excitatory-ChR2': 53}

Identical across subjects: ['245.464', '266.477', '430.420'] omission ; reward intervals (s) from 20.0 to 20.0 ; series: ['right_reward_times']

ok_loaded = loaded[~((loaded.arm == "Unknown") & loaded.duplicated(["subject", "start"], keep=False))].copy()
fp = sessions[sessions.subject_experiment == "FP"]
fp_subj = fp.drop_duplicates("subject").copy()
fp_subj["label"] = [seiler.note_field(events[a]["notes"], "Punishment Group") for a in fp_subj.asset_id]
fp_subj["behavior"] = [seiler.note_field(events[a]["notes"], "Behavior") for a in fp_subj.asset_id]
fp_subj["did_not_learn"] = [seiler.note_field(events[a]["notes"], "Did Not Learn") for a in fp_subj.asset_id]
PHENO = {"Punishment Resistant": "PR", "Delayed Punishment Resistant": "DPR", "Punishment Sensitive": "PS"}
fp_subj["pheno"] = np.where(fp_subj.behavior == "RI60", fp_subj.label.map(PHENO), "RR20")
print("Fiber photometry mice: file-name arm, stored punishment group and schedule")
print(pd.crosstab([fp_subj.subject_arm, fp_subj.behavior], fp_subj.label).to_string())
print("Labels consistent across each mouse's files:",
      all(len({seiler.note_field(events[a]["notes"], "Punishment Group") for a in fp[fp.subject == s].asset_id}) == 1
          for s in fp_subj.subject))
opto_subj = sessions[sessions.subject_experiment.str.startswith("Opto", na=False)].drop_duplicates("subject").copy()
opto_subj["label"] = [seiler.note_field(events[a]["notes"], "Punishment Group") for a in opto_subj.asset_id]
opto_subj["did_not_learn"] = [seiler.note_field(events[a]["notes"], "Did Not Learn") for a in opto_subj.asset_id]
print("\nOptogenetic mice, stored punishment group:", opto_subj.label.value_counts().to_dict())
print("Did Not Learn flags, all mice:", pd.concat([fp_subj, opto_subj]).groupby("subject_group").did_not_learn.value_counts().unstack(fill_value=0).to_string())

Fiber photometry mice: file-name arm, stored punishment group and schedule
label                 Delayed Punishment Resistant  Punishment Resistant  Punishment Sensitive
subject_arm behavior                                                                          
DPR         RI60                                 9                     0                     0
PR          RI60                                 0                     9                     0
PS          RI60                                 0                     0                    18
RR20        RR20                                 0                     0                     6
Labels consistent across each mouse's files: True

Optogenetic mice, stored punishment group: {'nan': 110}
Did Not Learn flags, all mice: did_not_learn                      False  True
subject_group                                 
FP-DPR                                 9     0
FP-PR                                  9     0
FP-PS                                 18     0
FP-RR20                                6     0
Opto-DLS-Excitatory-ChR2              11     0
Opto-DLS-Excitatory-ChR2Scrambled     12     0
Opto-DLS-Excitatory-EYFP              12     0
Opto-DMS-Excitatory-ChR2              14     0
Opto-DMS-Excitatory-ChR2Scrambled     12     0
Opto-DMS-Excitatory-EYFP               7     0
Opto-DMS-Inhibitory-EYFP              15     1
Opto-DMS-Inhibitory-NpHR              12     1
Opto-DMS-Inhibitory-NpHRScrambled     13     0

side_rows = []
for a, r in ok_loaded.set_index("asset_id").iterrows():
    if not isinstance(r.side, str) or r.stage not in ("FR1", "RI30", "RI60", "RR20", "shock"):
        continue
    ev = events[a]["events"]
    rw = seiler.rewards(ev)
    if len(rw) == 0:
        continue
    def frac(s):
        p = ev.get(f"{s}_nose_poke_times", np.array([]))
        return np.isin(np.round(rw, 2), np.round(p, 2)).mean() if len(p) else 0.0
    other = "left" if r.side == "right" else "right"
    series = "+".join(k.split("_")[0] for k in ("left_reward_times", "right_reward_times") if k in ev)
    side_rows.append({"stage": r.stage, "described_side": r.side, "reward_series": series,
                      "rewards_in_described_side_pokes": frac(r.side) > 0.9, "rewards_in_other_side_pokes": frac(other) > 0.9})
side_df = pd.DataFrame(side_rows)
print("Do reward times coincide with nosepokes on the side the session description names? (files with rewards)")
print(side_df.groupby(["stage", "described_side", "reward_series", "rewards_in_described_side_pokes",
                       "rewards_in_other_side_pokes"]).size().to_string())

Do reward times coincide with nosepokes on the side the session description names? (files with rewards)
stage  described_side  reward_series  rewards_in_described_side_pokes  rewards_in_other_side_pokes
FR1    left            left           True                             False                          242
       right           right          True                             False                          277
RI30   left            left           True                             False                           90
       right           right          True                             False                          102
RI60   left            left           False                            False                            1
                                      True                             False                          270
       right           right          False                            False                            1
                                      True                             False                          363
RR20   left            left           False                            False                           20
                                      True                             False                            8
       right           right          False                            False                           25
                                      True                             False                           11
shock  left            right          True                             False                          151
       right           right          True                             False                          155

probe_files = ok_loaded[ok_loaded.stage == "shock"]
chk = []
for a in probe_files.asset_id:
    ev = events[a]["events"]
    if "footshock_times" not in ev:
        continue
    sh, rw = ev["footshock_times"], seiler.rewards(ev)
    chk.append({"n_shock": len(sh), "n_rew": len(rw), "shock_at_reward": np.isin(np.round(sh, 2), np.round(rw, 2)).mean() if len(sh) else np.nan,
                "first5_unshocked": not np.isin(np.round(rw[:5], 2), np.round(sh, 2)).any(), "last_event_s": max(v.max() for v in ev.values() if len(v))})
chk = pd.DataFrame(chk)
print(f"Shock probe files with footshock_times: {len(chk)} of {len(probe_files)}")
print(f"Fraction of shocks at a reward time: median {chk.shock_at_reward.median():.2f}, min {chk.shock_at_reward.min():.2f}")
ratio = chk.n_shock / chk.n_rew.replace(0, np.nan)
print(f"Shocks per reward: median {ratio.median():.2f} (Methods: RR3, about 1/3); 5th to 95th percentile {ratio.quantile(.05):.2f} to {ratio.quantile(.95):.2f}")
print("First five rewards never shocked (Methods):", int(chk.first5_unshocked.sum()), "of", len(chk))
print(f"Last event time in probe files: up to {chk.last_event_s.max():.0f} s (Methods: 60 min or 10 min of inactivity)")
dur = []
for a, st in ok_loaded.set_index("asset_id").stage.items():
    ev = events[a]["events"]
    if ev and st in ("FR1", "RI30", "RI60", "RR20", "omission"):
        dur.append({"stage": st, "last_event_s": max((v.max() for v in ev.values() if len(v)), default=np.nan), "n_rewards": len(seiler.rewards(ev))})
dur = pd.DataFrame(dur)
print("\nLast event time (s) and reward count by stage:")
print(dur.groupby("stage").agg(files=("n_rewards", "size"), last_event_max=("last_event_s", "max"), last_event_median=("last_event_s", "median"),
                               rewards_max=("n_rewards", "max")).round(0).to_string())

Shock probe files with footshock_times: 285 of 311
Fraction of shocks at a reward time: median 1.00, min 1.00
Shocks per reward: median 0.31 (Methods: RR3, about 1/3); 5th to 95th percentile 0.21 to 0.33
First five rewards never shocked (Methods): 285 of 285
Last event time in probe files: up to 3601 s (Methods: 60 min or 10 min of inactivity)

Last event time (s) and reward count by stage:
          files  last_event_max  last_event_median  rewards_max
stage                                                          
FR1         625          3601.0             3536.0           50
RI30        194          3600.0             3160.0           50
RI60        641          5354.0             3469.0           50
RR20         84          3674.0             3560.0           49
omission    127          3601.0             1463.0           50

fp_probes = seiler.probe_table(fp, events)
print("Probe days per fiber photometry mouse:", fp_probes.groupby("subject").size().value_counts().sort_index().to_dict())
print("Probe days with more than one file:", int((fp_probes.n_files > 1).sum()), "; with more than one file holding footshock_times:",
      int((fp_probes.n_with_shocks > 1).sum()))
print(fp_probes[fp_probes.n_with_shocks > 1].to_string(index=False))
print("\nTraining sessions before the early and late probes (Results: 1 to 2 and 13 to 14 days):")
print(fp_probes[fp_probes.probe != ""].groupby("probe").train_days_before.describe().to_string())
fp_el = seiler.early_late(fp_probes).join(fp_subj.set_index("subject")[["pheno", "label", "behavior", "sex"]])
print("\nMice missing a probe:")
print(fp_el[fp_el[["early", "late"]].isna().any(axis=1)].to_string())
print("Shock counts, range:", fp_el[["early", "late"]].min().min(), "to", fp_el[["early", "late"]].max().max())

Probe days per fiber photometry mouse: {1: 5, 2: 27, 3: 10}
Probe days with more than one file: 8 ; with more than one file holding footshock_times: 2
subject       date  n_files  n_with_shocks  shocks  shocks_sum  train_days_before probe
 64.205 2018-11-13        2              2      35          62                 16  late
 87.239 2019-03-04        2              2       5          20                  2 early

Training sessions before the early and late probes (Results: 1 to 2 and 13 to 14 days):
       count       mean       std   min   25%   50%   75%   max
probe                                                          
early   40.0   1.575000  0.812956   0.0   1.0   1.0   2.0   4.0
late    39.0  13.435897  1.569359  11.0  12.0  13.0  14.0  18.0

Mice missing a probe:
         early  late pheno                 label behavior sex
subject                                                      
028.392   28.0   NaN    PR  Punishment Resistant     RI60   M
101.260    7.0   NaN  RR20  Punishment Sensitive     RR20   F
62.207    28.0   NaN    PR  Punishment Resistant     RI60   M
78.214     NaN   2.0    PS  Punishment Sensitive     RI60   F
89.247     NaN  23.0    PR  Punishment Resistant     RI60   F
Shock counts, range: 1.0 to 53.0

all_ph = fp[fp.asset_id.map(lambda a: events[a]["has_photometry"])]
print("Files with photometry by stage:", all_ph.stage.value_counts().to_dict())
fp_sessions = all_ph[all_ph.stage.isin(["RI60", "RR20"])].sort_values("start_ts").copy()
fp_first_probe = fp_probes.sort_values("date").groupby("subject").date.first()
ri60_dates = fp[fp.stage.isin(["RI60", "RR20"])].groupby("subject").date.apply(sorted)
fp_sessions["train_day"] = [sum(d < r.date for d in ri60_dates.get(r.subject, [])) + 1 for r in fp_sessions.itertuples()]
early_probe = fp_probes[fp_probes.probe == "early"].set_index("subject").date
fp_sessions["fp_stage"] = [
    ("early" if r.date < early_probe[r.subject] else "late") if r.subject in early_probe.index
    else ("early" if r.train_day <= 3 else "late") for r in fp_sessions.itertuples()]
fp_sessions = fp_sessions.merge(fp_subj[["subject", "pheno", "behavior"]], on="subject")
print("Photometry sessions:", len(fp_sessions), "from", fp_sessions.subject.nunique(), "mice")
print(pd.crosstab([fp_sessions.pheno, fp_sessions.stage], fp_sessions.fp_stage).to_string())
print("\nMice with photometry per phenotype and stage (paper Fig 3E: RI30 PR 6, DPR 6, PS 11; early PR 7, DPR 9, PS 15; late PR 5, DPR 6, PS 13):")
print(fp_sessions.groupby(["pheno", "fp_stage"]).subject.nunique().unstack().to_string())
print("\nTraining day of each photometry session (1 = first RI60 or RR20 session):")
print(fp_sessions.groupby("fp_stage").train_day.describe().to_string())


def process_photometry(row):
    rec = seiler.load_photometry(row.asset_id)
    d, fs, t0, ev = rec["data"], rec["fs"], rec["t0"], rec["events"]
    rew, unrew = seiler.split_nosepokes(ev, row.side)
    def inside(times):  # events whose PSTH window lies inside the recording (as seiler.psth)
        idx = np.round((np.asarray(times) - t0) * fs)
        return int(((idx - seiler.PSTH_PRE_S * fs >= 0) & (idx + seiler.PSTH_POST_S * fs < len(d))).sum())
    out = {"asset_id": row.asset_id, "fs": fs, "t0": t0, "n_samples": len(d), "duration_s": len(d) / fs,
           "last_event_s": max((v.max() for v in ev.values() if len(v)), default=np.nan),
           "n_rewards": len(seiler.rewards(ev)), "locations": ",".join(rec["locations"]),
           "sources": ",".join(s.split("_")[-1] for s in rec["sources"]),
           "n_rewarded": len(rew), "n_unrewarded": len(unrew),
           "n_rewarded_kept": inside(rew), "n_unrewarded_kept": inside(unrew), "coordinates": str(np.round(rec["coordinates"], 2).tolist()),
           "modulation": ";".join(f"{k.replace('commanded_voltage_series_', '')}={v:g}" for k, v in sorted(rec["modulation_hz"].items()))}
    for region, (si, ci) in {"DMS": (0, 1), "DLS": (2, 3)}.items():
        out[f"{region}_dropout_frac"] = seiler.dropout_mask(d[:, si], d[:, ci], fs).mean()
        z = seiler.zscored_dff(d[:, si], d[:, ci], fs)
        res = seiler.peak_score(z, fs, rew, unrew, t0)
        out[f"{region}_score"], out[f"{region}_rew_max"], out[f"{region}_unrew_min"] = res["score"], res["rew_max"], res["unrew_min"]
        out[f"{region}_psth_rew"], out[f"{region}_psth_unrew"] = res["psth_rew"], res["psth_unrew"]
        out[f"{region}_auc_2_10"] = res["auc_2_10"]
        # the authors' processing without the dropout step
        z0 = seiler.zscored_dff(d[:, si], d[:, ci], fs, mask_dropouts=False)
        out[f"{region}_score_nomask"] = seiler.peak_score(z0, fs, rew, unrew, t0)["score"]
        out[f"{region}_max_abs_z_nomask"] = np.nanmax(np.abs(z0))
        # clock check: the same events shifted by 37 s
        out[f"{region}_rew_max_shifted"] = seiler.peak_score(z, fs, rew + 37.0, unrew + 37.0, t0)["rew_max"]
        # column check: the channels the other way round
        zs = seiler.zscored_dff(d[:, ci], d[:, si], fs)
        out[f"{region}_rew_max_swapped"] = seiler.peak_score(zs, fs, rew, unrew, t0)["rew_max"]
    out["t_psth"] = res["t"]
    return out


t = time.time()
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor(6) as ex:
    ph = pd.DataFrame(list(ex.map(process_photometry, fp_sessions.itertuples())))
ph = fp_sessions.merge(ph, on="asset_id")
# Methods: recordings with excessive artifacts were excluded; here a region's session is excluded when
# more than 10% of its samples are dropouts (a choice here).
for region in ["DMS", "DLS"]:
    bad = ph[f"{region}_dropout_frac"] > 0.10
    ph.loc[bad, [f"{region}_score", f"{region}_rew_max", f"{region}_unrew_min", f"{region}_score_nomask"]] = np.nan
    print(f"{region} sessions excluded for dropouts over 10% of samples:", ph.loc[bad, ["subject", "date"]].to_dict("records"))
print(f"\nprocessed {len(ph)} photometry sessions in {time.time() - t:.0f} s")
print("Sampling rate (Hz):", ph.fs.unique(), "; start time:", ph.t0.unique())
print("Column locations:", list(ph.locations.unique()), "; excitation sources:", list(ph.sources.unique()))
print("Modulation frequencies (Hz) in the files:", ph.modulation.value_counts().to_dict())
print("Fiber coordinates (AP, ML, DV) per column in the files:", ph.coordinates.value_counts().to_dict(),
      "; Methods: DMS AP 0.8, ML 1.5, DV -2.8; DLS AP -0.1, ML 2.8, DV -3.5")
print(f"Recording length minus last behavioral event (s): {np.nanmin(ph.duration_s - ph.last_event_s):.0f} to {np.nanmax(ph.duration_s - ph.last_event_s):.0f}")
print(f"Rewarded / unrewarded nosepokes per session: {ph.n_rewarded.min()} to {ph.n_rewarded.max()} / {ph.n_unrewarded.min()} to {ph.n_unrewarded.max()}")
kept_r = ph.n_rewarded_kept / ph.n_rewarded.replace(0, np.nan)
kept_u = ph.n_unrewarded_kept / ph.n_unrewarded.replace(0, np.nan)
print(f"Events inside the recording (PSTH window complete): rewarded {ph.n_rewarded_kept.sum()} of {ph.n_rewarded.sum()}, "
      f"unrewarded {ph.n_unrewarded_kept.sum()} of {ph.n_unrewarded.sum()}; sessions losing any event: {(kept_r < 1).sum()} rewarded, {(kept_u < 1).sum()} unrewarded")
lost = ph[(kept_r < 0.9) | (kept_u < 0.9)]
print("Sessions keeping under 90% of either kind (subject, date, rewarded kept/all, unrewarded kept/all, recording minus last event in s):")
for r in lost.itertuples():
    print(f"  {r.subject} {r.date}: {r.n_rewarded_kept}/{r.n_rewarded}, {r.n_unrewarded_kept}/{r.n_unrewarded}, {r.duration_s - r.last_event_s:.0f}")
print(f"Fewest events kept in a session: rewarded {ph.n_rewarded_kept.min()}, unrewarded {ph.n_unrewarded_kept.min()}")
print("Sessions with rewards but no reward at a nosepoke time:", ph[(ph.n_rewarded == 0)][["subject", "date", "n_rewards"]].to_dict("records"))
for region in ["DMS", "DLS"]:
    m = ph[f"{region}_dropout_frac"] > 0
    fr = ph[f"{region}_dropout_frac"]
    print(f"{region}: percent of samples left out per session, median {100 * fr.median():.3f}, above 0.1% in {(fr > 0.001).sum()} sessions, "
          f"maximum {100 * fr.max():.2f}; largest |z| without the dropout step {ph[f'{region}_max_abs_z_nomask'].max():.0f} "
          f"(above 50 in {(ph[f'{region}_max_abs_z_nomask'] > 50).sum()} sessions)")

Files with photometry by stage: {'RI60': 53, 'RR20': 9, 'probetest': 1}
Photometry sessions: 62 from 40 mice
fp_stage     early  late
pheno stage             
DPR   RI60       9     7
PR    RI60       7     4
PS    RI60      15    11
RR20  RR20       5     4

Mice with photometry per phenotype and stage (paper Fig 3E: RI30 PR 6, DPR 6, PS 11; early PR 7, DPR 9, PS 15; late PR 5, DPR 6, PS 13):
fp_stage  early  late
pheno                
DPR           9     7
PR            7     4
PS           14    11
RR20          5     4

Training day of each photometry session (1 = first RI60 or RR20 session):

          count       mean       std   min   25%   50%   75%   max
fp_stage                                                          
early      36.0   1.305556  0.624246   1.0   1.0   1.0   1.0   3.0
late       26.0  13.307692  1.761119  12.0  12.0  13.0  14.0  18.0

DMS sessions excluded for dropouts over 10% of samples: []
DLS sessions excluded for dropouts over 10% of samples: [{'subject': '141.308', 'date': '2019-08-09'}]

processed 62 photometry sessions in 402 s
Sampling rate (Hz): [1017.25262451] ; start time: [0.]
Column locations: ['DMS,DMS,DLS,DLS'] ; excitation sources: ['signal,control,signal,control']
Modulation frequencies (Hz) in the files: {'dls_calcium_signal=450;dls_isosbestic_control=270;dms_calcium_signal=211;dms_isosbestic_control=330': 62}
Fiber coordinates (AP, ML, DV) per column in the files: {'[[0.8, 1.5, 2.8], [0.8, 1.5, 2.8], [0.1, 2.8, 3.5], [0.1, 2.8, 3.5]]': 62} ; Methods: DMS AP 0.8, ML 1.5, DV -2.8; DLS AP -0.1, ML 2.8, DV -3.5
Recording length minus last behavioral event (s): -2257 to 2516
Rewarded / unrewarded nosepokes per session: 0 to 87 / 3 to 383
Events inside the recording (PSTH window complete): rewarded 2336 of 2337, unrewarded 10892 of 10916; sessions losing any event: 1 rewarded, 4 unrewarded
Sessions keeping under 90% of either kind (subject, date, rewarded kept/all, unrewarded kept/all, recording minus last event in s):
  92.246 2019-02-27: 18/18, 69/85, -1391
Fewest events kept in a session: rewarded 0, unrewarded 3
Sessions with rewards but no reward at a nosepoke time: [{'subject': '332.393', 'date': '2020-07-28', 'n_rewards': 32}]
DMS: percent of samples left out per session, median 0.033, above 0.1% in 10 sessions, maximum 1.61; largest |z| without the dropout step 1730 (above 50 in 6 sessions)
DLS: percent of samples left out per session, median 0.033, above 0.1% in 10 sessions, maximum 99.97; largest |z| without the dropout step 1758 (above 50 in 6 sessions)

chk = pd.DataFrame({
    "region": np.repeat(["DMS", "DLS"], len(ph)),
    "as_labelled": np.r_[ph.DMS_rew_max, ph.DLS_rew_max],
    "events_shifted_37s": np.r_[ph.DMS_rew_max_shifted, ph.DLS_rew_max_shifted],
    "channels_swapped": np.r_[ph.DMS_rew_max_swapped, ph.DLS_rew_max_swapped]})
print("Maximum of the rewarded nosepoke PSTH in 0 to 1.5 s (z), median [5th, 95th percentile] over sessions:")
print(chk.groupby("region").agg(lambda x: f"{np.nanmedian(x):.2f} [{np.nanpercentile(x, 5):.2f}, {np.nanpercentile(x, 95):.2f}]").to_string())
for region in ["DMS", "DLS"]:
    a, b = ph[f"{region}_rew_max"], ph[f"{region}_rew_max_shifted"]
    ok = a.notna() & b.notna()
    print(f"{region}: as labelled above shifted in {(a[ok] > b[ok]).sum()} of {ok.sum()} sessions, Wilcoxon p = {stats.wilcoxon(a[ok], b[ok]).pvalue:.2g}")
    c = ph[f"{region}_rew_max_swapped"]
    print(f"{region}: as labelled above swapped in {(a[ok] > c[ok]).sum()} of {ok.sum()} sessions")
a_dms, a_dls = ph.DMS_auc_2_10, ph.DLS_auc_2_10
ok = a_dms.notna() & a_dls.notna()
print(f"Rewarded PSTH area 2 to 10 s after the nosepoke (z s), median DMS {a_dms.median():.2f}, DLS {a_dls.median():.2f}; "
      f"DLS above DMS in {(a_dls[ok] > a_dms[ok]).sum()} of {ok.sum()} sessions, Wilcoxon p = {stats.wilcoxon(a_dls[ok], a_dms[ok]).pvalue:.2g}")
print("Sessions whose rewarded maximum is below 0.3 z in both regions:", ph[(ph.DMS_rew_max < .3) & (ph.DLS_rew_max < .3)].subject.tolist())
fig, axes = plt.subplots(1, 2, figsize=(9, 3), sharey=True)
tt = ph.t_psth.iloc[0]
for ax, region in zip(axes, ["DMS", "DLS"]):
    for lab, c in [("PR", "k"), ("DPR", "m"), ("PS", "c"), ("RR20", "y")]:
        sel = ph[ph.pheno == lab]
        ax.plot(tt, np.nanmean(np.stack(sel[f"{region}_psth_rew"]), 0), c, label=f"{lab} rewarded (dashed: unrewarded)")
        ax.plot(tt, np.nanmean(np.stack(sel[f"{region}_psth_unrew"]), 0), c + "--")
    ax.axvline(0, color="grey", lw=.5); ax.set_title(region); ax.set_xlabel("s from nosepoke")
axes[0].set_ylabel("z"); axes[0].legend(fontsize=7)
plt.tight_layout(); plt.show()
allp = np.concatenate([np.stack(ph[f"{r}_psth_{k}"]).ravel() for r in ["DMS", "DLS"] for k in ["rew", "unrew"]])
print(f"Plotted session-mean PSTHs span {np.nanmin(allp):.2f} to {np.nanmax(allp):.2f} z before averaging over mice")

Maximum of the rewarded nosepoke PSTH in 0 to 1.5 s (z), median [5th, 95th percentile] over sessions:
               as_labelled  events_shifted_37s    channels_swapped
region                                                            
DLS      0.75 [0.03, 2.23]  0.09 [-0.14, 0.48]  0.25 [-0.11, 0.72]
DMS     0.75 [-0.01, 3.50]  0.14 [-0.20, 0.47]  0.28 [-0.12, 0.67]
DMS: as labelled above shifted in 54 of 61 sessions, Wilcoxon p = 1.6e-10
DMS: as labelled above swapped in 46 of 61 sessions
DLS: as labelled above shifted in 50 of 60 sessions, Wilcoxon p = 5.8e-09
DLS: as labelled above swapped in 45 of 60 sessions
Rewarded PSTH area 2 to 10 s after the nosepoke (z s), median DMS 0.70, DLS 1.13; DLS above DMS in 40 of 60 sessions, Wilcoxon p = 0.0082
Sessions whose rewarded maximum is below 0.3 z in both regions: ['249.391', '028.392', '272.396']

Plotted session-mean PSTHs span -1.97 to 4.98 z before averaging over mice

def mixed_anova(df, between, within="probe", dv="shocks", subject="subject"):
    """Two-way mixed ANOVA (one between, one within factor of two levels), complete cases."""
    w = df.pivot_table(index=[subject, between], columns=within, values=dv).dropna().reset_index()
    levels = [c for c in w.columns if c not in (subject, between)]
    a, b = w[levels[0]].to_numpy(), w[levels[1]].to_numpy()
    g = w[between].to_numpy()
    groups = pd.unique(g)
    n, k = len(w), len(groups)
    mean_subj = (a + b) / 2
    grand = mean_subj.mean()
    ss_between = 2 * sum(((mean_subj[g == x].mean() - grand) ** 2) * (g == x).sum() for x in groups)
    ss_subj_err = 2 * sum(((mean_subj[g == x] - mean_subj[g == x].mean()) ** 2).sum() for x in groups)
    d = b - a
    ss_time = n * (d.mean() ** 2) / 2
    ss_inter = sum((g == x).sum() * (d[g == x].mean() - d.mean()) ** 2 for x in groups) / 2
    ss_within_err = sum(((d[g == x] - d[g == x].mean()) ** 2).sum() for x in groups) / 2
    df_b, df_e = k - 1, n - k
    out = {}
    for name, ss, dfn, sse in [(between, ss_between, df_b, ss_subj_err), (within, ss_time, 1, ss_within_err),
                               (f"{between} x {within}", ss_inter, df_b, ss_within_err)]:
        F = (ss / dfn) / (sse / df_e)
        out[name] = {"F": F, "df": f"{dfn},{df_e}", "p": stats.f.sf(F, dfn, df_e)}
    return pd.DataFrame(out).T, w


def fmt(aov):
    return aov.assign(F=aov.F.map("{:.2f}".format), p=aov.p.map("{:.2g}".format)).to_string()


def tukey(values, groups):
    t = pairwise_tukeyhsd(np.asarray(values, float), np.asarray(groups))
    return pd.DataFrame(t._results_table.data[1:], columns=t._results_table.data[0])


fp_long = fp_probes[fp_probes.probe != ""].merge(fp_subj[["subject", "pheno", "sex", "behavior"]], on="subject")
fp_long_sum = fp_long.assign(shocks=fp_long.shocks_sum)
fp_el["schedule"] = fp_el.behavior

om = fp[fp.stage == "omission"].sort_values("start_ts").groupby("subject").asset_id.first()
fp_el["omission_min"] = [seiler.omission_time_min(events[om[s]]["events"]) if s in om.index else np.nan for s in fp_el.index]

late_date = fp_probes[fp_probes.probe == "late"].set_index("subject").date
train = fp[fp.stage.isin(["RI60", "RR20"])].sort_values("start_ts").copy()
train = train[[r.subject not in late_date.index or r.date < late_date[r.subject] for r in train.itertuples()]]
rows = []
for r in train.itertuples():
    ev = events[r.asset_id]["events"]
    pokes = seiler.active_pokes(ev, r.side)
    rm = seiler.reward_measures(ev, r.side)
    rows.append({"subject": r.subject, "date": r.date, "stage": r.stage, "n_rewards": len(seiler.rewards(ev)),
                 "n_active": len(pokes), "n_events": sum(len(v) for v in ev.values()), "series": len(ev),
                 "last_event_s": max((v.max() for v in ev.values() if len(v)), default=np.nan),
                 "size_mb": r.size / 1e6, "n_waits": rm["n_waits"], "series_names": ",".join(sorted(ev)),
                 "rate": len(pokes) / (seiler.session_length(ev) / 60.0), "wait": rm["wait"], "extra": rm["extra"],
                 "has_ports": rm["has_ports"], "iri": rm["iri"]})
tr_all = pd.DataFrame(rows).merge(fp_subj[["subject", "pheno", "sex"]], on="subject")
# Files with no active nosepoke and no reward are left out (listed in the next cell): files that hold
# port entries only, and a false start followed by the full session on the same day.
empty = tr_all[(tr_all.n_rewards == 0) & (tr_all.n_active == 0)]
tr = tr_all.drop(empty.index).reset_index(drop=True)
tr["day"] = tr.groupby("subject").cumcount() + 1
def mouse_means(t):
    return t.groupby("subject").agg(rate=("rate", "mean"), rewards=("n_rewards", "mean"), wait=("wait", "mean"),
                                    extra=("extra", "mean"), sessions=("rate", "size")).join(fp_subj.set_index("subject").pheno)
per_mouse = mouse_means(tr)
per_mouse_all = mouse_means(tr_all)  # variant with the empty files kept
print("Training sessions used (up to the late probe):", len(tr), "; per mouse", per_mouse.sessions.min(), "to", per_mouse.sessions.max())
print("Sessions without port entries (no wait or extra nosepokes):", int((~tr.has_ports).sum()))
print("Ranges: rate", tr.rate.min().round(2), "to", tr.rate.max().round(1), "per min; rewards", tr.n_rewards.min(), "to", tr.n_rewards.max())
print(per_mouse.groupby("pheno").size().to_string())

Training sessions used (up to the late probe): 551 ; per mouse 7 to 19
Sessions without port entries (no wait or extra nosepokes): 0
Ranges: rate 0.02 to 39.7 per min; rewards 1 to 50
pheno
DPR      9
PR       9
PS      18
RR20     6

print(f"Training files with 0 rewards: {(tr_all.n_rewards == 0).sum()}; with 0 active nosepokes: {(tr_all.n_active == 0).sum()}; "
      f"both, left out: {len(empty)} of {len(tr_all)}")
print(empty[["subject", "pheno", "stage", "date", "size_mb", "series_names", "n_events", "last_event_s"]]
      .round(2).sort_values(["pheno", "subject", "date"]).to_string(index=False))
print("\nFile size (MB) of all training files, median and range:", round(tr_all.size_mb.median(), 2), round(tr_all.size_mb.min(), 2), round(tr_all.size_mb.max(), 2))
print("Empty files by phenotype:", empty.pheno.value_counts().to_dict(), "; mice:", empty.subject.nunique())
same_day = tr_all.groupby(["subject", "date"]).size()
print("Mouse-days with more than one training file:", same_day[same_day > 1].to_dict())
print("Other files of those days:")
print(tr_all.set_index(["subject", "date"]).loc[same_day[same_day > 1].index][["stage", "series_names", "n_active", "n_rewards", "last_event_s"]].round(1).to_string())
print("Training sessions kept:", len(tr))
# variant: one file per mouse and day, the one with the most rewards
per_mouse_day1 = mouse_means(tr.sort_values("n_rewards").drop_duplicates(["subject", "date"], keep="last"))

Training files with 0 rewards: 21; with 0 active nosepokes: 21; both, left out: 21 of 572
subject pheno stage       date  size_mb                                 series_names  n_events  last_event_s
140.306    PS  RI60 2019-08-30     0.22 left_nose_poke_times,reward_port_entry_times        30        170.85
100.258  RR20  RR20 2019-04-29     0.22                      reward_port_entry_times       291       3563.65
100.258  RR20  RR20 2019-04-30     0.22                      reward_port_entry_times       200       3540.90
100.258  RR20  RR20 2019-05-01     0.22                      reward_port_entry_times       138       3589.75
100.258  RR20  RR20 2019-05-02     0.22                      reward_port_entry_times       192       3192.00
100.258  RR20  RR20 2019-05-03     0.22                      reward_port_entry_times       132       3413.25
 96.259  RR20  RR20 2019-04-29     0.23                      reward_port_entry_times       626       3590.35
 96.259  RR20  RR20 2019-04-30     0.22                      reward_port_entry_times       466       3594.55
 96.259  RR20  RR20 2019-05-01     0.22                      reward_port_entry_times       276       3508.20
 96.259  RR20  RR20 2019-05-02     0.22                      reward_port_entry_times       379       3373.05
 96.259  RR20  RR20 2019-05-03     0.22                      reward_port_entry_times       250       3220.95
 97.257  RR20  RR20 2019-04-29     0.22                      reward_port_entry_times       200       3556.80
 97.257  RR20  RR20 2019-04-30     0.22                      reward_port_entry_times       151       2422.10
 97.257  RR20  RR20 2019-05-01     0.22                      reward_port_entry_times        69       1736.70
 97.257  RR20  RR20 2019-05-02     0.22                      reward_port_entry_times       138       3588.95
 97.257  RR20  RR20 2019-05-03     0.22                      reward_port_entry_times        92       3271.15
 99.257  RR20  RR20 2019-04-29     0.22                      reward_port_entry_times       482       3480.85
 99.257  RR20  RR20 2019-04-30     0.22                      reward_port_entry_times       220       2830.80
 99.257  RR20  RR20 2019-05-01     0.22                      reward_port_entry_times       238       3404.60
 99.257  RR20  RR20 2019-05-02     0.22                      reward_port_entry_times       262       3291.30
 99.257  RR20  RR20 2019-05-03     0.22                      reward_port_entry_times       155       3391.80

File size (MB) of all training files, median and range: 0.24 0.22 591.44
Empty files by phenotype: {'RR20': 20, 'PS': 1} ; mice: 5
Mouse-days with more than one training file: {('112.283', '2019-06-07'): 2, ('114.273', '2019-06-12'): 2, ('115.273', '2019-06-28'): 2, ('140.306', '2019-08-30'): 2, ('271.396', '2020-07-07'): 2, ('348.393', '2020-07-23'): 2, ('81.236', '2019-02-06'): 3, ('88.239', '2019-02-19'): 2, ('88.239', '2019-02-28'): 2, ('88.239', '2019-03-01'): 2, ('92.246', '2019-03-11'): 2}
Other files of those days:
                   stage                                                                           series_names  n_active  n_rewards  last_event_s
subject date                                                                                                                                      
112.283 2019-06-07  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times        59         24        2300.8
        2019-06-07  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times       150         20        1388.6
114.273 2019-06-12  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times        81         17        1720.2
        2019-06-12  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times       383         36        3586.8
115.273 2019-06-28  RI60                       reward_port_entry_times,right_nose_poke_times,right_reward_times         7          3        1698.2
        2019-06-28  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times       101         42        3599.0
140.306 2019-08-30  RI60                                           left_nose_poke_times,reward_port_entry_times         0          0         170.8
        2019-08-30  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times       439         50        3556.8
271.396 2020-07-07  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times         2          2         719.6
        2020-07-07  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times       206         44        3661.7
348.393 2020-07-23  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times         4          3        3292.3
        2020-07-23  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times       164         50        3349.4
81.236  2019-02-06  RI60                       reward_port_entry_times,right_nose_poke_times,right_reward_times       317          6         362.0
        2019-02-06  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times       464         28        2734.4
        2019-02-06  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times      1370         42        3600.4
88.239  2019-02-19  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times         2          1         530.6
        2019-02-19  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times       772         44        4128.3
        2019-02-28  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times       196          9        1393.1
        2019-02-28  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times       912         23        3601.0
        2019-03-01  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times         1          1         366.9
        2019-03-01  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times      1128         49        3592.8
92.246  2019-03-11  RI60  left_nose_poke_times,reward_port_entry_times,right_nose_poke_times,right_reward_times         7          4         552.3
        2019-03-11  RI60   left_nose_poke_times,left_reward_times,reward_port_entry_times,right_nose_poke_times       144         42        3589.8
Training sessions kept: 551

def fig1d(long):
    aov, w = mixed_anova(long, "behavior")
    ri = w[w.behavior == "RI60"]
    t_ = stats.ttest_rel(ri["late"], ri["early"])
    return aov, w, t_

aov, w, t_ = fig1d(fp_long)
print(fmt(aov))
print(w.groupby("behavior")[["early", "late"]].agg(["mean", "std", "size"]).round(2).to_string())
print(f"RI60 late vs early: paired t = {t_.statistic:.2f}, p = {t_.pvalue:.2g}, Bonferroni x2 p = {min(1, 2 * t_.pvalue):.2g}")
rr = w[w.behavior == "RR20"]
print(f"RR20 late vs early: paired t p = {stats.ttest_rel(rr['late'], rr['early']).pvalue:.2g} (n = {len(rr)})")
aov_s, w_s, t_s = fig1d(fp_long_sum)
print("\nVariant, shocks summed over a probe day's files:"); print(fmt(aov_s)); print(
      f"RI60 late vs early Bonferroni p = {min(1, 2 * t_s.pvalue):.2g}")
fig, ax = plt.subplots(figsize=(3.5, 3))
for sched, c in [("RI60", "k"), ("RR20", "grey")]:
    d = w[w.behavior == sched]
    for _, r in d.iterrows():
        ax.plot([0, 1], [r.early, r.late], c, alpha=.3, lw=.8)
    ax.plot([0, 1], [d.early.mean(), d.late.mean()], c, lw=3, label=f"{sched} (n={len(d)})")
ax.set_xticks([0, 1], ["early", "late"]); ax.set_ylabel("shocks"); ax.legend(); plt.tight_layout(); plt.show()
print("Plotted counts span", w[["early", "late"]].min().min(), "to", w[["early", "late"]].max().max())

                     F    df      p
behavior          3.46  1,35  0.071
probe             5.35  1,35  0.027
behavior x probe  0.30  1,35   0.59
probe     early              late            
           mean   std size   mean    std size
behavior                                     
RI60      11.38  7.77   32  16.31  14.21   32
RR20       4.80  1.64    5   6.60   5.37    5
RI60 late vs early: paired t = 2.25, p = 0.032, Bonferroni x2 p = 0.063
RR20 late vs early: paired t p = 0.55 (n = 5)

Variant, shocks summed over a probe day's files:
                     F    df      p
behavior          3.48  1,35  0.071
probe             4.65  1,35  0.038
behavior x probe  0.29  1,35    0.6
RI60 late vs early Bonferroni p = 0.089

Plotted counts span 1.0 to 53.0

o = fp_el.dropna(subset=["omission_min"])
print(o.groupby("schedule").omission_min.agg(["mean", "std", "size", "min", "max"]).round(2).to_string())
t1 = stats.ttest_ind(o[o.schedule == "RI60"].omission_min, o[o.schedule == "RR20"].omission_min)
print(f"Unpaired t-test RI60 vs RR20: t = {t1.statistic:.2f}, p = {t1.pvalue:.3g}")
print("Variance ratio RI60 / RR20:", round(o[o.schedule == "RI60"].omission_min.var() / o[o.schedule == "RR20"].omission_min.var(), 1))
print(f"Welch variant: p = {stats.ttest_ind(o[o.schedule == 'RI60'].omission_min, o[o.schedule == 'RR20'].omission_min, equal_var=False).pvalue:.3g}")
print("Mice with fewer than 50 rewards (60 min):", (o.omission_min >= 60).groupby(o.schedule).sum().to_dict())
print("RI60 omission times by phenotype:", o[o.schedule == "RI60"].groupby("pheno").omission_min.agg(["mean", "size"]).round(1).to_dict("index"))

           mean    std  size    min   max
schedule                                 
RI60      33.97  16.68    19  17.10  60.0
RR20      17.72   1.05     4  17.08  19.3
Unpaired t-test RI60 vs RR20: t = 1.91, p = 0.0697
Variance ratio RI60 / RR20: 251.1
Welch variant: p = 0.000496
Mice with fewer than 50 rewards (60 min): {'RI60': 4, 'RR20': 0}
RI60 omission times by phenotype: {'DPR': {'mean': 31.7, 'size': 5}, 'PR': {'mean': 51.6, 'size': 4}, 'PS': {'mean': 28.1, 'size': 10}}

o1 = fp_el[(fp_el.schedule == "RI60")].dropna(subset=["omission_min", "late"])
r1 = stats.pearsonr(o1.late, o1.omission_min)
print(f"Pearson r = {r1.statistic:.3f}, p = {r1.pvalue:.3g}, n = {len(o1)}")
o1e = fp_el[(fp_el.schedule == "RI60")].dropna(subset=["omission_min", "early"])
print(f"Variant, early probe: r = {stats.pearsonr(o1e.early, o1e.omission_min).statistic:.3f}, p = {stats.pearsonr(o1e.early, o1e.omission_min).pvalue:.3g}, n = {len(o1e)}")
print(f"Spearman, late probe: rho = {stats.spearmanr(o1.late, o1.omission_min).statistic:.3f}, p = {stats.spearmanr(o1.late, o1.omission_min).pvalue:.3g}")

Pearson r = 0.335, p = 0.175, n = 18
Variant, early probe: r = 0.592, p = 0.0076, n = 19
Spearman, late probe: rho = 0.403, p = 0.097

ri = fp_el[fp_el.schedule == "RI60"].copy()
print("Stored labels of RI60 mice:", ri.pheno.value_counts().to_dict())
both = ri.dropna(subset=["early", "late"]).copy()
both["rule"] = [seiler.classify(e, l) for e, l in zip(both.early, both.late)]
both["change_pct"] = 100 * (both.late - both.early) / both.early
print(f"Mice with both probes: {len(both)}; rule agrees with the stored label for {(both.rule == both.pheno).sum()}")
print(pd.crosstab(both.pheno, both.rule).to_string())
print("\nDisagreements:")
print(both[both.rule != both.pheno][["early", "late", "change_pct", "pheno", "rule"]].round(1).to_string())
print(f"\nMedian early count of the RI60 mice: {ri.early.median():.1f} (Methods: median split at 13)")
print(f"75th percentile of the percent change: {both.change_pct.quantile(.75):.1f} (Methods: top quartile, above 85%)")
both["rule_ge13"] = ["DPR" if 100 * (l - e) / e > 85 else ("PR" if e >= 13 else "PS") for e, l in zip(both.early, both.late)]
print("Variant, PR when early >= 13: agrees for", int((both.rule_ge13 == both.pheno).sum()), "of", len(both))
bs = seiler.early_late(fp_probes, "shocks_sum").reindex(both.index)
both["rule_sum"] = [seiler.classify(e, l) for e, l in zip(bs.early, bs.late)]
print("Variant, counts summed over a day's files: agrees for", int((both.rule_sum == both.pheno).sum()), "of", len(both))
one = ri[ri[["early", "late"]].isna().any(axis=1)]
print("\nMice with one probe, label and available count:")
print(one[["early", "late", "pheno"]].to_string())

Stored labels of RI60 mice: {'PS': 18, 'PR': 9, 'DPR': 9}
Mice with both probes: 32; rule agrees with the stored label for 31
rule   DPR  PR  PS
pheno             
DPR      9   0   0
PR       0   5   1
PS       0   0  17

Disagreements:
         early  late  change_pct pheno rule
subject                                    
048.392   13.0  18.0        38.5    PR   PS

Median early count of the RI60 mice: 9.0 (Methods: median split at 13)
75th percentile of the percent change: 87.7 (Methods: top quartile, above 85%)
Variant, PR when early >= 13: agrees for 32 of 32
Variant, counts summed over a day's files: agrees for 30 of 32

Mice with one probe, label and available count:
         early  late pheno
subject                   
028.392   28.0   NaN    PR
62.207    28.0   NaN    PR
78.214     NaN   2.0    PS
89.247     NaN  23.0    PR

ri_long = fp_long[fp_long.behavior == "RI60"]
aov2, w2 = mixed_anova(ri_long, "pheno")
print(fmt(aov2))
print(w2.groupby("pheno")[["early", "late"]].agg(["mean", "std", "size"]).round(2).to_string())
for ph_ in ["PR", "DPR", "PS"]:
    d = w2[w2.pheno == ph_]
    p = stats.ttest_rel(d.late, d.early).pvalue
    print(f"{ph_}: late - early mean {np.mean(d.late - d.early):.2f}, paired t p = {p:.2g}, Bonferroni x3 p = {min(1, 3 * p):.2g}, n = {len(d)}")

                   F    df      p
pheno          25.54  2,29  4e-07
probe          13.35  1,29  0.001
pheno x probe  26.37  2,29  3e-07
probe  early              late            
        mean   std size   mean    std size
pheno                                     
DPR    11.56  5.81    9  32.11  11.34    9
PR     23.50  7.26    6  20.50  13.46    6
PS      7.00  2.85   17   6.47   4.54   17
PR: late - early mean -3.00, paired t p = 0.59, Bonferroni x3 p = 1, n = 6
DPR: late - early mean 20.56, paired t p = 0.00014, Bonferroni x3 p = 0.00041, n = 9
PS: late - early mean -0.53, paired t p = 0.56, Bonferroni x3 p = 1, n = 17

pm = per_mouse[per_mouse.pheno.isin(["PR", "DPR", "PS"])]
print(pm.groupby("pheno").rate.agg(["mean", "std", "size", "min", "max"]).round(2).to_string())
print(f"One-way ANOVA: p = {stats.f_oneway(*[pm[pm.pheno == g].rate for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
print(tukey(pm.rate, pm.pheno).to_string())
trr = tr[tr.pheno.isin(["PR", "DPR", "PS"])]
m1 = smf.mixedlm("rate ~ C(pheno) + day", trr, groups=trr.subject).fit(reml=False)
m0 = smf.mixedlm("rate ~ day", trr, groups=trr.subject).fit(reml=False)
lr = 2 * (m1.llf - m0.llf)
print(f"Variant, mixed model over {len(trr)} sessions: phenotype likelihood ratio {lr:.1f}, p = {stats.chi2.sf(lr, 2):.2g}")
pz = per_mouse_all[per_mouse_all.pheno.isin(["PR", "DPR", "PS"])]
print("Variant with the empty files kept:", pz.groupby("pheno").rate.mean().round(2).to_dict(),
      f"; ANOVA p = {stats.f_oneway(*[pz[pz.pheno == g].rate for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
print(tukey(pz.rate, pz.pheno).to_string())
p1 = per_mouse_day1[per_mouse_day1.pheno.isin(["PR", "DPR", "PS"])]
print("Variant, one file per mouse-day:", p1.groupby("pheno").rate.mean().round(2).to_dict(),
      f"; ANOVA p = {stats.f_oneway(*[p1[p1.pheno == g].rate for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
print(tukey(p1.rate, p1.pheno).to_string())
print("Mean rate by sex within PS (sex is unbalanced across phenotypes):", pm.join(fp_subj.set_index("subject").sex).groupby(["pheno", "sex"]).rate.agg(["mean", "size"]).round(2).to_dict("index"))

        mean   std  size   min    max
pheno                                
DPR    12.79  5.28     9  7.55  20.74
PR     14.20  3.61     9  8.44  18.59
PS      5.27  3.72    18  1.50  14.50
One-way ANOVA: p = 4.94e-06

  group1 group2  meandiff   p-adj    lower   upper  reject
0    DPR     PR    1.4102  0.7508  -3.3655  6.1860   False
1    DPR     PS   -7.5152  0.0003 -11.6511 -3.3793    True
2     PR     PS   -8.9255  0.0000 -13.0614 -4.7896    True

Variant, mixed model over 492 sessions: phenotype likelihood ratio 26.0, p = 2.3e-06
Variant with the empty files kept: {'DPR': 12.79, 'PR': 14.2, 'PS': 5.24} ; ANOVA p = 4.27e-06

  group1 group2  meandiff   p-adj    lower   upper  reject
0    DPR     PR    1.4102  0.7490  -3.3460  6.1665   False
1    DPR     PS   -7.5513  0.0002 -11.6704 -3.4323    True
2     PR     PS   -8.9616  0.0000 -13.0806 -4.8425    True
Variant, one file per mouse-day: {'DPR': 13.31, 'PR': 14.2, 'PS': 5.36} ; ANOVA p = 7.26e-06

  group1 group2  meandiff   p-adj    lower   upper  reject
0    DPR     PR    0.8860  0.9001  -4.0819  5.8538   False
1    DPR     PS   -7.9549  0.0002 -12.2572 -3.6526    True
2     PR     PS   -8.8409  0.0000 -13.1431 -4.5386    True
Mean rate by sex within PS (sex is unbalanced across phenotypes): {('DPR', 'F'): {'mean': 10.81, 'size': 4}, ('DPR', 'M'): {'mean': 14.37, 'size': 5}, ('PR', 'F'): {'mean': 8.44, 'size': 1}, ('PR', 'M'): {'mean': 14.92, 'size': 8}, ('PS', 'F'): {'mean': 5.28, 'size': 17}, ('PS', 'M'): {'mean': 5.18, 'size': 1}}

o2 = fp_el[(fp_el.schedule == "RI60")].dropna(subset=["omission_min"])
print(o2.groupby("pheno").omission_min.agg(["mean", "std", "size", "min", "max"]).round(2).to_string(), "(minutes; paper n PR 4, DPR 5, PS 11)")
an = stats.f_oneway(*[o2[o2.pheno == g].omission_min for g in ["PR", "DPR", "PS"]])
print(f"One-way ANOVA: F = {an.statistic:.2f}, p = {an.pvalue:.3g}, df = 2,{len(o2) - 3}")
print(tukey(o2.omission_min, o2.pheno).to_string())
print("Values per mouse (min):", o2.sort_values("pheno")[["pheno", "omission_min"]].round(1).to_dict("records"))

        mean    std  size    min   max
pheno                                 
DPR    31.67  11.90     5  18.72  45.5
PR     51.65  16.70     4  26.60  60.0
PS     28.06  14.83    10  17.10  60.0 (minutes; paper n PR 4, DPR 5, PS 11)
One-way ANOVA: F = 3.85, p = 0.0432, df = 2,16

  group1 group2  meandiff   p-adj    lower    upper  reject
0    DPR     PR   19.9833  0.1328  -5.1860  45.1527   False
1    DPR     PS   -3.6117  0.8936 -24.1624  16.9390   False
2     PR     PS  -23.5950  0.0365 -45.7923  -1.3977    True
Values per mouse (min): [{'pheno': 'DPR', 'omission_min': 18.7}, {'pheno': 'DPR', 'omission_min': 38.7}, {'pheno': 'DPR', 'omission_min': 35.6}, {'pheno': 'DPR', 'omission_min': 45.5}, {'pheno': 'DPR', 'omission_min': 19.8}, {'pheno': 'PR', 'omission_min': 60.0}, {'pheno': 'PR', 'omission_min': 26.6}, {'pheno': 'PR', 'omission_min': 60.0}, {'pheno': 'PR', 'omission_min': 60.0}, {'pheno': 'PS', 'omission_min': 19.2}, {'pheno': 'PS', 'omission_min': 19.9}, {'pheno': 'PS', 'omission_min': 24.4}, {'pheno': 'PS', 'omission_min': 23.8}, {'pheno': 'PS', 'omission_min': 48.4}, {'pheno': 'PS', 'omission_min': 17.7}, {'pheno': 'PS', 'omission_min': 60.0}, {'pheno': 'PS', 'omission_min': 17.1}, {'pheno': 'PS', 'omission_min': 33.0}, {'pheno': 'PS', 'omission_min': 17.2}]

print(pm.groupby("pheno").wait.agg(["mean", "std", "size", "min", "max"]).round(2).to_string())
print(f"One-way ANOVA: p = {stats.f_oneway(*[pm[pm.pheno == g].wait.dropna() for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
print(tukey(pm.wait, pm.pheno).to_string())
wm = tr[tr.pheno.isin(["PR", "DPR", "PS"])].groupby(["pheno", "subject"]).agg(
    mean_wait=("wait", "mean"), median_session_wait=("wait", "median"), max_session_wait=("wait", "max"),
    sessions=("wait", "count"), waits=("n_waits", "sum")).round(1)
print("Per mouse: mean and median of session mean waits (s), and the number of waits:")
print(wm.sort_values(["pheno", "mean_wait"]).to_string())
med = tr[tr.pheno.isin(["PR", "DPR", "PS"])].groupby(["subject", "pheno"]).wait.median().reset_index()
print("Variant, median over sessions per mouse:", med.groupby("pheno").wait.mean().round(2).to_dict())

        mean    std  size    min     max
pheno                                   
DPR    21.86   3.90     9  16.42   28.33
PR     21.05   3.55     9  15.32   25.95
PS     53.94  44.67    18  17.95  171.00
One-way ANOVA: p = 0.0174

  group1 group2  meandiff   p-adj    lower    upper  reject
0    DPR     PR   -0.8084  0.9984 -38.0156  36.3987   False
1    DPR     PS   32.0769  0.0512  -0.1454  64.2993   False
2     PR     PS   32.8854  0.0447   0.6630  65.1077    True
Per mouse: mean and median of session mean waits (s), and the number of waits:
               mean_wait  median_session_wait  max_session_wait  sessions  waits
pheno subject                                                                   
DPR   64.205        16.4                 16.6              24.0        16    761
      349.393       18.7                 15.7              49.9        13    609
      81.236        18.8                 18.8              39.1        15    655
      346.394       20.8                 18.2              41.0        12    588
      88.239        21.7                 20.0              36.6        16    614
      87.239        22.1                 23.3              31.2        11    386
      333.393       22.7                 21.2              37.4        12    567
      272.396       27.2                 24.1              44.9        12    550
      334.394       28.3                 20.6              95.1        13    606
PR    141.308       15.3                 15.8              40.4        12    488
      028.392       17.2                 14.6              47.6        15    685
      62.207        19.4                 18.4              37.7        19    866
      89.247        20.1                 21.3              30.7        14    495
      61.207        20.8                 19.3              50.2        17    761
      80.236        21.6                 23.4              31.1        12    551
      111.285       24.0                 23.4              38.9        14    622
      048.392       25.1                 18.0              70.1        13    606
      63.207        26.0                 23.6              43.4        14    551
PS    347.393       18.0                 16.9              28.7        12    573
      249.391       18.6                 17.9              37.6        12    576
      332.393       19.2                 20.0              30.8        12    552
      90.247        19.7                 19.9              26.8        13    590
      112.283       32.6                 27.6              64.8        15    515
      78.214        39.5                 39.0              62.4        12    500
      110.271       39.9                 33.7              81.3        12    421
      94.246        39.9                 37.4             102.4        15    452
      114.273       40.8                 38.6              73.3        14    596
      115.273       42.5                 26.8             235.5        14    579
      113.283       44.3                 43.8              78.3        13    402
      271.396       44.6                 32.2             159.4        13    544
      92.246        45.9                 40.1              84.0        14    518
      93.246        47.0                 39.9             114.2        13    373
      75.214        67.3                 60.6             140.9        16    461
      140.306       71.4                 22.3             655.3        13    520
      348.393      168.7                 39.2            1082.5        12    450
      139.298      171.0                 23.7            1893.0        14    375
Variant, median over sessions per mouse: {'DPR': 19.84, 'PR': 19.76, 'PS': 32.2}

print(pm.groupby("pheno").extra.agg(["mean", "std", "size", "min", "max"]).round(2).to_string())
print(f"One-way ANOVA: p = {stats.f_oneway(*[pm[pm.pheno == g].extra.dropna() for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
print(tukey(pm.extra, pm.pheno).to_string())

       mean   std  size   min   max
pheno                              
DPR    2.28  2.84     9  0.29  8.15
PR     3.33  1.81     9  0.36  5.82
PS     1.61  1.71    18  0.40  7.77
One-way ANOVA: p = 0.139

  group1 group2  meandiff   p-adj   lower   upper  reject
0    DPR     PR    1.0520  0.5317 -1.3334  3.4374   False
1    DPR     PS   -0.6702  0.7080 -2.7361  1.3956   False
2     PR     PS   -1.7222  0.1172 -3.7881  0.3436   False

iri = {g: np.concatenate(tr[(tr.pheno == g) & (tr.stage == "RI60")].iri.tolist()) for g in ["PR", "DPR", "PS"]}
for g, v in iri.items():
    print(f"{g}: mean {v.mean():.2f} s, SD {v.std(ddof=1):.2f}, median {np.median(v):.1f}, n intervals {len(v)}, range {v.min():.1f} to {v.max():.0f}")
for g in ["PR", "DPR"]:
    ks = stats.ks_2samp(iri["PS"], iri[g])
    print(f"KS PS vs {g}: D = {ks.statistic:.3f}, p = {ks.pvalue:.2g}")
mouse_iri = tr[tr.stage == "RI60"].groupby(["subject", "pheno"]).iri.apply(lambda x: np.concatenate(x.tolist()).mean()).reset_index()
print("Variant, per mouse mean interval:", mouse_iri.groupby("pheno").iri.agg(["mean", "size"]).round(1).to_dict("index"),
      f"; ANOVA p = {stats.f_oneway(*[mouse_iri[mouse_iri.pheno == g].iri for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
fig, ax = plt.subplots(figsize=(4, 3))
bins = np.arange(0, 301, 10)
for g, c in [("PR", "k"), ("DPR", "m"), ("PS", "c")]:
    ax.hist(iri[g], bins, histtype="step", density=True, color=c, label=g)
ax.set_xlabel("inter-reward interval (s)"); ax.legend(); plt.tight_layout(); plt.show()
print("Intervals above 300 s (outside the axes):", {g: int((v > 300).sum()) for g, v in iri.items()})

PR: mean 71.96 s, SD 71.39, median 60.7, n intervals 5682, range 5.4 to 1678
DPR: mean 69.57 s, SD 65.38, median 60.4, n intervals 5336, range 3.3 to 1954
PS: mean 81.75 s, SD 88.80, median 64.2, n intervals 9053, range 3.8 to 2165

KS PS vs PR: D = 0.099, p = 5e-30

KS PS vs DPR: D = 0.106, p = 3.3e-33
Variant, per mouse mean interval: {'DPR': {'mean': 70.6, 'size': 9}, 'PR': {'mean': 72.5, 'size': 9}, 'PS': {'mean': 83.7, 'size': 18}} ; ANOVA p = 0.0326

Intervals above 300 s (outside the axes): {'PR': 70, 'DPR': 46, 'PS': 184}

print(pm.groupby("pheno").rewards.agg(["mean", "std", "size", "min", "max"]).round(2).to_string())
an = stats.f_oneway(*[pm[pm.pheno == g].rewards for g in ["PR", "DPR", "PS"]])
print(f"One-way ANOVA: F = {an.statistic:.2f}, p = {an.pvalue:.3g}")
print(tukey(pm.rewards, pm.pheno).to_string())
print("Variant with the empty files kept:", pz.groupby("pheno").rewards.mean().round(2).to_dict(),
      f"; ANOVA p = {stats.f_oneway(*[pz[pz.pheno == g].rewards for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
print(tukey(pz.rewards, pz.pheno).to_string())
print("Variant, one file per mouse-day:", p1.groupby("pheno").rewards.mean().round(2).to_dict(),
      f"; ANOVA p = {stats.f_oneway(*[p1[p1.pheno == g].rewards for g in ['PR', 'DPR', 'PS']]).pvalue:.3g}")
print(tukey(p1.rewards, p1.pheno).to_string())

        mean   std  size    min    max
pheno                                 
DPR    44.98  5.39     9  35.44  50.00
PR     44.66  3.79     9  36.50  47.92
PS     39.12  7.10    18  28.86  48.92
One-way ANOVA: F = 4.02, p = 0.0275

  group1 group2  meandiff   p-adj    lower   upper  reject
0    DPR     PR   -0.3198  0.9931  -7.3072  6.6676   False
1    DPR     PS   -5.8619  0.0592 -11.9132  0.1894   False
2     PR     PS   -5.5421  0.0780 -11.5934  0.5092   False
Variant with the empty files kept: {'DPR': 44.98, 'PR': 44.66, 'PS': 38.95} ; ANOVA p = 0.0211

  group1 group2  meandiff   p-adj    lower   upper  reject
0    DPR     PR   -0.3198  0.9929  -7.2406  6.6010   False
1    DPR     PS   -6.0374  0.0480 -12.0310 -0.0439    True
2     PR     PS   -5.7176  0.0640 -11.7112  0.2759   False
Variant, one file per mouse-day: {'DPR': 46.16, 'PR': 44.66, 'PS': 39.98} ; ANOVA p = 0.033

  group1 group2  meandiff   p-adj    lower   upper  reject
0    DPR     PR   -1.4989  0.8588  -8.4795  5.4816   False
1    DPR     PS   -6.1853  0.0440 -12.2307 -0.1400    True
2     PR     PS   -4.6864  0.1540 -10.7317  1.3589   False

ri30 = []
for r in fp[fp.stage == "RI30"].itertuples():
    ev = events[r.asset_id]["events"]
    ri30.append({"subject": r.subject, "rate": len(seiler.active_pokes(ev, r.side)) / (seiler.session_length(ev) / 60.0)})
ri30 = pd.DataFrame(ri30).groupby("subject").rate.mean().to_frame("ri30_rate").join(fp_subj.set_index("subject").pheno)
slopes = tr[tr.stage == "RI60"].groupby("subject").apply(lambda d: np.polyfit(d.day, d.rate, 1)[0] if len(d) > 2 else np.nan).rename("ri60_slope")
esc = ri30.join(slopes, how="outer")
esc["pheno"] = esc.index.map(fp_subj.set_index("subject").pheno)
esc = esc[esc.pheno.isin(["PR", "DPR", "PS"])]
print(esc.groupby("pheno").agg(ri30_rate_mean=("ri30_rate", "mean"), ri30_n=("ri30_rate", "count"),
                               ri60_slope_median=("ri60_slope", "median"), ri60_slope_mean=("ri60_slope", "mean"), ri60_n=("ri60_slope", "count")).round(3).to_string())
for g in ["DPR", "PS"]:
    p = stats.mannwhitneyu(esc[esc.pheno == "PR"].ri30_rate.dropna(), esc[esc.pheno == g].ri30_rate.dropna()).pvalue
    print(f"RI30 rate PR vs {g}: Mann-Whitney p = {p:.3g}")
p = stats.mannwhitneyu(esc[esc.pheno == "DPR"].ri60_slope.dropna(), esc[esc.pheno == "PR"].ri60_slope.dropna()).pvalue
print(f"RI60 slope DPR vs PR: Mann-Whitney p = {p:.3g}")
fig, ax = plt.subplots(figsize=(5, 3))
for g, c in [("PR", "k"), ("DPR", "m"), ("PS", "c")]:
    d = tr[(tr.pheno == g) & (tr.stage == "RI60")].groupby("day").rate.mean()
    ax.plot(d.index, d.values, c, label=g)
ax.set_xlabel("RI60 day"); ax.set_ylabel("active nosepokes / min"); ax.legend(); plt.tight_layout(); plt.show()
print("Plotted day means span", tr[tr.stage == "RI60"].groupby(["pheno", "day"]).rate.mean().min().round(2), "to", tr[tr.stage == "RI60"].groupby(["pheno", "day"]).rate.mean().max().round(2))

       ri30_rate_mean  ri30_n  ri60_slope_median  ri60_slope_mean  ri60_n
pheno                                                                    
DPR             3.032       9              0.809            0.943       9
PR              4.850       8              0.494            0.557       9
PS              2.216      18              0.098            0.212      18
RI30 rate PR vs DPR: Mann-Whitney p = 0.114
RI30 rate PR vs PS: Mann-Whitney p = 0.00627
RI60 slope DPR vs PR: Mann-Whitney p = 0.289

Plotted day means span 1.82 to 24.75

pmr = per_mouse.assign(schedule=np.where(per_mouse.pheno == "RR20", "RR20", "RI60"))
print(pmr.groupby("schedule").rewards.agg(["mean", "std", "size", "min", "max"]).round(2).to_string())
t2 = stats.ttest_ind(pmr[pmr.schedule == "RI60"].rewards, pmr[pmr.schedule == "RR20"].rewards)
print(f"Unpaired t-test: t = {t2.statistic:.2f}, p = {t2.pvalue:.2g}")
pmz = per_mouse_all.assign(schedule=np.where(per_mouse_all.pheno == "RR20", "RR20", "RI60"))
print("Variant with the empty files kept:", pmz.groupby("schedule").rewards.mean().round(2).to_dict(),
      f"; t-test p = {stats.ttest_ind(pmz[pmz.schedule == 'RI60'].rewards, pmz[pmz.schedule == 'RR20'].rewards).pvalue:.2g}")
pmr["rate"] = per_mouse.rate
print("Active nosepokes per minute by schedule (Fig S1G, similar rates):", pmr.groupby("schedule").rate.agg(["mean", "std"]).round(2).to_dict("index"))

           mean    std  size    min    max
schedule                                  
RI60      41.97   6.54    36  28.86  50.00
RR20      22.63  11.55     6   6.86  42.14
Unpaired t-test: t = 5.96, p = 5.4e-07
Variant with the empty files kept: {'RI60': 41.89, 'RR20': 15.38} ; t-test p = 1.4e-11
Active nosepokes per minute by schedule (Fig S1G, similar rates): {'RI60': {'mean': 9.39, 'std': 5.81}, 'RR20': {'mean': 7.92, 'std': 3.97}}

aov3, w3 = mixed_anova(ri_long, "sex")
print(fmt(aov3))
print(w3.groupby("sex")[["early", "late"]].agg(["mean", "std", "size"]).round(2).to_string())
print("Sex by phenotype (RI60, all mice):", pd.crosstab(fp_subj[fp_subj.behavior == "RI60"].pheno, fp_subj[fp_subj.behavior == "RI60"].sex).to_dict())

                F    df       p
sex          9.81  1,30  0.0039
probe        4.96  1,30   0.034
sex x probe  0.37  1,30    0.55
probe  early              late            
        mean   std size   mean    std size
sex                                       
F       8.25  4.14   20  12.15  11.90   20
M      16.58  9.65   12  23.25  15.52   12
Sex by phenotype (RI60, all mice): {'F': {'DPR': 4, 'PR': 1, 'PS': 17}, 'M': {'DPR': 5, 'PR': 8, 'PS': 1}}

# One score per mouse and stage: the mean over its sessions in that stage (one mouse has two early sessions).
sc = ph[ph.pheno.isin(["PR", "DPR", "PS"])].groupby(["subject", "pheno", "fp_stage"])[
    ["DMS_score", "DLS_score", "DMS_score_nomask", "DLS_score_nomask"]].mean().reset_index()
print("Mice per phenotype and stage with a DMS score:")
print(sc.dropna(subset=["DMS_score"]).groupby(["pheno", "fp_stage"]).size().unstack().to_string())
print("Mice per phenotype and stage with a DLS score:")
print(sc.dropna(subset=["DLS_score"]).groupby(["pheno", "fp_stage"]).size().unstack().to_string())


def phenotype_model(df, col):
    d = df.dropna(subset=[col]).rename(columns={col: "y"})
    full = smf.mixedlm("y ~ C(pheno) * C(fp_stage)", d, groups=d.subject).fit(reml=False)
    no_ph = smf.mixedlm("y ~ C(fp_stage)", d, groups=d.subject).fit(reml=False)
    no_stage = smf.mixedlm("y ~ C(pheno)", d, groups=d.subject).fit(reml=False)
    lr_p = 2 * (full.llf - no_ph.llf)
    lr_s = 2 * (full.llf - no_stage.llf)
    out = {"n_obs": len(d), "n_mice": d.subject.nunique(),
           "phenotype_LR": lr_p, "phenotype_p": stats.chi2.sf(lr_p, 4),
           "stage_LR": lr_s, "stage_p": stats.chi2.sf(lr_s, 3)}
    for st in ["early", "late"]:
        s = d[d.fp_stage == st]
        out[f"anova_{st}_p"] = stats.f_oneway(*[s[s.pheno == g].y for g in ["PR", "DPR", "PS"] if (s.pheno == g).sum() > 1]).pvalue
    return out, d.groupby(["pheno", "fp_stage"]).y.agg(["mean", "std", "size"])

Mice per phenotype and stage with a DMS score:
fp_stage  early  late
pheno                
DPR           9     7
PR            7     4
PS           14    10
Mice per phenotype and stage with a DLS score:
fp_stage  early  late
pheno                
DPR           9     7
PR            6     4
PS           14    10

res3f, tab3f = phenotype_model(sc, "DMS_score")
print(tab3f.round(3).unstack().to_string())
print({k: (round(v, 4) if isinstance(v, float) else v) for k, v in res3f.items()})
print("Means over both stages:", sc.groupby("pheno").DMS_score.mean().round(3).to_dict())
d = sc.dropna(subset=["DMS_score"])
for st in ["early", "late"]:
    s = d[d.fp_stage == st]
    print(st, "Tukey:"); print(tukey(s.DMS_score, s.pheno).to_string())
d2 = d.pivot(index="subject", columns="fp_stage", values="DMS_score").dropna()
print(f"Early vs late, paired over {len(d2)} mice: mean change {np.mean(d2.late - d2.early):.3f}, paired t p = {stats.ttest_rel(d2.late, d2.early).pvalue:.3g}")
resn, tabn = phenotype_model(sc, "DMS_score_nomask")
print("Variant, the authors' processing without the dropout step:", {k: round(v, 4) for k, v in resn.items() if "p" in k})
print("Per-mouse scores span", round(d.DMS_score.min(), 2), "to", round(d.DMS_score.max(), 2))
fig, ax = plt.subplots(figsize=(4, 3))
for i, g in enumerate(["PR", "DPR", "PS"]):
    for j, st in enumerate(["early", "late"]):
        v = d[(d.pheno == g) & (d.fp_stage == st)].DMS_score
        ax.scatter(np.full(len(v), i * 3 + j) + np.random.default_rng(0).uniform(-.15, .15, len(v)), v, s=10, c="k" if st == "early" else "r")
        ax.plot([i * 3 + j - .3, i * 3 + j + .3], [v.mean()] * 2, "b")
ax.set_xticks([0.5, 3.5, 6.5], ["PR", "DPR", "PS"]); ax.set_ylabel("DMS score (z); black early, red late"); plt.tight_layout(); plt.show()

           mean           std         size     
fp_stage  early   late  early   late early late
pheno                                          
DPR       1.483  1.879  1.385  1.932     9    7
PR        2.106  2.784  1.880  1.300     7    4
PS        0.900  1.251  0.636  1.350    14   10
{'n_obs': 51, 'n_mice': 34, 'phenotype_LR': np.float64(7.7843), 'phenotype_p': np.float64(0.0998), 'stage_LR': np.float64(6.5535), 'stage_p': np.float64(0.0876), 'anova_early_p': np.float64(0.1242), 'anova_late_p': np.float64(0.2696)}
Means over both stages: {'DPR': 1.656, 'PR': 2.352, 'PS': 1.047}
early Tukey:

  group1 group2  meandiff   p-adj   lower   upper  reject
0    DPR     PR    0.6224  0.5880 -0.9329  2.1776   False
1    DPR     PS   -0.5831  0.5244 -1.9016  0.7354   False
2     PR     PS   -1.2055  0.1104 -2.6340  0.2231   False
late Tukey:

  group1 group2  meandiff   p-adj   lower   upper  reject
0    DPR     PR    0.9050  0.6320 -1.5922  3.4022   False
1    DPR     PS   -0.6275  0.6984 -2.5909  1.3360   False
2     PR     PS   -1.5325  0.2476 -3.8895  0.8246   False
Early vs late, paired over 17 mice: mean change 0.490, paired t p = 0.0458

Variant, the authors' processing without the dropout step: {'phenotype_LR': np.float64(4.1299), 'phenotype_p': np.float64(0.3887), 'stage_p': np.float64(0.2742), 'anova_early_p': np.float64(0.1879), 'anova_late_p': np.float64(0.5631)}
Per-mouse scores span -0.12 to 5.69

d3e = sc.dropna(subset=["DMS_score"]).pivot(index=["subject", "pheno"], columns="fp_stage", values="DMS_score").reset_index()
print("DMS score by phenotype and stage (all mice with that stage):", sc.groupby(["pheno", "fp_stage"]).DMS_score.mean().round(2).unstack().to_dict("index"))
for g in ["PR", "DPR", "PS"]:
    x = d3e[(d3e.pheno == g)].dropna(subset=["early", "late"])
    if len(x) > 1:
        print(f"{g}: {len(x)} mice with both stages, early {x.early.mean():.2f}, late {x.late.mean():.2f}, change {np.mean(x.late - x.early):.2f}, paired t p = {stats.ttest_rel(x.late, x.early).pvalue:.2g}")
    else:
        print(f"{g}: {len(x)} mice with both stages")

DMS score by phenotype and stage (all mice with that stage): {'DPR': {'early': 1.48, 'late': 1.88}, 'PR': {'early': 2.11, 'late': 2.78}, 'PS': {'early': 0.9, 'late': 1.25}}
PR: 3 mice with both stages, early 1.68, late 2.85, change 1.17, paired t p = 0.19
DPR: 7 mice with both stages, early 1.55, late 1.88, change 0.33, paired t p = 0.3
PS: 7 mice with both stages, early 0.77, late 1.13, change 0.36, paired t p = 0.4

pairs = sc.merge(fp_probes[fp_probes.probe != ""][["subject", "probe", "shocks"]].rename(columns={"probe": "fp_stage"}), on=["subject", "fp_stage"])
def corr(df, x, y="shocks"):
    d = df.dropna(subset=[x, y])
    r = stats.pearsonr(d[x], d[y])
    return f"r = {r.statistic:.3f}, p = {r.pvalue:.2g}, n = {len(d)}"
print("(a) pooled early and late:", corr(pairs, "DMS_score"))
print("(b) late only:            ", corr(pairs[pairs.fp_stage == "late"], "DMS_score"))
print("(c) early only:           ", corr(pairs[pairs.fp_stage == "early"], "DMS_score"))
print("Spearman, pooled:", stats.spearmanr(pairs.dropna(subset=["DMS_score"]).DMS_score, pairs.dropna(subset=["DMS_score"]).shocks))
pc = stats.pearsonr(*pairs[pairs.fp_stage == "early"].dropna(subset=["DMS_score"])[["DMS_score", "shocks"]].T.values).pvalue
print(f"Reading (c) with a Bonferroni correction for the three readings: p = {min(1, 3 * pc):.2g}")
print("Variant without the dropout step, pooled:", corr(pairs, "DMS_score_nomask"))
print("Variant without the dropout step, late only: ", corr(pairs[pairs.fp_stage == "late"], "DMS_score_nomask"))
print("Variant without the dropout step, early only:", corr(pairs[pairs.fp_stage == "early"], "DMS_score_nomask"))
used = ph.merge(pairs[["subject", "fp_stage"]], on=["subject", "fp_stage"])
print(f"\nPercent of samples left out by the dropout step in the {len(used)} sessions behind these pairs:")
for region in ["DMS", "DLS"]:
    fr_ = 100 * used[f"{region}_dropout_frac"]
    print(f"  {region}: median {fr_.median():.3f}, above 0.1% in {(fr_ > 0.1).sum()}, above 1% in {(fr_ > 1).sum()}, maximum {fr_.max():.2f}")
print(used.loc[(used.DMS_dropout_frac > 0.001) | (used.DLS_dropout_frac > 0.001),
               ["subject", "date", "fp_stage", "DMS_dropout_frac", "DLS_dropout_frac", "DMS_max_abs_z_nomask", "DLS_max_abs_z_nomask",
                "DMS_score", "DMS_score_nomask", "DLS_score", "DLS_score_nomask"]]
      .assign(DMS_dropout_frac=lambda x: 100 * x.DMS_dropout_frac, DLS_dropout_frac=lambda x: 100 * x.DLS_dropout_frac).round(2).to_string(index=False))
rrp = ph[ph.pheno == "RR20"].groupby(["subject", "fp_stage"]).DMS_score.mean().reset_index().merge(
    fp_probes[fp_probes.probe != ""][["subject", "probe", "shocks"]].rename(columns={"probe": "fp_stage"}), on=["subject", "fp_stage"])
print("Variant with RR20 mice added, pooled:", corr(pd.concat([pairs, rrp]), "DMS_score"))
print("Variant, each mouse's mean score against its mean shocks:", corr(pairs.groupby("subject")[["DMS_score", "shocks"]].mean().reset_index(), "DMS_score"))
fig, ax = plt.subplots(figsize=(3.5, 3))
for st, c in [("early", "k"), ("late", "r")]:
    s = pairs[pairs.fp_stage == st]
    ax.scatter(s.DMS_score, s.shocks, c=c, s=12, label=st)
ax.set_xlabel("DMS score (z)"); ax.set_ylabel("shocks"); ax.legend(); plt.tight_layout(); plt.show()
print("Plotted points span score", round(pairs.DMS_score.min(), 2), "to", round(pairs.DMS_score.max(), 2), "and shocks", pairs.shocks.min(), "to", pairs.shocks.max())

(a) pooled early and late: r = 0.170, p = 0.24, n = 49
(b) late only:             r = 0.014, p = 0.95, n = 20
(c) early only:            r = 0.375, p = 0.045, n = 29
Spearman, pooled: SignificanceResult(statistic=np.float64(0.156707623422631), pvalue=np.float64(0.2822423150071741))
Reading (c) with a Bonferroni correction for the three readings: p = 0.13
Variant without the dropout step, pooled: r = 0.100, p = 0.49, n = 49
Variant without the dropout step, late only:  r = -0.082, p = 0.73, n = 20
Variant without the dropout step, early only: r = 0.362, p = 0.053, n = 29

Percent of samples left out by the dropout step in the 51 sessions behind these pairs:
  DMS: median 0.033, above 0.1% in 9, above 1% in 1, maximum 1.61
  DLS: median 0.033, above 0.1% in 9, above 1% in 1, maximum 99.97
subject       date fp_stage  DMS_dropout_frac  DLS_dropout_frac  DMS_max_abs_z_nomask  DLS_max_abs_z_nomask  DMS_score  DMS_score_nomask  DLS_score  DLS_score_nomask
 63.207 2018-10-15    early              0.32              0.32                 11.07                  9.10       4.65              4.28       1.48              1.38
 80.236 2019-01-21    early              0.57              0.11                 10.18                 73.44       2.35              1.71       0.48              0.40
 88.239 2019-02-19    early              1.61              0.03               1521.97                  6.98       1.92              0.09       0.50              0.50
 92.246 2019-02-27    early              0.70              0.30               1484.95               1666.64       0.30              0.04       2.27              0.01
 87.239 2019-02-28    early              0.15              0.21                 15.85                 37.76       3.40              3.15       1.27              0.16
141.308 2019-08-09    early              0.04             99.97                  5.00                  3.28       2.68              2.68        NaN               NaN
140.306 2019-09-03     late              0.03              0.21                  5.17               1156.26       3.95              3.95       0.97              0.02
139.298 2019-09-12     late              0.11              0.11               1729.70               1758.49      -0.06             -0.00       0.51              0.02
346.394 2020-07-22     late              0.55              0.49               1073.11               1609.37       1.10              0.00       3.31              0.03
332.393 2020-07-28     late              0.15              0.15                959.31                 90.13        NaN               NaN        NaN               NaN
048.392 2020-07-28     late              0.12              0.04               1681.12                  4.70       1.81              0.02       4.44              4.44
Variant with RR20 mice added, pooled: r = 0.116, p = 0.39, n = 58
Variant, each mouse's mean score against its mean shocks: r = 0.210, p = 0.24, n = 33

Plotted points span score -0.12 to 5.69 and shocks 1 to 53

def r_ci(df, x, y):
    d = df.dropna(subset=[x, y])
    r = stats.pearsonr(d[x], d[y])
    ci = r.confidence_interval(0.95)
    return f"95% interval of r {ci.low:.2f} to {ci.high:.2f} (n = {len(d)})"

oms = sc.merge(fp_el[["omission_min"]].reset_index(), on="subject")
print("Late score vs omission time:", corr(oms[oms.fp_stage == "late"], "DMS_score", "omission_min"))
print("  ", r_ci(oms[oms.fp_stage == "late"], "DMS_score", "omission_min"))
print("Variant, pooled early and late:", corr(oms, "DMS_score", "omission_min"))
print("Variant, early score:", corr(oms[oms.fp_stage == "early"], "DMS_score", "omission_min"))

Late score vs omission time: r = 0.014, p = 0.96, n = 12
   95% interval of r -0.56 to 0.58 (n = 12)
Variant, pooled early and late: r = -0.050, p = 0.8, n = 29
Variant, early score: r = -0.092, p = 0.72, n = 17

res4f, tab4f = phenotype_model(sc, "DLS_score")
print(tab4f.round(3).unstack().to_string())
print({k: (round(v, 4) if isinstance(v, float) else v) for k, v in res4f.items()})
resn4, _ = phenotype_model(sc, "DLS_score_nomask")
print("Variant, without the dropout step:", {k: round(v, 4) for k, v in resn4.items() if "p" in k})

           mean           std         size     
fp_stage  early   late  early   late early late
pheno                                          
DPR       1.036  2.006  0.478  1.213     9    7
PR        1.248  1.794  0.904  1.900     6    4
PS        0.927  1.305  0.640  0.757    14   10
{'n_obs': 50, 'n_mice': 33, 'phenotype_LR': np.float64(3.475), 'phenotype_p': np.float64(0.4817), 'stage_LR': np.float64(8.6714), 'stage_p': np.float64(0.034), 'anova_early_p': np.float64(0.6118), 'anova_late_p': np.float64(0.4726)}

Variant, without the dropout step: {'phenotype_LR': np.float64(2.3049), 'phenotype_p': np.float64(0.6799), 'stage_p': np.float64(0.2536), 'anova_early_p': np.float64(0.4868), 'anova_late_p': np.float64(0.6589)}

print("(a) pooled early and late:", corr(pairs, "DLS_score"))
print("Variant, late only:", corr(pairs[pairs.fp_stage == "late"], "DLS_score"))
print("Variant, early only:", corr(pairs[pairs.fp_stage == "early"], "DLS_score"))
print("Variant without the dropout step, pooled:", corr(pairs, "DLS_score_nomask"))
print("Variant without the dropout step, late only: ", corr(pairs[pairs.fp_stage == "late"], "DLS_score_nomask"))
print("Variant without the dropout step, early only:", corr(pairs[pairs.fp_stage == "early"], "DLS_score_nomask"))

(a) pooled early and late: r = 0.326, p = 0.024, n = 48
Variant, late only: r = 0.374, p = 0.1, n = 20
Variant, early only: r = 0.103, p = 0.6, n = 28
Variant without the dropout step, pooled: r = 0.113, p = 0.45, n = 48
Variant without the dropout step, late only:  r = 0.033, p = 0.89, n = 20
Variant without the dropout step, early only: r = 0.182, p = 0.35, n = 28

print("Late score vs omission time:", corr(oms[oms.fp_stage == "late"], "DLS_score", "omission_min"))
print("  ", r_ci(oms[oms.fp_stage == "late"], "DLS_score", "omission_min"))
print("Variant, pooled early and late:", corr(oms, "DLS_score", "omission_min"))

Late score vs omission time: r = 0.311, p = 0.32, n = 12
   95% interval of r -0.32 to 0.75 (n = 12)
Variant, pooled early and late: r = 0.141, p = 0.47, n = 28

def opto_tables(experiment):
    s = sessions[sessions.subject_experiment == experiment]
    probes = seiler.probe_table(s, events, training_stages=("RI60",))
    el = seiler.early_late(probes).join(s.drop_duplicates("subject").set_index("subject")[["subject_arm", "sex"]], how="right")
    el["pheno"] = [seiler.classify(e, l) if pd.notna(e) and pd.notna(l) else None for e, l in zip(el.early, el.late)]
    omf = s[s.stage == "omission"].sort_values("start_ts").groupby("subject").asset_id.first()
    el["omission_min"] = [seiler.omission_time_min(events[omf[x]]["events"]) if x in omf.index else np.nan for x in el.index]
    el["n_ri60"] = s[s.stage == "RI60"].drop_duplicates(["subject", "start"]).groupby("subject").size().reindex(el.index).fillna(0).astype(int)
    el["did_not_learn"] = el.index.map(opto_subj.set_index("subject").did_not_learn)
    return s, probes, el

opto = {e: opto_tables(e) for e in ["Opto-DMS-Excitatory", "Opto-DLS-Excitatory", "Opto-DMS-Inhibitory"]}
for e, (s, probes, el) in opto.items():
    print(f"== {e}: {len(el)} mice")
    print("Mice per arm and sex:", el.groupby(["subject_arm", "sex"]).size().unstack(fill_value=0).to_dict("index"))
    print("Probe days per mouse:", probes.groupby("subject").size().value_counts().sort_index().to_dict(),
          "; days with two files holding shocks:", int((probes.n_with_shocks > 1).sum()))
    print("Training days before early / late probe:", probes[probes.probe == "early"].train_days_before.describe()[["min", "50%", "max"]].to_dict(),
          probes[probes.probe == "late"].train_days_before.describe()[["min", "50%", "max"]].to_dict())
    print("Mice with both probes:", int(el[["early", "late"]].notna().all(axis=1).sum()), "; with an omission probe:", int(el.omission_min.notna().sum()))
    print("Mice missing a probe:", el[el[["early", "late"]].isna().any(axis=1)][["subject_arm", "early", "late", "n_ri60", "did_not_learn"]].to_dict("index"))

== Opto-DMS-Excitatory: 33 mice
Mice per arm and sex: {'ChR2': {'F': 7, 'M': 7}, 'ChR2Scrambled': {'F': 6, 'M': 6}, 'EYFP': {'F': 3, 'M': 4}}
Probe days per mouse: {2: 32} ; days with two files holding shocks: 1
Training days before early / late probe: {'min': 0.0, '50%': 1.0, 'max': 2.0} {'min': 10.0, '50%': 12.0, 'max': 14.0}
Mice with both probes: 32 ; with an omission probe: 32
Mice missing a probe: {'126.286': {'subject_arm': 'ChR2', 'early': nan, 'late': nan, 'n_ri60': 0, 'did_not_learn': 'False'}}
== Opto-DLS-Excitatory: 35 mice
Mice per arm and sex: {'ChR2': {'F': 5, 'M': 6}, 'ChR2Scrambled': {'F': 6, 'M': 6}, 'EYFP': {'F': 6, 'M': 6}}
Probe days per mouse: {1: 4, 2: 31} ; days with two files holding shocks: 1
Training days before early / late probe: {'min': 0.0, '50%': 1.0, 'max': 2.0} {'min': 10.0, '50%': 12.0, 'max': 24.0}
Mice with both probes: 31 ; with an omission probe: 32
Mice missing a probe: {'128.424': {'subject_arm': 'ChR2Scrambled', 'early': 40.0, 'late': nan, 'n_ri60': 12, 'did_not_learn': 'False'}, '238.388': {'subject_arm': 'EYFP', 'early': 12.0, 'late': nan, 'n_ri60': 12, 'did_not_learn': 'False'}, '313.403': {'subject_arm': 'EYFP', 'early': 8.0, 'late': nan, 'n_ri60': 13, 'did_not_learn': 'False'}, '342.400': {'subject_arm': 'EYFP', 'early': 23.0, 'late': nan, 'n_ri60': 14, 'did_not_learn': 'False'}}
== Opto-DMS-Inhibitory: 42 mice
Mice per arm and sex: {'EYFP': {'F': 9, 'M': 7}, 'NpHR': {'F': 8, 'M': 5}, 'NpHRScrambled': {'F': 6, 'M': 7}}
Probe days per mouse: {1: 5, 2: 32, 3: 2} ; days with two files holding shocks: 0
Training days before early / late probe: {'min': 1.0, '50%': 1.0, 'max': 4.0} {'min': 10.0, '50%': 12.0, 'max': 14.0}
Mice with both probes: 34 ; with an omission probe: 37
Mice missing a probe: {'112.415': {'subject_arm': 'NpHR', 'early': 10.0, 'late': nan, 'n_ri60': 13, 'did_not_learn': 'False'}, '202.465': {'subject_arm': 'NpHR', 'early': 6.0, 'late': nan, 'n_ri60': 18, 'did_not_learn': 'False'}, '244.465': {'subject_arm': 'EYFP', 'early': nan, 'late': nan, 'n_ri60': 8, 'did_not_learn': 'True'}, '245.464': {'subject_arm': 'NpHRScrambled', 'early': 5.0, 'late': nan, 'n_ri60': 16, 'did_not_learn': 'False'}, '259.478': {'subject_arm': 'NpHRScrambled', 'early': 10.0, 'late': nan, 'n_ri60': 14, 'did_not_learn': 'False'}, '262.259': {'subject_arm': 'NpHRScrambled', 'early': nan, 'late': nan, 'n_ri60': 1, 'did_not_learn': 'False'}, '266.477': {'subject_arm': 'NpHR', 'early': nan, 'late': 3.0, 'n_ri60': 15, 'did_not_learn': 'False'}, '270.476': {'subject_arm': 'NpHR', 'early': nan, 'late': nan, 'n_ri60': 0, 'did_not_learn': 'True'}}

def escalation(experiment, ref="ChR2"):
    s = sessions[(sessions.subject_experiment == experiment) & sessions.stage.isin(["FR1", "RI30"])].sort_values("start_ts")
    s = s[~((s.arm == "Unknown") & s.duplicated(["subject", "start"], keep=False))].drop_duplicates(["subject", "start"])
    rows = []
    for r in s.itertuples():
        ev = events[r.asset_id]["events"]
        if isinstance(r.side, str):
            pokes = seiler.active_pokes(ev, r.side)
        else:  # FR1 on both nosepokes
            pokes = np.concatenate([seiler.active_pokes(ev, "left"), seiler.active_pokes(ev, "right")])
        rows.append({"subject": r.subject, "arm": r.subject_arm, "stage": r.stage, "rate": len(pokes) / (seiler.session_length(ev) / 60.0)})
    d = pd.DataFrame(rows)
    d["idx"] = d.groupby(["subject", "stage"]).cumcount() + 1
    out = {}
    for st in ["FR1", "RI30"]:
        x = d[d.stage == st]
        fit = smf.ols(f"rate ~ idx * C(arm, Treatment('{ref}'))", x).fit()
        slopes = {a: np.polyfit(x[x.arm == a].idx, x[x.arm == a].rate, 1)[0] for a in sorted(x.arm.unique())}
        inter = {k.split("T.")[1].rstrip("]"): (fit.params[k], fit.pvalues[k]) for k in fit.params.index if k.startswith("idx:")}
        others = [a for a in sorted(x.arm.unique()) if a != ref]
        f2 = smf.ols(f"rate ~ idx * C(arm, Treatment('{others[0]}'))", x).fit()
        k2 = [k for k in f2.params.index if k.startswith("idx:") and others[1] in k][0]
        out[st] = {"slopes": slopes, "interaction_vs_ref": inter, "sessions": len(x), "mice": x.subject.nunique(),
                   "other_pair": (f"{others[1]} vs {others[0]}", f2.params[k2], f2.pvalues[k2]),
                   "sessions_per_mouse": (x.groupby("subject").size().min(), x.groupby("subject").size().max())}
    return d, out


def escalation_mixed(d, ref="ChR2", re_formula="~idx"):
    """Variant: a mixed model with a random intercept and slope per mouse (or, with re_formula=None, an intercept only).
    The fit's warnings are returned with it, as they bear on how far its p values can be trusted."""
    out = {}
    for st in ["FR1", "RI30"]:
        x = d[d.stage == st]
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            try:
                m = smf.mixedlm(f"rate ~ idx * C(arm, Treatment('{ref}'))", x, groups=x.subject, re_formula=re_formula).fit(reml=False)
                out[st] = {k.split("T.")[1].rstrip("]"): (round(m.params[k], 3), round(m.pvalues[k], 4)) for k in m.params.index if k.startswith("idx:")}
                out[st]["converged"] = m.converged
                out[st]["random effects covariance"] = np.round(np.asarray(m.cov_re), 4).tolist()
                others = [a for a in sorted(x.arm.unique()) if a != ref]
                m2 = smf.mixedlm(f"rate ~ idx * C(arm, Treatment('{others[0]}'))", x, groups=x.subject, re_formula=re_formula).fit(reml=False)
                k2 = [k for k in m2.params.index if k.startswith("idx:") and others[1] in k][0]
                out[st][f"{others[1]} vs {others[0]}"] = (round(m2.params[k2], 3), round(m2.pvalues[k2], 4))
            except Exception as e:  # report a failed fit, do not hide it
                out[st] = {"failed": repr(e)}
        msgs = pd.Series([str(w.message).split("\n")[0] for w in caught if not issubclass(w.category, (RuntimeWarning, FutureWarning))])
        out[st]["warnings"] = msgs.value_counts().to_dict()
    return out

esc5, res5 = escalation("Opto-DMS-Excitatory")
for st, r in res5.items():
    print(st, f"{r['sessions']} sessions from {r['mice']} mice ({r['sessions_per_mouse'][0]} to {r['sessions_per_mouse'][1]} per mouse)")
    print("  slope (nosepokes/min per session):", {k: round(v, 3) for k, v in r["slopes"].items()})
    print("  arm x index interaction against ChR2 (difference in slope, p):", {k: (round(v[0], 3), round(v[1], 4)) for k, v in r["interaction_vs_ref"].items()})
for st, r in escalation_mixed(esc5).items():
    print(f"Variant, {st} mixed model with a random intercept and slope per mouse (difference in slope against ChR2, and the third pair, p):", r)
for st, r in escalation_mixed(esc5, re_formula=None).items():
    print(f"Variant, {st} mixed model with a random intercept per mouse only:", r)
print("Mean rate by stage and arm:", esc5.groupby(["stage", "arm"]).rate.mean().round(2).to_dict())
per_mouse_slope = esc5.groupby(["stage", "arm", "subject"]).apply(lambda x: np.polyfit(x.idx, x.rate, 1)[0] if len(x) > 1 else np.nan).rename("slope").reset_index()
print("Variant, per-mouse slopes (median):", per_mouse_slope.groupby(["stage", "arm"]).slope.median().round(3).to_dict())
for st in ["FR1", "RI30"]:
    x = per_mouse_slope[per_mouse_slope.stage == st]
    for a in ["EYFP", "ChR2Scrambled"]:
        print(f"  {st} ChR2 vs {a}, Mann-Whitney on per-mouse slopes p = {stats.mannwhitneyu(x[x.arm == 'ChR2'].slope.dropna(), x[x.arm == a].slope.dropna()).pvalue:.3g}")

FR1 170 sessions from 33 mice (1 to 13 per mouse)
  slope (nosepokes/min per session): {'ChR2': np.float64(0.011), 'ChR2Scrambled': np.float64(0.051), 'EYFP': np.float64(0.005)}
  arm x index interaction against ChR2 (difference in slope, p): {'ChR2Scrambled': (np.float64(0.04), np.float64(0.6292)), 'EYFP': (np.float64(-0.006), np.float64(0.9797))}
RI30 52 sessions from 29 mice (1 to 4 per mouse)
  slope (nosepokes/min per session): {'ChR2': np.float64(1.487), 'ChR2Scrambled': np.float64(0.957), 'EYFP': np.float64(4.529)}
  arm x index interaction against ChR2 (difference in slope, p): {'ChR2Scrambled': (np.float64(-0.531), np.float64(0.8233)), 'EYFP': (np.float64(3.042), np.float64(0.4897))}

Variant, FR1 mixed model with a random intercept and slope per mouse (difference in slope against ChR2, and the third pair, p): {'ChR2Scrambled': (np.float64(-0.419), np.float64(0.0343)), 'EYFP': (np.float64(-0.697), np.float64(0.0075)), 'converged': True, 'random effects covariance': [[0.0001, 0.0009], [0.0009, 0.1644]], 'EYFP vs ChR2Scrambled': (np.float64(-0.277), np.float64(0.2662)), 'warnings': {'Maximum Likelihood optimization failed to converge. Check mle_retvals': 2, 'Retrying MixedLM optimization with lbfgs': 2, 'The MLE may be on the boundary of the parameter space.': 2}}
Variant, RI30 mixed model with a random intercept and slope per mouse (difference in slope against ChR2, and the third pair, p): {'ChR2Scrambled': (np.float64(-0.954), np.float64(0.6817)), 'EYFP': (np.float64(1.128), np.float64(0.7001)), 'converged': True, 'random effects covariance': [[1.3576, -3.8858], [-3.8858, 11.122]], 'EYFP vs ChR2Scrambled': (np.float64(2.082), np.float64(0.5117)), 'warnings': {}}
Variant, FR1 mixed model with a random intercept per mouse only: {'ChR2Scrambled': (np.float64(-0.056), np.float64(0.4634)), 'EYFP': (np.float64(-0.175), np.float64(0.3261)), 'converged': True, 'random effects covariance': [[0.8636]], 'EYFP vs ChR2Scrambled': (np.float64(-0.119), np.float64(0.4963)), 'warnings': {}}
Variant, RI30 mixed model with a random intercept per mouse only: {'ChR2Scrambled': (np.float64(-0.966), np.float64(0.5634)), 'EYFP': (np.float64(2.052), np.float64(0.4769)), 'converged': True, 'random effects covariance': [[13.4514]], 'EYFP vs ChR2Scrambled': (np.float64(3.018), np.float64(0.3222)), 'warnings': {}}
Mean rate by stage and arm: {('FR1', 'ChR2'): 1.59, ('FR1', 'ChR2Scrambled'): 1.22, ('FR1', 'EYFP'): 1.46, ('RI30', 'ChR2'): 7.75, ('RI30', 'ChR2Scrambled'): 4.01, ('RI30', 'EYFP'): 4.21}
Variant, per-mouse slopes (median): {('FR1', 'ChR2'): 0.657, ('FR1', 'ChR2Scrambled'): 0.275, ('FR1', 'EYFP'): -0.021, ('RI30', 'ChR2'): 0.715, ('RI30', 'ChR2Scrambled'): 0.996, ('RI30', 'EYFP'): 1.417}
  FR1 ChR2 vs EYFP, Mann-Whitney on per-mouse slopes p = 0.0145
  FR1 ChR2 vs ChR2Scrambled, Mann-Whitney on per-mouse slopes p = 0.0208
  RI30 ChR2 vs EYFP, Mann-Whitney on per-mouse slopes p = 0.6
  RI30 ChR2 vs ChR2Scrambled, Mann-Whitney on per-mouse slopes p = 1

_, p5, el5 = opto["Opto-DMS-Excitatory"]
long5 = p5[p5.probe != ""].merge(el5[["subject_arm"]].reset_index(), on="subject")
aov5, w5 = mixed_anova(long5, "subject_arm")
print(fmt(aov5))
print(w5.groupby("subject_arm")[["early", "late"]].agg(["mean", "std", "size"]).round(2).to_string())
print("First probe, Tukey:"); print(tukey(w5.early, w5.subject_arm).to_string())
print("Second probe, Tukey:"); print(tukey(w5.late, w5.subject_arm).to_string())
e1 = el5.dropna(subset=["early"])
print("Variant, all mice with a first probe:", e1.groupby("subject_arm").early.agg(["mean", "size"]).round(2).to_dict("index"))
print(tukey(e1.early, e1.subject_arm).to_string())

                        F    df     p
subject_arm          0.25  2,29  0.78
probe                0.00  1,29  0.99
subject_arm x probe  1.95  2,29  0.16
probe          early               late            
                mean    std size   mean    std size
subject_arm                                        
ChR2           28.69  12.98   13  23.38  10.65   13
ChR2Scrambled  20.67  12.03   12  24.83  18.01   12
EYFP           21.86  14.37    7  24.43  15.55    7
First probe, Tukey:

          group1         group2  meandiff   p-adj    lower    upper  reject
0           ChR2  ChR2Scrambled   -8.0256  0.2832 -20.8141   4.7628   False
1           ChR2           EYFP   -6.8352  0.5055 -21.8114   8.1411   False
2  ChR2Scrambled           EYFP    1.1905  0.9796 -14.0026  16.3836   False
Second probe, Tukey:

          group1         group2  meandiff   p-adj    lower    upper  reject
0           ChR2  ChR2Scrambled    1.4487  0.9678 -13.2159  16.1133   False
1           ChR2           EYFP    1.0440  0.9877 -16.1294  18.2174   False
2  ChR2Scrambled           EYFP   -0.4048  0.9982 -17.8268  17.0173   False
Variant, all mice with a first probe: {'ChR2': {'mean': 28.69, 'size': 13}, 'ChR2Scrambled': {'mean': 20.67, 'size': 12}, 'EYFP': {'mean': 21.86, 'size': 7}}

          group1         group2  meandiff   p-adj    lower    upper  reject
0           ChR2  ChR2Scrambled   -8.0256  0.2832 -20.8141   4.7628   False
1           ChR2           EYFP   -6.8352  0.5055 -21.8114   8.1411   False
2  ChR2Scrambled           EYFP    1.1905  0.9796 -14.0026  16.3836   False

c5 = el5.dropna(subset=["pheno"])
tab5 = pd.crosstab(c5.subject_arm, c5.pheno)
print(tab5.to_string())
print("Fraction PR:", (c5.pheno == "PR").groupby(c5.subject_arm).mean().round(3).to_dict())
is_chr2 = c5.subject_arm == "ChR2"
ft = stats.fisher_exact([[((c5.pheno == "PR") & is_chr2).sum(), ((c5.pheno != "PR") & is_chr2).sum()],
                         [((c5.pheno == "PR") & ~is_chr2).sum(), ((c5.pheno != "PR") & ~is_chr2).sum()]])
print(f"Fisher exact, PR vs not, ChR2 vs both controls: odds ratio {ft.statistic:.2f}, p = {ft.pvalue:.3g}")
print("Fiber photometry RI60 mice, for reference:", (both.pheno.value_counts(normalize=True)).round(3).to_dict())

pheno          DPR  PR  PS
subject_arm               
ChR2             0  12   1
ChR2Scrambled    3   6   3
EYFP             0   5   2
Fraction PR: {'ChR2': 0.923, 'ChR2Scrambled': 0.5, 'EYFP': 0.714}
Fisher exact, PR vs not, ChR2 vs both controls: odds ratio 8.73, p = 0.0497
Fiber photometry RI60 mice, for reference: {'PS': 0.531, 'DPR': 0.281, 'PR': 0.188}

ch = c5[c5.subject_arm == "ChR2"]
print(pd.crosstab(ch.sex, ch.pheno).to_string())
print("Percent PR by sex:", (100 * (ch.pheno == "PR").groupby(ch.sex).mean()).round(1).to_dict(), "; n:", ch.sex.value_counts().to_dict())
print("ChR2 mice without both probes:", el5[(el5.subject_arm == "ChR2") & el5.pheno.isna()][["sex", "early", "late"]].to_dict("index"))

pheno  PR  PS
sex          
F       5   1
M       7   0
Percent PR by sex: {'F': 83.3, 'M': 100.0} ; n: {'M': 7, 'F': 6}
ChR2 mice without both probes: {'126.286': {'sex': 'F', 'early': nan, 'late': nan}}

def omission_anova(el):
    x = el.dropna(subset=["omission_min"])
    print(x.groupby("subject_arm").omission_min.agg(["mean", "std", "size", "min", "max"]).round(2).to_string())
    an = stats.f_oneway(*[g.omission_min for _, g in x.groupby("subject_arm")])
    kw = stats.kruskal(*[g.omission_min for _, g in x.groupby("subject_arm")])
    print(f"One-way ANOVA: F = {an.statistic:.2f}, p = {an.pvalue:.3g}; Kruskal-Wallis variant p = {kw.pvalue:.3g}")
    return an

an5 = omission_anova(el5)

                mean    std  size    min    max
subject_arm                                    
ChR2           25.79  13.67    13  17.06  60.00
ChR2Scrambled  35.53  15.97    12  17.08  60.00
EYFP           37.29  15.86     7  19.23  53.22
One-way ANOVA: F = 1.87, p = 0.172; Kruskal-Wallis variant p = 0.141

s6, p6, el6 = opto["Opto-DMS-Inhibitory"]
print(el6.groupby("subject_arm").agg(mice=("n_ri60", "size"), did_not_learn=("did_not_learn", lambda x: (x == "True").sum()),
                                     no_ri60=("n_ri60", lambda x: (x == 0).sum()), both_probes=("pheno", "count")).to_string())
print("Mice flagged or without RI60:", el6[(el6.did_not_learn == "True") | (el6.n_ri60 == 0)][["subject_arm", "sex", "n_ri60", "did_not_learn"]].to_dict("index"))
stages6 = s6.groupby(["subject", "stage"]).size().unstack(fill_value=0)
print("Sessions per stage of those mice:", stages6.loc[el6[(el6.did_not_learn == "True") | (el6.n_ri60 == 0)].index].to_dict("index"))
first_fr1 = s6[s6.stage.isin(["FR1", "FR1habit"])].groupby("subject").start_ts.min()
first_ri30 = s6[s6.stage == "RI30"].groupby("subject").start_ts.min()
print("Mice whose first session is RI30 (light from RI30, possibly group 2):", int((first_ri30.reindex(el6.index) < first_fr1.reindex(el6.index)).sum()))
paper_n = {"NpHR": (8, 8), "EYFP": (9, 7), "NpHRScrambled": (9, 6)}  # Figure 6F legend, groups 1 and 2
paper_drop = {"NpHR": "25% of 8 in group 1 (2) and 29% of group 2 (2)", "EYFP": "none", "NpHRScrambled": "22% of 9 in group 1 (2)"}
cnt = el6.groupby("subject_arm").size()
for a, (g1, g2) in paper_n.items():
    print(f"{a}: Figure 6F {g1 + g2} mice ({g1} + {g2}); dandiset {cnt.get(a, 0)}; shortfall {g1 + g2 - cnt.get(a, 0)}; "
          f"paper drops {paper_drop[a]}; flagged or without RI60 in the dandiset {int(((el6.subject_arm == a) & ((el6.did_not_learn == 'True') | (el6.n_ri60 == 0))).sum())}")
print("Session descriptions mentioning a group:", int(s6.description.str.contains("[Gg]roup").sum()))
print("FR1 sessions with optogenetic stimulation, by arm (all mice):", s6[s6.stage == "FR1"].groupby("subject_arm").size().to_dict())

               mice  did_not_learn  no_ri60  both_probes
subject_arm                                             
EYFP             16              1        0           15
NpHR             13              1        1            9
NpHRScrambled    13              0        0           10
Mice flagged or without RI60: {'244.465': {'subject_arm': 'EYFP', 'sex': 'F', 'n_ri60': 8, 'did_not_learn': 'True'}, '270.476': {'subject_arm': 'NpHR', 'sex': 'F', 'n_ri60': 0, 'did_not_learn': 'True'}}
Sessions per stage of those mice: {'244.465': {'FR1': 17, 'FR1habit': 4, 'RI30': 1, 'RI60': 8, 'omission': 0, 'probetest': 0, 'shock': 0}, '270.476': {'FR1': 14, 'FR1habit': 4, 'RI30': 0, 'RI60': 0, 'omission': 0, 'probetest': 0, 'shock': 0}}
Mice whose first session is RI30 (light from RI30, possibly group 2): 0
NpHR: Figure 6F 16 mice (8 + 8); dandiset 13; shortfall 3; paper drops 25% of 8 in group 1 (2) and 29% of group 2 (2); flagged or without RI60 in the dandiset 1
EYFP: Figure 6F 16 mice (9 + 7); dandiset 16; shortfall 0; paper drops none; flagged or without RI60 in the dandiset 1
NpHRScrambled: Figure 6F 15 mice (9 + 6); dandiset 13; shortfall 2; paper drops 22% of 9 in group 1 (2); flagged or without RI60 in the dandiset 0
Session descriptions mentioning a group: 0
FR1 sessions with optogenetic stimulation, by arm (all mice): {'EYFP': 115, 'NpHR': 105, 'NpHRScrambled': 96}

long6 = p6[p6.probe != ""].merge(el6[["subject_arm"]].reset_index(), on="subject")
aov6, w6 = mixed_anova(long6, "subject_arm")
print(w6.groupby("subject_arm")[["early", "late"]].agg(["mean", "std", "size"]).round(2).to_string())
print("Second probe, Tukey:"); print(tukey(w6.late, w6.subject_arm).to_string())
print("Variant, two-way mixed ANOVA:"); print(fmt(aov6))
l6 = el6.dropna(subset=["late"])
print("Variant, all mice with a second probe:", l6.groupby("subject_arm").late.agg(["mean", "size"]).round(2).to_dict("index"))
print(tukey(l6.late, l6.subject_arm).to_string())

probe          early              late            
                mean   std size   mean    std size
subject_arm                                       
EYFP           16.33  6.59   15  23.47  18.03   15
NpHR            9.89  7.98    9   9.33   9.39    9
NpHRScrambled  15.80  6.71   10  18.70  11.92   10
Second probe, Tukey:

  group1         group2  meandiff   p-adj    lower    upper  reject
0   EYFP           NpHR  -14.1333  0.0694 -29.1989   0.9323   False
1   EYFP  NpHRScrambled   -4.7667  0.7032 -19.3539   9.8206   False
2   NpHR  NpHRScrambled    9.3667  0.3511  -7.0507  25.7840   False
Variant, two-way mixed ANOVA:
                        F    df      p
subject_arm          3.16  2,31  0.056
probe                3.70  1,31  0.064
subject_arm x probe  1.27  2,31    0.3
Variant, all mice with a second probe: {'EYFP': {'mean': 23.47, 'size': 15}, 'NpHR': {'mean': 8.7, 'size': 10}, 'NpHRScrambled': {'mean': 18.7, 'size': 10}}

  group1         group2  meandiff   p-adj    lower    upper  reject
0   EYFP           NpHR  -14.7667  0.0431 -29.1414  -0.3919    True
1   EYFP  NpHRScrambled   -4.7667  0.6967 -19.1414   9.6081   False
2   NpHR  NpHRScrambled   10.0000  0.2771  -5.7468  25.7468   False

c6 = el6.dropna(subset=["pheno"])
print(pd.crosstab(c6.subject_arm, c6.pheno).to_string())
fr = pd.crosstab(c6.subject_arm, c6.pheno, normalize="index").round(3)
print(fr.to_string())
def fisher(df, arm_a, arm_b, ph_):
    a, b = df[df.subject_arm == arm_a], df[df.subject_arm == arm_b]
    return stats.fisher_exact([[(a.pheno == ph_).sum(), (a.pheno != ph_).sum()], [(b.pheno == ph_).sum(), (b.pheno != ph_).sum()]]).pvalue
print(f"Fisher PS, NpHR vs EYFP: p = {fisher(c6, 'NpHR', 'EYFP', 'PS'):.3g}; PR, NpHRScrambled vs EYFP: p = {fisher(c6, 'NpHRScrambled', 'EYFP', 'PR'):.3g}")
print("By sex:"); print(pd.crosstab([c6.subject_arm, c6.sex], c6.pheno).to_string())

pheno          DPR  PR  PS
subject_arm               
EYFP             4   6   5
NpHR             0   1   8
NpHRScrambled    1   6   3
pheno            DPR     PR     PS
subject_arm                       
EYFP           0.267  0.400  0.333
NpHR           0.000  0.111  0.889
NpHRScrambled  0.100  0.600  0.300
Fisher PS, NpHR vs EYFP: p = 0.0131; PR, NpHRScrambled vs EYFP: p = 0.428
By sex:
pheno              DPR  PR  PS
subject_arm   sex             
EYFP          F      2   3   3
              M      2   3   2
NpHR          F      0   1   3
              M      0   0   5
NpHRScrambled F      1   2   2
              M      0   4   1

an6 = omission_anova(el6)

                mean    std  size    min    max
subject_arm                                    
EYFP           28.74  12.08    15  16.99  53.52
NpHR           21.47   5.51    12  16.67  32.28
NpHRScrambled  27.82  13.46    10  16.67  60.00
One-way ANOVA: F = 1.66, p = 0.206; Kruskal-Wallis variant p = 0.252

esc7, res7 = escalation("Opto-DLS-Excitatory")
for st, r in res7.items():
    print(st, f"{r['sessions']} sessions from {r['mice']} mice")
    print("  slope:", {k: round(v, 3) for k, v in r["slopes"].items()})
    print("  arm x index interaction against ChR2 (difference in slope, p):", {k: (round(v[0], 3), round(v[1], 4)) for k, v in r["interaction_vs_ref"].items()})
    print("  the third pair:", r["other_pair"][0], "difference in slope", round(r["other_pair"][1], 3), "p =", round(r["other_pair"][2], 4))
for st, r in escalation_mixed(esc7).items():
    print(f"Variant, {st} mixed model with a random intercept and slope per mouse (difference in slope against ChR2, and the third pair, p):", r)
for st, r in escalation_mixed(esc7, re_formula=None).items():
    print(f"Variant, {st} mixed model with a random intercept per mouse only:", r)
for st in ["FR1", "RI30"]:
    x = esc7[esc7.stage == st]
    f1 = smf.ols("rate ~ idx * C(arm)", x).fit(); f0 = smf.ols("rate ~ idx + C(arm)", x).fit()
    print(f"{st}: F test of all arm x index terms p = {f1.compare_f_test(f0)[1]:.3g}")

FR1 197 sessions from 35 mice
  slope: {'ChR2': np.float64(0.223), 'ChR2Scrambled': np.float64(0.087), 'EYFP': np.float64(0.224)}
  arm x index interaction against ChR2 (difference in slope, p): {'ChR2Scrambled': (np.float64(-0.136), np.float64(0.1173)), 'EYFP': (np.float64(0.002), np.float64(0.9841))}
  the third pair: EYFP vs ChR2Scrambled difference in slope 0.137 p = 0.0129
RI30 49 sessions from 35 mice
  slope: {'ChR2': np.float64(0.771), 'ChR2Scrambled': np.float64(-1.198), 'EYFP': np.float64(0.925)}
  arm x index interaction against ChR2 (difference in slope, p): {'ChR2Scrambled': (np.float64(-1.969), np.float64(0.1194)), 'EYFP': (np.float64(0.153), np.float64(0.9209))}
  the third pair: EYFP vs ChR2Scrambled difference in slope 2.123 p = 0.0907

Variant, FR1 mixed model with a random intercept and slope per mouse (difference in slope against ChR2, and the third pair, p): {'ChR2Scrambled': (np.float64(-0.07), np.float64(0.9089)), 'EYFP': (np.float64(-0.173), np.float64(0.5861)), 'converged': True, 'random effects covariance': [[0.0288, 0.0328], [0.0328, 0.0374]], 'EYFP vs ChR2Scrambled': (np.float64(-0.103), np.float64(0.7553)), 'warnings': {'Random effects covariance is singular': 8}}
Variant, RI30 mixed model with a random intercept and slope per mouse (difference in slope against ChR2, and the third pair, p): {'ChR2Scrambled': (np.float64(1.295), np.float64(0.4281)), 'EYFP': (np.float64(0.057), np.float64(0.9608)), 'converged': True, 'random effects covariance': [[6.7192, -3.5321], [-3.5321, 3.3283]], 'EYFP vs ChR2Scrambled': (np.float64(-1.237), np.float64(0.4525)), 'warnings': {}}
Variant, FR1 mixed model with a random intercept per mouse only: {'ChR2Scrambled': (np.float64(-0.146), np.float64(0.057)), 'EYFP': (np.float64(-0.084), np.float64(0.2624)), 'converged': True, 'random effects covariance': [[0.3814]], 'EYFP vs ChR2Scrambled': (np.float64(0.062), np.float64(0.229)), 'warnings': {}}
Variant, RI30 mixed model with a random intercept per mouse only: {'ChR2Scrambled': (np.float64(-1.355), np.float64(0.1647)), 'EYFP': (np.float64(0.105), np.float64(0.9271)), 'converged': True, 'random effects covariance': [[1.6761]], 'EYFP vs ChR2Scrambled': (np.float64(1.46), np.float64(0.1343)), 'warnings': {}}
FR1: F test of all arm x index terms p = 0.0337
RI30: F test of all arm x index terms p = 0.117

_, p7, el7 = opto["Opto-DLS-Excitatory"]
long7 = p7[p7.probe != ""].merge(el7[["subject_arm"]].reset_index(), on="subject")
aov7, w7 = mixed_anova(long7, "subject_arm")
print(fmt(aov7))
print(w7.groupby("subject_arm")[["early", "late"]].agg(["mean", "std", "size"]).round(2).to_string())

                         F    df      p
subject_arm           3.23  2,28  0.055
probe                13.50  1,28  0.001
subject_arm x probe   2.96  2,28  0.068
probe          early               late            
                mean    std size   mean    std size
subject_arm                                        
ChR2           18.36  11.10   11  24.36  14.38   11
ChR2Scrambled  12.82   7.25   11  15.18  10.58   11
EYFP           19.11   9.24    9  32.89  17.14    9

c7 = el7.dropna(subset=["pheno"])
tab7 = pd.crosstab(c7.subject_arm, c7.pheno)
print(tab7.to_string())
chi = stats.chi2_contingency(tab7)
print(f"Chi-square: {chi.statistic:.2f}, df {chi.dof}, p = {chi.pvalue:.3g}; smallest expected count {chi.expected_freq.min():.1f}")

pheno          DPR  PR  PS
subject_arm               
ChR2             0   6   5
ChR2Scrambled    2   3   6
EYFP             3   5   1
Chi-square: 7.28, df 4, p = 0.122; smallest expected count 1.5

an7 = omission_anova(el7)

                mean    std  size    min    max
subject_arm                                    
ChR2           28.44  15.94    10  17.06  56.32
ChR2Scrambled  33.51  16.62    11  17.04  60.00
EYFP           36.08  14.16    11  18.55  60.00
One-way ANOVA: F = 0.64, p = 0.532; Kruskal-Wallis variant p = 0.303

zen = requests.get("https://zenodo.org/api/records/11266654", timeout=60).json()
print("Zenodo record:", zen["doi"], "version", zen.get("metadata", {}).get("version"), "; files:", [f["key"] for f in zen["files"]])
xl = requests.get(zen["files"][0]["links"]["self"], timeout=60).content
print("Densitometry.xlsx md5:", hashlib.md5(xl).hexdigest(), "(Zenodo checksum", zen["files"][0]["checksum"] + ")")
sheet = pd.read_excel(io.BytesIO(xl), header=None)
blocks = {}
for i, label in sheet[0].items():
    if isinstance(label, str) and label.startswith(("Female", "Male")):
        rows = sheet.iloc[i + 1:i + 8]
        blocks[label.strip()] = pd.DataFrame({"sample": rows[0].str.strip().values, "corrected": rows[7].astype(float).values,
                                              "sheet_ratio": rows[8].values})
print("Blocks in the table:", list(blocks))
print("Header row and the first block as read (rows i to i+7, columns 0 to 8):")
first = [i for i, l in sheet[0].items() if isinstance(l, str) and l.startswith(("Female", "Male"))]
print(sheet.iloc[:first[0]].to_string())
for i in first:
    print(sheet.iloc[i:i + 8, :9].to_string())
dat_rows = []
for region in ["DLS", "DMS"]:
    act, dat = blocks[f"Female {region} Actin"], blocks[f"Female {region} DAT"]
    ratio = dat.corrected.values / act.corrected.values
    geno = np.where(dat["sample"].str.startswith("WT"), "WT", "het")
    pct = 100 * ratio / ratio[geno == "WT"].mean()
    print(f"{region}: computed DAT/actin equals the sheet's ratio column: {np.allclose(ratio, dat.sheet_ratio.astype(float))}")
    for s_, g, p in zip(dat["sample"], geno, pct):
        dat_rows.append({"region": region, "sample": s_, "genotype": g, "pct_of_wt": p})
dd = pd.DataFrame(dat_rows)
print(dd.groupby(["region", "genotype"]).pct_of_wt.agg(["mean", "std", "size"]).round(1).to_string())
for region in ["DMS", "DLS"]:
    x = dd[dd.region == region]
    t3 = stats.ttest_ind(x[x.genotype == "het"].pct_of_wt, x[x.genotype == "WT"].pct_of_wt)
    mw = stats.mannwhitneyu(x[x.genotype == "het"].pct_of_wt, x[x.genotype == "WT"].pct_of_wt)
    print(f"{region}: het vs WT unpaired t = {t3.statistic:.2f}, p = {t3.pvalue:.3g}; Mann-Whitney p = {mw.pvalue:.3g}")
print("WT sample labels (one reads 'WT 4' where the others read 'WT F'):", sorted(dd[dd.genotype == "WT"]["sample"].unique()))
print("Male samples in the table:", int(dd["sample"].str.contains(" M ").sum()))

Zenodo record: 10.5281/zenodo.11266654 version None ; files: ['Densitometry.xlsx']

Densitometry.xlsx md5: 798a810141f05133b1105fd53ae532c7 (Zenodo checksum md5:798a810141f05133b1105fd53ae532c7)

Blocks in the table: ['Female DLS Actin', 'Female DLS DAT', 'Female DMS Actin', 'Female DMS DAT']
Header row and the first block as read (rows i to i+7, columns 0 to 8):
Empty DataFrame
Columns: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
Index: []
                  0     1          2    3      4   5             6                       7    8
0  Female DLS Actin  area       mean  min    max NaN  white-sample  corrected sample-blank  NaN
1        WT F 410.1  0.24  40536.569    0  52771 NaN     24998.431               10502.962  NaN
2       WT F 410.2   0.24   40759.65    0  52155 NaN      24775.35               10279.881  NaN
3       WT 4 410.3   0.24  41490.057    0  52151 NaN     24044.943                9549.474  NaN
4      DAT F 484.1   0.24  39085.262    0  51379 NaN     26449.738               11954.269  NaN
5      DAT F 484.2   0.24  37679.735    0  51403 NaN     27855.265               13359.796  NaN
6      DAT F 484.3   0.24   37732.15    0  51031 NaN      27802.85               13307.381  NaN
7      DAT F 484.4   0.24  38180.667    0  50289 NaN     27354.333               12858.864  NaN
                 0      1          2      3      4   5          6         7         8
9   Female DLS DAT   area       mean    min    max NaN        NaN       NaN   DAT/ACT
10      WT F 410.1  0.403  40176.011      0  47299 NaN  25358.989   3927.81  0.373972
11     WT F 410.2   0.403  41208.121   9886  45405 NaN  24326.879    2895.7  0.281686
12     WT 4 410.3   0.403  41631.266  17656  45151 NaN  23903.734  2472.555  0.258921
13    DAT F 484.1   0.403  41466.513   7454  45795 NaN  24068.487  2637.308  0.220616
14    DAT F 484.2   0.403  40425.958   2604  45275 NaN  25109.042  3677.863  0.275293
15    DAT F 484.3   0.403   40257.63   6972  45883 NaN   25277.37  3846.191  0.289027
16    DAT F 484.4   0.403   42477.53  23400  45917 NaN   23057.47  1626.291  0.126472
                   0      1          2    3      4   5          6          7    8
19  Female DMS Actin   area       mean  min    max NaN        NaN        NaN  NaN
20        WT F 410.1  0.248  38886.222    0  59133 NaN  26648.778  19786.179  NaN
21       WT F 410.2   0.248  44591.149    0  59107 NaN  20943.851  14081.252  NaN
22       WT 4 410.3   0.248  43165.932    0  58975 NaN  22369.068  15506.469  NaN
23      DAT F 484.1   0.248  40640.485    0  58775 NaN  24894.515  18031.916  NaN
24      DAT F 484.2   0.248  40396.025    0  58725 NaN  25138.975  18276.376  NaN
25      DAT F 484.3   0.248  39556.131    0  58661 NaN  25978.869   19116.27  NaN
26      DAT F 484.4   0.248   44629.67    0  58761 NaN   20905.33  14042.731  NaN
                 0      1          2      3      4   5          6         7         8
28  Female DMS DAT    NaN        NaN    NaN    NaN NaN        NaN       NaN       NaN
29      WT F 410.1  0.315  38797.736   7706  45887 NaN  26737.264  5395.135  0.272672
30     WT F 410.2   0.315  42319.147  19720  46131 NaN  23215.853  1873.724  0.133065
31     WT 4 410.3   0.315  37640.921   2748  45679 NaN  27894.079   6551.95   0.42253
32    DAT F 484.1   0.315  37336.337    728  45855 NaN  28198.663  6856.534  0.380244
33    DAT F 484.2   0.315  36479.241    374  45631 NaN  29055.759   7713.63  0.422055
34    DAT F 484.3   0.315  35672.696      0  45557 NaN  29862.304  8520.175  0.445703
35    DAT F 484.4   0.315  41356.669   9018  46675 NaN  24178.331  2836.202  0.201969
DLS: computed DAT/actin equals the sheet's ratio column: True
DMS: computed DAT/actin equals the sheet's ratio column: True
                  mean   std  size
region genotype                   
DLS    WT        100.0  20.0     3
       het        74.7  24.2     4
DMS    WT        100.0  52.4     3
       het       131.3  40.0     4
DMS: het vs WT unpaired t = 0.90, p = 0.408; Mann-Whitney p = 0.629
DLS: het vs WT unpaired t = -1.46, p = 0.203; Mann-Whitney p = 0.4
WT sample labels (one reads 'WT 4' where the others read 'WT F'): ['WT 4 410.3', 'WT F 410.1', 'WT F 410.2']
Male samples in the table: 0

print("Session stages over all files:", sessions.stage.value_counts().to_dict())
print("Acquisition types in the western blot files: images only;",
      "files whose description mentions fear, freezing, open field or histology:",
      int(sessions.description.str.contains("fear|freez|open field|histolog", case=False).sum()))
print(f"Total runtime {time.time() - started:.0f} s")

Session stages over all files: {'RI60': 2163, 'FR1': 937, 'shock': 324, 'RI30': 265, 'omission': 134, 'probetest': 104, 'FR1habit': 93, 'RR20': 84, 'unknown': 18, 'other': 16, 'RR5': 1}
Acquisition types in the western blot files: images only; files whose description mentions fear, freezing, open field or histology: 0
Total runtime 1424 s

Claim	Tested
fig1d-ri60-vs-rr20-shocks	tested
fig1e-omission-ri60-longer	tested
fig1f-omission-shock-no-correlation	tested
fig2a-phenotype-rule	tested
fig2b-only-dpr-increase	tested
fig2c-nosepoke-rate	tested
fig2d-omission-pr-vs-ps	tested
fig2k-ps-wait-longer	tested
fig2l-extra-nosepokes	tested
fig2m-inter-reward-interval	tested
fig2n-rewards-per-session	tested
fig2o-escalation-timing	tested
figs1i-ri60-more-rewards	tested
figs2b-males-more-shocks	tested
fig3e-ri30-dms-peaks	blocked: no RI30 photometry in the dandiset
fig3f-dms-score-phenotype	tested
fig3g-dms-score-shock-correlation	tested
fig3h-dms-score-omission-no-correlation	tested
fig4f-dls-score-no-phenotype	tested
fig4g-dls-score-shock-no-correlation	tested
fig4h-dls-score-omission-no-correlation	tested
fig5f-dms-chr2-escalation	tested
fig5g-dms-chr2-more-shocks-early	tested
fig5h-dms-chr2-mostly-pr	tested
figs5f-chr2-pr-by-sex	tested
fig5i-dms-stim-no-omission-effect	tested
fig6-nphr-learning-dropout	blocked: the dropped mice are not in the dandiset (counts in its section)
fig6g-nphr-fewer-shocks-late	tested
fig6h-nphr-more-ps	tested
fig6i-nphr-no-omission-effect	tested
fig7f-dls-escalation-similar	tested
fig7g-dls-stim-no-shock-effect	tested
fig7h-dls-similar-phenotypes	tested
fig7i-dls-no-omission-effect	tested
figs3ab-dat-protein-not-reduced	tested on Zenodo data
figs1a-shock-aversive	blocked: fear conditioning not in the dandiset
figs3cd-open-field	blocked: open field not in the dandiset
fig3c-gcamp-in-dopamine-neurons	blocked: histology not in the dandiset

Claim	Data	Outcome
fig1d-ri60-vs-rr20-shocks	dandiset	not reproduced
fig1e-omission-ri60-longer	dandiset	partially reproduced
fig1f-omission-shock-no-correlation	dandiset	reproduced
fig2a-phenotype-rule	dandiset	partially reproduced
fig2b-only-dpr-increase	dandiset	reproduced
fig2c-nosepoke-rate	dandiset	reproduced
fig2d-omission-pr-vs-ps	dandiset	reproduced
fig2k-ps-wait-longer	dandiset	partially reproduced
fig2l-extra-nosepokes	dandiset	partially reproduced
fig2m-inter-reward-interval	dandiset	reproduced
fig2n-rewards-per-session	dandiset	partially reproduced
fig2o-escalation-timing	dandiset	partially reproduced
figs1i-ri60-more-rewards	dandiset	reproduced
figs2b-males-more-shocks	dandiset	reproduced
fig3e-ri30-dms-peaks	dandiset	blocked: data missing
fig3f-dms-score-phenotype	dandiset	not reproduced
fig3g-dms-score-shock-correlation	dandiset	partially reproduced
fig3h-dms-score-omission-no-correlation	dandiset	reproduced
fig4f-dls-score-no-phenotype	dandiset	reproduced
fig4g-dls-score-shock-no-correlation	dandiset	not reproduced
fig4h-dls-score-omission-no-correlation	dandiset	reproduced
fig5f-dms-chr2-escalation	dandiset	not reproduced
fig5g-dms-chr2-more-shocks-early	dandiset	not reproduced
fig5h-dms-chr2-mostly-pr	dandiset	reproduced
figs5f-chr2-pr-by-sex	dandiset	partially reproduced
fig5i-dms-stim-no-omission-effect	dandiset	reproduced
fig6-nphr-learning-dropout	dandiset	blocked: data missing
fig6g-nphr-fewer-shocks-late	dandiset	partially reproduced
fig6h-nphr-more-ps	dandiset	reproduced
fig6i-nphr-no-omission-effect	dandiset	reproduced
fig7f-dls-escalation-similar	dandiset	not reproduced
fig7g-dls-stim-no-shock-effect	dandiset	reproduced
fig7h-dls-similar-phenotypes	dandiset	reproduced
fig7i-dls-no-omission-effect	dandiset	reproduced
figs3ab-dat-protein-not-reduced	elsewhere (Zenodo)	reproduced
figs1a-shock-aversive	dandiset	blocked: data missing
figs3cd-open-field	dandiset	blocked: data missing
fig3c-gcamp-in-dopamine-neurons	dandiset	blocked: data missing

Reproducing Seiler et al. (2022) from dandiset 000971¶

Data check¶

Fiber photometry cohort: behavior¶

fig1d-ri60-vs-rr20-shocks¶

fig1e-omission-ri60-longer¶

fig1f-omission-shock-no-correlation¶

fig2a-phenotype-rule¶

fig2b-only-dpr-increase¶

fig2c-nosepoke-rate¶

fig2d-omission-pr-vs-ps¶

fig2k-ps-wait-longer¶

fig2l-extra-nosepokes¶

fig2m-inter-reward-interval¶

fig2n-rewards-per-session¶

fig2o-escalation-timing¶

figs1i-ri60-more-rewards¶

figs2b-males-more-shocks¶

Fiber photometry cohort: dopamine axon signals¶

fig3e-ri30-dms-peaks¶

fig3f-dms-score-phenotype¶

fig3e-ri30-dms-peaks, the RI60 part¶

fig3g-dms-score-shock-correlation¶

fig3h-dms-score-omission-no-correlation¶

fig4f-dls-score-no-phenotype¶

fig4g-dls-score-shock-no-correlation¶

fig4h-dls-score-omission-no-correlation¶

Optogenetic cohorts¶

fig5f-dms-chr2-escalation¶

fig5g-dms-chr2-more-shocks-early¶

fig5h-dms-chr2-mostly-pr¶

figs5f-chr2-pr-by-sex¶

fig5i-dms-stim-no-omission-effect¶

fig6-nphr-learning-dropout¶

fig6g-nphr-fewer-shocks-late¶

fig6h-nphr-more-ps¶

fig6i-nphr-no-omission-effect¶

fig7f-dls-escalation-similar¶

fig7g-dls-stim-no-shock-effect¶

fig7h-dls-similar-phenotypes¶

fig7i-dls-no-omission-effect¶

Data outside the dandiset¶

figs3ab-dat-protein-not-reduced¶

Claims that cannot be tested¶

Summary¶