Source code for brainpy._src.connect.base

# -*- coding: utf-8 -*-

import abc
from typing import Union, List, Tuple

import jax.numpy as jnp
import numpy as onp

from brainpy import tools, math as bm
from brainpy.errors import ConnectorError

import textwrap

__all__ = [
  # the connection types
  'CONN_MAT',
  'PRE_IDS', 'POST_IDS',
  'PRE2POST', 'POST2PRE',
  'PRE2SYN', 'POST2SYN',
  'SUPPORTED_SYN_STRUCTURE',

  # the connection dtypes
  'set_default_dtype', 'MAT_DTYPE', 'IDX_DTYPE', 'get_idx_type',

  # brainpy_object class
  'Connector', 'TwoEndConnector', 'OneEndConnector',

  # methods
  'mat2coo', 'mat2csc', 'mat2csr',
  'csr2csc', 'csr2mat', 'csr2coo',
  'coo2csr', 'coo2csc', 'coo2mat',
  'coo2mat_num', 'mat2mat_num',

  # visualize
  'visualizeMat',
]

CONN_MAT = 'conn_mat'
PRE_IDS = 'pre_ids'
POST_IDS = 'post_ids'
PRE2POST = 'pre2post'
POST2PRE = 'post2pre'
PRE2SYN = 'pre2syn'
POST2SYN = 'post2syn'
PRE_SLICE = 'pre_slice'
POST_SLICE = 'post_slice'
COO = 'coo'
CSR = 'csr'
CSC = 'csc'

SUPPORTED_SYN_STRUCTURE = [CONN_MAT,
                           PRE_IDS, POST_IDS,
                           PRE2POST, POST2PRE,
                           PRE2SYN, POST2SYN,
                           PRE_SLICE, POST_SLICE,
                           COO, CSR, CSC]

MAT_DTYPE = jnp.bool_
IDX_DTYPE = jnp.int32


def get_idx_type():
  return IDX_DTYPE


[docs] def set_default_dtype(mat_dtype=None, idx_dtype=None): """Set the default dtype. Use this method, you can set the default dtype for connetion matrix and connection index. For examples: >>> import numpy as np >>> import brainpy as bp >>> >>> conn = bp.conn.GridFour()(4, 4) >>> conn.require('conn_mat') Array([[False, True, False, False], [ True, False, True, False], [False, True, False, True], [False, False, True, False]], dtype=bool) >>> bp.connect.set_default_dtype(mat_dtype=np.float32) >>> conn = bp.conn.GridFour()(4, 4) >>> conn.require('conn_mat') Array([[0., 1., 0., 0.], [1., 0., 1., 0.], [0., 1., 0., 1.], [0., 0., 1., 0.]], dtype=float32) Parameters ---------- mat_dtype : type The default dtype for connection matrix. idx_dtype : type The default dtype for connection index. """ if mat_dtype is not None: global MAT_DTYPE MAT_DTYPE = mat_dtype if idx_dtype is not None: global IDX_DTYPE IDX_DTYPE = idx_dtype
[docs] class Connector(abc.ABC): """Base Synaptic Connector Class.""" pass
[docs] class TwoEndConnector(Connector): """Synaptic connector to build connections between two neuron groups. If users want to customize their `Connector`, there are two ways: 1. Implementing ``build_conn(self)`` function, which returns one of the connection data ``csr`` (CSR sparse data, a tuple of <post_ids, inptr>), ``coo`` (COO sparse data, a tuple of <pre_ids, post_ids>), or ``mat`` (a binary connection matrix). For instance, .. code-block:: python import brainpy as bp class MyConnector(bp.conn.TwoEndConnector): def build_conn(self): return dict(csr=, mat=, coo=) 2. Implementing functions ``build_mat()``, ``build_csr()``, and ``build_coo()``. Users can provide all three functions, or one of them. .. code-block:: python import brainpy as bp class MyConnector(bp.conn.TwoEndConnector): def build_mat(self, ): return conn_matrix def build_csr(self, ): return post_ids, inptr def build_coo(self, ): return pre_ids, post_ids """ def __init__( self, pre: Union[int, Tuple[int, ...]] = None, post: Union[int, Tuple[int, ...]] = None, ): self.pre_size = None self.post_size = None self.pre_num = None self.post_num = None if pre is not None: if isinstance(pre, int): pre = (pre,) else: pre = tuple(pre) self.pre_size = pre self.pre_num = tools.size2num(self.pre_size) if post is not None: if isinstance(post, int): post = (post,) else: post = tuple(post) self.post_size = post self.post_num = tools.size2num(self.post_size) def __repr__(self): return self.__class__.__name__ def __call__(self, pre_size, post_size): """Create the concrete connections between two end objects. Parameters ---------- pre_size : int, tuple of int, list of int The size of the pre-synaptic group. post_size : int, tuple of int, list of int The size of the post-synaptic group. Returns ------- conn : TwoEndConnector Return the self. """ if isinstance(pre_size, int): pre_size = (pre_size,) else: pre_size = tuple(pre_size) if isinstance(post_size, int): post_size = (post_size,) else: post_size = tuple(post_size) self.pre_size, self.post_size = pre_size, post_size self.pre_num = tools.size2num(self.pre_size) self.post_num = tools.size2num(self.post_size) return self def _reset_conn(self, pre_size, post_size): """Reset connection attributes. Parameters ---------- pre_size : int, tuple of int, list of int The size of the pre-synaptic group. post_size : int, tuple of int, list of int The size of the post-synaptic group. """ self.__call__(pre_size, post_size) @property def is_version2_style(self): if ((hasattr(self.build_coo, 'not_customized') and self.build_coo.not_customized) and (hasattr(self.build_csr, 'not_customized') and self.build_csr.not_customized) and (hasattr(self.build_mat, 'not_customized') and self.build_mat.not_customized)): return False else: return True def _check(self, structures: Union[Tuple, List, str]): # check synaptic structures if isinstance(structures, str): structures = [structures] if structures is None or len(structures) == 0: raise ConnectorError('No synaptic structure is received.') for n in structures: if n not in SUPPORTED_SYN_STRUCTURE: raise ConnectorError(f'Unknown synapse structure "{n}". ' f'Only {SUPPORTED_SYN_STRUCTURE} is supported.') def _return_by_mat(self, structures, mat, all_data: dict): assert mat.ndim == 2 if (CONN_MAT in structures) and (CONN_MAT not in all_data): all_data[CONN_MAT] = bm.as_jax(mat, dtype=MAT_DTYPE) if len([s for s in structures if s not in [CONN_MAT]]) > 0: ij = mat2coo(mat) self._return_by_coo(structures, coo=ij, all_data=all_data) def _return_by_csr(self, structures, csr: tuple, all_data: dict): indices, indptr = csr np = onp if isinstance(indices, onp.ndarray) else bm assert self.pre_num == indptr.size - 1 if (CONN_MAT in structures) and (CONN_MAT not in all_data): conn_mat = csr2mat((indices, indptr), self.pre_num, self.post_num) all_data[CONN_MAT] = bm.as_jax(conn_mat, dtype=MAT_DTYPE) if (PRE_IDS in structures) and (PRE_IDS not in all_data): pre_ids = np.repeat(np.arange(self.pre_num), np.diff(indptr)) all_data[PRE_IDS] = bm.as_jax(pre_ids, dtype=get_idx_type()) if (POST_IDS in structures) and (POST_IDS not in all_data): all_data[POST_IDS] = bm.as_jax(indices, dtype=get_idx_type()) if (COO in structures) and (COO not in all_data): pre_ids = np.repeat(np.arange(self.pre_num), np.diff(indptr)) all_data[COO] = (bm.as_jax(pre_ids, dtype=get_idx_type()), bm.as_jax(indices, dtype=get_idx_type())) if (PRE2POST in structures) and (PRE2POST not in all_data): all_data[PRE2POST] = (bm.as_jax(indices, dtype=get_idx_type()), bm.as_jax(indptr, dtype=get_idx_type())) if (CSR in structures) and (CSR not in all_data): all_data[CSR] = (bm.as_jax(indices, dtype=get_idx_type()), bm.as_jax(indptr, dtype=get_idx_type())) if (POST2PRE in structures) and (POST2PRE not in all_data): indc, indptrc = csr2csc((indices, indptr), self.post_num) all_data[POST2PRE] = (bm.as_jax(indc, dtype=get_idx_type()), bm.as_jax(indptrc, dtype=get_idx_type())) if (CSC in structures) and (CSC not in all_data): indc, indptrc = csr2csc((indices, indptr), self.post_num) all_data[CSC] = (bm.as_jax(indc, dtype=get_idx_type()), bm.as_jax(indptrc, dtype=get_idx_type())) if (PRE2SYN in structures) and (PRE2SYN not in all_data): syn_seq = np.arange(indices.size, dtype=get_idx_type()) all_data[PRE2SYN] = (bm.as_jax(syn_seq, dtype=get_idx_type()), bm.as_jax(indptr, dtype=get_idx_type())) if (POST2SYN in structures) and (POST2SYN not in all_data): syn_seq = np.arange(indices.size, dtype=get_idx_type()) _, indptrc, syn_seqc = csr2csc((indices, indptr), self.post_num, syn_seq) all_data[POST2SYN] = (bm.as_jax(syn_seqc, dtype=get_idx_type()), bm.as_jax(indptrc, dtype=get_idx_type())) def _return_by_coo(self, structures, coo: tuple, all_data: dict): pre_ids, post_ids = coo if (CONN_MAT in structures) and (CONN_MAT not in all_data): all_data[CONN_MAT] = bm.as_jax(coo2mat(coo, self.pre_num, self.post_num), dtype=MAT_DTYPE) if (PRE_IDS in structures) and (PRE_IDS not in all_data): all_data[PRE_IDS] = bm.as_jax(pre_ids, dtype=get_idx_type()) if (POST_IDS in structures) and (POST_IDS not in all_data): all_data[POST_IDS] = bm.as_jax(post_ids, dtype=get_idx_type()) if (COO in structures) and (COO not in all_data): all_data[COO] = (bm.as_jax(pre_ids, dtype=get_idx_type()), bm.as_jax(post_ids, dtype=get_idx_type())) if CSC in structures and CSC not in all_data: csc = coo2csc(coo, self.post_num) all_data[CSC] = (bm.as_jax(csc[0], dtype=get_idx_type()), bm.as_jax(csc[1], dtype=get_idx_type())) if POST2PRE in structures and POST2PRE not in all_data: csc = coo2csc(coo, self.post_num) all_data[POST2PRE] = (bm.as_jax(csc[0], dtype=get_idx_type()), bm.as_jax(csc[1], dtype=get_idx_type())) if (len([s for s in structures if s not in [CONN_MAT, PRE_IDS, POST_IDS, COO, CSC, POST2PRE]]) > 0): csr = coo2csr(coo, self.pre_num) self._return_by_csr(structures, csr=csr, all_data=all_data) def _make_returns(self, structures, conn_data): """Make the desired synaptic structures and return them. """ csr = None mat = None coo = None if isinstance(conn_data, dict): csr = conn_data.get('csr', None) mat = conn_data.get('mat', None) coo = conn_data.get('coo', None) or conn_data.get('ij', None) 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] in ['coo', 'ij']: coo = conn_data[1] else: raise ConnectorError(f'Must provide one of "csr", "mat" or "coo". Got "{conn_data[0]}" instead.') else: raise ConnectorError('Unknown type') # checking if (csr is None) and (mat is None) and (coo is None): raise ConnectorError('Must provide one of "csr", "mat" or "coo".') structures = (structures,) if isinstance(structures, str) else structures assert isinstance(structures, (tuple, list)) all_data = dict() # "csr" structure if csr is not None: if (PRE2POST in structures) and (PRE2POST not in all_data): all_data[PRE2POST] = (bm.as_jax(csr[0], dtype=get_idx_type()), bm.as_jax(csr[1], dtype=get_idx_type())) self._return_by_csr(structures, csr=csr, all_data=all_data) # "mat" structure if mat is not None: assert mat.ndim == 2 if (CONN_MAT in structures) and (CONN_MAT not in all_data): all_data[CONN_MAT] = bm.as_jax(mat, dtype=MAT_DTYPE) self._return_by_mat(structures, mat=mat, all_data=all_data) # "coo" structure if coo is not None: if (PRE_IDS in structures) and (PRE_IDS not in structures): all_data[PRE_IDS] = bm.as_jax(coo[0], dtype=get_idx_type()) if (POST_IDS in structures) and (POST_IDS not in structures): all_data[POST_IDS] = bm.as_jax(coo[1], dtype=get_idx_type()) self._return_by_coo(structures, coo=coo, all_data=all_data) # return if len(structures) == 1: return all_data[structures[0]] else: return tuple([all_data[n] for n in structures])
[docs] def require(self, *structures): """Require all the connection data needed. Examples -------- >>> import brainpy as bp >>> conn = bp.connect.FixedProb(0.1) >>> mat = conn.require(10, 20, 'conn_mat') >>> mat.shape (10, 20) """ if len(structures) > 0: pre_size = None post_size = None if not isinstance(structures[0], str): pre_size = structures[0] structures = structures[1:] if len(structures) > 0: if not isinstance(structures[0], str): post_size = structures[0] structures = structures[1:] if pre_size is not None: self.__call__(pre_size, post_size) else: return tuple() if self.pre_num is None or self.post_num is None: raise ConnectorError(f'self.pre_num or self.post_num is not defined. ' f'Please use "self.require(pre_size, post_size, DATA1, DATA2, ...)" ') _has_coo_imp = not hasattr(self.build_coo, 'not_customized') _has_csr_imp = not hasattr(self.build_csr, 'not_customized') _has_mat_imp = not hasattr(self.build_mat, 'not_customized') self._check(structures) if (_has_coo_imp or _has_csr_imp or _has_mat_imp): if len(structures) == 1: if PRE2POST in structures and _has_csr_imp: r = self.build_csr() return bm.as_jax(r[0], dtype=get_idx_type()), bm.as_jax(r[1], dtype=get_idx_type()) elif CSR in structures and _has_csr_imp: r = self.build_csr() return bm.as_jax(r[0], dtype=get_idx_type()), bm.as_jax(r[1], dtype=get_idx_type()) elif CONN_MAT in structures and _has_mat_imp: return bm.as_jax(self.build_mat(), dtype=MAT_DTYPE) elif PRE_IDS in structures and _has_coo_imp: return bm.as_jax(self.build_coo()[0], dtype=get_idx_type()) elif POST_IDS in structures and _has_coo_imp: return bm.as_jax(self.build_coo()[1], dtype=get_idx_type()) elif COO in structures and _has_coo_imp: r = self.build_coo() return bm.as_jax(r[0], dtype=get_idx_type()), bm.as_jax(r[1], dtype=get_idx_type()) elif len(structures) == 2: if (PRE_IDS in structures and POST_IDS in structures and _has_coo_imp): r = self.build_coo() if structures[0] == PRE_IDS: return bm.as_jax(r[0], dtype=get_idx_type()), bm.as_jax(r[1], dtype=get_idx_type()) else: return bm.as_jax(r[1], dtype=get_idx_type()), bm.as_jax(r[0], dtype=get_idx_type()) if ((CSR in structures or PRE2POST in structures) and _has_csr_imp and COO in structures and _has_coo_imp): csr = self.build_csr() csr = (bm.as_jax(csr[0], dtype=get_idx_type()), bm.as_jax(csr[1], dtype=get_idx_type())) coo = self.build_coo() coo = (bm.as_jax(coo[0], dtype=get_idx_type()), bm.as_jax(coo[1], dtype=get_idx_type())) if structures[0] == COO: return coo, csr else: return csr, coo if ((CSR in structures or PRE2POST in structures) and _has_csr_imp and CONN_MAT in structures and _has_mat_imp): csr = self.build_csr() csr = (bm.as_jax(csr[0], dtype=get_idx_type()), bm.as_jax(csr[1], dtype=get_idx_type())) mat = bm.as_jax(self.build_mat(), dtype=MAT_DTYPE) if structures[0] == CONN_MAT: return mat, csr else: return csr, mat if (COO in structures and _has_coo_imp and CONN_MAT in structures and _has_mat_imp): coo = self.build_coo() coo = (bm.as_jax(coo[0], dtype=get_idx_type()), bm.as_jax(coo[1], dtype=get_idx_type())) mat = bm.as_jax(self.build_mat(), dtype=MAT_DTYPE) if structures[0] == COO: return coo, mat else: return mat, coo conn_data = dict(csr=None, ij=None, mat=None) if _has_coo_imp: conn_data['coo'] = self.build_coo() # if (CSR in structures or PRE2POST in structures) and _has_csr_imp: # conn_data['csr'] = self.build_csr() # if CONN_MAT in structures and _has_mat_imp: # conn_data['mat'] = self.build_mat() elif _has_csr_imp: conn_data['csr'] = self.build_csr() # if COO in structures and _has_coo_imp: # conn_data['coo'] = self.build_coo() # if CONN_MAT in structures and _has_mat_imp: # conn_data['mat'] = self.build_mat() elif _has_mat_imp: conn_data['mat'] = self.build_mat() # if COO in structures and _has_coo_imp: # conn_data['coo'] = self.build_coo() # if (CSR in structures or PRE2POST in structures) and _has_csr_imp: # conn_data['csr'] = self.build_csr() else: raise ValueError else: conn_data = self.build_conn() return self._make_returns(structures, conn_data)
[docs] def requires(self, *structures): """Require all the connection data needed.""" return self.require(*structures)
[docs] @tools.not_customized def build_conn(self): """build connections with certain data type. If users want to customize their connections, please provide one of the following functions: - ``build_mat()``: build a matrix binary connection matrix. - ``build_csr()``: build a csr sparse connection data. - ``build_coo()``: build a coo sparse connection data. - ``build_conn()``: deprecated. Returns ------- conn: tuple, dict A tuple with two elements: connection type (str) and connection data. For example: ``return 'csr', (ind, indptr)`` Or a dict with three elements: csr, mat and coo. For example: ``return dict(csr=(ind, indptr), mat=None, coo=None)`` """ pass
[docs] @tools.not_customized def build_mat(self): """Build a binary matrix connection data. If users want to customize their connections, please provide one of the following functions: - ``build_mat()``: build a matrix binary connection matrix. - ``build_csr()``: build a csr sparse connection data. - ``build_coo()``: build a coo sparse connection data. - ``build_conn()``: deprecated. Returns ------- conn: Array A binary matrix with the shape ``(num_pre, num_post)``. """ pass
[docs] @tools.not_customized def build_csr(self): """Build a csr sparse connection data. Returns ------- conn: tuple A tuple denoting the ``(indices, indptr)``. """ pass
[docs] @tools.not_customized def build_coo(self): """Build a coo sparse connection data. Returns ------- conn: tuple A tuple denoting the ``(pre_ids, post_ids)``. """ pass
[docs] class OneEndConnector(TwoEndConnector): """Synaptic connector to build synapse connections within a population of neurons.""" def __init__(self, *args, **kwargs): super(OneEndConnector, self).__init__(*args, **kwargs) def __call__(self, pre_size, post_size=None): if post_size is None: post_size = pre_size try: assert pre_size == post_size except AssertionError: raise ConnectorError( f'The shape of pre-synaptic group should be the same with the post group. ' f'But we got {pre_size} != {post_size}.') if isinstance(pre_size, int): pre_size = (pre_size,) else: pre_size = tuple(pre_size) if isinstance(post_size, int): post_size = (post_size,) else: post_size = tuple(post_size) self.pre_size, self.post_size = pre_size, post_size self.pre_num = tools.size2num(self.pre_size) self.post_num = tools.size2num(self.post_size) return self def _reset_conn(self, pre_size, post_size=None): self.__init__() self.__call__(pre_size, post_size)
[docs] def mat2csr(dense): """convert a dense matrix to (indices, indptr).""" if isinstance(dense, onp.ndarray): pre_ids, post_ids = onp.where(dense > 0) else: pre_ids, post_ids = jnp.where(bm.as_jax(dense) > 0) return coo2csr((pre_ids, post_ids), dense.shape[0])
[docs] def mat2coo(dense): if isinstance(dense, onp.ndarray): pre_ids, post_ids = onp.where(dense > 0) else: pre_ids, post_ids = jnp.where(bm.as_jax(dense) > 0) return pre_ids.astype(dtype=get_idx_type()), post_ids.astype(dtype=get_idx_type())
[docs] def mat2csc(dense): if isinstance(dense, onp.ndarray): pre_ids, post_ids = onp.where(dense > 0) else: pre_ids, post_ids = jnp.where(bm.as_jax(dense) > 0) return coo2csr((post_ids, pre_ids), dense.shape[1])
[docs] def csr2mat(csr, num_pre, num_post): """convert (indices, indptr) to a dense matrix.""" indices, indptr = csr if isinstance(indices, onp.ndarray): d = onp.zeros((num_pre, num_post), dtype=MAT_DTYPE) # num_pre, num_post pre_ids = onp.repeat(onp.arange(indptr.size - 1), onp.diff(indptr)) d[pre_ids, indices] = True return d else: d = bm.zeros((num_pre, num_post), dtype=MAT_DTYPE) # num_pre, num_post pre_ids = jnp.repeat(jnp.arange(indptr.size - 1), jnp.diff(indptr)) d[pre_ids, indices] = True return d.value
[docs] def csr2csc(csr, post_num, data=None): """Convert csr to csc.""" return coo2csc(csr2coo(csr), post_num, data)
[docs] def csr2coo(csr): np = onp if isinstance(csr[0], onp.ndarray) else jnp indices, indptr = csr pre_ids = np.repeat(np.arange(indptr.size - 1), np.diff(indptr)) return pre_ids, indices
[docs] def coo2mat(ij, num_pre, num_post): """convert (indices, indptr) to a dense matrix.""" pre_ids, post_ids = ij if isinstance(pre_ids, onp.ndarray): d = onp.zeros((num_pre, num_post), dtype=MAT_DTYPE) # num_pre, num_post d[pre_ids, post_ids] = True return d else: d = bm.zeros((num_pre, num_post), dtype=MAT_DTYPE) d[pre_ids, post_ids] = True return d.value
[docs] def coo2csr(coo, num_pre): """convert pre_ids, post_ids to (indices, indptr) when'jax_platform_name' = 'gpu'""" pre_ids, post_ids = coo if isinstance(pre_ids, onp.ndarray): sort_ids = onp.argsort(pre_ids) post_ids = onp.asarray(post_ids) post_ids = post_ids[sort_ids] indices = post_ids unique_pre_ids, pre_count = onp.unique(pre_ids, return_counts=True) final_pre_count = onp.zeros(num_pre, dtype=jnp.uint32) final_pre_count[unique_pre_ids] = pre_count else: sort_ids = onp.argsort(bm.as_jax(pre_ids)) post_ids = bm.as_jax(post_ids) post_ids = post_ids[sort_ids] indices = post_ids unique_pre_ids, pre_count = jnp.unique(pre_ids, return_counts=True) final_pre_count = bm.zeros(num_pre, dtype=jnp.uint32) final_pre_count[unique_pre_ids] = pre_count final_pre_count = bm.as_jax(final_pre_count) indptr = final_pre_count.cumsum() indptr = onp.insert(indptr, 0, 0) return indices.astype(get_idx_type()), indptr.astype(get_idx_type())
[docs] def coo2csc(coo, post_num, data=None): """Convert csr to csc.""" pre_ids, indices = coo if isinstance(indices, onp.ndarray): # to maintain the original order of the elements with the same value sort_ids = onp.argsort(indices) pre_ids_new = onp.asarray(pre_ids[sort_ids], dtype=get_idx_type()) unique_post_ids, count = onp.unique(indices, return_counts=True) post_count = onp.zeros(post_num, dtype=get_idx_type()) post_count[unique_post_ids] = count indptr_new = post_count.cumsum() indptr_new = onp.insert(indptr_new, 0, 0) indptr_new = onp.asarray(indptr_new, dtype=get_idx_type()) else: pre_ids = bm.as_jax(pre_ids) indices = bm.as_jax(indices) # to maintain the original order of the elements with the same value sort_ids = jnp.argsort(indices) pre_ids_new = jnp.asarray(pre_ids[sort_ids], dtype=get_idx_type()) unique_post_ids, count = jnp.unique(indices, return_counts=True) post_count = bm.zeros(post_num, dtype=get_idx_type()) post_count[unique_post_ids] = count indptr_new = post_count.value.cumsum() indptr_new = jnp.insert(indptr_new, 0, 0) indptr_new = jnp.asarray(indptr_new, dtype=get_idx_type()) if data is None: return pre_ids_new, indptr_new else: data_new = data[sort_ids] return pre_ids_new, indptr_new, data_new
[docs] def coo2mat_num(ij, num_pre, num_post, num, seed=0): """ convert (indices, indptr) to a dense connection number matrix.\n Specific for FixedTotalNum. """ rng = bm.random.RandomState(seed) mat = coo2mat(ij, num_pre, num_post) # get nonzero indices and number nonzero_idx = jnp.nonzero(mat) nonzero_num = jnp.count_nonzero(mat) # get multi connection number multi_conn_num = num - nonzero_num # alter the element type to int mat = mat.astype(jnp.int32) # 随机在mat中选取nonzero_idx的元素,将其值加1 index = rng.choice(nonzero_num, size=(multi_conn_num,), replace=False) for i in index: mat = mat.at[nonzero_idx[0][i], nonzero_idx[1][i]].set(mat[nonzero_idx[0][i], nonzero_idx[1][i]] + 1) return mat
[docs] def mat2mat_num(mat, num, seed=0): """ Convert boolean matrix to a dense connection number matrix.\n Specific for FixedTotalNum. """ rng = bm.random.RandomState(seed) # get nonzero indices and number nonzero_idx = jnp.nonzero(mat) nonzero_num = jnp.count_nonzero(mat) # get multi connection number multi_conn_num = num - nonzero_num # alter the element type to int mat = mat.astype(jnp.int32) # 随机在mat中选取nonzero_idx的元素,将其值加1 index = rng.choice(nonzero_num, size=(multi_conn_num,), replace=False) for i in index: mat = mat.at[nonzero_idx[0][i], nonzero_idx[1][i]].set(mat[nonzero_idx[0][i], nonzero_idx[1][i]] + 1) return mat
[docs] def visualizeMat(mat, description='Untitled'): """ Visualize the matrix. (Need seaborn and matplotlib) parameters ---------- mat : jnp.ndarray The matrix to be visualized. description : str The title of the figure. """ try: import seaborn as sns import matplotlib.pyplot as plt except (ModuleNotFoundError, ImportError): print('Please install seaborn and matplotlib for this function') return sns.heatmap(mat, cmap='viridis') warpped_title = textwrap.fill(description, width=60) plt.title(warpped_title) plt.show()