brainstate 0.1.0.post20250211__py2.py3-none-any.whl → 0.1.0.post20250216__py2.py3-none-any.whl

brainstate/surrogate.py

@@ -21,6 +21,8 @@ import jax.numpy as jnp
 import jax.scipy as sci
 from jax.interpreters import batching, ad, mlir
 
+from brainstate.util._pretty_pytree import PrettyObject
+
 if jax.__version_info__ < (0, 4, 38):
     from jax.core import Primitive
 else:
@@ -77,7 +79,10 @@ def _heaviside_imp(x, dx):
 
 
 def _heaviside_batching(args, axes):
-    return heaviside_p.bind(*args), [axes[0]]
+    x, dx = args
+    if axes[0] != axes[1]:
+        dx = batching.moveaxis(dx, axes[1], axes[0])
+    return heaviside_p.bind(x, dx), tuple([axes[0]])
 
 
 def _heaviside_jvp(primals, tangents):
@@ -97,7 +102,7 @@ ad.primitive_jvps[heaviside_p] = _heaviside_jvp
 mlir.register_lowering(heaviside_p, mlir.lower_fun(_heaviside_imp, multiple_results=True))
 
 
-class Surrogate(object):
+class Surrogate(PrettyObject):
     """The base surrograte gradient function.
 
     To customize a surrogate gradient function, you can inherit this class and
@@ -142,9 +147,20 @@ class Surrogate(object):
 class Sigmoid(Surrogate):
     """Spike function with the sigmoid-shaped surrogate gradient.
 
+    This class implements a spiking neuron activation with a sigmoid-shaped
+    surrogate gradient for backpropagation. It can be used in spiking neural
+    networks to approximate the non-differentiable step function during training.
+
+    Parameters
+    ----------
+    alpha : float, optional
+        A parameter controlling the steepness of the sigmoid curve in the
+        surrogate gradient. Higher values make the transition sharper.
+        Default is 4.0.
+
     See Also
     --------
-    sigmoid
+    sigmoid : Function version of this class.
 
     """
 
@@ -153,9 +169,33 @@ class Sigmoid(Surrogate):
         self.alpha = alpha
 
     def surrogate_fun(self, x):
+        """Compute the surrogate function.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The output of the surrogate function.
+        """
         return sci.special.expit(self.alpha * x)
 
     def surrogate_grad(self, x):
+        """Compute the gradient of the surrogate function.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The gradient of the surrogate function.
+        """
         sgax = sci.special.expit(x * self.alpha)
         dx = (1. - sgax) * sgax * self.alpha
         return dx
@@ -171,7 +211,12 @@ def sigmoid(
     x: jax.Array,
     alpha: float = 4.,
 ):
-    r"""Spike function with the sigmoid-shaped surrogate gradient.
+    r"""
+    Compute a spike function with a sigmoid-shaped surrogate gradient.
+
+    This function implements a spiking neuron activation with a sigmoid-shaped
+    surrogate gradient for backpropagation. It can be used in spiking neural
+    networks to approximate the non-differentiable step function during training.
 
     If `origin=False`, return the forward function:
 
@@ -210,16 +255,28 @@ def sigmoid(
 
     Parameters
     ----------
-    x: jax.Array, Array
-      The input data.
-    alpha: float
-      Parameter to control smoothness of gradient
-
+    x : jax.Array
+        The input array representing the neuron's membrane potential.
+    alpha : float, optional
+        A parameter controlling the steepness of the sigmoid curve in the
+        surrogate gradient. Higher values make the transition sharper.
+        Default is 4.0.
 
     Returns
     -------
-    out: jax.Array
-      The spiking state.
+    jax.Array
+        An array of the same shape as the input, containing binary values (0 or 1)
+        representing the spiking state of each neuron.
+        
+    Notes
+    -----
+    The forward pass uses a step function (1 for x >= 0, 0 for x < 0),
+    while the backward pass uses a sigmoid-shaped surrogate gradient for
+    smooth optimization.
+
+    The surrogate gradient is defined as:
+    g'(x) = alpha * (1 - sigmoid(alpha * x)) * sigmoid(alpha * x)
+
     """
     return Sigmoid(alpha=alpha)(x)
 
@@ -238,11 +295,15 @@ class PiecewiseQuadratic(Surrogate):
         self.alpha = alpha
 
     def surrogate_fun(self, x):
-        z = jnp.where(x < -1 / self.alpha,
-                      0.,
-                      jnp.where(x > 1 / self.alpha,
-                                1.,
-                                (-self.alpha * jnp.abs(x) / 2 + 1) * self.alpha * x + 0.5))
+        z = jnp.where(
+            x < -1 / self.alpha,
+            0.,
+            jnp.where(
+                x > 1 / self.alpha,
+                1.,
+                (-self.alpha * jnp.abs(x) / 2 + 1) * self.alpha * x + 0.5
+            )
+        )
         return z
 
     def surrogate_grad(self, x):
@@ -260,7 +321,12 @@ def piecewise_quadratic(
     x: jax.Array,
     alpha: float = 1.,
 ):
-    r"""Judge spiking state with a piecewise quadratic function [1]_ [2]_ [3]_ [4]_ [5]_.
+    r"""
+    Judge spiking state with a piecewise quadratic function [1]_ [2]_ [3]_ [4]_ [5]_.
+
+    This function implements a surrogate gradient method for spiking neural networks
+    using a piecewise quadratic function. It provides a differentiable approximation
+    of the step function used in the forward pass of spiking neurons.
 
     If `origin=False`, computes the forward function:
 
@@ -306,18 +372,29 @@ def piecewise_quadratic(
        >>> plt.legend()
        >>> plt.show()
 
-    Parameters
+  Parameters
     ----------
-    x: jax.Array, Array
-      The input data.
-    alpha: float
-      Parameter to control smoothness of gradient
-
+    x : jax.Array
+        The input array representing the neuron's membrane potential.
+    alpha : float, optional
+        A parameter controlling the steepness of the surrogate gradient.
+        Higher values result in a steeper gradient. Default is 1.0.
 
     Returns
     -------
-    out: jax.Array
-      The spiking state.
+    jax.Array
+        An array of the same shape as the input, containing binary values (0 or 1)
+        representing the spiking state of each neuron.
+
+    Notes
+    -----
+    The function uses different computations for forward and backward passes:
+    - Forward: Step function (1 for x >= 0, 0 for x < 0)
+    - Backward: Piecewise quadratic function for smooth gradient
+
+    The surrogate gradient is defined as:
+    g'(x) = 0 if |x| > 1/alpha
+            -alpha^2|x| + alpha if |x| <= 1/alpha
 
     References
     ----------
@@ -331,11 +408,22 @@ def piecewise_quadratic(
 
 
 class PiecewiseExp(Surrogate):
-    """Judge spiking state with a piecewise exponential function.
+    """
+    Judge spiking state with a piecewise exponential function.
+
+    This class implements a surrogate gradient method for spiking neural networks
+    using a piecewise exponential function. It provides a differentiable approximation
+    of the step function used in the forward pass of spiking neurons.
+
+    Parameters
+    ----------
+    alpha : float, optional
+        A parameter controlling the steepness of the surrogate gradient.
+        Higher values result in a steeper gradient. Default is 1.0.
 
     See Also
     --------
-    piecewise_exp
+    piecewise_exp : Function version of this class.
     """
 
     def __init__(self, alpha: float = 1.):
@@ -343,16 +431,62 @@ class PiecewiseExp(Surrogate):
         self.alpha = alpha
 
     def surrogate_grad(self, x):
+        """
+        Compute the surrogate gradient.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The surrogate gradient.
+        """
         dx = (self.alpha / 2) * jnp.exp(-self.alpha * jnp.abs(x))
         return dx
 
     def surrogate_fun(self, x):
-        return jnp.where(x < 0, jnp.exp(self.alpha * x) / 2, 1 - jnp.exp(-self.alpha * x) / 2)
+        """
+        Compute the surrogate function.
+
+        Parameters
+        ----------
+        x : jax.Array
+            The input array.
+
+        Returns
+        -------
+        jax.Array
+            The output of the surrogate function.
+        """
+        return jnp.where(
+            x < 0,
+            jnp.exp(self.alpha * x) / 2,
+            1 - jnp.exp(-self.alpha * x) / 2
+        )
 
     def __repr__(self):
+        """
+        Return a string representation of the PiecewiseExp instance.
+
+        Returns
+        -------
+        str
+            A string representation of the instance.
+        """
         return f'{self.__class__.__name__}(alpha={self.alpha})'
 
     def __hash__(self):
+        """
+        Compute a hash value for the PiecewiseExp instance.
+
+        Returns
+        -------
+        int
+            A hash value for the instance.
+        """
         return hash((self.__class__, self.alpha))
 
 
@@ -363,6 +497,10 @@ def piecewise_exp(
 ):
     r"""Judge spiking state with a piecewise exponential function [1]_.
 
+    This function implements a surrogate gradient method for spiking neural networks
+    using a piecewise exponential function. It provides a differentiable approximation
+    of the step function used in the forward pass of spiking neurons.
+
     If `origin=False`, computes the forward function:
 
     .. math::
@@ -403,16 +541,26 @@ def piecewise_exp(
 
     Parameters
     ----------
-    x: jax.Array, Array
-      The input data.
-    alpha: float
-      Parameter to control smoothness of gradient
-
+    x : jax.Array
+        The input array representing the neuron's membrane potential.
+    alpha : float, optional
+        A parameter controlling the steepness of the surrogate gradient.
+        Higher values result in a steeper gradient. Default is 1.0.
 
     Returns
     -------
-    out: jax.Array
-      The spiking state.
+    jax.Array
+        An array of the same shape as the input, containing binary values (0 or 1)
+        representing the spiking state of each neuron.
+
+    Notes
+    -----
+    The function uses different computations for forward and backward passes:
+    - Forward: Step function (1 for x >= 0, 0 for x < 0)
+    - Backward: Piecewise exponential function for smooth gradient
+
+    The surrogate gradient is defined as:
+    g'(x) = (alpha / 2) * exp(-alpha * |x|)
 
     References
     ----------

brainstate/transform.py

@@ -15,6 +15,3 @@
 
 # alias for compilation and augmentation functions
 
-from .compile import *
-from .augment import *
-

brainstate/typing.py

@@ -16,13 +16,17 @@
 from __future__ import annotations
 
 import builtins
+
+import brainunit as u
 import functools as ft
 import importlib
 import inspect
-
-import brainunit as u
 import jax
 import numpy as np
+from typing import (
+    Any, Callable, Hashable, List, Protocol, Tuple, TypeVar, Union,
+    runtime_checkable, TYPE_CHECKING, Generic, Sequence
+)
 
 tp = importlib.import_module("typing")
 
@@ -41,35 +45,35 @@ __all__ = [
     'Missing',
 ]
 
-K = tp.TypeVar('K')
+K = TypeVar('K')
 
 
-@tp.runtime_checkable
-class Key(tp.Hashable, tp.Protocol):
+@runtime_checkable
+class Key(Hashable, Protocol):
     def __lt__(self: K, value: K, /) -> bool:
         ...
 
 
-Ellipsis = builtins.ellipsis if tp.TYPE_CHECKING else tp.Any
+Ellipsis = builtins.ellipsis if TYPE_CHECKING else Any
 
-PathParts = tp.Tuple[Key, ...]
-Predicate = tp.Callable[[PathParts, tp.Any], bool]
-FilterLiteral = tp.Union[type, str, Predicate, bool, Ellipsis, None]
-Filter = tp.Union[FilterLiteral, tp.Tuple['Filter', ...], tp.List['Filter']]
+PathParts = Tuple[Key, ...]
+Predicate = Callable[[PathParts, Any], bool]
+FilterLiteral = Union[type, str, Predicate, bool, Ellipsis, None]
+Filter = Union[FilterLiteral, Tuple['Filter', ...], List['Filter']]
 
-_T = tp.TypeVar("_T")
+_T = TypeVar("_T")
 
-_Annotation = tp.TypeVar("_Annotation")
+_Annotation = TypeVar("_Annotation")
 
 
-class _Array(tp.Generic[_Annotation]):
+class _Array(Generic[_Annotation]):
     pass
 
 
 _Array.__module__ = "builtins"
 
 
-def _item_to_str(item: tp.Union[str, type, slice]) -> str:
+def _item_to_str(item: Union[str, type, slice]) -> str:
     if isinstance(item, slice):
         if item.step is not None:
             raise NotImplementedError
@@ -83,7 +87,7 @@ def _item_to_str(item: tp.Union[str, type, slice]) -> str:
 
 
 def _maybe_tuple_to_str(
-    item: tp.Union[str, type, slice, tp.Tuple[tp.Union[str, type, slice], ...]]
+    item: Union[str, type, slice, Tuple[Union[str, type, slice], ...]]
 ) -> str:
     if isinstance(item, tuple):
         if len(item) == 0:
@@ -113,7 +117,7 @@ class Array:
 Array.__module__ = "builtins"
 
 
-class _FakePyTree(tp.Generic[_T]):
+class _FakePyTree(Generic[_T]):
     pass
 
 
@@ -255,11 +259,10 @@ f. A structure can end with a `...`, to denote that the PyTree must be a prefix
     cases, all named pieces must already have been seen and their structures bound.
 """  # noqa: E501
 
-Size = tp.Union[int, tp.Sequence[int]]
-Axes = tp.Union[int, tp.Sequence[int]]
-SeedOrKey = tp.Union[int, jax.Array, np.ndarray]
-Shape = tp.Sequence[int]
-
+Size = Union[int, Sequence[int]]
+Axes = Union[int, Sequence[int]]
+SeedOrKey = Union[int, jax.Array, np.ndarray]
+Shape = Sequence[int]
 
 # --- Array --- #
 
@@ -267,7 +270,7 @@ Shape = tp.Sequence[int]
 # standard JAX array (i.e. not including future non-standard array types like
 # KeyArray and BInt). It's different than np.typing.ArrayLike in that it doesn't
 # accept arbitrary sequences, nor does it accept string data.
-ArrayLike = tp.Union[
+ArrayLike = Union[
     jax.Array,  # JAX array type
     np.ndarray,  # NumPy array type
     np.bool_, np.number,  # NumPy scalar types
@@ -281,7 +284,7 @@ ArrayLike = tp.Union[
 DType = np.dtype
 
 
-class SupportsDType(tp.Protocol):
+class SupportsDType(Protocol):
     @property
     def dtype(self) -> DType: ...
 
@@ -291,9 +294,9 @@ class SupportsDType(tp.Protocol):
 # because JAX doesn't support objects or structured dtypes.
 # Unlike np.typing.DTypeLike, we exclude None, and instead require
 # explicit annotations when None is acceptable.
-DTypeLike = tp.Union[
+DTypeLike = Union[
     str,  # like 'float32', 'int32'
-    type[tp.Any],  # like np.float32, np.int32, float, int
+    type[Any],  # like np.float32, np.int32, float, int
     np.dtype,  # like np.dtype('float32'), np.dtype('int32')
     SupportsDType,  # like jnp.float32, jnp.int32
 ]

brainstate/util/__init__.py

@@ -13,36 +13,38 @@
 # limitations under the License.
 # ==============================================================================
 
-from ._dict import *
-from ._dict import __all__ as _mapping_all
+from . import filter
 from ._error import *
 from ._error import __all__ as _error_all
-from ._filter import *
-from ._filter import __all__ as _filter_all
 from ._others import *
 from ._others import __all__ as _others_all
+from ._pretty_pytree import *
+from ._pretty_pytree import __all__ as _mapping_all
 from ._pretty_repr import *
 from ._pretty_repr import __all__ as _pretty_repr_all
+from ._pretty_table import *
+from ._pretty_table import __all__ as _table_all
 from ._scaling import *
 from ._scaling import __all__ as _mem_scale_all
 from ._struct import *
 from ._struct import __all__ as _struct_all
 
 __all__ = (
-    _others_all
+    ['filter']
+    + _others_all
     + _mem_scale_all
-    + _filter_all
     + _pretty_repr_all
     + _struct_all
     + _error_all
     + _mapping_all
+    + _table_all
 )
 del (
     _others_all,
     _mem_scale_all,
-    _filter_all,
     _pretty_repr_all,
     _struct_all,
     _error_all,
     _mapping_all,
+    _table_all,
 )

brainstate/util/_caller.py

@@ -18,9 +18,8 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Any, TypeVar, Protocol, Generic
-
 import jax
+from typing import Any, TypeVar, Protocol, Generic
 
 __all__ = [
     'DelayedAccessor',

brainstate/util/_error.py

@@ -21,8 +21,35 @@ __all__ = [
 
 
 class BrainStateError(Exception):
+    """
+    A custom exception class for BrainState-related errors.
+
+    This exception is raised when a BrainState-specific error occurs during
+    the execution of the program. It serves as a base class for more specific
+    BrainState exceptions.
+
+    Attributes:
+        Inherits all attributes from the built-in Exception class.
+
+    Usage::
+
+        raise BrainStateError("A BrainState-specific error occurred.")
+    """
     pass
 
 
 class TraceContextError(BrainStateError):
+    """
+    A custom exception class for trace context-related errors in BrainState.
+
+    This exception is raised when an error occurs specifically related to
+    trace context operations or manipulations within the BrainState framework.
+
+    Attributes:
+        Inherits all attributes from the BrainStateError class.
+
+    Usage::
+
+        raise TraceContextError("An error occurred while handling trace context.")
+    """
     pass

brainstate/util/_others.py

@@ -15,20 +15,21 @@
 
 from __future__ import annotations
 
+import gc
+
 import copy
 import functools
-import gc
+import jax
 import threading
 import types
 from collections.abc import Iterable
-from typing import Any, Callable, Tuple, Union, Dict
-
-import jax
 from jax.lib import xla_bridge
+from typing import Any, Callable, Tuple, Union, Dict
 
 from brainstate._utils import set_module_as
 
 __all__ = [
+    'split_total',
     'clear_buffer_memory',
     'not_instance_eval',
     'is_instance_eval',
@@ -37,6 +38,61 @@ __all__ = [
 ]
 
 
+def split_total(
+    total: int,
+    fraction: Union[int, float],
+) -> int:
+    """
+   Calculate the number of epochs for simulation based on a total and a fraction.
+
+   This function determines the number of epochs to simulate given a total number
+   of epochs and either a fraction or a specific number of epochs to run.
+
+   Parameters:
+   -----------
+   total : int
+       The total number of epochs. Must be a positive integer.
+   fraction : Union[int, float]
+       If ``float``: A value between 0 and 1 representing the fraction of total epochs to run.
+       If ``int``: The specific number of epochs to run, must not exceed the total.
+
+   Returns:
+   --------
+   int
+       The calculated number of epochs to simulate.
+
+   Raises:
+   -------
+   ValueError
+       If total is not positive, fraction is negative, or if fraction as float is > 1
+       or as int is > total.
+   AssertionError
+       If total is not an integer.
+   """
+    assert isinstance(total, int), "Length must be an integer."
+    if total <= 0:
+        raise ValueError("'total' must be a positive integer.")
+    if fraction < 0:
+        raise ValueError("'fraction' value cannot be negative.")
+
+    if isinstance(fraction, float):
+        if fraction < 0:
+            raise ValueError("'fraction' value cannot be negative.")
+        if fraction > 1:
+            raise ValueError("'fraction' value cannot be greater than 1.")
+        return int(total * fraction)
+
+    elif isinstance(fraction, int):
+        if fraction < 0:
+            raise ValueError("'fraction' value cannot be negative.")
+        if fraction > total:
+            raise ValueError("'fraction' value cannot be greater than total.")
+        return fraction
+
+    else:
+        raise ValueError("'fraction' must be an integer or float.")
+
+
 class NameContext(threading.local):
     def __init__(self):
         self.typed_names: Dict[str, int] = {}
@@ -249,17 +305,6 @@ class DictManager(dict):
         else:
             raise ValueError(f'Unsupported method: {by}')
 
-    def union_by_value_ids(self, other: dict):
-        """
-        Union the stack by the value ids.
-
-        Args:
-          other:
-
-        Returns:
-
-        """
-
     def __add__(self, other: dict):
         """
         Compose other instance of dict.