123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131 |
- import multiprocessing
- import numpy as np
- from brian2.units import Hz
- from pypet import Trajectory
- from pypet.brian2 import Brian2MonitorResult, Brian2Result
- from tqdm import tqdm
- from scripts.spatial_network.head_direction_index_over_noise_scale import calculate_rates
- from scripts.spatial_network.run_synaptic_strength_scan_orientation_map_small_scale import DATA_FOLDER, TRAJ_NAME
- n_exc_test_print = 465
- def get_spike_train_dictionary(number_of_neurons, spike_times, neuron_indices):
- spike_train_dict = {}
- for neuron_idx in range(number_of_neurons):
- spike_train_dict[neuron_idx] = []
- for neuron_idx, t in zip(neuron_indices, spike_times):
- spike_train_dict[neuron_idx].append(t)
- return spike_train_dict
- def get_firing_rate_dict_per_cell_and_direction(traj, run_name):
- traj.f_set_crun(run_name)
- firing_rate_dict = {}
- direction_names = ["dir{:d}".format(idx) for idx in range(traj.input.number_of_directions)]
- for idx in range(traj.N_E):
- firing_rate_dict[idx] = []
- for direction in direction_names:
- number_of_neurons = traj.N_E
- all_spike_times = traj.results.runs[run_name][direction].spikes.e.t
- neuron_indices = traj.results.runs[run_name][direction].spikes.e.i
- ex_spike_trains = get_spike_train_dictionary(number_of_neurons, all_spike_times,
- neuron_indices)
- ex_spike_rates = calculate_rates(ex_spike_trains.values())
- for idx, spike_rate in enumerate(ex_spike_rates):
- firing_rate_dict[idx].append(spike_rate)
- traj.f_restore_default()
- return firing_rate_dict
- def get_firing_rate_array_per_cell_and_direction(traj, run_name):
- traj.f_set_crun(run_name)
- firing_rate_array = np.ndarray((traj.N_E, traj.input.number_of_directions))
- direction_names = ["dir{:d}".format(idx) for idx in range(traj.input.number_of_directions)]
- for dir_idx, direction in enumerate(direction_names):
- number_of_neurons = traj.N_E
- all_spike_times = traj.results.runs[run_name][direction].spikes.e.t
- neuron_indices = traj.results.runs[run_name][direction].spikes.e.i
- ex_spike_trains = get_spike_train_dictionary(number_of_neurons, all_spike_times,
- neuron_indices)
- ex_spike_rates = calculate_rates(ex_spike_trains.values())
- for n_idx, spike_rate in enumerate(ex_spike_rates):
- #TODO: Why on earth does the unit vanish?
- firing_rate_array[n_idx, dir_idx] = spike_rate
- traj.f_restore_default()
- return firing_rate_array
- def get_head_direction_indices(directions, firing_rate_array):
- n_exc_neurons = firing_rate_array.shape[0]
- n_directions = len(directions)
- tuning_vectors = np.zeros((n_exc_neurons,n_directions,2))
- rate_sums = np.zeros((n_exc_neurons,))
- for ex_id, ex_rates in enumerate(firing_rate_array):
- rate_sum = 0.
- for dir_id, dir in enumerate(directions):
- tuning_vectors[ex_id, dir_id] = np.array([np.cos(dir), np.sin(dir)]) * ex_rates[dir_id]
- rate_sum += ex_rates[dir_id]
- rate_sums[ex_id] = rate_sum
- tuning_vectors_return = tuning_vectors.copy()
- for ex_id in range(n_exc_neurons):
- if rate_sums[ex_id] != 0.:
- tuning_vectors[ex_id, :, :] /= rate_sums[ex_id]
- tuning_vectors_summed = np.sum(tuning_vectors, axis=1)
- head_direction_indices = np.array([np.linalg.norm(v) for v in tuning_vectors_summed])
- return head_direction_indices, tuning_vectors_return
- def get_runs_with_circular_morphology(traj):
- filtered_indices = traj.f_find_idx(('parameters.long_axis', 'parameters.short_axis'), lambda r1, r2: r1 == r2)
- return filtered_indices
- def analyse_single_run(run_name):
- print(run_name)
- firing_rate_array = get_firing_rate_array_per_cell_and_direction(traj, run_name)
- # firing_rate_dict = get_firing_rate_dict_per_cell_and_direction(traj, run_name)
- # traj.f_add_result('runs.$.firing_rate_array', firing_rate_array,
- # comment='The firing rates of the excitatory population')
- head_direction_indices, tuning_vectors = get_head_direction_indices(directions, firing_rate_array)
- # traj.f_add_result('runs.$.head_direction_indices', head_direction_indices,
- # comment='The HDIs of the excitatory population')
- # traj.f_add_result('runs.$.tuning_vectors', tuning_vectors,
- # comment='The tuning vectors of the excitatory population')
- return firing_rate_array, head_direction_indices, tuning_vectors
- if __name__ == "__main__":
- traj = Trajectory(TRAJ_NAME, add_time=False, dynamic_imports=Brian2MonitorResult)
- NO_LOADING = 0
- FULL_LOAD = 2
- traj.f_load(filename=DATA_FOLDER + TRAJ_NAME + ".hdf5", load_parameters=FULL_LOAD, load_results=FULL_LOAD)
- # traj.v_auto_load = True
- correlation_lengths = traj.f_get('correlation_length').f_get_range()
- long_axis = traj.f_get('long_axis').f_get_range()
- short_axis = traj.f_get('short_axis').f_get_range()
- directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions, endpoint=False)
- circular_indices = list(get_runs_with_circular_morphology(traj))
- pool = multiprocessing.Pool()
- # for idx, run_name in enumerate(tqdm(traj.f_get_run_names())):
- multi_proc_result = pool.map(analyse_single_run, traj.f_get_run_names())
- print(type(multi_proc_result),len(multi_proc_result))
- for idx, run_name in enumerate(traj.f_get_run_names()):
- traj.f_set_crun(run_name)
- traj.f_add_result('runs.$.firing_rate_array', multi_proc_result[idx][0],
- comment='The firing rates of the excitatory population')
- traj.f_add_result('runs.$.head_direction_indices', multi_proc_result[idx][1],
- comment='The HDIs of the excitatory population')
- traj.f_add_result('runs.$.tuning_vectors', multi_proc_result[idx][2],
- comment='The tuning vectors of the excitatory population')
- traj.f_restore_default()
- traj.f_store()
|