GridCellsCopy/s_run.py at main · tiny-infinity/GridCellsCopy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""
Main simulation run.
1. Initializes network based on params by calling
    network_initialize from network_init.py
2. Runs simulation.
3. Saves data in hdf5 format.
"""

import numpy as np
from neuron import h
from neuron.units import ms, mV
import time
import sim_utils as s_utils
import h5py
import os
h5py.get_config().track_order = True
import logging


h.nrnmpi_init()
pc = h.ParallelContext()

tstart = time.perf_counter()
args, unknown = s_utils.sim_run_arg_parser().parse_known_args()
sim_id = args.sim_id

if pc.id() == 0:
    logging.basicConfig(handlers=[logging.FileHandler(f"logs/sim_{sim_id}.log",mode='w'),
            logging.StreamHandler()], encoding='utf-8', level=args.verbose,
            format=f'%(asctime)s:%(levelname)s:{sim_id}:Sim 0: %(message)s')
    logger = logging.getLogger()
else:
    logger = None
#load params file
params = s_utils.json_read(f"cache/params_{sim_id}.json")

#initialize network
t2 = time.perf_counter()
s_utils.log_from_rank_0(logger,pc.id(),"Initializing network",level=logging.DEBUG)
network = s_utils.network_intialize(params)
tinit = time.perf_counter()
s_utils.log_from_rank_0(logger,pc.id(),f"Network intialized in {round(tinit-t2,2)}s",level=logging.DEBUG)

#load params
sim_dur = network.params["sim_dur"]
sim_num = str(network.params["sim_num"])
data_root = s_utils.process_data_root(network.params["data_root"])
data_loc = data_root+f"{sim_id}/"

#run simulation
t3 = time.perf_counter()
h.celsius = 37
t = h.Vector().record(h._ref_t)
pc.set_maxstep(10 * ms)
s_utils.log_from_rank_0(logger,pc.id(),f"Simulation started")
h.finitialize(-65 * mV)
if not params["show_progress_bar"]:
    pc.psolve(sim_dur * ms)
elif os.getenv("WHERE_AM_I")=="ADA":
    #check progress every 5s
    sim_chunks = s_utils.get_multiples_with_remainder(sim_dur,5000)
    for chunk in sim_chunks:
        pc.psolve(chunk * ms)
        progress=int((chunk/sim_dur)*100)
        if pc.id()==0:
            print(f"{progress}%",end=" ",flush=True)
else:
    pbar=s_utils.ProgressBar(total=int(sim_dur),pc=pc)
    for i in range(int(sim_dur)+1):
        pc.psolve(i * ms)
        pbar.increment(int(pc.t(0)),pc)
    pbar.finish(pc)

tsim = round(time.perf_counter()-t3, 2)
s_utils.log_from_rank_0(logger,pc.id(),f"Simulation completed in {tsim}s")

#save data for stellate cells
for param_to_record,states in network.params['record_handle_stell'].items():
    if states['state']==True:
        local_data={}
        for cell in network.stellate_cells:
            if cell.recorder.get(f'{param_to_record}',None) is not None:
                local_data[cell._gid]= list(cell.recorder[f'{param_to_record}'])
        all_data = pc.py_alltoall([local_data] + [None] * (pc.nhost() - 1))
        pc.barrier()
        if pc.id() == 0:
            data = {}
            for process_data in all_data:
                data.update(process_data)
            data = dict(sorted(data.items()))
            data_arr = s_utils.list_to_numpy(data.values())
            with h5py.File(data_loc + f"{param_to_record}_{sim_id}.hdf5", "w") as file:
                group = file.create_group(sim_num)
                group.create_dataset(f"{param_to_record}", data=data_arr, compression="gzip",dtype=np.float32)

#save data for internuerons
for param_to_record,states in network.params['record_handle_intrnrn'].items():
    if states['state']==True:
        local_data={}
        for cell in network.interneurons:
            if cell.recorder.get(f'{param_to_record}',None) is not None:
                local_data[cell._gid]= list(cell.recorder[f'{param_to_record}'])
        all_data = pc.py_alltoall([local_data] + [None] * (pc.nhost() - 1))
        pc.barrier()
        if pc.id() == 0:
            data = {}
            for process_data in all_data:
                data.update(process_data)
            data = dict(sorted(data.items()))
            data_arr = s_utils.list_to_numpy(data.values())
            with h5py.File(data_loc + f"{param_to_record}_{sim_id}.hdf5", "w") as file:
                group = file.create_group(sim_num)
                group.create_dataset(f"{param_to_record}", data=data_arr, compression="gzip",dtype=np.float32)

s_utils.log_from_rank_0(logger,pc.id(),f"Data saved in {data_loc}",level=logging.DEBUG)
pc.barrier()
pc.done()
h.quit()