brainpy.math.share as the global context for sharing data across all modules/nodes

[chaoming0625]

brainpy.math.share is designed to be a global context to communicate with all modules/nodes in a BrainPy program.

A snippet for its appilication:

import brainpy as bp
import brainpy.math as bm
import brainpy.connect as C
from brainpy.math.context import share


class EINetv2(bp.DynamicalSystem):
  def __init__(self, scale=1.0, e_input=20., i_input=20., delay=None):
    super().__init__()

    self.bg_exc = e_input
    self.bg_inh = i_input

    # network size
    num_exc = int(3200 * scale)
    num_inh = int(800 * scale)

    # neurons
    pars = dict(V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.,
                V_initializer=bp.init.Normal(-55., 2.))
    self.E = bp.experimental.LIF(num_exc, **pars)
    self.I = bp.experimental.LIF(num_inh, **pars)

    # synapses
    we = 0.6 / scale  # excitatory synaptic weight (voltage)
    wi = 6.7 / scale  # inhibitory synaptic weight
    self.E2E = bp.experimental.Exponential(
      C.FixedProb(0.02, pre=self.E.size, post=self.E.size),
      g_max=we, tau=5., out=bp.experimental.COBA(self.E.V, E=0.)
    )
    self.E2I = bp.experimental.Exponential(
      C.FixedProb(0.02, pre=self.E.size, post=self.I.size, ),
      g_max=we, tau=5., out=bp.experimental.COBA(self.I.V, E=0.)
    )
    self.I2E = bp.experimental.Exponential(
      C.FixedProb(0.02, pre=self.I.size, post=self.E.size),
      g_max=wi, tau=10., out=bp.experimental.COBA(self.E.V, E=-80.)
    )
    self.I2I = bp.experimental.Exponential(
      C.FixedProb(0.02, pre=self.I.size, post=self.I.size),
      g_max=wi, tau=10., out=bp.experimental.COBA(self.I.V, E=-80.)
    )
    delayE = bp.experimental.Delay(self.E.spike, entries={'E': delay})
    delayI = bp.experimental.Delay(self.I.spike, entries={'I': delay})

    share.save('t', 2.)
    share.save('dt', 0.1)
    share.save('E-spike', delayE)
    share.save('I-spike', delayI)

  @bp.not_pass_shargs
  def update(self):
    e_spike = share.load('E-spike').at('E')
    i_spike = share.load('I-spike').at('I')
    e_inp = self.E2E(e_spike) + self.I2E(i_spike) + self.bg_exc
    i_inp = self.I2I(i_spike) + self.E2I(e_spike) + self.bg_inh
    self.E(e_inp)
    self.I(i_inp)