Support array indexing for multi-synaptic indices

brian2/synapses/synapses.py

@@ -474,7 +474,7 @@ def slice_to_test(x):
         pass
 
     if isinstance(x, slice):
-        if isinstance(x, slice) and x == slice(None):
+        if x == slice(None):
             # No need for testing
             return lambda y: np.repeat(True, len(y))
         start, stop, step = x.start, x.stop, x.step
@@ -668,20 +668,18 @@ def __call__(self, index=None, index_var="_idx"):
                         "'multisynaptic_index' when you create "
                         "the Synapses object."
                     )
+
+                # We want to access the raw arrays here, not go through the Variable
+                synapse_numbers = self.synapse_number.get_value()[matching_synapses]
+
                 if isinstance(k_indices, (numbers.Integral, slice)):
                     test_k = slice_to_test(k_indices)
+                    final_indices = matching_synapses[test_k(synapse_numbers)]
                 else:
-                    raise NotImplementedError(
-                        "Indexing synapses with arrays notimplemented yet"
-                    )
+                    final_indices = matching_synapses[
+                        np.in1d(synapse_numbers, k_indices)
+                    ]
 
-                # We want to access the raw arrays here, not go through the Variable
-                synapse_numbers = self.synapse_number.get_value()[matching_synapses]
-                final_indices = np.intersect1d(
-                    matching_synapses,
-                    np.flatnonzero(test_k(synapse_numbers)),
-                    assume_unique=True,
-                )
         else:
             raise IndexError(f"Unsupported index type {type(index)}")
 

brian2/tests/test_synapses.py

@@ -721,9 +721,11 @@ def test_state_variable_indexing():
     assert len(S.w[:]) == len(S.w[np.arange(len(G1) * len(G2) * 2)])
     assert S.w[3] == S.w[np.int32(3)] == S.w[np.int64(3)]  # See issue #888
 
-    # Array-indexing (not yet supported for synapse index)
+    # Array-indexing
     assert_equal(S.w[:, 0:3], S.w[:, [0, 1, 2]])
     assert_equal(S.w[:, 0:3], S.w[np.arange(len(G1)), [0, 1, 2]])
+    assert_equal(S.w[:, 0:3], S.w[:, [0, 1, 2], [0, 1]])
+    assert_equal(S.w[:, 0:3, 0], S.w[:, [0, 1, 2], [0]])
 
     # string-based indexing
     assert_equal(S.w[0:3, :], S.w["i<3"])
@@ -780,6 +782,8 @@ def test_state_variable_indexing_with_subgroups():
     # Array-indexing (not yet supported for synapse index)
     assert_equal(S[:, 0:3].w[:], S[:, [0, 1, 2]].w[:])
     assert_equal(S[:, 0:3].w[:], S[np.arange(len(G1)), [0, 1, 2]].w[:])
+    assert_equal(S[:, 0:3].w[:], S[:, [0, 1, 2], [0, 1]].w[:])
+    assert_equal(S[:, 0:3, 0].w[:], S[:, [0, 1, 2], [0]].w[:])
 
     # string-based indexing
     assert_equal(S[0:3, :].w[:], S["i<3"].w[:])