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

brainstate/nn/_dynamics.py

@@ -33,7 +33,7 @@ For handling the delays:
 
 """
 
-from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, TYPE_CHECKING, Tuple
+from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, Tuple
 
 import jax
 import numpy as np
@@ -41,23 +41,26 @@ import numpy as np
 from brainstate import environ
 from brainstate._state import State
 from brainstate.graph import Node
-from brainstate.mixin import ParamDescriber
 from brainstate.typing import Size, ArrayLike
 from ._delay import StateWithDelay, Delay
 from ._module import Module
 
+T = TypeVar('T')
+
 __all__ = [
-    'DynamicsGroup',
     'Dynamics',
+
+    'receive_update_output',
+    'not_receive_update_output',
+    'receive_update_input',
+    'not_receive_update_input',
+
     'Prefetch',
     'PrefetchDelay',
     'PrefetchDelayAt',
     'OutputDelayAt',
 ]
 
-T = TypeVar('T')
-_max_order = 10
-
 
 class Dynamics(Module):
     """
@@ -122,7 +125,7 @@ class Dynamics(Module):
 
     __module__ = 'brainstate.nn'
 
-    graph_invisible_attrs = ('_before_updates', '_after_updates', '_current_inputs', '_delta_inputs')
+    graph_invisible_attrs = ()
 
     # before updates
     _before_updates: Optional[Dict[Hashable, Callable]]
@@ -136,11 +139,7 @@ class Dynamics(Module):
     # delta inputs
     _delta_inputs: Optional[Dict[str, ArrayLike | Callable]]
 
-    def __init__(
-        self,
-        in_size: Size,
-        name: Optional[str] = None,
-    ):
+    def __init__(self, in_size: Size, name: Optional[str] = None):
         # initialize
         super().__init__(name=name)
 
@@ -157,12 +156,6 @@ class Dynamics(Module):
             raise ValueError(f'"in_size" must be int, or a tuple/list of int. But we got {type(in_size)}')
         self.in_size = in_size
 
-        # current inputs
-        self._current_inputs = None
-
-        # delta inputs
-        self._delta_inputs = None
-
         # before updates
         self._before_updates = None
 
@@ -172,14 +165,6 @@ class Dynamics(Module):
         # in-/out- size of neuron population
         self.out_size = self.in_size
 
-    # def __pretty_repr_item__(self, name, value):
-    #     if name in [
-    #         '_before_updates', '_after_updates', '_current_inputs', '_delta_inputs',
-    #         '_in_size', '_out_size', '_name', '_mode',
-    #     ]:
-    #         return (name, value) if value is None else (name[1:], value)  # skip the first `_`
-    #     return super().__pretty_repr_item__(name, value)
-
     @property
     def varshape(self):
         """
@@ -200,327 +185,72 @@ class Dynamics(Module):
         """
         return self.in_size
 
-    @property
-    def current_inputs(self):
-        """
-        Get the dictionary of current inputs registered with this dynamics model.
-
-        Current inputs represent direct input currents that flow into the model.
-
-        Returns
-        -------
-        dict or None
-            A dictionary mapping keys to current input functions or values,
-            or None if no current inputs have been registered.
-
-        See Also
-        --------
-        add_current_input : Register a new current input
-        sum_current_inputs : Apply and sum all current inputs
-        delta_inputs : Dictionary of instantaneous change inputs
-        """
-        return self._current_inputs
-
-    @property
-    def delta_inputs(self):
-        """
-        Get the dictionary of delta inputs registered with this dynamics model.
-
-        Delta inputs represent instantaneous changes to state variables (dX/dt).
-
-        Returns
-        -------
-        dict or None
-            A dictionary mapping keys to delta input functions or values,
-            or None if no delta inputs have been registered.
-
-        See Also
-        --------
-        add_delta_input : Register a new delta input
-        sum_delta_inputs : Apply and sum all delta inputs
-        current_inputs : Dictionary of direct current inputs
-        """
-        return self._delta_inputs
-
-    def add_current_input(
-        self,
-        key: str,
-        inp: Union[Callable, ArrayLike],
-        label: Optional[str] = None
-    ):
-        """
-        Add a current input function or array to the dynamics model.
-
-        Current inputs represent direct input currents that can be accessed during
-        model updates through the `sum_current_inputs()` method.
-
-        Parameters
-        ----------
-        key : str
-            Unique identifier for this current input. Used to retrieve or reference
-            the input later.
-        inp : Union[Callable, ArrayLike]
-            The input data or function that generates input data.
-            - If callable: Will be called during updates with arguments passed to `sum_current_inputs()`
-            - If array-like: Will be applied once and then automatically removed from available inputs
-        label : Optional[str], default=None
-            Optional grouping label for the input. When provided, allows selective
-            processing of inputs by label in `sum_current_inputs()`.
-
-        Raises
-        ------
-        ValueError
-            If the key has already been used for a different current input.
-
-        Notes
-        -----
-        - Inputs with the same label can be processed together using the `label`
-          parameter in `sum_current_inputs()`.
-        - Non-callable inputs are consumed when used (removed after first use).
-        - Callable inputs persist and can be called repeatedly.
-
-        See Also
-        --------
-        sum_current_inputs : Sum all current inputs matching a given label
-        add_delta_input : Add a delta input function or array
-        """
-        key = _input_label_repr(key, label)
-        if self._current_inputs is None:
-            self._current_inputs = dict()
-        if key in self._current_inputs:
-            if id(self._current_inputs[key]) != id(inp):
-                raise ValueError(f'Key "{key}" has been defined and used in the current inputs of {self}.')
-        self._current_inputs[key] = inp
-
-    def add_delta_input(
-        self,
-        key: str,
-        inp: Union[Callable, ArrayLike],
-        label: Optional[str] = None
-    ):
-        """
-        Add a delta input function or array to the dynamics model.
-
-        Delta inputs represent instantaneous changes to the model state (i.e., dX/dt contributions).
-        This method registers a function or array that provides delta inputs which will be
-        accessible during model updates through the `sum_delta_inputs()` method.
-
-        Parameters
-        ----------
-        key : str
-            Unique identifier for this delta input. Used to retrieve or reference
-            the input later.
-        inp : Union[Callable, ArrayLike]
-            The input data or function that generates input data.
-            - If callable: Will be called during updates with arguments passed to `sum_delta_inputs()`
-            - If array-like: Will be applied once and then automatically removed from available inputs
-        label : Optional[str], default=None
-            Optional grouping label for the input. When provided, allows selective
-            processing of inputs by label in `sum_delta_inputs()`.
-
-        Raises
-        ------
-        ValueError
-            If the key has already been used for a different delta input.
-
-        Notes
-        -----
-        - Inputs with the same label can be processed together using the `label`
-          parameter in `sum_delta_inputs()`.
-        - Non-callable inputs are consumed when used (removed after first use).
-        - Callable inputs persist and can be called repeatedly.
-
-        See Also
-        --------
-        sum_delta_inputs : Sum all delta inputs matching a given label
-        add_current_input : Add a current input function or array
-        """
-        key = _input_label_repr(key, label)
-        if self._delta_inputs is None:
-            self._delta_inputs = dict()
-        if key in self._delta_inputs:
-            if id(self._delta_inputs[key]) != id(inp):
-                raise ValueError(f'Key "{key}" has been defined and used.')
-        self._delta_inputs[key] = inp
-
-    def get_input(self, key: str):
+    def prefetch(self, item: str) -> 'Prefetch':
         """
-        Get a registered input function by its key.
+        Create a reference to a state or variable that may not be initialized yet.
 
-        Retrieves either a current input or a delta input function that was previously
-        registered with the given key. This method checks both current_inputs and
-        delta_inputs dictionaries for the specified key.
+        This method allows accessing module attributes or states before they are
+        fully defined, acting as a placeholder that will be resolved when called.
+        Particularly useful for creating references to variables that will be defined
+        during initialization or runtime.
 
         Parameters
         ----------
-        key : str
-            The unique identifier used when the input function was registered.
+        item : str
+            The name of the attribute or state to reference.
 
         Returns
         -------
-        Callable or ArrayLike
-            The input function or array associated with the given key.
-
-        Raises
-        ------
-        ValueError
-            If no input function is found with the specified key in either
-            current_inputs or delta_inputs.
-
-        See Also
-        --------
-        add_current_input : Register a current input function
-        add_delta_input : Register a delta input function
+        Prefetch
+            A Prefetch object that provides access to the referenced item.
 
         Examples
         --------
-        >>> model = Dynamics(10)
-        >>> model.add_current_input('stimulus', lambda t: np.sin(t))
-        >>> input_func = model.get_input('stimulus')
-        >>> input_func(0.5)  # Returns sin(0.5)
+        >>> import brainstate
+        >>> import brainunit as u
+        >>> neuron = brainstate.nn.LIF(...)
+        >>> v_ref = neuron.prefetch('V')  # Reference to voltage
+        >>> v_value = v_ref()  # Get current value
+        >>> delayed_v = v_ref.delay.at(5.0 * u.ms)  # Get delayed value
         """
-        if self._current_inputs is not None and key in self._current_inputs:
-            return self._current_inputs[key]
-        elif self._delta_inputs is not None and key in self._delta_inputs:
-            return self._delta_inputs[key]
-        else:
-            raise ValueError(f'Input key {key} is not in current/delta inputs of the module {self}.')
+        return Prefetch(self, item)
 
-    def sum_current_inputs(
-        self,
-        init: Any,
-        *args,
-        label: Optional[str] = None,
-        **kwargs
-    ):
+    def prefetch_delay(self, state: str, delay_time, init: Callable = None) -> 'PrefetchDelayAt':
         """
-        Summarize all current inputs by applying and summing all registered current input functions.
-
-        This method iterates through all registered current input functions (from `.current_inputs`)
-        and applies them to calculate the total input current for the dynamics model. It adds all results
-        to the initial value provided.
+        Create a reference to a delayed state or variable in the module.
 
-        Parameters
-        ----------
-        init : Any
-            The initial value to which all current inputs will be added.
-        *args
-            Variable length argument list passed to each current input function.
-        label : Optional[str], default=None
-            If provided, only process current inputs with this label prefix.
-            When None, process all current inputs regardless of label.
-        **kwargs
-            Arbitrary keyword arguments passed to each current input function.
+        This method simplifies the process of accessing a delayed version of a state or variable
+        within the module. It first creates a prefetch reference to the specified state,
+        then specifies the delay time for accessing this state.
 
-        Returns
-        -------
-        Any
-            The initial value plus all applicable current inputs summed together.
+        Args:
+            state (str): The name of the state or variable to reference.
+            delay_time (ArrayLike): The amount of time to delay the variable access,
+                typically in time units (e.g., milliseconds).
+            init (Callable, optional): An optional initialization function to provide
+                a default value if the delayed state is not yet available.
 
-        Notes
-        -----
-        - Non-callable current inputs are applied once and then automatically removed from
-          the current_inputs dictionary.
-        - Callable current inputs remain registered for subsequent calls.
-        - When a label is provided, only current inputs with keys starting with that label
-          are applied.
-        """
-        if self._current_inputs is None:
-            return init
-        if label is None:
-            filter_fn = lambda k: True
-        else:
-            label_repr = _input_label_start(label)
-            filter_fn = lambda k: k.startswith(label_repr)
-        for key in tuple(self._current_inputs.keys()):
-            if filter_fn(key):
-                out = self._current_inputs[key]
-                if callable(out):
-                    try:
-                        init = init + out(*args, **kwargs)
-                    except Exception as e:
-                        raise ValueError(
-                            f'Error in current input value {key}: {out}\n'
-                            f'Error: {e}'
-                        ) from e
-                else:
-                    try:
-                        init = init + out
-                    except Exception as e:
-                        raise ValueError(
-                            f'Error in current input value {key}: {out}\n'
-                            f'Error: {e}'
-                        ) from e
-                    self._current_inputs.pop(key)
-        return init
-
-    def sum_delta_inputs(
-        self,
-        init: Any,
-        *args,
-        label: Optional[str] = None,
-        **kwargs
-    ):
+        Returns:
+            PrefetchDelayAt: An object that provides access to the variable at the specified delay time.
         """
-        Summarize all delta inputs by applying and summing all registered delta input functions.
+        return PrefetchDelayAt(self, state, delay_time, init=init)
 
-        This method iterates through all registered delta input functions (from `.delta_inputs`)
-        and applies them to calculate instantaneous changes to model states. It adds all results
-        to the initial value provided.
+    def output_delay(self, *delay_time) -> 'OutputDelayAt':
+        """
+        Create a reference to the delayed output of the module.
 
-        Parameters
-        ----------
-        init : Any
-            The initial value to which all delta inputs will be added.
-        *args
-            Variable length argument list passed to each delta input function.
-        label : Optional[str], default=None
-            If provided, only process delta inputs with this label prefix.
-            When None, process all delta inputs regardless of label.
-        **kwargs
-            Arbitrary keyword arguments passed to each delta input function.
+        This method simplifies the process of accessing a delayed version of the module's output.
+        It instantiates an `OutputDelayAt` object, which can be used to retrieve the output value
+        at the specified delay time.
 
-        Returns
-        -------
-        Any
-            The initial value plus all applicable delta inputs summed together.
+        Args:
+            delay (Optional[ArrayLike]): The amount of time to delay the output access,
+                typically in time units (e.g., milliseconds). Defaults to None.
 
-        Notes
-        -----
-        - Non-callable delta inputs are applied once and then automatically removed from
-          the delta_inputs dictionary.
-        - Callable delta inputs remain registered for subsequent calls.
-        - When a label is provided, only delta inputs with keys starting with that label
-          are applied.
+        Returns:
+            OutputDelayAt: An object that provides access to the module's output at the specified delay time.
         """
-        if self._delta_inputs is None:
-            return init
-        if label is None:
-            filter_fn = lambda k: True
-        else:
-            label_repr = _input_label_start(label)
-            filter_fn = lambda k: k.startswith(label_repr)
-        for key in tuple(self._delta_inputs.keys()):
-            if filter_fn(key):
-                out = self._delta_inputs[key]
-                if callable(out):
-                    try:
-                        init = init + out(*args, **kwargs)
-                    except Exception as e:
-                        raise ValueError(
-                            f'Error in delta input function {key}: {out}\n'
-                            f'Error: {e}'
-                        ) from e
-                else:
-                    try:
-                        init = init + out
-                    except Exception as e:
-                        raise ValueError(
-                            f'Error in delta input value {key}: {out}\n'
-                            f'Error: {e}'
-                        ) from e
-                    self._delta_inputs.pop(key)
-        return init
+        return OutputDelayAt(self, delay_time)
 
     @property
     def before_updates(self):
@@ -559,7 +289,7 @@ class Dynamics(Module):
         """
         return self._after_updates
 
-    def _add_before_update(self, key: Any, fun: Callable):
+    def add_before_update(self, key: Any, fun: Callable):
         """
         Register a function to be executed before the module's update.
 
@@ -585,7 +315,7 @@ class Dynamics(Module):
             raise KeyError(f'{key} has been registered in before_updates of {self}')
         self.before_updates[key] = fun
 
-    def _add_after_update(self, key: Any, fun: Callable):
+    def add_after_update(self, key: Any, fun: Callable):
         """
         Register a function to be executed after the module's update.
 
@@ -611,7 +341,7 @@ class Dynamics(Module):
             raise KeyError(f'{key} has been registered in after_updates of {self}')
         self.after_updates[key] = fun
 
-    def _get_before_update(self, key: Any):
+    def get_before_update(self, key: Any):
         """
         Retrieve a registered before-update function by its key.
 
@@ -636,7 +366,7 @@ class Dynamics(Module):
             raise KeyError(f'{key} is not registered in before_updates of {self}')
         return self.before_updates.get(key)
 
-    def _get_after_update(self, key: Any):
+    def get_after_update(self, key: Any):
         """
         Retrieve a registered after-update function by its key.
 
@@ -661,7 +391,7 @@ class Dynamics(Module):
             raise KeyError(f'{key} is not registered in after_updates of {self}')
         return self.after_updates.get(key)
 
-    def _has_before_update(self, key: Any):
+    def has_before_update(self, key: Any):
         """
         Check if a before-update function is registered with the given key.
 
@@ -679,7 +409,7 @@ class Dynamics(Module):
             return False
         return key in self.before_updates
 
-    def _has_after_update(self, key: Any):
+    def has_after_update(self, key: Any):
         """
         Check if an after-update function is registered with the given key.
 
@@ -722,120 +452,6 @@ class Dynamics(Module):
                     model(ret)
         return ret
 
-    def prefetch(self, item: str) -> 'Prefetch':
-        """
-        Create a reference to a state or variable that may not be initialized yet.
-
-        This method allows accessing module attributes or states before they are
-        fully defined, acting as a placeholder that will be resolved when called.
-        Particularly useful for creating references to variables that will be defined
-        during initialization or runtime.
-
-        Parameters
-        ----------
-        item : str
-            The name of the attribute or state to reference.
-
-        Returns
-        -------
-        Prefetch
-            A Prefetch object that provides access to the referenced item.
-
-        Examples
-        --------
-        >>> import brainstate
-        >>> import brainunit as u
-        >>> neuron = brainstate.nn.LIF(...)
-        >>> v_ref = neuron.prefetch('V')  # Reference to voltage
-        >>> v_value = v_ref()  # Get current value
-        >>> delayed_v = v_ref.delay.at(5.0 * u.ms)  # Get delayed value
-        """
-        return Prefetch(self, item)
-
-    def align_pre(self, dyn: Union[ParamDescriber[T], T]) -> T:
-        """
-        Registers a dynamics module to execute after this module.
-
-        This method establishes a sequential execution relationship where the specified
-        dynamics module will be called after this module completes its update. This
-        creates a feed-forward connection in the computational graph.
-
-        Parameters
-        ----------
-        dyn : Union[ParamDescriber[T], T]
-            The dynamics module to be executed after this module. Can be either:
-            - An instance of Dynamics
-            - A ParamDescriber that can instantiate a Dynamics object
-
-        Returns
-        -------
-        T
-            The dynamics module that was registered, allowing for method chaining.
-
-        Raises
-        ------
-        TypeError
-            If the input is not a Dynamics instance or a ParamDescriber that creates
-            a Dynamics instance.
-
-        Examples
-        --------
-        >>> import brainstate
-        >>> n1 = brainstate.nn.LIF(10)
-        >>> n1.align_pre(brainstate.nn.Expon.desc(n1.varshape))  # n2 will run after n1
-        """
-        if isinstance(dyn, Dynamics):
-            self._add_after_update(id(dyn), dyn)
-            return dyn
-        elif isinstance(dyn, ParamDescriber):
-            if not issubclass(dyn.cls, Dynamics):
-                raise TypeError(f'The input {dyn} should be an instance of {Dynamics}.')
-            if not self._has_after_update(dyn.identifier):
-                self._add_after_update(
-                    dyn.identifier,
-                    dyn() if ('in_size' in dyn.kwargs or len(dyn.args) > 0) else dyn(in_size=self.varshape)
-                )
-            return self._get_after_update(dyn.identifier)
-        else:
-            raise TypeError(f'The input {dyn} should be an instance of {Dynamics} or a delayed initializer.')
-
-    def prefetch_delay(self, state: str, delay_time, init: Callable = None) -> 'PrefetchDelayAt':
-        """
-        Create a reference to a delayed state or variable in the module.
-
-        This method simplifies the process of accessing a delayed version of a state or variable
-        within the module. It first creates a prefetch reference to the specified state,
-        then specifies the delay time for accessing this state.
-
-        Args:
-            state (str): The name of the state or variable to reference.
-            delay_time (ArrayLike): The amount of time to delay the variable access,
-                typically in time units (e.g., milliseconds).
-            init (Callable, optional): An optional initialization function to provide
-                a default value if the delayed state is not yet available.
-
-        Returns:
-            PrefetchDelayAt: An object that provides access to the variable at the specified delay time.
-        """
-        return PrefetchDelayAt(self, state, delay_time, init=init)
-
-    def output_delay(self, *delay_time) -> 'OutputDelayAt':
-        """
-        Create a reference to the delayed output of the module.
-
-        This method simplifies the process of accessing a delayed version of the module's output.
-        It instantiates an `OutputDelayAt` object, which can be used to retrieve the output value
-        at the specified delay time.
-
-        Args:
-            delay (Optional[ArrayLike]): The amount of time to delay the output access,
-                typically in time units (e.g., milliseconds). Defaults to None.
-
-        Returns:
-            OutputDelayAt: An object that provides access to the module's output at the specified delay time.
-        """
-        return OutputDelayAt(self, delay_time)
-
 
 class Prefetch(Node):
     """
@@ -1047,14 +663,14 @@ class PrefetchDelayAt(Node):
         self.delay_time = delay_time
         if len(delay_time) > 0:
             key = _get_prefetch_delay_key(item)
-            if not module._has_after_update(key):
-                module._add_after_update(
+            if not module.has_after_update(key):
+                module.add_after_update(
                     key,
                     not_receive_update_output(
                         StateWithDelay(module, item, init=init)
                     )
                 )
-            self.state_delay: StateWithDelay = module._get_after_update(key)
+            self.state_delay: StateWithDelay = module.get_after_update(key)
             self.delay_info = self.state_delay.register_delay(*delay_time)
 
     def __call__(self, *args, **kwargs):
@@ -1108,10 +724,10 @@ class OutputDelayAt(Node):
         assert isinstance(module, Dynamics), 'The module should be an instance of Dynamics.'
         self.module = module
         key = _get_output_delay_key()
-        if not module._has_after_update(key):
+        if not module.has_after_update(key):
             delay = Delay(jax.ShapeDtypeStruct(module.out_size, dtype=environ.dftype()), take_aware_unit=True)
-            module._add_after_update(key, receive_update_output(delay))
-        self.out_delay: Delay = module._get_after_update(key)
+            module.add_after_update(key, receive_update_output(delay))
+        self.out_delay: Delay = module.get_after_update(key)
         self.delay_info = self.out_delay.register_delay(*delay_time)
 
     def __call__(self, *args, **kwargs):
@@ -1138,7 +754,7 @@ def _get_prefetch_item_delay(target: Union[Prefetch, PrefetchDelay, PrefetchDela
         f'The target module should be an instance '
         f'of Dynamics. But got {target.module}.'
     )
-    delay = target.module._get_after_update(_get_prefetch_delay_key(target.item))
+    delay = target.module.get_after_update(_get_prefetch_delay_key(target.item))
     if not isinstance(delay, StateWithDelay):
         raise TypeError(f'The prefetch target should be a {StateWithDelay.__name__} when accessing '
                         f'its delay. But got {delay}.')
@@ -1187,9 +803,6 @@ def maybe_init_prefetch(target, *args, **kwargs):
         # delay.register_delay(*target.delay_time)
 
 
-DynamicsGroup = Module
-
-
 def receive_update_output(cls: object):
     """
     The decorator to mark the object (as the after updates) to receive the output of the update function.
@@ -1255,13 +868,3 @@ def not_receive_update_input(cls: object):
     if hasattr(cls, '_receive_update_input'):
         delattr(cls, '_receive_update_input')
     return cls
-
-
-def _input_label_start(label: str):
-    # unify the input label repr.
-    return f'{label} // '
-
-
-def _input_label_repr(name: str, label: Optional[str] = None):
-    # unify the input label repr.
-    return name if label is None else (_input_label_start(label) + str(name))