forked from zcqsntr/worm_neural_nets
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathrun_worms_sim.py
176 lines (122 loc) · 6.03 KB
/
run_worms_sim.py
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
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
import math
import os
import numpy as np
import sys
import time
from wormSimulator import WormSimulator
import copy
import argparse
import matplotlib.pyplot as plt
def load_params(param_file):
if param_file[-3:] == 'csv':
params = np.loadtxt(open(param_file, "rb"), delimiter=",")
elif param_file[-3:] == 'npy':
params = np.load(param_file)
else:
raise ValueError('parameter file not recognised')
return params
def run_behaviour_assay(simulator, weights):
simulator.set_mode('B')
all_sectors, all_sols = simulator.run_experiment_par(weights, n_worms)
data_dir = os.path.join(out_dir, 'data')
os.makedirs(data_dir, exist_ok=True)
np.save(os.path.join(data_dir, 'all_sectors.npy'), all_sectors)
np.save(os.path.join(data_dir, 'all_sols.npy'), all_sols)
if plot:
ncols = 10
nrows = math.ceil(len(all_sectors) / ncols)
fig, axs = plt.subplots(nrows=nrows, ncols=ncols, figsize=(15, 7.5))
for i, sectors in enumerate(all_sectors):
ax = axs[i]
ax.set_ylim(bottom=-6.1, top=6.1)
ax.violinplot(list(map(sum, sectors)))
plot_dir = os.path.join(out_dir, 'plots')
os.makedirs(plot_dir, exist_ok=True)
plt.savefig(os.path.join(plot_dir, 'violin_plots.pdf'))
for weights, sols in enumerate(all_sols):
weights_dir = os.path.join(plot_dir, 'sols', 'weights' + str(weights))
os.makedirs(weights_dir, exist_ok=True)
for worm, sol in enumerate(sols):
simulator.plot_sol(sol, save_path=os.path.join(weights_dir, 'worm' + str(worm)))
def run_calcium_experiments(simulator, weights):
simulator.set_mode('C')
_, calcium_sims = simulator.run_experiment_par(weights, n_worms=1)
data_dir = os.path.join(out_dir, 'data')
os.makedirs(data_dir, exist_ok=True)
np.save(os.path.join(data_dir, 'calcium_sims.npy'), calcium_sims)
if plot:
ncols = 10
fig, axs = plt.subplots(nrows=10, ncols=ncols, figsize=(15, 7.5))
plot_dir = os.path.join(out_dir, 'plots')
os.makedirs(plot_dir, exist_ok=True)
for i, c_sim in enumerate(calcium_sims):
c_sim = c_sim[0] #just one worm
# add to the summary plot
r = i // ncols
c = i % ncols
ax = axs[r, c]
alpha = 0.5
lw = 1
max_t = simulator.t_span[-1]
dt = simulator.dt
ax.plot(np.arange(0, max_t, dt), c_sim[0, 1:-1], label='AWC', alpha=alpha, lw=lw)
ax.plot(np.arange(0, max_t, dt), c_sim[3, 1:-1], label='AIB', alpha=alpha, lw=lw)
ax.plot(np.arange(0, max_t, dt), c_sim[5, 1:-1], label='AIY', alpha=alpha, lw=lw)
# ax.plot(np.arange(0, max_t, dt), c_sim[4, 1:-1], label='AIA', alpha = alpha, lw = lw) #uncomment this line to plot AIA
ax.set_title('worm ' + str(i + 1), fontsize=10)
if r == 0 and c == 0:
ax.legend(loc=2, bbox_to_anchor=(0, 2.2))
if c != 0:
ax.set_yticklabels([])
if r != 9:
ax.set_xticklabels([])
# create the detailed plot
d_fig, d_axs = plt.subplots(2, 2, figsize=(12.0, 8.0))
p_c = plt.rcParams['axes.prop_cycle']
colours = p_c.by_key()['color']
inds = [0, 3, 5, 4]
labels = ['AWC', 'AIB', 'AIY', 'AIA']
for j in range(4):
d_axs[j // 2, j % 2].plot(np.arange(0, max_t, dt), c_sim[inds[j], 1:-1], label=labels[j], c=colours[j])
d_axs[j // 2, j % 2].legend()
d_axs[j // 2, j % 2].set_xlabel('Time (s)')
d_axs[j // 2, j % 2].set_ylabel('Neuron voltage')
d_fig.suptitle('worm ' + str(i + 1))
os.makedirs(os.path.join(plot_dir, 'detailed_calcium'), exist_ok=True)
d_fig.savefig(os.path.join(plot_dir, 'detailed_calcium', 'weights' + str(i + 1) + '.pdf'))
fig.savefig(os.path.join(plot_dir, 'calcium_plots.pdf'))
dir_path = os.path.dirname(os.path.realpath(__file__))
simulator = WormSimulator(dt=0.005)
worm_trapped = False
conc_interval = None
parser = argparse.ArgumentParser(description='run worm simulations')
parser.add_argument('--weights_file', type=str, help='input parameter file, can be csv or saved numpy array')
parser.add_argument('--weights', type=str, help='can be used to quickly simulate a set of parameters, either this or --in_file must be specified, if both specified --in_file will be used')
parser.add_argument('--out_dir', type=str, help='directory to save results in, default is ./working_dir')
parser.add_argument('--plot', type=str, help='1 to plot 0 to not, default is 1')
parser.add_argument('--opt', type=str, help='A to run behaviour assay, C to run calcium plot, B to run both, default is B')
parser.add_argument('--n_worms', type=str, help='number of worms to simulate in each experiment for the violin plots, default is 100. If --calcium=1 this argument is ignored as only one worm is required')
if __name__ == '__main__':
args = parser.parse_args()
if args.opt == 'A':
opt = 'A'
elif args.opt == 'C':
opt = 'C'
else:
opt = 'B'
plot = (args.plot == '1' or args.plot is None)
out_dir = 'working_dir' if args.out_dir is None else args.out_dir
n_worms = int(args.n_worms) if args.n_worms is not None else 100
if args.weights_file is not None:
weights = load_params(args.weights_file)
elif args.weights is not None:
weights = np.array([list(map(float, args.weights.split(',')))])
else:
raise ValueError('One of --param_file or --parameters must be given')
if weights.shape[1] != 9:
raise ValueError('weights must have shape (-1, 9) current weights shape = {}'.format(weights.shape))
if opt in ['A', 'B']:
run_behaviour_assay(simulator, weights)
if opt in ['C', 'B']:
run_calcium_experiments(simulator, weights)
#np.save('/results/final_results/230111_mock/fitting_data/final_calcium_sims.npy', calcium_sims)