import os
import time
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy import stats
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score

import rsg

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

started = time.time()
data = rsg.load()
S, T = data["train"], data["test"]
trials, spikes = S["trials"], S["spikes"]
good = rsg.completed(trials)
units, unit_table = rsg.select_units(S, good)
DS = [(e, th, rsg.dataset_name(e, th)) for e, th in rsg.DATASETS]
print(f"loaded dandiset {rsg.DANDISET_ID} version {rsg.DANDISET_VERSION} in {time.time() - started:.0f} s")

loaded dandiset 000130 version 0.241017.1448 in 5 s

for key, f in [("train", S), ("test", T)]:
    print(f"{key}: {data['urls'][rsg.TRAIN if key == 'train' else rsg.TEST]}")
    for k, v in f["meta"].items():
        print(f"   {k}: {v}")
    print(f"   trials: {len(f['trials'])}; columns: {list(f['trials'].columns)}")
    print(f"   units: {len(f['unit_ids'])}; held out: {int(f['heldout'].sum())}")
print("\nTest-file units that are in the train file:", int(np.isin(T["unit_ids"], S["unit_ids"]).sum()), "of", len(T["unit_ids"]),
      "; train units held out:", int(S["heldout"].sum()), "; held-out train units in the test file:",
      int(np.isin(S["unit_ids"][S["heldout"]], T["unit_ids"]).sum()))
print("\nTrain trials by NLB split:", trials.split.value_counts().to_dict())
print("Train trials with Set and Go:", int((trials.set_time.notna() & trials.go_time.notna()).sum()),
      "; of these flagged is_outlier:", int(trials[trials.set_time.notna() & trials.go_time.notna()].is_outlier.sum()),
      "; completed and not outliers (used here):", len(good))
print("Used trials by split:", good.split.value_counts().to_dict())
print("\nUsed trials per condition (rows: effector, direction, prior; columns: ts in ms):")
print(pd.crosstab([good.is_eye.map({True: 'Eye', False: 'Hand'}), good.theta.map({0.0: 'Right', 180.0: 'Left'}),
                   good.is_short.map({True: 'Short', False: 'Long'})], good.ts).to_string())
print("\nColumn descriptions (train trial table):")
for c, d in S["descriptions"].items():
    print(f"   {c}: {d}")

train: https://dandiarchive.s3.amazonaws.com/blobs/73a/6f6/73a6f648-8e97-4c90-b323-a15844e4e3d4
   session_description: Data from monkey Haydn performing ready-set-go time interval reproduction task. This file contains continuous segments of the full session on 2016-12-11 that can be used for training models for the Neural Latents Benchmark.
   session_start_time: 2016-12-11T00:00:00-05:00
   session_id: 20161211
   subject_id: Haydn
   species: Macaca mulatta
   age: P4Y
   related_publications: ['http://dx.doi.org/10.1016/j.neuron.2019.06.012']
   electrode_location: ['DMFC']
   electrode_x_all_nan: True
   n_electrodes: 72
   trials: 1160; columns: ['start_time', 'stop_time', 'fix_on_time', 'fix_time', 'target_on_time', 'ready_time', 'set_time', 'go_time', 'target_acq_time', 'reward_time', 'bad_time', 'is_short', 'is_eye', 'theta', 'ts', 'tp', 'fix_time_dur', 'target_time_dur', 'iti', 'reward_dur', 'is_outlier', 'split']
   units: 54; held out: 14
test: https://dandiarchive.s3.amazonaws.com/blobs/eaf/1fa/eaf1fa97-ce58-477d-b79d-82425cea818f
   session_description: Data from monkey Haydn performing ready-set-go time interval reproduction task. This file contains segments of trials from the full session on 2016-12-11 that are to be used for evaluating models for the Neural Latents Benchmark.
   session_start_time: 2016-12-11T00:00:00-05:00
   session_id: 20161211
   subject_id: Haydn
   species: Macaca mulatta
   age: P4Y
   related_publications: ['http://dx.doi.org/10.1016/j.neuron.2019.06.012']
   electrode_location: ['DMFC']
   electrode_x_all_nan: True
   n_electrodes: 72
   trials: 283; columns: ['start_time', 'stop_time', 'go_time', 'split']
   units: 40; held out: 0

Test-file units that are in the train file: 40 of 40 ; train units held out: 14 ; held-out train units in the test file: 0

Train trials by NLB split: {'train': 748, 'val': 258, 'none': 154}
Train trials with Set and Go: 1029 ; of these flagged is_outlier: 46 ; completed and not outliers (used here): 983
Used trials by split: {'train': 720, 'val': 246, 'none': 17}

Used trials per condition (rows: effector, direction, prior; columns: ts in ms):
ts                     480.0   560.0   640.0   720.0   800.0   900.0   1000.0  1100.0  1200.0
is_eye theta is_short                                                                        
Eye    Left  Long           0       0       0       0      36      29      30      29      31
             Short         21      28      17      17      22       0       0       0       0
       Right Long           0       0       0       0      38      34      23      25      27
             Short         24      27      31      22      25       0       0       0       0
Hand   Left  Long           0       0       0       0      14      24      16      16      18
             Short         33      21      25      31      24       0       0       0       0
       Right Long           0       0       0       0      19      17      30      21      25
             Short         21      35      14      21      22       0       0       0       0

Column descriptions (train trial table):
   start_time: Start time of epoch, in seconds
   stop_time: Stop time of epoch, in seconds
   fix_on_time: Time when the fixation cue on the display turns on
   fix_time: Time when the animal fixates on the fixation cue
   target_on_time: Time of target presentation
   ready_time: Time of ready cue delivery
   set_time: Time of set cue presentation
   go_time: Time of monkey response, either with an eye saccade or joystick movement
   target_acq_time: Time of target acquisition
   reward_time: Time of monkey reward, if there was one
   bad_time: Time of trial failure, if the trial was unsuccessful
   is_short: Whether short prior (True) or long prior (False)
   is_eye: Whether motion to target was with hand (False) or eye (True)
   theta: Whether the target was on the right (0) or left (180)
   ts: Time between ready and set stimuli, in ms
   tp: Time between set stimulus and monkey go response, in ms
   fix_time_dur: Time between animal fixation and target onset, in ms
   target_time_dur: Time between target onset and ready, in ms
   iti: Inter-trial interval, in ms. Equal to 500 for trials sampled from short prior and 1220 for trials sampled from long prior
   reward_dur: Duration of reward, in ms
   is_outlier: Whether a trial's tp is considered an outlier using probabilistic mixture model
   split: Trial split that the trial belongs to for the Neural Latents Benchmark. Can be "train", "val", "test", or "none"

c = trials[trials.set_time.notna() & trials.go_time.notna()].copy()
d_ts = c.ts - (c.set_time - c.ready_time) * 1000
d_tp = c.tp - (c.go_time - c.set_time) * 1000
print("ts minus (set_time - ready_time), ms: median %.1f, range %.1f to %.1f; within 5 ms: %.0f%%"
      % (d_ts.median(), d_ts.min(), d_ts.max(), 100 * (abs(d_ts) <= 5).mean()))
for flag, g in d_tp.groupby(c.is_outlier):
    print(f"tp minus (go_time - set_time), ms, is_outlier={flag}: median {g.median():.1f}, 5th to 95th percentile "
          f"{g.quantile(.05):.1f} to {g.quantile(.95):.1f}, range {g.min():.1f} to {g.max():.1f}")
print("Distinct ts by prior:", {('Short' if k else 'Long'): sorted(v) for k, v in trials.groupby('is_short').ts.unique().items()})
print("iti by prior (ms):", {('Short' if k else 'Long'): sorted(v) for k, v in trials.groupby('is_short').iti.unique().items()})
print("Paper: inter-trial intervals of 1220 ms (Short) and 500 ms (Long).")
print("theta values:", sorted(trials.theta.unique()), "; is_eye values:", sorted(trials.is_eye.unique()))
print("tp of used trials, ms: %.0f to %.0f; of uncompleted trials (no Go):" % (good.tp.min(), good.tp.max()),
      sorted(trials[trials.go_time.isna()].tp.round().unique())[:6], "...")

# Refit of the Methods' mixture model, per prior, ts, effector and direction, on completed trials.
flags = []
for key, g in c[c.tp > 0].groupby(["is_short", "ts", "is_eye", "theta"]):
    r = rsg.outlier_mixture(g.tp.to_numpy(), key[1])
    flags.append(pd.Series(r > 0.5, index=g.index))
refit = pd.concat(flags).reindex(c.index, fill_value=True)
print("\nMixture-model refit against is_outlier (completed trials):")
print(pd.crosstab(c.is_outlier.rename("is_outlier"), refit.rename("refit lapse")).to_string())
print(f"Outlier fraction of completed trials: file {c.is_outlier.mean():.2%}, refit {refit.mean():.2%}; paper, monkey H: 3.84%")
only_file = c[c.is_outlier & ~refit]
print("tp / ts of trials flagged by the file but kept by the refit: %.2f to %.2f (%d trials); of trials flagged by both: %s"
      % ((only_file.tp / only_file.ts).min(), (only_file.tp / only_file.ts).max(), len(only_file),
         np.round(np.sort((c[c.is_outlier & refit].tp / c[c.is_outlier & refit].ts).to_numpy()), 2).tolist()))

ts minus (set_time - ready_time), ms: median 0.0, range -23.5 to 23.5; within 5 ms: 69%
tp minus (go_time - set_time), ms, is_outlier=False: median 17.9, 5th to 95th percentile 1.1 to 29.8, range 0.1 to 200.6
tp minus (go_time - set_time), ms, is_outlier=True: median 17.5, 5th to 95th percentile -1151.9 to 38.2, range -5322.0 to 162.4
Distinct ts by prior: {'Long': [np.float64(800.0), np.float64(900.0), np.float64(1000.0), np.float64(1100.0), np.float64(1200.0)], 'Short': [np.float64(480.0), np.float64(560.0), np.float64(640.0), np.float64(720.0), np.float64(800.0)]}
iti by prior (ms): {'Long': [np.float64(500.0)], 'Short': [np.float64(1220.0)]}
Paper: inter-trial intervals of 1220 ms (Short) and 500 ms (Long).
theta values: [np.float64(0.0), np.float64(180.0)] ; is_eye values: [np.False_, np.True_]
tp of used trials, ms: 407 to 1463; of uncompleted trials (no Go): [np.float64(-2400.0), np.float64(-2200.0), np.float64(-2000.0), np.float64(-1800.0), np.float64(-1600.0), np.float64(-1440.0)] ...

Mixture-model refit against is_outlier (completed trials):

refit lapse  False  True 
is_outlier               
False          983      0
True            36     10
Outlier fraction of completed trials: file 4.47%, refit 0.97%; paper, monkey H: 3.84%
tp / ts of trials flagged by the file but kept by the refit: 0.53 to 1.60 (36 trials); of trials flagged by both: [-2.0, -2.0, -2.0, 0.14, 0.21, 0.37, 0.59, 0.62, 2.21, 2.49]

print("Units: rate over the observed intervals, spikes/s")
print(unit_table.rate.describe().round(2).to_string())
print(f"Kept by the Methods' rule (>= {rsg.MIN_RATE} spike/s and >= {rsg.MIN_TRIALS} trials per condition): "
      f"{len(units)} of {len(unit_table)} (held-out units among the kept: {int(unit_table.heldout[unit_table.kept].sum())})")
print("Paper: GPFA on 'simultaneously recorded neurons in a single session (N = 48 in H)'.")
obs = S["obs_intervals"]
same = all(np.array_equal(o, obs[0]) for o in obs)
print("\nAll units share the same observation intervals:", same, "; intervals:", np.round(obs[0], 2).tolist())
gaps = obs[0][1:, 0] - obs[0][:-1, 1]
inside = [((obs[0][:, 0] <= a) & (obs[0][:, 1] >= b)).any() for a, b in zip(trials.start_time, trials.stop_time)]
print(f"Gaps between intervals: {np.round(gaps, 3).tolist()} s; trials inside one interval: {sum(inside)} of {len(trials)}")
print("Spike times sorted within every unit:", all(np.all(np.diff(s) >= 0) for s in spikes),
      "; spikes outside the observation intervals:",
      int(sum((~((s[:, None] >= obs[0][:, 0]) & (s[:, None] <= obs[0][:, 1])).any(1)).sum() for s in spikes)))

# Recordings deposited twice: does any test trial reappear in the train file? Compare the
# inter-spike-interval sequence of each test trial with the whole train recording, for 3 units.
hits, checked = 0, 0
for uid in T["unit_ids"][:3]:
    s_tr = np.round(np.diff(spikes[int(np.flatnonzero(S["unit_ids"] == uid)[0])]), 4)
    s_te = T["spikes"][int(np.flatnonzero(T["unit_ids"] == uid)[0])]
    for a, b in zip(T["trials"].start_time, T["trials"].stop_time):
        pat = np.round(np.diff(s_te[(s_te >= a) & (s_te < b)]), 4)
        if len(pat) < 3:
            continue
        checked += 1
        hits += any(np.array_equal(s_tr[i:i + len(pat)], pat) for i in np.flatnonzero(s_tr == pat[0]))
print(f"\nTest trials whose spike pattern recurs in the train file: {hits} of {checked} checked (3 units)")
seg = rsg.counts(spikes, units, good.ready_time.to_numpy(), 0, 400, 400)[:, 0, :]
dup = pd.DataFrame(seg).duplicated().sum()
print("Train trials with identical spike counts in all kept units over Ready to Ready + 400 ms:", int(dup))

Units: rate over the observed intervals, spikes/s
count    54.00
mean      6.57
std       5.66
min       0.39
25%       2.47
50%       4.18
75%      11.13
max      21.69
Kept by the Methods' rule (>= 1.0 spike/s and >= 5 trials per condition): 48 of 54 (held-out units among the kept: 13)
Paper: GPFA on 'simultaneously recorded neurons in a single session (N = 48 in H)'.

All units share the same observation intervals: True ; intervals: [[0.0, 740.01], [740.1, 1519.46], [1519.6, 2277.3], [2277.4, 3054.63], [3054.7, 3792.03], [3792.1, 4809.33]]
Gaps between intervals: [0.09, 0.138, 0.105, 0.072, 0.072] s; trials inside one interval: 1160 of 1160
Spike times sorted within every unit: True ; spikes outside the observation intervals: 0

Test trials whose spike pattern recurs in the train file: 0 of 674 checked (3 units)
Train trials with identical spike counts in all kept units over Ready to Ready + 400 ms: 0

# Do spikes share the clock of the trial table? Decode "before" against "after" an event from the
# population counts in 200 ms windows, with the event time shifted by an offset. A response
# locked to the event gives the best decoding at zero offset.
rng_clock = np.random.default_rng(1)
offsets = [-400, -200, -100, -50, 0, 50, 100, 200, 400]
rows = []
for event in ["ready_time", "set_time", "go_time"]:
    for off in offsets:
        t0 = good[event].to_numpy() + off / 1000
        pre = rsg.counts(spikes, units, t0, -200, 0, 200)[:, 0, :]
        post = rsg.counts(spikes, units, t0, 0, 200, 200)[:, 0, :]
        X = np.sqrt(np.vstack([pre, post]))
        y = np.r_[np.zeros(len(pre)), np.ones(len(post))]
        perm = rng_clock.permutation(len(y))
        acc = cross_val_score(LogisticRegression(max_iter=2000), X[perm], y[perm], cv=5).mean()
        rows.append({"event": event, "offset_ms": off, "accuracy": acc})
clock = pd.DataFrame(rows).pivot(index="offset_ms", columns="event", values="accuracy")
print("Cross-validated accuracy, before against after the event (chance 0.5):")
print(clock.round(3).to_string())
print("Offset of best decoding:", clock.idxmax().to_dict())

# Contrast between conditions: decode the prior from the counts in the 400 ms before Ready, with
# the trial labels as they are and shifted by one block of trials.
pre_ready = np.sqrt(rsg.counts(spikes, units, good.ready_time.to_numpy(), -400, 0, 400)[:, 0, :])
for shift in [0, 25, 50]:
    y = np.roll(good.is_short.to_numpy(), shift)
    perm = rng_clock.permutation(len(y))
    acc = cross_val_score(LogisticRegression(max_iter=2000), pre_ready[perm], y[perm], cv=5).mean()
    print(f"Prior decoded before Ready, labels shifted by {shift} trials: accuracy {acc:.3f}")

Cross-validated accuracy, before against after the event (chance 0.5):
event      go_time  ready_time  set_time
offset_ms                               
-400         0.728       0.555     0.679
-200         0.865       0.555     0.673
-100         0.900       0.588     0.690
-50          0.920       0.626     0.708
 0           0.941       0.766     0.791
 50          0.957       0.890     0.884
 100         0.974       0.932     0.929
 200         0.949       0.928     0.917
 400         0.769       0.913     0.884
Offset of best decoding: {'go_time': 100, 'ready_time': 100, 'set_time': 100}
Prior decoded before Ready, labels shifted by 0 trials: accuracy 0.641
Prior decoded before Ready, labels shifted by 25 trials: accuracy 0.558
Prior decoded before Ready, labels shifted by 50 trials: accuracy 0.520

rows = []
for e, th, name in DS:
    for short in [True, False]:
        g = good[(good.is_eye == e) & (good.theta == th) & (good.is_short == short)]
        b, se, ci = rsg.weighted_slope(g.ts, g.tp)
        b_ols = np.polyfit(g.ts, g.tp, 1)[0]
        rows.append({"dataset": name, "prior": "Short" if short else "Long", "trials": len(g), "slope": b,
                     "ci_low": ci[0], "ci_high": ci[1], "slope_unweighted": b_ols})
slopes = pd.DataFrame(rows)
print(slopes.round(3).to_string(index=False))
inside = ((slopes.slope > 0) & (slopes.slope < 1)).sum()
strict = ((slopes.ci_low > 0) & (slopes.ci_high < 1)).sum()
print(f"\nSlopes between 0 and 1: {inside} of 8; with the 95% interval excluding 0 and 1: {strict} of 8")
print("Slope range: %.3f to %.3f" % (slopes.slope.min(), slopes.slope.max()))

   dataset prior  trials  slope  ci_low  ci_high  slope_unweighted
 Eye Right Short     129  1.110   1.014    1.206             1.112
 Eye Right  Long     147  0.903   0.769    1.038             0.894
  Eye Left Short     105  1.011   0.896    1.126             1.060
  Eye Left  Long     155  0.845   0.743    0.948             0.850
Hand Right Short     113  0.861   0.728    0.993             0.857
Hand Right  Long     112  0.612   0.488    0.736             0.617
 Hand Left Short     134  0.612   0.499    0.725             0.616
 Hand Left  Long      88  0.502   0.369    0.636             0.496

Slopes between 0 and 1: 6 of 8; with the 95% interval excluding 0 and 1: 5 of 8
Slope range: 0.502 to 1.110

fig, axes = plt.subplots(1, 4, figsize=(14, 3.4), sharey=True)
for ax, (e, th, name) in zip(axes, DS):
    for short, col in [(True, "tab:red"), (False, "tab:blue")]:
        g = good[(good.is_eye == e) & (good.theta == th) & (good.is_short == short)]
        m = g.groupby("ts").tp.agg(["mean", "std"])
        ax.errorbar(m.index, m["mean"], m["std"], color=col, marker="o", capsize=2)
    ax.plot([450, 1250], [450, 1250], "k:", lw=1)
    ax.set_title(name)
    ax.set_xlabel("ts (ms)")
axes[0].set_ylabel("tp (ms), mean and SD")
plt.tight_layout()
plt.show()
print("Mean tp per prior and ts: %.0f to %.0f ms; all points within the axes."
      % (good.groupby(['is_eye', 'theta', 'is_short', 'ts']).tp.mean().min(), good.groupby(['is_eye', 'theta', 'is_short', 'ts']).tp.mean().max()))

Mean tp per prior and ts: 497 to 1177 ms; all points within the axes.

o = good[good.ts == 800]
s_tp, l_tp = o[o.is_short].tp, o[~o.is_short].tp
p = stats.ranksums(s_tp, l_tp).pvalue
print(f"ts = 800 ms: Short n = {len(s_tp)}, median tp - 800 = {s_tp.median() - 800:.1f} ms; "
      f"Long n = {len(l_tp)}, median tp - 800 = {l_tp.median() - 800:.1f} ms; rank-sum p = {p:.2g}")
for e, th, name in DS:
    g = o[(o.is_eye == e) & (o.theta == th)]
    a, b = g[g.is_short].tp, g[~g.is_short].tp
    print(f"   {name}: Short {a.median() - 800:+.1f} (n={len(a)}), Long {b.median() - 800:+.1f} (n={len(b)}), "
          f"p = {stats.ranksums(a, b).pvalue:.2g}")

ts = 800 ms: Short n = 93, median tp - 800 = 28.3 ms; Long n = 107, median tp - 800 = 47.6 ms; rank-sum p = 0.073
   Eye Right: Short +40.9 (n=25), Long +36.5 (n=38), p = 0.81
   Eye Left: Short +79.1 (n=22), Long +18.6 (n=36), p = 0.051
   Hand Right: Short +29.7 (n=22), Long +96.8 (n=19), p = 0.00099
   Hand Left: Short -61.6 (n=24), Long -2.8 (n=14), p = 0.069

rows = []
for e, th, name in DS:
    for short in [True, False]:
        g = good[(good.is_eye == e) & (good.theta == th) & (good.is_short == short)]
        ext, mid = rsg.increments(g.groupby("ts").tp.mean().loc[list(rsg.prior_ts(short))])
        rows.append({"dataset": name, "prior": "Short" if short else "Long", "extreme": ext, "middle": mid})
inc = pd.DataFrame(rows)
inc["middle_minus_extreme"] = inc.middle - inc.extreme
print(inc.round(1).to_string(index=False))
w = stats.wilcoxon(inc.middle, inc.extreme)
print(f"\nMiddle above extreme in {(inc.middle_minus_extreme > 0).sum()} of 8; median difference "
      f"{inc.middle_minus_extreme.median():.1f} ms; two-sided signed-rank p = {w.pvalue:.4f}")

   dataset prior  extreme  middle  middle_minus_extreme
 Eye Right Short     65.7   104.1                  38.4
 Eye Right  Long     49.1   117.0                  67.8
  Eye Left Short     66.2    96.1                  29.9
  Eye Left  Long     78.0    89.3                  11.3
Hand Right Short     47.9    80.6                  32.7
Hand Right  Long     36.8    80.3                  43.5
 Hand Left Short     56.1    45.3                 -10.7
 Hand Left  Long     40.8    56.8                  16.0

Middle above extreme in 7 of 8; median difference 31.3 ms; two-sided signed-rank p = 0.0156

wide = slopes.pivot(index="dataset", columns="prior", values="slope")
wide["long_below_short"] = wide.Long < wide.Short
print(wide.round(3).to_string())
print(f"\nLong slope below Short slope in {int(wide.long_below_short.sum())} of 4 datasets; "
      f"mean Short {wide.Short.mean():.3f}, mean Long {wide.Long.mean():.3f}")
wu = slopes.pivot(index="dataset", columns="prior", values="slope_unweighted")
print(f"Unweighted slopes: Long below Short in {int((wu.Long < wu.Short).sum())} of 4")

prior        Long  Short  long_below_short
dataset                                   
Eye Left    0.845  1.011              True
Eye Right   0.903  1.110              True
Hand Left   0.502  0.612              True
Hand Right  0.612  0.861              True

Long slope below Short slope in 4 of 4 datasets; mean Short 0.898, mean Long 0.716
Unweighted slopes: Long below Short in 4 of 4

models, rows, pred_rows = {}, [], []
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th)]
    m = rsg.BayesObserver.fit(g)
    models[name] = m
    obs_bias, pred_bias, obs_sd, pred_sd = [], [], [], []
    for short in [True, False]:
        ts = np.array(rsg.prior_ts(short), float)
        _, mean_tp, sd_tp = m.moments(ts, short)
        gg = g[g.is_short == short].groupby("ts").tp.agg(["mean", "std"]).loc[ts]
        obs_bias += list(gg["mean"] - ts)
        pred_bias += list(mean_tp - ts)
        obs_sd += list(gg["std"])
        pred_sd += list(sd_tp)
        for t, ob, pb, os_, ps in zip(ts, gg["mean"] - ts, mean_tp - ts, gg["std"], sd_tp):
            pred_rows.append({"dataset": name, "prior": "Short" if short else "Long", "ts": t,
                              "bias_obs": ob, "bias_model": pb, "sd_obs": os_, "sd_model": ps})
    ob, pb = np.array(obs_bias), np.array(pred_bias)
    r2 = 1 - np.sum((ob - pb) ** 2) / np.sum((ob - ob.mean()) ** 2)
    r_sd = stats.pearsonr(obs_sd, pred_sd)[0]
    rows.append({"dataset": name, "wm": m.wm, "wp": m.wp, "b_short": m.b[True], "b_long": m.b[False],
                 "converged": m.converged, "bias_R2": r2, "sd_r": r_sd})
fits = pd.DataFrame(rows)
print(fits.round(3).to_string(index=False))
print(f"\nBias R2 >= 0.8 in {(fits.bias_R2 >= 0.8).sum()} of 4; SD correlation > 0 in {(fits.sd_r > 0).sum()} of 4")
pred = pd.DataFrame(pred_rows)
print("\nObserved and model bias and SD (ms):")
print(pred.round(1).to_string(index=False))

   dataset    wm    wp  b_short  b_long  converged  bias_R2  sd_r
 Eye Right 0.013 0.106    6.411  24.391       True    0.122 0.909
  Eye Left 0.025 0.103   40.838   1.691       True    0.671 0.807
Hand Right 0.074 0.104   36.932  50.359       True    0.865 0.277
 Hand Left 0.117 0.090    7.749 -81.479       True    0.978 0.749

Bias R2 >= 0.8 in 2 of 4; SD correlation > 0 in 4 of 4

Observed and model bias and SD (ms):
   dataset prior     ts  bias_obs  bias_model  sd_obs  sd_model
 Eye Right Short  480.0      16.9        11.8    55.1      51.6
 Eye Right Short  560.0     -16.3         6.6    48.5      60.0
 Eye Right Short  640.0      -9.4         6.6    71.6      68.5
 Eye Right Short  720.0      32.0         6.7    62.7      77.1
 Eye Right Short  800.0      36.6        -2.5    63.3      84.2
 Eye Right  Long  800.0      44.3        33.6    86.3      86.1
 Eye Right  Long  900.0      22.3        24.7    85.8      96.4
 Eye Right  Long 1000.0       2.4        24.7    95.0     107.1
 Eye Right  Long 1100.0      56.2        24.8   127.8     117.8
 Eye Right  Long 1200.0     -23.5        10.9   165.2     126.2
  Eye Left Short  480.0      59.5        51.8    45.1      50.7
  Eye Left Short  560.0      20.0        41.6    52.7      59.3
  Eye Left Short  640.0      26.9        41.7    81.6      67.8
  Eye Left Short  720.0      52.1        41.7   100.4      76.2
  Eye Left Short  800.0      64.0        23.7    79.3      80.9
  Eye Left  Long  800.0      21.3        20.1    93.0      84.5
  Eye Left  Long  900.0      16.0         2.9    91.9      95.3
  Eye Left  Long 1000.0       6.7         3.0   103.5     105.9
  Eye Left  Long 1100.0      -5.3         2.6   112.2     116.3
  Eye Left  Long 1200.0     -44.0       -24.1    91.9     121.3
Hand Right Short  480.0      80.5        72.7    78.1      56.2
Hand Right Short  560.0      33.0        47.2    61.1      69.7
Hand Right Short  640.0      21.1        41.8   106.2      80.5
Hand Right Short  720.0      34.2        27.4   116.4      83.3
Hand Right Short  800.0      17.4       -10.6    80.1      82.4
Hand Right  Long  800.0     117.2       110.2    93.1      93.6
Hand Right  Long  900.0      70.0        72.0    90.7     109.2
Hand Right  Long 1000.0      63.8        55.2    93.4     121.9
Hand Right  Long 1100.0      30.5        29.2    93.6     124.6
Hand Right  Long 1200.0     -48.7       -21.0    87.8     123.5
 Hand Left Short  480.0      68.7        69.6    70.6      54.6
 Hand Left Short  560.0      34.5        35.3    63.4      68.8
 Hand Left Short  640.0      21.5        11.5    65.7      78.0
 Hand Left Short  720.0     -34.9       -19.2    71.5      78.6
 Hand Left Short  800.0     -48.5       -63.7    91.3      75.7
 Hand Left  Long  800.0       8.3        18.5    66.6      89.0
 Hand Left  Long  900.0     -27.4       -34.3    87.7     101.7
 Hand Left  Long 1000.0     -94.2       -79.2    86.5     111.1
 Hand Left  Long 1100.0    -113.8      -128.1   109.7     113.6
 Hand Left  Long 1200.0    -196.4      -187.7   103.4     112.1

fig, axes = plt.subplots(1, 2, figsize=(8, 3.5))
for short, col in [(True, "tab:red"), (False, "tab:blue")]:
    p_ = pred[pred.prior == ("Short" if short else "Long")]
    axes[0].scatter(p_.bias_model, p_.bias_obs, color=col, s=15)
    axes[1].scatter(p_.sd_model, p_.sd_obs, color=col, s=15)
for ax, lab in zip(axes, ["bias (ms)", "SD of tp (ms)"]):
    lim = [min(ax.get_xlim()[0], ax.get_ylim()[0]), max(ax.get_xlim()[1], ax.get_ylim()[1])]
    ax.plot(lim, lim, "k:", lw=1)
    ax.set_xlabel("model " + lab)
    ax.set_ylabel("observed " + lab)
plt.tight_layout()
plt.show()
print("Bias: observed %.0f to %.0f, model %.0f to %.0f ms; SD: observed %.0f to %.0f, model %.0f to %.0f ms; all within the axes."
      % (pred.bias_obs.min(), pred.bias_obs.max(), pred.bias_model.min(), pred.bias_model.max(),
         pred.sd_obs.min(), pred.sd_obs.max(), pred.sd_model.min(), pred.sd_model.max()))

Bias: observed -196 to 117, model -188 to 110 ms; SD: observed 45 to 165, model 51 to 126 ms; all within the axes.

def glm_rows(trs, ts_values):
    # One row per trial and time point t (ms after Ready) for which Set came after t (attrition).
    elapsed = ((trs.set_time - trs.ready_time) * 1000).to_numpy()
    cnt = np.stack([rsg.counts(spikes, units, trs.ready_time.to_numpy() + t / 1000, -rsg.GLM_WINDOW_MS, 0,
                               rsg.GLM_WINDOW_MS)[:, 0, :] for t in ts_values], axis=1)
    rows, ys = [], []
    for i, (_, r) in enumerate(trs.iterrows()):
        for k, t in enumerate(ts_values):
            if elapsed[i] >= t - 1e-6:
                rows.append({"trial": i, "t": k, "prior": float(not r.is_short), "eff": float(not r.is_eye), "dir": float(r.theta == 180)})
                ys.append(cnt[i, k])
    return pd.DataFrame(rows), np.array(ys)


def design(df, with_t=True, with_prior=True):
    cols = [np.ones(len(df)), df.eff.to_numpy(), df["dir"].to_numpy()]
    if with_prior:
        cols.append(df.prior.to_numpy())
    if with_t:
        cols += [(df.t == k).to_numpy().astype(float) for k in range(1, 5)]
    return np.column_stack(cols)


def fit(y, X):
    r = sm.GLM(y, X, family=sm.families.Poisson()).fit()
    return r.llf, X.shape[1]


def bic(llf, k, n):
    return -2 * llf + k * np.log(n)


res = []
df_all, Y_all = glm_rows(good, rsg.SHORT_TS)
per_prior = {short: glm_rows(good[good.is_short == short], rsg.prior_ts(short)) for short in [True, False]}
for j, u in enumerate(units):
    y = Y_all[:, j]
    l_full, k_full = fit(y, design(df_all))
    l_red, k_red = fit(y, design(df_all, with_prior=False))
    n = len(y)
    row = {"unit": int(S["unit_ids"][u]), "prior_dBIC": bic(l_red, k_red, n) - bic(l_full, k_full, n),
           "prior_LR_p": stats.chi2.sf(2 * (l_full - l_red), k_full - k_red)}
    for short, (dfp, Yp) in per_prior.items():
        yp = Yp[:, j]
        lf, kf = fit(yp, design(dfp, with_prior=False))
        lr, kr = fit(yp, design(dfp, with_t=False, with_prior=False))
        row[f"ts_dBIC_{'short' if short else 'long'}"] = bic(lr, kr, len(yp)) - bic(lf, kf, len(yp))
        row[f"ts_LR_p_{'short' if short else 'long'}"] = stats.chi2.sf(2 * (lf - lr), kf - kr)
    res.append(row)
glm = pd.DataFrame(res)
glm["prior_dep"] = glm.prior_dBIC > 0
glm["ts_dep"] = (glm.ts_dBIC_short > 0) | (glm.ts_dBIC_long > 0)


def wilson(k, n):
    lo, hi = sm.stats.proportion_confint(k, n, method="wilson")
    return f"{100 * k / n:.1f}% ({k} of {n}; 95% interval {100 * lo:.1f} to {100 * hi:.1f}%)"


print("Rows (trial x time point): all trials", len(df_all), "; Short", len(per_prior[True][0]), "; Long", len(per_prior[False][0]))
print("Prior-dependent (BIC):", wilson(int(glm.prior_dep.sum()), len(glm)), "; paper, H: 65%")
print("ts-dependent (BIC):   ", wilson(int(glm.ts_dep.sum()), len(glm)), "; paper, H: 27%")
print("Both:", int((glm.prior_dep & glm.ts_dep).sum()), "; prior only:", int((glm.prior_dep & ~glm.ts_dep).sum()),
      "; ts only:", int((~glm.prior_dep & glm.ts_dep).sum()), "; neither:", int((~glm.prior_dep & ~glm.ts_dep).sum()))
print("\nVariants, likelihood-ratio test in place of BIC:")
for a in [0.05, 0.01]:
    pd_ = (glm.prior_LR_p < a).sum()
    td_ = ((glm.ts_LR_p_short < a) | (glm.ts_LR_p_long < a)).sum()
    print(f"   p < {a}: prior {100 * pd_ / len(glm):.1f}%, ts {100 * td_ / len(glm):.1f}%")
print("ts-dependent in Short only:", int(((glm.ts_dBIC_short > 0) & ~(glm.ts_dBIC_long > 0)).sum()),
      "; Long only:", int((~(glm.ts_dBIC_short > 0) & (glm.ts_dBIC_long > 0)).sum()),
      "; both:", int(((glm.ts_dBIC_short > 0) & (glm.ts_dBIC_long > 0)).sum()))

Rows (trial x time point): all trials 3690 ; Short 1198 ; Long 1416
Prior-dependent (BIC): 75.0% (36 of 48; 95% interval 61.2 to 85.1%) ; paper, H: 65%
ts-dependent (BIC):    45.8% (22 of 48; 95% interval 32.6 to 59.7%) ; paper, H: 27%
Both: 19 ; prior only: 17 ; ts only: 3 ; neither: 9

Variants, likelihood-ratio test in place of BIC:
   p < 0.05: prior 87.5%, ts 77.1%
   p < 0.01: prior 77.1%, ts 66.7%
ts-dependent in Short only: 6 ; Long only: 7 ; both: 9

# Chance levels: the same BIC rule with the prior labels shuffled across trials within each
# effector and direction, and with the time points shuffled among the rows of each trial.
rng_glm = np.random.default_rng(7)
N_SHUFFLES = 5


def shuffled_prior(df, trs):
    lab = trs.is_short.to_numpy().copy()
    for _, idx in trs.reset_index(drop=True).groupby(["is_eye", "theta"]).groups.items():
        lab[idx] = rng_glm.permutation(lab[idx])
    out = df.copy()
    out["prior"] = (~lab[df.trial.to_numpy()]).astype(float)
    return out


def shuffled_time(df):
    out = df.copy()
    t = out.t.to_numpy().copy()
    for idx in out.groupby("trial").indices.values():
        t[idx] = rng_glm.permutation(t[idx])
    out["t"] = t
    return out


null_rows = []
for s in range(N_SHUFFLES):
    dfs = shuffled_prior(df_all, good)
    dft = {short: shuffled_time(dfp) for short, (dfp, _) in per_prior.items()}
    n_prior = n_ts_s = n_ts_l = n_ts = 0
    for j in range(len(units)):
        y = Y_all[:, j]
        lf, kf = fit(y, design(dfs))
        lr, kr = fit(y, design(dfs, with_prior=False))
        n_prior += bic(lr, kr, len(y)) - bic(lf, kf, len(y)) > 0
        dep = {}
        for short, (_, Yp) in per_prior.items():
            yp = Yp[:, j]
            lf, kf = fit(yp, design(dft[short], with_prior=False))
            lr, kr = fit(yp, design(dft[short], with_t=False, with_prior=False))
            dep[short] = bic(lr, kr, len(yp)) - bic(lf, kf, len(yp)) > 0
        n_ts_s += dep[True]
        n_ts_l += dep[False]
        n_ts += dep[True] or dep[False]
    null_rows.append({"shuffle": s, "prior_dep_pct": 100 * n_prior / len(units), "ts_dep_short_pct": 100 * n_ts_s / len(units),
                      "ts_dep_long_pct": 100 * n_ts_l / len(units), "ts_dep_either_pct": 100 * n_ts / len(units)})
glm_null = pd.DataFrame(null_rows)
print("Fractions of units passing the BIC rule with shuffled labels (%):")
print(glm_null.round(1).to_string(index=False))
print(f"\nObserved: prior {100 * glm.prior_dep.mean():.1f}%, ts in Short {100 * (glm.ts_dBIC_short > 0).mean():.1f}%, "
      f"ts in Long {100 * (glm.ts_dBIC_long > 0).mean():.1f}%, ts in either {100 * glm.ts_dep.mean():.1f}%")
print(f"Shuffled, maximum over {N_SHUFFLES}: prior {glm_null.prior_dep_pct.max():.1f}%, ts in either {glm_null.ts_dep_either_pct.max():.1f}%")

# Slow drift can mimic a prior effect when the priors come in blocks. Shift the prior labels
# along the session (as in the clock check) so that the block structure is kept.
print("\nPrior labels shifted along the session, block structure kept:")
for shift in [50, 100, 200]:
    lab = np.roll(good.is_short.to_numpy(), shift)
    dfs = df_all.copy()
    dfs["prior"] = (~lab[df_all.trial.to_numpy()]).astype(float)
    n_prior = 0
    for j in range(len(units)):
        y = Y_all[:, j]
        lf, kf = fit(y, design(dfs))
        lr, kr = fit(y, design(dfs, with_prior=False))
        n_prior += bic(lr, kr, len(y)) - bic(lf, kf, len(y)) > 0
    print(f"   shifted by {shift} trials: labels unchanged on {100 * (lab == good.is_short.to_numpy()).mean():.0f}% of trials; "
          f"prior-dependent {100 * n_prior / len(units):.1f}%")
runs = (good.is_short.to_numpy()[1:] != good.is_short.to_numpy()[:-1]).sum() + 1
print(f"Prior blocks in the session (runs of one prior over used trials): {runs}; median length "
      f"{np.median(np.diff(np.r_[0, np.flatnonzero(np.diff(good.is_short.to_numpy().astype(int))) + 1, len(good)])):.0f} trials")

Fractions of units passing the BIC rule with shuffled labels (%):
 shuffle  prior_dep_pct  ts_dep_short_pct  ts_dep_long_pct  ts_dep_either_pct
       0            0.0               2.1              0.0                2.1
       1            4.2               0.0              0.0                0.0
       2            0.0               0.0              0.0                0.0
       3            8.3               0.0              0.0                0.0
       4            0.0               2.1              0.0                2.1

Observed: prior 75.0%, ts in Short 31.2%, ts in Long 33.3%, ts in either 45.8%
Shuffled, maximum over 5: prior 8.3%, ts in either 2.1%

Prior labels shifted along the session, block structure kept:

   shifted by 50 trials: labels unchanged on 56% of trials; prior-dependent 10.4%

   shifted by 100 trials: labels unchanged on 53% of trials; prior-dependent 14.6%

   shifted by 200 trials: labels unchanged on 48% of trials; prior-dependent 2.1%
Prior blocks in the session (runs of one prior over used trials): 87; median length 9 trials

rng_dfr = np.random.default_rng(3)
edges = np.arange(-320, 1440 + 1, 160)
centres = (edges[:-1] + edges[1:]) / 2


def rms_dfr(cnt, short):
    d = cnt[short].mean(0) - cnt[~short].mean(0)          # bins x units, counts per 160 ms
    return np.sqrt(np.mean((d * 1000 / 160) ** 2, axis=1))  # spikes/s


per_ds, nulls = {}, []
cnts = {}
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th) & (good.ts == 800)]
    cnts[name] = (rsg.counts(spikes, units, g.ready_time.to_numpy(), -320, 1440, 160), g.is_short.to_numpy())
    per_ds[name] = rms_dfr(*cnts[name])
mean_rms = np.mean(list(per_ds.values()), axis=0)
for _ in range(200):
    nulls.append(np.mean([rms_dfr(cnt, rng_dfr.permutation(lab)) for cnt, lab in cnts.values()], axis=0))
nulls = np.array(nulls)
tab = pd.DataFrame({"bin_start": edges[:-1], "bin_end": edges[1:], **{k: v for k, v in per_ds.items()},
                    "mean": mean_rms, "shuffle_mean": nulls.mean(0), "shuffle_95": np.percentile(nulls, 95, axis=0)})
print("RMS over units of the Short minus Long rate difference at ts = 800 ms (spikes/s), by 160 ms bin from Ready:")
print(tab.round(2).to_string(index=False))
k = int(np.argmax(mean_rms))
thr = np.percentile(nulls.max(1), 95)
print(f"\nLargest mean RMS: {mean_rms[k]:.2f} spikes/s in the bin {edges[k]} to {edges[k + 1]} ms after Ready; "
      f"95th percentile of the shuffled maximum: {thr:.2f}")
print("Largest bin per dataset:", {n: f"{edges[np.argmax(v)]}-{edges[np.argmax(v) + 1]}" for n, v in per_ds.items()})
print("Used trials at ts = 800 ms per dataset (Short, Long):",
      {n: (int(l.sum()), int((~l).sum())) for n, (c_, l) in cnts.items()})

RMS over units of the Short minus Long rate difference at ts = 800 ms (spikes/s), by 160 ms bin from Ready:
 bin_start  bin_end  Eye Right  Eye Left  Hand Right  Hand Left  mean  shuffle_mean  shuffle_95
      -320     -160       2.16      2.24        2.30       2.87  2.39          2.09        2.42
      -160        0       1.97      1.85        2.59       2.73  2.29          2.05        2.32
         0      160       2.58      2.77        2.78       2.58  2.68          2.14        2.44
       160      320       3.51      4.36        3.75       4.73  4.09          2.60        3.01
       320      480       8.92      7.84        2.91       4.19  5.97          2.38        2.78
       480      640       9.14      9.49        3.94       4.88  6.86          2.45        2.91
       640      800       4.59      5.97        4.39       4.61  4.89          2.33        2.60
       800      960       6.81      9.73        4.24       3.49  6.07          2.45        2.80
       960     1120       5.15      4.83        3.82       4.58  4.59          2.58        2.97
      1120     1280       4.41      4.46        3.44       2.61  3.73          2.41        2.76
      1280     1440       1.88      2.70        2.15       3.04  2.45          2.22        2.61

Largest mean RMS: 6.86 spikes/s in the bin 480 to 640 ms after Ready; 95th percentile of the shuffled maximum: 3.11
Largest bin per dataset: {'Eye Right': '480-640', 'Eye Left': '800-960', 'Hand Right': '640-800', 'Hand Left': '480-640'}
Used trials at ts = 800 ms per dataset (Short, Long): {'Eye Right': (25, 38), 'Eye Left': (22, 36), 'Hand Right': (22, 19), 'Hand Left': (24, 14)}

plt.figure(figsize=(6, 3.2))
for n, v in per_ds.items():
    plt.plot(centres, v, color="0.6", lw=1)
plt.plot(centres, mean_rms, "k", lw=2, label="mean of datasets")
plt.fill_between(centres, np.percentile(nulls, 5, axis=0), np.percentile(nulls, 95, axis=0), color="tab:green", alpha=.2, label="shuffle 5-95%")
plt.axvspan(480, 800, color="tab:red", alpha=.1, label="Short support")
plt.axvline(0, color="k", ls=":")
plt.axvline(800, color="k", ls="--")
plt.xlabel("time from Ready (ms); Set at 800")
plt.ylabel("RMS rate difference (spikes/s)")
plt.legend(fontsize=7)
plt.show()
print("Plotted values: %.2f to %.2f spikes/s; all within the axes." % (min(v.min() for v in per_ds.values()), max(v.max() for v in per_ds.values())))

Plotted values: 1.85 to 9.73 spikes/s; all within the axes.

# Stability of the peak: resample trials with replacement within each dataset and prior, and
# find the bin of the largest mean RMS difference.
rng_dfr_boot = np.random.default_rng(8)
peaks = []
for _ in range(200):
    per = []
    for cnt, lab in cnts.values():
        i = np.concatenate([rng_dfr_boot.choice(np.flatnonzero(lab == v), (lab == v).sum()) for v in (True, False)])
        per.append(rms_dfr(cnt[i], lab[i]))
    peaks.append(int(np.argmax(np.mean(per, axis=0))))
peaks = pd.Series(peaks).map(lambda k: f"{edges[k]} to {edges[k + 1]}")
print("Bin of the largest mean RMS over 200 bootstrap samples of trials (ms after Ready):")
print(peaks.value_counts().to_string())
in_support = peaks.isin(["480 to 640", "640 to 800"]).mean()
print(f"Fraction of samples with the peak between 480 and 800 ms: {in_support:.2f}")

Bin of the largest mean RMS over 200 bootstrap samples of trials (ms after Ready):
480 to 640    192
800 to 960      8
Fraction of samples with the peak between 480 and 800 ms: 0.96

EST, PROD, rows = {}, {}, []
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th)]
    est = {}
    for short in [True, False]:
        gg = g[g.is_short == short]
        rate, ends, n = rsg.estimation_rates(spikes, units, gg, max(rsg.prior_ts(short)))
        est[short] = {"rate": rate, "t": ends, "n": n}
    pca_e = rsg.PCA([est[True]["rate"], est[False]["rate"]])
    prod = {}
    for short in [True, False]:
        for ts in rsg.prior_ts(short):
            gg = g[(g.is_short == short) & (g.ts == ts)]
            prod[(short, ts)] = {"rate": rsg.production_rates(spikes, units, gg), "min_tp": gg.tp.min(),
                                 "mean_tp": gg.tp.mean(), "n": len(gg)}
    pca_p = rsg.PCA([v["rate"] for v in prod.values()])
    EST[name] = {"pca": pca_e, **est}
    PROD[name] = {"pca": pca_p, "traj": prod}
    cum_e, cum_p = np.cumsum(pca_e.explained), np.cumsum(pca_p.explained)
    rows.append({"dataset": name, "est_PCs_75": pca_e.n, "prod_PCs_75": pca_p.n,
                 "est_PCs_70": int(np.searchsorted(cum_e, .7) + 1), "est_PCs_80": int(np.searchsorted(cum_e, .8) + 1),
                 "prod_PCs_70": int(np.searchsorted(cum_p, .7) + 1), "prod_PCs_80": int(np.searchsorted(cum_p, .8) + 1),
                 "est_var_3PC": cum_e[2], "prod_var_5PC": cum_p[4]})
dims = pd.DataFrame(rows)
print(dims.round(3).to_string(index=False))
n_ok = dims.est_PCs_75.between(3, 4).sum() + dims.prod_PCs_75.between(5, 11).sum()
print(f"\nIn range (estimation 3 or 4, production 5 to 11): {n_ok} of 8")
print("Shortest tp per ts, used to cut production trajectories (ms):",
      {f"{'S' if k[0] else 'L'}{k[1]}": round(min(PROD[n]['traj'][k]['min_tp'] for n in PROD)) for k in PROD[DS[0][2]]['traj']})

   dataset  est_PCs_75  prod_PCs_75  est_PCs_70  est_PCs_80  prod_PCs_70  prod_PCs_80  est_var_3PC  prod_var_5PC
 Eye Right           2            2           2           2            2            2        0.902         0.916
  Eye Left           2            2           2           2            2            2        0.899         0.911
Hand Right           2            5           2           3            4            6        0.869         0.771
 Hand Left           2            4           2           3            3            5        0.863         0.810

In range (estimation 3 or 4, production 5 to 11): 1 of 8
Shortest tp per ts, used to cut production trajectories (ms): {'S480': 407, 'S560': 455, 'S640': 447, 'S720': 546, 'S800': 602, 'L800': 637, 'L900': 711, 'L1000': 736, 'L1100': 794, 'L1200': 797}

# Variants: rates normalised per unit before PCA (soft normalisation, and z-scoring), with the
# components reaching 75% of the variance. Estimation and production scales are taken from the
# trajectories that enter each PCA.
rows = []
for e, th, name in DS:
    E, P = EST[name], PROD[name]
    est_m = [E[True]["rate"], E[False]["rate"]]
    prod_m = [v["rate"] for v in P["traj"].values()]
    for how in ["soft", "z"]:
        se, sp = rsg.unit_scale(est_m, how), rsg.unit_scale(prod_m, how)
        pe = rsg.PCA([m / se for m in est_m])
        pp = rsg.PCA([m / sp for m in prod_m])
        E[f"pca_{how}"], E[f"scale_{how}"] = pe, se
        rows.append({"dataset": name, "normalisation": how, "est_PCs_75": pe.n, "prod_PCs_75": pp.n})
dims_norm = pd.DataFrame(rows)
print(dims_norm.to_string(index=False))
for how, g in dims_norm.groupby("normalisation"):
    print(f"{how}: in range (estimation 3 or 4, production 5 to 11) {g.est_PCs_75.between(3, 4).sum() + g.prod_PCs_75.between(5, 11).sum()} of 8")

   dataset normalisation  est_PCs_75  prod_PCs_75
 Eye Right          soft           3            4
 Eye Right             z           3            6
  Eye Left          soft           3            4
  Eye Left             z           3            6
Hand Right          soft           3            7
Hand Right             z           4           10
 Hand Left          soft           4            7
 Hand Left             z           4            9
soft: in range (estimation 3 or 4, production 5 to 11) 6 of 8
z: in range (estimation 3 or 4, production 5 to 11) 8 of 8

rng_axes = np.random.default_rng(4)
PROJ, rows = {}, []
for e, th, name in DS:
    E = EST[name]
    for short in [True, False]:
        lo, hi = min(rsg.prior_ts(short)), max(rsg.prior_ts(short))
        X = E["pca"](E[short]["rate"])
        t = E[short]["t"]
        u = rsg.encoding_axis(X, t, lo, hi)
        sel = (t >= lo - 1e-6) & (t <= hi + 1e-6)
        xu, tt, w = X[sel] @ u, t[sel], E[short]["n"][sel]
        te = models[name].moments(tt, short)[0]
        _, _, r2_b, rmse_b = rsg.weighted_fit(te, xu, w)
        _, slope_ts, r2_l, rmse_l = rsg.weighted_fit(tt, xu, w)
        rand = []
        for _ in range(1000):
            v = rng_axes.standard_normal(X.shape[1])
            rand.append(rsg.weighted_fit(te, X[sel] @ (v / np.linalg.norm(v)), w)[2])
        rand = np.array(rand)
        PROJ[(name, short)] = {"t": tt, "xu": xu, "w": w, "te": te, "u": u, "X": X, "tall": t}
        rows.append({"dataset": name, "prior": "Short" if short else "Long", "points": int(sel.sum()), "PCs": X.shape[1],
                     "R2_bayes": r2_b, "RMSE_bayes": rmse_b, "R2_linear": r2_l, "RMSE_linear": rmse_l,
                     "slope_ts": slope_ts, "random_R2_median": np.median(rand),
                     "random_R2_95": np.percentile(rand, 95), "frac_random_above": np.mean(rand >= r2_b)})
enc = pd.DataFrame(rows)
print(enc.round(4).to_string(index=False))
hl = enc[enc.dataset == "Hand Left"].set_index("prior").R2_bayes
mean_r2 = enc.groupby("prior").R2_bayes.mean()
print(f"\nReading (a), Hand Left: Short R2 = {hl['Short']:.4f}, Long R2 = {hl['Long']:.4f}")
print(f"Reading (b), mean of 4 datasets: Short R2 = {mean_r2['Short']:.4f}, Long R2 = {mean_r2['Long']:.4f}")
print("Tolerance: R2 >= 0.894 (Short) and >= 0.896 (Long)")

   dataset prior  points  PCs  R2_bayes  RMSE_bayes  R2_linear  RMSE_linear  slope_ts  random_R2_median  random_R2_95  frac_random_above
 Eye Right Short      17    2    0.9581      3.1506     0.9533       3.3274    0.1899            0.9433        0.9965              0.396
 Eye Right  Long      21    2    0.9766      2.5824     0.9717       2.7761    0.1203            0.9732        0.9780              0.286
  Eye Left Short      17    2    0.9097      6.2773     0.8975       6.7255    0.2288            0.9053        0.9793              0.475
  Eye Left  Long      21    2    0.9889      1.5583     0.9798       2.1310    0.1249            0.9722        0.9940              0.265
Hand Right Short      17    2    0.9699      0.9907     0.9723       0.8045    0.0497            0.9135        0.9896              0.284
Hand Right  Long      21    2    0.9939      0.4332     0.9904       0.5557    0.0484            0.9833        0.9938              0.046
 Hand Left Short      17    2    0.9982      0.3730     0.9975       0.4362    0.0803            0.8963        0.9975              0.025
 Hand Left  Long      21    2    0.9960      0.2070     0.9936       0.2521    0.0272            0.7347        0.9942              0.028

Reading (a), Hand Left: Short R2 = 0.9982, Long R2 = 0.9960
Reading (b), mean of 4 datasets: Short R2 = 0.9590, Long R2 = 0.9888
Tolerance: R2 >= 0.894 (Short) and >= 0.896 (Long)

fig, axes = plt.subplots(1, 4, figsize=(14, 3.2))
for ax, (e, th, name) in zip(axes, DS):
    for short, col in [(True, "tab:red"), (False, "tab:blue")]:
        p_ = PROJ[(name, short)]
        a, b, _, _ = rsg.weighted_fit(p_["te"], p_["xu"], p_["w"])
        ax.plot(p_["t"], (p_["xu"] - a) / b, "o", color=col, ms=3)   # projections mapped onto te by the fit
        ax.plot(p_["t"], p_["te"], "-", color=col)
    ax.plot([450, 1250], [450, 1250], "k:", lw=1)
    ax.set_title(name)
    ax.set_xlabel("time from Ready = ts (ms)")
axes[0].set_ylabel("te (line) and mapped Xu (dots), ms")
plt.tight_layout()
plt.show()
vals = np.concatenate([(PROJ[k]["xu"] - rsg.weighted_fit(PROJ[k]["te"], PROJ[k]["xu"], PROJ[k]["w"])[0]) / rsg.weighted_fit(PROJ[k]["te"], PROJ[k]["xu"], PROJ[k]["w"])[1] for k in PROJ])
print("Mapped projections: %.0f to %.0f ms; all within the axes." % (vals.min(), vals.max()))

Mapped projections: 515 to 1159 ms; all within the axes.

# Variant: the encoding axis in the PC spaces of normalised rates (pca-dimensionality).
rng_norm = np.random.default_rng(9)
enc_rows = []
for e, th, name in DS:
    E = EST[name]
    for how in ["soft", "z"]:
        pe, se = E[f"pca_{how}"], E[f"scale_{how}"]
        for short in [True, False]:
            lo, hi = min(rsg.prior_ts(short)), max(rsg.prior_ts(short))
            X, t = pe(E[short]["rate"] / se), E[short]["t"]
            u = rsg.encoding_axis(X, t, lo, hi)
            sel = (t >= lo - 1e-6) & (t <= hi + 1e-6)
            w, te = E[short]["n"][sel], PROJ[(name, short)]["te"]
            r2 = rsg.weighted_fit(te, X[sel] @ u, w)[2]
            rand = []
            for _ in range(1000):
                v = rng_norm.standard_normal(X.shape[1])
                rand.append(rsg.weighted_fit(te, X[sel] @ (v / np.linalg.norm(v)), w)[2])
            enc_rows.append({"dataset": name, "prior": "Short" if short else "Long", "normalisation": how, "PCs": X.shape[1],
                             "R2_bayes": r2, "random_R2_median": np.median(rand), "frac_random_above": np.mean(np.array(rand) >= r2)})
enc_norm = pd.DataFrame(enc_rows)
print("Encoding axis with normalised rates (75% of variance):")
print(enc_norm.round(4).to_string(index=False))
for how, g in enc_norm.groupby("normalisation"):
    hl = g[g.dataset == "Hand Left"].set_index("prior").R2_bayes
    print(f"{how}: Hand Left R2 Short {hl['Short']:.4f}, Long {hl['Long']:.4f}; mean of 4 Short {g[g.prior == 'Short'].R2_bayes.mean():.4f}, "
          f"Long {g[g.prior == 'Long'].R2_bayes.mean():.4f}; above 95% of random axes in {(g.frac_random_above < 0.05).sum()} of 8")

Encoding axis with normalised rates (75% of variance):
   dataset prior normalisation  PCs  R2_bayes  random_R2_median  frac_random_above
 Eye Right Short          soft    3    0.9656            0.9105              0.296
 Eye Right  Long          soft    3    0.9815            0.9207              0.144
 Eye Right Short             z    3    0.9558            0.8894              0.297
 Eye Right  Long             z    3    0.9775            0.8829              0.108
  Eye Left Short          soft    3    0.9428            0.9067              0.372
  Eye Left  Long          soft    3    0.9904            0.9230              0.049
  Eye Left Short             z    3    0.9402            0.9089              0.395
  Eye Left  Long             z    3    0.9858            0.8762              0.076
Hand Right Short          soft    3    0.9759            0.8432              0.108
Hand Right  Long          soft    3    0.9967            0.8978              0.031
Hand Right Short             z    4    0.9731            0.7605              0.092
Hand Right  Long             z    4    0.9514            0.8029              0.175
 Hand Left Short          soft    4    0.9842            0.8982              0.185
 Hand Left  Long          soft    4    0.9749            0.7946              0.074
 Hand Left Short             z    4    0.9900            0.8990              0.142
 Hand Left  Long             z    4    0.9596            0.7884              0.093
soft: Hand Left R2 Short 0.9842, Long 0.9749; mean of 4 Short 0.9671, Long 0.9859; above 95% of random axes in 2 of 8
z: Hand Left R2 Short 0.9900, Long 0.9596; mean of 4 Short 0.9648, Long 0.9686; above 95% of random axes in 0 of 8

enc["dRMSE_bayes_minus_linear"] = enc.RMSE_bayes - enc.RMSE_linear
print(enc[["dataset", "prior", "RMSE_bayes", "RMSE_linear", "dRMSE_bayes_minus_linear"]].round(3).to_string(index=False))
for prior, g in enc.groupby("prior", sort=False):
    print(f"{prior}: Bayesian lower in {(g.dRMSE_bayes_minus_linear < 0).sum()} of 4; "
          f"signed-rank p = {stats.wilcoxon(g.RMSE_bayes, g.RMSE_linear).pvalue:.3f} (smallest possible with 4: 0.125)")

   dataset prior  RMSE_bayes  RMSE_linear  dRMSE_bayes_minus_linear
 Eye Right Short       3.151        3.327                    -0.177
 Eye Right  Long       2.582        2.776                    -0.194
  Eye Left Short       6.277        6.726                    -0.448
  Eye Left  Long       1.558        2.131                    -0.573
Hand Right Short       0.991        0.805                     0.186
Hand Right  Long       0.433        0.556                    -0.122
 Hand Left Short       0.373        0.436                    -0.063
 Hand Left  Long       0.207        0.252                    -0.045
Short: Bayesian lower in 3 of 4; signed-rank p = 0.625 (smallest possible with 4: 0.125)
Long: Bayesian lower in 4 of 4; signed-rank p = 0.125 (smallest possible with 4: 0.125)

# Variant: RMSE weighted by the number of trials in each bin, as in the fit.
def wrmse(x, y, w):
    a, b, _, _ = rsg.weighted_fit(x, y, w)
    return float(np.sqrt(np.sum(w * (y - a - b * x) ** 2) / np.sum(w)))


enc["dWRMSE"] = [wrmse(PROJ[(r.dataset, r.prior == "Short")]["te"], PROJ[(r.dataset, r.prior == "Short")]["xu"], PROJ[(r.dataset, r.prior == "Short")]["w"])
                 - wrmse(PROJ[(r.dataset, r.prior == "Short")]["t"], PROJ[(r.dataset, r.prior == "Short")]["xu"], PROJ[(r.dataset, r.prior == "Short")]["w"])
                 for r in enc.itertuples()]
for prior, g in enc.groupby("prior", sort=False):
    print(f"Weighted RMSE, {prior}: Bayesian lower in {(g.dWRMSE < 0).sum()} of 4; differences " + ", ".join(f"{x:+.3f}" for x in g.dWRMSE))

Weighted RMSE, Short: Bayesian lower in 3 of 4; differences -0.192, -0.424, +0.031, -0.052
Weighted RMSE, Long: Bayesian lower in 4 of 4; differences -0.204, -0.509, -0.103, -0.050

rows = []
for e, th, name in DS:
    p_ = PROJ[(name, False)]
    X, t = p_["X"], p_["tall"]
    u = rsg.encoding_axis(X, t, 480, 800)
    sel = (t >= 480 - 1e-6) & (t <= 800 + 1e-6)
    xu, tt = X[sel] @ u, t[sel]
    w = EST[name][False]["n"][sel]
    te = models[name].moments(tt, True)[0]
    rows.append({"dataset": name, "RMSE_bayes": rsg.weighted_fit(te, xu, w)[3], "RMSE_linear": rsg.weighted_fit(tt, xu, w)[3],
                 "R2_bayes": rsg.weighted_fit(te, xu, w)[2], "R2_linear": rsg.weighted_fit(tt, xu, w)[2]})
sil = pd.DataFrame(rows)
sil["dRMSE"] = sil.RMSE_bayes - sil.RMSE_linear
print(sil.round(4).to_string(index=False))
p_sil = stats.wilcoxon(sil.RMSE_bayes, sil.RMSE_linear).pvalue
print(f"\nShort In Long: Bayesian lower in {(sil.dRMSE < 0).sum()} of 4; median dRMSE {sil.dRMSE.median():.3f}; "
      f"signed-rank p = {p_sil:.3f} (smallest possible with 4 datasets: 0.125)")
sh = enc[enc.prior == "Short"]
print(f"For comparison, Short prior: median dRMSE {sh.dRMSE_bayes_minus_linear.median():.3f}; "
      f"relative dRMSE (dRMSE / RMSE_linear) Short In Long {np.median(sil.dRMSE / sil.RMSE_linear):.3f}, "
      f"Short {np.median(sh.dRMSE_bayes_minus_linear / sh.RMSE_linear):.3f}")

   dataset  RMSE_bayes  RMSE_linear  R2_bayes  R2_linear   dRMSE
 Eye Right      1.1392       1.2534    0.9931     0.9916 -0.1142
  Eye Left      1.4967       1.7048    0.9886     0.9853 -0.2082
Hand Right      1.0873       1.1897    0.8714     0.8461 -0.1024
 Hand Left      0.2877       0.2268    0.9952     0.9970  0.0609

Short In Long: Bayesian lower in 3 of 4; median dRMSE -0.108; signed-rank p = 0.250 (smallest possible with 4 datasets: 0.125)
For comparison, Short prior: median dRMSE -0.120; relative dRMSE (dRMSE / RMSE_linear) Short In Long -0.089, Short -0.060

sl = enc.pivot(index="dataset", columns="prior", values="slope_ts")
sl["long_below_short"] = sl.Long < sl.Short
print("Slope of Xu on ts (spikes/s per ms, shared PC space of each dataset):")
print(sl.round(4).to_string())
print(f"\nLong below Short in {int(sl.long_below_short.sum())} of 4 datasets; ratio Long / Short: "
      + ", ".join(f"{r:.2f}" for r in sl.Long / sl.Short))

Slope of Xu on ts (spikes/s per ms, shared PC space of each dataset):
prior         Long   Short  long_below_short
dataset                                     
Eye Left    0.1249  0.2288              True
Eye Right   0.1203  0.1899              True
Hand Left   0.0272  0.0803              True
Hand Right  0.0484  0.0497              True

Long below Short in 4 of 4 datasets; ratio Long / Short: 0.55, 0.63, 0.34, 0.97

rows = []
for e, th, name in DS:
    P = PROD[name]
    for (short, ts), v in P["traj"].items():
        states = P["pca"](v["rate"])
        p_ = PROJ[(name, short)]
        rows.append({"dataset": name, "prior": "Short" if short else "Long", "ts": ts,
                     "xu_at_set": p_["xu"][np.argmin(abs(p_["t"] - ts))], "speed": rsg.step_speed(states),
                     "duration_ms": len(states) * rsg.BIN_MS, "mean_tp": v["mean_tp"], "trials": v["n"]})
spd = pd.DataFrame(rows)
spd["speed_per_tp"] = spd.speed / spd.mean_tp
print(spd.round(3).to_string(index=False))
for prior in ["Short", "Long"]:
    g = spd[spd.prior == prior]
    r, p = stats.pearsonr(g.xu_at_set, g.speed)
    rb, pb = stats.pearsonr(g.xu_at_set, g.speed_per_tp)
    z = g.groupby("dataset").xu_at_set.transform(lambda x: (x - x.mean()) / x.std())
    rz, pz = stats.pearsonr(z, g.speed)
    within = [stats.pearsonr(h.xu_at_set, h.speed)[0] for _, h in g.groupby("dataset")]
    print(f"{prior}: rho = {r:.3f} (p = {p:.2g}); variant speed / mean tp: rho = {rb:.3f} (p = {pb:.2g}); "
          f"Xu z-scored per dataset: rho = {rz:.3f} (p = {pz:.2g}); within datasets: " + ", ".join(f"{x:.2f}" for x in within))

   dataset prior   ts  xu_at_set  speed  duration_ms  mean_tp  trials  speed_per_tp
 Eye Right Short  480    -41.591  7.926          400  496.868      24         0.016
 Eye Right Short  560    -37.181  7.723          440  543.739      27         0.014
 Eye Right Short  640    -19.017  6.358          500  630.550      31         0.010
 Eye Right Short  720      1.280  5.299          640  752.037      22         0.007
 Eye Right Short  800     13.480  4.379          720  836.602      25         0.005
 Eye Right  Long  800    -15.897  5.908          700  844.335      38         0.007
 Eye Right  Long  900     -7.900  4.806          760  922.279      34         0.005
 Eye Right  Long 1000      8.071  4.158          840 1002.392      23         0.004
 Eye Right  Long 1100     19.829  3.902          920 1156.215      25         0.003
 Eye Right  Long 1200     23.605  3.392          880 1176.536      27         0.003
  Eye Left Short  480    -44.640  9.128          460  539.534      21         0.017
  Eye Left Short  560    -44.994  7.389          460  579.976      28         0.013
  Eye Left Short  640    -23.645  6.881          540  666.884      17         0.010
  Eye Left Short  720      8.176  5.069          580  772.095      17         0.007
  Eye Left Short  800     16.639  5.195          700  864.015      22         0.006
  Eye Left  Long  800    -18.508  6.807          620  821.290      36         0.008
  Eye Left  Long  900    -10.591  5.139          740  916.020      29         0.006
  Eye Left  Long 1000      2.436  4.471          760 1006.664      30         0.004
  Eye Left  Long 1100     19.204  4.111          840 1094.700      29         0.004
  Eye Left  Long 1200     24.193  3.691          960 1155.967      31         0.003
Hand Right Short  480      5.492  6.303          400  560.489      21         0.011
Hand Right Short  560      7.982  5.331          460  593.034      35         0.009
Hand Right Short  640     12.336  5.168          440  661.104      14         0.008
Hand Right Short  720     16.562  4.196          540  754.223      21         0.006
Hand Right Short  800     22.484  3.855          660  817.430      22         0.005
Hand Right  Long  800      3.034  4.939          740  917.209      19         0.005
Hand Right  Long  900      7.330  4.211          820  969.975      17         0.004
Hand Right  Long 1000     13.597  3.621          900 1063.784      30         0.003
Hand Right  Long 1100     16.990  3.146         1000 1130.492      21         0.003
Hand Right  Long 1200     21.276  3.288          980 1151.263      25         0.003
 Hand Left Short  480     -2.382  6.056          420  548.726      33         0.011
 Hand Left Short  560      3.404  5.269          440  594.515      21         0.009
 Hand Left Short  640      9.793  4.758          520  661.548      25         0.007
 Hand Left Short  720     17.029  3.769          560  685.139      31         0.006
 Hand Left Short  800     21.626  3.645          600  751.458      24         0.005
 Hand Left  Long  800      0.049  4.821          680  808.309      14         0.006
 Hand Left  Long  900      2.108  3.793          700  872.599      24         0.004
 Hand Left  Long 1000      5.146  3.954          720  905.788      16         0.004
 Hand Left  Long 1100      7.728  3.817          780  986.246      16         0.004
 Hand Left  Long 1200      9.737  3.683          780 1003.612      18         0.004
Short: rho = -0.941 (p = 6.2e-10); variant speed / mean tp: rho = -0.917 (p = 1.4e-08); Xu z-scored per dataset: rho = -0.759 (p = 0.00011); within datasets: -0.92, -1.00, -0.99, -0.96
Long: rho = -0.878 (p = 3.7e-07); variant speed / mean tp: rho = -0.923 (p = 6.5e-09); Xu z-scored per dataset: rho = -0.779 (p = 5.3e-05); within datasets: -0.92, -0.96, -0.75, -0.95

# Both quantities depend on ts. Variant: residuals of each on ts (linear, within dataset and
# prior), pooled, so that only the variation not explained by ts is correlated.
def resid_on_ts(d, col):
    return d.groupby("dataset")[[col, "ts"]].apply(
        lambda h: pd.Series(h[col] - np.polyval(np.polyfit(h.ts, h[col], 1), h.ts), index=h.index)).droplevel(0)


for prior in ["Short", "Long"]:
    g = spd[spd.prior == prior]
    r, p = stats.pearsonr(resid_on_ts(g, "xu_at_set").loc[g.index], resid_on_ts(g, "speed").loc[g.index])
    print(f"{prior}: rho of residuals on ts = {r:.3f} (nominal p = {p:.2g}, not adjusted for the 8 fitted lines)")

Short: rho of residuals on ts = -0.195 (nominal p = 0.41, not adjusted for the 8 fitted lines)
Long: rho of residuals on ts = 0.069 (nominal p = 0.77, not adjusted for the 8 fitted lines)

rows, dropped = [], 0
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th)]
    for short in [True, False]:
        ts_v = np.array(rsg.prior_ts(short), float)
        sp = []
        for ts in ts_v:
            gg = g[(g.is_short == short) & (g.ts == ts)]
            dropped += int((gg.tp < 400).sum())
            gg = gg[gg.tp >= 400]
            sp.append(rsg.step_speed(PROD[name]["pca"](rsg.production_rates(spikes, units, gg, 400))))
        sp = np.array(sp)
        unbiased = sp[2] * ts_v[2] / ts_v
        rows.append({"dataset": name, "prior": "Short" if short else "Long", **{f"speed_{int(t)}": s for t, s in zip(ts_v, sp)},
                     "slope_measured": np.polyfit(ts_v, sp, 1)[0], "slope_unbiased": np.polyfit(ts_v, unbiased, 1)[0]})
ini = pd.DataFrame(rows)
ini["measured_minus_unbiased"] = ini.slope_measured - ini.slope_unbiased
print(ini.round(5).to_string(index=False))
w = stats.wilcoxon(ini.slope_measured, ini.slope_unbiased)
print(f"\nMeasured shallower than unbiased in {(ini.measured_minus_unbiased > 0).sum()} of 8; "
      f"two-sided signed-rank p = {w.pvalue:.4f}; trials dropped for tp < 400 ms: {dropped}")
print("Measured slope negative in", int((ini.slope_measured < 0).sum()), "of 8")

   dataset prior  speed_480  speed_560  speed_640  speed_720  speed_800  slope_measured  slope_unbiased  speed_900  speed_1000  speed_1100  speed_1200  measured_minus_unbiased
 Eye Right Short    7.92627    7.76494    6.71340    6.16218    4.80529        -0.00981        -0.01108        NaN         NaN         NaN         NaN                  0.00128
 Eye Right  Long        NaN        NaN        NaN        NaN    8.25311        -0.00778        -0.00670    7.05230     6.46757     5.90688     4.93791                 -0.00108
  Eye Left Short    9.18311    7.59299    7.40464    5.96920    6.25443        -0.00935        -0.01222        NaN         NaN         NaN         NaN                  0.00287
  Eye Left  Long        NaN        NaN        NaN        NaN    8.95276        -0.00563        -0.00662    7.27091     6.39416     6.48285     6.53402                  0.00099
Hand Right Short    6.30300    5.61327    5.57478    4.99242    4.80451        -0.00452        -0.00920        NaN         NaN         NaN         NaN                  0.00468
Hand Right  Long        NaN        NaN        NaN        NaN    7.01812        -0.00482        -0.00565    5.86106     5.45497     4.80427     5.13790                  0.00083
 Hand Left Short    6.37733    5.33394    5.15677    4.12045    4.10442        -0.00720        -0.00851        NaN         NaN         NaN         NaN                  0.00131
 Hand Left  Long        NaN        NaN        NaN        NaN    6.19997        -0.00215        -0.00533    4.86244     5.14397     5.14086     4.98508                  0.00317

Measured shallower than unbiased in 7 of 8; two-sided signed-rank p = 0.0391; trials dropped for tp < 400 ms: 0
Measured slope negative in 8 of 8

for prior in ["Short", "Long"]:
    g = spd[spd.prior == prior]
    r, p = stats.pearsonr(g.speed, g.mean_tp)
    rb, pb = stats.pearsonr(g.speed_per_tp, g.mean_tp)
    within = [stats.pearsonr(h.speed, h.mean_tp)[0] for _, h in g.groupby("dataset")]
    print(f"{prior}: rho(speed, mean tp) = {r:.3f} (p = {p:.2g}), 20 points; variant speed / mean tp: {rb:.3f} (p = {pb:.2g}); "
          "within datasets: " + ", ".join(f"{x:.2f}" for x in within))

# Variant: the Neural Latents Benchmark's measure, trial by trial within prior, effector and
# direction; single-trial rates smoothed as above and projected on the production PCs, from Set
# to the trial's own Go.
rows, trial_rows = [], []
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th)]
    for short in [True, False]:
        gg = g[g.is_short == short]
        c_ = rsg.counts(spikes, units, gg.set_time.to_numpy(), 0, 2400)
        sp = []
        for i, tp in enumerate(gg.tp.to_numpy()):
            nb = int(tp // rsg.BIN_MS)
            sp.append(rsg.step_speed(PROD[name]["pca"](rsg.smooth(c_[i, :nb] * 1000 / rsg.BIN_MS))))
        r, p = stats.pearsonr(sp, gg.tp)
        trial_rows += [{"dataset": name, "prior": "Short" if short else "Long", "ts": a, "tp": b, "speed": s_}
                       for a, b, s_ in zip(gg.ts, gg.tp, sp)]
        rows.append({"dataset": name, "prior": "Short" if short else "Long", "trials": len(gg), "r": r, "p": p})
nlb = pd.DataFrame(rows)
print("\nTrial-by-trial speed against tp (NLB-style):")
print(nlb.round(3).to_string(index=False))
print(f"Mean r over the 8 conditions: {nlb.r.mean():.3f}")

Short: rho(speed, mean tp) = -0.733 (p = 0.00024), 20 points; variant speed / mean tp: -0.861 (p = 1.1e-06); within datasets: -0.93, -0.99, -0.96, -0.96
Long: rho(speed, mean tp) = -0.744 (p = 0.00017), 20 points; variant speed / mean tp: -0.858 (p = 1.3e-06); within datasets: -0.96, -0.94, -0.80, -0.97

Trial-by-trial speed against tp (NLB-style):
   dataset prior  trials      r     p
 Eye Right Short     129 -0.740 0.000
 Eye Right  Long     147 -0.539 0.000
  Eye Left Short     105 -0.696 0.000
  Eye Left  Long     155 -0.578 0.000
Hand Right Short     113 -0.546 0.000
Hand Right  Long     112 -0.420 0.000
 Hand Left Short     134 -0.442 0.000
 Hand Left  Long      88 -0.195 0.069
Mean r over the 8 conditions: -0.519

for prior in ["Short", "Long"]:
    g = spd[spd.prior == prior]
    r, p = stats.pearsonr(resid_on_ts(g, "speed").loc[g.index], resid_on_ts(g, "mean_tp").loc[g.index])
    print(f"{prior}: rho of residuals on ts (speed, mean tp) = {r:.3f} (nominal p = {p:.2g})")

# Variant: the paper's control for co-dependence on ts (Methods), trial by trial; speed and tp
# z-scored within each dataset, prior and ts, then correlated.
tt_ = pd.DataFrame(trial_rows)
for col in ["speed", "tp"]:
    tt_[f"z_{col}"] = tt_.groupby(["dataset", "prior", "ts"])[col].transform(lambda x: (x - x.mean()) / x.std())
zr = tt_.groupby(["dataset", "prior"], sort=False).apply(lambda h: pd.Series(stats.pearsonr(h.z_speed, h.z_tp), index=["r", "p"]))
print("Trial by trial, z-scored within ts:")
print(zr.round(3).to_string())
for prior, g in tt_.groupby("prior", sort=False):
    r, p = stats.pearsonr(g.z_speed, g.z_tp)
    print(f"{prior}, pooled over datasets: r = {r:.3f} (p = {p:.2g}, {len(g)} trials); negative in "
          f"{int((zr.loc[(slice(None), prior), 'r'] < 0).sum())} of 4, p < 0.05 and negative in "
          f"{int(((zr.loc[(slice(None), prior), 'r'] < 0) & (zr.loc[(slice(None), prior), 'p'] < 0.05)).sum())}")

# Null for the trial-by-trial measure: units fire as Poisson processes at a constant rate (each
# unit's mean from Set to Go in the dataset and prior), each trial keeps its own duration tp, and
# speed is measured as above. Shorter trials average fewer bins, which alone could tie speed to tp.
rng_tt = np.random.default_rng(11)
null_r = {"Short": [], "Long": []}
for _ in range(20):
    sims = []
    for e, th, name in DS:
        g = good[(good.is_eye == e) & (good.theta == th)]
        for short in [True, False]:
            gg = g[g.is_short == short]
            nbs = (gg.tp.to_numpy() // rsg.BIN_MS).astype(int)
            c_ = rsg.counts(spikes, units, gg.set_time.to_numpy(), 0, 2400)
            lam = np.concatenate([c_[i, :n] for i, n in enumerate(nbs)]).mean(0)
            for ts, tp, n in zip(gg.ts, gg.tp, nbs):
                sim = rng_tt.poisson(lam, size=(n, len(units))).astype(float)
                sims.append({"dataset": name, "prior": "Short" if short else "Long", "ts": ts, "tp": tp,
                             "speed": rsg.step_speed(PROD[name]["pca"](rsg.smooth(sim * 1000 / rsg.BIN_MS)))})
    sims = pd.DataFrame(sims)
    for col in ["speed", "tp"]:
        sims[f"z_{col}"] = sims.groupby(["dataset", "prior", "ts"])[col].transform(lambda x: (x - x.mean()) / x.std())
    for prior, g in sims.groupby("prior"):
        null_r[prior].append(stats.pearsonr(g.z_speed, g.z_tp)[0])
for prior in ["Short", "Long"]:
    g = tt_[tt_.prior == prior]
    r = stats.pearsonr(g.z_speed, g.z_tp)[0]
    v = np.array(null_r[prior])
    print(f"Poisson null with matched durations, {prior}: null r mean {v.mean():.3f}, range {v.min():.3f} to {v.max():.3f} "
          f"(20 simulations); observed {r:.3f}, below every null value: {bool(r < v.min())}")

Short: rho of residuals on ts (speed, mean tp) = 0.084 (nominal p = 0.72)
Long: rho of residuals on ts (speed, mean tp) = -0.283 (nominal p = 0.23)
Trial by trial, z-scored within ts:
                      r      p
dataset    prior              
Eye Right  Short -0.437  0.000
           Long  -0.251  0.002
Eye Left   Short -0.377  0.000
           Long  -0.232  0.004
Hand Right Short -0.261  0.005
           Long  -0.173  0.068
Hand Left  Short -0.203  0.018
           Long  -0.096  0.372
Short, pooled over datasets: r = -0.317 (p = 1.1e-12, 481 trials); negative in 4 of 4, p < 0.05 and negative in 4
Long, pooled over datasets: r = -0.201 (p = 5.5e-06, 502 trials); negative in 4 of 4, p < 0.05 and negative in 2

Poisson null with matched durations, Short: null r mean -0.071, range -0.144 to -0.010 (20 simulations); observed -0.317, below every null value: True
Poisson null with matched durations, Long: null r mean -0.033, range -0.103 to 0.028 (20 simulations); observed -0.201, below every null value: True

fig, axes = plt.subplots(1, 2, figsize=(9, 3.4))
for prior, col in [("Short", "tab:red"), ("Long", "tab:blue")]:
    g = spd[spd.prior == prior]
    for n, h in g.groupby("dataset"):
        axes[0].plot(h.xu_at_set, h.speed, "o-", color=col, ms=3, lw=.7)
        axes[1].plot(h.speed, h.mean_tp, "o-", color=col, ms=3, lw=.7)
axes[0].set_xlabel("Xu at Set (spikes/s)")
axes[0].set_ylabel("speed (spikes/s per 20 ms)")
axes[1].set_xlabel("speed (spikes/s per 20 ms)")
axes[1].set_ylabel("mean tp (ms)")
plt.tight_layout()
plt.show()
print("Xu %.1f to %.1f; speed %.2f to %.2f; mean tp %.0f to %.0f; all within the axes."
      % (spd.xu_at_set.min(), spd.xu_at_set.max(), spd.speed.min(), spd.speed.max(), spd.mean_tp.min(), spd.mean_tp.max()))

Xu -45.0 to 24.2; speed 3.15 to 9.13; mean tp 497 to 1177; all within the axes.

rng_boot = np.random.default_rng(5)


def support_slope(rate, ends, short):
    lo, hi = min(rsg.prior_ts(short)), max(rsg.prior_ts(short))
    sp = np.linalg.norm(np.diff(rate, axis=0), axis=1)
    t = ends[1:]
    sel = (t > lo) & (t <= hi + 1e-6)
    return np.polyfit(t[sel], sp[sel], 1)[0] * 100, sp[sel].mean()   # per 100 ms


rows = []
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th)]
    for short in [True, False]:
        gg = g[g.is_short == short]
        t_max = max(rsg.prior_ts(short))
        c_ = rsg.counts(spikes, units, gg.ready_time.to_numpy(), 0, t_max)
        el = ((gg.set_time - gg.ready_time) * 1000).to_numpy()
        rate, ends, _ = rsg.attrition_rates(c_, el)
        slope, mean_sp = support_slope(rate, ends, short)
        boots = []
        for _ in range(200):
            i = rng_boot.integers(0, len(gg), len(gg))
            r_, e_, _ = rsg.attrition_rates(c_[i], el[i])
            boots.append(support_slope(r_, e_, short)[0])
        rows.append({"dataset": name, "prior": "Short" if short else "Long", "mean_speed": mean_sp, "slope_per_100ms": slope,
                     "slope_rel_pct": 100 * slope / mean_sp, "boot_low": np.percentile(boots, 2.5), "boot_high": np.percentile(boots, 97.5)})
asp = pd.DataFrame(rows)
print("Speed in the full space of", len(units), "units (spikes/s per 20 ms step); slope per 100 ms, with 95% bootstrap interval over trials:")
print(asp.round(3).to_string(index=False))
for prior, g in asp.groupby("prior", sort=False):
    print(f"{prior}: signed-rank p = {stats.wilcoxon(g.slope_per_100ms).pvalue:.3f} (smallest possible with 4: 0.125); "
          f"slopes {', '.join(f'{x:+.3f}' for x in g.slope_per_100ms)}; intervals excluding 0: "
          f"{int(((g.boot_low > 0) | (g.boot_high < 0)).sum())} of 4")

Speed in the full space of 48 units (spikes/s per 20 ms step); slope per 100 ms, with 95% bootstrap interval over trials:
   dataset prior  mean_speed  slope_per_100ms  slope_rel_pct  boot_low  boot_high
 Eye Right Short       5.397            1.008         18.683     0.724      2.216
 Eye Right  Long       3.676            0.579         15.743     0.583      1.430
  Eye Left Short       6.312            1.270         20.126     1.024      3.149
  Eye Left  Long       3.590            0.478         13.303     0.389      1.256
Hand Right Short       3.748            0.890         23.753     0.835      2.632
Hand Right  Long       3.067            0.578         18.845     0.551      1.547
 Hand Left Short       4.284            0.826         19.276     0.787      2.132
 Hand Left  Long       2.861            0.668         23.356     0.682      1.855
Short: signed-rank p = 0.125 (smallest possible with 4: 0.125); slopes +1.008, +1.270, +0.890, +0.826; intervals excluding 0: 4 of 4
Long: signed-rank p = 0.125 (smallest possible with 4: 0.125); slopes +0.579, +0.478, +0.578, +0.668; intervals excluding 0: 4 of 4

# Attrition adds noise: later bins average fewer trials, and noisier rates give larger steps.
# Variant 1: only the trials with the longest ts of the prior, which last the whole support, so
# the number of trials is the same in every bin. Null: units firing as Poisson processes at a
# constant rate (each unit's mean over the support before Set), with the same trials and
# Set times, measured both ways. The point estimate of the main measure is also compared with
# the median of its bootstrap.
rng_null = np.random.default_rng(10)
rows = []
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th)]
    for short in [True, False]:
        gg = g[g.is_short == short]
        lo, hi = min(rsg.prior_ts(short)), max(rsg.prior_ts(short))
        c_ = rsg.counts(spikes, units, gg.ready_time.to_numpy(), 0, hi)
        el = ((gg.set_time - gg.ready_time) * 1000).to_numpy()
        ends = np.arange(1, c_.shape[1] + 1) * rsg.BIN_MS
        full = gg.ts.to_numpy() == hi
        fixed = support_slope(rsg.smooth(c_[full].mean(0) * 1000 / rsg.BIN_MS), ends, short)[0]
        use = (el[:, None] >= ends[None, :] - 1e-6) & (ends[None, :] > lo)
        lam = c_[use].mean(0)                          # counts per bin and unit, before Set, in the support
        null_main, null_fixed, null_speed = [], [], []
        for _ in range(100):
            sim = rng_null.poisson(lam, size=c_.shape).astype(float)
            sl_, sp_ = support_slope(rsg.attrition_rates(sim, el)[0], ends, short)
            null_main.append(sl_)
            null_speed.append(sp_)
            null_fixed.append(support_slope(rsg.smooth(sim[full].mean(0) * 1000 / rsg.BIN_MS), ends, short)[0])
        main = asp[(asp.dataset == name) & (asp.prior == ("Short" if short else "Long"))].iloc[0]
        rows.append({"dataset": name, "prior": "Short" if short else "Long", "speed_obs": main.mean_speed,
                     "speed_null_mean": np.mean(null_speed), "slope_main": main.slope_per_100ms,
                     "null_main_mean": np.mean(null_main), "null_main_95": np.percentile(null_main, 95),
                     "trials_full": int(full.sum()), "slope_fixed": fixed, "null_fixed_mean": np.mean(null_fixed),
                     "null_fixed_2.5": np.percentile(null_fixed, 2.5), "null_fixed_97.5": np.percentile(null_fixed, 97.5)})
spn = pd.DataFrame(rows)
spn["excess_main"] = spn.slope_main - spn.null_main_mean
print("Slopes of speed per 100 ms of the support; Poisson null of constant rates, 100 simulations:")
print(spn.round(3).to_string(index=False))
for prior, g in spn.groupby("prior", sort=False):
    print(f"{prior}: main slope above the null's 95th percentile in {(g.slope_main > g.null_main_95).sum()} of 4; "
          f"fixed trials, slope outside the null's 95% range in {((g.slope_fixed < g['null_fixed_2.5']) | (g.slope_fixed > g['null_fixed_97.5'])).sum()} of 4, "
          f"positive in {(g.slope_fixed > 0).sum()} of 4; signed-rank p (fixed trials) = {stats.wilcoxon(g.slope_fixed).pvalue:.3f}")
print(f"\nMean speed over the support: null above observed in {(spn.speed_null_mean > spn.speed_obs).sum()} of 8 "
      f"(null / observed {(spn.speed_null_mean / spn.speed_obs).min():.2f} to {(spn.speed_null_mean / spn.speed_obs).max():.2f})")
print("\nBootstrap of the main measure: point estimate below the bootstrap 2.5th percentile in",
      int((asp.slope_per_100ms < asp.boot_low).sum()), "of 8, so the bootstrap is biased upward and its intervals are not used.")

Slopes of speed per 100 ms of the support; Poisson null of constant rates, 100 simulations:
   dataset prior  speed_obs  speed_null_mean  slope_main  null_main_mean  null_main_95  trials_full  slope_fixed  null_fixed_mean  null_fixed_2.5  null_fixed_97.5  excess_main
 Eye Right Short      5.397            3.689       1.008           1.683         2.262           25        0.497            0.694           0.067            1.476       -0.675
 Eye Right  Long      3.676            3.027       0.579           0.999         1.275           27        0.256            0.389           0.020            0.853       -0.420
  Eye Left Short      6.312            4.085       1.270           1.813         2.278           22        0.529            0.801           0.012            1.492       -0.543
  Eye Left  Long      3.590            2.783       0.478           0.905         1.192           31        0.097            0.359           0.017            0.785       -0.428
Hand Right Short      3.748            3.347       0.890           1.430         1.872           22        0.357            0.555           0.067            1.074       -0.540
Hand Right  Long      3.067            2.853       0.578           0.949         1.318           25       -0.106            0.345          -0.040            0.648       -0.371
 Hand Left Short      4.284            3.162       0.826           1.461         1.940           24        0.288            0.591          -0.021            1.149       -0.636
 Hand Left  Long      2.861            3.369       0.668           1.130         1.531           18        0.137            0.418          -0.103            0.835       -0.462
Short: main slope above the null's 95th percentile in 0 of 4; fixed trials, slope outside the null's 95% range in 0 of 4, positive in 4 of 4; signed-rank p (fixed trials) = 0.125
Long: main slope above the null's 95th percentile in 0 of 4; fixed trials, slope outside the null's 95% range in 1 of 4, positive in 3 of 4; signed-rank p (fixed trials) = 0.375

Mean speed over the support: null above observed in 1 of 8 (null / observed 0.65 to 1.18)

Bootstrap of the main measure: point estimate below the bootstrap 2.5th percentile in 2 of 8, so the bootstrap is biased upward and its intervals are not used.

def transient_slopes(kept=rsg.VAR_KEPT, norm=None):
    rows = []
    for e, th, name in DS:
        g = good[(good.is_eye == e) & (good.theta == th)]
        traj = {}
        for short in [True, False]:
            for ts in rsg.prior_ts(short):
                gg = g[(g.is_short == short) & (g.ts == ts)]
                traj[(short, ts)] = rsg.aligned_rates(spikes, units, gg.set_time.to_numpy(), -100, 300)[5:15]  # Set to Set + 200 ms
        if norm:
            sc = rsg.unit_scale(list(traj.values()), norm)
            traj = {k: v / sc for k, v in traj.items()}
        pca_t = rsg.PCA(list(traj.values()), kept) if kept else (lambda x: x)
        t = np.arange(1, 11) * rsg.BIN_MS
        for short in [True, False]:
            ts_v = rsg.prior_ts(short)
            ref = pca_t(traj[(short, ts_v[2])])
            for which, ts in [("shortest", ts_v[0]), ("longest", ts_v[-1])]:
                dist = np.linalg.norm(pca_t(traj[(short, ts)]) - ref, axis=1)
                rows.append({"dataset": name, "prior": "Short" if short else "Long", "ts": which,
                             "PCs": getattr(pca_t, "n", len(units)), "mean_distance": dist.mean(),
                             "slope_per_100ms": np.polyfit(t, dist, 1)[0] * 100})
    out = pd.DataFrame(rows)
    out["slope_rel_pct"] = 100 * out.slope_per_100ms / out.mean_distance
    return out


tr_ = transient_slopes()
print(tr_.round(3).to_string(index=False))
for which, g in tr_.groupby("ts"):
    w = stats.wilcoxon(g.slope_per_100ms)
    print(f"{which} ts: signed-rank p = {w.pvalue:.3f}; positive slopes {(g.slope_per_100ms > 0).sum()} of 8; "
          f"median relative slope {g.slope_rel_pct.median():+.1f}% per 100 ms")
print("\nVariants (signed-rank p for the shortest and the longest ts; PCs kept per dataset):")
for kept, norm, lab in [(0.60, None, "60% of variance"), (0.90, None, "90% of variance"), (None, None, "full space"),
                        (0.75, "z", "z-scored rates, 75%"), (0.75, "soft", "soft-normalised rates, 75%")]:
    v = transient_slopes(kept, norm)
    ps = {w: stats.wilcoxon(g.slope_per_100ms).pvalue for w, g in v.groupby("ts")}
    pos = {w: int((g.slope_per_100ms > 0).sum()) for w, g in v.groupby("ts")}
    print(f"   {lab}: shortest p = {ps['shortest']:.3f} ({pos['shortest']} of 8 positive), longest p = {ps['longest']:.3f} "
          f"({pos['longest']} of 8 positive); PCs {v.groupby('dataset', sort=False).PCs.first().tolist()}")

   dataset prior       ts  PCs  mean_distance  slope_per_100ms  slope_rel_pct
 Eye Right Short shortest    2         33.900           -5.552        -16.379
 Eye Right Short  longest    2         20.800           -6.914        -33.242
 Eye Right  Long shortest    2         24.513            5.228         21.329
 Eye Right  Long  longest    2         22.422            0.347          1.547
  Eye Left Short shortest    1         36.737           -3.630         -9.881
  Eye Left Short  longest    1         30.574           -8.238        -26.944
  Eye Left  Long shortest    1         35.824            3.716         10.374
  Eye Left  Long  longest    1         21.763           -3.087        -14.186
Hand Right Short shortest    3         18.999            6.966         36.668
Hand Right Short  longest    3         17.968           -2.200        -12.244
Hand Right  Long shortest    3         20.095            7.824         38.938
Hand Right  Long  longest    3          7.772           -2.207        -28.402
 Hand Left Short shortest    4         20.994            6.727         32.045
 Hand Left Short  longest    4         15.348           -2.504        -16.318
 Hand Left  Long shortest    4         14.481            8.189         56.549
 Hand Left  Long  longest    4         15.418            0.311          2.019
longest ts: signed-rank p = 0.039; positive slopes 2 of 8; median relative slope -15.3% per 100 ms
shortest ts: signed-rank p = 0.078; positive slopes 6 of 8; median relative slope +26.7% per 100 ms

Variants (signed-rank p for the shortest and the longest ts; PCs kept per dataset):
   60% of variance: shortest p = 0.078 (6 of 8 positive), longest p = 0.016 (1 of 8 positive); PCs [1, 1, 2, 2]

   90% of variance: shortest p = 0.109 (5 of 8 positive), longest p = 0.039 (2 of 8 positive); PCs [5, 4, 6, 8]

   full space: shortest p = 0.078 (6 of 8 positive), longest p = 0.023 (1 of 8 positive); PCs [48, 48, 48, 48]
   z-scored rates, 75%: shortest p = 0.742 (5 of 8 positive), longest p = 0.039 (1 of 8 positive); PCs [6, 5, 7, 8]

   soft-normalised rates, 75%: shortest p = 0.312 (5 of 8 positive), longest p = 0.055 (2 of 8 positive); PCs [4, 4, 5, 6]

def go_curves(kept, min_n=1):
    rows, curves = [], {}
    for e, th, name in DS:
        g = good[(good.is_eye == e) & (good.theta == th)]
        traj = {}
        for short in [True, False]:
            for ts in rsg.prior_ts(short):
                gg = g[(g.is_short == short) & (g.ts == ts)]
                traj[(short, ts)] = rsg.aligned_rates(spikes, units, gg.go_time.to_numpy(), -900, 100)[5:45]  # Go - 800 to Go
        pca_g = rsg.PCA(list(traj.values()), kept, min_n) if kept else (lambda x: x)
        t = np.arange(-39, 1) * rsg.BIN_MS   # bin end times relative to Go
        for short in [True, False]:
            ts_v = rsg.prior_ts(short)
            ref = pca_g(traj[(short, ts_v[2])])
            for ts in ts_v[:2] + ts_v[3:]:
                diff = pca_g(traj[(short, ts)]) - ref
                dist = np.linalg.norm(diff, axis=1)
                curves[(name, short, ts)] = dist
                k = np.argmin(dist)
                # With one component, a minimum is a crossing when the signed difference changes sign there.
                crossing = diff.shape[1] == 1 and np.any(np.sign(diff[max(k - 1, 0):k + 2, 0]) != np.sign(diff[k, 0]))
                rows.append({"dataset": name, "prior": "Short" if short else "Long", "ts": ts,
                             "PCs": getattr(pca_g, "n", len(units)), "t_min": t[np.argmin(dist)],
                             "d_at_go_over_min": dist[-1] / dist.min(), "min_at_crossing": bool(crossing), "d_start": dist[0], "d_go": dist[-1]})
    return pd.DataFrame(rows), curves, t


thr_tab, thr_curves, t_go = go_curves(0.60)
print(thr_tab.round(3).to_string(index=False))
print(f"\nMain (60% of variance): median time of the minimum {thr_tab.t_min.median():.0f} ms (bin end, relative to Go); "
      f"minimum in the last 40 ms in {(thr_tab.t_min >= -20).sum()} of {len(thr_tab)} curves; in the last 200 ms in "
      f"{(thr_tab.t_min >= -180).sum()}; distance at Go below that 800 ms earlier in {(thr_tab.d_go < thr_tab.d_start).sum()}")
for kept, lab in [(0.75, "75% of variance"), (None, "full space")]:
    v, _, _ = go_curves(kept)
    print(f"Variant, {lab}: median time of the minimum {v.t_min.median():.0f} ms; in the last 40 ms {(v.t_min >= -20).sum()} of {len(v)}")

   dataset prior   ts  PCs  t_min  d_at_go_over_min  min_at_crossing  d_start   d_go
 Eye Right Short  480    1   -600             2.265             True   30.771  1.455
 Eye Right Short  560    1    -20             1.320             True    4.027  0.123
 Eye Right Short  720    1   -240           177.672             True    9.028  3.734
 Eye Right Short  800    1      0             1.000             True   19.518  0.240
 Eye Right  Long  800    1   -180             9.758             True    1.435  3.711
 Eye Right  Long  900    1   -120            18.362             True    5.267  2.153
 Eye Right  Long 1100    1    -80             1.749            False   34.278  3.511
 Eye Right  Long 1200    1   -280            13.242             True   38.929  3.458
  Eye Left Short  480    1   -540             1.795             True   13.723  1.485
  Eye Left Short  560    1    -80            29.558             True    7.761  2.051
  Eye Left Short  720    1      0             1.000            False    9.958  1.583
  Eye Left Short  800    1      0             1.000            False   20.924  0.457
  Eye Left  Long  800    1   -760             3.855             True    6.222  3.160
  Eye Left  Long  900    1      0             1.000            False    1.695  1.666
  Eye Left  Long 1100    1   -220             3.244             True   18.422  0.732
  Eye Left  Long 1200    1   -200            63.014             True   28.217  0.237
Hand Right Short  480    3   -260             2.342            False   32.141 15.212
Hand Right Short  560    3   -260             6.502            False   11.257  7.488
Hand Right Short  720    3   -460             2.783            False    5.860  9.533
Hand Right Short  800    3   -480             1.924            False   10.949 11.672
Hand Right  Long  800    3   -200             3.617            False   17.831  2.996
Hand Right  Long  900    3   -180             5.847            False    5.973  5.835
Hand Right  Long 1100    3    -40             3.873            False    6.387  4.641
Hand Right  Long 1200    3   -620             3.473            False    5.762  7.041
 Hand Left Short  480    2      0             1.000            False   10.654  0.536
 Hand Left Short  560    2   -620             4.091            False    9.414  2.317
 Hand Left Short  720    2   -700           126.054            False    2.526  4.982
 Hand Left Short  800    2   -740             4.630            False    1.405  5.769
 Hand Left  Long  800    2   -140             1.350            False   15.874  3.555
 Hand Left  Long  900    2   -100             1.266            False    4.264  0.861
 Hand Left  Long 1100    2   -740             2.987            False    6.204  5.125
 Hand Left  Long 1200    2   -280             5.283            False   11.355  6.166

Main (60% of variance): median time of the minimum -210 ms (bin end, relative to Go); minimum in the last 40 ms in 6 of 32 curves; in the last 200 ms in 14; distance at Go below that 800 ms earlier in 26

Variant, 75% of variance: median time of the minimum -80 ms; in the last 40 ms 11 of 32

Variant, full space: median time of the minimum -40 ms; in the last 40 ms 14 of 32

for kept, min_n, lab in [(0.60, 2, "60% of variance, at least 2 components"), (0.60, 3, "60% of variance, at least 3 components")]:
    v, _, _ = go_curves(kept, min_n)
    print(f"Variant, {lab}: median time of the minimum {v.t_min.median():.0f} ms; in the last 40 ms {(v.t_min >= -20).sum()} of {len(v)}; "
          f"in the last 200 ms {(v.t_min >= -180).sum()}")
one = thr_tab[thr_tab.PCs == 1]
print(f"Main, datasets with 1 component: {len(one)} curves, median time of the minimum {one.t_min.median():.0f} ms, "
      f"distance at Go over minimum above 10 in {(one.d_at_go_over_min > 10).sum()}; minimum at Go in {(one.t_min == 0).sum()}; "
      f"minimum where the signed difference changes sign (a crossing) in {one.min_at_crossing.sum()}")

Variant, 60% of variance, at least 2 components: median time of the minimum -230 ms; in the last 40 ms 8 of 32; in the last 200 ms 14

Variant, 60% of variance, at least 3 components: median time of the minimum -70 ms; in the last 40 ms 13 of 32; in the last 200 ms 23
Main, datasets with 1 component: 16 curves, median time of the minimum -150 ms, distance at Go over minimum above 10 in 5; minimum at Go in 4; minimum where the signed difference changes sign (a crossing) in 12

fig, axes = plt.subplots(1, 2, figsize=(9, 3.2), sharey=True)
for (name, short, ts), d in thr_curves.items():
    axes[0 if short else 1].plot(t_go, d, color="tab:red" if short else "tab:blue", lw=.6, alpha=.6)
for ax, lab in zip(axes, ["Short", "Long"]):
    ax.set_title(lab)
    ax.set_xlabel("time from Go (ms)")
axes[0].set_ylabel("distance to middle ts (spikes/s)")
plt.tight_layout()
plt.show()
allv = np.concatenate(list(thr_curves.values()))
print("Distances %.1f to %.1f; all within the axes." % (allv.min(), allv.max()))

Distances 0.0 to 58.4; all within the axes.

rng_gpfa = np.random.default_rng(6)
t_gpfa = time.time()


def gpfa_rows(name, g, lat):
    rows = []
    for short in [True, False]:
        mask = g.is_short.to_numpy() == short
        L = [lat[i] for i in np.flatnonzero(mask)]
        nb = max(x.shape[1] for x in L)
        mean = np.array([np.mean([x[:, k] for x in L if x.shape[1] > k], axis=0) for k in range(nb)])
        tt = np.arange(1, nb + 1) * rsg.BIN_MS
        ts_v = rsg.prior_ts(short)
        u = rsg.encoding_axis(mean, tt, min(ts_v), max(ts_v))
        xu = np.array([x[:, -1] @ u for x in L])
        ts_tr = g.ts.to_numpy()[mask]
        z = (xu - xu.mean()) / xu.std()
        var = np.array([z[ts_tr == t].var(ddof=1) for t in ts_v])
        mxu = np.array([xu[ts_tr == t].mean() for t in ts_v])
        rows.append({"dataset": name, "prior": "Short" if short else "Long", **{f"var_{t}": v for t, v in zip(ts_v, var)},
                     "quad_coef": np.polyfit(np.array(ts_v, float), var, 2)[0],
                     "mean_xu": mxu, "slope_xu": np.polyfit(np.array(ts_v, float), mxu, 1)[0]})
    return rows


GP, rows = {}, []
for e, th, name in DS:
    g = good[(good.is_eye == e) & (good.theta == th)]
    sts = rsg.estimation_spiketrains(spikes, units, g)
    model, dim, scores, folds = rsg.fit_gpfa(sts, rng_gpfa)
    lat = rsg.gpfa_transform(model, sts)
    GP[name] = {"lat": lat, "trials": g, "dim": dim, "sts": sts, "folds": folds, "scores": scores, "model": model}
    print(f"{name}: {len(g)} trials, latent dimensionality {dim} (cross-validated log-likelihood "
          + ", ".join(f"{k}: {v:.0f}" for k, v in scores.items()) + ")")
    rows += gpfa_rows(name, g, lat)
gp = pd.DataFrame(rows)
print(f"GPFA fits took {time.time() - t_gpfa:.0f} s")
print(gp.drop(columns=["mean_xu"]).to_string(index=False, float_format=lambda x: f"{x:.3g}"))
w = stats.wilcoxon(gp.quad_coef, alternative="less")
print(f"\nQuadratic coefficient negative in {(gp.quad_coef < 0).sum()} of 8; one-tailed signed-rank p = {w.pvalue:.4f}")

Eye Right: 276 trials, latent dimensionality 8 (cross-validated log-likelihood 4: -174266, 8: -173776)

Eye Left: 260 trials, latent dimensionality 8 (cross-validated log-likelihood 4: -164318, 8: -163780)

Hand Right: 225 trials, latent dimensionality 8 (cross-validated log-likelihood 4: -124348, 8: -123808)

Hand Left: 222 trials, latent dimensionality 8 (cross-validated log-likelihood 4: -115304, 8: -114842)
GPFA fits took 268 s
   dataset prior  var_480  var_560  var_640  var_720  var_800  quad_coef  slope_xu  var_900  var_1000  var_1100  var_1200
 Eye Right Short    0.764    0.525    0.645    0.799    0.262  -6.28e-06   0.00224      NaN       NaN       NaN       NaN
 Eye Right  Long      NaN      NaN      NaN      NaN    0.988  -6.21e-07    0.0014    0.732     0.934     0.353     0.445
  Eye Left Short    0.484    0.228    0.308    0.381    0.189   1.37e-06    0.0027      NaN       NaN       NaN       NaN
  Eye Left  Long      NaN      NaN      NaN      NaN    0.558  -6.67e-06   0.00155    0.917     0.735     0.477     0.407
Hand Right Short    0.842    0.954     1.04    0.882    0.611  -1.13e-05  0.000555      NaN       NaN       NaN       NaN
Hand Right  Long      NaN      NaN      NaN      NaN    0.635  -2.23e-05  0.000375     1.65      1.38     0.518     0.272
 Hand Left Short    0.909     1.22    0.696    0.488    0.694   1.17e-06  0.000939      NaN       NaN       NaN       NaN
 Hand Left  Long      NaN      NaN      NaN      NaN     1.27   5.99e-06  0.000453     1.44     0.468     0.702     0.689

Quadratic coefficient negative in 5 of 8; one-tailed signed-rank p = 0.1250

var_cols = {True: [f"var_{t}" for t in rsg.SHORT_TS], False: [f"var_{t}" for t in rsg.LONG_TS]}
fig, axes = plt.subplots(1, 2, figsize=(8, 3.2), sharey=True)
allv = []
for ax, short in zip(axes, [True, False]):
    for _, r in gp[gp.prior == ("Short" if short else "Long")].iterrows():
        v = r[var_cols[short]].to_numpy(float)
        allv += list(v)
        ax.plot(rsg.prior_ts(short), v, "o-", color="tab:red" if short else "tab:blue", lw=.8, ms=3)
    ax.set_xlabel("ts (ms)")
axes[0].set_ylabel("variance of z-scored Xu")
plt.tight_layout()
plt.show()
print("Variances %.2f to %.2f; all within the axes." % (min(allv), max(allv)))

Variances 0.19 to 1.65; all within the axes.

# Convergence: the training log-likelihood of each final model over its EM iterations
# (computed every fifth iteration), and the relative gain over the last 25 iterations.
print("Training log-likelihood of the final models:")
for name, G in GP.items():
    ll = np.array(G["model"].fit_info["log_likelihoods"], float)
    it = np.flatnonzero(~np.isnan(ll)) + 1
    ll = ll[~np.isnan(ll)]
    late = ll[it >= it[-1] - 25][0]
    print(f"   {name}: {len(G['model'].fit_info['log_likelihoods'])} iterations; at iteration 50 {ll[it == 50][0]:.0f}, "
          f"75 {ll[it == 75][0]:.0f}, {it[-1]} {ll[-1]:.0f}; gain over the last 25 iterations {100 * (ll[-1] - late) / abs(ll[-1]):.3f}%")

# Variant: a wider grid; 12 latents scored on the same folds. Where 12 wins, the variance test
# is repeated with the 12-latent model.
t_var = time.time()
rows_12, wins = [], []
for name, G in GP.items():
    s12 = rsg.gpfa_cv_score(G["sts"], G["folds"], 12)
    best = max({**G["scores"], 12: s12}.items(), key=lambda kv: kv[1])[0]
    print(f"{name}: cross-validated log-likelihood with 12 latents {s12:.0f} (chosen {G['dim']}: {G['scores'][G['dim']]:.0f}); best of 4, 8, 12: {best}")
    if best == 12:
        wins.append(name)
        lat12 = rsg.gpfa_transform(rsg.gpfa_model(G["sts"], 12), G["sts"])
        rows_12 += gpfa_rows(name, G["trials"], lat12)
    else:
        rows_12 += [r.to_dict() for _, r in gp[gp.dataset == name].iterrows()]
gp12 = pd.DataFrame(rows_12)
w12 = stats.wilcoxon(gp12.quad_coef, alternative="less")
print(f"With the grid 4, 8, 12 ({len(wins)} datasets refit with 12): quadratic coefficient negative in "
      f"{(gp12.quad_coef < 0).sum()} of 8; one-tailed signed-rank p = {w12.pvalue:.4f}")
inc12 = np.array([np.subtract(*rsg.increments(m)[::-1]) for m in gp12.mean_xu])
sg12 = gp12.pivot(index="dataset", columns="prior", values="slope_xu")
print(f"Mean projections with this grid: middle increment above extreme in {(inc12 > 0).sum()} of 8 "
      f"(two-sided p = {stats.wilcoxon(inc12).pvalue:.4f}); Short slope above Long in {int((sg12.Short > sg12.Long).sum())} of 4")
print(f"Variant took {time.time() - t_var:.0f} s")

Training log-likelihood of the final models:
   Eye Right: 100 iterations; at iteration 50 -207406, 75 -207154, 100 -206982; gain over the last 25 iterations 0.083%
   Eye Left: 100 iterations; at iteration 50 -193544, 75 -193328, 100 -193179; gain over the last 25 iterations 0.077%
   Hand Right: 100 iterations; at iteration 50 -147714, 75 -147556, 100 -147441; gain over the last 25 iterations 0.078%
   Hand Left: 100 iterations; at iteration 50 -140688, 75 -140567, 100 -140474; gain over the last 25 iterations 0.066%

Eye Right: cross-validated log-likelihood with 12 latents -173646 (chosen 8: -173776); best of 4, 8, 12: 12

Eye Left: cross-validated log-likelihood with 12 latents -163590 (chosen 8: -163780); best of 4, 8, 12: 12

Hand Right: cross-validated log-likelihood with 12 latents -123685 (chosen 8: -123808); best of 4, 8, 12: 12

Hand Left: cross-validated log-likelihood with 12 latents -114724 (chosen 8: -114842); best of 4, 8, 12: 12

With the grid 4, 8, 12 (4 datasets refit with 12): quadratic coefficient negative in 5 of 8; one-tailed signed-rank p = 0.1250
Mean projections with this grid: middle increment above extreme in 8 of 8 (two-sided p = 0.0078); Short slope above Long in 4 of 4
Variant took 234 s

inc_g = pd.DataFrame([{"dataset": r.dataset, "prior": r.prior, **dict(zip(["extreme", "middle"], rsg.increments(r.mean_xu)))}
                      for _, r in gp.iterrows()])
print(inc_g.round(4).to_string(index=False))
w = stats.wilcoxon(inc_g.middle, inc_g.extreme)
print(f"\nMiddle increment above extreme in {(inc_g.middle > inc_g.extreme).sum()} of 8; two-sided signed-rank p = {w.pvalue:.4f}")
sg = gp.pivot(index="dataset", columns="prior", values="slope_xu")
print("\nSlope of mean Xu on ts (latent units per ms):")
print(sg.round(5).to_string())
print(f"Short slope above Long in {int((sg.Short > sg.Long).sum())} of 4; one-tailed signed-rank p = "
      f"{stats.wilcoxon(sg.Short, sg.Long, alternative='greater').pvalue:.4f}")
print("Mean Xu rises with ts (positive slope) in", int((gp.slope_xu > 0).sum()), "of 8")
print(f"\nTotal run time {time.time() - started:.0f} s")

   dataset prior  extreme  middle
 Eye Right Short   0.0693  0.2527
 Eye Right  Long  -0.0029  0.2355
  Eye Left Short   0.1249  0.2761
  Eye Left  Long   0.0583  0.2201
Hand Right Short   0.0247  0.0575
Hand Right  Long  -0.0148  0.0724
 Hand Left Short   0.0280  0.1065
 Hand Left  Long   0.0363  0.0512

Middle increment above extreme in 8 of 8; two-sided signed-rank p = 0.0078

Slope of mean Xu on ts (latent units per ms):
prior          Long    Short
dataset                     
Eye Left    0.00155  0.00270
Eye Right   0.00140  0.00224
Hand Left   0.00045  0.00094
Hand Right  0.00038  0.00056
Short slope above Long in 4 of 4; one-tailed signed-rank p = 0.0625
Mean Xu rises with ts (positive slope) in 8 of 8

Total run time 539 s

Claim	Source	Tested here
beh-slopes: slopes of tp on ts between 0 and 1	Fig. 1E	Tested
beh-overlap-prior: opposite biases at ts = 800 ms	Fig. 1E	Tested
beh-sigmoid-increments: smaller tp increments at the extremes	Fig. 2F	Tested
beh-long-more-bias: smaller slopes for the Long prior	Fig. 2H	Tested
beh-bayes-model: the Bayesian observer captures bias and variance	Fig. 2C, D	Tested
glm-prior-ts: 65% prior-dependent and 27% ts-dependent units (H)	Fig. 3E	Tested
dfr-overlap-support: prior difference at 800 ms peaks during the support	Fig. 3D	Tested
pca-dimensionality: 3-4 PCs (estimation), 5-10 (production)	Fig. S3	Tested
encoding-r2: projections match the Bayesian estimate (R2 0.993, 0.996)	Fig. 4C	Tested
encoding-bayes-beats-linear: Bayesian better than linear	Fig. 4D	Tested
short-in-long-control: no sigmoid in the matched Long period	Fig. 4D	Tested
neural-long-more-bias: smaller neural slope for Long	Fig. 4D	Tested
speed-vs-projection: larger projection, slower speed	Fig. 4E	Tested
initial-speed-biased: initial speed biased toward the mean	Fig. 4F	Tested
speed-predicts-tp: speed predicts tp	Fig. 4G	Tested
alt-speed-unmodulated: speed constant during the support	Fig. 5B	Tested
alt-transient-parallel: parallel trajectories after Set	Fig. 5D	Tested
alt-threshold-monotonic: distances fall steadily up to Go	Fig. 5F	Tested
gpfa-variance-inverted-u: inverted-U variance of single-trial projections	Fig. 6C	Tested
gpfa-mean-sigmoid: sigmoidal mean, higher Short slope	Fig. 6D	Tested
rnn-curved-manifold, rnn-perturbation: recurrent network models	Fig. 7	Blocked, data missing (no model in the dandiset; code on request)
monkey-g: the results hold in monkey G	throughout	Blocked (monkey G not in the dandiset)

Claim	Outcome
beh-slopes	Not reproduced (6 of 8 slopes between 0 and 1)
beh-overlap-prior	Not reproduced (both biases positive, p = 0.073)
beh-sigmoid-increments	Reproduced (7 of 8, p = 0.0156)
beh-long-more-bias	Reproduced (4 of 4)
beh-bayes-model	Partially reproduced (SD in 4 of 4, bias R2 >= 0.8 in 2 of 4)
glm-prior-ts	Not reproduced (75.0% and 45.8% against 65% and 27%; chance at most 14.6%)
dfr-overlap-support	Reproduced (largest at 480 to 640 ms)
pca-dimensionality	Not reproduced (2 PCs in the estimation epoch; 8 of 8 in range with z-scored rates, a variant)
encoding-r2	Reproduced (0.998 and 0.996, Hand Left); above 95% of random axes in only 3 of 8
encoding-bayes-beats-linear	Partially reproduced (Long 4 of 4, Short 3 of 4)
short-in-long-control	Reproduced by the paper's test, which cannot detect an effect with 4 datasets
neural-long-more-bias	Reproduced (4 of 4)
speed-vs-projection	Partially reproduced (rho -0.94 and -0.88, stronger than reported)
initial-speed-biased	Reproduced (7 of 8, p = 0.039)
speed-predicts-tp	Partially reproduced (rho -0.73 and -0.74, stronger than reported)
alt-speed-unmodulated	Reproduced by the paper's test, which cannot detect an effect; slopes no larger than noise gives
alt-transient-parallel	Not reproduced (longest ts, p = 0.039)
alt-threshold-monotonic	Not reproduced (median minimum 210 ms before Go with 60% of variance)
gpfa-variance-inverted-u	Not reproduced (5 of 8, p = 0.125)
gpfa-mean-sigmoid	Reproduced (p = 0.0078 and 0.0625, the paper's values)
rnn-curved-manifold, rnn-perturbation	Blocked: data missing (the trained networks and their code are not public)
monkey-g	Blocked: data missing

Reproducing Sohn et al. (2019) from dandiset 000130¶

Data check¶

Behaviour¶

beh-slopes¶

beh-overlap-prior¶

beh-sigmoid-increments¶

beh-long-more-bias¶

beh-bayes-model¶

Single units¶

glm-prior-ts¶

dfr-overlap-support¶

Population trajectories¶

pca-dimensionality¶

encoding-r2¶

encoding-bayes-beats-linear¶

short-in-long-control¶

neural-long-more-bias¶

speed-vs-projection¶

initial-speed-biased¶

speed-predicts-tp¶

Alternative mechanisms¶

alt-speed-unmodulated¶

alt-transient-parallel¶

alt-threshold-monotonic¶

Single-trial trajectories (GPFA)¶

gpfa-variance-inverted-u¶

gpfa-mean-sigmoid¶

Summary¶