forked from NeoNeuron/2aRNN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdata.py
175 lines (144 loc) · 5.54 KB
/
data.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
"""
Generate training data for context-dependent decision-making tasks.
Reference:
Mante, V., Sussillo, D., Shenoy, K. V. & Newsome, W. T. Context-dependent computation by recurrent dynamics in prefrontal cortex. Nature 503, 78-84 (2013).
"""
import numpy as np
default_timing = [
300, # fixation
1000, # stimulus
900, # delay
500, # response
]
def gen_data(n_trials, dt=20, timing=default_timing, noise=0.05, seed=0):
"""
Generate training data for context-dependent decision-making tasks.
Args:
n_trials (int): Number of trials.
dt (float): Size of time step.
timing (list): Timing of different task period.
seed (int): Random seed.
Returns:
x (np.ndarray): Input data of shape (n_trials, time_steps, 5).
y (np.ndarray): Output data of shape (n_trials, time_steps, 2).
metadata (dict): Metadata including the context, stimulus, action, and timing.
"""
np.random.seed(seed)
rng = np.random.RandomState(seed)
# Generate context data
context = rng.randint(2, size=n_trials)
# Generate stimulus data
stim1_coh = rng.uniform(-1, 1, n_trials).astype(np.float32)
stim2_coh = rng.uniform(-1, 1, n_trials).astype(np.float32)
stim_coh = np.where(context, stim2_coh, stim1_coh)
# Generate action data
action = np.sign(stim_coh)
# timing
n_steps_per_period = (np.asarray(timing) / dt).astype(int)
n_steps_cumsum = np.cumsum(n_steps_per_period)[:-1]
n_timing = {
'fixation': slice(n_steps_cumsum[-1]),
'stimulus': slice(n_steps_cumsum[0], n_steps_cumsum[1]),
'delay': slice(n_steps_cumsum[1], n_steps_cumsum[2]),
'response': slice(n_steps_cumsum[-1], None),
}
# Generate input data
n_steps = np.sum(n_steps_per_period)
x = np.zeros((n_trials, n_steps, 5), dtype=np.float32)
# add fixation
x[:, n_timing['fixation'], 0] = 1
# add stimulus
x[:, n_timing['stimulus'], 1] = stim1_coh.reshape(-1, 1)
x[:, n_timing['stimulus'], 2] = stim2_coh.reshape(-1, 1)
# add noise
x[:, n_timing['stimulus'], 1:3] += rng.normal(
0, noise*np.sqrt(dt), size=(n_trials, n_steps_per_period[1], 2)).astype(np.float32)
# add context cue
x[context==0, :, 3] = 1
x[context==1, :, 4] = 1
# Generate output data
y = np.zeros((n_trials, n_steps, 2), dtype=np.float32)
# add fixation
y[:, n_timing['fixation'], 0] = 1
y[:, n_timing['response'], 1] = action.reshape(-1, 1)
metadata = {
'stim1_coh': stim1_coh,
'stim2_coh': stim2_coh,
'stim_coh': stim_coh,
'ctx': context,
'action': action,
'timing': n_timing,
}
return x, y, metadata
def gen_data_fixed_stim(n_trials, fixed_stim, ctx, dt=20, timing=[300, 1000, 900, 500], noise=0.05, seed=0):
np.random.seed(seed)
rng = np.random.RandomState(seed)
context = np.full(n_trials, ctx)
if ctx == 0:
stim1_coh = np.full(n_trials, fixed_stim)
stim2_coh = rng.uniform(-1, 1, n_trials).astype(np.float32)
else:
stim1_coh = rng.uniform(-1, 1, n_trials).astype(np.float32)
stim2_coh = np.full(n_trials, fixed_stim)
stim_coh = np.where(context, stim2_coh, stim1_coh)
action = np.sign(stim_coh)
n_steps_per_period = (np.asarray(timing) / dt).astype(int)
n_steps_cumsum = np.cumsum(n_steps_per_period)[:-1]
n_timing = {
'fixation': slice(n_steps_cumsum[-1]),
'stimulus': slice(n_steps_cumsum[0], n_steps_cumsum[1]),
'delay': slice(n_steps_cumsum[1], n_steps_cumsum[2]),
'response': slice(n_steps_cumsum[-1], None),
}
n_steps = np.sum(n_steps_per_period)
x = np.zeros((n_trials, n_steps, 5), dtype=np.float32)
x[:, n_timing['fixation'], 0] = 1
x[:, n_timing['stimulus'], 1] = stim1_coh.reshape(-1, 1)
x[:, n_timing['stimulus'], 2] = stim2_coh.reshape(-1, 1)
x[:, n_timing['stimulus'], 1:3] += rng.normal(0, noise * np.sqrt(dt),
size=(n_trials, n_steps_per_period[1], 2)).astype(np.float32)
x[context == 0, :, 3] = 1
x[context == 1, :, 4] = 1
y = np.zeros((n_trials, n_steps, 2), dtype=np.float32)
y[:, n_timing['fixation'], 0] = 1
y[:, n_timing['response'], 1] = action.reshape(-1, 1)
metadata = {
'stim1_coh': stim1_coh,
'stim2_coh': stim2_coh,
'stim_coh': stim_coh,
'ctx': context,
'action': action,
'timing': n_timing,
}
return x, y, metadata
def get_common_limits(h_list, beta_list):
all_x = []
all_y = []
for h in h_list:
for beta in beta_list:
proj = h @ beta
all_x.extend(proj)
all_y.extend(proj)
x_min, x_max = min(all_x), max(all_x)
y_min, y_max = min(all_y), max(all_y)
margin = 0.1
x_range = x_max - x_min
y_range = y_max - y_min
return [x_min - margin * x_range, x_max + margin * x_range], [y_min - margin * y_range, y_max + margin * y_range]
if __name__ == '__main__':
import matplotlib.pyplot as plt
x, y, metadata = gen_data(100)
print(x.shape, y.shape)
print(metadata.keys())
fig, ax = plt.subplots(1, 2, figsize=(12, 2.5))
cb1 = ax[0].pcolormesh(x[2].T)
cb2 = ax[1].pcolormesh(y[2].T)
plt.colorbar(cb1, ax=ax[0])
plt.colorbar(cb2, ax=ax[1])
ax[0].set_title('Input data')
ax[1].set_title('Output data')
ax[0].set_xlabel('Time steps')
ax[1].set_xlabel('Time steps')
ax[0].set_ylabel('Input features')
ax[1].set_ylabel('Output features')
fig.savefig('data.png', bbox_inches='tight')