Emulator tracks spikes individually

Will be required to allow population spikes (atoms).

nengo_loihi/builder.py

@@ -576,7 +576,7 @@ def build_connection(model, conn):
         weights = weights * neuron_type.amplitude
 
     mid_cx = pre_cx
-    mid_axon_inds = slice(None)
+    mid_axon_inds = None
     post_tau = tau_s
     if needs_interneurons and not isinstance(conn.post_obj, Neurons):
         # --- add interneurons
@@ -719,7 +719,7 @@ def build_connection(model, conn):
 
         mid_ax = CxAxons(mid_cx.n, label="encoders")
         mid_ax.target = post_cx.named_synapses['inter_encoders']
-        mid_ax.target_inds = mid_axon_inds
+        mid_ax.set_axon_map(mid_axon_inds)
         mid_cx.add_axons(mid_ax)
         model.objs[conn]['mid_axons'] = mid_ax
 

nengo_loihi/loihi_cx.py

@@ -423,14 +423,51 @@ def format(self, **kwargs):
 
 
 class CxAxons(object):
+    """A group of axons, targeting a specific CxSynapses object.
+
+    Parameters
+    ----------
+    n_axons : int
+        The number of outgoing axons.
+    group : CxGroup
+        Parent CxGroup for this object (set in `CxGroup.add_axons`).
+    target : CxSynapses
+        Target synapses for these axons.
+    cx_to_axon_map : list of length `group.n` (optional, default: None)
+        Index of the axon in `target` targeted by each group compartment.
+    cx_atoms : list of length `group.n` (optional, default: None)
+        Atom (weight index) associated with each group compartment.
+    """
+
+    class Spike(object):
+        """A spike, targeting a particular axon within a CxSynapses object.
+
+        The CxSynapses target is implicit, given by the CxAxons object that
+        creates this Spike.
+
+        Parameters
+        ----------
+        axon_id : int
+            The index of the axon within the targeted CxSynapses object.
+        atom : int
+            An index into the target CxSynapses weights. This allows spikes
+            targeting a particular axon to use different weights.
+        """
+
+        __slots__ = ['axon_id', 'atom']
+
+        def __init__(self, axon_id, atom=0):
+            self.axon_id = axon_id
+            self.atom = atom
+
     def __init__(self, n_axons, label=None):
         self.n_axons = n_axons
         self.label = label
         self.group = None
 
         self.target = None
-        self.target_inds = slice(None)  # which synapse inputs are targeted
-        # ^ TODO: this does not allow multiple pre-cx per axon, loihi does
+        self.cx_to_axon_map = None
+        self.cx_atoms = None
 
     def __str__(self):
         return "%s(%s)" % (
@@ -445,6 +482,24 @@ def axon_slots(self):
         """The total number of axonCfg slots used by all axons."""
         return self.slots_per_axon * self.n_axons
 
+    def set_axon_map(self, cx_to_axon_map, cx_atoms=None):
+        self.cx_to_axon_map = cx_to_axon_map
+        self.cx_atoms = cx_atoms
+
+    def map_cx_axons(self, cx_idxs):
+        return (self.cx_to_axon_map[cx_idxs]
+                if self.cx_to_axon_map is not None else cx_idxs)
+
+    def map_cx_atoms(self, cx_idxs):
+        return (self.cx_atoms[cx_idxs] if self.cx_atoms is not None else
+                [0 for _ in cx_idxs])
+
+    def map_cx_spikes(self, cx_idxs):
+        axon_ids = self.map_cx_axons(cx_idxs)
+        atoms = self.map_cx_atoms(cx_idxs)
+        return [self.Spike(axon_id, atom=atom) if axon_id >= 0 else None
+                for axon_id, atom in zip(axon_ids, atoms)]
+
 
 class CxProbe(object):
     _slice = slice
@@ -649,11 +704,11 @@ def overflow(x, bits, name=None):
         self.ref = np.hstack([group.refractDelay for group in self.groups])
 
         # --- allocate synapse memory
-        self.a_in = {synapses: np.zeros(synapses.n_axons, dtype=np.int32)
-                     for group in self.groups for synapses in group.synapses}
+        self.axons_in = {synapses: [] for group in self.groups
+                         for synapses in group.synapses}
         self.z = {synapses: np.zeros(synapses.n_axons, dtype=np.float64)
                   for group in self.groups for synapses in group.synapses
-                  if synapses.tracing}
+                  if synapses.tracing}  # synapse traces
 
         # --- noise
         enableNoise = np.hstack([
@@ -760,39 +815,37 @@ def step(self):  # noqa: C901
         self.q[:-1] = self.q[1:]  # advance delays
         self.q[-1] = 0
 
-        for a_in in self.a_in.values():
-            a_in[:] = 0
+        # --- clear spikes going in to each synapse
+        for axons_in_spikes in self.axons_in.values():
+            axons_in_spikes.clear()
 
+        # --- inputs pass spikes to synapses
         if self.t >= 1:  # input spikes take one time-step to arrive
             for input in self.inputs:
                 for axons in input.axons:
-                    synapses = axons.target
-                    assert axons.target_inds == slice(None)
-                    self.a_in[synapses] += input.spikes[self.t - 1]
+                    cx_idxs = input.spikes[self.t - 1].nonzero()[0]
+                    spikes = axons.map_cx_spikes(cx_idxs)
+                    self.axons_in[axons.target].extend(spikes)
 
+        # --- axons pass spikes to synapses
         for group in self.groups:
             for axons in group.axons:
-                synapses = axons.target
-                s_in = self.a_in[synapses]
-
-                a_slice = self.group_cxs[axons.group]
-                sa = self.s[a_slice]
-                np.add.at(s_in, axons.target_inds, sa)  # allows repeat inds
+                cx_idxs = self.s[self.group_cxs[axons.group]].nonzero()[0]
+                spikes = axons.map_cx_spikes(cx_idxs)
+                self.axons_in[axons.target].extend(spikes)
 
+        # --- synapse spikes use weights to modify compartment input
         for group in self.groups:
             for synapses in group.synapses:
-                s_in = self.a_in[synapses]
-
                 b_slice = self.group_cxs[synapses.group]
                 weights = synapses.weights
                 indices = synapses.indices
                 qb = self.q[:, b_slice]
                 # delays = np.zeros(qb.shape[1], dtype=np.int32)
 
-                for i in s_in.nonzero()[0]:
-                    for _ in range(s_in[i]):  # faster than mult since likely 1
-                        qb[0, indices[i]] += weights[i]
-                    # qb[delays[indices[i]], indices[i]] += weights[i]
+                for spike in self.axons_in[synapses]:
+                    assert spike.atom == 0
+                    qb[0, indices[spike.axon_id]] += weights[spike.axon_id]
 
                 if synapses.tracing:
                     z = self.z[synapses]
@@ -802,7 +855,8 @@ def step(self):  # noqa: C901
                     decay = np.exp(-1.0 / tau)
                     z *= decay
 
-                    z += mag * s_in
+                    for spike in self.axons_in[synapses]:
+                        z[spike.axon_id] += mag
 
         # --- updates
         q0 = self.q[0, :]

nengo_loihi/loihi_interface.py

@@ -343,16 +343,18 @@ def build_synapses(n2core, core, group, synapses, cx_idxs):
             )
 
 
-def build_axons(n2core, core, group, axons, cx_idxs):
+def build_axons(n2core, core, group, axons, cx_ids):
     tchip_idx, tcore_idx, tsyn_idxs = core.board.find_synapses(axons.target)
-    taxon_idxs = np.asarray(tsyn_idxs)[axons.target_inds]
     n2board = n2core.parent.parent
     tchip_id = n2board.n2Chips[tchip_idx].id
     tcore_id = n2board.n2Chips[tchip_idx].n2Cores[tcore_idx].id
-    assert axons.n_axons == len(cx_idxs) == len(taxon_idxs)
-    for i in range(axons.n_axons):
-        n2core.createDiscreteAxon(
-            cx_idxs[i], tchip_id, tcore_id, int(taxon_idxs[i]))
+
+    cx_idxs = np.arange(len(cx_ids))
+    spikes = axons.map_cx_spikes(cx_idxs)
+    for cx_id, spike in zip(cx_ids, spikes):
+        taxon_idx = int(spike.axon_id)
+        taxon_id = int(tsyn_idxs[taxon_idx])
+        n2core.createDiscreteAxon(cx_id, tchip_id, tcore_id, taxon_id)
 
 
 def build_probe(n2core, core, group, probe, cx_idxs):