brainstate 0.0.1__py2.py3-none-any.whl

brainstate/nn/_projection/_delta.py

@@ -0,0 +1,241 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Optional, Union
+
+from brainstate._module import (Module, Dynamics, Projection, ReceiveInputProj,
+                                UpdateReturn, register_delay_of_target)
+from brainstate._utils import set_module_as
+from brainstate.mixin import (Mode, BindCondData)
+from ._utils import is_instance
+
+__all__ = [
+  'HalfProjDelta', 'FullProjDelta',
+]
+
+
+class _Delta(BindCondData):
+  def __init__(self):
+    self._cond = None
+
+  def bind_cond(self, cond):
+    self._cond = cond
+
+  def __call__(self, *args, **kwargs):
+    r = self._cond
+    return r
+
+
+@set_module_as('brainstate.nn')
+class HalfProjDelta(Projection):
+  """Defining the half-part of the synaptic projection for the Delta synapse model.
+
+  The synaptic projection requires the input is the spiking data, otherwise
+  the synapse is not the Delta synapse model.
+
+  The ``half-part`` means that the model only includes ``comm`` -> ``syn`` -> ``out`` -> ``post``.
+  Therefore, the model's ``update`` function needs the manual providing of the spiking input.
+
+  **Model Descriptions**
+
+  .. math::
+
+      I_{syn} (t) = \sum_{j\in C} g_{\mathrm{max}} * \delta(t-t_j-D)
+
+  where :math:`g_{\mathrm{max}}` denotes the chemical synaptic strength,
+  :math:`t_j` the spiking moment of the presynaptic neuron :math:`j`,
+  :math:`C` the set of neurons connected to the post-synaptic neuron,
+  and :math:`D` the transmission delay of chemical synapses.
+  For simplicity, the rise and decay phases of post-synaptic currents are
+  omitted in this model.
+
+
+  **Code Examples**
+
+  .. code-block::
+
+      import brainstate as bp
+      import brainstate.math as bm
+
+      class Net(bp.DynamicalSystem):
+        def __init__(self):
+          super().__init__()
+
+          self.pre = bp.dyn.PoissonGroup(10, 100.)
+          self.post = bp.dyn.LifRef(1)
+          self.syn = bp.dyn.HalfProjDelta(bp.dnn.Linear(10, 1, bp.init.OneInit(2.)), self.post)
+
+        def update(self):
+          self.syn(self.pre())
+          self.post()
+          return self.post.V.value
+
+      net = Net()
+      indices = bm.arange(1000).to_numpy()
+      vs = bm.for_loop(net.step_run, indices, progress_bar=True)
+      bp.visualize.line_plot(indices, vs, show=True)
+
+  Args:
+    comm: DynamicalSystem. The synaptic communication.
+    post: DynamicalSystem. The post-synaptic neuron group.
+    name: str. The projection name.
+    mode: Mode. The computing mode.
+  """
+
+  _invisible_nodes = ['post']
+
+  def __init__(
+      self,
+      comm: Module,
+      post: ReceiveInputProj,
+      name: Optional[str] = None,
+      mode: Optional[Mode] = None,
+  ):
+    super().__init__(name=name, mode=mode)
+
+    # synaptic models
+    is_instance(comm, Module)
+    is_instance(post, ReceiveInputProj)
+    self.comm = comm
+
+    # output initialization
+    out = _Delta()
+    post.add_input_fun(self.name, out, category='delta')
+    self.out = out
+
+    # references
+    self.post = post
+
+  def update(self, x):
+    # call the communication
+    current = self.comm(x)
+    # bind the output
+    self.out.bind_cond(current)
+    # return the current, if needed
+    return current
+
+
+@set_module_as('brainstate.nn')
+class FullProjDelta(Projection):
+  """Full-chain of the synaptic projection for the Delta synapse model.
+
+  The synaptic projection requires the input is the spiking data, otherwise
+  the synapse is not the Delta synapse model.
+
+  The ``full-chain`` means that the model needs to provide all information needed for a projection,
+  including ``pre`` -> ``delay`` -> ``comm`` -> ``post``.
+
+  **Model Descriptions**
+
+  .. math::
+
+      I_{syn} (t) = \sum_{j\in C} g_{\mathrm{max}} * \delta(t-t_j-D)
+
+  where :math:`g_{\mathrm{max}}` denotes the chemical synaptic strength,
+  :math:`t_j` the spiking moment of the presynaptic neuron :math:`j`,
+  :math:`C` the set of neurons connected to the post-synaptic neuron,
+  and :math:`D` the transmission delay of chemical synapses.
+  For simplicity, the rise and decay phases of post-synaptic currents are
+  omitted in this model.
+
+
+  **Code Examples**
+
+  .. code-block::
+
+      import brainstate as bp
+      import brainstate.math as bm
+
+
+      class Net(bp.DynamicalSystem):
+        def __init__(self):
+          super().__init__()
+
+          self.pre = bp.dyn.PoissonGroup(10, 100.)
+          self.post = bp.dyn.LifRef(1)
+          self.syn = bp.dyn.FullProjDelta(self.pre, 0., bp.dnn.Linear(10, 1, bp.init.OneInit(2.)), self.post)
+
+        def update(self):
+          self.syn()
+          self.pre()
+          self.post()
+          return self.post.V.value
+
+
+      net = Net()
+      indices = bm.arange(1000).to_numpy()
+      vs = bm.for_loop(net.step_run, indices, progress_bar=True)
+      bp.visualize.line_plot(indices, vs, show=True)
+
+
+  Args:
+    pre: The pre-synaptic neuron group.
+    delay: The synaptic delay.
+    comm: DynamicalSystem. The synaptic communication.
+    post: DynamicalSystem. The post-synaptic neuron group.
+    name: str. The projection name.
+    mode: Mode. The computing mode.
+  """
+
+  _invisible_nodes = ['pre', 'post', 'delay']
+
+  def __init__(
+      self,
+      pre: UpdateReturn,
+      delay: Union[None, int, float],
+      comm: Module,
+      post: Dynamics,
+      name: Optional[str] = None,
+      mode: Optional[Mode] = None,
+  ):
+    super().__init__(name=name, mode=mode)
+
+    # synaptic models
+    is_instance(pre, UpdateReturn)
+    is_instance(comm, Module)
+    is_instance(post, Dynamics)
+    self.comm = comm
+
+    # delay initialization
+    if delay is not None and delay > 0.:
+      delay_cls = register_delay_of_target(pre)
+      delay_cls.register_entry(self.name, delay)
+      self.delay = delay_cls
+      self.has_delay = True
+    else:
+      self.delay = None
+      self.has_delay = False
+
+    # output initialization
+    out = _Delta()
+    post.add_input_fun(self.name, out, category='delta')
+    self.out = out
+
+    # references
+    self.pre = pre
+    self.post = post
+
+  def update(self):
+    # get delay
+    if self.has_delay:
+      x = self.delay.at(self.name)
+    else:
+      x = self.pre.update_return()
+    # call the communication
+    current = self.comm(x)
+    # bind the output
+    self.out.bind_cond(current)
+    # return the current, if needed
+    return current

brainstate/nn/_projection/_utils.py

@@ -0,0 +1,17 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def is_instance(x, cls):
+  assert isinstance(x, cls), 'The input should be an instance of {}!'.format(cls)

brainstate/nn/_projection/_vanilla.py

@@ -0,0 +1,101 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Optional
+
+from brainstate._module import (Module, Projection, Dynamics, ReceiveInputProj)
+from brainstate._utils import set_module_as
+from brainstate.mixin import (BindCondData, Mode)
+from ._utils import is_instance
+
+__all__ = [
+  'VanillaProj',
+]
+
+
+@set_module_as('brainstate.nn')
+class VanillaProj(Projection):
+  """Synaptic projection which defines the synaptic computation with the dimension of pre-synaptic neuron group.
+
+  **Code Examples**
+
+  To simulate an E/I balanced network model:
+
+  .. code-block::
+
+      class EINet(bp.DynSysGroup):
+        def __init__(self):
+          super().__init__()
+          self.N = bp.dyn.LifRef(4000, V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.,
+                                 V_initializer=bp.init.Normal(-55., 2.))
+          self.delay = bp.VarDelay(self.N.spike, entries={'I': None})
+          self.syn1 = bp.dyn.Expon(size=3200, tau=5.)
+          self.syn2 = bp.dyn.Expon(size=800, tau=10.)
+          self.E = bp.dyn.VanillaProj(comm=bp.dnn.JitFPHomoLinear(3200, 4000, prob=0.02, weight=0.6),
+                                      out=bp.dyn.COBA(E=0.),
+                                      post=self.N)
+          self.I = bp.dyn.VanillaProj(comm=bp.dnn.JitFPHomoLinear(800, 4000, prob=0.02, weight=6.7),
+                                      out=bp.dyn.COBA(E=-80.),
+                                      post=self.N)
+
+        def update(self, input):
+          spk = self.delay.at('I')
+          self.E(self.syn1(spk[:3200]))
+          self.I(self.syn2(spk[3200:]))
+          self.delay(self.N(input))
+          return self.N.spike.value
+
+      model = EINet()
+      indices = bm.arange(1000)
+      spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices)
+      bp.visualize.raster_plot(indices, spks, show=True)
+
+
+  Args:
+    comm: The synaptic communication.
+    out: The synaptic output.
+    post: The post-synaptic neuron group.
+    name: str. The projection name.
+    mode: Mode. The computing mode.
+  """
+
+  _invisible_nodes = ['post', ]
+
+  def __init__(
+      self,
+      comm: Module,
+      out: BindCondData,
+      post: ReceiveInputProj,
+      name: Optional[str] = None,
+      mode: Optional[Mode] = None,
+  ):
+    super().__init__(name=name, mode=mode)
+
+    # synaptic models
+    is_instance(comm, Module)
+    is_instance(out, BindCondData)
+    is_instance(post, Dynamics)
+    self.comm = comm
+
+    # output initialization
+    post.add_input_fun(self.name, out)
+    self.out = out
+
+    self.post = post
+
+  def update(self, x):
+    current = self.comm(x)
+    self.out.bind_cond(current)
+    return current