Source code for brainpy._src.math.pre_syn_post

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


import jax.numpy as jnp
from jax import vmap, jit, ops as jops

from brainpy._src.math.interoperability import as_jax
from brainpy._src.math import event
from brainpy.errors import MathError

__all__ = [
  # pre-to-post
  'pre2post_sum',
  'pre2post_prod',
  'pre2post_max',
  'pre2post_min',
  'pre2post_mean',

  # pre-to-post event operator
  'pre2post_event_sum',
  'pre2post_csr_event_sum',

  # pre-to-syn
  'pre2syn',

  # syn-to-post
  'syn2post_sum', 'syn2post',
  'syn2post_prod',
  'syn2post_max',
  'syn2post_min',
  'syn2post_mean',
  'syn2post_softmax',
]


def _raise_pre_ids_is_none(pre_ids):
  if pre_ids is None:
    raise MathError(f'pre2post synaptic computation needs "pre_ids" '
                    f'when providing heterogeneous "pre_values" '
                    f'(brainpy.math.ndim(pre_values) != 0).')


[docs] def pre2post_event_sum(events, pre2post, post_num: int, values=1.): """The pre-to-post event-driven synaptic summation with `CSR` synapse structure. When ``values`` is a scalar, this function is equivalent to .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) post_ids, idnptr = pre2post for i in range(pre_num): if events[i]: for j in range(idnptr[i], idnptr[i+1]): post_val[post_ids[i]] += values When ``values`` is a vector (with the length of ``len(post_ids)``), this function is equivalent to .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) post_ids, idnptr = pre2post for i in range(pre_num): if events[i]: for j in range(idnptr[i], idnptr[i+1]): post_val[post_ids[i]] += values[j] Parameters ---------- events: ArrayType The events, must be bool. pre2post: tuple of ArrayType, tuple of ArrayType A tuple contains the connection information of pre-to-post. post_num: int The number of post-synaptic group. values: float, ArrayType The value to make summation. Returns ------- out: ArrayType A tensor with the shape of ``post_num``. """ indices, idnptr = pre2post events = as_jax(events) indices = as_jax(indices) idnptr = as_jax(idnptr) values = as_jax(values) return event.csrmv(values, indices, idnptr, events, shape=(events.shape[0], post_num), transpose=True)
pre2post_csr_event_sum = pre2post_event_sum
[docs] def pre2post_sum(pre_values, post_num, post_ids, pre_ids=None): """The pre-to-post synaptic summation. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for i, j in zip(pre_ids, post_ids): post_val[j] += pre_values[pre_ids[i]] Parameters ---------- pre_values: float, ArrayType The pre-synaptic values. post_ids: ArrayType The connected post-synaptic neuron ids. post_num: int Output dimension. The number of post-synaptic neurons. pre_ids: optional, ArrayType The connected pre-synaptic neuron ids. Returns ------- post_val: ArrayType The value with the size of post-synaptic neurons. """ out = jnp.zeros(post_num) pre_values = as_jax(pre_values) post_ids = as_jax(post_ids) if jnp.ndim(pre_values) != 0: _raise_pre_ids_is_none(pre_ids) pre_ids = as_jax(pre_ids) pre_values = pre_values[pre_ids] return out.at[post_ids].add(pre_values)
[docs] def pre2post_prod(pre_values, post_num, post_ids, pre_ids=None): """The pre-to-post synaptic production. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for i, j in zip(pre_ids, post_ids): post_val[j] *= pre_values[pre_ids[i]] Parameters ---------- pre_values: float, ArrayType The pre-synaptic values. pre_ids: ArrayType The connected pre-synaptic neuron ids. post_ids: ArrayType The connected post-synaptic neuron ids. post_num: int Output dimension. The number of post-synaptic neurons. Returns ------- post_val: ArrayType The value with the size of post-synaptic neurons. """ out = jnp.zeros(post_num) pre_values = as_jax(pre_values) post_ids = as_jax(post_ids) if jnp.ndim(pre_values) != 0: _raise_pre_ids_is_none(pre_ids) pre_ids = as_jax(pre_ids) pre_values = pre_values[pre_ids] return out.at[post_ids].multiply(pre_values)
[docs] def pre2post_min(pre_values, post_num, post_ids, pre_ids=None): """The pre-to-post synaptic minimization. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for i, j in zip(pre_ids, post_ids): post_val[j] = np.minimum(post_val[j], pre_values[pre_ids[i]]) Parameters ---------- pre_values: float, ArrayType The pre-synaptic values. pre_ids: ArrayType The connected pre-synaptic neuron ids. post_ids: ArrayType The connected post-synaptic neuron ids. post_num: int Output dimension. The number of post-synaptic neurons. Returns ------- post_val: ArrayType The value with the size of post-synaptic neurons. """ out = jnp.zeros(post_num) pre_values = as_jax(pre_values) post_ids = as_jax(post_ids) if jnp.ndim(pre_values) != 0: _raise_pre_ids_is_none(pre_ids) pre_ids = as_jax(pre_ids) pre_values = pre_values[pre_ids] return out.at[post_ids].min(pre_values)
[docs] def pre2post_max(pre_values, post_num, post_ids, pre_ids=None): """The pre-to-post synaptic maximization. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for i, j in zip(pre_ids, post_ids): post_val[j] = np.maximum(post_val[j], pre_values[pre_ids[i]]) Parameters ---------- pre_values: float, ArrayType The pre-synaptic values. pre_ids: ArrayType The connected pre-synaptic neuron ids. post_ids: ArrayType The connected post-synaptic neuron ids. post_num: int Output dimension. The number of post-synaptic neurons. Returns ------- post_val: ArrayType The value with the size of post-synaptic neurons. """ out = jnp.zeros(post_num) pre_values = as_jax(pre_values) post_ids = as_jax(post_ids) if jnp.ndim(pre_values) != 0: _raise_pre_ids_is_none(pre_ids) pre_ids = as_jax(pre_ids) pre_values = pre_values[pre_ids] return out.at[post_ids].max(pre_values)
[docs] def pre2post_mean(pre_values, post_num, post_ids, pre_ids=None): """The pre-to-post synaptic mean computation. Parameters ---------- pre_values: float, ArrayType The pre-synaptic values. pre_ids: ArrayType The connected pre-synaptic neuron ids. post_ids: ArrayType The connected post-synaptic neuron ids. post_num: int Output dimension. The number of post-synaptic neurons. Returns ------- post_val: ArrayType The value with the size of post-synaptic neurons. """ out = jnp.zeros(post_num) pre_values = as_jax(pre_values) post_ids = as_jax(post_ids) if jnp.ndim(pre_values) == 0: return out.at[post_ids].set(pre_values) # return out.at[jnp.unique(post_ids)].set(pre_values) else: _raise_pre_ids_is_none(pre_ids) pre_ids = as_jax(pre_ids) pre_values = pre2syn(pre_values, pre_ids) return syn2post_mean(pre_values, post_ids, post_num)
_pre2syn = vmap(lambda pre_id, pre_vs: pre_vs[pre_id], in_axes=(0, None))
[docs] def pre2syn(pre_values, pre_ids): """The pre-to-syn computation. Change the pre-synaptic data to the data with the dimension of synapses. This function is equivalent to: .. highlight:: python .. code-block:: python syn_val = np.zeros(len(pre_ids)) for syn_i, pre_i in enumerate(pre_ids): syn_val[i] = pre_values[pre_i] Parameters ---------- pre_values: float, ArrayType The pre-synaptic value. pre_ids: ArrayType The pre-synaptic neuron index. Returns ------- syn_val: ArrayType The synaptic value. """ pre_values = as_jax(pre_values) pre_ids = as_jax(pre_ids) if jnp.ndim(pre_values) == 0: return jnp.ones(len(pre_ids), dtype=pre_values.dtype) * pre_values else: return _pre2syn(pre_ids, pre_values)
_jit_seg_sum = jit(jops.segment_sum, static_argnums=(2, 3)) _jit_seg_prod = jit(jops.segment_prod, static_argnums=(2, 3)) _jit_seg_max = jit(jops.segment_max, static_argnums=(2, 3)) _jit_seg_min = jit(jops.segment_min, static_argnums=(2, 3))
[docs] def syn2post_sum(syn_values, post_ids, post_num: int, indices_are_sorted=False): """The syn-to-post summation computation. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for syn_i, post_i in enumerate(post_ids): post_val[post_i] += syn_values[syn_i] Parameters ---------- syn_values: ArrayType The synaptic values. post_ids: ArrayType The post-synaptic neuron ids. post_num: int The number of the post-synaptic neurons. Returns ------- post_val: ArrayType The post-synaptic value. """ post_ids = as_jax(post_ids) syn_values = as_jax(syn_values) if syn_values.dtype == jnp.bool_: syn_values = jnp.asarray(syn_values, dtype=jnp.int32) return _jit_seg_sum(syn_values, post_ids, post_num, indices_are_sorted)
syn2post = syn2post_sum
[docs] def syn2post_prod(syn_values, post_ids, post_num: int, indices_are_sorted=False): """The syn-to-post product computation. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for syn_i, post_i in enumerate(post_ids): post_val[post_i] *= syn_values[syn_i] Parameters ---------- syn_values: ArrayType The synaptic values. post_ids: ArrayType The post-synaptic neuron ids. If ``post_ids`` is generated by ``brainpy.conn.TwoEndConnector``, then it has sorted indices. Otherwise, this function cannot guarantee indices are sorted. You's better set ``indices_are_sorted=False``. post_num: int The number of the post-synaptic neurons. indices_are_sorted: whether ``post_ids`` is known to be sorted. Returns ------- post_val: ArrayType The post-synaptic value. """ post_ids = as_jax(post_ids) syn_values = as_jax(syn_values) if syn_values.dtype == jnp.bool_: syn_values = jnp.asarray(syn_values, dtype=jnp.int32) return _jit_seg_prod(syn_values, post_ids, post_num, indices_are_sorted)
[docs] def syn2post_max(syn_values, post_ids, post_num: int, indices_are_sorted=False): """The syn-to-post maximum computation. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for syn_i, post_i in enumerate(post_ids): post_val[post_i] = np.maximum(post_val[post_i], syn_values[syn_i]) Parameters ---------- syn_values: ArrayType The synaptic values. post_ids: ArrayType The post-synaptic neuron ids. If ``post_ids`` is generated by ``brainpy.conn.TwoEndConnector``, then it has sorted indices. Otherwise, this function cannot guarantee indices are sorted. You's better set ``indices_are_sorted=False``. post_num: int The number of the post-synaptic neurons. indices_are_sorted: whether ``post_ids`` is known to be sorted. Returns ------- post_val: ArrayType The post-synaptic value. """ post_ids = as_jax(post_ids) syn_values = as_jax(syn_values) if syn_values.dtype == jnp.bool_: syn_values = jnp.asarray(syn_values, dtype=jnp.int32) return _jit_seg_max(syn_values, post_ids, post_num, indices_are_sorted)
[docs] def syn2post_min(syn_values, post_ids, post_num: int, indices_are_sorted=False): """The syn-to-post minimization computation. This function is equivalent to: .. highlight:: python .. code-block:: python post_val = np.zeros(post_num) for syn_i, post_i in enumerate(post_ids): post_val[post_i] = np.minimum(post_val[post_i], syn_values[syn_i]) Parameters ---------- syn_values: ArrayType The synaptic values. post_ids: ArrayType The post-synaptic neuron ids. If ``post_ids`` is generated by ``brainpy.conn.TwoEndConnector``, then it has sorted indices. Otherwise, this function cannot guarantee indices are sorted. You's better set ``indices_are_sorted=False``. post_num: int The number of the post-synaptic neurons. indices_are_sorted: whether ``post_ids`` is known to be sorted. Returns ------- post_val: ArrayType The post-synaptic value. """ post_ids = as_jax(post_ids) syn_values = as_jax(syn_values) if syn_values.dtype == jnp.bool_: syn_values = jnp.asarray(syn_values, dtype=jnp.int32) return _jit_seg_min(syn_values, post_ids, post_num, indices_are_sorted)
[docs] def syn2post_mean(syn_values, post_ids, post_num: int, indices_are_sorted=False): """The syn-to-post mean computation. Parameters ---------- syn_values: ArrayType The synaptic values. post_ids: ArrayType The post-synaptic neuron ids. If ``post_ids`` is generated by ``brainpy.conn.TwoEndConnector``, then it has sorted indices. Otherwise, this function cannot guarantee indices are sorted. You's better set ``indices_are_sorted=False``. post_num: int The number of the post-synaptic neurons. indices_are_sorted: whether ``post_ids`` is known to be sorted. Returns ------- post_val: ArrayType The post-synaptic value. """ post_ids = as_jax(post_ids) syn_values = as_jax(syn_values) if syn_values.dtype == jnp.bool_: syn_values = jnp.asarray(syn_values, dtype=jnp.int32) nominator = _jit_seg_sum(syn_values, post_ids, post_num, indices_are_sorted) denominator = _jit_seg_sum(jnp.ones_like(syn_values), post_ids, post_num, indices_are_sorted) return jnp.nan_to_num(nominator / denominator)
[docs] def syn2post_softmax(syn_values, post_ids, post_num: int, indices_are_sorted=False): """The syn-to-post softmax computation. Parameters ---------- syn_values: ArrayType The synaptic values. post_ids: ArrayType The post-synaptic neuron ids. If ``post_ids`` is generated by ``brainpy.conn.TwoEndConnector``, then it has sorted indices. Otherwise, this function cannot guarantee indices are sorted. You's better set ``indices_are_sorted=False``. post_num: int The number of the post-synaptic neurons. indices_are_sorted: whether ``post_ids`` is known to be sorted. Returns ------- post_val: ArrayType The post-synaptic value. """ post_ids = as_jax(post_ids) syn_values = as_jax(syn_values) if syn_values.dtype == jnp.bool_: syn_values = jnp.asarray(syn_values, dtype=jnp.int32) syn_maxs = _jit_seg_max(syn_values, post_ids, post_num, indices_are_sorted) syn_values = syn_values - syn_maxs[post_ids] syn_values = jnp.exp(syn_values) normalizers = _jit_seg_sum(syn_values, post_ids, post_num, indices_are_sorted) softmax = syn_values / normalizers[post_ids] return jnp.nan_to_num(softmax)