Split inhib/exc populations

Author: C-Earl

scripts/Chris/DQN/Eval.ipynb

@@ -7,57 +7,104 @@
 
 
 class Reservoir(Network):
-  def __init__(self, in_size, res_size, hyper_params,
-               w_in_res, w_res_res, device='cpu'):
+  def __init__(self, in_size, exc_size, inh_size, hyper_params,
+               w_in_exc, w_in_inh, w_exc_exc, w_exc_inh, w_inh_exc, w_inh_inh,
+               device='cpu'):
     super().__init__()
 
     ## Layers ##
     input = Input(n=in_size)
-    res = AdaptiveLIFNodes(
-      n=res_size,
-      thresh=hyper_params['thresh'],
-      theta_plus=hyper_params['theta_plus'],
-      refrac=hyper_params['refrac'],
-      reset=hyper_params['reset'],
-      tc_theta_decay=hyper_params['tc_theta_decay'],
-      tc_decay=hyper_params['tc_decay'],
+    res_exc = AdaptiveLIFNodes(
+      n=exc_size,
+      thresh=hyper_params['thresh_exc'],
+      theta_plus=hyper_params['theta_plus_exc'],
+      refrac=hyper_params['refrac_exc'],
+      reset=hyper_params['reset_exc'],
+      tc_theta_decay=hyper_params['tc_theta_decay_exc'],
+      tc_decay=hyper_params['tc_decay_exc'],
       traces=True,
     )
-    res_monitor = Monitor(res, ["s"], device=device)
-    self.add_monitor(res_monitor, name='res_monitor')
-    self.res_monitor = res_monitor
+    exc_monitor = Monitor(res_exc, ["s"], device=device)
+    self.add_monitor(exc_monitor, name='res_monitor_exc')
+    self.exc_monitor = exc_monitor
+    res_inh = AdaptiveLIFNodes(
+      n=inh_size,
+      thresh=hyper_params['thresh_inh'],
+      theta_plus=hyper_params['theta_plus_inh'],
+      refrac=hyper_params['refrac_inh'],
+      reset=hyper_params['reset_inh'],
+      tc_theta_decay=hyper_params['tc_theta_decay_inh'],
+      tc_decay=hyper_params['tc_decay_inh'],
+      traces=True,
+    )
+    inh_monitor = Monitor(res_inh, ["s"], device=device)
+    self.add_monitor(inh_monitor, name='res_monitor_inh')
+    self.inh_monitor = inh_monitor
     self.add_layer(input, name='input')
-    self.add_layer(res, name='res')
+    self.add_layer(res_exc, name='res_exc')
+    self.add_layer(res_inh, name='res_inh')
 
     ## Connections ##
-    in_res_wfeat = Weight(name='in_res_weight_feature', value=w_in_res,)
-    in_res_conn = MulticompartmentConnection(
-      source=input, target=res,
-      device=device, pipeline=[in_res_wfeat],
+    in_exc_wfeat = Weight(name='in_exc_weight_feature', value=w_in_exc,)
+    in_exc_conn = MulticompartmentConnection(
+      source=input, target=res_exc,
+      device=device, pipeline=[in_exc_wfeat],
+    )
+    in_inh_wfeat = Weight(name='in_inh_weight_feature', value=w_in_inh,)
+    in_inh_conn = MulticompartmentConnection(
+      source=input, target=res_inh,
+      device=device, pipeline=[in_inh_wfeat],
+    )
+
+    exc_exc_wfeat = Weight(name='exc_exc_weight_feature', value=w_exc_exc,)
+                           # learning_rule=MSTDP,
+                           # nu=hyper_params['nu_exc_exc'], range=hyper_params['range_exc_exc'], decay=hyper_params['decay_exc_exc'])
+    exc_exc_conn = MulticompartmentConnection(
+      source=res_exc, target=res_exc,
+      device=device, pipeline=[exc_exc_wfeat],
+    )
+    exc_inh_wfeat = Weight(name='exc_inh_weight_feature', value=w_exc_inh,)
+                           # learning_rule=MSTDP,
+                           # nu=hyper_params['nu_exc_inh'], range=hyper_params['range_exc_inh'], decay=hyper_params['decay_exc_inh'])
+    exc_inh_conn = MulticompartmentConnection(
+      source=res_exc, target=res_inh,
+      device=device, pipeline=[exc_inh_wfeat],
+    )
+    inh_exc_wfeat = Weight(name='inh_exc_weight_feature', value=w_inh_exc,)
+                           # learning_rule=MSTDP,
+                           # nu=hyper_params['nu_inh_exc'], range=hyper_params['range_inh_exc'], decay=hyper_params['decay_inh_exc'])
+    inh_exc_conn = MulticompartmentConnection(
+      source=res_inh, target=res_exc,
+      device=device, pipeline=[inh_exc_wfeat],
     )
-    res_res_wfeat = Weight(name='res_res_weight_feature', value=w_res_res,
+    inh_inh_wfeat = Weight(name='inh_inh_weight_feature', value=w_inh_inh,)
                            # learning_rule=MSTDP,
-                           nu=hyper_params['nu'], range=hyper_params['range'], decay=hyper_params['decay'])
-    res_res_conn = MulticompartmentConnection(
-      source=res, target=res,
-      device=device, pipeline=[res_res_wfeat],
+                           # nu=hyper_params['nu_inh_inh'], range=hyper_params['range_inh_inh'], decay=hyper_params['decay_inh_inh'])
+    inh_inh_conn = MulticompartmentConnection(
+      source=res_inh, target=res_inh,
+      device=device, pipeline=[inh_inh_wfeat],
     )
-    self.add_connection(in_res_conn, source='input', target='res')
-    self.add_connection(res_res_conn, source='res', target='res')
-    self.res_res_conn = res_res_conn
+    self.add_connection(in_exc_conn, source='input', target='res_exc')
+    self.add_connection(in_inh_conn, source='input', target='res_inh')
+    self.add_connection(exc_exc_conn, source='res_exc', target='res_exc')
+    self.add_connection(exc_inh_conn, source='res_exc', target='res_inh')
+    self.add_connection(inh_exc_conn, source='res_inh', target='res_exc')
+    self.add_connection(inh_inh_conn, source='res_inh', target='res_inh')
 
     ## Migrate ##
     self.to(device)
 
   def store(self, spike_train, sim_time):
     self.learning = True
     self.run(inputs={'input': spike_train}, time=sim_time, reward=1)
-    res_spikes = self.res_monitor.get('s')
+    exc_spikes = self.exc_monitor.get('s')
+    inh_spikes = self.inh_monitor.get('s')
     self.learning = False
-    return res_spikes
+    return exc_spikes, inh_spikes
 
   def recall(self, spike_train, sim_time):
     self.learning = False
     self.run(inputs={'input': spike_train}, time=sim_time,)
-    res_spikes = self.res_monitor.get('s')
-    return res_spikes
+    exc_spikes = self.exc_monitor.get('s')
+    inh_spikes = self.inh_monitor.get('s')
+    return exc_spikes, inh_spikes

scripts/Chris/DQN/Reservoir.py

@@ -11,53 +11,64 @@
   ## Constants ##
   WIDTH = 5
   HEIGHT = 5
-  SAMPLES_PER_POS = 1000
+  SAMPLES_PER_POS = 10
   NOISE = 0.1   # Noise in sampling
   WINDOW_FREQ = 10
   WINDOW_SIZE = 10
   NUM_CELLS = 20
   X_RANGE = (0, 5)
   Y_RANGE = (0, 5)
   SIM_TIME = 50
-  MAX_SPIKE_FREQ = 0.3
+  MAX_SPIKE_FREQ = 0.8
   GC_MULTIPLES = 1
-  RES_SIZE = 250
+  EXC_SIZE = 250
+  INH_SIZE = 50
   STORE_SAMPLES = 100
-  hyper_params = {
-    'thresh': -55,
-    'theta_plus': 0,
-    'refrac': 1,
-    'reset': -65,
-    'tc_theta_decay': 500,
-    'tc_decay': 30,
-    'nu': (0.01, -0.01),
-    'range': [-1, 1],
-    'decay': None,
-  }
   WINDOW_FREQ = 10
   WINDOW_SIZE = 10
   OUT_DIM = 2
   TRAIN_RATIO = 0.8
   BATCH_SIZE = 10
   PLOT = True
+  exc_hyper_params = {
+    'thresh_exc': -55,
+    'theta_plus_exc': 0,
+    'refrac_exc': 1,
+    'reset_exc': -65,
+    'tc_theta_decay_exc': 500,
+    'tc_decay_exc': 30,
+    # 'nu': (0.01, -0.01),
+    # 'range': [-1, 1],
+    # 'decay': None,
+  }
+  inh_hyper_params = {
+    'thresh_inh': -55,
+    'theta_plus_inh': 0,
+    'refrac_inh': 1,
+    'reset_inh': -65,
+    'tc_theta_decay_inh': 500,
+    'tc_decay_inh': 30,
+  }
+  hyper_params = exc_hyper_params | inh_hyper_params
 
   ## Sample Generation ##
-  x_offsets = np.random.uniform(-1, 1, NUM_CELLS)
-  y_offsets = np.random.uniform(-1, 1, NUM_CELLS)
-  offsets = list(zip(x_offsets, y_offsets))           # Grid Cell x & y offsets
-  scales = [np.random.uniform(1.7, 5) for i in range(NUM_CELLS)]   # Dist. between Grid Cell peaks
-  vars = [.85] * NUM_CELLS              # Variance of Grid Cell activity
+  # x_offsets = np.random.uniform(-1, 1, NUM_CELLS)
+  #
+  # y_offsets = np.random.uniform(-1, 1, NUM_CELLS)
+  # offsets = list(zip(x_offsets, y_offsets))           # Grid Cell x & y offsets
+  # scales = [np.random.uniform(1.7, 5) for i in range(NUM_CELLS)]   # Dist. between Grid Cell peaks
+  # vars = [.85] * NUM_CELLS              # Variance of Grid Cell activity
   # samples, labels, sorted_samples = sample_generator(scales, offsets, vars, X_RANGE, Y_RANGE, SAMPLES_PER_POS,
   #                                                    noise=NOISE, padding=1, plot=PLOT)
   #
-  # ## Spike Train Generation ##
+  # # Spike Train Generation ##
   # spike_trains, labels, sorted_spike_trains = spike_train_generator(samples, labels, SIM_TIME, GC_MULTIPLES, MAX_SPIKE_FREQ)
 
-  ## Association (Store) ##
-  store_reservoir(RES_SIZE, STORE_SAMPLES, NUM_CELLS, GC_MULTIPLES, SIM_TIME, hyper_params, PLOT)
+  # ## Association (Store) ##
+  store_reservoir(EXC_SIZE, INH_SIZE, STORE_SAMPLES, NUM_CELLS, GC_MULTIPLES, SIM_TIME, hyper_params, PLOT)
 
-  ## Association (Recall) ##
-  # recall_reservoir(RES_SIZE, SIM_TIME, PLOT)
+  # ## Association (Recall) ##
+  recall_reservoir(EXC_SIZE, INH_SIZE, SIM_TIME, PLOT)
 
   # # Preprocess Recalls ##
   # recalled_mem_preprocessing(WINDOW_FREQ, WINDOW_SIZE, PLOT)

scripts/Chris/DQN/pipeline_executor.py

@@ -3,7 +3,7 @@
 import torch
 from matplotlib import pyplot as plt
 
-def recall_reservoir(res_size, sim_time, plot=False):
+def recall_reservoir(exc_size, inh_size, sim_time, plot=False):
   print("Recalling memories...")
 
   ## Load memory module and memory keys ##
@@ -13,16 +13,18 @@ def recall_reservoir(res_size, sim_time, plot=False):
     memory_keys, labels = pkl.load(f)
 
   ## Recall memories ##
-  recalled_memories = np.zeros((len(memory_keys), sim_time, res_size))
+  # TODO: Plot output spikes according to inh exc populations
+  recalled_memories = np.zeros((len(memory_keys), sim_time, exc_size + inh_size))
   recalled_memories_sorted = {}
   for i, (key, label) in enumerate(zip(memory_keys, labels)):
-    res_spike_train = res_module.recall(torch.tensor(key.reshape(sim_time, -1)), sim_time=sim_time)  # Recall the sample
-    recalled_memories[i] = res_spike_train.squeeze()  # Store the recalled memory
+    exc_spikes, inh_spikes = res_module.recall(torch.tensor(key.reshape(sim_time, -1)), sim_time=sim_time)  # Recall the sample
+    all_spikes = torch.cat((exc_spikes, inh_spikes), dim=2).squeeze()
+    recalled_memories[i] = all_spikes  # Store the recalled memory
     label = tuple(label.round())
     if label not in recalled_memories_sorted:
-      recalled_memories_sorted[label] = [res_spike_train.squeeze()]
+      recalled_memories_sorted[label] = [all_spikes]
     else:
-      recalled_memories_sorted[label].append(res_spike_train.squeeze())
+      recalled_memories_sorted[label].append(all_spikes)
 
   ## Save recalled memories ##
   with open('Data/recalled_memories.pkl', 'wb') as f:

scripts/Chris/DQN/recall_reservoir.py

@@ -5,37 +5,45 @@
 import numpy as np
 from matplotlib import pyplot as plt
 
-def store_reservoir(res_size, num_samples, num_grid_cells, gc_multiples, sim_time,
+def store_reservoir(exc_size, inh_size, num_samples, num_grid_cells, gc_multiples, sim_time,
                     hyper_params, plot=False):
   print("Storing memories...")
 
   ## Create synaptic weights ##
   in_size = num_grid_cells * gc_multiples
-  w_in_res = torch.rand(in_size, res_size)
-  w_res_res = torch.rand(res_size, res_size)
-  w_in_res = sparsify(w_in_res, 0.85)
-  w_res_res = sparsify(w_res_res, 0.85)
-  w_res_res = assign_inhibition(w_res_res, 0.2, 1)
-  res = Reservoir(in_size, res_size, hyper_params, w_in_res, w_res_res)
+  w_in_exc = torch.rand(in_size, exc_size)    # Initialize weights
+  w_in_inh = torch.rand(in_size, inh_size)
+  w_exc_exc = torch.rand(exc_size, exc_size)
+  w_exc_inh = torch.rand(exc_size, inh_size)
+  w_inh_exc = -torch.rand(inh_size, exc_size)
+  w_inh_inh = torch.rand(inh_size, inh_size)
+  w_in_exc = sparsify(w_in_exc, 0.85)   # 0 x% of weights
+  w_in_inh = sparsify(w_in_inh, 0.85)
+  w_exc_exc = sparsify(w_exc_exc, 0.85)
+  w_exc_inh = sparsify(w_exc_inh, 0.85)
+  w_inh_exc = sparsify(w_inh_exc, 0.85)
+  w_inh_inh = sparsify(w_inh_inh, 0.85)
+  res = Reservoir(in_size, exc_size, inh_size, hyper_params,
+                  w_in_exc, w_in_inh, w_exc_exc, w_exc_inh, w_inh_exc, w_inh_inh)
 
   ## Load grid cell spike-train samples ##
   with open('Data/grid_cell_spk_trains.pkl', 'rb') as f:
     grid_cell_data, labels = pkl.load(f)  # (samples, time, num_cells)
 
   ## Store memories ##
   # -> STDP active
-  if plot:
-    fig, ax = plt.subplots(2, 2, figsize=(10, 5))
-    im = ax[0, 0].imshow(w_in_res)
-    ax[0, 0].set_title("Initial Input-to-Res")
-    plt.colorbar(im, ax=ax[0, 0])
-    ax[0, 0].set_xlabel("Res Neuron")
-    ax[0, 0].set_ylabel("Input Neuron")
-    im = ax[0, 1].imshow(w_res_res)
-    ax[0, 1].set_title("Initial Res-to-Res")
-    plt.colorbar(im, ax=ax[0, 1])
-    ax[0, 1].set_xlabel("Res Neuron")
-    ax[0, 1].set_ylabel("Res Neuron")
+  # if plot:
+  #   fig, ax = plt.subplots(2, 2, figsize=(10, 5))
+  #   im = ax[0, 0].imshow(w_in_res)
+  #   ax[0, 0].set_title("Initial Input-to-Res")
+  #   plt.colorbar(im, ax=ax[0, 0])
+  #   ax[0, 0].set_xlabel("Res Neuron")
+  #   ax[0, 0].set_ylabel("Input Neuron")
+  #   im = ax[0, 1].imshow(w_res_res)
+  #   ax[0, 1].set_title("Initial Res-to-Res")
+  #   plt.colorbar(im, ax=ax[0, 1])
+  #   ax[0, 1].set_xlabel("Res Neuron")
+  #   ax[0, 1].set_ylabel("Res Neuron")
 
   # Store samples
   sample_inds = np.random.choice(len(grid_cell_data), num_samples, replace=False)
@@ -46,19 +54,19 @@ def store_reservoir(res_size, num_samples, num_grid_cells, gc_multiples, sim_tim
     res.store(torch.tensor(s.reshape(sim_time, -1)), sim_time=sim_time)
     res.reset_state_variables()
 
-  if plot:
-    im = ax[1, 0].imshow(w_in_res)
-    ax[1, 0].set_title("Final Input-to-Res")
-    plt.colorbar(im, ax=ax[1, 0])
-    ax[1, 0].set_xlabel("Res Neuron")
-    ax[1, 0].set_ylabel("Input Neuron")
-    im = ax[1, 1].imshow(w_res_res)
-    ax[1, 1].set_title("Final Res-to-Res")
-    plt.colorbar(im, ax=ax[1, 1])
-    ax[1, 1].set_xlabel("Res Neuron")
-    ax[1, 1].set_ylabel("Res Neuron")
-    plt.tight_layout()
-    plt.show()
+  # if plot:
+  #   im = ax[1, 0].imshow(w_in_res)
+  #   ax[1, 0].set_title("Final Input-to-Res")
+  #   plt.colorbar(im, ax=ax[1, 0])
+  #   ax[1, 0].set_xlabel("Res Neuron")
+  #   ax[1, 0].set_ylabel("Input Neuron")
+  #   im = ax[1, 1].imshow(w_res_res)
+  #   ax[1, 1].set_title("Final Res-to-Res")
+  #   plt.colorbar(im, ax=ax[1, 1])
+  #   ax[1, 1].set_xlabel("Res Neuron")
+  #   ax[1, 1].set_ylabel("Res Neuron")
+  #   plt.tight_layout()
+  #   plt.show()
 
   ## Save ##
   with open('Data/reservoir_module.pkl', 'wb') as f: