brainstate 0.1.7__py2.py3-none-any.whl → 0.1.9__py2.py3-none-any.whl

brainstate/nn/_synouts.py

@@ -1,162 +1,162 @@
-# 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 brainunit as u
-import jax.numpy as jnp
-
-from brainstate.mixin import BindCondData
-from brainstate.typing import ArrayLike
-from ._module import Module
-
-__all__ = [
-    'SynOut', 'COBA', 'CUBA', 'MgBlock',
-]
-
-
-class SynOut(Module, BindCondData):
-    """
-    Base class for synaptic outputs.
-
-    :py:class:`~.SynOut` is also subclass of :py:class:`~.ParamDesc` and :py:class:`~.BindCondData`.
-    """
-
-    __module__ = 'brainstate.nn'
-
-    def __init__(self, ):
-        super().__init__()
-        self._conductance = None
-
-    def __call__(self, *args, **kwargs):
-        if self._conductance is None:
-            raise ValueError(f'Please first pack conductance data at the current step using '
-                             f'".{BindCondData.bind_cond.__name__}(data)". {self}')
-        ret = self.update(self._conductance, *args, **kwargs)
-        return ret
-
-    def update(self, conductance, potential):
-        raise NotImplementedError
-
-
-class COBA(SynOut):
-    r"""
-    Conductance-based synaptic output.
-
-    Given the synaptic conductance, the model output the post-synaptic current with
-
-    .. math::
-
-       I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t))
-
-    Parameters
-    ----------
-    E: ArrayLike
-      The reversal potential.
-
-    See Also
-    --------
-    CUBA
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(self, E: ArrayLike):
-        super().__init__()
-
-        self.E = E
-
-    def update(self, conductance, potential):
-        return conductance * (self.E - potential)
-
-
-class CUBA(SynOut):
-    r"""Current-based synaptic output.
-
-    Given the conductance, this model outputs the post-synaptic current with a identity function:
-
-    .. math::
-
-       I_{\mathrm{syn}}(t) = g_{\mathrm{syn}}(t)
-
-    Parameters
-    ----------
-    scale: ArrayLike
-      The scaling factor for the conductance. Default 1. [mV]
-
-    See Also
-    --------
-    COBA
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(self, scale: ArrayLike = u.volt):
-        super().__init__()
-        self.scale = scale
-
-    def update(self, conductance, potential=None):
-        return conductance * self.scale
-
-
-class MgBlock(SynOut):
-    r"""Synaptic output based on Magnesium blocking.
-
-    Given the synaptic conductance, the model output the post-synaptic current with
-
-    .. math::
-
-       I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t)) g_{\infty}(V,[{Mg}^{2+}]_{o})
-
-    where The fraction of channels :math:`g_{\infty}` that are not blocked by magnesium can be fitted to
-
-    .. math::
-
-       g_{\infty}(V,[{Mg}^{2+}]_{o}) = (1+{e}^{-\alpha V} \frac{[{Mg}^{2+}]_{o}} {\beta})^{-1}
-
-    Here :math:`[{Mg}^{2+}]_{o}` is the extracellular magnesium concentration.
-
-    Parameters
-    ----------
-    E: ArrayLike
-      The reversal potential for the synaptic current. [mV]
-    alpha: ArrayLike
-      Binding constant. Default 0.062
-    beta: ArrayLike
-      Unbinding constant. Default 3.57
-    cc_Mg: ArrayLike
-      Concentration of Magnesium ion. Default 1.2 [mM].
-    V_offset: ArrayLike
-      The offset potential. Default 0. [mV]
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        E: ArrayLike = 0.,
-        cc_Mg: ArrayLike = 1.2,
-        alpha: ArrayLike = 0.062,
-        beta: ArrayLike = 3.57,
-        V_offset: ArrayLike = 0.,
-    ):
-        super().__init__()
-
-        self.E = E
-        self.V_offset = V_offset
-        self.cc_Mg = cc_Mg
-        self.alpha = alpha
-        self.beta = beta
-
-    def update(self, conductance, potential):
-        norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - potential)))
-        return conductance * (self.E - potential) / norm
+# 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 brainunit as u
+import jax.numpy as jnp
+
+from brainstate.mixin import BindCondData
+from brainstate.typing import ArrayLike
+from ._module import Module
+
+__all__ = [
+    'SynOut', 'COBA', 'CUBA', 'MgBlock',
+]
+
+
+class SynOut(Module, BindCondData):
+    """
+    Base class for synaptic outputs.
+
+    :py:class:`~.SynOut` is also subclass of :py:class:`~.ParamDesc` and :py:class:`~.BindCondData`.
+    """
+
+    __module__ = 'brainstate.nn'
+
+    def __init__(self, ):
+        super().__init__()
+        self._conductance = None
+
+    def __call__(self, *args, **kwargs):
+        if self._conductance is None:
+            raise ValueError(f'Please first pack conductance data at the current step using '
+                             f'".{BindCondData.bind_cond.__name__}(data)". {self}')
+        ret = self.update(self._conductance, *args, **kwargs)
+        return ret
+
+    def update(self, conductance, potential):
+        raise NotImplementedError
+
+
+class COBA(SynOut):
+    r"""
+    Conductance-based synaptic output.
+
+    Given the synaptic conductance, the model output the post-synaptic current with
+
+    .. math::
+
+       I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t))
+
+    Parameters
+    ----------
+    E: ArrayLike
+      The reversal potential.
+
+    See Also
+    --------
+    CUBA
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(self, E: ArrayLike):
+        super().__init__()
+
+        self.E = E
+
+    def update(self, conductance, potential):
+        return conductance * (self.E - potential)
+
+
+class CUBA(SynOut):
+    r"""Current-based synaptic output.
+
+    Given the conductance, this model outputs the post-synaptic current with a identity function:
+
+    .. math::
+
+       I_{\mathrm{syn}}(t) = g_{\mathrm{syn}}(t)
+
+    Parameters
+    ----------
+    scale: ArrayLike
+      The scaling factor for the conductance. Default 1. [mV]
+
+    See Also
+    --------
+    COBA
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(self, scale: ArrayLike = u.volt):
+        super().__init__()
+        self.scale = scale
+
+    def update(self, conductance, potential=None):
+        return conductance * self.scale
+
+
+class MgBlock(SynOut):
+    r"""Synaptic output based on Magnesium blocking.
+
+    Given the synaptic conductance, the model output the post-synaptic current with
+
+    .. math::
+
+       I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t)) g_{\infty}(V,[{Mg}^{2+}]_{o})
+
+    where The fraction of channels :math:`g_{\infty}` that are not blocked by magnesium can be fitted to
+
+    .. math::
+
+       g_{\infty}(V,[{Mg}^{2+}]_{o}) = (1+{e}^{-\alpha V} \frac{[{Mg}^{2+}]_{o}} {\beta})^{-1}
+
+    Here :math:`[{Mg}^{2+}]_{o}` is the extracellular magnesium concentration.
+
+    Parameters
+    ----------
+    E: ArrayLike
+      The reversal potential for the synaptic current. [mV]
+    alpha: ArrayLike
+      Binding constant. Default 0.062
+    beta: ArrayLike
+      Unbinding constant. Default 3.57
+    cc_Mg: ArrayLike
+      Concentration of Magnesium ion. Default 1.2 [mM].
+    V_offset: ArrayLike
+      The offset potential. Default 0. [mV]
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        E: ArrayLike = 0.,
+        cc_Mg: ArrayLike = 1.2,
+        alpha: ArrayLike = 0.062,
+        beta: ArrayLike = 3.57,
+        V_offset: ArrayLike = 0.,
+    ):
+        super().__init__()
+
+        self.E = E
+        self.V_offset = V_offset
+        self.cc_Mg = cc_Mg
+        self.alpha = alpha
+        self.beta = beta
+
+    def update(self, conductance, potential):
+        norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - potential)))
+        return conductance * (self.E - potential) / norm

brainstate/nn/_synouts_test.py

@@ -1,57 +1,57 @@
-# 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 brainunit as u
-import jax.numpy as jnp
-import numpy as np
-
-import brainstate
-
-
-class TestSynOutModels(unittest.TestCase):
-    def setUp(self):
-        self.conductance = jnp.array([0.5, 1.0, 1.5])
-        self.potential = jnp.array([-70.0, -65.0, -60.0])
-        self.E = jnp.array([-70.0])
-        self.alpha = jnp.array([0.062])
-        self.beta = jnp.array([3.57])
-        self.cc_Mg = jnp.array([1.2])
-        self.V_offset = jnp.array([0.0])
-
-    def test_COBA(self):
-        model = brainstate.nn.COBA(E=self.E)
-        output = model.update(self.conductance, self.potential)
-        expected_output = self.conductance * (self.E - self.potential)
-        np.testing.assert_array_almost_equal(output, expected_output)
-
-    def test_CUBA(self):
-        model = brainstate.nn.CUBA()
-        output = model.update(self.conductance)
-        expected_output = self.conductance * model.scale
-        self.assertTrue(u.math.allclose(output, expected_output))
-
-    def test_MgBlock(self):
-        model = brainstate.nn.MgBlock(E=self.E, cc_Mg=self.cc_Mg, alpha=self.alpha, beta=self.beta, V_offset=self.V_offset)
-        output = model.update(self.conductance, self.potential)
-        norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - self.potential)))
-        expected_output = self.conductance * (self.E - self.potential) / norm
-        np.testing.assert_array_almost_equal(output, expected_output)
-
-
-if __name__ == '__main__':
-    unittest.main()
+# 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 brainunit as u
+import jax.numpy as jnp
+import numpy as np
+
+import brainstate
+
+
+class TestSynOutModels(unittest.TestCase):
+    def setUp(self):
+        self.conductance = jnp.array([0.5, 1.0, 1.5])
+        self.potential = jnp.array([-70.0, -65.0, -60.0])
+        self.E = jnp.array([-70.0])
+        self.alpha = jnp.array([0.062])
+        self.beta = jnp.array([3.57])
+        self.cc_Mg = jnp.array([1.2])
+        self.V_offset = jnp.array([0.0])
+
+    def test_COBA(self):
+        model = brainstate.nn.COBA(E=self.E)
+        output = model.update(self.conductance, self.potential)
+        expected_output = self.conductance * (self.E - self.potential)
+        np.testing.assert_array_almost_equal(output, expected_output)
+
+    def test_CUBA(self):
+        model = brainstate.nn.CUBA()
+        output = model.update(self.conductance)
+        expected_output = self.conductance * model.scale
+        self.assertTrue(u.math.allclose(output, expected_output))
+
+    def test_MgBlock(self):
+        model = brainstate.nn.MgBlock(E=self.E, cc_Mg=self.cc_Mg, alpha=self.alpha, beta=self.beta, V_offset=self.V_offset)
+        output = model.update(self.conductance, self.potential)
+        norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - self.potential)))
+        expected_output = self.conductance * (self.E - self.potential) / norm
+        np.testing.assert_array_almost_equal(output, expected_output)
+
+
+if __name__ == '__main__':
+    unittest.main()

brainstate/nn/_utils.py

@@ -1,89 +1,89 @@
-# 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 Union, Tuple
-
-from brainstate._state import ParamState
-from brainstate.util import PrettyTable
-from ._module import Module
-
-__all__ = [
-    "count_parameters",
-]
-
-
-def _format_parameter_count(num_params, precision=2):
-    if num_params < 1000:
-        return str(num_params)
-
-    suffixes = ['', 'K', 'M', 'B', 'T', 'P', 'E']
-    magnitude = 0
-    while abs(num_params) >= 1000:
-        magnitude += 1
-        num_params /= 1000.0
-
-    format_string = '{:.' + str(precision) + 'f}{}'
-    formatted_value = format_string.format(num_params, suffixes[magnitude])
-
-    # 检查是否接近 1000，如果是，尝试使用更大的基数
-    if magnitude < len(suffixes) - 1 and num_params >= 1000 * (1 - 10 ** (-precision)):
-        magnitude += 1
-        num_params /= 1000.0
-        formatted_value = format_string.format(num_params, suffixes[magnitude])
-
-    return formatted_value
-
-
-def count_parameters(
-    module: Module,
-    precision: int = 2,
-    return_table: bool = False,
-) -> Union[Tuple[PrettyTable, int], int]:
-    """
-    Count and display the number of trainable parameters in a neural network model.
-
-    This function iterates through all the parameters of the given model,
-    counts the number of parameters for each module, and displays them in a table.
-    It also calculates and returns the total number of trainable parameters.
-
-    Parameters:
-    -----------
-    model : brainstate.nn.Module
-        The neural network model for which to count parameters.
-
-    Returns:
-    --------
-    int
-        The total number of trainable parameters in the model.
-
-    Prints:
-    -------
-    A pretty-formatted table showing the number of parameters for each module,
-    followed by the total number of trainable parameters.
-    """
-    assert isinstance(module, Module), "Input must be a neural network module"  # noqa: E501
-    table = PrettyTable(["Modules", "Parameters"])
-    total_params = 0
-    for name, parameter in module.states(ParamState).items():
-        param = parameter.numel()
-        table.add_row([name, _format_parameter_count(param, precision=precision)])
-        total_params += param
-    table.add_row(["Total", _format_parameter_count(total_params, precision=precision)])
-    print(table)
-    if return_table:
-        return table, total_params
-    return total_params
+# 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 Union, Tuple
+
+from brainstate._state import ParamState
+from brainstate.util import PrettyTable
+from ._module import Module
+
+__all__ = [
+    "count_parameters",
+]
+
+
+def _format_parameter_count(num_params, precision=2):
+    if num_params < 1000:
+        return str(num_params)
+
+    suffixes = ['', 'K', 'M', 'B', 'T', 'P', 'E']
+    magnitude = 0
+    while abs(num_params) >= 1000:
+        magnitude += 1
+        num_params /= 1000.0
+
+    format_string = '{:.' + str(precision) + 'f}{}'
+    formatted_value = format_string.format(num_params, suffixes[magnitude])
+
+    # 检查是否接近 1000，如果是，尝试使用更大的基数
+    if magnitude < len(suffixes) - 1 and num_params >= 1000 * (1 - 10 ** (-precision)):
+        magnitude += 1
+        num_params /= 1000.0
+        formatted_value = format_string.format(num_params, suffixes[magnitude])
+
+    return formatted_value
+
+
+def count_parameters(
+    module: Module,
+    precision: int = 2,
+    return_table: bool = False,
+) -> Union[Tuple[PrettyTable, int], int]:
+    """
+    Count and display the number of trainable parameters in a neural network model.
+
+    This function iterates through all the parameters of the given model,
+    counts the number of parameters for each module, and displays them in a table.
+    It also calculates and returns the total number of trainable parameters.
+
+    Parameters:
+    -----------
+    model : brainstate.nn.Module
+        The neural network model for which to count parameters.
+
+    Returns:
+    --------
+    int
+        The total number of trainable parameters in the model.
+
+    Prints:
+    -------
+    A pretty-formatted table showing the number of parameters for each module,
+    followed by the total number of trainable parameters.
+    """
+    assert isinstance(module, Module), "Input must be a neural network module"  # noqa: E501
+    table = PrettyTable(["Modules", "Parameters"])
+    total_params = 0
+    for name, parameter in module.states(ParamState).items():
+        param = parameter.numel()
+        table.add_row([name, _format_parameter_count(param, precision=precision)])
+        total_params += param
+    table.add_row(["Total", _format_parameter_count(total_params, precision=precision)])
+    print(table)
+    if return_table:
+        return table, total_params
+    return total_params