# In[ ]: N_Outputs = 50 N_Exc = len(exc_trajectory_inputs) N_Inh = len(inh_trajectory_inputs) seed = 19 srf_network = PRF(exc_input_func=partial(trajectory_input, exc_trajectory_inputs), inh_input_func=partial(trajectory_input, inh_trajectory_inputs), n_excitatory=N_Exc, n_inhibitory=N_Inh, n_outputs=N_Outputs, sigma_scale_E=0.005, p_EE=0.03, p_E=0.1, isp_target_rate=0.5, learning_rate_I=0.1, learning_rate_E=0.01, learning_rate_EE=0.001, label="Spatial receptive field network", seed=seed) with srf_network: p_output_spikes = nengo.Probe(srf_network.output.neurons) p_inh_weights = nengo.Probe(srf_network.conn_I, 'weights', sample_every=1.0) p_exc_weights = nengo.Probe(srf_network.conn_E, 'weights',
def build_network(params, inputs, dimensions, input_encoders=None, direct_input=True, presentation_time=1., pause_time=0.1, coords=None, n_outputs=None, n_exc=None, n_inh=None, n_inh_decoder=None, seed=0, t_learn_exc=None, t_learn_inh=None, sample_weights_every=10.): local_random = np.random.RandomState(seed) if coords is None: coords = {} srf_output_coords = coords.get('srf_output', None) srf_exc_coords = coords.get('srf_exc', None) srf_inh_coords = coords.get('srf_inh', None) decoder_coords = coords.get('decoder', None) decoder_inh_coords = coords.get('decoder_inh', None) if type(inputs) == np.ndarray: n_inputs = np.product(inputs.shape[1:]) else: n_inputs = len(inputs) if srf_output_coords is not None: n_outputs = srf_output_coords.shape[0] if srf_exc_coords is not None: n_exc = srf_exc_coords.shape[0] if srf_inh_coords is not None: n_inh = srf_inh_coords.shape[0] if decoder_inh_coords is not None: n_inh_decoder = decoder_inh_coords.shape[0] if n_outputs is None: raise RuntimeError( "n_outputs is not provided and srf_output coordinates are not provided" ) if n_exc is None: raise RuntimeError( "n_exc is not provided and srf_exc coordinates are not provided") if n_inh is None: raise RuntimeError( "n_exc is not provided and srf_inh coordinates are not provided") if n_inh_decoder is None: n_inh_decoder = n_inh if srf_exc_coords is None: srf_exc_coords = np.asarray(range(n_exc)).reshape((n_exc, 1)) / n_exc if srf_output_coords is None: srf_output_coords = np.asarray(range(n_outputs)).reshape( (n_outputs, 1)) / n_outputs if srf_inh_coords is None: srf_inh_coords = np.asarray(range(n_inh)).reshape((n_inh, 1)) / n_inh if decoder_coords is None: decoder_coords = np.asarray(range(n_exc)).reshape((n_exc, 1)) / n_exc if decoder_inh_coords is None: decoder_inh_coords = np.asarray(range(n_inh_decoder)).reshape( (n_inh_decoder, 1)) / n_inh_decoder autoencoder_network = nengo.Network( label="Learning with spatial receptive fields", seed=seed) exc_input_process = None if type(inputs) == np.ndarray: exc_input_process = PresentInputWithPause(inputs, presentation_time, pause_time) else: exc_input_process = inputs with autoencoder_network as model: learning_rate_E = params.get('learning_rate_E', 1e-5) learning_rate_EE = params.get('learning_rate_EE', 1e-5) learning_rate_I = params.get('learning_rate_I', 1e-4) learning_rate_E_func = ( lambda t: learning_rate_E if t <= t_learn_exc else 0.0 ) if t_learn_exc is not None and learning_rate_E is not None else None learning_rate_EE_func = ( lambda t: learning_rate_EE if t <= t_learn_exc else 0.0 ) if t_learn_exc is not None and learning_rate_EE is not None else None learning_rate_I_func = ( lambda t: learning_rate_I if t <= t_learn_inh else 0.0 ) if t_learn_inh is not None and learning_rate_E is not None else None srf_network = PRF(dimensions=dimensions, exc_input_process=exc_input_process, connect_exc_inh_input=True, connect_out_out=True if ('w_initial_EE' in params) and (params['w_initial_EE'] is not None) else False, connect_inh_inh=True, n_excitatory=n_exc, n_inhibitory=n_inh, n_outputs=n_outputs, output_coordinates=srf_output_coords, exc_coordinates=srf_exc_coords, inh_coordinates=srf_inh_coords, w_initial_E=params['w_initial_E'], w_initial_I=params['w_initial_I'], w_initial_EI=params['w_initial_EI'], w_initial_EE=params.get('w_initial_EE', None), w_EI_Ext=params['w_EI_Ext'], p_E=params['p_E_srf'], p_I=params['p_I_srf'], p_EE=params['p_EE'], p_EI_Ext=params['p_EI_Ext'], p_EI=params['p_EI'], tau_E=params['tau_E'], tau_I=params['tau_I'], tau_input=params['tau_input'], learning_rate_I=learning_rate_I, learning_rate_E=learning_rate_E, learning_rate_EE=learning_rate_EE, learning_rate_E_func=learning_rate_E_func, learning_rate_EE_func=learning_rate_EE_func, learning_rate_I_func=learning_rate_I_func, sigma_scale_E=params['sigma_scale_E'], sigma_scale_EI=params['sigma_scale_EI'], sigma_scale_EI_Ext=params['sigma_scale_EI_Ext'], sigma_scale_EE=params['sigma_scale_EE'], sigma_scale_I=params['sigma_scale_I'], isp_target_rate=params['isp_target_rate'], direct_input=direct_input, label="Spatial receptive field network", seed=seed) if input_encoders is not None: srf_network.exc.encoders = input_encoders decoder = nengo.Ensemble(n_exc, dimensions=dimensions, neuron_type=nengo.LIF(), radius=1, intercepts=nengo.dists.Choice([0.1]), max_rates=nengo.dists.Choice([40])) decoder_inh = nengo.Ensemble(n_inh_decoder, dimensions=dimensions, neuron_type=nengo.LIF(tau_rc=0.005), radius=1, intercepts=nengo.dists.Choice([0.1]), max_rates=nengo.dists.Choice([100])) tau_E = params['tau_E'] w_DEC_E = params['w_DEC_E'] p_DEC = params['p_DEC'] model.weights_initial_DEC_E = local_random.uniform( size=n_outputs * n_exc).reshape((n_exc, n_outputs)) * w_DEC_E for i in range(n_exc): dist = cdist(decoder_coords[i, :].reshape((1, -1)), srf_network.output_coordinates).flatten() sigma = 0.1 prob = distance_probs(dist, sigma) sources = np.asarray(local_random.choice(n_outputs, round(p_DEC * n_outputs), replace=False, p=prob), dtype=np.int32) model.weights_initial_DEC_E[ i, np.logical_not(np.in1d(range(n_outputs), sources))] = 0. learning_rate_DEC_E = params['learning_rate_D_Exc'] model.conn_DEC_E = nengo.Connection( srf_network.output.neurons, decoder.neurons, transform=model.weights_initial_DEC_E, synapse=nengo.Alpha(tau_E), learning_rule_type=HSP(directed=True)) model.node_learning_rate_DEC_E = nengo.Node( lambda t: learning_rate_DEC_E) model.conn_learning_rate_DEC_E = nengo.Connection( model.node_learning_rate_DEC_E, model.conn_DEC_E.learning_rule) w_DEC_I_ff = params['w_DEC_I'] weights_initial_DEC_I_ff = local_random.uniform( size=n_inh * n_outputs).reshape((n_inh, n_outputs)) * w_DEC_I_ff for i in range(n_inh_decoder): dist = cdist(decoder_inh_coords[i, :].reshape((1, -1)), srf_network.output_coordinates).flatten() sigma = 1.0 prob = distance_probs(dist, sigma) sources = np.asarray(local_random.choice(n_outputs, round(p_DEC * n_outputs), replace=False, p=prob), dtype=np.int32) weights_initial_DEC_I_ff[ i, np.logical_not(np.in1d(range(n_outputs), sources))] = 0. conn_DEC_I_ff = nengo.Connection(srf_network.output.neurons, decoder_inh.neurons, transform=weights_initial_DEC_I_ff, synapse=nengo.Alpha(tau_E)) w_DEC_I_fb = params['w_initial_I_DEC_fb'] weights_initial_DEC_I_fb = local_random.uniform( size=n_inh_decoder * n_exc).reshape( (n_exc, n_inh_decoder)) * w_DEC_I_fb for i in range(n_exc): dist = cdist(decoder_coords[i, :].reshape((1, -1)), decoder_inh_coords).flatten() sigma = 1.0 prob = distance_probs(dist, sigma) sources = np.asarray(local_random.choice(n_inh_decoder, round(0.2 * n_inh_decoder), replace=False, p=prob), dtype=np.int32) weights_initial_DEC_I_fb[ i, np.logical_not(np.in1d(range(n_inh_decoder), sources))] = 0. conn_DEC_I_fb = nengo.Connection( decoder_inh.neurons, decoder.neurons, transform=weights_initial_DEC_I_fb, synapse=nengo.Lowpass(params['tau_I']), learning_rule_type=CDISP(learning_rate=params['learning_rate_D'])) w_SRF_I_bp = params['w_initial_I'] weights_initial_SRF_I_bp = local_random.uniform( size=n_inh_decoder * n_outputs).reshape( (n_outputs, n_inh_decoder)) * w_SRF_I_bp for i in range(n_outputs): dist = cdist(srf_network.output_coordinates[i, :].reshape((1, -1)), decoder_inh_coords).flatten() sigma = 0.1 prob = distance_probs(dist, sigma) sources = np.asarray(local_random.choice(n_inh_decoder, round(0.05 * n_inh_decoder), replace=False, p=prob), dtype=np.int32) weights_initial_SRF_I_bp[ i, np.logical_not(np.in1d(range(n_inh_decoder), sources))] = 0. conn_SRF_I_bp = nengo.Connection(decoder_inh.neurons, srf_network.output.neurons, transform=weights_initial_SRF_I_bp, synapse=nengo.Alpha(params['tau_I'])) coincidence_detection = nengo.Node( size_in=2 * n_exc, size_out=n_exc, output=lambda t, x: np.subtract(x[:n_exc], x[n_exc:])) nengo.Connection(coincidence_detection, conn_DEC_I_fb.learning_rule) nengo.Connection(srf_network.exc.neurons, coincidence_detection[n_exc:]) nengo.Connection(decoder.neurons, coincidence_detection[:n_exc]) p_srf_rec_weights = None with srf_network: p_srf_output_spikes = nengo.Probe(srf_network.output.neurons, 'output', synapse=None) p_srf_exc_spikes = nengo.Probe(srf_network.exc.neurons, 'output') p_srf_inh_spikes = nengo.Probe(srf_network.inh.neurons, 'output') p_srf_inh_weights = nengo.Probe(srf_network.conn_I, 'weights', sample_every=sample_weights_every) p_srf_exc_weights = nengo.Probe(srf_network.conn_E, 'weights', sample_every=sample_weights_every) if srf_network.conn_EE is not None: p_srf_rec_weights = nengo.Probe( srf_network.conn_EE, 'weights', sample_every=sample_weights_every) p_decoder_spikes = nengo.Probe(decoder.neurons, synapse=None) p_decoder_inh_spikes = nengo.Probe(decoder_inh.neurons, synapse=None) p_decoder_weights = nengo.Probe(model.conn_DEC_E, 'weights', sample_every=sample_weights_every) model.srf_network = srf_network model.decoder_ens = decoder model.decoder_inh_ens = decoder_inh return { 'network': autoencoder_network, 'neuron_probes': { 'srf_output_spikes': p_srf_output_spikes, 'srf_exc_spikes': p_srf_exc_spikes, 'srf_inh_spikes': p_srf_inh_spikes, 'decoder_spikes': p_decoder_spikes, 'decoder_inh_spikes': p_decoder_inh_spikes, }, 'weight_probes': { 'srf_exc_weights': p_srf_exc_weights, 'srf_rec_weights': p_srf_rec_weights, 'decoder_weights': p_decoder_weights, 'srf_inh_weights': p_srf_inh_weights, # 'srf_exc_weights': p_srf_exc_weights, } }
dimensions=np.prod(env_iface.state_dim), radius=sensor_radius) gabor_size = (5, 5) # Size of the gabor filter # Generate the encoders for the sensory ensemble sensor_encoders = Gabor().generate(n_input, gabor_size, rng=rng) sensor_encoders = Mask(image_shape).populate(sensor_encoders, rng=rng, flatten=True) sensor_net.encoders = sensor_encoders srf_net = PRF(n_excitatory=n_input, n_inhibitory=n_inhibitory, connect_exc_inh_input=True, n_outputs=n_place, dimensions=env_iface.n_actions, label="Spatial receptive field network", seed=seed, **srf_params) actor_net = nengo.Ensemble(n_neurons=n_actor, dimensions=env_iface.n_actions, radius=actor_radius) sensor_conn = nengo.Connection(sensor, sensor_net) sensor_srf_conn = nengo.Connection(sensor_net.neurons, srf_net.exc.neurons, synapse=nengo.Lowpass(tau), transform=np.eye(n_input) * srf_params['w_input'])
srf_network = PRF(exc_input_func=partial(array_input, train_data, dt), connect_exc_inh_input = True, n_excitatory = N_exc_srf, n_inhibitory = N_inh_srf, n_outputs = N_outputs_srf, isp_target_rate = 1.0, tau_input = params['tau_input'], w_initial_E = params['w_initial_E'], w_initial_I = params['w_initial_I'], w_initial_EI = params['w_initial_EI'], w_EI_Ext = params['w_EI_Ext'], p_E = params['p_E_srf'], p_EE = params['p_EE'], p_EI_Ext = params['p_EI_Ext'], p_EI = params['p_EI'], tau_E = params['tau_E'], tau_I = params['tau_I'], learning_rate_I=params['learning_rate_I'], learning_rate_E=params['learning_rate_E'], exc_coordinates = srf_exc_coords, inh_coordinates = srf_inh_coords, output_coordinates = srf_output_coords, label="Spatial receptive field network", seed=srf_seed) decoder = nengo.Ensemble(N_exc_srf, dimensions=1,
def eval_srf_net(input_matrix, params): N_inputs = input_matrix.shape[0] N_outputs_srf = params['N_outputs_srf'] N_exc_srf = N_inputs N_inh_srf = params['N_inh_srf'] N_exc_place = params['N_exc_place'] N_inh_place = params['N_inh_place'] dt = 0.001 t_end = input_matrix.shape[1] * dt srf_place_network = nengo.Network( label="Learning with spatial receptive fields", seed=params['seed']) rng = np.random.RandomState(seed=params['seed']) with srf_place_network as model: srf_network = PRF(exc_input_func=partial(array_input, input_matrix, params['dt']), connect_exc_inh_input=True, n_excitatory=N_exc_srf, n_inhibitory=params['N_inh_srf'], n_outputs=params['N_outputs_srf'], w_initial_E=params['w_initial_E'], w_initial_I=params['w_initial_I'], w_initial_EI=params['w_initial_EI'], w_EI_Ext=params['w_EI_Ext'], p_E=params['p_E_srf'], p_EE=params['p_EE'], p_EI_Ext=params['p_EI_Ext'], p_EI=params['p_EI'], tau_E=params['tau_E'], tau_I=params['tau_I'], tau_input=params['tau_input'], isp_target_rate=params['isp_target_rate'], learning_rate_I=params['learning_rate_I'], learning_rate_E=params['learning_rate_E'], label="Spatial receptive field network", seed=params['seed']) place_network = EI(n_excitatory=params['N_exc_place'], n_inhibitory=params['N_inh_place'], isp_target_rate=params['isp_target_rate'], learning_rate_I=params['learning_rate_I'], learning_rate_E=params['learning_rate_E'], p_E=params['p_E_place'], p_EE=params['p_EE'], p_EI=params['p_EI'], tau_E=params['tau_E'], tau_I=params['tau_I'], tau_input=params['tau_input'], connect_exc_inh=True, label="Place network", seed=params['seed']) w_PV_E = params['w_PV_E'] p_PV_E = params['p_PV_E'] weights_dist_PV_E = rng.normal(size=N_outputs_srf * N_exc_place).reshape( (N_exc_place, N_outputs_srf)) weights_initial_PV_E = (weights_dist_PV_E - weights_dist_PV_E.min() ) / (weights_dist_PV_E.max() - weights_dist_PV_E.min()) * w_PV_E for i in range(N_exc_place): dist = i - np.asarray(range(N_outputs_srf)) sigma = p_PV_E * N_outputs_srf weights = np.exp(-dist / sigma**2) prob = weights / weights.sum(axis=0) sources = np.asarray(rng.choice(N_outputs_srf, round(p_PV_E * N_outputs_srf), replace=False, p=prob), dtype=np.int32) weights_initial_PV_E[ i, np.logical_not(np.in1d(range(N_outputs_srf), sources))] = 0. conn_PV_E = nengo.Connection( srf_network.output.neurons, place_network.exc.neurons, transform=weights_initial_PV_E, synapse=nengo.Alpha(params['tau_E']), learning_rule_type=HSP(learning_rate=params['learning_rate_E'])) w_SV_E = params['w_PV_E'] p_SV_E = params['p_SV_E'] weights_dist_SV_E = rng.normal(size=N_exc_srf * N_exc_place).reshape( (N_exc_place, N_exc_srf)) weights_initial_SV_E = (weights_dist_SV_E - weights_dist_SV_E.min() ) / (weights_dist_SV_E.max() - weights_dist_SV_E.min()) * w_SV_E for i in range(N_exc_place): dist = i - np.asarray(range(N_exc_srf)) sigma = p_SV_E * N_exc_srf weights = np.exp(-dist / sigma**2) prob = weights / weights.sum(axis=0) sources = np.asarray(rng.choice(N_exc_srf, round(p_SV_E * N_exc_srf), replace=False, p=prob), dtype=np.int32) weights_initial_SV_E[ i, np.logical_not(np.in1d(range(N_exc_srf), sources))] = 0. conn_SV_E = nengo.Connection( srf_network.exc.neurons, place_network.exc.neurons, transform=weights_initial_SV_E, synapse=nengo.Alpha(params['tau_E']), learning_rule_type=HSP(learning_rate=params['learning_rate_E'])) w_PV_I = params['w_PV_I'] weights_initial_PV_I = rng.uniform( size=N_inh_place * N_outputs_srf).reshape( (N_inh_place, N_outputs_srf)) * w_PV_I conn_PV_I = nengo.Connection(srf_network.output.neurons, place_network.inh.neurons, transform=weights_initial_PV_I, synapse=nengo.Alpha(params['tau_I'])) with place_network: p_place_output_spikes = nengo.Probe(place_network.exc.neurons, synapse=None) with srf_network: p_output_spikes = nengo.Probe(srf_network.output.neurons, synapse=None) with nengo.Simulator(srf_place_network, optimize=True) as sim: sim.run(t_end) srf_output_spikes = sim.data[p_output_spikes] srf_output_rates = rates_kernel(sim.trange(), srf_output_spikes, tau=0.1) place_output_spikes = sim.data[p_place_output_spikes] place_output_rates = rates_kernel(sim.trange(), place_output_spikes, tau=0.1) return np.asarray([ obj_modulation_depth(srf_output_rates), obj_modulation_depth(place_output_rates), obj_fraction_active(srf_output_rates), obj_fraction_active(place_output_rates) ])
# Create input nodes representing the sine and cosine x = nengo.Node(output=np.sin) y = nengo.Node(output=np.cos) with model: # The indices in neurons define which dimension the input will project to nengo.Connection(x, input_ens[0]) nengo.Connection(y, input_ens[1]) with model: srf_network = PRF(n_excitatory=N_Exc, n_inhibitory=N_Inh, n_outputs=N_Outputs, connect_exc_inh_input=True, p_E=0.02, p_EE=0.01, learning_rate_E=0.05, learning_rate_I=0.01, tau_E=0.005, tau_I=0.020, label="Spatial receptive field network") nengo.Connection(input_ens.neurons, srf_network.exc.neurons, synapse=nengo.Alpha(tau_input), transform=np.eye(N_Exc) * w_input) with model: x_probe = nengo.Probe(x, "output") y_probe = nengo.Probe(y, "output") input_probe = nengo.Probe(input_ens, "decoded_output", synapse=0.01)
seed = 21 srf_rmrl_network = nengo.Network( label="Reward-modulated replay learning with spatial receptive fields", seed=seed) with srf_rmrl_network as model: rng = np.random.RandomState(seed=seed) place_network = PRF(exc_input_func=partial( input_matrix_at, trajectory_input_matrix(exc_trajectory_inputs, trj_t)), inh_input_func=partial( input_matrix_at, trajectory_input_matrix(inh_trajectory_inputs, trj_t)), n_excitatory=n_excitatory, n_inhibitory=n_inhibitory, n_outputs=n_outputs, label="Spatial receptive field network", seed=seed) value_network = PRF(n_excitatory=n_excitatory_value, n_inhibitory=n_inhibitory_value, n_outputs=n_outputs, label="Value network", seed=seed) weights_dist_PV_E = rng.normal(size=n_outputs * n_excitatory_value).reshape( (n_excitatory_value, n_outputs))