brainstate 0.1.0__py2.py3-none-any.whl → 0.1.0.post20241125__py2.py3-none-any.whl

benchmark/COBA_2005.py

@@ -0,0 +1,125 @@
+# 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.
+# ==============================================================================
+
+#
+# Implementation of the paper:
+#
+# - Brette, R., Rudolph, M., Carnevale, T., Hines, M., Beeman, D., Bower, J. M., et al. (2007),
+#   Simulation of networks of spiking neurons: a review of tools and strategies., J. Comput. Neurosci., 23, 3, 349–98
+#
+# which is based on the balanced network proposed by:
+#
+# - Vogels, T. P. and Abbott, L. F. (2005), Signal propagation and logic gating in networks of integrate-and-fire neurons., J. Neurosci., 25, 46, 10786–95
+#
+import os
+import sys
+
+sys.path.append('../')
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.99'
+os.environ['JAX_TRACEBACK_FILTERING'] = 'off'
+
+
+import jax
+import brainunit as u
+import time
+import brainstate as bst
+
+
+class EINet(bst.nn.DynamicsGroup):
+    def __init__(self, scale):
+        super().__init__()
+        self.n_exc = int(3200 * scale)
+        self.n_inh = int(800 * scale)
+        self.num = self.n_exc + self.n_inh
+        self.N = bst.nn.LIFRef(self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
+                               tau=20. * u.ms, tau_ref=5. * u.ms,
+                               V_initializer=bst.init.Normal(-55., 2., unit=u.mV))
+        self.E = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=0.6 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=5. * u.ms),
+            out=bst.nn.COBA.desc(E=0. * u.mV),
+            post=self.N
+        )
+        self.I = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=6.7 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=10. * u.ms),
+            out=bst.nn.COBA.desc(E=-80. * u.mV),
+            post=self.N
+        )
+
+    def init_state(self, *args, **kwargs):
+        self.rate = bst.ShortTermState(u.math.zeros(self.num))
+
+    def update(self, t, inp):
+        with bst.environ.context(t=t):
+            spk = self.N.get_spike() != 0.
+            self.E(spk[:self.n_exc])
+            self.I(spk[self.n_exc:])
+            self.N(inp)
+            self.rate.value += self.N.get_spike()
+
+
+@bst.compile.jit(static_argnums=0)
+def run(scale: float):
+    # network
+    net = EINet(scale)
+    bst.nn.init_all_states(net)
+
+    duration = 1e4 * u.ms
+    # simulation
+    with bst.environ.context(dt=0.1 * u.ms):
+        times = u.math.arange(0. * u.ms, duration, bst.environ.get_dt())
+        bst.compile.for_loop(lambda t: net.update(t, 20. * u.mA), times)
+
+    return net.num, net.rate.value.sum() / net.num / duration.to_decimal(u.second)
+
+
+for s in [1, 2, 4, 6, 8, 10, 20, 40, 60, 80, 100]:
+    jax.block_until_ready(run(s))
+
+    t0 = time.time()
+    n, rate = jax.block_until_ready(run(s))
+    t1 = time.time()
+    print(f'scale={s}, size={n}, time = {t1 - t0} s, firing rate = {rate} Hz')
+
+
+# A6000 NVIDIA GPU
+
+# scale=1, size=4000, time = 2.659956455230713 s, firing rate = 50.62445068359375 Hz
+# scale=2, size=8000, time = 2.7318649291992188 s, firing rate = 50.613040924072266 Hz
+# scale=4, size=16000, time = 2.807222604751587 s, firing rate = 50.60573959350586 Hz
+# scale=6, size=24000, time = 3.026782512664795 s, firing rate = 50.60918045043945 Hz
+# scale=8, size=32000, time = 3.1258811950683594 s, firing rate = 50.607574462890625 Hz
+# scale=10, size=40000, time = 3.172346353530884 s, firing rate = 50.60942840576172 Hz
+# scale=20, size=80000, time = 3.751189947128296 s, firing rate = 50.612369537353516 Hz
+# scale=40, size=160000, time = 5.0217814445495605 s, firing rate = 50.617958068847656 Hz
+# scale=60, size=240000, time = 7.002646207809448 s, firing rate = 50.61948776245117 Hz
+# scale=80, size=320000, time = 9.384576320648193 s, firing rate = 50.618499755859375 Hz
+# scale=100, size=400000, time = 11.69654369354248 s, firing rate = 50.61605453491211 Hz
+
+
+# AMD Ryzen 7 7840HS
+
+# scale=1, size=4000, time = 4.436027526855469 s, firing rate = 50.6119270324707 Hz
+# scale=2, size=8000, time = 8.349745273590088 s, firing rate = 50.612266540527344 Hz
+# scale=4, size=16000, time = 16.39163303375244 s, firing rate = 50.61349105834961 Hz
+# scale=6, size=24000, time = 15.725558042526245 s, firing rate = 50.6125602722168 Hz
+# scale=8, size=32000, time = 21.31995177268982 s, firing rate = 50.61244583129883 Hz
+# scale=10, size=40000, time = 27.811061143875122 s, firing rate = 50.61423873901367 Hz
+# scale=20, size=80000, time = 45.54235219955444 s, firing rate = 50.61320877075195 Hz
+# scale=40, size=160000, time = 82.22228026390076 s, firing rate = 50.61309814453125 Hz
+# scale=60, size=240000, time = 125.44037556648254 s, firing rate = 50.613094329833984 Hz
+# scale=80, size=320000, time = 171.20458459854126 s, firing rate = 50.613365173339844 Hz
+# scale=100, size=400000, time = 215.4547393321991 s, firing rate = 50.6129150390625 Hz

benchmark/CUBA_2005.py

@@ -0,0 +1,149 @@
+# 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.
+# ==============================================================================
+
+#
+# Implementation of the paper:
+#
+# - Brette, R., Rudolph, M., Carnevale, T., Hines, M., Beeman, D., Bower, J. M., et al. (2007),
+#   Simulation of networks of spiking neurons: a review of tools and strategies., J. Comput. Neurosci., 23, 3, 349–98
+#
+# which is based on the balanced network proposed by:
+#
+# - Vogels, T. P. and Abbott, L. F. (2005), Signal propagation and logic gating in networks of integrate-and-fire neurons., J. Neurosci., 25, 46, 10786–95
+#
+
+import os
+import sys
+
+sys.path.append('../')
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.99'
+os.environ['JAX_TRACEBACK_FILTERING'] = 'off'
+
+
+import jax
+import time
+
+import brainunit as u
+
+import brainstate as bst
+
+
+
+class FixedProb(bst.nn.Module):
+    def __init__(self, n_pre, n_post, prob, weight):
+        super().__init__()
+        self.prob = prob
+        self.weight = weight
+        self.n_pre = n_pre
+        self.n_post = n_post
+
+        self.mask = bst.random.rand(n_pre, n_post) < prob
+
+    def update(self, x):
+        return (x @ self.mask) * self.weight
+
+
+class EINet(bst.nn.DynamicsGroup):
+    def __init__(self, scale=1.0):
+        super().__init__()
+        self.n_exc = int(3200 * scale)
+        self.n_inh = int(800 * scale)
+        self.num = self.n_exc + self.n_inh
+        self.N = bst.nn.LIFRef(
+            self.num, V_rest=-49. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
+            tau=20. * u.ms, tau_ref=5. * u.ms,
+            V_initializer=bst.init.Normal(-55., 2., unit=u.mV)
+        )
+        self.E = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=1.62 * u.mS),
+            # comm=FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=1.62 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=5. * u.ms),
+            out=bst.nn.CUBA.desc(scale=u.volt),
+            post=self.N
+        )
+        self.I = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=-9.0 * u.mS),
+            # comm=FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=-9.0 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=10. * u.ms),
+            out=bst.nn.CUBA.desc(scale=u.volt),
+            post=self.N
+        )
+
+    def init_state(self, *args, **kwargs):
+        self.rate = bst.ShortTermState(u.math.zeros(self.num))
+
+    def update(self, t, inp):
+        with bst.environ.context(t=t):
+            spk = self.N.get_spike()
+            self.E(spk[:self.n_exc])
+            self.I(spk[self.n_exc:])
+            self.N(inp)
+            self.rate.value += self.N.get_spike()
+
+
+@bst.compile.jit(static_argnums=0)
+def run(scale: float):
+    # network
+    net = EINet(scale)
+    bst.nn.init_all_states(net)
+
+    duration = 1e4 * u.ms
+    # simulation
+    with bst.environ.context(dt=0.1 * u.ms):
+        times = u.math.arange(0. * u.ms, duration, bst.environ.get_dt())
+        bst.compile.for_loop(lambda t: net.update(t, 20. * u.mA), times,
+                             # pbar=bst.compile.ProgressBar(100)
+                             )
+
+    return net.num, net.rate.value.sum() / net.num / duration.to_decimal(u.second)
+
+
+for s in [1, 2, 4, 6, 8, 10, 20, 40, 60, 80, 100]:
+    jax.block_until_ready(run(s))
+
+    t0 = time.time()
+    n, rate = jax.block_until_ready(run(s))
+    t1 = time.time()
+    print(f'scale={s}, size={n}, time = {t1 - t0} s, firing rate = {rate} Hz')
+
+
+# A6000 NVIDIA GPU
+
+# scale=1, size=4000, time = 2.6354849338531494 s, firing rate = 24.982027053833008 Hz
+# scale=2, size=8000, time = 2.6781561374664307 s, firing rate = 23.719463348388672 Hz
+# scale=4, size=16000, time = 2.7448785305023193 s, firing rate = 24.592931747436523 Hz
+# scale=6, size=24000, time = 2.8237478733062744 s, firing rate = 24.159996032714844 Hz
+# scale=8, size=32000, time = 2.9344418048858643 s, firing rate = 24.956790924072266 Hz
+# scale=10, size=40000, time = 3.042517900466919 s, firing rate = 23.644424438476562 Hz
+# scale=20, size=80000, time = 3.6727631092071533 s, firing rate = 24.226743698120117 Hz
+# scale=40, size=160000, time = 4.857396602630615 s, firing rate = 24.329742431640625 Hz
+# scale=60, size=240000, time = 6.812030792236328 s, firing rate = 24.370006561279297 Hz
+# scale=80, size=320000, time = 9.227966547012329 s, firing rate = 24.41067886352539 Hz
+# scale=100, size=400000, time = 11.405697584152222 s, firing rate = 24.32524871826172 Hz
+
+
+# AMD Ryzen 7 7840HS
+
+# scale=1, size=4000, time = 1.1661601066589355 s, firing rate = 22.438201904296875 Hz
+# scale=2, size=8000, time = 3.3255884647369385 s, firing rate = 23.868364334106445 Hz
+# scale=4, size=16000, time = 6.950139999389648 s, firing rate = 24.21693229675293 Hz
+# scale=6, size=24000, time = 10.011993169784546 s, firing rate = 24.240270614624023 Hz
+# scale=8, size=32000, time = 13.027734518051147 s, firing rate = 24.753198623657227 Hz
+# scale=10, size=40000, time = 16.449942350387573 s, firing rate = 24.7176570892334 Hz
+# scale=20, size=80000, time = 30.754598140716553 s, firing rate = 24.119956970214844 Hz
+# scale=40, size=160000, time = 63.6387836933136 s, firing rate = 24.72784996032715 Hz
+# scale=60, size=240000, time = 78.58532166481018 s, firing rate = 24.402742385864258 Hz
+# scale=80, size=320000, time = 102.4250214099884 s, firing rate = 24.59092140197754 Hz
+# scale=100, size=400000, time = 145.35173273086548 s, firing rate = 24.33751106262207 Hz

brainstate/_state.py

@@ -679,7 +679,7 @@ class StateTraceStack(Generic[A]):
         """
         for st, val in zip(self.states, self._original_state_values):
             # internal use
-            st._value = val
+            st.restore_value(val)
 
     def merge(self, *traces) -> 'StateTraceStack':
         """