brainstate 0.1.0.post20250420__py2.py3-none-any.whl → 0.1.0.post20250422__py2.py3-none-any.whl

brainstate/nn/_dyn_impl/_dynamics_synapse.py

@@ -36,6 +36,25 @@ __all__ = [
 class Synapse(Dynamics):
     """
     Base class for synapse dynamics.
+
+    This class serves as the foundation for all synapse models in the library,
+    providing a common interface for implementing various types of synaptic
+    connectivity and transmission mechanisms.
+
+    Synapses are responsible for modeling the transmission of signals between
+    neurons, including temporal dynamics, plasticity, and neurotransmitter effects.
+    All specific synapse implementations (like Expon, Alpha, AMPA, etc.) should
+    inherit from this class.
+
+    See Also
+    --------
+    Expon : Simple first-order exponential decay synapse model
+    Alpha : Alpha function synapse model
+    DualExpon : Dual exponential synapse model
+    STP : Synapse with short-term plasticity
+    STD : Synapse with short-term depression
+    AMPA : AMPA receptor synapse model
+    GABAa : GABAa receptor synapse model
     """
     __module__ = 'brainstate.nn'
 
@@ -44,9 +63,41 @@ class Expon(Synapse, AlignPost):
     r"""
     Exponential decay synapse model.
 
-    Args:
-      tau: float. The time constant of decay. [ms]
-      %s
+    This class implements a simple first-order exponential decay synapse model where
+    the synaptic conductance g decays exponentially with time constant tau:
+
+    $$
+    dg/dt = -g/\tau + \text{input}
+    $$
+
+    The model is widely used for basic synaptic transmission modeling.
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    tau : ArrayLike, default=8.0*u.ms
+        Time constant of decay in milliseconds.
+    g_initializer : ArrayLike or Callable, default=init.ZeroInit(unit=u.mS)
+        Initial value or initializer for synaptic conductance.
+
+    Attributes
+    ----------
+    g : HiddenState
+        Synaptic conductance state variable.
+    tau : Parameter
+        Time constant of decay.
+
+    Notes
+    -----
+    The implementation uses an exponential Euler integration method.
+    The output of this synapse is the conductance value.
+
+    This class inherits from :py:class:`AlignPost`, which means it can be used in projection patterns
+    where synaptic variables are aligned with post-synaptic neurons, enabling event-driven
+    computation and more efficient handling of sparse connectivity patterns.
     """
     __module__ = 'brainstate.nn'
 
@@ -77,6 +128,62 @@ class Expon(Synapse, AlignPost):
 
 
 class DualExpon(Synapse, AlignPost):
+    r"""
+    Dual exponential synapse model.
+
+    This class implements a synapse model with separate rise and decay time constants,
+    which produces a more biologically realistic conductance waveform than a single
+    exponential model. The model is characterized by the differential equation system:
+
+    dg_rise/dt = -g_rise/tau_rise
+    dg_decay/dt = -g_decay/tau_decay
+    g = a * (g_decay - g_rise)
+
+    where $a$ is a normalization factor that ensures the peak conductance reaches
+    the desired amplitude.
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    tau_decay : ArrayLike, default=10.0*u.ms
+        Time constant of decay in milliseconds.
+    tau_rise : ArrayLike, default=1.0*u.ms
+        Time constant of rise in milliseconds.
+    A : ArrayLike, optional
+        Amplitude scaling factor. If None, a scaling factor is automatically
+        calculated to normalize the peak amplitude.
+    g_initializer : ArrayLike or Callable, default=init.ZeroInit(unit=u.mS)
+        Initial value or initializer for synaptic conductance.
+
+    Attributes
+    ----------
+    g_rise : HiddenState
+        Rise component of synaptic conductance.
+    g_decay : HiddenState
+        Decay component of synaptic conductance.
+    tau_rise : Parameter
+        Time constant of rise phase.
+    tau_decay : Parameter
+        Time constant of decay phase.
+    a : Parameter
+        Normalization factor calculated from tau_rise, tau_decay, and A.
+
+    Notes
+    -----
+    The dual exponential model produces a conductance waveform that is more
+    physiologically realistic than a simple exponential decay, with a finite
+    rise time followed by a slower decay.
+
+    The implementation uses an exponential Euler integration method.
+    The output of this synapse is the normalized difference between decay and rise components.
+
+    This class inherits from :py:class:`AlignPost`, which means it can be used in projection patterns
+    where synaptic variables are aligned with post-synaptic neurons, enabling event-driven
+    computation and more efficient handling of sparse connectivity patterns.
+    """
     __module__ = 'brainstate.nn'
 
     def __init__(
@@ -127,6 +234,45 @@ class DualExpon(Synapse, AlignPost):
 
 
 class Alpha(Synapse):
+    r"""
+    Alpha synapse model.
+
+    This class implements the alpha function synapse model, which produces
+    a smooth, biologically realistic synaptic conductance waveform.
+    The model is characterized by the differential equation system:
+
+    dh/dt = -h/tau
+    dg/dt = -g/tau + h/tau
+
+    This produces a response that rises and then falls with a characteristic
+    time constant $\tau$, with peak amplitude occurring at time $t = \tau$.
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    tau : ArrayLike, default=8.0*u.ms
+        Time constant of the alpha function in milliseconds.
+    g_initializer : ArrayLike or Callable, default=init.ZeroInit(unit=u.mS)
+        Initial value or initializer for synaptic conductance.
+
+    Attributes
+    ----------
+    g : HiddenState
+        Synaptic conductance state variable.
+    h : HiddenState
+        Auxiliary state variable for implementing the alpha function.
+    tau : Parameter
+        Time constant of the alpha function.
+
+    Notes
+    -----
+    The alpha function is defined as g(t) = (t/tau) * exp(1-t/tau) for t ≥ 0.
+    This implementation uses an exponential Euler integration method.
+    The output of this synapse is the conductance value.
+    """
     __module__ = 'brainstate.nn'
 
     def __init__(
@@ -160,15 +306,71 @@ class Alpha(Synapse):
 
 
 class STP(Synapse):
-    r"""Synaptic output with short-term plasticity.
-
-    %s
-
-    Args:
-      tau_f: float, ArrayType, Callable. The time constant of short-term facilitation.
-      tau_d: float, ArrayType, Callable. The time constant of short-term depression.
-      U: float, ArrayType, Callable. The fraction of resources used per action potential.
-      %s
+    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'
 
@@ -217,14 +419,59 @@ class STP(Synapse):
 
 
 class STD(Synapse):
-    r"""Synaptic output with short-term depression.
-
-    %s
-
-    Args:
-      tau: float, ArrayType, Callable. The time constant of recovery of the synaptic vesicles.
-      U: float, ArrayType, Callable. The fraction of resources used per action potential.
-      %s
+    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'
 
@@ -261,57 +508,76 @@ class STD(Synapse):
 
 
 class AMPA(Synapse):
-    r"""AMPA synapse model.
-
-    **Model Descriptions**
-
-    AMPA receptor is an ionotropic receptor, which is an ion channel.
-    When it is bound by neurotransmitters, it will immediately open the
-    ion channel, causing the change of membrane potential of postsynaptic neurons.
-
-    A classical model is to use the Markov process to model ion channel switch.
-    Here :math:`g` represents the probability of channel opening, :math:`1-g`
-    represents the probability of ion channel closing, and :math:`\alpha` and
-    :math:`\beta` are the transition probability. Because neurotransmitters can
-    open ion channels, the transfer probability from :math:`1-g` to :math:`g`
-    is affected by the concentration of neurotransmitters. We denote the concentration
-    of neurotransmitters as :math:`[T]` and get the following Markov process.
-
-    .. image:: ../../_static/synapse_markov.png
-        :align: center
-
-    We obtained the following formula when describing the process by a differential equation.
-
-    .. math::
-
-        \frac{ds}{dt} =\alpha[T](1-g)-\beta g
-
-    where :math:`\alpha [T]` denotes the transition probability from state :math:`(1-g)`
-    to state :math:`(g)`; and :math:`\beta` represents the transition probability of
-    the other direction. :math:`\alpha` is the binding constant. :math:`\beta` is the
-    unbinding constant. :math:`[T]` is the neurotransmitter concentration, and
-    has the duration of 0.5 ms.
-
-    Moreover, the post-synaptic current on the post-synaptic neuron is formulated as
-
-    .. math::
-
-        I_{syn} = g_{max} g (V-E)
-
-    where :math:`g_{max}` is the maximum conductance, and `E` is the reverse potential.
-
-    This module can be used with interface ``brainpy.dyn.ProjAlignPreMg2``, as shown in the following example:
-
-    .. [1] Vijayan S, Kopell N J. Thalamic model of awake alpha oscillations
-           and implications for stimulus processing[J]. Proceedings of the
-           National Academy of Sciences, 2012, 109(45): 18553-18558.
-
-    Args:
-      alpha: float, ArrayType, Callable. Binding constant. [ms^-1 mM^-1]
-      beta: float, ArrayType, Callable. Unbinding constant. [ms^-1 mM^-1]
-      T: float, ArrayType, Callable. Transmitter concentration when synapse is triggered by
-        a pre-synaptic spike. Default 0.5 [mM].
-      T_dur: float, ArrayType, Callable. Transmitter concentration duration time after being triggered. Default 1 [ms]
+    r"""AMPA receptor synapse model.
+
+    This class implements a kinetic model of AMPA (α-amino-3-hydroxy-5-methyl-4-isoxazolepropionic acid)
+    receptor-mediated synaptic transmission. AMPA receptors are ionotropic glutamate receptors that mediate
+    fast excitatory synaptic transmission in the central nervous system.
+
+    The model uses a Markov process approach to describe the state transitions of AMPA receptors
+    between closed and open states, governed by neurotransmitter binding:
+
+    $$
+    \frac{dg}{dt} = \alpha [T] (1-g) - \beta g
+    $$
+
+    $$
+    I_{syn} = g_{max} \cdot g \cdot (V - E)
+    $$
+
+    where:
+    - $g$ represents the fraction of receptors in the open state
+    - $\alpha$ is the binding rate constant [ms^-1 mM^-1]
+    - $\beta$ is the unbinding rate constant [ms^-1]
+    - $[T]$ is the neurotransmitter concentration [mM]
+    - $I_{syn}$ is the resulting synaptic current
+    - $g_{max}$ is the maximum conductance
+    - $V$ is the membrane potential
+    - $E$ is the reversal potential
+
+    The neurotransmitter concentration $[T]$ follows a square pulse of amplitude T and
+    duration T_dur after each presynaptic spike.
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    alpha : ArrayLike, default=0.98/(u.ms*u.mM)
+        Binding rate constant [ms^-1 mM^-1].
+    beta : ArrayLike, default=0.18/u.ms
+        Unbinding rate constant [ms^-1].
+    T : ArrayLike, default=0.5*u.mM
+        Peak neurotransmitter concentration when released [mM].
+    T_dur : ArrayLike, default=0.5*u.ms
+        Duration of neurotransmitter presence in the synaptic cleft [ms].
+    g_initializer : ArrayLike or Callable, default=init.ZeroInit()
+        Initial value or initializer for the synaptic conductance.
+
+    Attributes
+    ----------
+    g : HiddenState
+        Fraction of receptors in the open state.
+    spike_arrival_time : ShortTermState
+        Time of the most recent presynaptic spike.
+
+    Notes
+    -----
+    - The model captures the fast-rising and relatively fast-decaying excitatory currents
+      characteristic of AMPA receptor-mediated transmission.
+    - The time course of the synaptic conductance is determined by both the binding and
+      unbinding rate constants and the duration of transmitter presence.
+    - This implementation uses an exponential Euler integration method.
+
+    References
+    ----------
+    .. [1] Destexhe, A., Mainen, Z. F., & Sejnowski, T. J. (1994). Synthesis of models for
+           excitable membranes, synaptic transmission and neuromodulation using a common
+           kinetic formalism. Journal of computational neuroscience, 1(3), 195-230.
+    .. [2] Vijayan, S., & Kopell, N. J. (2012). Thalamic model of awake alpha oscillations
+           and implications for stimulus processing. Proceedings of the National Academy
+           of Sciences, 109(45), 18553-18558.
     """
 
     def __init__(
@@ -353,38 +619,74 @@ class AMPA(Synapse):
 
 
 class GABAa(AMPA):
-    r"""GABAa synapse model.
-
-    **Model Descriptions**
-
-    GABAa synapse model has the same equation with the `AMPA synapse <./brainmodels.synapses.AMPA.rst>`_,
-
-    .. math::
-
-        \frac{d g}{d t}&=\alpha[T](1-g) - \beta g \\
-        I_{syn}&= - g_{max} g (V - E)
-
-    but with the difference of:
-
-    - Reversal potential of synapse :math:`E` is usually low, typically -80. mV
-    - Activating rate constant :math:`\alpha=0.53`
-    - De-activating rate constant :math:`\beta=0.18`
-    - Transmitter concentration :math:`[T]=1\,\mu ho(\mu S)` when synapse is
-      triggered by a pre-synaptic spike, with the duration of 1. ms.
-
-    This module can be used with interface ``brainpy.dyn.ProjAlignPreMg2``, as shown in the following example:
-
-    .. [1] Destexhe, Alain, and Denis Paré. "Impact of network activity
-           on the integrative properties of neocortical pyramidal neurons
-           in vivo." Journal of neurophysiology 81.4 (1999): 1531-1547.
-
-    Args:
-      alpha: float, ArrayType, Callable. Binding constant. Default 0.062
-      beta: float, ArrayType, Callable. Unbinding constant. Default 3.57
-      T: float, ArrayType, Callable. Transmitter concentration when synapse is triggered by
-        a pre-synaptic spike.. Default 1 [mM].
-      T_dur: float, ArrayType, Callable. Transmitter concentration duration time
-        after being triggered. Default 1 [ms]
+    r"""GABAa receptor synapse model.
+
+    This class implements a kinetic model of GABAa (gamma-aminobutyric acid type A)
+    receptor-mediated synaptic transmission. GABAa receptors are ionotropic chloride channels
+    that mediate fast inhibitory synaptic transmission in the central nervous system.
+
+    The model uses the same Markov process approach as the AMPA model but with different
+    kinetic parameters appropriate for GABAa receptors:
+
+    $$
+    \frac{dg}{dt} = \alpha [T] (1-g) - \beta g
+    $$
+
+    $$
+    I_{syn} = - g_{max} \cdot g \cdot (V - E)
+    $$
+
+    where:
+    - $g$ represents the fraction of receptors in the open state
+    - $\alpha$ is the binding rate constant [ms^-1 mM^-1], typically slower than AMPA
+    - $\beta$ is the unbinding rate constant [ms^-1]
+    - $[T]$ is the neurotransmitter (GABA) concentration [mM]
+    - $I_{syn}$ is the resulting synaptic current (note the negative sign indicating inhibition)
+    - $g_{max}$ is the maximum conductance
+    - $V$ is the membrane potential
+    - $E$ is the reversal potential (typically around -80 mV for chloride)
+
+    The neurotransmitter concentration $[T]$ follows a square pulse of amplitude T and
+    duration T_dur after each presynaptic spike.
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    alpha : ArrayLike, default=0.53/(u.ms*u.mM)
+        Binding rate constant [ms^-1 mM^-1]. Typically slower than AMPA receptors.
+    beta : ArrayLike, default=0.18/u.ms
+        Unbinding rate constant [ms^-1].
+    T : ArrayLike, default=1.0*u.mM
+        Peak neurotransmitter concentration when released [mM]. Higher than AMPA.
+    T_dur : ArrayLike, default=1.0*u.ms
+        Duration of neurotransmitter presence in the synaptic cleft [ms]. Longer than AMPA.
+    g_initializer : ArrayLike or Callable, default=init.ZeroInit()
+        Initial value or initializer for the synaptic conductance.
+
+    Attributes
+    ----------
+    Inherits all attributes from AMPA class.
+
+    Notes
+    -----
+    - GABAa receptors typically produce slower-rising and longer-lasting currents compared to AMPA receptors.
+    - The inhibitory nature of GABAa receptors is reflected in the convention of using a negative sign in the
+      synaptic current equation.
+    - The reversal potential for GABAa receptors is typically around -80 mV (due to chloride), making them
+      inhibitory for neurons with resting potentials more positive than this value.
+    - This model does not include desensitization, which can be significant for prolonged GABA exposure.
+
+    References
+    ----------
+    .. [1] Destexhe, A., Mainen, Z. F., & Sejnowski, T. J. (1994). Synthesis of models for
+           excitable membranes, synaptic transmission and neuromodulation using a common
+           kinetic formalism. Journal of computational neuroscience, 1(3), 195-230.
+    .. [2] Destexhe, A., & Paré, D. (1999). Impact of network activity on the integrative
+           properties of neocortical pyramidal neurons in vivo. Journal of neurophysiology,
+           81(4), 1531-1547.
     """
 
     def __init__(

brainstate/nn/_dyn_impl/_dynamics_synapse_test.py

@@ -15,10 +15,11 @@
 
 from __future__ import annotations
 
+import unittest
+
 import brainunit as u
 import jax.numpy as jnp
 import pytest
-import unittest
 
 import brainstate as bst
 from brainstate.nn import Expon, STP, STD