Neurons to ensemble connections with interneurons

.nengobones.yml

@@ -282,8 +282,6 @@ setup_cfg:
         inaccurate
 
       # builder inconsistencies
-      test_connection.py:test_neurons_to_ensemble*:
-        transform shape not implemented
       test_connection.py:test_transform_probe:
         transform shape not implemented
       test_connection.py:test_list_indexing*:

CHANGES.rst

@@ -29,6 +29,8 @@ Release history
   in a round-robin format. This allocator can be selected using the
   ``hardware_options`` argument when creating ``nengo_loihi.Simulator``.
   (`#197 <https://github.com/nengo/nengo-loihi/pull/197>`__)
+- Added support for ``Ensemble.neurons -> Ensemble`` connections.
+  (`#156 <https://github.com/nengo/nengo-loihi/pull/156>`__)
 
 **Changed**
 

nengo_loihi/builder/connection.py

@@ -375,6 +375,23 @@ def build_no_solver(model, solver, conn, rng, sampled_transform):
     return _build_no_solver(*args)
 
 
+def expand_to_2d(weights, pre_size, post_size):
+    if weights.ndim == 0:
+        assert pre_size == post_size
+        weights2d = weights * np.eye(pre_size)
+    elif weights.ndim == 1:
+        assert pre_size == post_size
+        assert weights.size == pre_size
+        weights2d = np.diag(weights)
+    else:
+        assert weights.ndim == 2
+        weights2d = weights
+
+    assert weights2d.shape[0] == post_size
+    assert weights2d.shape[1] == pre_size
+    return weights2d
+
+
 def build_chip_connection(model, conn):  # noqa: C901
     if nengo_transforms is not None:
         if isinstance(conn.transform, nengo_transforms.Convolution):
@@ -409,24 +426,16 @@ def build_chip_connection(model, conn):  # noqa: C901
 
     needs_decode_neurons = False
     target_encoders = None
-    if isinstance(conn.pre_obj, Node):
+    if (isinstance(conn.pre_obj, Node)
+            and not isinstance(conn.pre_obj, ChipReceiveNeurons)):
         assert conn.pre_slice == slice(None)
 
-        if np.array_equal(transform, np.array(1.)):
-            # TODO: this identity transform may be avoidable
-            transform = np.eye(conn.pre.size_out)
-        else:
-            assert transform.ndim == 2, "transform shape not handled yet"
-            assert transform.shape[1] == conn.pre.size_out
+        transform = expand_to_2d(transform, conn.pre.size_out,
+                                 conn.post.size_in)
 
-        assert transform.shape[1] == conn.pre.size_out
-        if isinstance(conn.pre_obj, ChipReceiveNeurons):
-            weights = transform / model.dt
-            neuron_type = conn.pre_obj.neuron_type
-        else:
-            # input is on-off neuron encoded, so double/flip transform
-            weights = np.column_stack([transform, -transform])
-            target_encoders = 'node_encoders'
+        # input is on-off neuron encoded, so double/flip transform
+        weights = np.column_stack([transform, -transform])
+        target_encoders = 'node_encoders'
     elif (isinstance(conn.pre_obj, Ensemble)
           and isinstance(conn.pre_obj.neuron_type, nengo.Direct)):
         raise NotImplementedError()
@@ -459,11 +468,16 @@ def build_chip_connection(model, conn):  # noqa: C901
             post_slice = None
         else:
             needs_decode_neurons = True
-    elif isinstance(conn.pre_obj, Neurons):
+    elif isinstance(conn.pre_obj, (Neurons, ChipReceiveNeurons)):
         assert conn.pre_slice == slice(None)
-        assert transform.ndim == 2, "transform shape not handled yet"
-        weights = transform / model.dt
-        neuron_type = conn.pre_obj.ensemble.neuron_type
+        weights = expand_to_2d(transform, conn.pre.size_out, conn.post.size_in)
+        weights = weights / model.dt
+        neuron_type = (conn.pre_obj.neuron_type
+                       if isinstance(conn.pre_obj, ChipReceiveNeurons)
+                       else conn.pre_obj.ensemble.neuron_type)
+
+        if isinstance(conn.post_obj, Ensemble):
+            needs_decode_neurons = True
     else:
         raise NotImplementedError("Connection from type %r" % (
             type(conn.pre_obj),))

nengo_loihi/builder/tests/test_connection.py

@@ -5,6 +5,8 @@
 import numpy as np
 import pytest
 
+from nengo_loihi.builder.connection import expand_to_2d
+
 
 @pytest.mark.skipif(LooseVersion(nengo.__version__) <= LooseVersion('2.8.0'),
                     reason="requires more recent Nengo version")
@@ -66,3 +68,9 @@ def test_manual_decoders(
         assert np.mean(sim.data[pre_probe]) > 100
         # But that post has no activity due to the zero weights
         assert np.all(sim.data[post_probe] == 0)
+
+
+def test_expand_to_2d():
+    assert np.allclose(expand_to_2d(np.array(2.0), 3, 3), np.eye(3) * 2)
+    assert np.allclose(expand_to_2d(np.arange(3), 3, 3), np.diag(np.arange(3)))
+    assert np.allclose(expand_to_2d(np.ones((2, 3)), 3, 2), np.ones((2, 3)))

nengo_loihi/tests/test_connection.py

@@ -172,6 +172,61 @@ def test_ensemble_to_neurons(Simulator, seed, allclose, plt):
                     atol=5)
 
 
+@pytest.mark.parametrize("pre_on_chip, post_ensemble", [(True, True),
+                                                        (True, False),
+                                                        (False, True)])
+def test_neurons_to_ensemble_transform(
+        pre_on_chip, post_ensemble, Simulator, seed, rng, allclose, plt):
+    with nengo.Network(seed=seed) as net:
+        add_params(net)
+
+        stim = nengo.Node(lambda t: [np.sin(t * 2 * np.pi)])
+
+        n_pre = 50
+        pre_encoders = np.ones((n_pre, 1))
+        pre_encoders[n_pre // 2:] *= -1
+        pre = nengo.Ensemble(n_pre, 1, encoders=pre_encoders)
+        net.config[pre].on_chip = pre_on_chip
+        nengo.Connection(stim, pre, synapse=None)
+
+        n_post = 51
+        pre_decoders = pre_encoders.T / (100 * n_pre / 2)
+        post = (nengo.Ensemble(n_post, 1) if post_ensemble
+                else nengo.Node(size_in=1))
+
+        nengo.Connection(pre.neurons, post, transform=pre_decoders,
+                         synapse=0.005)
+
+        p_pre = nengo.Probe(pre, synapse=nengo.synapses.Alpha(0.03))
+        p_post = nengo.Probe(post, synapse=nengo.synapses.Alpha(0.03))
+
+    with nengo.Simulator(net) as nengosim:
+        nengosim.run(1.0)
+
+    with Simulator(net) as sim:
+        sim.run(1.0)
+
+    y0 = nengo.synapses.Lowpass(0.01).filt(nengosim.data[p_post].sum(axis=1))
+    y1 = sim.data[p_post].sum(axis=1)
+
+    t = sim.trange()
+    plt.subplot(2, 1, 1)
+    plt.plot(t, nengosim.data[p_pre], c='k')
+    plt.plot(t, sim.data[p_pre], c='g')
+    plt.ylim([-1, 1])
+    plt.ylabel("Decoded pre value")
+    plt.xlabel("Time (s)")
+
+    plt.subplot(2, 1, 2)
+    plt.plot(t, y0, c='k')
+    plt.plot(t, y1, c='g')
+    plt.ylim([-1, 1])
+    plt.ylabel("Decoded post value")
+    plt.xlabel("Time (s)")
+
+    assert allclose(y1, y0, rtol=1e-1, atol=0.1 * y0.max())
+
+
 def test_dists(Simulator, seed):
     # check that distributions on connection transforms are handled correctly
 

setup.cfg

@@ -225,8 +225,6 @@ nengo_test_unsupported =
         "inaccurate"
     test_actionselection.py:test_thalamus
         "inaccurate"
-    test_connection.py:test_neurons_to_ensemble*
-        "transform shape not implemented"
     test_connection.py:test_transform_probe
         "transform shape not implemented"
     test_connection.py:test_list_indexing*