brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/nn/_projection/_delta.py

@@ -1,241 +0,0 @@
-# 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/_vanilla.py

@@ -1,101 +0,0 @@
-# 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