brainstate 0.1.10__py2.py3-none-any.whl → 0.2.0__py2.py3-none-any.whl

brainstate/nn/_rate_rnns_test.py

@@ -1,63 +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.
-# ==============================================================================
-
-
-import unittest
-
-import jax.numpy as jnp
-
-import brainstate
-
-
-class TestRateRNNModels(unittest.TestCase):
-    def setUp(self):
-        self.num_in = 3
-        self.num_out = 3
-        self.batch_size = 4
-        self.x = jnp.ones((self.batch_size, self.num_in))
-
-    def test_ValinaRNNCell(self):
-        model = brainstate.nn.ValinaRNNCell(num_in=self.num_in, num_out=self.num_out)
-        model.init_state(batch_size=self.batch_size)
-        output = model.update(self.x)
-        self.assertEqual(output.shape, (self.batch_size, self.num_out))
-
-    def test_GRUCell(self):
-        model = brainstate.nn.GRUCell(num_in=self.num_in, num_out=self.num_out)
-        model.init_state(batch_size=self.batch_size)
-        output = model.update(self.x)
-        self.assertEqual(output.shape, (self.batch_size, self.num_out))
-
-    def test_MGUCell(self):
-        model = brainstate.nn.MGUCell(num_in=self.num_in, num_out=self.num_out)
-        model.init_state(batch_size=self.batch_size)
-        output = model.update(self.x)
-        self.assertEqual(output.shape, (self.batch_size, self.num_out))
-
-    def test_LSTMCell(self):
-        model = brainstate.nn.LSTMCell(num_in=self.num_in, num_out=self.num_out)
-        model.init_state(batch_size=self.batch_size)
-        output = model.update(self.x)
-        self.assertEqual(output.shape, (self.batch_size, self.num_out))
-
-    def test_URLSTMCell(self):
-        model = brainstate.nn.URLSTMCell(num_in=self.num_in, num_out=self.num_out)
-        model.init_state(batch_size=self.batch_size)
-        output = model.update(self.x)
-        self.assertEqual(output.shape, (self.batch_size, self.num_out))
-
-
-if __name__ == '__main__':
-    unittest.main()

brainstate/nn/_readout.py

@@ -1,209 +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.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-
-import numbers
-from typing import Callable
-
-import brainunit as u
-import jax
-
-from brainstate import environ, init, surrogate
-from brainstate._state import HiddenState, ParamState
-from brainstate.typing import Size, ArrayLike
-from ._exp_euler import exp_euler_step
-from ._module import Module
-from ._neuron import Neuron
-
-__all__ = [
-    'LeakyRateReadout',
-    'LeakySpikeReadout',
-]
-
-
-class LeakyRateReadout(Module):
-    r"""
-    Leaky dynamics for the read-out module.
-
-    This module implements a leaky integrator with the following dynamics:
-
-    .. math::
-        r_{t} = \alpha r_{t-1} + x_{t} W
-
-    where:
-      - :math:`r_{t}` is the output at time t
-      - :math:`\alpha = e^{-\Delta t / \tau}` is the decay factor
-      - :math:`x_{t}` is the input at time t
-      - :math:`W` is the weight matrix
-
-    The leaky integrator acts as a low-pass filter, allowing the network
-    to maintain memory of past inputs with an exponential decay determined
-    by the time constant tau.
-
-    Parameters
-    ----------
-    in_size : int or sequence of int
-        Size of the input dimension(s)
-    out_size : int or sequence of int
-        Size of the output dimension(s)
-    tau : ArrayLike, optional
-        Time constant of the leaky dynamics, by default 5ms
-    w_init : Callable, optional
-        Weight initialization function, by default KaimingNormal()
-    name : str, optional
-        Name of the module, by default None
-
-    Attributes
-    ----------
-    decay : float
-        Decay factor computed as exp(-dt/tau)
-    weight : ParamState
-        Weight matrix connecting input to output
-    r : HiddenState
-        Hidden state representing the output values
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        out_size: Size,
-        tau: ArrayLike = 5. * u.ms,
-        w_init: Callable = init.KaimingNormal(),
-        name: str = None,
-    ):
-        super().__init__(name=name)
-
-        # parameters
-        self.in_size = (in_size,) if isinstance(in_size, numbers.Integral) else tuple(in_size)
-        self.out_size = (out_size,) if isinstance(out_size, numbers.Integral) else tuple(out_size)
-        self.tau = init.param(tau, self.in_size)
-        self.decay = u.math.exp(-environ.get_dt() / self.tau)
-
-        # weights
-        self.weight = ParamState(init.param(w_init, (self.in_size[0], self.out_size[0])))
-
-    def init_state(self, batch_size=None, **kwargs):
-        self.r = HiddenState(init.param(init.Constant(0.), self.out_size, batch_size))
-
-    def reset_state(self, batch_size=None, **kwargs):
-        self.r.value = init.param(init.Constant(0.), self.out_size, batch_size)
-
-    def update(self, x):
-        self.r.value = self.decay * self.r.value + x @ self.weight.value
-        return self.r.value
-
-
-class LeakySpikeReadout(Neuron):
-    r"""
-    Integrate-and-fire neuron model with leaky dynamics for readout functionality.
-
-    This class implements a spiking neuron with the following dynamics:
-
-    .. math::
-        \frac{dV}{dt} = \frac{-V + I_{in}}{\tau}
-
-    where:
-      - :math:`V` is the membrane potential
-      - :math:`\tau` is the membrane time constant
-      - :math:`I_{in}` is the input current
-
-    Spike generation occurs when :math:`V > V_{th}` according to:
-
-    .. math::
-        S_t = \text{surrogate}\left(\frac{V - V_{th}}{V_{th}}\right)
-
-    After spiking, the membrane potential is reset according to the reset mode:
-      - Soft reset: :math:`V \leftarrow V - V_{th} \cdot S_t`
-      - Hard reset: :math:`V \leftarrow V - V_t \cdot S_t` (where :math:`V_t` is detached)
-
-    Parameters
-    ----------
-    in_size : Size
-        Size of the input dimension
-    tau : ArrayLike, optional
-        Membrane time constant, by default 5ms
-    V_th : ArrayLike, optional
-        Spike threshold, by default 1mV
-    w_init : Callable, optional
-        Weight initialization function, by default KaimingNormal(unit=mV)
-    V_initializer : ArrayLike, optional
-        Initial membrane potential, by default ZeroInit(unit=mV)
-    spk_fun : Callable, optional
-        Surrogate gradient function for spike generation, by default ReluGrad()
-    spk_reset : str, optional
-        Reset mechanism after spike ('soft' or 'hard'), by default 'soft'
-    name : str, optional
-        Name of the module, by default None
-
-    Attributes
-    ----------
-    V : HiddenState
-        Membrane potential state variable
-    weight : ParamState
-        Synaptic weight matrix
-    """
-
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        tau: ArrayLike = 5. * u.ms,
-        V_th: ArrayLike = 1. * u.mV,
-        w_init: Callable = init.KaimingNormal(unit=u.mV),
-        V_initializer: ArrayLike = init.ZeroInit(unit=u.mV),
-        spk_fun: Callable = surrogate.ReluGrad(),
-        spk_reset: str = 'soft',
-        name: str = None,
-    ):
-        super().__init__(in_size, name=name, spk_fun=spk_fun, spk_reset=spk_reset)
-
-        # parameters
-        self.tau = init.param(tau, self.varshape)
-        self.V_th = init.param(V_th, self.varshape)
-        self.V_initializer = V_initializer
-
-        # weights
-        self.weight = ParamState(init.param(w_init, (self.in_size[-1], self.out_size[-1])))
-
-    def init_state(self, batch_size, **kwargs):
-        self.V = HiddenState(init.param(self.V_initializer, self.varshape, batch_size))
-
-    def reset_state(self, batch_size, **kwargs):
-        self.V.value = init.param(self.V_initializer, self.varshape, batch_size)
-
-    @property
-    def spike(self):
-        return self.get_spike(self.V.value)
-
-    def get_spike(self, V):
-        v_scaled = (V - self.V_th) / self.V_th
-        return self.spk_fun(v_scaled)
-
-    def update(self, spk):
-        # reset
-        last_V = self.V.value
-        last_spike = self.get_spike(last_V)
-        V_th = self.V_th if self.spk_reset == 'soft' else jax.lax.stop_gradient(last_V)
-        V = last_V - V_th * last_spike
-        # membrane potential
-        x = spk @ self.weight.value
-        dv = lambda v: (-v + self.sum_current_inputs(x, v)) / self.tau
-        V = exp_euler_step(dv, V)
-        self.V.value = self.sum_delta_inputs(V)
-        return self.get_spike(V)

brainstate/nn/_readout_test.py

@@ -1,53 +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.
-# ==============================================================================
-
-
-import unittest
-
-import jax.numpy as jnp
-
-import brainstate
-
-
-class TestReadoutModels(unittest.TestCase):
-    def setUp(self):
-        self.in_size = 3
-        self.out_size = 3
-        self.batch_size = 4
-        self.tau = 5.0
-        self.V_th = 1.0
-        self.x = jnp.ones((self.batch_size, self.in_size))
-
-    def test_LeakyRateReadout(self):
-        with brainstate.environ.context(dt=0.1):
-            model = brainstate.nn.LeakyRateReadout(in_size=self.in_size, out_size=self.out_size, tau=self.tau)
-            model.init_state(batch_size=self.batch_size)
-            output = model.update(self.x)
-            self.assertEqual(output.shape, (self.batch_size, self.out_size))
-
-    def test_LeakySpikeReadout(self):
-        with brainstate.environ.context(dt=0.1):
-            model = brainstate.nn.LeakySpikeReadout(in_size=self.in_size, tau=self.tau, V_th=self.V_th,
-                                                    V_initializer=brainstate.init.ZeroInit(),
-                                                    w_init=brainstate.init.KaimingNormal())
-            model.init_state(batch_size=self.batch_size)
-            with brainstate.environ.context(t=0.):
-                output = model.update(self.x)
-            self.assertEqual(output.shape, (self.batch_size, self.out_size))
-
-
-if __name__ == '__main__':
-    with brainstate.environ.context(dt=0.1):
-        unittest.main()

brainstate/nn/_stp.py

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