brainstate 0.1.2__py2.py3-none-any.whl → 0.1.4__py2.py3-none-any.whl

brainstate/nn/{_event/_fixedprob_mv.py → _fixedprob.py}

@@ -16,19 +16,20 @@
 
 from typing import Union, Callable, Optional
 
+import brainevent
 import brainunit as u
 import jax
 import jax.numpy as jnp
 import numpy as np
 
 from brainstate import random, augment, environ, init
-from brainstate._compatible_import import brainevent
-from brainstate._state import ParamState
+from brainstate._state import ParamState, FakeState
 from brainstate.compile import for_loop
-from brainstate.nn._module import Module
 from brainstate.typing import Size, ArrayLike
+from ._module import Module
 
 __all__ = [
+    'FixedNumConn',
     'EventFixedNumConn',
     'EventFixedProb',
 ]
@@ -44,12 +45,11 @@ def init_indices_without_replace(
     rng = random.default_rng(seed)
 
     if method == 'vmap':
-        @augment.vmap
-        def rand_indices(key):
-            rng.set_key(key)
+        @augment.vmap(axis_size=n_pre)
+        def rand_indices():
             return rng.choice(n_post, size=(conn_num,), replace=False)
 
-        return rand_indices(rng.split_key(n_pre))
+        return rand_indices()
 
     elif method == 'for_loop':
         return for_loop(
@@ -61,9 +61,9 @@ def init_indices_without_replace(
         raise ValueError(f"Unknown method: {method}")
 
 
-class EventFixedNumConn(Module):
+class FixedNumConn(Module):
     """
-    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
+    The ``FixedNumConn`` module implements a fixed probability connection with CSR sparse data structure.
 
     Parameters
     ----------
@@ -77,7 +77,7 @@ class EventFixedNumConn(Module):
         If it is an integer, representing the number of connections.
     conn_weight : float or callable or jax.Array or brainunit.Quantity
         Maximum synaptic conductance, i.e., synaptic weight.
-    conn_target : str, optional
+    efferent_target : str, optional
         The target of the connection. Default is 'post', meaning that each pre-synaptic neuron connects to
         a fixed number of post-synaptic neurons. The connection number is determined by the value of ``n_conn``.
 
@@ -104,7 +104,8 @@ class EventFixedNumConn(Module):
         out_size: Size,
         conn_num: Union[int, float],
         conn_weight: Union[Callable, ArrayLike],
-        conn_target: str = 'post',  # 'pre' or 'post'
+        efferent_target: str = 'post',  # 'pre' or 'post'
+        afferent_ratio: Union[int, float] = 1.,
         allow_multi_conn: bool = True,
         seed: Optional[int] = None,
         name: Optional[str] = None,
@@ -116,11 +117,14 @@ class EventFixedNumConn(Module):
         # network parameters
         self.in_size = in_size
         self.out_size = out_size
-        self.conn_target = conn_target
-        assert conn_target in ('pre', 'post'), 'The target of the connection must be either "pre" or "post".'
+        self.efferent_target = efferent_target
+        assert efferent_target in ('pre', 'post'), 'The target of the connection must be either "pre" or "post".'
+        assert 0. <= afferent_ratio <= 1., 'Afferent ratio must be in [0, 1].'
         if isinstance(conn_num, float):
             assert 0. <= conn_num <= 1., 'Connection probability must be in [0, 1].'
-            conn_num = int(self.out_size[-1] * conn_num) if conn_target == 'post' else int(self.in_size[-1] * conn_num)
+            conn_num = (int(self.out_size[-1] * conn_num)
+                        if efferent_target == 'post' else
+                        int(self.in_size[-1] * conn_num))
         assert isinstance(conn_num, int), 'Connection number must be an integer.'
         self.conn_num = conn_num
         self.seed = seed
@@ -128,14 +132,13 @@ class EventFixedNumConn(Module):
 
         # connections
         if self.conn_num >= 1:
-            if self.conn_target == 'post':
+            if self.efferent_target == 'post':
                 n_post = self.out_size[-1]
                 n_pre = self.in_size[-1]
             else:
                 n_post = self.in_size[-1]
                 n_pre = self.out_size[-1]
 
-            # indices of post connected neurons
             with jax.ensure_compile_time_eval():
                 if allow_multi_conn:
                     rng = np.random if seed is None else np.random.RandomState(seed)
@@ -143,15 +146,83 @@ class EventFixedNumConn(Module):
                 else:
                     indices = init_indices_without_replace(self.conn_num, n_pre, n_post, seed, conn_init)
                 indices = u.math.asarray(indices, dtype=environ.ditype())
-            conn_weight = init.param(conn_weight, (n_pre, self.conn_num), allow_none=False)
-            conn_weight = u.math.asarray(conn_weight)
-            self.weight = param_type(conn_weight)
-            csr = (
-                brainevent.FixedPostNumConn((conn_weight, indices), shape=(n_pre, n_post))
-                if self.conn_target == 'post' else
-                brainevent.FixedPreNumConn((conn_weight, indices), shape=(n_pre, n_post))
-            )
-            self.conn = csr
+
+            if afferent_ratio == 1.:
+                conn_weight = u.math.asarray(init.param(conn_weight, (n_pre, self.conn_num), allow_none=False))
+                self.weight = param_type(conn_weight)
+                csr = (
+                    brainevent.FixedPostNumConn((conn_weight, indices), shape=(n_pre, n_post))
+                    if self.efferent_target == 'post' else
+                    brainevent.FixedPreNumConn((conn_weight, indices), shape=(n_pre, n_post))
+                )
+                self.conn = csr
+
+            else:
+                self.pre_selected = np.random.random(n_pre) < afferent_ratio
+                indices = indices[self.pre_selected].flatten()
+                conn_weight = u.math.asarray(init.param(conn_weight, (indices.size,), allow_none=False))
+                self.weight = param_type(conn_weight)
+                indptr = (jnp.arange(1, n_pre + 1) * self.conn_num -
+                          jnp.cumsum(~self.pre_selected) * self.conn_num)
+                indptr = jnp.insert(indptr, 0, 0)  # insert 0 at the beginning
+                csr = (
+                    brainevent.CSR((conn_weight, indices, indptr), shape=(n_pre, n_post))
+                    if self.efferent_target == 'post' else
+                    brainevent.CSC((conn_weight, indices, indptr), shape=(n_pre, n_post))
+                )
+                self.conn = csr
+
+        else:
+            conn_weight = u.math.asarray(init.param(conn_weight, (), allow_none=False))
+            self.weight = FakeState(conn_weight)
+
+    def update(self, x: jax.Array) -> Union[jax.Array, u.Quantity]:
+        if self.conn_num >= 1:
+            csr = self.conn.with_data(self.weight.value)
+            return x @ csr
+        else:
+            weight = self.weight.value
+            r = u.math.zeros(x.shape[:-1] + (self.out_size[-1],), dtype=weight.dtype)
+            r = u.maybe_decimal(u.Quantity(r, unit=u.get_unit(weight)))
+            return u.math.asarray(r, dtype=environ.dftype())
+
+
+class EventFixedNumConn(FixedNumConn):
+    """
+    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
+
+    Parameters
+    ----------
+    in_size : Size
+        Number of pre-synaptic neurons, i.e., input size.
+    out_size : Size
+        Number of post-synaptic neurons, i.e., output size.
+    conn_num : float, int
+        If it is a float, representing the probability of connection, i.e., connection probability.
+
+        If it is an integer, representing the number of connections.
+    conn_weight : float or callable or jax.Array or brainunit.Quantity
+        Maximum synaptic conductance, i.e., synaptic weight.
+    conn_target : str, optional
+        The target of the connection. Default is 'post', meaning that each pre-synaptic neuron connects to
+        a fixed number of post-synaptic neurons. The connection number is determined by the value of ``n_conn``.
+
+        If 'pre', each post-synaptic neuron connects to a fixed number of pre-synaptic neurons.
+    conn_init : str, optional
+        The initialization method of the connection weight. Default is 'vmap', meaning that the connection weight
+        is initialized by parallelized across multiple threads.
+
+        If 'for_loop', the connection weight is initialized by a for loop.
+    allow_multi_conn : bool, optional
+        Whether multiple connections are allowed from a single pre-synaptic neuron.
+        Default is True, meaning that a value of ``a`` can be selected multiple times.
+    seed: int, optional
+        Random seed. Default is None. If None, the default random seed will be used.
+    name : str, optional
+        Name of the module.
+    """
+
+    __module__ = 'brainstate.nn'
 
     def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
         if self.conn_num >= 1:

brainstate/nn/{_dyn_impl/_inputs.py → _inputs.py}

@@ -20,12 +20,11 @@ import jax
 import numpy as np
 
 from brainstate import environ, init, random
-from brainstate._state import ShortTermState
-from brainstate._state import State, maybe_state
+from brainstate._state import ShortTermState, State, maybe_state
 from brainstate.compile import while_loop
-from brainstate.nn._dynamics._dynamics_base import Dynamics, Prefetch
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size, DTypeLike
+from ._dynamics import Dynamics, Prefetch
+from ._module import Module
 
 __all__ = [
     'SpikeTime',
@@ -134,7 +133,7 @@ class PoissonSpike(Dynamics):
         self.freqs = init.param(freqs, self.varshape, allow_none=False)
 
     def update(self):
-        spikes = random.rand(self.varshape) <= (self.freqs * environ.get_dt())
+        spikes = random.rand(*self.varshape) <= (self.freqs * environ.get_dt())
         spikes = u.math.asarray(spikes, dtype=self.spk_type)
         return spikes
 

brainstate/nn/{_interaction/_linear.py → _linear.py}

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 
 from brainstate import init, functional
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size
+from ._module import Module
 
 __all__ = [
     'Linear',
@@ -350,10 +350,7 @@ class OneToOne(Module):
         self.weight = param_type(param)
 
     def update(self, pre_val):
-        pre_val, pre_unit = u.get_mantissa(pre_val), u.get_unit(pre_val)
-        w_val, w_unit = u.get_mantissa(self.weight.value['weight']), u.get_unit(self.weight.value['weight'])
-        post_val = pre_val * w_val
-        post_val = u.maybe_decimal(u.Quantity(post_val, unit=w_unit * pre_unit))
+        post_val = pre_val * self.weight.value['weight']
         if 'bias' in self.weight.value:
             post_val = post_val + self.weight.value['bias']
         return post_val

brainstate/nn/{_event/_linear_mv.py → _linear_mv.py}

@@ -19,10 +19,10 @@ import brainunit as u
 import jax
 
 from brainstate import init
-from brainstate._compatible_import import brainevent
+import brainevent
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import Size, ArrayLike
+from ._module import Module
 
 __all__ = [
     'EventLinear',

brainstate/nn/{_event/__init__.py → _ltp.py}

@@ -16,11 +16,13 @@
 # -*- coding: utf-8 -*-
 
 
-from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
-from ._linear_mv import EventLinear
+from ._synapse import Synapse
 
 __all__ = [
-    'EventLinear',
-    'EventFixedProb',
-    'EventFixedNumConn',
+    'LongTermPlasticity',
 ]
+
+
+class LongTermPlasticity(Synapse):
+    pass
+

brainstate/nn/_module.py

@@ -34,7 +34,7 @@ import numpy as np
 from brainstate._state import State
 from brainstate.graph import Node, states, nodes, flatten
 from brainstate.mixin import ParamDescriber, ParamDesc
-from brainstate.typing import PathParts
+from brainstate.typing import PathParts, Size
 from brainstate.util import FlattedDict, NestedDict, BrainStateError
 
 # maximum integer
@@ -62,8 +62,8 @@ class Module(Node, ParamDesc):
 
     __module__ = 'brainstate.nn'
 
-    _in_size: Optional[Tuple[int, ...]]
-    _out_size: Optional[Tuple[int, ...]]
+    _in_size: Optional[Size]
+    _out_size: Optional[Size]
     _name: Optional[str]
 
     if not TYPE_CHECKING:
@@ -87,7 +87,7 @@ class Module(Node, ParamDesc):
         raise AttributeError('The name of the model is read-only.')
 
     @property
-    def in_size(self) -> Tuple[int, ...]:
+    def in_size(self) -> Size:
         return self._in_size
 
     @in_size.setter
@@ -98,7 +98,7 @@ class Module(Node, ParamDesc):
         self._in_size = tuple(in_size)
 
     @property
-    def out_size(self) -> Tuple[int, ...]:
+    def out_size(self) -> Size:
         return self._out_size
 
     @out_size.setter

brainstate/nn/{_dyn_impl/_dynamics_neuron.py → _neuron.py}

@@ -22,9 +22,9 @@ import jax
 
 from brainstate import init, surrogate, environ
 from brainstate._state import HiddenState, ShortTermState
-from brainstate.nn._dynamics._dynamics_base import Dynamics
-from brainstate.nn._exp_euler import exp_euler_step
 from brainstate.typing import ArrayLike, Size
+from ._dynamics import Dynamics
+from ._exp_euler import exp_euler_step
 
 __all__ = [
     'Neuron', 'IF', 'LIF', 'LIFRef', 'ALIF',

brainstate/nn/{_interaction/_normalizations.py → _normalizations.py}

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 
 from brainstate import environ, init
 from brainstate._state import ParamState, BatchState
-from brainstate.nn._module import Module
 from brainstate.typing import DTypeLike, ArrayLike, Size, Axes
+from ._module import Module
 
 __all__ = [
     'BatchNorm0d',

brainstate/nn/{_interaction/_poolings.py → _poolings.py}

@@ -25,8 +25,8 @@ import jax.numpy as jnp
 import numpy as np
 
 from brainstate import environ
-from brainstate.nn._module import Module
 from brainstate.typing import Size
+from ._module import Module
 
 __all__ = [
     'Flatten', 'Unflatten',

brainstate/nn/{_interaction/_poolings_test.py → _poolings_test.py}

@@ -103,7 +103,7 @@ class TestPool(parameterized.TestCase):
         for target_size in [10, 9, 8, 7, 6]
     )
     def test_adaptive_pool1d(self, target_size):
-        from brainstate.nn._interaction._poolings import _adaptive_pool1d
+        from brainstate.nn._poolings import _adaptive_pool1d
 
         arr = brainstate.random.rand(100)
         op = jax.numpy.mean