wald = np.array(valid_sess) > 12
#%%
elm_clr = '#EC7063'
wld_clr = '#3498DB'
plt.errorbar(trn_err[elmo, :].mean(1),
tst_err[elmo, :].mean(1),
xerr=trn_err[elmo, :].std(1),
yerr=tst_err[elmo, :].std(1),
ls='none',
ecolor=elm_clr,
marker='.',
c=elm_clr,
markersize=10)
plt.errorbar(trn_err[wald, :].mean(1),
tst_err[wald, :].mean(1),
xerr=trn_err[wald, :].std(1),
yerr=tst_err[wald, :].std(1),
ls='none',
ecolor=wld_clr,
marker='.',
c=wld_clr,
markersize=10)
plt.legend(['Monkey E', 'Monkey W'])
# plt.scatter(trn_err.mean(1), tst_err.mean(1), c=np.arange(num_pop)>12)
dicplt.square_axis()
plt.plot([0.5, 0.5], plt.ylim(), color=(0.5, 0.5, 0.5), alpha=0.5)
plt.plot(plt.xlim(), [0.5, 0.5], color=(0.5, 0.5, 0.5), alpha=0.5)
plt.plot(plt.xlim(), plt.xlim(), 'k--')
fake_grads = neur_basis[:, nidx, :] @ (x_ @ (corr_dir[nidx, :] * fake_fz).T)
# corr_grad = fake_fz@((x_.T@x_)*corr_dir[nidx]**2).sum(0)
# anticorr_grad = fake_fz@((x_.T@x_)*(anticorr_dir[nidx]*corr_dir[nidx])).sum(0)
# corr_grad = (corr_dir[nidx,:]*fake_fz)@(x_.T@x_)@corr_dir[nidx]
# anticorr_grad = (corr_dir[nidx,:]*fake_fz)@(x_.T@x_)@anticorr_dir[nidx]
plt.quiver(fake_W[:, 0],
fake_W[:, 1],
fake_grads[0, :],
fake_grads[1, :],
color=(0.5, 0.5, 0.5))
# plt.quiver(fake_W[:,0],fake_W[:,1], anticorr_grad,corr_grad, color=(0.5,0.5,0.5))
dicplt.square_axis()
cols = cm.viridis(np.unique(grp) / 7)
for this_grp in np.unique(grp[these_neur]):
plt.scatter(neur_weights[0, grp == this_grp, 0],
neur_weights[0, grp == this_grp, 1],
marker='o',
c=cols[this_grp])
plt.plot(neur_weights[:, grp == this_grp, 0],
neur_weights[:, grp == this_grp, 1],
color=cols[this_grp])
#%%
# basins = list(itt.chain(*([combinations(range(num_cond), i) for i in range(1,3)])))
# basin_prototype = np.stack([x_[:,p].sum(1) for p in basins])
# which_basin = (weights[-1,...]@basin_prototype.T).argmax(1)