2036151/scan_code/network_fitting.py at main · ModelDBRepository/2036151 · GitHub

History

337 lines (270 loc) · 14.7 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

from sNMO.optimizer import snmOptimizer

from brian2.units import *

import nevergrad as ng

import network_CADEX as cadex

import numpy as np

from spike_train_utils import build_isi_from_spike_train, plot_patterns

from sNMO.error.spikeTrainErrors import emd_pdist_spk, isi_swasserstein_2d, isi_wasserstein_dd, biemd

from matplotlib.pyplot import savefig

import time as time

import pandas as pd

from joblib import Parallel, delayed, dump, load

import sbi

import torch

import os

import pyabf

##synaptic params

pexc_input_conn = ng.p.Dict(Nexc=ng.p.TransitionChoice(np.array([10, 30, 50,100, 150, 200, 300, 400, 500]).astype(int)),

Exc_p = ng.p.Log(lower=0.001, upper=1))

default_input_conn = ng.p.Dict(Nexc=500, Exc_p=None) #

p_conn = np.array([0.01, 0.02, 0.04, 0.05, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5])

var_dict = ng.p.Dict(ee=ng.p.Scalar(lower=0.02, upper=15),

ii=ng.p.Scalar(lower=0.02, upper=15),

wcrh=ng.p.Scalar(lower=0.002, upper=15),

der= 1, #ng.p.Scalar(lower=0.01, upper=1),

de=ng.p.Scalar(lower=0.1, upper=20),

di=ng.p.Scalar(lower=0.1, upper=20),

dcrh=ng.p.Scalar(lower=10, upper=500),

input_hz=ng.p.Log(lower=10,upper=500),

exti_in=0,#ng.p.Log(lower=0.001,upper=500),

wexti=0,#ng.p.Log(lower=0.00002, upper=4),

##network params

p_ei = ng.p.TransitionChoice(p_conn.copy()),

p_ie = ng.p.TransitionChoice(p_conn.copy()),

##connectivity param

connectivity='random',

exc_input=default_input_conn, #ng.p.Choice([pexc_input_conn, default_input_conn]),

#ext_to_gaba = ng.p.TransitionChoice([True, False]),

i_pr = 0.05,

b_change=ng.p.Scalar(lower=0.75, upper=7),

)

OUTPUT_FOLDER = "/media/smestern/sgbackup/aoi_paper_2/output/"

rounds = 400

workers = 6

#Load the data to fit to

burst_ref = np.load("unit_3_burst_match1.npy")

non_burst_ref = np.load("unit_3_nonburst_match.npy")

#some error scale params

SNM_SF = 0.05

DT = 0.1*ms

RUN_TIME = 40

BASELINE_TIMES = [10, 15]

BASELINE_TIMES_END = [25, 40]

OPT_TONIC = False

def parse_params(params, all_fast=False, mono_synapse=False, threeway=False, slow_gaba=False, slow_crh=False, feedforward=False):

#here we will pop_out the params into their respective variables

#first bring the exc connectivity params up to the top level

params.update(params['exc_input'])

network_args = {'p_ei': params['p_ei'], 'p_ie': params['p_ie']}

synapse_params = {'taue': params['de']*ms,

'taui': params['di']*ms, 'taucrh': params['dcrh']*ms,

'wcrh': params['wcrh']*nS, 'we': params['ee']*nS, 'wi': params['ii']*nS, 'exte_in': params['input_hz']*Hz,

'exti_in': params['exti_in']*Hz,

'wexti': params['wexti']*nS, 'Nexc':params['Nexc'],

}

if threeway:

network_args.update({'p_ee': params['p_ee']})

synapse_params.update({'we_e': params['we_e']*nS})

if params['Exc_p'] is not None:

network_args['ECRH_params'] = {'p': params['Exc_p']}

else:

network_args['ECRH_params'] = None

if all_fast:

network_args.update(cadex.ALL_FAST_SYNAPSE_MODELS)

if mono_synapse:

network_args.update(cadex.MONO_EXP_SYNAPSE_MODELS)

if slow_gaba and not slow_crh:

network_args.update(cadex.SLOW_GABA_SYNAPSE_MODELS)

#network_args["ECRH_model"] = cadex.MONO_EXP_SYNAPSE_MODELS["ECRH_model"]

if slow_gaba and slow_crh:

network_args.update(cadex.SLOW_GABA_SYNAPSE_MODELS)

network_args["EI_model"] = cadex.DEFAULT_SYNAPSE_MODELS["EI_model"]

synapse_params['taup'] = 9999*second

synapse_params['taubr'] = 9999*second

#here we introduce

if params['connectivity'] != 'random' and params['connectivity'] != 'threeway':

network_args['connectivity_params'] = {'NUM_CLUSTERS': params['connectivity']['NUM_CLUSTERS'], 'inter_cluster_num': int(params['connectivity']['NUM_CLUSTERS']*params['connectivity']['inter_cluster_num'])+1}

network_args['connectivity_params'].update({'p_ei': params['p_ei'], 'p_ie': params['p_ie']})

network_args['connectivity'] = "custom_circ_graph"

else:

network_args['connectivity'] = params['connectivity']

if feedforward:

network_args['p_ei'] = 0.0

network_args['p_ie'] = 0.0

base_str = "eval: "

#base_str += f"poisson_exc.rates = poisson_exc.rates/Hz*{params['epsp_rate_switch']}*Hz; "

base_str += f"P[:{500}].pr = {params['i_pr']}; "

base_str += f" P[:{500}].p_e = {params['e_pr']}; " if 'e_pr' in params else ""

base_str += f"P[:{500}].p_exti = {params['exti_p']}; " if 'exti_p' in params else ""

base_str += f"CRH[:{500}].bd=(CRH.b[:{500}]/nS * {params['b_change']})*nS; CRH[:{500}].tauw = (CRH[:{500}].tauw/ms * {1.5})*ms;"

network_args['network_events'] = {30: base_str}

return network_args, synapse_params

def error_func(params):

network_args, synapse_params = parse_params(params)

network_events = network_args.pop('network_events')

#run the network and get the output

output = cadex.cadex_network(run_time=RUN_TIME, GABA_param_file=cadex.CRH_PARAM_FILE, network_args=network_args, synapse_params=synapse_params, network_events=network_events)

#also run the network in #tonic mode

if OPT_TONIC:

network_events= {5:'disconnect_GABA_increase_adapt',20: 'double_EPSP'}

synapse_params_dis = synapse_params.copy()

synapse_params_dis['taup'] = 9e9*second

synapse_params_dis['taubr'] = 9e9*second

output_tonic = cadex.cadex_network(run_time=RUN_TIME, GABA_param_file=cadex.CRH_PARAM_FILE, network_args=network_args, synapse_params=synapse_params_dis, network_events=network_events)

spikes = output['spikes']

volt = output['states']

#now we need to calculate the error

sim_isi_strt = build_isi_from_spike_train(spikes, low_cut=BASELINE_TIMES[0], high_cut=BASELINE_TIMES[1], indiv=True)[:500]

sim_isi_end = build_isi_from_spike_train(spikes, low_cut=BASELINE_TIMES_END[0], high_cut=BASELINE_TIMES_END[1], indiv=True)[:500]

#output_tonic

if OPT_TONIC:

spikes_tonic = output_tonic['spikes']

sim_isi_strt_tonic = build_isi_from_spike_train(spikes_tonic, low_cut=BASELINE_TIMES[0], high_cut=BASELINE_TIMES[1], indiv=True)[:500]

sim_isi_end_tonic = build_isi_from_spike_train(spikes_tonic, low_cut=BASELINE_TIMES_END[0], high_cut=BASELINE_TIMES_END[1],indiv=True)[:500]

else:

sim_isi_strt_tonic = None

sim_isi_end_tonic = None

return sim_isi_strt, sim_isi_end, spikes, volt, sim_isi_strt_tonic, sim_isi_end_tonic

def leakyrelu_error(x, target=0.0, thres=550, scale=0.1):

#this is an error function that replicates a leaky relu like function.

#here it measures the error between x, and target. but if x is below thres, it is scaled by scale

error = np.abs(x - target)

error = np.where(x < thres, error*scale, error)

return error

def ng_error_func(params):

sim_isi_strt, sim_isi_end, spikes, volt, sim_isi_strt_tonic, sim_isi_end_tonic = error_func(params)

return _ng_error_func(params, sim_isi_strt, sim_isi_end, spikes, volt, sim_isi_strt_tonic, sim_isi_end_tonic)

def _ng_error_func(params, sim_isi_strt, sim_isi_end, spikes, volt, sim_isi_strt_tonic, sim_isi_end_tonic, parse_params_func=parse_params):

#params in this case is an array of the parameters

run_id = np.random.randint(0,1000000) + time.time()

#iter through and compute the error for each unit

burst_error = []

nonburst_error = []

burst_pre_fr = []

nonburst_pre_fr = []

burst_post_fr = []

nonburst_post_fr = []

for unit in np.arange(500):

if len(sim_isi_strt[unit]) > 2:

burst_error.append(isi_wasserstein_dd(sim_isi_strt[unit], burst_ref))

nonburst_error.append(isi_wasserstein_dd(sim_isi_strt[unit], non_burst_ref))

burst_pre_fr.append(len(sim_isi_strt[unit])/np.diff(BASELINE_TIMES)[0])

burst_post_fr.append(len(sim_isi_end[unit])/np.diff(BASELINE_TIMES_END)[0])

if OPT_TONIC:

nonburst_pre_fr.append(len(sim_isi_strt_tonic[unit])/np.diff(BASELINE_TIMES)[0])

nonburst_post_fr.append(len(sim_isi_end_tonic[unit])/np.diff(BASELINE_TIMES_END)[0])

else:

nonburst_pre_fr.append(0)

nonburst_post_fr.append(0)

else:

burst_error.append(999)

nonburst_error.append(999)

burst_pre_fr.append(0)

nonburst_pre_fr.append(0)

burst_post_fr.append(0)

nonburst_post_fr.append(0)

#compute out the burst errors

burst_error = np.array(burst_error)

nonburst_error = np.array(nonburst_error)

burst_error = burst_error[~np.isnan(burst_error)]

nonburst_error = nonburst_error[~np.isnan(nonburst_error)]

burst_error_full = np.copy(burst_error)

nonburst_error_full = np.copy(nonburst_error)

#get the mean of the lowest 10% of the errors

burst_error = np.mean(np.sort(burst_error)[:int(len(burst_error)*0.1)])

#in this case we also want to see if any units are nonbursting, we want heterogeneity. but this error should be smaller

nonburst_error = np.mean(np.sort(nonburst_error)[:int(len(nonburst_error)*0.1)]) / 50 if OPT_TONIC else 0

#examine the mean pre-post for burst and nonburst

burst_pre_fr_mean = np.mean(burst_pre_fr)

nonburst_pre_fr_mean = np.mean(nonburst_pre_fr)

burst_post_fr_mean = np.mean(burst_post_fr)

nonburst_post_fr_mean = np.mean(nonburst_post_fr)

burst_pre_max = np.max(burst_pre_fr)

nonburst_pre_max = np.max(nonburst_pre_fr)

burst_post_max = np.max(burst_post_fr)

nonburst_post_max = np.max(nonburst_post_fr)

pre_post_dict = {'burst_pre_fr_mean': burst_pre_fr_mean,'nonburst_pre_fr_mean': nonburst_pre_fr_mean,'burst_post_fr_mean': burst_post_fr_mean,'nonburst_post_fr_mean': nonburst_post_fr_mean,'burst_pre_max': burst_pre_max,'nonburst_pre_max': nonburst_pre_max,'burst_post_max': burst_post_max,'nonburst_post_max': nonburst_post_max}

#get the idx of the burst_pre_fr

burst_pre_fr_idx = np.argsort(burst_pre_fr)

burst_pre_fr_idx = burst_pre_fr_idx[burst_pre_fr_idx <=10]

# c22_error = pycatch22.catch22_all(volt.d_I[burst_pre_fr_idx[0], :]/pA )

# c22_error = c22_scaler.transform(np.array(c22_error['values']).reshape(1,-1))

# c22_error = np.mean(np.abs(c22_error - c22_feat_baseline))

c22_error = 0

#some other error metrics, not used for the optimization

mean_volt = np.mean(volt.v[:500, :]/mV)

mean_I = np.std(volt.d_I[:500, :]/pA)

#SNM error, this is a bit more complicated

#if there is no voltage, or no current, then we will just set the error to 999

if np.isnan(mean_volt) or np.isnan(mean_I):

return 999

else:

#ensure volt is below -50mV most of the time

volt_error = np.mean(np.where(volt.v[:500, :]/mV > -10, 1, 0))

#ensure the peak positive current is low

#measure how long the current is above 250pA

current_error = np.mean(np.where(volt.d_I[:500, :]/pA > 400, 1, 0)) * ( RUN_TIME)

#peak negative is a bit more complicated, we need to find the peak negative current, however this is skewed by the action potentials

#so we need to find the peak negative current between spikes

#current_error = leakyrelu_error(np.amax(volt.d_I[:500, :]/pA), target=0, thres=550, scale=0.1)

#just clip to 0

#current_error = np.clip(current_error, 0, None)/550 #divide by 250 to normalize

snm_error = (volt_error + current_error) * SNM_SF

#plot and save

error_id = str(np.round(burst_error + nonburst_error + snm_error,3)).ljust(5, '0')

plot_patterns(spikes.spike_trains(), 500);savefig(os.path.join(OUTPUT_FOLDER, f"spikes_{error_id}_{int(run_id)}.png"))

#save the params

params_out = {}

network_args, synapse_params = parse_params_func(params)

params_out.update(network_args)

params_out.update(synapse_params)

params_out.update(network_args['ECRH_params']) if network_args['ECRH_params'] is not None else None

params_out['burst_error'] = burst_error

params_out['nonburst_error'] = nonburst_error

params_out['mean_volt'] = mean_volt

params_out['mean_I'] = mean_I

params_out['snm_error'] = snm_error

params_out['runtime'] = RUN_TIME

params_out['max_D_I'] = np.amax(volt.d_I[:500, :]/pA)

params_out.update(pre_post_dict)

params_df = pd.DataFrame(params_out, index=[0])

params_df['connectivity'] = str(params['connectivity'])

params_df.to_csv(os.path.join(OUTPUT_FOLDER, f"params_{error_id}_{int(run_id)}.csv"))

#save the spikes and volts with heavy compression

states = volt.get_states(['v', 'd_I', 'ge', 'gi'])

for key in states:

states[key] = states[key][:, :10]

#add the burst and nonburst error to the states

states['burst_error'] = burst_error_full

states['nonburst_error'] = nonburst_error_full

#add in the pre-post firing rates

states['burst_pre_fr'] = burst_pre_fr

states['nonburst_pre_fr'] = nonburst_pre_fr

states['burst_post_fr'] = burst_post_fr

states['nonburst_post_fr'] = nonburst_post_fr

#add in the c22 error

states['c22_error'] = c22_error

dump([spikes.spike_trains(), states], os.path.join(OUTPUT_FOLDER, f"spikes_{error_id}_{int(run_id)}.joblib"), compress=('lzma', 9))

#also dump the exact params for later use

dump({'network_args': network_args, 'synapse_params': synapse_params, 'network_events': {}}, os.path.join(OUTPUT_FOLDER, f"params_{error_id}_{int(run_id)}.joblib"),

compress=('lzma', 9))

return burst_error + nonburst_error + snm_error + c22_error

def nevergrad():

optimizer = ng.optimizers.ScrHammersleySearch(parametrization=var_dict, budget=int(workers*rounds), num_workers=workers, )

optimizer.enable_pickling()

#optimizer = ng.optimizers.ParaPortfolio.load(f"optimizer.joblib")

for _ in np.arange((optimizer.budget - optimizer.num_ask) // workers):

points_list = []

for x in np.arange(workers):

while True:

try:

points_list.append(optimizer.ask())#

break

except:

print("failed to ask")

error = Parallel(n_jobs=workers)(delayed(ng_error_func)(p.value) for p in points_list)

for points, er in zip(points_list, error):

optimizer.tell(points, er)

optimizer.dump(f"optimizer.joblib")

if __name__ == '__main__':

nevergrad()

2036151/scan_code/network_fitting.py at main · ModelDBRepository/2036151 · GitHub

Uh oh!

Expand file tree

Breadcrumbs

network_fitting.py

Latest commit

History

Breadcrumbs

network_fitting.py

File metadata and controls