import os
from os.path import join as osjoin
from pathlib import Path
import warnings
from IPython.display import clear_output
from copy import deepcopy
import numpy as np
import matplotlib.pyplot as plt
import neurogym as ngym
from neurogym.utils import TruncExp, plot_env
from stable_baselines3.common.env_checker import check_env
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from joblib import Parallel, delayed
import hetero_utils as hu
# Suppress warnings
warnings.filterwarnings("ignore")
clear_output()
# DEMO_NAME = 'neuroscience'
# SAVE_PATH = osjoin(Path(__file__).resolve().parents[0], DEMO_NAME)
SAVE_PATH = osjoin(Path(__file__).resolve().parents[0], 'results')
Path(SAVE_PATH).mkdir(parents=True, exist_ok=True)
[docs]
def make_dataset(
seq_len, # number of timestamps to be generated
dim_ring = 2, # the number of choices the monkey has (laied on a ring)
stim_period = 100,
delay_period= [50, 30, 300],
decison_period= 10,
sigma=1, # stimulus noise (see below)
dt=1, # the simulation dt (in milisecond)
n_trials_to_plot=0, # set to positive if you want to plot environment's output
seed=1, # added to the ngym for reproducibility
verbose = False, # report the environment specification or not
):
# we use fixation to reset the dynamic. Thus, there is no reason
# for this to last more than one dt.
timing = { "fixation": 1, "decision": decison_period}
if isinstance(stim_period, (int, float)):
timing['stimulus'] = stim_period
else:
timing['stimulus'] = TruncExp(*stim_period) # mean, min, max
if isinstance(delay_period, (int, float)):
timing['delay'] = delay_period
else:
timing['delay'] = TruncExp(*delay_period) # mean, min, max
# sanity check
for key, val in timing.items():
if np.isscalar(val):
assert dt <= val, f'{key}={val} in timing is smaller tha dt={dt}.'
task = "ContextDecisionMaking-v0"
kwargs = {
"use_expl_context": True, # Use explicit context mode
"dt": dt, # simulation timestep
"dim_ring": dim_ring, # Number of choices (laid out on a ring, hence the name)
"timing": timing, # stimulus timing
"sigma": sigma, # Standard deviation of the Gaussian noise in the ring representation
"abort": False, # Whether to allow aborting the trial if the agent does not fixate
}
env = ngym.make(task, **kwargs)
check_env(env)
dataset = ngym.Dataset(task, batch_size=1,
env_kwargs=kwargs,
seq_len=seq_len,
#seed=seed # add this manually to the ngym source code before uncommenting
)
ob_size = env.observation_space.shape[0]
act_size = env.action_space.n
print(f'ob_size = {ob_size}, act_size = {act_size}')
# Print environment specifications
if verbose:
print("Trial timing (in milliseconds):")
print(env.timing)
print("\nObservation space structure:")
print(env.observation_space)
print("Observation components:")
print(env.observation_space.name)
print("\nAction space structure:")
print(env.action_space)
print("Action mapping:")
print(env.action_space.name)
if n_trials_to_plot>0:
obs, info = env.reset()
# Visualize example trials
fig = plot_env(
env,
name='Context , envDecision Making',
num_trials=n_trials_to_plot,
def_act=0,
fig_kwargs={'figsize':(10,8)}
)
plt.savefig(osjoin(SAVE_PATH, f'env.png'))
plt.close('all')
return dataset , env
[docs]
def correct_targets(u, y):
# two contexts, one fixation, one go cue
dim_ring = (u.shape[1]-4)//2
# a copy of y (to be adapted)
target_timeseries = np.copy(y)
# labels = []
trial_onsets = get_trial_onset_idxs(u)
go_onsets = get_go_onset_idxs(u)
go_offsets= get_go_offset_idxs(u)
for t_onset, g_onset, g_offset in zip(trial_onsets, go_onsets, go_offsets):
u1 = u[t_onset:g_onset, 2:2+dim_ring].mean(axis=0) # mode 1 (color)
u2 = u[t_onset:g_onset, 2+dim_ring:2+2*dim_ring].mean(axis=0) # mode 2 (motion)
c = u[t_onset:g_onset, 2+2*dim_ring:].mean(axis=0) # context
choice_color = np.argmax(u1).squeeze()
choice_motion= np.argmax(u2).squeeze()
ctx = np.argmax(c).squeeze()
if ctx == 0:
target_timeseries[g_onset:g_offset] = choice_color
# labels.append(choice_color)
else:
target_timeseries[g_onset:g_offset] = choice_motion
# labels.append(choice_motion)
return target_timeseries#, np.array(labels)
[docs]
def draw_uyl(dataset, verbose=True, return_label=True):
"""
Draws new samples from the dataset (psychological environment) and
adds the go-cue to the first component of stimulus as it is not
provided in ngym. See: https://github.com/neurogym/neurogym/issues/279
Note that due to stochasticity of the delay period, the
number of samples (fixation + stimulus + delay + go-cue) is not the
same between different calls. Thus, we also return the length of
each trial.
"""
inputs, targets = dataset()
# copy is needed to avoid overwriting upon further calls
inputs = np.copy(inputs.squeeze())
targets = np.copy(targets.squeeze())
# add the go cue as the first component
inputs = np.vstack([targets>0, inputs.T]).T
# keep only complete trials (staring with fixation and ending with go que)
fixation_idxs = np.where(inputs[:, 1]>0)[0]
start = fixation_idxs[0]
end = fixation_idxs[-1]
inputs = inputs[start:end]
targets= targets[start:end]
# It seems to me that neurogym doesn't set the labels correctly. So I
# make fix them manually here.
# targets, labels = correct_targets(inputs, targets)
targets = correct_targets(inputs, targets)
# extract the length of each trial
# lengths = np.diff(np.concat((fixation_idxs, [len(inputs)])))
lengths = np.diff(fixation_idxs)
_, offs = hu.get_trial_start_end(lengths)
labels = np.array([targets[off-1] for off in offs])
if verbose:
# _dt = dataset.env.spec.kwargs['dt'] # in ms
# _dur_decision = dataset.env.spec.kwargs['timing']['decision'] # in ms
# steps_in_decision = int(_dur_decision//_dt)
# ntrials = int(inputs.sum(axis=0)[0] // steps_in_decision)
# onset_idxs = get_trial_onset_idxs(inputs)
print(f'There are {len(lengths)} trials in this sequence.')
if return_label:
return inputs, labels, lengths
else:
return inputs, targets, lengths
[docs]
def get_trial_onset_idxs(u, fix_column = 1):
return np.argwhere(np.diff(u[:, fix_column])<0).squeeze()
[docs]
def get_go_idxs(u, go_column=0):
return np.argwhere(u[:, go_column]>0).squeeze()
[docs]
def get_go_onset_idxs(u, go_column = 0):
return np.argwhere(np.diff(u[:, go_column], prepend=u[0, go_column])>0).squeeze()
[docs]
def get_go_offset_idxs(u, go_column = 0):
tmp = np.argwhere(np.diff(u[:, go_column], prepend=u[0, go_column])<0).squeeze()
tmp = np.concatenate([tmp, [len(u)]]) # append last index
return tmp
[docs]
def get_go_interval(u, arr=None):
go_onset = get_go_onset_idxs(u)
go_offset = get_go_offset_idxs(u)
slices = [slice(on, off) for on, off in zip(go_onset, go_offset)]
if arr is not None:
sliced = np.stack([arr[s] for s in slices])
else:
sliced = None
return slices, sliced
[docs]
def get_train_test(dataset, center=False, unit_std=False, normed_trial=True,
orthgonalize=False, whiten=False, zero_phase=False):
"""
Generates two independent seqeunce from the dataset, used as the train
and test sets. The number of samples in each set is roughly the same
for long sequences (to be defined upon instantiation of the `dataset`
object). Stimulus channels are centered using the statistics of the
train set.
Returns train and test tuples containing:
- stimulus `u` (nseq x dim-u)
- labels `y` (nseq x num-classes)
- go-cue indices `idx` (nseq)
"""
us_trn, ys_trn, lens_trn = draw_uyl(dataset)
us_tst, ys_tst, lens_tst = draw_uyl(dataset)
print(ys_trn, set(ys_trn))
if center:
ss = StandardScaler(with_std=unit_std)
us_trn = ss.fit_transform(us_trn)
us_tst = ss.transform(us_tst)
if orthgonalize:
if zero_phase:
whiten = False
pca = PCA(whiten=whiten)
trn_t = pca.fit_transform(us_trn[:,2:])
tst_t = pca.transform(us_tst[:, 2:])
if zero_phase:
trn_t = trn_t@pca.components_.T
tst_t = tst_t@pca.components_.T
us_trn = np.c_[us_trn[:,:2], trn_t]
us_tst = np.c_[us_tst[:,:2], tst_t]
if normed_trial:
ons, offs = hu.get_trial_start_end(lens_trn)
for on, off in zip(ons, offs):
us_trn[on:off, 2:] /= np.linalg.norm(us_trn[on:off, 2:])
ons, offs = hu.get_trial_start_end(lens_tst)
for on, off in zip(ons, offs):
us_tst[on:off, 2:] /= np.linalg.norm(us_tst[on:off, 2:])
udim = us_trn.shape[1] # go cue + fixation + ... + ctx 1 + ctx 2
dim_ring = (udim - 2 - 2)//2
channles_to_keep = np.arange(1, 1+2*dim_ring+2)
us_trn = us_trn[:,channles_to_keep] # dropping the redundant context
us_tst = us_tst[:,channles_to_keep] # dropping the redundant context
return (us_trn, ys_trn, lens_trn), (us_tst, ys_tst, lens_tst)
[docs]
def run_single_net(h='paper', h_is_multiplicative=True, tau0=None,
N=250, p=.1, seed=2026,
seq_len=200000, dt=1, dim_ring=6, af='tanh',
pooling='snapshot', nsnapshot = 10, golength=10,
unit_std=False, center=False,
orthgonalize=True, zero_phase=True,
filt=False, first_unstream=True, flattening_method='flat',
):
ds, env = make_dataset(dt=dt, seq_len=seq_len,
dim_ring=dim_ring, seed=seed,
) # creates random actions
trn, tst = get_train_test(ds,
unit_std=unit_std, center=center,
orthgonalize=orthgonalize, zero_phase=zero_phase)
us_trn, ys_trn, lens_trn = trn
us_tst, ys_tst, lens_tst = tst
h = hu.set_het_level(value=h, lens = lens_trn)
# extract tau
if tau0 is None:
tau0, fmaxs = hu.extract_base_tau(us_trn, dt=dt,
sample_lengths=lens_trn,
filt=filt, first_unstream=first_unstream,
flattening_method=flattening_method, center=False,
savepath=savepath, savename='taus'
)
# generate net
net_settings = dict(N=N, p=p, N_input=us_trn.shape[1], savepath=savepath, savename='net')
A_adj, w_in, w_rec = hu.generate_net(**net_settings)
w_in /= us_trn.std(axis=0) # scale the input weights
# samples taus
tau_settings = dict(size=N, mu=tau0, clip_below=dt, savepath=savepath, savename='tau')
std = h*tau0 if h_is_multiplicative else h*dt
tau_hom = hu.sample_taus(std = 0*tau0, **tau_settings)
tau_het = hu.sample_taus(std = std , **tau_settings)
# integrate
int_settings = dict(A_adj=A_adj, w_in=w_in, w_rec=w_rec, af=af, dt=dt, reset_index=1)
# integrating homog
vs_trn_hom = hu.integrate(u=us_trn, taus=tau_hom, **int_settings)
vs_tst_hom = hu.integrate(u=us_tst, taus=tau_hom, **int_settings)
xs_trn_hom = hu.activation_func(vs_trn_hom, af)
xs_tst_hom = hu.activation_func(vs_tst_hom, af)
# integrating heterog
vs_trn_het = hu.integrate(u=us_trn, taus=tau_het, **int_settings)
vs_tst_het = hu.integrate(u=us_tst, taus=tau_het, **int_settings)
xs_trn_het = hu.activation_func(vs_trn_het, af)
xs_tst_het = hu.activation_func(vs_tst_het, af)
# classes distribution
hu.plot_dataset_hist(ys_trn, savepath=savepath, savename='trn')
hu.plot_dataset_hist(ys_tst, savepath=savepath, savename='tst')
# uvx
hu.plot_uvx_stream(us_trn, vs_trn_hom, xs_trn_hom, lengths=lens_trn,
input_labels = get_input_channel_names(us_trn),
savename='uvx_homog', savepath=savepath,)
hu.plot_uvx_stream(us_trn, vs_trn_het, xs_trn_het, lengths=lens_trn,
input_labels = get_input_channel_names(us_trn),
savename='uvx_heter', savepath=savepath,)
# manifolds
hu.plot_manifold(us_trn, vs_trn_hom, xs_trn_hom, ys_trn, lens_trn,
savename='manifold_homog', savepath=savepath,)
hu.plot_manifold(us_trn, vs_trn_het, xs_trn_het, ys_trn, lens_trn,
savename='manifold_heter', savepath=savepath,)
# pooling states
pooling_config = dict(flatten=True, nsnapshot=nsnapshot, golength=golength,)
Xs_trn_hom = hu.pool(xs_trn_hom, trial_lengths=lens_trn, func=pooling, **pooling_config)
Xs_trn_het = hu.pool(xs_trn_het, trial_lengths=lens_trn, func=pooling, **pooling_config)
Xs_tst_hom = hu.pool(xs_tst_hom, trial_lengths=lens_tst, func=pooling, **pooling_config)
Xs_tst_het = hu.pool(xs_tst_het, trial_lengths=lens_tst, func=pooling, **pooling_config)
clf = hu.classifier()
clf.fit(Xs_trn_hom, ys_trn)
trn_hom = clf.score(Xs_trn_hom, ys_trn)
tst_hom = clf.score(Xs_tst_hom, ys_tst)
clf.fit(Xs_trn_het, ys_trn)
trn_het = clf.score(Xs_trn_het, ys_trn)
tst_het = clf.score(Xs_tst_het, ys_tst)
# # extracting training slice
# if train_on_go:
# _, Xs_trn_hom = get_go_interval(us_trn, xs_trn_hom)
# _, Xs_trn_het = get_go_interval(us_trn, xs_trn_het)
# _, Xs_tst_hom = get_go_interval(us_tst, xs_tst_hom)
# _, Xs_tst_het = get_go_interval(us_tst, xs_tst_het)
# else:
# snap_config = dict(nsnapshot=num_snapshot, flatten=True)
# _, Xs_trn_hom = hu.get_snapshot(lens_trn, arr=xs_trn_hom, **snap_config)
# _, Xs_trn_het = hu.get_snapshot(lens_trn, arr=xs_trn_het, **snap_config)
# _, Xs_tst_hom = hu.get_snapshot(lens_tst, arr=xs_tst_hom, **snap_config)
# _, Xs_tst_het = hu.get_snapshot(lens_tst, arr=xs_tst_het, **snap_config)
# trn_decision_onset = get_go_onset_idxs(us_trn)
# tst_decision_onset = get_go_onset_idxs(us_tst)
# trn_decision_offset = get_go_offset_idxs(us_trn)
# tst_decision_offset = get_go_offset_idxs(us_tst)
# trn_slices = [slice(onset, offset) for onset, offset in zip(trn_decision_onset, trn_decision_offset)]
# tst_slices = [slice(onset, offset) for onset, offset in zip(tst_decision_onset, tst_decision_offset)]
# # forming the design matrix
# Xs_trn_hom = np.stack([xs_trn_hom[s].ravel() for s in trn_slices])
# Xs_trn_het = np.stack([xs_trn_het[s].ravel() for s in trn_slices])
# Xs_tst_hom = np.stack([xs_tst_hom[s].ravel() for s in tst_slices])
# Xs_tst_het = np.stack([xs_tst_het[s].ravel() for s in tst_slices])
# classifying
# clf_hom = hu.classifier()
# clf_het = hu.classifier()
# clf_hom.fit(Xs_trn_hom, ys_trn)
# clf_het.fit(Xs_trn_het, ys_trn)
# report
print(f'Homog : {trn_hom:.3f} | {tst_hom:.3f}')
print(f'Hetero: {trn_het:.3f} | {tst_het:.3f}')
# def run_multinet_v1(
# Nnet = 50, # number of MC realization
# h='paper', h_is_multiplicative=True, # how to set heterogeneity's value
# N=250, p=.1, af = 'tanh', # nework parameters
# seq_len=200000, dt=1, dim_ring=6, # task parameters
# delay_period= [50, 30, 300], # delay in decision making
# stim_period = 100, # stimulation duration
# unit_std=False, center=False, # input normalization
# orthgonalize=False, zero_phase=False, whiten=False, # input orthogonaliziation
# filt=False, # filtering before tau finding
# first_unstream=True, flattening_method='flat', # to stream or not
# pooling_funcs=['go', 'avg', 'snapshot'], # pooling methods
# nsnapshot = 10, golength=10, # pooling args
# prefix='', # simulation name
# seed=2026, n_jobs=-1,
# plot_uvx_man=True
# ):
# if prefix == '':
# mult = 'M' if h_is_multiplicative else ''
# prefix = f'N{N}_MC{Nnet}_d{dim_ring}_h{h}{mult}'
# prefix+= f'_c{int(center)}_std{int(unit_std)}'
# prefix+= f'_o{int(orthgonalize)}_z{int(zero_phase)}_w{int(whiten)}'
# print(prefix)
# savepath = osjoin(SAVE_PATH, prefix)
# Path(savepath).mkdir(parents=True, exist_ok=True)
# ds, env = make_dataset(dt=dt, seq_len=seq_len,
# dim_ring=dim_ring, seed=seed,
# delay_period=delay_period,
# stim_period=stim_period
# ) # creates random actions
# trn, tst = get_train_test(ds,
# unit_std=unit_std, center=center,
# orthgonalize=orthgonalize,
# zero_phase=zero_phase,
# whiten=whiten,)
# us_trn, ys_trn, lens_trn = trn
# us_tst, ys_tst, lens_tst = tst
# h = hu.set_het_level(value=h, lens = lens_trn)
# # extract tau
# tau0, fmaxs = hu.extract_base_tau(us_trn, dt=dt,
# sample_lengths=lens_trn,
# filt=filt, first_unstream=first_unstream,
# flattening_method=flattening_method, center=False,
# savepath=savepath, savename='taus_multinet'
# )
# # Check which results already exist and load them
# taus_het_path = osjoin(savepath, f'{prefix}_taus_het.npy')
# completed_net_indices = set()
# taus_het = np.zeros((Nnet, N))
# score_tst_hom = np.zeros((Nnet, len(pooling_funcs)))
# score_trn_hom = np.zeros((Nnet, len(pooling_funcs)))
# score_tst_het = np.zeros((Nnet, len(pooling_funcs)))
# score_trn_het = np.zeros((Nnet, len(pooling_funcs)))
# # print(taus_het.shape, score_tst_hom.shape, )
# # Load existing tau results if available
# if os.path.exists(taus_het_path):
# taus_het = np.load(taus_het_path)
# # Check which realizations and pools have all networks completed
# for net_idx in range(Nnet):
# all_completed = True
# for pool_idx, pooling_func in enumerate(pooling_funcs):
# poolname = f'{prefix}_{pooling_func}'
# # check it the score arrays exist at all for that pool
# score_path = osjoin(savepath, f'{poolname}_score_tst_hom.npy')
# if os.path.exists(score_path):
# loaded_scores = np.load(score_path)
# # if so, check if the results for net_idx exist
# if loaded_scores.shape[0] > net_idx and np.any(loaded_scores[net_idx] != 0):
# score_tst_hom[net_idx, pool_idx] = loaded_scores[net_idx]
# else:
# all_completed = False
# break
# # if the results exist, fill up the score arrays
# loaded_scores = np.load(osjoin(savepath, f'{poolname}_score_trn_hom.npy'))
# score_trn_hom[net_idx, pool_idx] = loaded_scores[net_idx]
# loaded_scores = np.load(osjoin(savepath, f'{poolname}_score_tst_het.npy'))
# score_tst_het[net_idx, pool_idx] = loaded_scores[net_idx]
# loaded_scores = np.load(osjoin(savepath, f'{poolname}_score_trn_het.npy'))
# score_trn_het[net_idx, pool_idx] = loaded_scores[net_idx]
# else:
# all_completed = False
# break
# if all_completed:
# completed_net_indices.add(net_idx)
# print(f'Network realization {net_idx} already completed, skipping.')
# def save_net_results(net_idx):
# # saves network realization (broken down based on the pooling strategy)
# # overrides previous ones, as it contains more results now
# np.save(osjoin(savepath, f'{prefix}_taus_het.npy'), taus_het)
# for pool_idx, pooling_func in enumerate(pooling_funcs):
# poolname = f'{prefix}_{pooling_func}'
# np.save(osjoin(savepath, f'{poolname}_score_tst_hom.npy'), score_tst_hom[..., pool_idx])
# np.save(osjoin(savepath, f'{poolname}_score_trn_hom.npy'), score_trn_hom[..., pool_idx])
# np.save(osjoin(savepath, f'{poolname}_score_tst_het.npy'), score_tst_het[..., pool_idx])
# np.save(osjoin(savepath, f'{poolname}_score_trn_het.npy'), score_trn_het[..., pool_idx])
# def run_net(net_idx, pooling_funcs):
# print(f'====== net number {net_idx} ======')
# np.random.seed(seed + net_idx)
# # sample taus
# tau_settings = dict(size=N, clip_below=dt, savename=None)
# std = h*tau0 if h_is_multiplicative else h*dt
# tau_hom = hu.sample_taus(mu=tau0, std= 0 , **tau_settings)
# tau_het = hu.sample_taus(mu=tau0, std=std, **tau_settings)
# # initialize net
# net_settings = dict(N=N, p=p, N_input=us_trn.shape[1], savename=None)
# A_adj, w_in, w_rec = hu.generate_net(**net_settings)
# w_in /= us_trn.std(axis=0) # scale the network
# # integrating
# int_settings = dict(A_adj=A_adj, w_in=w_in, w_rec=w_rec, af=af, dt=dt, reset_index=1)
# ## homog
# vs_trn_hom = hu.integrate(u=us_trn, taus=tau_hom, **int_settings).astype(np.float32)
# vs_tst_hom = hu.integrate(u=us_tst, taus=tau_hom, **int_settings).astype(np.float32)
# xs_trn_hom = hu.activation_func(vs_trn_hom, af)
# xs_tst_hom = hu.activation_func(vs_tst_hom, af)
# # heterog
# vs_trn_het = hu.integrate(u=us_trn, taus=tau_het, **int_settings).astype(np.float32)
# vs_tst_het = hu.integrate(u=us_tst, taus=tau_het, **int_settings).astype(np.float32)
# xs_trn_het = hu.activation_func(vs_trn_het, af)
# xs_tst_het = hu.activation_func(vs_tst_het, af)
# # plotting
# if plot_uvx_man:
# if net_idx == 0:
# input_labels = get_input_channel_names(us_trn)
# uvx_name = f'{prefix}_net{net_idx}'
# hu.plot_uvx_stream(us_trn, vs_trn_hom, xs_trn_hom, lengths=lens_trn,
# tmax=1500, input_labels = input_labels,
# savepath=savepath, savename=f'{uvx_name}_homog')
# hu.plot_uvx_stream(us_trn, vs_trn_het, xs_trn_het, lengths=lens_trn,
# tmax=1500, input_labels = input_labels,
# savepath=savepath, savename=f'{uvx_name}_heter')
# man_name = f'{prefix}_net{net_idx}'
# hu.plot_manifold(us_trn, vs_trn_hom, xs_trn_hom, ys_trn, lens_trn,
# savepath=savepath, savename=f'{man_name}_homog')
# hu.plot_manifold(us_trn, vs_trn_het, xs_trn_het, ys_trn, lens_trn,
# savepath=savepath, savename=f'{man_name}_heter', )
# # classification
# trn_homs = np.zeros(len(pooling_funcs))
# tst_homs = np.zeros(len(pooling_funcs))
# trn_hets = np.zeros(len(pooling_funcs))
# tst_hets = np.zeros(len(pooling_funcs))
# for pool_idx, func in enumerate(pooling_funcs):
# # pooling states
# pooling_config = dict(flatten=True, func=func,
# nsnapshot=nsnapshot, golength=golength,)
# Xs_trn_hom = hu.pool(xs_trn_hom, trial_lengths=lens_trn, **pooling_config)
# Xs_trn_het = hu.pool(xs_trn_het, trial_lengths=lens_trn, **pooling_config)
# Xs_tst_hom = hu.pool(xs_tst_hom, trial_lengths=lens_tst, **pooling_config)
# Xs_tst_het = hu.pool(xs_tst_het, trial_lengths=lens_tst, **pooling_config)
# # classify
# # Since we are already having serveral realizations in different core we
# # limit the classifications to few cores. Otherwise, this easily fills
# # up the ram. The number of cores here is to be adjusted depending on the
# # availabel cores and the number of concurrent realized nets.
# clf = hu.classifier(n_jobs=2)
# print(f'startin to classify {func}-{net_idx}')
# ## homog
# clf.fit(Xs_trn_hom, ys_trn)
# trn_hom = clf.score(Xs_trn_hom, ys_trn)
# tst_hom = clf.score(Xs_tst_hom, ys_tst)
# ## hetero
# clf.fit(Xs_trn_het, ys_trn)
# trn_het = clf.score(Xs_trn_het, ys_trn)
# tst_het = clf.score(Xs_tst_het, ys_tst)
# # report
# print(f'\t {func}-{net_idx} -- Homog: {tst_hom:.3f} | Hetero: {tst_het:.3f}')
# trn_homs[pool_idx] = trn_hom
# tst_homs[pool_idx] = tst_hom
# trn_hets[pool_idx] = trn_het
# tst_hets[pool_idx] = tst_het
# return net_idx, tau_het, trn_homs, tst_homs, trn_hets, tst_hets
# # Skip if this tau0 index is already completed
# net_idxs = [net_id for net_id in range(Nnet) if net_id not in completed_net_indices]
# par = Parallel(n_jobs=n_jobs, max_nbytes='1M', backend='loky')
# delayed_func = delayed(run_net)
# for presult in par(delayed_func(net_idx, pooling_funcs) for net_idx in net_idxs):
# net_idx, tau_het, trn_hom, tst_hom, trn_het, tst_het = presult
# taus_het[net_idx] = tau_het
# score_trn_hom[net_idx, :] = trn_hom
# score_tst_hom[net_idx, :] = tst_hom
# score_trn_het[net_idx, :] = trn_het
# score_tst_het[net_idx, :] = tst_het
# # Save results immediately after each network completes
# save_net_results(net_idx)
# hu.plot_score_summary(score_trn_hom, score_trn_het,
# score_tst_hom, score_tst_het,
# savepath=savepath, savename=prefix,
# chance_score=1./dim_ring,
# pooling_funcs=pooling_funcs,
# )
# def run_spaced_tau0_multinet_v1(
# Nnet = 50, # number of MC realization
# h='paper', h_is_multiplicative=True, # how to set heterogeneity's value
# N=250, p=.1, af = 'tanh', # nework parameters
# seq_len=200000, dt=1, dim_ring=6, # task parameters
# delay_period= [50, 30, 300], # delay in decision making
# stim_period = 100, # stimulation duration
# unit_std=False, center=False, # input normalization
# orthgonalize=True, zero_phase=True, whiten=False, # input orthogonaliziation
# num_tau0=10, # number of distinct E[taus] to use
# pooling_funcs=['go', 'avg', 'snapshot'], # pooling methods
# nsnapshot = 10, golength=10, # pooling args
# prefix='', # simulation name
# seed=2026, n_jobs=-1,
# plot_uvx=True, plot_manifold=False,
# ):
# if prefix == '':
# mult = 'M' if h_is_multiplicative else ''
# prefix = f'SPACED_N{N}_MC{Nnet}_d{dim_ring}_h{h}{mult}'
# print(prefix)
# savepath = osjoin(SAVE_PATH, prefix)
# Path(savepath).mkdir(parents=True, exist_ok=True)
# ds, env = make_dataset(dt=dt, seq_len=seq_len,
# dim_ring=dim_ring, seed=seed,
# delay_period=delay_period,
# stim_period=stim_period
# ) # creates random actions
# trn, tst = get_train_test(ds)
# us_trn, ys_trn, lens_trn = trn
# us_tst, ys_tst, lens_tst = tst
# h = hu.set_het_level(value=h, lens = lens_trn)
# # mean taus values
# spaced_tau = hu.get_spaced_hom_taus(lens_trn, dt=dt, num=num_tau0)
# # Check which results already exist and load them
# taus_hom_path = osjoin(savepath, f'{prefix}_taus_hom.npy')
# taus_het_path = osjoin(savepath, f'{prefix}_taus_het.npy')
# completed_tau_indices = set()
# spaced_taus_het = np.zeros((num_tau0, Nnet, N))
# score_tst_hom = np.zeros((num_tau0, Nnet, len(pooling_funcs)))
# score_trn_hom = np.zeros((num_tau0, Nnet, len(pooling_funcs)))
# score_tst_het = np.zeros((num_tau0, Nnet, len(pooling_funcs)))
# score_trn_het = np.zeros((num_tau0, Nnet, len(pooling_funcs)))
# # Load existing tau results if available
# if os.path.exists(taus_het_path):
# spaced_taus_het = np.load(taus_het_path)
# # Check which tau0 indices have all networks completed
# for tau0_idx in range(num_tau0):
# all_completed = True
# for pool_idx, pooling_func in enumerate(pooling_funcs):
# poolname = f'{prefix}_{pooling_func}'
# score_path = osjoin(savepath, f'{poolname}_score_tst_hom.npy')
# if os.path.exists(score_path):
# loaded_scores = np.load(score_path)
# if loaded_scores.shape[0] > tau0_idx and np.any(loaded_scores[tau0_idx] != 0):
# score_tst_hom[tau0_idx, :, pool_idx] = loaded_scores[tau0_idx]
# else:
# all_completed = False
# break
# loaded_scores = np.load(osjoin(savepath, f'{poolname}_score_trn_hom.npy'))
# score_trn_hom[tau0_idx, :, pool_idx] = loaded_scores[tau0_idx]
# loaded_scores = np.load(osjoin(savepath, f'{poolname}_score_tst_het.npy'))
# score_tst_het[tau0_idx, :, pool_idx] = loaded_scores[tau0_idx]
# loaded_scores = np.load(osjoin(savepath, f'{poolname}_score_trn_het.npy'))
# score_trn_het[tau0_idx, :, pool_idx] = loaded_scores[tau0_idx]
# else:
# all_completed = False
# break
# if all_completed:
# completed_tau_indices.add(tau0_idx)
# print(f'Tau0 index {tau0_idx} already completed, skipping.')
# def save_tau0_results(tau0_idx):
# if tau0_idx == 0 :
# np.save(osjoin(savepath, f'{prefix}_taus_hom.npy'), spaced_taus_het)
# np.save(osjoin(savepath, f'{prefix}_taus_het.npy'), spaced_taus_het)
# for pool_idx, pooling_func in enumerate(pooling_funcs):
# poolname = f'{prefix}_{pooling_func}'
# np.save(osjoin(savepath, f'{poolname}_score_tst_hom.npy'), score_tst_hom[..., pool_idx])
# np.save(osjoin(savepath, f'{poolname}_score_trn_hom.npy'), score_trn_hom[..., pool_idx])
# np.save(osjoin(savepath, f'{poolname}_score_tst_het.npy'), score_tst_het[..., pool_idx])
# np.save(osjoin(savepath, f'{poolname}_score_trn_het.npy'), score_trn_het[..., pool_idx])
# def run_net(tau0_idx, tau0, net_idx, pooling_funcs):
# print(f' ----- net number {net_idx} -------')
# np.random.seed(seed + net_idx)
# # sample taus
# tau_settings = dict(size=N, clip_below=dt, savename=None)
# std = h*tau0 if h_is_multiplicative else h*dt
# tau_hom = hu.sample_taus(mu=tau0, std= 0 , **tau_settings)
# tau_het = hu.sample_taus(mu=tau0, std=std, **tau_settings)
# # initialize net
# net_settings = dict(N=N, p=p, N_input=us_trn.shape[1], savename=None)
# A_adj, w_in, w_rec = hu.generate_net(**net_settings)
# w_in /= us_trn.std(axis=0) # scale the network
# # integrating
# int_settings = dict(A_adj=A_adj, w_in=w_in, w_rec=w_rec, af=af, dt=dt, reset_index=1)
# ## homog
# vs_trn_hom = hu.integrate(u=us_trn, taus=tau_hom, **int_settings).astype(np.float32)
# vs_tst_hom = hu.integrate(u=us_tst, taus=tau_hom, **int_settings).astype(np.float32)
# xs_trn_hom = hu.activation_func(vs_trn_hom, af)
# xs_tst_hom = hu.activation_func(vs_tst_hom, af)
# # heterog
# vs_trn_het = hu.integrate(u=us_trn, taus=tau_het, **int_settings).astype(np.float32)
# vs_tst_het = hu.integrate(u=us_tst, taus=tau_het, **int_settings).astype(np.float32)
# xs_trn_het = hu.activation_func(vs_trn_het, af)
# xs_tst_het = hu.activation_func(vs_tst_het, af)
# # plotting
# if plot_uvx:
# if net_idx == 0:
# input_labels = get_input_channel_names(us_trn)
# uvx_name = f'{prefix}_tau{tau0_idx}_net{net_idx}'
# hu.plot_uvx_stream(us_trn, vs_trn_hom, xs_trn_hom, lengths=lens_trn,
# tmax=1500, input_labels = input_labels,
# savepath=savepath, savename=f'{uvx_name}_homog')
# hu.plot_uvx_stream(us_trn, vs_trn_het, xs_trn_het, lengths=lens_trn,
# tmax=1500, input_labels = input_labels,
# savepath=savepath, savename=f'{uvx_name}_heter')
# if plot_man:
# if tau0_idx == 0 or tau0_idx == 9:
# man_name = f'{prefix}_tau{tau0_idx}_net{net_idx}'
# hu.plot_manifold(us_trn, vs_trn_hom, xs_trn_hom, ys_trn, lens_trn,
# savepath=savepath, savename=f'{man_name}_homog')
# hu.plot_manifold(us_trn, vs_trn_het, xs_trn_het, ys_trn, lens_trn,
# savepath=savepath, savename=f'{man_name}_heter', )
# # classification
# trn_homs, tst_homs, trn_hets, tst_hets = [], [], [], []
# for pool_idx, func in enumerate(pooling_funcs):
# # pooling states
# pooling_config = dict(flatten=True, func=func,
# nsnapshot=nsnapshot, golength=golength,)
# Xs_trn_hom = hu.pool(xs_trn_hom, trial_lengths=lens_trn, **pooling_config)
# Xs_trn_het = hu.pool(xs_trn_het, trial_lengths=lens_trn, **pooling_config)
# Xs_tst_hom = hu.pool(xs_tst_hom, trial_lengths=lens_tst, **pooling_config)
# Xs_tst_het = hu.pool(xs_tst_het, trial_lengths=lens_tst, **pooling_config)
# # classify
# # Since we are already having serveral realizations in different core we
# # limit the classifications to few cores. Otherwise, this easily fills
# # up the ram. The number of cores here is to be adjusted depending on the
# # availabel cores and the number of concurrent realized nets.
# clf = hu.classifier(n_jobs=2)
# print(f'startin to classify {func}-{net_idx}')
# ## homog
# clf.fit(Xs_trn_hom, ys_trn)
# trn_hom = clf.score(Xs_trn_hom, ys_trn)
# tst_hom = clf.score(Xs_tst_hom, ys_tst)
# ## hetero
# clf.fit(Xs_trn_het, ys_trn)
# trn_het = clf.score(Xs_trn_het, ys_trn)
# tst_het = clf.score(Xs_tst_het, ys_tst)
# # report
# print(f'\t {func}-{net_idx} -- Homog: {tst_hom:.3f} | Hetero: {tst_het:.3f}')
# trn_homs.append(trn_hom)
# tst_homs.append(tst_hom)
# trn_hets.append(trn_het)
# tst_hets.append(tst_het)
# trn_homs = np.array(trn_homs)
# tst_homs = np.array(tst_homs)
# trn_hets = np.array(trn_hets)
# tst_hets = np.array(tst_hets)
# return net_idx, tau_het, trn_homs, tst_homs, trn_hets, tst_hets
# def save_net_results_immediate(tau0_idx, results):
# """Save results immediately after each network completes to avoid waiting for all workers."""
# # Save tau_het for this specific network
# spaced_taus_het[tau0_idx, results['net_idx']] = results['tau_het']
# # Save scores for this specific network
# score_trn_hom[tau0_idx, results['net_idx'], :] = results['trn_homs']
# score_tst_hom[tau0_idx, results['net_idx'], :] = results['tst_homs']
# score_trn_het[tau0_idx, results['net_idx'], :] = results['trn_hets']
# score_tst_het[tau0_idx, results['net_idx'], :] = results['tst_hets']
# # Save results immediately after each network completes
# save_tau0_results(tau0_idx)
# for tau0_idx, tau0 in enumerate(spaced_tau):
# # Skip if this tau0 index is already completed
# if tau0_idx in completed_tau_indices:
# continue
# print(f'====== {tau0_idx+1}: tau0 = {tau0} ======')
# par = Parallel(n_jobs=n_jobs, max_nbytes='1M', backend='loky')
# delayed_func = delayed(run_net)
# for presult in par(delayed_func(tau0_idx, tau0, net_idx, pooling_funcs) for net_idx in range(Nnet)):
# net_idx, tau_het, trn_homs, tst_homs, trn_hets, tst_hets = presult
# spaced_taus_het[tau0_idx, net_idx] = tau_het
# score_trn_hom[tau0_idx, net_idx, :] = trn_homs
# score_tst_hom[tau0_idx, net_idx, :] = tst_homs
# score_trn_het[tau0_idx, net_idx, :] = trn_hets
# score_tst_het[tau0_idx, net_idx, :] = tst_hets
# save_tau0_results(net_idx)
# # The results are now saved inside run_net immediately after computation
# # No need to save here again, as it's already handled in run_net
# # pass
# hu.plot_score_summary_of_spaced_hom(score_trn_hom, score_trn_het,
# score_tst_hom, score_tst_het,
# spaced_taus_hom = spaced_tau,
# savepath=savepath, savename=prefix,
# chance_score=1./dim_ring,
# pooling_funcs=pooling_funcs,
# )
# # taus_hom = np.load(osjoin(SAVE_PATH, f'{prefix}_taus_hom.npy'))
# # taus_het = np.load(osjoin(SAVE_PATH, f'{prefix}_taus_het.npy'))
# # score_tst_hom = np.load(osjoin(SAVE_PATH, f'{prefix}_score_tst_hom.npy'))
# # score_trn_hom = np.load(osjoin(SAVE_PATH, f'{prefix}_score_trn_hom.npy'))
# # score_tst_het = np.load(osjoin(SAVE_PATH, f'{prefix}_score_tst_het.npy'))
# # score_trn_het = np.load(osjoin(SAVE_PATH, f'{prefix}_score_trn_het.npy'))
# def run_multinet_v0(
# Nnet = 50, # number of MC realization
# h='paper', h_is_multiplicative=True, # how to set heterogeneity's value
# N=250, p=.1, af = 'tanh', # nework parameters
# seq_len=200000, dt=1, dim_ring=6, # task parameters
# delay_period= [50, 30, 300], # delay in decision making
# stim_period = 100, # stimulation duration
# unit_std=False, center=False, # input normalization
# orthgonalize=False, zero_phase=True, whiten=False, # input orthogonaliziation
# filt=False, # filtering before tau finding
# first_unstream=True, flattening_method='flat', # to stream or not
# pooling_funcs=['go', 'avg', 'snapshot'], # pooling methods
# golength=10, nsnapshot = 10, # pooling args
# seed=2026, n_jobs=-1,
# prefix='', # simulation name
# plot_uvx=True, plot_man=False,
# ):
# if prefix == '':
# mult = 'M' if h_is_multiplicative else ''
# prefix = f'N{N}_MC{Nnet}_d{dim_ring}_h{h}{mult}'
# prefix+= f'_c{int(center)}_std{int(unit_std)}'
# prefix+= f'_o{int(orthgonalize)}_z{int(zero_phase)}_w{int(whiten)}'
# print(prefix)
# savepath = osjoin(SAVE_PATH, prefix)
# Path(savepath).mkdir(parents=True, exist_ok=True)
# ds, env = make_dataset(dt=dt, seq_len=seq_len,
# dim_ring=dim_ring, seed=seed,
# delay_period=delay_period,
# stim_period=stim_period
# ) # creates random actions
# trn, tst = get_train_test(ds,
# unit_std=unit_std, center=center,
# orthgonalize=orthgonalize,
# zero_phase=zero_phase,
# whiten=whiten,)
# us_trn, ys_trn, lens_trn = trn
# us_tst, ys_tst, lens_tst = tst
# Nclasses = len(np.unique(ys_tst))
# reset_dim = 1
# h = hu.set_het_level(value=h, lens = lens_trn)
# # extract tau
# tau0, fmaxs = hu.extract_base_tau(us_trn, dt=dt,
# sample_lengths=lens_trn,
# filt=filt, first_unstream=first_unstream,
# flattening_method=flattening_method, center=False,
# savepath=savepath, savename='taus_multinet'
# )
# # needed to make sample_integrate_classify_save utility happy :)
# tau0s = [tau0]
# num_tau0 = 1
# pooling_kwargs = dict(nsnapshot = nsnapshot, golength=golength, flatten=True)
# for tau0_idx, tau0 in enumerate(tau0s):
# std = h*tau0 if h_is_multiplicative else h*dt
# for net_idx in range(Nnet):
# hu.sample_integrate_classify_save(
# tau0=tau0, std=std, dt=dt, N=N, p=p,
# us_trn=us_trn, us_tst=us_tst,
# lens_trn=lens_trn, lens_tst=lens_tst,
# ys_trn=ys_trn, ys_tst=ys_tst,
# pooling_funcs=pooling_funcs,
# pooling_kwargs=pooling_kwargs, af=af,
# prefix=prefix, savepath=savepath,
# plot_uvx=plot_uvx, plot_man=plot_man,
# net_idx=net_idx, tau0_idx=tau0_idx,
# reset_dim=reset_dim)
# # summarize
# trn_homs, tst_homs, trn_hets, tst_hets = hu.collect_scores(
# savepath=savepath, Ntau0=num_tau0,
# Nnet=Nnet, Npool=len(pooling_funcs))
# hu.plot_score_summary(
# score_trn_hom = trn_homs,
# score_trn_het = trn_hets,
# score_tst_hom = tst_homs,
# score_tst_het = tst_hets,
# savepath=savepath, savename=prefix,
# chance_score=1./Nclasses,
# pooling_funcs=pooling_funcs,
# )
# def run_spaced_tau0_multinet_v0(
# Nnet = 50, # number of MC realization
# h='paper', h_is_multiplicative=True, # how to set heterogeneity's value
# N=250, p=.1, af = 'tanh', # nework parameters
# seq_len=200000, dt=1, dim_ring=6, # task parameters
# delay_period= [50, 30, 300], # delay in decision making
# stim_period = 100, # stimulation duration
# unit_std=False, center=False, # input normalization
# orthgonalize=False, zero_phase=True, whiten=False, # input orthogonaliziation
# num_tau0=10, # number of distinct E[taus] to use
# pooling_funcs=['go', 'avg', 'snapshot'], # pooling methods
# golength=10, nsnapshot = 10, # pooling args
# prefix='', # simulation name
# seed=2026, n_jobs=-1,
# plot_uvx=True, plot_man=True,
# ):
# if prefix == '':
# mult = 'M' if h_is_multiplicative else ''
# prefix = f'SPACED_N{N}_MC{Nnet}_d{dim_ring}_h{h}{mult}'
# print(prefix)
# savepath = osjoin(SAVE_PATH, prefix)
# Path(savepath).mkdir(parents=True, exist_ok=True)
# ds, env = make_dataset(dt=dt, seq_len=seq_len,
# dim_ring=dim_ring, seed=seed,
# delay_period=delay_period,
# stim_period=stim_period
# ) # creates random actions
# trn, tst = get_train_test(ds)
# us_trn, ys_trn, lens_trn = trn
# us_tst, ys_tst, lens_tst = tst
# Nclasses = len(np.unique(ys_tst))
# reset_dim = 1
# h = hu.set_het_level(value=h, lens = lens_trn)
# # mean taus values
# spaced_tau = hu.get_spaced_hom_taus(lens_trn, dt=dt, num=num_tau0)
# pooling_kwargs = dict(nsnapshot = nsnapshot, golength=golength, flatten=True)
# for tau0_idx, tau0 in enumerate(spaced_tau):
# std = h*tau0 if h_is_multiplicative else h*dt
# for net_idx in range(Nnet):
# hu.sample_integrate_classify_save(
# tau0=tau0, std=std, dt=dt, N=N, p=p,
# us_trn=us_trn, us_tst=us_tst,
# lens_trn=lens_trn, lens_tst=lens_tst,
# ys_trn=ys_trn, ys_tst=ys_tst,
# pooling_funcs=pooling_funcs,
# pooling_kwargs=pooling_kwargs, af=af,
# prefix=prefix, savepath=savepath,
# plot_uvx=plot_uvx, plot_man=plot_man,
# net_idx=net_idx, tau0_idx=tau0_idx,
# reset_dim=reset_dim)
# # summarize
# trn_homs, tst_homs, trn_hets, tst_hets = hu.collect_scores(
# savepath=savepath, Ntau0=num_tau0,
# Nnet=Nnet, Npool=len(pooling_funcs))
# hu.plot_score_summary_of_spaced_hom(trn_homs, trn_hets, tst_homs, tst_hets,
# spaced_taus_hom = spaced_tau,
# savepath=savepath, savename=prefix,
# chance_score=1./Nclasses,
# pooling_funcs=pooling_funcs,
# )
[docs]
def run_multinet(
seed, Nnet, # random seed and number of MC realization
N, p, af, # nework parameters
h, h_is_multiplicative, # how to set heterogeneity's value
seq_len, dt, dim_ring, # task parameters
delay_period, # delay in decision making
stim_period, # stimulation duration
unit_std, center, # input normalization
orthgonalize, zero_phase, whiten, # input orthogonaliziation
filt, # filtering before tau finding
first_unstream, flattening_method, # to stream or not and how to flatten
pooling_funcs, # pooling methods
golength, nsnapshot, # pooling args
n_jobs, cv, Creg, # training
prefix, # simulation name
plot_uvx, plot_man, # plotting
**kwargs
):
if prefix == '':
mult = 'M' if h_is_multiplicative else ''
prefix = f'N{N}_MC{Nnet}_d{dim_ring}_h{h}{mult}'
prefix+= f'_c{int(center)}_std{int(unit_std)}'
prefix+= f'_o{int(orthgonalize)}_z{int(zero_phase)}_w{int(whiten)}'
print(prefix)
savepath = osjoin(SAVE_PATH, prefix)
Path(savepath).mkdir(parents=True, exist_ok=True)
ds, env = make_dataset(dt=dt, seq_len=seq_len,
dim_ring=dim_ring, seed=seed,
delay_period=delay_period,
stim_period=stim_period
) # creates random actions
trn, tst = get_train_test(ds,
unit_std=unit_std, center=center,
orthgonalize=orthgonalize,
zero_phase=zero_phase,
whiten=whiten,)
us_trn, ys_trn, lens_trn = trn
us_tst, ys_tst, lens_tst = tst
Nclasses = dim_ring
reset_dim = 1
# h = hu.set_het_level(value=h, lens = lens_trn)
# extract tau
tau0, fmaxs = hu.extract_base_tau(us_trn, dt=dt,
sample_lengths=lens_trn,
filt=filt, first_unstream=first_unstream,
flattening_method=flattening_method, center=False,
savepath=savepath, savename='base_extraction'
)
# needed to make sample_integrate_classify_save utility happy :)
tau0s = [tau0]
num_tau0 = 1
pooling_kwargs = dict(nsnapshot = nsnapshot, golength=golength, flatten=True)
memmap_temp = osjoin(savepath, 'joblib_memmap')
Path(memmap_temp).mkdir(parents=True, exist_ok=True)
Parallel(
n_jobs=n_jobs,
backend='loky',
batch_size=1,
pre_dispatch='n_jobs',
max_nbytes=0,
mmap_mode='r',
temp_folder=memmap_temp,
)(
delayed(hu.sample_integrate_classify_save)(
dt=dt, tau0=tau0,
# std=(h*tau0 if h_is_multiplicative else h*dt),
h_value=h, h_is_multip = h_is_multiplicative,
N=N, p=p,
us_trn=us_trn, us_tst=us_tst,
lens_trn=lens_trn, lens_tst=lens_tst,
ys_trn=ys_trn, ys_tst=ys_tst,
pooling_funcs=pooling_funcs,
pooling_kwargs=deepcopy(pooling_kwargs),
af=af, cv=cv, Creg=Creg,
prefix=prefix, savepath=savepath,
plot_uvx=plot_uvx, plot_man=plot_man,
net_idx=net_idx, tau0_idx=tau0_idx,
reset_dim=reset_dim)
for tau0_idx, tau0 in enumerate(tau0s)
for net_idx in range(Nnet)
)
# summarize
trn_homs, tst_homs, trn_hets, tst_hets = hu.collect_scores(
savepath=savepath, Ntau0=num_tau0,
Nnet=Nnet, Npool=len(pooling_funcs))
hu.plot_score_summary(
score_trn_hom = trn_homs,
score_trn_het = trn_hets,
score_tst_hom = tst_homs,
score_tst_het = tst_hets,
savepath=savepath, savename=prefix,
chance_score=1./Nclasses,
pooling_funcs=pooling_funcs,
)
[docs]
def run_spaced_tau0_multinet(
seed, Nnet, # random seed and number of MC realization
N, p, af, J, Ju, # nework parameters
h, h_is_multiplicative, # how to set heterogeneity's value
# tau_min, tau_max,
seq_len, dt, dim_ring, # task parameters
delay_period, # delay in decision making
stim_period, # stimulation duration
num_tau0, # number of distinct E[taus] to use
pooling_funcs, # pooling methods
golength, nsnapshot, # pooling args
n_jobs, cv, Creg, clf_name, # training
prefix, # simulation name
plot_uvx, plot_man, # plotting
sigma,
**kwargs
):
if prefix == '':
mult = 'M' if h_is_multiplicative else ''
prefix = f'SPACED_N{N}_MC{Nnet}_d{dim_ring}_h{h}{mult}'
# prefix+= f'_c{int(center)}_std{int(unit_std)}'
# prefix+= f'_o{int(orthgonalize)}_z{int(zero_phase)}_w{int(whiten)}'
print(prefix)
savepath = osjoin(SAVE_PATH, prefix)
Path(savepath).mkdir(parents=True, exist_ok=True)
ds, env = make_dataset(dt=dt, seq_len=seq_len,
dim_ring=dim_ring, seed=seed,
delay_period=delay_period,
stim_period=stim_period,
sigma=sigma,
) # creates random actions
trn, tst = get_train_test(ds)
us_trn, ys_trn, lens_trn = trn
us_tst, ys_tst, lens_tst = tst
Nclasses = dim_ring
reset_dim = 0
# h = hu.set_het_level(value=h, lens = lens_trn)
# mean taus values
spaced_tau = hu.get_spaced_hom_taus(lens_trn, dt=dt, num=num_tau0,
# tau_min=tau_min, tau_max=tau_max,
)
pooling_kwargs = dict(nsnapshot = nsnapshot, golength=golength, flatten=True)
memmap_temp = osjoin(savepath, 'joblib_memmap')
Path(memmap_temp).mkdir(parents=True, exist_ok=True)
Parallel(
n_jobs=n_jobs,
backend='loky',
batch_size=1,
pre_dispatch='n_jobs',
max_nbytes=0,
mmap_mode='r',
temp_folder=memmap_temp,
)(
delayed(hu.sample_integrate_classify_save)(
dt=dt, tau0=tau0,
h_value=h, h_is_multip = h_is_multiplicative,
N=N, p=p, J=J, Ju=Ju,
us_trn=us_trn, us_tst=us_tst,
lens_trn=lens_trn, lens_tst=lens_tst,
ys_trn=ys_trn, ys_tst=ys_tst,
pooling_funcs=pooling_funcs,
pooling_kwargs=deepcopy(pooling_kwargs),
af=af, cv=cv, Creg=Creg, clf_name=clf_name,
prefix=prefix, savepath=savepath,
plot_uvx=plot_uvx, plot_man=plot_man,
net_idx=net_idx, tau0_idx=tau0_idx,
reset_dim=reset_dim)
for tau0_idx, tau0 in enumerate(spaced_tau)
for net_idx in range(Nnet)
)
# summarize
trn_homs, tst_homs, trn_hets, tst_hets = hu.collect_scores(
savepath=savepath, Ntau0=num_tau0,
Nnet=Nnet, Npool=len(pooling_funcs))
hu.plot_score_summary_of_spaced_hom(trn_homs, trn_hets, tst_homs, tst_hets,
spaced_taus_hom = spaced_tau,
savepath=savepath, savename=prefix,
chance_score=1./Nclasses,
pooling_funcs=pooling_funcs,
Cregs=Creg,
)
het_taus = hu.collect_taus(savepath=savepath, Ntau0=len(spaced_tau), Nnet=Nnet)
hu.plot_taus_summary(het_taus=het_taus, hom_taus=spaced_tau, savepath=savepath)
# if __name__ == "__main__":
# # run_spaced_tau0_multinet(Nnet=3, N=10, num_tau0=2, seq_len=10000, prefix='test')
# # run_multinet(Nnet=3, N=10, seq_len=10000, prefix='test2')
# # run_multinet(n_jobs=6, Nnet=20, N=250)
# # run_multinet(n_jobs=6, Nnet=20, N=250, stim_period=[50, 10, 250], delay_period=100, prefix='variable_stim')
# # run_multinet(n_jobs=6, Nnet=5, N=20, stim_period=[50, 10, 250], delay_period=100, prefix='variable_stim_test')
# run_multinet(delay_period=[50, 10, 250], stim_period=100 , prefix='delay', )
# run_multinet(stim_period=[50, 10, 250], delay_period=100 , prefix='stim', )
# run_spaced_tau0_multinet(delay_period=[50, 10, 250], stim_period=100 , prefix='delay_multi_tau',)
# run_spaced_tau0_multinet(stim_period=[50, 10, 250], delay_period=100 , prefix='stim_multi_tau',)