maxidx = fr_roll.argmax()

log_range = np.logspace(0, np.log10(spike_time.size - maxidx),

n_sample).astype(np.int) - 1

sample_range = maxidx + log_range

time, stress, rec_spike_time, rec_fr_roll,

**kwargs):

sample_spike_time, sample_fr_roll = get_log_sample_after_peak(

rec_spike_time, rec_fr_roll, 50)

mod_spike_time, mod_fr_inst = get_mod_spike(lmpars, groups, time, stress)

mod_fr_inst_interp = np.interp(sample_spike_time,

mod_spike_time, mod_fr_inst)

residual = mod_fr_inst_interp - sample_fr_roll

print((residual**2).sum())

params = lmpars_to_params(lmpars)

mod_spike_time, mod_fr_inst = stress_to_fr_inst(time, stress,

groups, **params)

lmpars_dict = lmpars.valuesdict()

# Export parameters to separate dicts and use indices as keys

separate_dict = {'tau': {}, 'k': {}}

for var, val in lmpars_dict.items():

for param in separate_dict.keys():

if param in var:

index = int(var.split(param)[-1])

separate_dict[param][index] = val

# Convert to final format of parameter dict

params = {}

for param, indexed_dict in separate_dict.items():

params[param + '_arr'] = np.array(

np.array([indexed_dict[index]

for index in sorted(indexed_dict.keys())]))

def get_fname(animal, datatype):

return os.path.join('data', 'rec', '%s_%s.csv' % (animal, datatype))

fname_dict = {datatype: get_fname(animal, datatype)

for datatype in ['spike', 'displ']}

displ_arr = np.genfromtxt(fname_dict['displ'], delimiter=',')

static_displ_list = np.round(displ_arr[-1], 2).tolist()

spike_arr = np.genfromtxt(fname_dict['spike'], delimiter=',')

spike_time_list = [spike.nonzero()[0] / FS for spike in spike_arr.T]

fr_inst_list = [spike_time_to_fr_inst(spike_time)

for spike_time in spike_time_list]

fr_roll_list, max_time_list, max_fr_roll_list = [], [], []

for spike_time in spike_time_list:

fr_roll = spike_time_to_fr_roll(spike_time, WINDOW)

fr_roll_list.append(fr_roll)

max_time_list.append(spike_time[fr_roll.argmax()])

max_fr_roll_list.append(fr_roll.max())

rec_dict = {

'static_displ_list': static_displ_list,

'spike_time_list': spike_time_list,

'fr_inst_list': fr_inst_list,

'fr_roll_list': fr_roll_list,

'max_time_list': max_time_list,

'max_fr_roll_list': max_fr_roll_list}

"""

Stretch the ramp phase of the stress such that the ramp time matches the

`max_time_target`. The hold phase of the stress is unchanged.

Parameters

----------

time : 1xN array

stress : 1xN array

max_time_spike : float

When does the peak firing happens for the recording

stretch_coeff : float

The actual `max_time_target` is stretched by this coeff.

Returns

-------

stress_new : 1XN array

"""

# Clean the zeros in the beginning of the stress trace, if any

start_idx = (stress == 0).nonzero()[0][-1]

time = time[start_idx:] - time[start_idx]

stress = stress[start_idx:]

# Do the stretch

max_idx = stress.argmax()

max_time_original = time[max_idx]

time_new = time.copy()

max_time_target = max_time_spike * stretch_coeff

time_new[:max_idx + 1] *= max_time_target / max_time_original

time_new[max_idx + 1:] += max_time_target - max_time_original

stress_new = np.interp(time, time_new, stress)

if rec_dict is None:

rec_dict = load_rec(animal)

# Read recording data

rec_fr_inst = rec_dict['fr_inst_list'][stim]

rec_spike_time = rec_dict['spike_time_list'][stim]

rec_fr_roll = rec_dict['fr_roll_list'][stim]

static_displ = rec_dict['static_displ_list'][stim]

rec_data_dict = {

'rec_spike_time': rec_spike_time,

'rec_fr_inst': rec_fr_inst,

'rec_fr_roll': rec_fr_roll}

# Read model data

time, stress = get_interp_stress(static_displ)

max_time = rec_dict['max_time_list'][stim]

# stretch_coeff = 1 + 0.25 * static_displ / REF_DISPL

stretch_coeff = 1 + 0.4 * static_displ / REF_DISPL

stress = adjust_stress_ramp_time(time, stress, max_time, stretch_coeff)

mod_data_dict = {

'groups': MC_GROUPS,

'time': time,

'stress': stress}

fit_data_dict = dict(list(rec_data_dict.items()) +

list(mod_data_dict.items()))

data_dicts = {

'rec_data_dict': rec_data_dict,

'mod_data_dict': mod_data_dict,

'fit_data_dict': fit_data_dict}

rec_spike_time, plot_kws={}, roll=True,

plot_rec=True, plot_mod=True,

fig=None, axs=None, save_data=False, fname=None,

**kwargs):

if roll:

rec_fr = kwargs['rec_fr_roll']

else:

rec_fr = kwargs['rec_fr_inst']

if fig is None and axs is None:

fig, axs = plt.subplots()

axs0 = axs

axs1 = axs

elif isinstance(axs, np.ndarray):

axs0 = axs[0]

axs1 = axs[1]

else:

axs0 = axs

axs1 = axs

if plot_mod:

mod_spike_time, mod_fr_inst = get_mod_spike(lmpars_fit, groups,

time, stress)

axs0.plot(mod_spike_time, mod_fr_inst * 1e3, '-', **plot_kws)

axs0.set_xlim(0, 5000)

axs0.set_xlabel('Time (msec)')

axs0.set_ylabel('Instantaneous firing (Hz)')

if save_data:

np.savetxt('./data/%s_mod_spike_time.csv' % fname,

mod_spike_time, delimiter=',')

np.savetxt('./data/%s_mod_fr_inst.csv' % fname,

mod_fr_inst * 1e3, delimiter=',')

if plot_rec:

axs1.plot(rec_spike_time, rec_fr * 1e3, '.', **plot_kws)

axs1.set_xlim(0, 5000)

axs1.set_xlabel('Time (msec)')

axs1.set_ylabel('Instantaneous firing (Hz)')

if save_data:

np.savetxt('./data/%s_rec_spike_time.csv' % fname,

rec_spike_time, delimiter=',')

np.savetxt('./data/%s_rec_fr_inst.csv' % fname,

rec_fr * 1e3, delimiter=',')

fig.tight_layout()

if isinstance(lmpar_list, Parameters):

return lmpar_list

lmpar_dict_list = [lmpar.valuesdict() for lmpar in lmpar_list]

all_param_dict = defaultdict(list)

for lmpar_dict in lmpar_dict_list:

for key, value in lmpar_dict.items():

all_param_dict[key].append(value)

mean_lmpar = copy.deepcopy(lmpar_list[0])

for key, value in all_param_dict.items():

mean_lmpar[key].set(value=np.mean(value))

lmpars_cko = copy.deepcopy(lmpars)

for k, scale in k_scale_dict.items():

lmpars_cko[k].value *= scale

"""

Generate the copy-pasteable table for paper writing

"""

params_ser = pd.Series(lmpars.valuesdict())

params_paper = pd.Series()

params_paper['tau1'] = params_ser['tau1']

params_paper['tau2'] = params_ser['tau2']

if 'tau4' in params_ser.keys():

params_paper['tau3'] = params_ser['tau3']

params_paper['knr'] = params_ser['k1']

if 'tau4' in params_ser.keys():

params_paper['kmc'] = params_ser['k2'] + params_ser['k4']

params_paper['kmc1'] = params_ser['k2'] / params_paper['kmc']

params_paper['kmc2'] = params_ser['k4'] / params_paper['kmc']

params_paper['kusa'] = params_ser['k3']

else:

params_paper['kmc'] = params_ser['k2'] + params_ser['k3']

params_paper['kmc1'] = params_ser['k2'] / params_paper['kmc']

params_paper['kmc2'] = params_ser['k3'] / params_paper['kmc']

"""

If the parameters are stored in `data/fit`, then will load from file;

Otherwise, fitting will be performed.

"""

def __init__(self, lmpars_init, label=None):

self.lmpars_init = lmpars_init

if label is None:

self.label = get_time_stamp()

else:

self.label = label

# Load data

self.load_rec_dicts()

self.load_data_dicts_dicts()

self.get_ref_fit()

def get_ref_fit(self):

pname = os.path.join('data', 'fit', self.label)

if os.path.exists(pname):

with open(os.path.join(pname, 'ref_mean_lmpars.pkl'), 'rb') as f:

self.ref_mean_lmpars = pickle.load(f)

self.ref_result_list = []

for fname in os.listdir(pname):

if fname.startswith('ref_fit') and fname.endswith('.pkl'):

with open(os.path.join(pname, fname), 'rb') as f:

self.ref_result_list.append(pickle.load(f))

else:

self.fit_ref()

def load_rec_dicts(self):

self.rec_dicts = {animal: load_rec(animal) for animal in ANIMAL_LIST}

def get_data_dicts(self, animal, stim):

data_dicts = get_data_dicts(stim, rec_dict=self.rec_dicts[animal])

return data_dicts

def load_data_dicts_dicts(self):

self.data_dicts_dicts = {}

for animal in ANIMAL_LIST:

self.data_dicts_dicts[animal] = {}

for stim in range(STIM_NUM):

self.data_dicts_dicts[animal][stim] = self.get_data_dicts(

animal, stim)

def fit_ref(self, export=True):

data_dicts_dict = self.data_dicts_dicts[REF_ANIMAL]

# Prepare data for multiprocessing

fit_mp_list = []

for stim in REF_STIM_LIST:

fit_mp_list.append([self.lmpars_init,

data_dicts_dict[stim]['fit_data_dict']])

with Pool(5) as p:

self.ref_result_list = p.map(fit_single_rec_mp, fit_mp_list)

lmpar_list = [result.params for result in self.ref_result_list]

self.ref_mean_lmpars = get_mean_lmpar(lmpar_list)

# Plot the fit for multiple displacements

if export:

self.export_ref_fit()

def export_ref_fit(self):

pname = os.path.join('data', 'fit', self.label)

os.mkdir(pname)

for stim, result in zip(REF_STIM_LIST, self.ref_result_list):

fname_report = 'ref_fit_%d.txt' % stim

fname_pickle = 'ref_fit_%d.pkl' % stim

with open(os.path.join(pname, fname_report), 'w') as f:

f.write(fit_report(result))

with open(os.path.join(pname, fname_pickle), 'wb') as f:

pickle.dump(result, f)

with open(os.path.join(pname, 'ref_mean_lmpars.pkl'), 'wb') as f:

pickle.dump(self.ref_mean_lmpars, f)

# Plot

fig, axs = self.plot_ref_fit(roll=True)

fig.savefig(os.path.join(pname, 'ref_fit_roll.png'), dpi=300)

plt.close(fig)

fig, axs = self.plot_ref_fit(roll=False)

fig.savefig(os.path.join(pname, 'ref_fit_inst.png'), dpi=300)

plt.close(fig)

def plot_ref_fit(self, roll=True):

fig, axs = plt.subplots(2, 1, figsize=(3.5, 6))

for stim, ref_result in zip(REF_STIM_LIST, self.ref_result_list):

lmpars_fit = ref_result.params

color = COLOR_LIST[stim]

plot_single_fit(

lmpars_fit, fig=fig, axs=axs[0], roll=roll,

plot_kws={'color': color},

**self.data_dicts_dicts[REF_ANIMAL][stim]['fit_data_dict'])

plot_single_fit(

self.ref_mean_lmpars, fig=fig, axs=axs[1], roll=roll,

plot_kws={'color': color},

**self.data_dicts_dicts[REF_ANIMAL][stim]['fit_data_dict'])

axs[0].set_title('Individual fitting parameters')

axs[1].set_title('Using the average fitting parameter')

fig.tight_layout()

return fig, axs

def plot_cko_customized(self, k_scale_dict,

animal_rec=None, animal_mod=None,

fig=None, axs=None,

close_fig=False, save_fig=False, show_label=False,

save_data=False, fname=''):

lmpars_cko = copy.deepcopy(self.ref_mean_lmpars)

for k, scale in k_scale_dict.items():

lmpars_cko[k].value *= scale

label = str(k_scale_dict).translate({ord(c): None for c in '{}\': .,'})

if fig is None and axs is None:

fig, axs = plt.subplots()

close_fig = True

save_fig = True

for stim in REF_STIM_LIST:

color = COLOR_LIST[stim]

if animal_rec is not None:

plot_single_fit(

lmpars_cko, fig=fig, axs=axs, roll=False,

plot_rec=True, plot_mod=False,

plot_kws={'color': color}, save_data=save_data,

fname='%s_stim_%s' % (fname, {0: 'high', 2: 'low'}[stim]),

**self.data_dicts_dicts[animal_rec][stim]['fit_data_dict'])

if animal_mod is not None:

plot_single_fit(

lmpars_cko, fig=fig, axs=axs, roll=False,

plot_rec=False, plot_mod=True, save_data=save_data,

fname='%s_stim_%s' % (fname, {0: 'high', 2: 'low'}[stim]),

plot_kws={'color': color},

**self.data_dicts_dicts[animal_mod][stim]['fit_data_dict'])

if show_label:

axs.set_title('Method: %s Rec: %s Mod: %s' %

(label, animal_rec, animal_mod))

axs.set_ylim(0, 200)

if save_fig:

fig.tight_layout()

fig.savefig('./data/output/method_%s_rec_%s_mod_%s.png' %

(label, animal_rec, animal_mod))

if close_fig:

plt.close(fig)

lmpars_new = copy.deepcopy(lmpars_old)

ksa = lmpars_old['k2'].value + lmpars_old['k4'].value

lmpars_new['k2'].value = ksa * ksa1

lmpars_new['k4'].value = ksa * (1 - ksa1)

k_scale_dict = {

key: lmpars_new[key].value / lmpars_old[key].value

for key in ['k2', 'k4']}

fine_time=fine_time, fine_stress=fine_stress,

fitApproach=fitApproach,

**plotkws):

# Initialization

lmpars = copy.deepcopy(fitApproach.ref_mean_lmpars)

lmpars[key].value = val

params_dict_new = lmpars_to_params(lmpars)

# Formatting data

if key == 'k2':

lmpars['k4'].value = k4_list[k2_list.index(val)]

params_dict_new = lmpars_to_params(lmpars)

keyname = 'kmc1'

params_paper = get_params_paper(lmpars)

intval = int(params_paper[keyname] * 100)

elif key == 'tau1':

keyname = 'tau_nr'

intval = int(val)

elif key == 'tau2':

keyname = 'tau_mc'

intval = int(val)

# Calculation

single_current = stress_to_current(fine_time, fine_stress,

**params_dict_new).sum(axis=1)

spike_time, fr_inst = get_mod_spike(

lmpars,

**fitApproach.data_dicts_dicts[animal][stim]['mod_data_dict'])

# Plotting

axs[0].plot(fine_time, single_current, '-',

label='%s=%d' % (keyname, intval),

**plotkws)

axs[1].plot(spike_time, fr_inst * 1e3, '.',

label='%s=%d' % (keyname, intval),

**plotkws)

# Saving data

np.savetxt('./data/GregCSVs/ParamTuning/current_%s_%d.csv' %

(keyname, intval), single_current, delimiter=',')

np.savetxt('./data/GregCSVs/ParamTuning/time_%s_%d.csv' %

(keyname, intval), fine_time, delimiter=',')

np.savetxt('./data/GregCSVs/ParamTuning/fr_inst_%s_%d.csv' %

(keyname, intval), fr_inst, delimiter=',')

np.savetxt('./data/GregCSVs/ParamTuning/spike_time_%s_%d.csv' %

(keyname, intval), spike_time, delimiter=',')