brainstate 0.1.1__py2.py3-none-any.whl → 0.1.3__py2.py3-none-any.whl

brainstate/nn/{_dyn_impl/_dynamics_synapse_test.py → _synapse_test.py}

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
@@ -21,7 +20,7 @@ import brainunit as u
 import jax.numpy as jnp
 import pytest
 
-import brainstate as bst
+import brainstate
 from brainstate.nn import Expon, STP, STD
 
 
@@ -32,7 +31,7 @@ class TestSynapse(unittest.TestCase):
         self.time_steps = 100
 
     def generate_input(self):
-        return bst.random.randn(self.time_steps, self.batch_size, self.in_size) * u.mS
+        return brainstate.random.randn(self.time_steps, self.batch_size, self.in_size) * u.mS
 
     def test_expon_synapse(self):
         tau = 20.0 * u.ms
@@ -46,8 +45,8 @@ class TestSynapse(unittest.TestCase):
 
         # Test forward pass
         state = synapse.init_state(self.batch_size)
-        call = bst.compile.jit(synapse)
-        with bst.environ.context(dt=0.1 * u.ms):
+        call = brainstate.compile.jit(synapse)
+        with brainstate.environ.context(dt=0.1 * u.ms):
             for t in range(self.time_steps):
                 out = call(inputs[t])
                 self.assertEqual(out.shape, (self.batch_size, self.in_size))
@@ -75,7 +74,7 @@ class TestSynapse(unittest.TestCase):
 
         # Test forward pass
         state = synapse.init_state(self.batch_size)
-        call = bst.compile.jit(synapse)
+        call = brainstate.compile.jit(synapse)
         for t in range(self.time_steps):
             out = call(inputs[t])
             self.assertEqual(out.shape, (self.batch_size, self.in_size))
@@ -118,15 +117,15 @@ class TestSynapse(unittest.TestCase):
             self.assertEqual(synapse.in_size, in_size)
             self.assertEqual(synapse.out_size, in_size)
 
-            inputs = bst.random.randn(self.time_steps, self.batch_size, *in_size) * u.mS
+            inputs = brainstate.random.randn(self.time_steps, self.batch_size, *in_size) * u.mS
             state = synapse.init_state(self.batch_size)
-            call = bst.compile.jit(synapse)
-            with bst.environ.context(dt=0.1 * u.ms):
+            call = brainstate.compile.jit(synapse)
+            with brainstate.environ.context(dt=0.1 * u.ms):
                 for t in range(self.time_steps):
                     out = call(inputs[t])
                     self.assertEqual(out.shape, (self.batch_size, *in_size))
 
 
 if __name__ == '__main__':
-    with bst.environ.context(dt=0.1):
+    with brainstate.environ.context(dt=0.1):
         unittest.main()

brainstate/nn/_synaptic_projection.py

@@ -0,0 +1,133 @@
+# 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 Callable, Union
+
+import brainunit as u
+
+from brainstate._compatible_import import brainevent
+from brainstate.mixin import ParamDescriber, AlignPost, UpdateReturn
+from ._dynamics import Dynamics, Projection
+from ._projection import AlignPostProj, RawProj
+from ._stp import ShortTermPlasticity
+from ._synapse import Synapse
+from ._synouts import SynOut
+
+__all__ = [
+    'align_pre_projection',
+    'align_post_projection',
+]
+
+
+class align_pre_projection(Projection):
+    """
+    Represents a pre-synaptic alignment projection mechanism.
+
+    This class inherits from the `Projection` base class and is designed to
+    manage the pre-synaptic alignment process in neural network simulations.
+    It takes into account pre-synaptic dynamics, synaptic properties, delays,
+    communication functions, synaptic outputs, post-synaptic dynamics, and
+    short-term plasticity.
+
+    Attributes:
+        pre (Dynamics): The pre-synaptic dynamics object.
+        syn (Synapse): The synaptic object after pre-synaptic alignment.
+        delay (u.Quantity[u.second]): The output delay from the synapse.
+        projection (RawProj): The raw projection object handling communication,
+            output, and post-synaptic dynamics.
+        stp (ShortTermPlasticity, optional): The short-term plasticity object,
+            defaults to None.
+    """
+
+    def __init__(
+        self,
+        pre: Dynamics,
+        syn: Union[Synapse, ParamDescriber[Synapse]],
+        delay: u.Quantity[u.second] | None,
+        comm: Callable,
+        out: SynOut,
+        post: Dynamics,
+        stp: ShortTermPlasticity = None,
+    ):
+        super().__init__()
+        pre = pre
+        syn: Synapse = pre.align_pre(syn)
+        assert isinstance(syn, UpdateReturn), "Synapse must implement UpdateReturn interface"
+        # require "syn" implement the "update_return()" function
+        self.delay = syn.output_delay(delay)
+        self.projection = RawProj(comm=comm, out=out, post=post)
+        self.stp = stp
+
+    def update(self):
+        x = self.delay()
+        if self.stp is not None:
+            x = self.stp(x)
+        return self.projection(x)
+
+
+class align_post_projection(Projection):
+    """
+    Represents a post-synaptic alignment projection mechanism.
+
+    This class inherits from the `Projection` base class and is designed to
+    manage the post-synaptic alignment process in neural network simulations.
+    It takes into account spike generators, communication functions, synaptic
+    properties, synaptic outputs, post-synaptic dynamics, and short-term plasticity.
+
+    Args:
+        *spike_generator: Callable(s) that generate spike events or transform input spikes.
+        comm (Callable): Communication function for the projection.
+        syn (Union[AlignPost, ParamDescriber[AlignPost]]): The post-synaptic alignment object or its parameter describer.
+        out (Union[SynOut, ParamDescriber[SynOut]]): The synaptic output object or its parameter describer.
+        post (Dynamics): The post-synaptic dynamics object.
+        stp (ShortTermPlasticity, optional): The short-term plasticity object, defaults to None.
+
+    """
+    def __init__(
+        self,
+        *spike_generator,
+        comm: Callable,
+        syn: Union[AlignPost, ParamDescriber[AlignPost]],
+        out: Union[SynOut, ParamDescriber[SynOut]],
+        post: Dynamics,
+        stp: ShortTermPlasticity = None,
+    ):
+        super().__init__()
+        self.spike_generator = spike_generator
+        self.projection = AlignPostProj(comm=comm, syn=syn, out=out, post=post)
+        self.stp = stp
+
+    def update(self, *x):
+        for fun in self.spike_generator:
+            x = fun(*x)
+            if isinstance(x, (tuple, list)):
+                x = tuple(x)
+            else:
+                x = (x,)
+        assert len(x) == 1, "Spike generator must return a single value or a tuple/list of values"
+        x = brainevent.BinaryArray(x[0])  # Ensure input is a BinaryFloat for spike generation
+        if self.stp is not None:
+            x = brainevent.MaskedFloat(self.stp(x))  # Ensure STP output is a MaskedFloat
+        return self.projection(x)
+
+
+class align_pre_ltp(Projection):
+    pass
+
+
+class align_post_ltp(Projection):
+    pass

brainstate/nn/{_dynamics/_synouts.py → _synouts.py}

@@ -19,8 +19,8 @@ import brainunit as u
 import jax.numpy as jnp
 
 from brainstate.mixin import BindCondData
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike
+from ._module import Module
 
 __all__ = [
     'SynOut', 'COBA', 'CUBA', 'MgBlock',
@@ -47,6 +47,9 @@ class SynOut(Module, BindCondData):
         ret = self.update(self._conductance, *args, **kwargs)
         return ret
 
+    def update(self, conductance, potential):
+        raise NotImplementedError
+
 
 class COBA(SynOut):
     r"""

brainstate/nn/{_dynamics/_synouts_test.py → _synouts_test.py}

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
@@ -21,7 +20,7 @@ import brainunit as u
 import jax.numpy as jnp
 import numpy as np
 
-import brainstate as bst
+import brainstate
 
 
 class TestSynOutModels(unittest.TestCase):
@@ -35,19 +34,19 @@ class TestSynOutModels(unittest.TestCase):
         self.V_offset = jnp.array([0.0])
 
     def test_COBA(self):
-        model = bst.nn.COBA(E=self.E)
+        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 = bst.nn.CUBA()
+        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 = bst.nn.MgBlock(E=self.E, cc_Mg=self.cc_Mg, alpha=self.alpha, beta=self.beta, V_offset=self.V_offset)
+        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

brainstate/optim/_lr_scheduler_test.py

@@ -19,12 +19,12 @@ import unittest
 
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 
 
 class TestMultiStepLR(unittest.TestCase):
     def test1(self):
-        lr = bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1)
+        lr = brainstate.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1)
         for i in range(40):
             r = lr(i)
             if i < 10:
@@ -37,7 +37,7 @@ class TestMultiStepLR(unittest.TestCase):
                 self.assertTrue(jnp.allclose(r, 0.0001))
 
     def test2(self):
-        lr = bst.compile.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
+        lr = brainstate.compile.jit(brainstate.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
         for i in range(40):
             r = lr(i)
             if i < 10:

brainstate/optim/_optax_optimizer_test.py

@@ -13,39 +13,38 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax
 import optax
 
-import brainstate as bst
+import brainstate
 
 
 class TestOptaxOptimizer(unittest.TestCase):
     def test1(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.linear1 = bst.nn.Linear(2, 3)
-                self.linear2 = bst.nn.Linear(3, 4)
+                self.linear1 = brainstate.nn.Linear(2, 3)
+                self.linear2 = brainstate.nn.Linear(3, 4)
 
             def __call__(self, x):
                 return self.linear2(self.linear1(x))
 
-        x = bst.random.randn(1, 2)
+        x = brainstate.random.randn(1, 2)
         y = jax.numpy.ones((1, 4))
 
         model = Model()
         tx = optax.adam(1e-3)
-        optimizer = bst.optim.OptaxOptimizer(tx)
-        optimizer.register_trainable_weights(model.states(bst.ParamState))
+        optimizer = brainstate.optim.OptaxOptimizer(tx)
+        optimizer.register_trainable_weights(model.states(brainstate.ParamState))
 
         loss_fn = lambda: ((model(x) - y) ** 2).mean()
         prev_loss = loss_fn()
 
-        grads = bst.augment.grad(loss_fn, model.states(bst.ParamState))()
+        grads = brainstate.augment.grad(loss_fn, model.states(brainstate.ParamState))()
         optimizer.update(grads)
 
         new_loss = loss_fn()