feat: compatible with dict and tuple return types

Author: ztqakita

brainpy/connect/base.py

@@ -228,20 +228,22 @@ def _return_by_ij(self, structures, ij: tuple, all_data: dict):
       csr = ij2csr(pre_ids, post_ids, self.pre_num)
       self._return_by_csr(structures, csr=csr, all_data=all_data)
 
-  def make_returns(self, structures, conn_type, conn_data):
+  def make_returns(self, structures, conn_data, csr=None, mat=None, ij=None):
     """Make the desired synaptic structures and return them.
     """
-    csr = None
-    mat = None
-    ij = None
-    if conn_type == 'csr':
-      csr = conn_data
-    elif conn_type == 'mat':
-      mat = conn_data
-    elif conn_type == 'ij':
-      ij = conn_data
-    else:
-      raise ConnectorError(f'conn_type must be one of "csr", "mat" or "ij", but we got "{conn_type}" instead.')
+    if isinstance(conn_data, dict):
+      csr = conn_data['csr']
+      mat = conn_data['mat']
+      ij = conn_data['ij']
+    elif isinstance(conn_data, tuple):
+      if conn_data[0] == 'csr':
+        csr = conn_data[1]
+      elif conn_data[0] == 'mat':
+        mat = conn_data[1]
+      elif conn_data[0] == 'ij':
+        ij = conn_data[1]
+      else:
+        raise ConnectorError(f'Must provide one of "csr", "mat" or "ij". Got "{conn_data[0]}" instead.')
 
     # checking
     all_data = dict()
@@ -281,12 +283,21 @@ def make_returns(self, structures, conn_type, conn_data):
       return tuple([all_data[n] for n in structures])
 
   def build_conn(self):
+    """build connections with certain data type.
+
+    Returns
+    -------
+    A tuple with two elements: connection type (str) and connection data.
+      example: return 'csr', (ind, indptr)
+    Or a dict with three elements: csr, mat and ij.
+      example: return dict(csr=(ind, indptr), mat=None, ij=None)
+    """
     raise NotImplementedError
 
   def require(self, *structures):
     self.check(structures)
-    conn_type, conn_data = self.build_conn()
-    return self.make_returns(structures, conn_type, conn_data)
+    conn_data = self.build_conn()
+    return self.make_returns(structures, conn_data)
 
   def requires(self, *structures):
     return self.require(*structures)

brainpy/connect/regular_conn.py

@@ -39,7 +39,7 @@ def build_conn(self):
     ind = np.arange(self.pre_num)
     indptr = np.arange(self.pre_num + 1)
 
-    return 'csr', (ind, indptr)
+    return dict(csr=(ind, indptr), mat=None, ij=None)
 
 
 one2one = One2One()
@@ -60,7 +60,7 @@ def build_conn(self):
     if not self.include_self:
       np.fill_diagonal(mat, False)
 
-    return 'mat', mat
+    return dict(csr=None, mat=mat, ij=None)
 
 
 all2all = All2All(include_self=True)