brainstate 0.0.2.post20241009__py2.py3-none-any.whl → 0.1.0__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (175) hide show
  1. brainstate/__init__.py +31 -11
  2. brainstate/_state.py +760 -316
  3. brainstate/_state_test.py +41 -12
  4. brainstate/_utils.py +31 -4
  5. brainstate/augment/__init__.py +40 -0
  6. brainstate/augment/_autograd.py +608 -0
  7. brainstate/augment/_autograd_test.py +1193 -0
  8. brainstate/augment/_eval_shape.py +102 -0
  9. brainstate/augment/_eval_shape_test.py +40 -0
  10. brainstate/augment/_mapping.py +525 -0
  11. brainstate/augment/_mapping_test.py +210 -0
  12. brainstate/augment/_random.py +99 -0
  13. brainstate/{transform → compile}/__init__.py +25 -13
  14. brainstate/compile/_ad_checkpoint.py +204 -0
  15. brainstate/compile/_ad_checkpoint_test.py +51 -0
  16. brainstate/compile/_conditions.py +259 -0
  17. brainstate/compile/_conditions_test.py +221 -0
  18. brainstate/compile/_error_if.py +94 -0
  19. brainstate/compile/_error_if_test.py +54 -0
  20. brainstate/compile/_jit.py +314 -0
  21. brainstate/compile/_jit_test.py +143 -0
  22. brainstate/compile/_loop_collect_return.py +516 -0
  23. brainstate/compile/_loop_collect_return_test.py +59 -0
  24. brainstate/compile/_loop_no_collection.py +185 -0
  25. brainstate/compile/_loop_no_collection_test.py +51 -0
  26. brainstate/compile/_make_jaxpr.py +756 -0
  27. brainstate/compile/_make_jaxpr_test.py +134 -0
  28. brainstate/compile/_progress_bar.py +111 -0
  29. brainstate/compile/_unvmap.py +159 -0
  30. brainstate/compile/_util.py +147 -0
  31. brainstate/environ.py +408 -381
  32. brainstate/environ_test.py +34 -32
  33. brainstate/{nn/event → event}/__init__.py +6 -6
  34. brainstate/event/_csr.py +308 -0
  35. brainstate/event/_csr_test.py +118 -0
  36. brainstate/event/_fixed_probability.py +271 -0
  37. brainstate/event/_fixed_probability_test.py +128 -0
  38. brainstate/event/_linear.py +219 -0
  39. brainstate/event/_linear_test.py +112 -0
  40. brainstate/{nn/event → event}/_misc.py +7 -7
  41. brainstate/functional/_activations.py +521 -511
  42. brainstate/functional/_activations_test.py +300 -300
  43. brainstate/functional/_normalization.py +43 -43
  44. brainstate/functional/_others.py +15 -15
  45. brainstate/functional/_spikes.py +49 -49
  46. brainstate/graph/__init__.py +33 -0
  47. brainstate/graph/_graph_context.py +443 -0
  48. brainstate/graph/_graph_context_test.py +65 -0
  49. brainstate/graph/_graph_convert.py +246 -0
  50. brainstate/graph/_graph_node.py +300 -0
  51. brainstate/graph/_graph_node_test.py +75 -0
  52. brainstate/graph/_graph_operation.py +1746 -0
  53. brainstate/graph/_graph_operation_test.py +724 -0
  54. brainstate/init/_base.py +28 -10
  55. brainstate/init/_generic.py +175 -172
  56. brainstate/init/_random_inits.py +470 -415
  57. brainstate/init/_random_inits_test.py +150 -0
  58. brainstate/init/_regular_inits.py +66 -69
  59. brainstate/init/_regular_inits_test.py +51 -0
  60. brainstate/mixin.py +236 -244
  61. brainstate/mixin_test.py +44 -46
  62. brainstate/nn/__init__.py +26 -51
  63. brainstate/nn/_collective_ops.py +199 -0
  64. brainstate/nn/_dyn_impl/__init__.py +46 -0
  65. brainstate/nn/_dyn_impl/_dynamics_neuron.py +290 -0
  66. brainstate/nn/_dyn_impl/_dynamics_neuron_test.py +162 -0
  67. brainstate/nn/_dyn_impl/_dynamics_synapse.py +320 -0
  68. brainstate/nn/_dyn_impl/_dynamics_synapse_test.py +132 -0
  69. brainstate/nn/_dyn_impl/_inputs.py +154 -0
  70. brainstate/nn/{_projection/__init__.py → _dyn_impl/_projection_alignpost.py} +6 -13
  71. brainstate/nn/_dyn_impl/_rate_rnns.py +400 -0
  72. brainstate/nn/_dyn_impl/_rate_rnns_test.py +64 -0
  73. brainstate/nn/_dyn_impl/_readout.py +128 -0
  74. brainstate/nn/_dyn_impl/_readout_test.py +54 -0
  75. brainstate/nn/_dynamics/__init__.py +37 -0
  76. brainstate/nn/_dynamics/_dynamics_base.py +631 -0
  77. brainstate/nn/_dynamics/_dynamics_base_test.py +79 -0
  78. brainstate/nn/_dynamics/_projection_base.py +346 -0
  79. brainstate/nn/_dynamics/_state_delay.py +453 -0
  80. brainstate/nn/_dynamics/_synouts.py +161 -0
  81. brainstate/nn/_dynamics/_synouts_test.py +58 -0
  82. brainstate/nn/_elementwise/__init__.py +22 -0
  83. brainstate/nn/_elementwise/_dropout.py +418 -0
  84. brainstate/nn/_elementwise/_dropout_test.py +100 -0
  85. brainstate/nn/_elementwise/_elementwise.py +1122 -0
  86. brainstate/nn/_elementwise/_elementwise_test.py +171 -0
  87. brainstate/nn/_exp_euler.py +97 -0
  88. brainstate/nn/_exp_euler_test.py +36 -0
  89. brainstate/nn/_interaction/__init__.py +32 -0
  90. brainstate/nn/_interaction/_connections.py +726 -0
  91. brainstate/nn/_interaction/_connections_test.py +254 -0
  92. brainstate/nn/_interaction/_embedding.py +59 -0
  93. brainstate/nn/_interaction/_normalizations.py +388 -0
  94. brainstate/nn/_interaction/_normalizations_test.py +75 -0
  95. brainstate/nn/_interaction/_poolings.py +1179 -0
  96. brainstate/nn/_interaction/_poolings_test.py +219 -0
  97. brainstate/nn/_module.py +328 -0
  98. brainstate/nn/_module_test.py +211 -0
  99. brainstate/nn/metrics.py +309 -309
  100. brainstate/optim/__init__.py +14 -2
  101. brainstate/optim/_base.py +66 -0
  102. brainstate/optim/_lr_scheduler.py +363 -400
  103. brainstate/optim/_lr_scheduler_test.py +25 -24
  104. brainstate/optim/_optax_optimizer.py +103 -176
  105. brainstate/optim/_optax_optimizer_test.py +41 -1
  106. brainstate/optim/_sgd_optimizer.py +950 -1025
  107. brainstate/random/_rand_funs.py +3269 -3268
  108. brainstate/random/_rand_funs_test.py +568 -0
  109. brainstate/random/_rand_seed.py +149 -117
  110. brainstate/random/_rand_seed_test.py +50 -0
  111. brainstate/random/_rand_state.py +1360 -1318
  112. brainstate/random/_random_for_unit.py +13 -13
  113. brainstate/surrogate.py +1262 -1243
  114. brainstate/{nn/_projection/_utils.py → transform.py} +1 -2
  115. brainstate/typing.py +157 -130
  116. brainstate/util/__init__.py +52 -0
  117. brainstate/util/_caller.py +100 -0
  118. brainstate/util/_dict.py +734 -0
  119. brainstate/util/_dict_test.py +160 -0
  120. brainstate/util/_error.py +28 -0
  121. brainstate/util/_filter.py +178 -0
  122. brainstate/util/_others.py +497 -0
  123. brainstate/util/_pretty_repr.py +208 -0
  124. brainstate/util/_scaling.py +260 -0
  125. brainstate/util/_struct.py +524 -0
  126. brainstate/util/_tracers.py +75 -0
  127. brainstate/{_visualization.py → util/_visualization.py} +16 -16
  128. {brainstate-0.0.2.post20241009.dist-info → brainstate-0.1.0.dist-info}/METADATA +11 -11
  129. brainstate-0.1.0.dist-info/RECORD +135 -0
  130. brainstate/_module.py +0 -1637
  131. brainstate/_module_test.py +0 -207
  132. brainstate/nn/_base.py +0 -251
  133. brainstate/nn/_connections.py +0 -686
  134. brainstate/nn/_dynamics.py +0 -426
  135. brainstate/nn/_elementwise.py +0 -1438
  136. brainstate/nn/_embedding.py +0 -66
  137. brainstate/nn/_misc.py +0 -133
  138. brainstate/nn/_normalizations.py +0 -389
  139. brainstate/nn/_others.py +0 -101
  140. brainstate/nn/_poolings.py +0 -1229
  141. brainstate/nn/_poolings_test.py +0 -231
  142. brainstate/nn/_projection/_align_post.py +0 -546
  143. brainstate/nn/_projection/_align_pre.py +0 -599
  144. brainstate/nn/_projection/_delta.py +0 -241
  145. brainstate/nn/_projection/_vanilla.py +0 -101
  146. brainstate/nn/_rate_rnns.py +0 -410
  147. brainstate/nn/_readout.py +0 -136
  148. brainstate/nn/_synouts.py +0 -166
  149. brainstate/nn/event/csr.py +0 -312
  150. brainstate/nn/event/csr_test.py +0 -118
  151. brainstate/nn/event/fixed_probability.py +0 -276
  152. brainstate/nn/event/fixed_probability_test.py +0 -127
  153. brainstate/nn/event/linear.py +0 -220
  154. brainstate/nn/event/linear_test.py +0 -111
  155. brainstate/random/random_test.py +0 -593
  156. brainstate/transform/_autograd.py +0 -585
  157. brainstate/transform/_autograd_test.py +0 -1181
  158. brainstate/transform/_conditions.py +0 -334
  159. brainstate/transform/_conditions_test.py +0 -220
  160. brainstate/transform/_error_if.py +0 -94
  161. brainstate/transform/_error_if_test.py +0 -55
  162. brainstate/transform/_jit.py +0 -265
  163. brainstate/transform/_jit_test.py +0 -118
  164. brainstate/transform/_loop_collect_return.py +0 -502
  165. brainstate/transform/_loop_no_collection.py +0 -170
  166. brainstate/transform/_make_jaxpr.py +0 -739
  167. brainstate/transform/_make_jaxpr_test.py +0 -131
  168. brainstate/transform/_mapping.py +0 -109
  169. brainstate/transform/_progress_bar.py +0 -111
  170. brainstate/transform/_unvmap.py +0 -143
  171. brainstate/util.py +0 -746
  172. brainstate-0.0.2.post20241009.dist-info/RECORD +0 -87
  173. {brainstate-0.0.2.post20241009.dist-info → brainstate-0.1.0.dist-info}/LICENSE +0 -0
  174. {brainstate-0.0.2.post20241009.dist-info → brainstate-0.1.0.dist-info}/WHEEL +0 -0
  175. {brainstate-0.0.2.post20241009.dist-info → brainstate-0.1.0.dist-info}/top_level.txt +0 -0
@@ -1,241 +0,0 @@
1
- # Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # http://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- # ==============================================================================
15
-
16
- from typing import Optional, Union
17
-
18
- from brainstate._module import (Module, Dynamics, Projection, ReceiveInputProj,
19
- UpdateReturn, register_delay_of_target)
20
- from brainstate._utils import set_module_as
21
- from brainstate.mixin import (Mode, BindCondData)
22
- from ._utils import is_instance
23
-
24
- __all__ = [
25
- 'HalfProjDelta', 'FullProjDelta',
26
- ]
27
-
28
-
29
- class _Delta(BindCondData):
30
- def __init__(self):
31
- self._cond = None
32
-
33
- def bind_cond(self, cond):
34
- self._cond = cond
35
-
36
- def __call__(self, *args, **kwargs):
37
- r = self._cond
38
- return r
39
-
40
-
41
- @set_module_as('brainstate.nn')
42
- class HalfProjDelta(Projection):
43
- """Defining the half-part of the synaptic projection for the Delta synapse model.
44
-
45
- The synaptic projection requires the input is the spiking data, otherwise
46
- the synapse is not the Delta synapse model.
47
-
48
- The ``half-part`` means that the model only includes ``comm`` -> ``syn`` -> ``out`` -> ``post``.
49
- Therefore, the model's ``update`` function needs the manual providing of the spiking input.
50
-
51
- **Model Descriptions**
52
-
53
- .. math::
54
-
55
- I_{syn} (t) = \sum_{j\in C} g_{\mathrm{max}} * \delta(t-t_j-D)
56
-
57
- where :math:`g_{\mathrm{max}}` denotes the chemical synaptic strength,
58
- :math:`t_j` the spiking moment of the presynaptic neuron :math:`j`,
59
- :math:`C` the set of neurons connected to the post-synaptic neuron,
60
- and :math:`D` the transmission delay of chemical synapses.
61
- For simplicity, the rise and decay phases of post-synaptic currents are
62
- omitted in this model.
63
-
64
-
65
- **Code Examples**
66
-
67
- .. code-block::
68
-
69
- import brainstate as bp
70
- import brainstate.math as bm
71
-
72
- class Net(bp.DynamicalSystem):
73
- def __init__(self):
74
- super().__init__()
75
-
76
- self.pre = bp.dyn.PoissonGroup(10, 100.)
77
- self.post = bp.dyn.LifRef(1)
78
- self.syn = bp.dyn.HalfProjDelta(bp.dnn.Linear(10, 1, bp.init.OneInit(2.)), self.post)
79
-
80
- def update(self):
81
- self.syn(self.pre())
82
- self.post()
83
- return self.post.V.value
84
-
85
- net = Net()
86
- indices = bm.arange(1000).to_numpy()
87
- vs = bm.for_loop(net.step_run, indices, progress_bar=True)
88
- bp.visualize.line_plot(indices, vs, show=True)
89
-
90
- Args:
91
- comm: DynamicalSystem. The synaptic communication.
92
- post: DynamicalSystem. The post-synaptic neuron group.
93
- name: str. The projection name.
94
- mode: Mode. The computing mode.
95
- """
96
-
97
- _invisible_nodes = ['post']
98
-
99
- def __init__(
100
- self,
101
- comm: Module,
102
- post: ReceiveInputProj,
103
- name: Optional[str] = None,
104
- mode: Optional[Mode] = None,
105
- ):
106
- super().__init__(name=name, mode=mode)
107
-
108
- # synaptic models
109
- is_instance(comm, Module)
110
- is_instance(post, ReceiveInputProj)
111
- self.comm = comm
112
-
113
- # output initialization
114
- out = _Delta()
115
- post.add_input_fun(self.name, out, category='delta')
116
- self.out = out
117
-
118
- # references
119
- self.post = post
120
-
121
- def update(self, x):
122
- # call the communication
123
- current = self.comm(x)
124
- # bind the output
125
- self.out.bind_cond(current)
126
- # return the current, if needed
127
- return current
128
-
129
-
130
- @set_module_as('brainstate.nn')
131
- class FullProjDelta(Projection):
132
- """Full-chain of the synaptic projection for the Delta synapse model.
133
-
134
- The synaptic projection requires the input is the spiking data, otherwise
135
- the synapse is not the Delta synapse model.
136
-
137
- The ``full-chain`` means that the model needs to provide all information needed for a projection,
138
- including ``pre`` -> ``delay`` -> ``comm`` -> ``post``.
139
-
140
- **Model Descriptions**
141
-
142
- .. math::
143
-
144
- I_{syn} (t) = \sum_{j\in C} g_{\mathrm{max}} * \delta(t-t_j-D)
145
-
146
- where :math:`g_{\mathrm{max}}` denotes the chemical synaptic strength,
147
- :math:`t_j` the spiking moment of the presynaptic neuron :math:`j`,
148
- :math:`C` the set of neurons connected to the post-synaptic neuron,
149
- and :math:`D` the transmission delay of chemical synapses.
150
- For simplicity, the rise and decay phases of post-synaptic currents are
151
- omitted in this model.
152
-
153
-
154
- **Code Examples**
155
-
156
- .. code-block::
157
-
158
- import brainstate as bp
159
- import brainstate.math as bm
160
-
161
-
162
- class Net(bp.DynamicalSystem):
163
- def __init__(self):
164
- super().__init__()
165
-
166
- self.pre = bp.dyn.PoissonGroup(10, 100.)
167
- self.post = bp.dyn.LifRef(1)
168
- self.syn = bp.dyn.FullProjDelta(self.pre, 0., bp.dnn.Linear(10, 1, bp.init.OneInit(2.)), self.post)
169
-
170
- def update(self):
171
- self.syn()
172
- self.pre()
173
- self.post()
174
- return self.post.V.value
175
-
176
-
177
- net = Net()
178
- indices = bm.arange(1000).to_numpy()
179
- vs = bm.for_loop(net.step_run, indices, progress_bar=True)
180
- bp.visualize.line_plot(indices, vs, show=True)
181
-
182
-
183
- Args:
184
- pre: The pre-synaptic neuron group.
185
- delay: The synaptic delay.
186
- comm: DynamicalSystem. The synaptic communication.
187
- post: DynamicalSystem. The post-synaptic neuron group.
188
- name: str. The projection name.
189
- mode: Mode. The computing mode.
190
- """
191
-
192
- _invisible_nodes = ['pre', 'post', 'delay']
193
-
194
- def __init__(
195
- self,
196
- pre: UpdateReturn,
197
- delay: Union[None, int, float],
198
- comm: Module,
199
- post: Dynamics,
200
- name: Optional[str] = None,
201
- mode: Optional[Mode] = None,
202
- ):
203
- super().__init__(name=name, mode=mode)
204
-
205
- # synaptic models
206
- is_instance(pre, UpdateReturn)
207
- is_instance(comm, Module)
208
- is_instance(post, Dynamics)
209
- self.comm = comm
210
-
211
- # delay initialization
212
- if delay is not None and delay > 0.:
213
- delay_cls = register_delay_of_target(pre)
214
- delay_cls.register_entry(self.name, delay)
215
- self.delay = delay_cls
216
- self.has_delay = True
217
- else:
218
- self.delay = None
219
- self.has_delay = False
220
-
221
- # output initialization
222
- out = _Delta()
223
- post.add_input_fun(self.name, out, category='delta')
224
- self.out = out
225
-
226
- # references
227
- self.pre = pre
228
- self.post = post
229
-
230
- def update(self):
231
- # get delay
232
- if self.has_delay:
233
- x = self.delay.at(self.name)
234
- else:
235
- x = self.pre.update_return()
236
- # call the communication
237
- current = self.comm(x)
238
- # bind the output
239
- self.out.bind_cond(current)
240
- # return the current, if needed
241
- return current
@@ -1,101 +0,0 @@
1
- # Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # http://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- # ==============================================================================
15
-
16
- from typing import Optional
17
-
18
- from brainstate._module import (Module, Projection, Dynamics, ReceiveInputProj)
19
- from brainstate._utils import set_module_as
20
- from brainstate.mixin import (BindCondData, Mode)
21
- from ._utils import is_instance
22
-
23
- __all__ = [
24
- 'VanillaProj',
25
- ]
26
-
27
-
28
- @set_module_as('brainstate.nn')
29
- class VanillaProj(Projection):
30
- """Synaptic projection which defines the synaptic computation with the dimension of pre-synaptic neuron group.
31
-
32
- **Code Examples**
33
-
34
- To simulate an E/I balanced network model:
35
-
36
- .. code-block::
37
-
38
- class EINet(bp.DynSysGroup):
39
- def __init__(self):
40
- super().__init__()
41
- self.N = bp.dyn.LifRef(4000, V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.,
42
- V_initializer=bp.init.Normal(-55., 2.))
43
- self.delay = bp.VarDelay(self.N.spike, entries={'I': None})
44
- self.syn1 = bp.dyn.Expon(size=3200, tau=5.)
45
- self.syn2 = bp.dyn.Expon(size=800, tau=10.)
46
- self.E = bp.dyn.VanillaProj(comm=bp.dnn.JitFPHomoLinear(3200, 4000, prob=0.02, weight=0.6),
47
- out=bp.dyn.COBA(E=0.),
48
- post=self.N)
49
- self.I = bp.dyn.VanillaProj(comm=bp.dnn.JitFPHomoLinear(800, 4000, prob=0.02, weight=6.7),
50
- out=bp.dyn.COBA(E=-80.),
51
- post=self.N)
52
-
53
- def update(self, input):
54
- spk = self.delay.at('I')
55
- self.E(self.syn1(spk[:3200]))
56
- self.I(self.syn2(spk[3200:]))
57
- self.delay(self.N(input))
58
- return self.N.spike.value
59
-
60
- model = EINet()
61
- indices = bm.arange(1000)
62
- spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices)
63
- bp.visualize.raster_plot(indices, spks, show=True)
64
-
65
-
66
- Args:
67
- comm: The synaptic communication.
68
- out: The synaptic output.
69
- post: The post-synaptic neuron group.
70
- name: str. The projection name.
71
- mode: Mode. The computing mode.
72
- """
73
-
74
- _invisible_nodes = ['post', ]
75
-
76
- def __init__(
77
- self,
78
- comm: Module,
79
- out: BindCondData,
80
- post: ReceiveInputProj,
81
- name: Optional[str] = None,
82
- mode: Optional[Mode] = None,
83
- ):
84
- super().__init__(name=name, mode=mode)
85
-
86
- # synaptic models
87
- is_instance(comm, Module)
88
- is_instance(out, BindCondData)
89
- is_instance(post, Dynamics)
90
- self.comm = comm
91
-
92
- # output initialization
93
- post.add_input_fun(self.name, out)
94
- self.out = out
95
-
96
- self.post = post
97
-
98
- def update(self, x):
99
- current = self.comm(x)
100
- self.out.bind_cond(current)
101
- return current