from copy import deepcopy
import argparse
import numpy as np
# import argparse
from time import time, ctime
# from hetero.parser import get_parser, create_param_dicts
from hetero.sim import Simulation
from hetero.config import CFG, VALID_ACTIVATIONS
[docs]
def get_options(arg):
return list(CFG[arg].keys())
[docs]
def get_parser():
"""
Creates a parser that is able to capture all variables necessary to construct the paramter dictionaries.
"""
parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
io = parser.add_argument_group("Input/ouput options")
het = parser.add_argument_group("Heterogeneity options")
net = parser.add_argument_group("Network options")
dyn = parser.add_argument_group("Dynamics options")
stim = parser.add_argument_group("Stimulation options")
task = parser.add_argument_group("Task options")
ro = parser.add_argument_group("Readout options")
optim = parser.add_argument_group("Optimization options")
data = parser.add_argument_group("Dataset options")
uq = parser.add_argument_group("Uncertainty quantification options")
io.add_argument('--root', type=str, default=CFG['root'],
help='Root directory to store results')
io.add_argument('--stype', type=str, default=CFG['stype'], help=(
"Data type for saving the network state. Set to float64 \n"
"for accurate forecasting benchmark, otherwise leave the \n"
"default (float16).""")
)
io.add_argument('--ncore', type=int, default=-1, help='Number of CPUs to use. -1 means use all.')
io.add_argument('--noint', action='store_true', default=False, help='Skip integration. Default False.')
io.add_argument('--nobm', action='store_true', default=False, help='Skip benchmarking. Default False.')
io.add_argument('--nosvd', action='store_true', default=False, help='Skip SVD calculation (needed for singular values and angles). Default False.')
io.add_argument('--noconn', action='store_true', default=False, help='Skip the connectivity plot. Default False.')
io.add_argument('--noscore', action='store_true', default=False, help='Skip the score vs. homogeneous score plot. Default False.')
io.add_argument('--noangle', action='store_true', default=False, help='Skip the angle-vs-score summary plot. Default False.')
io.add_argument('--nosumplot', action='store_true', default=False, help=(
"Skip the summary figures produced once all realizations are \n"
"benchmarked: the angle CDF and the singular-value/effective-rank \n"
"ones. Both aggregate every state on disk, so in a parameter sweep \n"
"they are re-rendered from all the previously computed cases after \n"
"each single run; use this to sweep first and plot once at the end. \n"
"Default False."))
io.add_argument('--notau', action='store_true', default=False, help='Skip the tau distribution plot. Default False.')
io.add_argument('--nouvx', action='store_true', default=False, help='Skip the u-v-x plot. Default False.')
io.add_argument('--noviz', action='store_true', default=False, help=(
"Skip every figure of the run: the complexity/udf maps, the connectivity, tau\n"
"and uvx traces, the readout scores, the angle-vs-score figure and the\n"
"angle-CDF / singular-value / effective-rank summaries. Implies --noconn,\n"
"--notau, --nouvx, --noscore, --noangle and --nosumplot. The svals/erank/angle\n"
"csv files are still computed (only --nosvd skips those). The figures of a\n"
"finished run are re-rendered with --replot, which overrides this flag.\n"
"Default False."))
io.add_argument('--replot', action='store_true', default=False, help=(
"Replot all visualizations of the simulation (connectivity, tau,\n"
"uvx, scores, angles, angle CDF, complexity, udf). Results needed by\n"
"each figure are loaded from the disk when already computed, and\n"
"computed (synthesis / integration / benchmark) only when missing.\n"
"Overrides the --noviz/--no* plot-skip flags. Default False.")
)
io.add_argument('--seed', type=int, default=CFG['seed'], help='Random number seed.')
# Heterogeneiy
het.add_argument('--distro', type=str, default=CFG['distro'], help='Heterogeniety distribution (e.g., lognormal, gamma, uniform)')
het.add_argument('--n_means', type=int, default=CFG['n_means'], help='Number of mean timescales')
het.add_argument('--n_vars', type=int, default=CFG['n_vars'], help='Number of heterogeneity levels')
het.add_argument('--log_distance', type=float, default=CFG['log_distance'], help='Log-10 distance between mean levels')
# network
net.add_argument('--N', type=int, default=CFG['N'], help='Number of neurons in the network')
net.add_argument('--p', type=float, default=CFG['p'], help='Sparsity of the connectivity matrix')
net.add_argument('--f', type=float, default=CFG['f'], help='Fraction of excitatory neurons')
net.add_argument('--mue', type=float, default=CFG['mue'], help='Mean excitatory weight')
net.add_argument('--sig0', type=float, default=CFG['sig0'], help='Excitatory and inhibitory weight dispersion (std)')
net.add_argument('--autapse', type=bool, default=CFG['autapse'], help='Are autapses allowed?')
net.add_argument('--delay', type=int, default=CFG['delay'], help='Synaptic delay (in unit of dt)')
# gain
net.add_argument('--Jn', type=float, default=CFG['Jn'], help='Gain of the white noise')
net.add_argument('--J', type=float, default=CFG['J'], help='Gain of the recurrent weights')
net.add_argument('--Ju', type=float, default=CFG['Ju'], help='Gain of the feedforward weights')
# stim
stim.add_argument('--stim', action='extend', nargs='+', default=[], help=(
"The label of the inputs to the network. More than one label \n"
"can be given to concatenate different timeseries.")
)
# task
task.add_argument('--task', type=str, choices=get_options('task'), default='tayloramus',
help="The task network is supposed to perform.")
# dataset
data.add_argument('--n_trn', type=float, default=CFG['n_trn'],
help=("Length of the training set (in unit of the input base timescale). For example,\n"
"for a sine input with period of 5 timesteps, n_trn=2 means a training set\n"
"containing 10 samples.")
)
data.add_argument('--n_tst', type=float, default=CFG['n_tst'],
help=("Length of the test set (in unit of the input base timescale). For example,\n"
"for a sine input with period of 5 timesteps, n_tst=2 means a test set\n"
"containing 10 samples.")
)
data.add_argument('--disable_scale_by_size', action='store_true', default= not CFG['scale_by_size'],
help=("Whether to scale the training set size by the network size. Default is True.\n"
"When true, enlarging the network by a factor of x expands the size of training\n"
"set by the same factor. Test set is unaffected.")
)
dyn.add_argument('--dyn', type=str, choices=get_options('dyn'), default='LI', help='The name of network dynamic.')
dyn.add_argument('--solver', type=str, default=CFG['solver'], help='Solver for integration')
# readout
ro.add_argument('--readout_activation', type=str, choices=get_options('readout_activation'), default='possig1',
help='The activation function used for the readout.')
# optimization
optim.add_argument('--alg', type=str, default=CFG['alg'],
help='Optimization algorithm, Ridge or ordinary least-squares (direct)')
optim.add_argument('--reg', type=float, default=CFG['reg'],
help='Optimization regularization value (relevant if needed by the algorithm)')
optim.add_argument('--fit_bias', type=bool, default=CFG['fit_bias'], help='Whether to fit a bias term? Default True.')
uq.add_argument('--n_trials', type=int, default=CFG['n_trials'], help=(
"The number of non-overlapping contiguous intervals over which a readout is trained.\n"
"These test performance corresponding to these independent readouts can be used to \n"
"estimate performance uncertainty.")
)
uq.add_argument('--n_realizations', type=int, default=CFG['n_realizations'], help=(
"The number of independent network (topology and weights) realizations. These score variation\n"
"acorss realizations is used as an estimate of performance uncertainty.")
)
return parser
# Every flag of the parser that only suppresses a *figure*, i.e. the flags
# ``--noviz`` sets at once (see ``main``). This is the single source of truth:
# a new plot-skip flag must be added here too, otherwise it silently escapes
# ``--noviz`` and gets re-rendered on every run of a parameter sweep.
# ``--nosvd`` is deliberately NOT in this list: it suppresses the *computation*
# of the svals/erank/angle csvs, not a figure, so ``--noviz`` leaves the
# numbers on disk intact. Likewise ``--noint``/--nobm`` are not plot flags.
PLOT_SKIP_FLAGS = ('noconn', 'notau', 'nouvx', 'noscore', 'noangle', 'nosumplot')
[docs]
def create_param_dicts(args, verbose=True):
"""
Creates proper dictionaries from args
"""
hetero_params = {
'distro': args.distro,
'n_means': args.n_means,
'n_vars': args.n_vars,
'log_distance': args.log_distance,
}
net_params = {
'N': args.N,
'p': args.p,
'f': args.f,
'mue': args.mue,
'sig0': args.sig0,
'autapse': args.autapse,
'topology': 'random',
'delay': args.delay
}
gain_params = {
'Jn': args.Jn,
'J': args.J,
'Ju': args.Ju
}
task_params = {
args.task : deepcopy(CFG['task'][args.task])
}
stim_params = {}
for stim in args.stim:
stim_params[stim] = deepcopy(CFG['stim'][stim])
dataset_params = {
'n_trn': args.n_trn,
'n_tst': args.n_tst,
'scale_by_size': not args.disable_scale_by_size,
}
dyn_params = {
args.dyn : deepcopy(CFG['dyn'][args.dyn])
}
# ensure that the readout activation and the dynamic are compatible
# LIF --> covolution
# LI --> activaition function
if args.readout_activation not in VALID_ACTIVATIONS[args.dyn]:
if args.dyn =='LI':
print(f'Warning: The readout of a rate neuron shall not be done via {args.readout_activation}. Switched back to possig1.')
args.readout_activation = 'possig1'
else:
print(f'Warning: The readout of a spiking neuron shall not be done {args.readout_activation} activation function. Switched back to d10.')
args.readout_activation = 'd10'
readout_params = {
'readout_activation': {
args.readout_activation: deepcopy(CFG['readout_activation'][args.readout_activation])
}
}
optim_params = {
'alg': args.alg,
'reg': args.reg,
'fit_bias': args.fit_bias,
}
if np.isscalar(args.n_trials):
args.n_trials = [args.n_trials]
uq_params = {
'n_trials': args.n_trials,
'n_realizations': args.n_realizations,
}
params = {
'heterog': hetero_params,
'net': net_params,
'set': dataset_params,
'gain': gain_params,
'readout': readout_params,
'task': task_params,
'dyn': dyn_params,
'stim': stim_params,
'optim': optim_params,
'uq': uq_params,
'seed': args.seed,
}
return params
[docs]
def main():
parser = get_parser()
args = parser.parse_args()
params = create_param_dicts(args)
# --noviz turns off *every* figure: the complexity/udf maps of the init
# plot, the connectivity, tau and uvx traces, the readout scores, the
# angle-vs-score figure and the angle-CDF / svals-erank summaries.
# The svals/erank/angle csvs are still computed (only --nosvd skips those).
if args.noviz:
for _flag in PLOT_SKIP_FLAGS:
setattr(args, _flag, True)
if args.replot:
if args.noviz or any(getattr(args, _f) for _f in PLOT_SKIP_FLAGS):
print('Note: --replot overrides the --noviz/--no* plot-skip flags; all figures will be re-rendered.')
if args.noint or args.nobm:
print('Note: --replot recomputes missing results (integration/benchmark) even when --noint/--nobm are set.')
print(f'\nSimulation parameters:')
print(f'- net: \t {params["net"]}')
print(f'- heterog: \t {params["heterog"]}')
print(f'- task: \t {params["task"]}')
print(f'- stim: \t {params["stim"]}')
print(f'- dataset: \t {params["set"]}')
print(f'- dyn: \t {params["dyn"]}')
print(f'- gains: \t {params["gain"]}')
print(f'- readout: \t {params["readout"]}')
print(f'- optimization:\t {params["optim"]}')
print(f'- UQ: \t {params["uq"]}')
print(f'- random seed:\t {args.seed}')
print('\n')
# rng = np.random.default_rng(args.seed)
dyn, dyn_params = next(iter(params['dyn'].items()))
tic = time()
print(f'Run started at {ctime()}')
# Initialize Simulation object once
sim = Simulation(
net_params=params['net'],
task_params=params['task'],
stim_params=params['stim'],
heterog_params=params['heterog'],
set_params=params['set'],
uq_params=params['uq'],
root=args.root,
save_type=args.stype,
seed=args.seed,
skip_init_plot=args.noviz,
)
# Call integrate_and_benchmark to handle all realizations
sim.integrate_and_benchmark(
n_realizations=args.n_realizations,
skip_integration=args.noint,
skip_benchmark=args.nobm,
dyn=dyn,
dyn_params=dyn_params,
gain_params=params['gain'],
readout_params=params['readout'],
optim_params=params['optim'],
solver=args.solver,
ncpu=args.ncore,
skip_svd=args.nosvd,
skip_connectivity_plot=args.noconn,
skip_score_plot=args.noscore,
skip_tau_plot=args.notau,
skip_uvx_plot=args.nouvx,
skip_angle_plot=args.noangle,
skip_summary_plot=args.nosumplot,
replot=args.replot,
)
print(f'Run ended at {ctime()}.')
print(f'Total runtime: {time()-tic:.1f} seconds.')
import sys
if __name__ == '__main__':
# try:
# njobs = 6
# sys.argv = [
# "hetero", # argv[0] is the script name
# "--N", "8",
# "--stim", "lorenz",
# "--dyn", "LI",
# "--ncore", str(njobs),
# "--n_realizations", "2",
# "--n_trials", "1"
# ]
# main()
# except Exception as e:
# print(e)
# for N in [50, 250]:
# for dyn in ['LI']:
# try:
# njobs = -1 #if N<500 else 4
# sys.argv = [
# "hetero", # argv[0] is the script name
# "--N", str(N),
# "--stim", "lorenz",
# "--dyn", f"{dyn}",
# "--root", f"/home/arash/Desktop/umg-codes/hetero/results/uq/{dyn}/",
# "--ncore", str(njobs),
# "--n_realizations", "10",
# "--n_trials", "1",
# "--n_trn", str(250 * 20), # 20 * N (but with a fixed N to coerce fixed horizon)
# "--no_scale_by_size",
# ]
# main()
# except Exception as e:
# print(e)
# for J in [0, .25, .5, .75, 1, 2, 5, 10]:
# try:
# sys.argv = [
# "run.py", # argv[0] is the script name
# "--J", str(J),
# "--stim", "lorenz",
# "--dyn", "LI",
# "--readout_activation", "tanh1",
# "--root", "/home/arash/Desktop/umg-code/hetero/tanh/"
# ]
# main()
# except Exception as e:
# print(e)
# for Ju in [0, .25, .5, .75, 1, 2, 5, 10]:
# try:
# sys.argv = [
# "run.py", # argv[0] is the script name
# "--Ju", str(Ju),
# "--stim", "lorenz",
# "--dyn", "LI",
# "--readout_activation", "tanh1",
# "--root", "/home/arash/Desktop/umg-code/hetero/tanh/"
# ]
# main()
# except Exception as e:
# print(e)
# main()
try:
main()
except Exception as e:
print(e)