brainstate 0.1.2__py2.py3-none-any.whl → 0.1.4__py2.py3-none-any.whl

brainstate/nn/_stp.py

@@ -0,0 +1,236 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+from typing import Optional
+
+import brainunit as u
+
+from brainstate import init
+from brainstate._state import HiddenState
+from brainstate.typing import ArrayLike, Size
+from ._exp_euler import exp_euler_step
+from ._synapse import Synapse
+
+__all__ = [
+    'ShortTermPlasticity', 'STP', 'STD',
+]
+
+
+class ShortTermPlasticity(Synapse):
+    pass
+
+
+class STP(ShortTermPlasticity):
+    r"""
+    Synapse with short-term plasticity.
+
+    This class implements a synapse model with short-term plasticity (STP), which captures
+    activity-dependent changes in synaptic efficacy that occur over milliseconds to seconds.
+    The model simultaneously accounts for both short-term facilitation and depression
+    based on the formulation by Tsodyks & Markram (1998).
+
+    The model is characterized by the following equations:
+
+    $$
+    \frac{du}{dt} = -\frac{u}{\tau_f} + U \cdot (1 - u) \cdot \delta(t - t_{spike})
+    $$
+
+    $$
+    \frac{dx}{dt} = \frac{1 - x}{\tau_d} - u \cdot x \cdot \delta(t - t_{spike})
+    $$
+
+    $$
+    g_{syn} = u \cdot x
+    $$
+
+    where:
+    - $u$ represents the utilization of synaptic efficacy (facilitation variable)
+    - $x$ represents the available synaptic resources (depression variable)
+    - $\tau_f$ is the facilitation time constant
+    - $\tau_d$ is the depression time constant
+    - $U$ is the baseline utilization parameter
+    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
+    - $g_{syn}$ is the effective synaptic conductance
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    U : ArrayLike, default=0.15
+        Baseline utilization parameter (fraction of resources used per action potential).
+    tau_f : ArrayLike, default=1500.*u.ms
+        Time constant of short-term facilitation in milliseconds.
+    tau_d : ArrayLike, default=200.*u.ms
+        Time constant of short-term depression (recovery of synaptic resources) in milliseconds.
+
+    Attributes
+    ----------
+    u : HiddenState
+        Utilization of synaptic efficacy (facilitation variable).
+    x : HiddenState
+        Available synaptic resources (depression variable).
+
+    Notes
+    -----
+    - Larger values of tau_f produce stronger facilitation effects.
+    - Larger values of tau_d lead to slower recovery from depression.
+    - The parameter U controls the initial release probability.
+    - The effective synaptic strength is the product of u and x.
+
+    References
+    ----------
+    .. [1] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
+           pyramidal neurons depends on neurotransmitter release probability.
+           Proceedings of the National Academy of Sciences, 94(2), 719-723.
+    .. [2] Tsodyks, M., Pawelzik, K., & Markram, H. (1998). Neural networks with dynamic
+           synapses. Neural computation, 10(4), 821-835.
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        U: ArrayLike = 0.15,
+        tau_f: ArrayLike = 1500. * u.ms,
+        tau_d: ArrayLike = 200. * u.ms,
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.tau_f = init.param(tau_f, self.varshape)
+        self.tau_d = init.param(tau_d, self.varshape)
+        self.U = init.param(U, self.varshape)
+
+    def init_state(self, batch_size: int = None, **kwargs):
+        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
+        self.u = HiddenState(init.param(init.Constant(self.U), self.varshape, batch_size))
+
+    def reset_state(self, batch_size: int = None, **kwargs):
+        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+        self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
+
+    def update(self, pre_spike):
+        u = exp_euler_step(lambda u: - u / self.tau_f, self.u.value)
+        x = exp_euler_step(lambda x: (1 - x) / self.tau_d, self.x.value)
+
+        # --- original code:
+        #   if pre_spike.dtype == jax.numpy.bool_:
+        #     u = bm.where(pre_spike, u + self.U * (1 - self.u), u)
+        #     x = bm.where(pre_spike, x - u * self.x, x)
+        #   else:
+        #     u = pre_spike * (u + self.U * (1 - self.u)) + (1 - pre_spike) * u
+        #     x = pre_spike * (x - u * self.x) + (1 - pre_spike) * x
+
+        # --- simplified code:
+        u = u + pre_spike * self.U * (1 - self.u.value)
+        x = x - pre_spike * u * self.x.value
+
+        self.u.value = u
+        self.x.value = x
+        return u * x * pre_spike
+
+
+class STD(ShortTermPlasticity):
+    r"""
+    Synapse with short-term depression.
+
+    This class implements a synapse model with short-term depression (STD), which captures
+    activity-dependent reduction in synaptic efficacy, typically caused by depletion of
+    neurotransmitter vesicles following repeated stimulation.
+
+    The model is characterized by the following equation:
+
+    $$
+    \frac{dx}{dt} = \frac{1 - x}{\tau} - U \cdot x \cdot \delta(t - t_{spike})
+    $$
+
+    $$
+    g_{syn} = x
+    $$
+
+    where:
+    - $x$ represents the available synaptic resources (depression variable)
+    - $\tau$ is the depression recovery time constant
+    - $U$ is the utilization parameter (fraction of resources depleted per spike)
+    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
+    - $g_{syn}$ is the effective synaptic conductance
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    tau : ArrayLike, default=200.*u.ms
+        Time constant governing recovery of synaptic resources in milliseconds.
+    U : ArrayLike, default=0.07
+        Utilization parameter (fraction of resources used per action potential).
+
+    Attributes
+    ----------
+    x : HiddenState
+        Available synaptic resources (depression variable).
+
+    Notes
+    -----
+    - Larger values of tau lead to slower recovery from depression.
+    - Larger values of U cause stronger depression with each spike.
+    - This model is a simplified version of the STP model that only includes depression.
+
+    References
+    ----------
+    .. [1] Abbott, L. F., Varela, J. A., Sen, K., & Nelson, S. B. (1997). Synaptic
+           depression and cortical gain control. Science, 275(5297), 220-224.
+    .. [2] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
+           pyramidal neurons depends on neurotransmitter release probability.
+           Proceedings of the National Academy of Sciences, 94(2), 719-723.
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        # synapse parameters
+        tau: ArrayLike = 200. * u.ms,
+        U: ArrayLike = 0.07,
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.tau = init.param(tau, self.varshape)
+        self.U = init.param(U, self.varshape)
+
+    def init_state(self, batch_size: int = None, **kwargs):
+        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
+
+    def reset_state(self, batch_size: int = None, **kwargs):
+        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+
+    def update(self, pre_spike):
+        x = exp_euler_step(lambda x: (1 - x) / self.tau, self.x.value)
+
+        # --- original code:
+        # self.x.value = bm.where(pre_spike, x - self.U * self.x, x)
+
+        # --- simplified code:
+        self.x.value = x - pre_spike * self.U * self.x.value
+
+        return self.x.value * pre_spike

brainstate/nn/{_dyn_impl/_dynamics_synapse.py → _synapse.py}

@@ -23,12 +23,12 @@ import brainunit as u
 from brainstate import init, environ
 from brainstate._state import ShortTermState, HiddenState
 from brainstate.mixin import AlignPost
-from brainstate.nn._dynamics._dynamics_base import Dynamics
-from brainstate.nn._exp_euler import exp_euler_step
-from brainstate.typing import ArrayLike, Size
+from brainstate.typing import ArrayLike, Size, PyTree
+from ._dynamics import Dynamics
+from ._exp_euler import exp_euler_step
 
 __all__ = [
-    'Synapse', 'Expon', 'DualExpon', 'Alpha', 'STP', 'STD', 'AMPA', 'GABAa',
+    'Synapse', 'Expon', 'DualExpon', 'Alpha', 'AMPA', 'GABAa',
 ]
 
 
@@ -123,6 +123,9 @@ class Expon(Synapse, AlignPost):
         g = exp_euler_step(lambda g: self.sum_current_inputs(-g) / self.tau, self.g.value)
         self.g.value = self.sum_delta_inputs(g)
         if x is not None: self.g.value += x
+        return self.update_return()
+
+    def update_return(self) -> PyTree:
         return self.g.value
 
 
@@ -229,6 +232,9 @@ class DualExpon(Synapse, AlignPost):
         if x is not None:
             self.g_rise.value += x
             self.g_decay.value += x
+        return self.update_return()
+
+    def update_return(self) -> PyTree:
         return self.a * (self.g_decay.value - self.g_rise.value)
 
 
@@ -304,208 +310,6 @@ class Alpha(Synapse):
         return self.g.value
 
 
-class STP(Synapse):
-    r"""
-    Synapse with short-term plasticity.
-
-    This class implements a synapse model with short-term plasticity (STP), which captures
-    activity-dependent changes in synaptic efficacy that occur over milliseconds to seconds.
-    The model simultaneously accounts for both short-term facilitation and depression
-    based on the formulation by Tsodyks & Markram (1998).
-
-    The model is characterized by the following equations:
-
-    $$
-    \frac{du}{dt} = -\frac{u}{\tau_f} + U \cdot (1 - u) \cdot \delta(t - t_{spike})
-    $$
-
-    $$
-    \frac{dx}{dt} = \frac{1 - x}{\tau_d} - u \cdot x \cdot \delta(t - t_{spike})
-    $$
-
-    $$
-    g_{syn} = u \cdot x
-    $$
-
-    where:
-    - $u$ represents the utilization of synaptic efficacy (facilitation variable)
-    - $x$ represents the available synaptic resources (depression variable)
-    - $\tau_f$ is the facilitation time constant
-    - $\tau_d$ is the depression time constant
-    - $U$ is the baseline utilization parameter
-    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
-    - $g_{syn}$ is the effective synaptic conductance
-
-    Parameters
-    ----------
-    in_size : Size
-        Size of the input.
-    name : str, optional
-        Name of the synapse instance.
-    U : ArrayLike, default=0.15
-        Baseline utilization parameter (fraction of resources used per action potential).
-    tau_f : ArrayLike, default=1500.*u.ms
-        Time constant of short-term facilitation in milliseconds.
-    tau_d : ArrayLike, default=200.*u.ms
-        Time constant of short-term depression (recovery of synaptic resources) in milliseconds.
-
-    Attributes
-    ----------
-    u : HiddenState
-        Utilization of synaptic efficacy (facilitation variable).
-    x : HiddenState
-        Available synaptic resources (depression variable).
-
-    Notes
-    -----
-    - Larger values of tau_f produce stronger facilitation effects.
-    - Larger values of tau_d lead to slower recovery from depression.
-    - The parameter U controls the initial release probability.
-    - The effective synaptic strength is the product of u and x.
-
-    References
-    ----------
-    .. [1] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
-           pyramidal neurons depends on neurotransmitter release probability.
-           Proceedings of the National Academy of Sciences, 94(2), 719-723.
-    .. [2] Tsodyks, M., Pawelzik, K., & Markram, H. (1998). Neural networks with dynamic
-           synapses. Neural computation, 10(4), 821-835.
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        name: Optional[str] = None,
-        U: ArrayLike = 0.15,
-        tau_f: ArrayLike = 1500. * u.ms,
-        tau_d: ArrayLike = 200. * u.ms,
-    ):
-        super().__init__(name=name, in_size=in_size)
-
-        # parameters
-        self.tau_f = init.param(tau_f, self.varshape)
-        self.tau_d = init.param(tau_d, self.varshape)
-        self.U = init.param(U, self.varshape)
-
-    def init_state(self, batch_size: int = None, **kwargs):
-        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
-        self.u = HiddenState(init.param(init.Constant(self.U), self.varshape, batch_size))
-
-    def reset_state(self, batch_size: int = None, **kwargs):
-        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
-        self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
-
-    def update(self, pre_spike):
-        u = exp_euler_step(lambda u: - u / self.tau_f, self.u.value)
-        x = exp_euler_step(lambda x: (1 - x) / self.tau_d, self.x.value)
-
-        # --- original code:
-        #   if pre_spike.dtype == jax.numpy.bool_:
-        #     u = bm.where(pre_spike, u + self.U * (1 - self.u), u)
-        #     x = bm.where(pre_spike, x - u * self.x, x)
-        #   else:
-        #     u = pre_spike * (u + self.U * (1 - self.u)) + (1 - pre_spike) * u
-        #     x = pre_spike * (x - u * self.x) + (1 - pre_spike) * x
-
-        # --- simplified code:
-        u = u + pre_spike * self.U * (1 - self.u.value)
-        x = x - pre_spike * u * self.x.value
-
-        self.u.value = u
-        self.x.value = x
-        return u * x * pre_spike
-
-
-class STD(Synapse):
-    r"""
-    Synapse with short-term depression.
-
-    This class implements a synapse model with short-term depression (STD), which captures
-    activity-dependent reduction in synaptic efficacy, typically caused by depletion of
-    neurotransmitter vesicles following repeated stimulation.
-
-    The model is characterized by the following equation:
-
-    $$
-    \frac{dx}{dt} = \frac{1 - x}{\tau} - U \cdot x \cdot \delta(t - t_{spike})
-    $$
-
-    $$
-    g_{syn} = x
-    $$
-
-    where:
-    - $x$ represents the available synaptic resources (depression variable)
-    - $\tau$ is the depression recovery time constant
-    - $U$ is the utilization parameter (fraction of resources depleted per spike)
-    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
-    - $g_{syn}$ is the effective synaptic conductance
-
-    Parameters
-    ----------
-    in_size : Size
-        Size of the input.
-    name : str, optional
-        Name of the synapse instance.
-    tau : ArrayLike, default=200.*u.ms
-        Time constant governing recovery of synaptic resources in milliseconds.
-    U : ArrayLike, default=0.07
-        Utilization parameter (fraction of resources used per action potential).
-
-    Attributes
-    ----------
-    x : HiddenState
-        Available synaptic resources (depression variable).
-
-    Notes
-    -----
-    - Larger values of tau lead to slower recovery from depression.
-    - Larger values of U cause stronger depression with each spike.
-    - This model is a simplified version of the STP model that only includes depression.
-
-    References
-    ----------
-    .. [1] Abbott, L. F., Varela, J. A., Sen, K., & Nelson, S. B. (1997). Synaptic
-           depression and cortical gain control. Science, 275(5297), 220-224.
-    .. [2] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
-           pyramidal neurons depends on neurotransmitter release probability.
-           Proceedings of the National Academy of Sciences, 94(2), 719-723.
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        name: Optional[str] = None,
-        # synapse parameters
-        tau: ArrayLike = 200. * u.ms,
-        U: ArrayLike = 0.07,
-    ):
-        super().__init__(name=name, in_size=in_size)
-
-        # parameters
-        self.tau = init.param(tau, self.varshape)
-        self.U = init.param(U, self.varshape)
-
-    def init_state(self, batch_size: int = None, **kwargs):
-        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
-
-    def reset_state(self, batch_size: int = None, **kwargs):
-        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
-
-    def update(self, pre_spike):
-        x = exp_euler_step(lambda x: (1 - x) / self.tau, self.x.value)
-
-        # --- original code:
-        # self.x.value = bm.where(pre_spike, x - self.U * self.x, x)
-
-        # --- simplified code:
-        self.x.value = x - pre_spike * self.U * self.x.value
-
-        return self.x.value * pre_spike
-
-
 class AMPA(Synapse):
     r"""AMPA receptor synapse model.
 
@@ -587,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)
 
@@ -606,14 +410,16 @@ 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:
+        """Return the synaptic conductance value."""
         return self.g.value
 
 
@@ -696,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,
@@ -707,3 +513,4 @@ class GABAa(AMPA):
             in_size=in_size,
             g_initializer=g_initializer
         )
+