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

brainstate/nn/_dyn_impl/_dynamics_synapse.py

@@ -0,0 +1,320 @@
+# 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from typing import Optional
+
+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
+
+__all__ = [
+    'Synapse', 'Expon', 'STP', 'STD', 'AMPA', 'GABAa',
+]
+
+
+class Synapse(Dynamics):
+    """
+    Base class for synapse dynamics.
+    """
+    __module__ = 'brainstate.nn'
+
+
+class Expon(Synapse, AlignPost):
+    r"""Exponential decay synapse model.
+
+    Args:
+      tau: float. The time constant of decay. [ms]
+      %s
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        tau: ArrayLike = 8.0 * u.ms,
+        g_initializer: ArrayLike = init.ZeroInit(unit=u.mS),
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.tau = init.param(tau, self.varshape)
+        self.g_initializer = g_initializer
+
+    def init_state(self, batch_size: int = None, **kwargs):
+        self.g = HiddenState(init.param(self.g_initializer, self.varshape, batch_size))
+
+    def reset_state(self, batch_size: int = None, **kwargs):
+        self.g.value = init.param(self.g_initializer, self.varshape, batch_size)
+
+    def update(self, x=None):
+        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.g.value
+
+
+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
+    """
+    __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):
+        du = lambda u: self.U - u / self.tau_f
+        dx = lambda x: (1 - x) / self.tau_d
+        u = exp_euler_step(du, self.u.value)
+        x = exp_euler_step(dx, 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
+
+
+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
+    """
+    __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):
+        dx = lambda x: (1 - x) / self.tau
+        x = exp_euler_step(dx, 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
+
+
+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]
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        alpha: ArrayLike = 0.98 / (u.ms * u.mM),
+        beta: ArrayLike = 0.18 / u.ms,
+        T: ArrayLike = 0.5 * u.mM,
+        T_dur: ArrayLike = 0.5 * u.ms,
+        g_initializer: ArrayLike = init.ZeroInit(),
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.alpha = init.param(alpha, self.varshape)
+        self.beta = init.param(beta, self.varshape)
+        self.T = init.param(T, self.varshape)
+        self.T_duration = init.param(T_dur, self.varshape)
+        self.g_initializer = g_initializer
+
+    def init_state(self, batch_size=None):
+        self.g = HiddenState(init.param(self.g_initializer, self.varshape, batch_size))
+        self.spike_arrival_time = ShortTermState(init.param(init.Constant(-1e7 * u.ms), self.varshape, batch_size))
+
+    def reset_state(self, batch_or_mode=None, **kwargs):
+        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)
+        return self.g.value
+
+
+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]
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        alpha: ArrayLike = 0.53 / (u.ms * u.mM),
+        beta: ArrayLike = 0.18 / u.ms,
+        T: ArrayLike = 1.0 * u.mM,
+        T_dur: ArrayLike = 1.0 * u.ms,
+    ):
+        super().__init__(alpha=alpha, beta=beta, T=T,
+                         T_dur=T_dur, name=name,
+                         in_size=in_size)

brainstate/nn/_dyn_impl/_dynamics_synapse_test.py

@@ -0,0 +1,132 @@
+# 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 __future__ import annotations
+
+import unittest
+
+import brainunit as u
+import jax.numpy as jnp
+import pytest
+
+import brainstate as bst
+from brainstate.nn import Expon, STP, STD
+
+
+class TestSynapse(unittest.TestCase):
+    def setUp(self):
+        self.in_size = 10
+        self.batch_size = 5
+        self.time_steps = 100
+
+    def generate_input(self):
+        return bst.random.randn(self.time_steps, self.batch_size, self.in_size) * u.mS
+
+    def test_expon_synapse(self):
+        tau = 20.0 * u.ms
+        synapse = Expon(self.in_size, tau=tau)
+        inputs = self.generate_input()
+
+        # Test initialization
+        self.assertEqual(synapse.in_size, (self.in_size,))
+        self.assertEqual(synapse.out_size, (self.in_size,))
+        self.assertEqual(synapse.tau, tau)
+
+        # Test forward pass
+        state = synapse.init_state(self.batch_size)
+        call = bst.compile.jit(synapse)
+        with bst.environ.context(dt=0.1 * u.ms):
+            for t in range(self.time_steps):
+                out = call(inputs[t])
+                self.assertEqual(out.shape, (self.batch_size, self.in_size))
+
+        # Test exponential decay
+        constant_input = jnp.ones((self.batch_size, self.in_size)) * u.mS
+        out1 = call(constant_input)
+        out2 = call(constant_input)
+        self.assertTrue(jnp.all(out2 > out1))  # Output should increase with constant input
+
+    @pytest.mark.skip(reason="Not implemented yet")
+    def test_stp_synapse(self):
+        tau_d = 200.0 * u.ms
+        tau_f = 20.0 * u.ms
+        U = 0.2
+        synapse = STP(self.in_size, tau_d=tau_d, tau_f=tau_f, U=U)
+        inputs = self.generate_input()
+
+        # Test initialization
+        self.assertEqual(synapse.in_size, (self.in_size,))
+        self.assertEqual(synapse.out_size, (self.in_size,))
+        self.assertEqual(synapse.tau_d, tau_d)
+        self.assertEqual(synapse.tau_f, tau_f)
+        self.assertEqual(synapse.U, U)
+
+        # Test forward pass
+        state = synapse.init_state(self.batch_size)
+        call = bst.compile.jit(synapse)
+        for t in range(self.time_steps):
+            out = call(inputs[t])
+            self.assertEqual(out.shape, (self.batch_size, self.in_size))
+
+        # Test short-term plasticity
+        constant_input = jnp.ones((self.batch_size, self.in_size)) * u.mS
+        out1 = call(constant_input)
+        out2 = call(constant_input)
+        self.assertTrue(jnp.any(out2 != out1))  # Output should change due to STP
+
+    @pytest.mark.skip(reason="Not implemented yet")
+    def test_std_synapse(self):
+        tau = 200.0
+        U = 0.2
+        synapse = STD(self.in_size, tau=tau, U=U)
+        inputs = self.generate_input()
+
+        # Test initialization
+        self.assertEqual(synapse.in_size, (self.in_size,))
+        self.assertEqual(synapse.out_size, (self.in_size,))
+        self.assertEqual(synapse.tau, tau)
+        self.assertEqual(synapse.U, U)
+
+        # Test forward pass
+        state = synapse.init_state(self.batch_size)
+        for t in range(self.time_steps):
+            out = synapse(inputs[t])
+            self.assertEqual(out.shape, (self.batch_size, self.in_size))
+
+        # Test short-term depression
+        constant_input = jnp.ones((self.batch_size, self.in_size))
+        out1 = synapse(constant_input)
+        out2 = synapse(constant_input)
+        self.assertTrue(jnp.all(out2 < out1))  # Output should decrease due to STD
+
+    def test_keep_size(self):
+        in_size = (2, 3)
+        for SynapseClass in [Expon, ]:
+            synapse = SynapseClass(in_size)
+            self.assertEqual(synapse.in_size, in_size)
+            self.assertEqual(synapse.out_size, in_size)
+
+            inputs = bst.random.randn(self.time_steps, self.batch_size, *in_size) * u.mS
+            state = synapse.init_state(self.batch_size)
+            call = bst.compile.jit(synapse)
+            with bst.environ.context(dt=0.1 * u.ms):
+                for t in range(self.time_steps):
+                    out = call(inputs[t])
+                    self.assertEqual(out.shape, (self.batch_size, *in_size))
+
+
+if __name__ == '__main__':
+    with bst.environ.context(dt=0.1):
+        unittest.main()

brainstate/nn/_dyn_impl/_inputs.py

@@ -0,0 +1,154 @@
+# 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 __future__ import annotations
+
+from typing import Union, Optional, Sequence, Callable
+
+import brainunit as u
+
+from brainstate import environ, init, random
+from brainstate._state import ShortTermState
+from brainstate.compile import while_loop
+from brainstate.nn._dynamics._dynamics_base import Dynamics
+from brainstate.typing import ArrayLike, Size, DTypeLike
+
+__all__ = [
+    'SpikeTime',
+    'PoissonSpike',
+    'PoissonEncoder',
+]
+
+
+class SpikeTime(Dynamics):
+    """The input neuron group characterized by spikes emitting at given times.
+
+    >>> # Get 2 neurons, firing spikes at 10 ms and 20 ms.
+    >>> SpikeTime(2, times=[10, 20])
+    >>> # or
+    >>> # Get 2 neurons, the neuron 0 fires spikes at 10 ms and 20 ms.
+    >>> SpikeTime(2, times=[10, 20], indices=[0, 0])
+    >>> # or
+    >>> # Get 2 neurons, neuron 0 fires at 10 ms and 30 ms, neuron 1 fires at 20 ms.
+    >>> SpikeTime(2, times=[10, 20, 30], indices=[0, 1, 0])
+    >>> # or
+    >>> # Get 2 neurons; at 10 ms, neuron 0 fires; at 20 ms, neuron 0 and 1 fire;
+    >>> # at 30 ms, neuron 1 fires.
+    >>> SpikeTime(2, times=[10, 20, 20, 30], indices=[0, 0, 1, 1])
+
+    Parameters
+    ----------
+    in_size : int, tuple, list
+        The neuron group geometry.
+    indices : list, tuple, ArrayType
+        The neuron indices at each time point to emit spikes.
+    times : list, tuple, ArrayType
+        The time points which generate the spikes.
+    name : str, optional
+        The name of the dynamic system.
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        indices: Union[Sequence, ArrayLike],
+        times: Union[Sequence, ArrayLike],
+        spk_type: DTypeLike = bool,
+        name: Optional[str] = None,
+        need_sort: bool = True,
+    ):
+        super().__init__(in_size=in_size, name=name)
+
+        # parameters
+        if len(indices) != len(times):
+            raise ValueError(f'The length of "indices" and "times" must be the same. '
+                             f'However, we got {len(indices)} != {len(times)}.')
+        self.num_times = len(times)
+        self.spk_type = spk_type
+
+        # data about times and indices
+        self.times = u.math.asarray(times)
+        self.indices = u.math.asarray(indices, dtype=environ.ditype())
+        if need_sort:
+            sort_idx = u.math.argsort(self.times)
+            self.indices = self.indices[sort_idx]
+            self.times = self.times[sort_idx]
+
+    def init_state(self, *args, **kwargs):
+        self.i = ShortTermState(-1)
+
+    def reset_state(self, batch_size=None, **kwargs):
+        self.i.value = -1
+
+    def update(self):
+        t = environ.get('t')
+
+        def _cond_fun(spikes):
+            i = self.i.value
+            return u.math.logical_and(i < self.num_times, t >= self.times[i])
+
+        def _body_fun(spikes):
+            i = self.i.value
+            spikes = spikes.at[..., self.indices[i]].set(True)
+            self.i.value += 1
+            return spikes
+
+        spike = u.math.zeros(self.varshape, dtype=self.spk_type)
+        spike = while_loop(_cond_fun, _body_fun, spike)
+        return spike
+
+
+class PoissonSpike(Dynamics):
+    """
+    Poisson Neuron Group.
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        freqs: Union[ArrayLike, Callable],
+        spk_type: DTypeLike = bool,
+        name: Optional[str] = None,
+    ):
+        super().__init__(in_size=in_size, name=name)
+
+        self.spk_type = spk_type
+
+        # parameters
+        self.freqs = init.param(freqs, self.varshape, allow_none=False)
+
+    def update(self):
+        spikes = random.rand(self.varshape) <= (self.freqs * environ.get_dt())
+        spikes = u.math.asarray(spikes, dtype=self.spk_type)
+        return spikes
+
+
+class PoissonEncoder(Dynamics):
+    """
+    Poisson Neuron Group.
+    """
+
+    def __init__(
+        self,
+        in_size: Size,
+        spk_type: DTypeLike = bool,
+        name: Optional[str] = None,
+    ):
+        super().__init__(in_size=in_size, name=name)
+        self.spk_type = spk_type
+
+    def update(self, freqs: ArrayLike):
+        spikes = random.rand(*self.varshape) <= (freqs * environ.get_dt())
+        spikes = u.math.asarray(spikes, dtype=self.spk_type)
+        return spikes

brainstate/nn/{_projection/__init__.py → _dyn_impl/_projection_alignpost.py}

@@ -13,20 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 
-"""
+from __future__ import annotations
 
-This module defines the basic classes for synaptic projections.
+from brainstate.nn._dynamics._dynamics_base import Projection
 
-"""
+__all__ = [
+]
 
-from ._align_post import *
-from ._align_post import __all__ as align_post_all
-from ._align_pre import *
-from ._align_pre import __all__ as align_pre_all
-from ._delta import *
-from ._delta import __all__ as delta_all
-from ._vanilla import *
-from ._vanilla import __all__ as vanilla_all
 
-__all__ = align_post_all + align_pre_all + delta_all + vanilla_all
-del align_post_all, align_pre_all, delta_all, vanilla_all
+class ExponentialSynapse(Projection):
+    pass