brainstate 0.1.3__py2.py3-none-any.whl → 0.1.5__py2.py3-none-any.whl

brainstate/nn/_linear_mv.py

@@ -19,7 +19,7 @@ import brainunit as u
 import jax
 
 from brainstate import init
-from brainstate._compatible_import import brainevent
+import brainevent
 from brainstate._state import ParamState
 from brainstate.typing import Size, ArrayLike
 from ._module import Module

brainstate/nn/_module.py

@@ -34,7 +34,7 @@ import numpy as np
 from brainstate._state import State
 from brainstate.graph import Node, states, nodes, flatten
 from brainstate.mixin import ParamDescriber, ParamDesc
-from brainstate.typing import PathParts
+from brainstate.typing import PathParts, Size
 from brainstate.util import FlattedDict, NestedDict, BrainStateError
 
 # maximum integer
@@ -62,8 +62,8 @@ class Module(Node, ParamDesc):
 
     __module__ = 'brainstate.nn'
 
-    _in_size: Optional[Tuple[int, ...]]
-    _out_size: Optional[Tuple[int, ...]]
+    _in_size: Optional[Size]
+    _out_size: Optional[Size]
     _name: Optional[str]
 
     if not TYPE_CHECKING:
@@ -87,7 +87,7 @@ class Module(Node, ParamDesc):
         raise AttributeError('The name of the model is read-only.')
 
     @property
-    def in_size(self) -> Tuple[int, ...]:
+    def in_size(self) -> Size:
         return self._in_size
 
     @in_size.setter
@@ -98,7 +98,7 @@ class Module(Node, ParamDesc):
         self._in_size = tuple(in_size)
 
     @property
-    def out_size(self) -> Tuple[int, ...]:
+    def out_size(self) -> Size:
         return self._out_size
 
     @out_size.setter

brainstate/nn/_projection.py

@@ -13,25 +13,31 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Union, Callable, Optional
+from typing import Callable, Union
+from typing import Optional
+
+import brainevent
+import brainunit as u
 
 from brainstate._state import State
-from brainstate.mixin import AlignPost, ParamDescriber, BindCondData, JointTypes
-from brainstate.util._others import get_unique_name
+from brainstate.mixin import BindCondData, JointTypes
+from brainstate.mixin import ParamDescriber, AlignPost
+from brainstate.util.others import get_unique_name
 from ._collective_ops import call_order
-from ._dynamics import Dynamics, maybe_init_prefetch, Prefetch, PrefetchDelayAt
+from ._dynamics import Dynamics, Projection, maybe_init_prefetch, Prefetch, PrefetchDelayAt
 from ._module import Module
+from ._stp import ShortTermPlasticity
+from ._synapse import Synapse
 from ._synouts import SynOut
 
 __all__ = [
     'AlignPostProj',
     'DeltaProj',
     'CurrentProj',
-]
-
 
-class Interaction(Module):
-    __module__ = 'brainstate.nn'
+    'align_pre_projection',
+    'align_post_projection',
+]
 
 
 def _check_modules(*modules):
@@ -101,7 +107,7 @@ class _AlignPost(Module):
         self.out.bind_cond(self.syn(*args, **kwargs))
 
 
-class AlignPostProj(Interaction):
+class AlignPostProj(Projection):
     """
     Align-post projection of the neural network.
 
@@ -113,7 +119,6 @@ class AlignPostProj(Interaction):
     Note that this projection needs the manual input of pre-synaptic spikes.
 
     >>> import brainstate
-    >>> import brainevent
     >>> import brainunit as u
     >>> n_exc = 3200
     >>> n_inh = 800
@@ -124,16 +129,14 @@ class AlignPostProj(Interaction):
     ...        tau=20. * u.ms, tau_ref=5. * u.ms,
     ...        V_initializer=brainstate.init.Normal(-55., 2., unit=u.mV)
     ... )
-    >>> pop.reset_state()
+    >>> pop.init_state()
     >>> E = brainstate.nn.AlignPostProj(
-    ...        comm=brainevent.nn.FixedProb(n_exc, num, prob=80 / num, weight=1.62 * u.mS),
+    ...        comm=brainstate.nn.FixedNumConn(n_exc, num, prob=80 / num, weight=1.62 * u.mS),
     ...        syn=brainstate.nn.Expon.desc(num, tau=5. * u.ms),
     ...        out=brainstate.nn.CUBA.desc(scale=u.volt),
     ...        post=pop
     ... )
-    >>> exe_current = E(pop.spike.value)
-
-
+    >>> exe_current = E(pop.get_spike())
 
     """
     __module__ = 'brainstate.nn'
@@ -226,48 +229,47 @@ class AlignPostProj(Interaction):
             self.out.bind_cond(conductance)
 
 
-class DeltaProj(Interaction):
-    r"""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.
-
-    #     brainstate.nn.DeltaInteraction(
-    #       LIF().prefetch('V'), bst.surrogate.ReluGrad(), comm, post
-    #     )
-
-    Args:
-      pre: The pre-synaptic neuron group.
-      delay: The synaptic delay.
-      comm: DynamicalSystem. The synaptic communication.
-      post: DynamicalSystem. The post-synaptic neuron group.
+class DeltaProj(Projection):
     """
+    Delta-based projection of the neural network.
+
+    This projection directly applies delta inputs to post-synaptic neurons without intervening
+    synaptic dynamics. It processes inputs through optional prefetch modules, applies a communication model,
+    and adds the result directly as a delta input to the post-synaptic population.
+
+    Parameters
+    ----------
+    *prefetch : State or callable
+        Optional prefetch modules to process input before communication.
+    comm : callable
+        Communication model that determines how signals are transmitted.
+    post : Dynamics
+        Post-synaptic neural population to receive the delta inputs.
+    label : Optional[str], default=None
+        Optional label for the projection to identify it in the post-synaptic population.
 
+    Examples
+    --------
+    >>> import brainstate
+    >>> import brainunit as u
+    >>> n_neurons = 100
+    >>> pop = brainstate.nn.LIF(n_neurons, V_rest=-70*u.mV, V_threshold=-50*u.mV)
+    >>> pop.init_state()
+    >>> delta_input = brainstate.nn.DeltaProj(
+    ...     comm=lambda x: x * 10.0*u.mV,
+    ...     post=pop
+    ... )
+    >>> delta_input(1.0)  # Apply voltage increment directly
+    """
     __module__ = 'brainstate.nn'
 
-    def __init__(self, *modules, comm: Callable, post: Dynamics, label=None):
+    def __init__(self, *prefetch, comm: Callable, post: Dynamics, label=None):
         super().__init__(name=get_unique_name(self.__class__.__name__))
 
         self.label = label
 
         # checking modules
-        self.modules = _check_modules(*modules)
+        self.prefetches = _check_modules(*prefetch)
 
         # checking communication model
         if not callable(comm):
@@ -285,58 +287,69 @@ class DeltaProj(Interaction):
 
     @call_order(2)
     def init_state(self, *args, **kwargs):
-        for module in self.modules:
-            maybe_init_prefetch(module, *args, **kwargs)
+        for prefetch in self.prefetches:
+            maybe_init_prefetch(prefetch, *args, **kwargs)
 
     def update(self, *x):
-        for module in self.modules:
+        for module in self.prefetches:
             x = (call_module(module, *x),)
         assert len(x) == 1, f'The output of the modules should be a single value, but got {x}.'
         x = self.comm(x[0])
         self.post.add_delta_input(self.name, x, label=self.label)
 
 
-class CurrentProj(Interaction):
+class CurrentProj(Projection):
     """
-    Full-chain synaptic projection with the align-pre reduction and delay+synapse updating and merging.
-
-    The ``full-chain`` means that the model needs to provide all information needed for a projection,
-    including ``pre`` -> ``delay`` -> ``syn`` -> ``comm`` -> ``out`` -> ``post``.
-    Note here, compared to ``FullProjAlignPreSD``, the ``delay`` and ``syn`` are exchanged.
-
-    The ``align-pre`` means that the synaptic variables have the same dimension as the pre-synaptic neuron group.
-
-    The ``delay+synapse updating`` means that the projection first delivers the pre neuron output (usually the
-    spiking)  to the delay model, then computes the synapse states, and finally computes the synaptic current.
-
-    The ``merging`` means that the same delay model is shared by all synapses, and the synapse model with same
-    parameters (such like time constants) will also share the same synaptic variables.
+    Current-based projection of the neural network.
+
+    This projection directly modulates post-synaptic currents without separate synaptic dynamics.
+    It processes inputs through optional prefetch modules, applies a communication model,
+    and binds the result to the output model which is then added as a current input to the post-synaptic population.
+
+    Parameters
+    ----------
+    *prefetch : State or callable
+        Optional prefetch modules to process input before communication.
+        The last element must be an instance of Prefetch or PrefetchDelayAt if any are provided.
+    comm : callable
+        Communication model that determines how signals are transmitted.
+    out : SynOut
+        Output model that converts communication results to post-synaptic currents.
+    post : Dynamics
+        Post-synaptic neural population to receive the currents.
 
-    Neither ``FullProjAlignPreDS`` nor ``FullProjAlignPreSD`` facilitates the event-driven computation.
-    This is because the ``comm`` is computed after the synapse state, which is a floating-point number, rather
-    than the spiking. To facilitate the event-driven computation, please use align post projections.
-
-    Args:
-      prefetch: The synaptic dynamics.
-      comm: The synaptic communication.
-      out: The synaptic output.
-      post: The post-synaptic neuron group.
+    Examples
+    --------
+    >>> import brainstate
+    >>> import brainunit as u
+    >>> n_neurons = 100
+    >>> pop = brainstate.nn.LIF(n_neurons, V_rest=-70*u.mV, V_threshold=-50*u.mV)
+    >>> pop.init_state()
+    >>> current_input = brainstate.nn.CurrentProj(
+    ...     comm=lambda x: x * 0.5,
+    ...     out=brainstate.nn.CUBA(scale=1.0*u.nA),
+    ...     post=pop
+    ... )
+    >>> current_input(0.2)  # Apply external current
     """
     __module__ = 'brainstate.nn'
 
     def __init__(
         self,
-        prefetch: Union[Prefetch, PrefetchDelayAt],
+        *prefetch,
         comm: Callable,
         out: SynOut,
         post: Dynamics,
     ):
         super().__init__(name=get_unique_name(self.__class__.__name__))
 
-        # pre-synaptic neuron group
-        if not isinstance(prefetch, (Prefetch, PrefetchDelayAt)):
-            raise TypeError(f'The pre should be a Prefetch or PrefetchDelayAt, but got {prefetch}.')
+        # check prefetch
         self.prefetch = prefetch
+        if len(self.prefetch) > 0 and not isinstance(prefetch[-1], (Prefetch, PrefetchDelayAt)):
+            raise TypeError(
+                f'The last element of prefetch should be an instance of {Prefetch} or {PrefetchDelayAt}, '
+                f'but got {prefetch[-1]}.'
+            )
 
         # check out
         if not isinstance(out, SynOut):
@@ -354,41 +367,120 @@ class CurrentProj(Interaction):
 
     @call_order(2)
     def init_state(self, *args, **kwargs):
-        maybe_init_prefetch(self.prefetch, *args, **kwargs)
+        for prefetch in self.prefetch:
+            maybe_init_prefetch(prefetch, *args, **kwargs)
 
     def update(self, *x):
-        x = self.prefetch(*x)
-        x = self.comm(x)
+        for prefetch in self.prefetch:
+            x = (call_module(prefetch, *x),)
+        x = self.comm(*x)
         self.out.bind_cond(x)
 
 
-class RawProj(Interaction):
+class align_pre_projection(Projection):
     """
+    Represents a pre-synaptic alignment projection mechanism.
+
+    This class inherits from the `Projection` base class and is designed to
+    manage the pre-synaptic alignment process in neural network simulations.
+    It takes into account pre-synaptic dynamics, synaptic properties, delays,
+    communication functions, synaptic outputs, post-synaptic dynamics, and
+    short-term plasticity.
+
+    Attributes:
+        pre (Dynamics): The pre-synaptic dynamics object.
+        syn (Synapse): The synaptic object after pre-synaptic alignment.
+        delay (u.Quantity[u.second]): The output delay from the synapse.
+        projection (CurrentProj): The current projection object handling communication,
+            output, and post-synaptic dynamics.
+        stp (ShortTermPlasticity, optional): The short-term plasticity object,
+            defaults to None.
     """
-    __module__ = 'brainstate.nn'
 
     def __init__(
         self,
+        *spike_generator,
+        syn: Dynamics,
         comm: Callable,
         out: SynOut,
         post: Dynamics,
+        stp: ShortTermPlasticity = None,
     ):
-        super().__init__(name=get_unique_name(self.__class__.__name__))
+        super().__init__()
 
-        # check out
-        if not isinstance(out, SynOut):
-            raise TypeError(f'The out should be a SynOut, but got {out}.')
-        self.out = out
+        self.spike_generator = _check_modules(*spike_generator)
+        self.projection = CurrentProj(comm=comm, out=out, post=post)
+        self.syn = syn
+        self.stp = stp
 
-        # check post
-        if not isinstance(post, Dynamics):
-            raise TypeError(f'The post should be a Dynamics, but got {post}.')
-        self.post = post
-        post.add_current_input(self.name, out)
+    @call_order(2)
+    def init_state(self, *args, **kwargs):
+        for module in self.spike_generator:
+            maybe_init_prefetch(module, *args, **kwargs)
 
-        # output initialization
-        self.comm = comm
+    def update(self, *x):
+        for fun in self.spike_generator:
+            x = fun(*x)
+            if isinstance(x, (tuple, list)):
+                x = tuple(x)
+            else:
+                x = (x,)
+        assert len(x) == 1, "Spike generator must return a single value or a tuple/list of values"
+        x = brainevent.BinaryArray(x[0])  # Ensure input is a BinaryFloat for spike generation
+        if self.stp is not None:
+            x = brainevent.MaskedFloat(self.stp(x))  # Ensure STP output is a MaskedFloat
+        x = self.syn(x)  # Apply pre-synaptic alignment
+        return self.projection(x)
+
+
+class align_post_projection(Projection):
+    """
+    Represents a post-synaptic alignment projection mechanism.
 
-    def update(self, x):
-        x = self.comm(x)
-        self.out.bind_cond(x)
+    This class inherits from the `Projection` base class and is designed to
+    manage the post-synaptic alignment process in neural network simulations.
+    It takes into account spike generators, communication functions, synaptic
+    properties, synaptic outputs, post-synaptic dynamics, and short-term plasticity.
+
+    Args:
+        *spike_generator: Callable(s) that generate spike events or transform input spikes.
+        comm (Callable): Communication function for the projection.
+        syn (Union[AlignPost, ParamDescriber[AlignPost]]): The post-synaptic alignment object or its parameter describer.
+        out (Union[SynOut, ParamDescriber[SynOut]]): The synaptic output object or its parameter describer.
+        post (Dynamics): The post-synaptic dynamics object.
+        stp (ShortTermPlasticity, optional): The short-term plasticity object, defaults to None.
+
+    """
+
+    def __init__(
+        self,
+        *spike_generator,
+        comm: Callable,
+        syn: Union[AlignPost, ParamDescriber[AlignPost]],
+        out: Union[SynOut, ParamDescriber[SynOut]],
+        post: Dynamics,
+        stp: ShortTermPlasticity = None,
+    ):
+        super().__init__()
+
+        self.spike_generator = _check_modules(*spike_generator)
+        self.projection = AlignPostProj(comm=comm, syn=syn, out=out, post=post)
+        self.stp = stp
+
+    @call_order(2)
+    def init_state(self, *args, **kwargs):
+        for module in self.spike_generator:
+            maybe_init_prefetch(module, *args, **kwargs)
+
+    def update(self, *x):
+        for fun in self.spike_generator:
+            x = fun(*x)
+            if isinstance(x, (tuple, list)):
+                x = tuple(x)
+            else:
+                x = (x,)
+        assert len(x) == 1, "Spike generator must return a single value or a tuple/list of values"
+        x = brainevent.BinaryArray(x[0])  # Ensure input is a BinaryFloat for spike generation
+        if self.stp is not None:
+            x = brainevent.MaskedFloat(self.stp(x))  # Ensure STP output is a MaskedFloat
+        return self.projection(x)

brainstate/nn/_synapse.py

@@ -307,9 +307,6 @@ class Alpha(Synapse):
         self.h.value = self.sum_delta_inputs(h)
         if x is not None:
             self.h.value += x
-        return self.update_return()
-
-    def update_return(self) -> PyTree:
         return self.g.value
 
 
@@ -394,7 +391,7 @@ class AMPA(Synapse):
         beta: ArrayLike = 0.18 / u.ms,
         T: ArrayLike = 0.5 * u.mM,
         T_dur: ArrayLike = 0.5 * u.ms,
-        g_initializer: ArrayLike | Callable = init.ZeroInit(),
+        g_initializer: ArrayLike | Callable = init.ZeroInit(unit=u.mS),
     ):
         super().__init__(name=name, in_size=in_size)
 
@@ -413,14 +410,12 @@ class AMPA(Synapse):
         self.g.value = init.param(self.g_initializer, self.varshape, batch_or_mode)
         self.spike_arrival_time.value = init.param(init.Constant(-1e7 * u.ms), self.varshape, batch_or_mode)
 
-    def dg(self, g, t, TT):
-        return self.alpha * TT * (1 - g) - self.beta * g
-
     def update(self, pre_spike):
         t = environ.get('t')
         self.spike_arrival_time.value = u.math.where(pre_spike, t, self.spike_arrival_time.value)
         TT = ((t - self.spike_arrival_time.value) < self.T_duration) * self.T
-        self.g.value = exp_euler_step(self.dg, self.g.value, t, TT)
+        dg = lambda g: self.alpha * TT * (1 - g) - self.beta * g
+        self.g.value = exp_euler_step(dg, self.g.value)
         return self.update_return()
 
     def update_return(self) -> PyTree:
@@ -507,7 +502,7 @@ class GABAa(AMPA):
         beta: ArrayLike = 0.18 / u.ms,
         T: ArrayLike = 1.0 * u.mM,
         T_dur: ArrayLike = 1.0 * u.ms,
-        g_initializer: ArrayLike | Callable = init.ZeroInit(),
+        g_initializer: ArrayLike | Callable = init.ZeroInit(unit=u.mS),
     ):
         super().__init__(
             alpha=alpha,
@@ -518,3 +513,4 @@ class GABAa(AMPA):
             in_size=in_size,
             g_initializer=g_initializer
         )
+