brainstate 0.1.9__py2.py3-none-any.whl → 0.2.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 (163) hide show
  1. brainstate/__init__.py +130 -19
  2. brainstate/_compatible_import.py +201 -9
  3. brainstate/_compatible_import_test.py +681 -0
  4. brainstate/_deprecation.py +210 -0
  5. brainstate/_deprecation_test.py +2319 -0
  6. brainstate/{util/error.py → _error.py} +10 -20
  7. brainstate/_state.py +94 -47
  8. brainstate/_state_test.py +1 -1
  9. brainstate/_utils.py +1 -1
  10. brainstate/environ.py +1279 -347
  11. brainstate/environ_test.py +1187 -26
  12. brainstate/graph/__init__.py +6 -13
  13. brainstate/graph/_node.py +240 -0
  14. brainstate/graph/_node_test.py +589 -0
  15. brainstate/graph/{_graph_operation.py → _operation.py} +632 -746
  16. brainstate/graph/_operation_test.py +1147 -0
  17. brainstate/mixin.py +1209 -141
  18. brainstate/mixin_test.py +991 -51
  19. brainstate/nn/__init__.py +74 -72
  20. brainstate/nn/_activations.py +587 -295
  21. brainstate/nn/_activations_test.py +109 -86
  22. brainstate/nn/_collective_ops.py +393 -274
  23. brainstate/nn/_collective_ops_test.py +746 -15
  24. brainstate/nn/_common.py +114 -66
  25. brainstate/nn/_common_test.py +154 -0
  26. brainstate/nn/_conv.py +1652 -143
  27. brainstate/nn/_conv_test.py +838 -227
  28. brainstate/nn/_delay.py +95 -29
  29. brainstate/nn/_delay_test.py +25 -20
  30. brainstate/nn/_dropout.py +359 -167
  31. brainstate/nn/_dropout_test.py +429 -52
  32. brainstate/nn/_dynamics.py +14 -90
  33. brainstate/nn/_dynamics_test.py +1 -12
  34. brainstate/nn/_elementwise.py +492 -313
  35. brainstate/nn/_elementwise_test.py +806 -145
  36. brainstate/nn/_embedding.py +369 -19
  37. brainstate/nn/_embedding_test.py +156 -0
  38. brainstate/nn/{_fixedprob.py → _event_fixedprob.py} +10 -16
  39. brainstate/nn/{_fixedprob_test.py → _event_fixedprob_test.py} +6 -5
  40. brainstate/nn/{_linear_mv.py → _event_linear.py} +2 -2
  41. brainstate/nn/{_linear_mv_test.py → _event_linear_test.py} +6 -5
  42. brainstate/nn/_exp_euler.py +200 -38
  43. brainstate/nn/_exp_euler_test.py +350 -8
  44. brainstate/nn/_linear.py +391 -71
  45. brainstate/nn/_linear_test.py +427 -59
  46. brainstate/nn/_metrics.py +1070 -0
  47. brainstate/nn/_metrics_test.py +611 -0
  48. brainstate/nn/_module.py +10 -3
  49. brainstate/nn/_module_test.py +1 -1
  50. brainstate/nn/_normalizations.py +688 -329
  51. brainstate/nn/_normalizations_test.py +663 -37
  52. brainstate/nn/_paddings.py +1020 -0
  53. brainstate/nn/_paddings_test.py +723 -0
  54. brainstate/nn/_poolings.py +1404 -342
  55. brainstate/nn/_poolings_test.py +828 -92
  56. brainstate/nn/{_rate_rnns.py → _rnns.py} +446 -54
  57. brainstate/nn/_rnns_test.py +593 -0
  58. brainstate/nn/_utils.py +132 -5
  59. brainstate/nn/_utils_test.py +402 -0
  60. brainstate/{init/_random_inits.py → nn/init.py} +301 -45
  61. brainstate/{init/_random_inits_test.py → nn/init_test.py} +51 -20
  62. brainstate/random/__init__.py +247 -1
  63. brainstate/random/_rand_funs.py +668 -346
  64. brainstate/random/_rand_funs_test.py +74 -1
  65. brainstate/random/_rand_seed.py +541 -76
  66. brainstate/random/_rand_seed_test.py +1 -1
  67. brainstate/random/_rand_state.py +601 -393
  68. brainstate/random/_rand_state_test.py +551 -0
  69. brainstate/transform/__init__.py +59 -0
  70. brainstate/transform/_ad_checkpoint.py +176 -0
  71. brainstate/{compile → transform}/_ad_checkpoint_test.py +1 -1
  72. brainstate/{augment → transform}/_autograd.py +360 -113
  73. brainstate/{augment → transform}/_autograd_test.py +2 -2
  74. brainstate/transform/_conditions.py +316 -0
  75. brainstate/{compile → transform}/_conditions_test.py +11 -11
  76. brainstate/{compile → transform}/_error_if.py +22 -20
  77. brainstate/{compile → transform}/_error_if_test.py +1 -1
  78. brainstate/transform/_eval_shape.py +145 -0
  79. brainstate/{augment → transform}/_eval_shape_test.py +1 -1
  80. brainstate/{compile → transform}/_jit.py +99 -46
  81. brainstate/{compile → transform}/_jit_test.py +3 -3
  82. brainstate/{compile → transform}/_loop_collect_return.py +219 -80
  83. brainstate/{compile → transform}/_loop_collect_return_test.py +1 -1
  84. brainstate/{compile → transform}/_loop_no_collection.py +133 -34
  85. brainstate/{compile → transform}/_loop_no_collection_test.py +2 -2
  86. brainstate/transform/_make_jaxpr.py +2016 -0
  87. brainstate/transform/_make_jaxpr_test.py +1510 -0
  88. brainstate/transform/_mapping.py +529 -0
  89. brainstate/transform/_mapping_test.py +194 -0
  90. brainstate/{compile → transform}/_progress_bar.py +78 -25
  91. brainstate/{augment → transform}/_random.py +65 -45
  92. brainstate/{compile → transform}/_unvmap.py +102 -5
  93. brainstate/transform/_util.py +286 -0
  94. brainstate/typing.py +594 -61
  95. brainstate/typing_test.py +780 -0
  96. brainstate/util/__init__.py +9 -32
  97. brainstate/util/_others.py +1025 -0
  98. brainstate/util/_others_test.py +962 -0
  99. brainstate/util/_pretty_pytree.py +1301 -0
  100. brainstate/util/_pretty_pytree_test.py +675 -0
  101. brainstate/util/{pretty_repr.py → _pretty_repr.py} +161 -27
  102. brainstate/util/_pretty_repr_test.py +696 -0
  103. brainstate/util/filter.py +557 -81
  104. brainstate/util/filter_test.py +912 -0
  105. brainstate/util/struct.py +769 -382
  106. brainstate/util/struct_test.py +602 -0
  107. {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/METADATA +34 -17
  108. brainstate-0.2.0.dist-info/RECORD +111 -0
  109. brainstate/augment/__init__.py +0 -30
  110. brainstate/augment/_eval_shape.py +0 -99
  111. brainstate/augment/_mapping.py +0 -1060
  112. brainstate/augment/_mapping_test.py +0 -597
  113. brainstate/compile/__init__.py +0 -38
  114. brainstate/compile/_ad_checkpoint.py +0 -204
  115. brainstate/compile/_conditions.py +0 -256
  116. brainstate/compile/_make_jaxpr.py +0 -888
  117. brainstate/compile/_make_jaxpr_test.py +0 -156
  118. brainstate/compile/_util.py +0 -147
  119. brainstate/functional/__init__.py +0 -27
  120. brainstate/graph/_graph_node.py +0 -244
  121. brainstate/graph/_graph_node_test.py +0 -73
  122. brainstate/graph/_graph_operation_test.py +0 -563
  123. brainstate/init/__init__.py +0 -26
  124. brainstate/init/_base.py +0 -52
  125. brainstate/init/_generic.py +0 -244
  126. brainstate/init/_regular_inits.py +0 -105
  127. brainstate/init/_regular_inits_test.py +0 -50
  128. brainstate/nn/_inputs.py +0 -608
  129. brainstate/nn/_ltp.py +0 -28
  130. brainstate/nn/_neuron.py +0 -705
  131. brainstate/nn/_neuron_test.py +0 -161
  132. brainstate/nn/_others.py +0 -46
  133. brainstate/nn/_projection.py +0 -486
  134. brainstate/nn/_rate_rnns_test.py +0 -63
  135. brainstate/nn/_readout.py +0 -209
  136. brainstate/nn/_readout_test.py +0 -53
  137. brainstate/nn/_stp.py +0 -236
  138. brainstate/nn/_synapse.py +0 -505
  139. brainstate/nn/_synapse_test.py +0 -131
  140. brainstate/nn/_synaptic_projection.py +0 -423
  141. brainstate/nn/_synouts.py +0 -162
  142. brainstate/nn/_synouts_test.py +0 -57
  143. brainstate/nn/metrics.py +0 -388
  144. brainstate/optim/__init__.py +0 -38
  145. brainstate/optim/_base.py +0 -64
  146. brainstate/optim/_lr_scheduler.py +0 -448
  147. brainstate/optim/_lr_scheduler_test.py +0 -50
  148. brainstate/optim/_optax_optimizer.py +0 -152
  149. brainstate/optim/_optax_optimizer_test.py +0 -53
  150. brainstate/optim/_sgd_optimizer.py +0 -1104
  151. brainstate/random/_random_for_unit.py +0 -52
  152. brainstate/surrogate.py +0 -1957
  153. brainstate/transform.py +0 -23
  154. brainstate/util/caller.py +0 -98
  155. brainstate/util/others.py +0 -540
  156. brainstate/util/pretty_pytree.py +0 -945
  157. brainstate/util/pretty_pytree_test.py +0 -159
  158. brainstate/util/pretty_table.py +0 -2954
  159. brainstate/util/scaling.py +0 -258
  160. brainstate-0.1.9.dist-info/RECORD +0 -130
  161. {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/WHEEL +0 -0
  162. {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/licenses/LICENSE +0 -0
  163. {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/top_level.txt +0 -0
@@ -1,4 +1,4 @@
1
- # Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
1
+ # Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -17,19 +17,361 @@
17
17
  import unittest
18
18
 
19
19
  import brainunit as u
20
+ import jax.numpy as jnp
21
+ import numpy as np
20
22
 
21
23
  import brainstate
22
24
 
23
25
 
24
- class TestExpEuler(unittest.TestCase):
25
- def test1(self):
26
- def fun(x, tau):
26
+ class TestExpEulerODE(unittest.TestCase):
27
+ """Test cases for ODE integration using exp_euler_step."""
28
+
29
+ def test_exponential_decay(self):
30
+ """Test simple exponential decay: dx/dt = -x"""
31
+ def drift(x, t):
32
+ return -x
33
+
34
+ with brainstate.environ.context(dt=0.01):
35
+ x0 = jnp.array(1.0)
36
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
37
+
38
+ # Expected: x(t+dt) ≈ x(t) * exp(-dt) ≈ 0.99004983
39
+ expected = np.exp(-0.01)
40
+ np.testing.assert_allclose(x1, expected)
41
+
42
+ def test_exponential_decay_with_time_constant(self):
43
+ """Test exponential decay with time constant: dx/dt = -x/tau"""
44
+ def drift(x, tau):
27
45
  return -x / tau
28
46
 
47
+ with brainstate.environ.context(dt=1.0 * u.ms):
48
+ x0 = u.math.asarray(1.0 * u.mV)
49
+ tau = 10.0 * u.ms
50
+ x1 = brainstate.nn.exp_euler_step(drift, x0, tau)
51
+
52
+ # Expected: x(t+dt) ≈ x(t) * exp(-dt/tau)
53
+ expected = np.exp(-0.1) * u.mV
54
+ assert u.math.allclose(x1, expected)
55
+
56
+ def test_linear_growth(self):
57
+ """Test linear growth: dx/dt = a"""
58
+ def drift(x):
59
+ return jnp.asarray([2.0])
60
+
61
+ with brainstate.environ.context(dt=0.1):
62
+ x0 = jnp.asarray([1.0])
63
+ x1 = brainstate.nn.exp_euler_step(drift, x0)
64
+
65
+ # For constant derivative, result should be x0 + a*dt
66
+ expected = 1.0 + 2.0 * 0.1
67
+ np.testing.assert_allclose(x1, expected)
68
+
69
+ def test_multidimensional_system(self):
70
+ """Test multi-dimensional ODE: harmonic oscillator"""
71
+ def drift(x, t):
72
+ # dx/dt = [x1, -x0] (circular motion)
73
+ return jnp.array([x[1], -x[0]])
74
+
75
+ with brainstate.environ.context(dt=0.01):
76
+ x0 = jnp.array([1.0, 0.0])
77
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
78
+
79
+ # Check that energy is approximately conserved
80
+ energy0 = np.sum(x0 ** 2)
81
+ energy1 = np.sum(x1 ** 2)
82
+ np.testing.assert_allclose(energy1, energy0, rtol=0.1)
83
+
84
+ def test_stiff_equation(self):
85
+ """Test stiff equation where exponential Euler should be stable"""
86
+ def drift(x, t):
87
+ # Stiff equation: dx/dt = -100*x
88
+ return -100.0 * x
89
+
29
90
  with brainstate.environ.context(dt=0.1):
91
+ x0 = jnp.array(1.0)
92
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
93
+
94
+ # Should remain stable and decay
95
+ expected = np.exp(-10.0)
96
+ np.testing.assert_allclose(x1, expected, rtol=1e-2)
97
+ self.assertGreater(x1, 0.0) # Should not become negative
98
+
99
+
100
+ class TestExpEulerSDE(unittest.TestCase):
101
+ """Test cases for SDE integration using exp_euler_step."""
102
+
103
+ def test_simple_sde_with_constant_diffusion(self):
104
+ """Test SDE with constant diffusion: dx = -x*dt + sigma*dW"""
105
+ def drift(x, t):
106
+ return -x
107
+
108
+ def diffusion(x, t):
109
+ return jnp.array(0.1)
110
+
111
+ with brainstate.environ.context(dt=0.01):
112
+ brainstate.random.seed(42)
113
+ x0 = jnp.array(1.0)
114
+ x1 = brainstate.nn.exp_euler_step(drift, diffusion, x0, None)
115
+
116
+ # Result should have both drift and diffusion components
117
+ # Cannot test exact value due to randomness, but check it's reasonable
118
+ self.assertIsInstance(x1, (jnp.ndarray, float))
119
+
120
+ def test_ornstein_uhlenbeck_process(self):
121
+ """Test Ornstein-Uhlenbeck process: dx = -theta*x*dt + sigma*dW"""
122
+ theta = 0.5
123
+ sigma = 0.3
124
+
125
+ def drift(x, t):
126
+ return -theta * x
127
+
128
+ def diffusion(x, t):
129
+ return jnp.full_like(x, sigma)
130
+
131
+ with brainstate.environ.context(dt=0.01):
132
+ brainstate.random.seed(123)
133
+ x0 = jnp.array(1.0)
134
+ x1 = brainstate.nn.exp_euler_step(drift, diffusion, x0, None)
135
+
136
+ # Mean should decrease (drift dominates initially)
137
+ # Run multiple steps and check statistics
138
+ x = x0
139
+ results = []
140
+ for _ in range(100):
141
+ x = brainstate.nn.exp_euler_step(drift, diffusion, x, None)
142
+ results.append(x)
143
+
144
+ # Mean should converge toward 0
145
+ final_mean = np.mean(results[-10:])
146
+ self.assertLess(abs(final_mean), 0.5)
147
+
148
+ def test_sde_multidimensional(self):
149
+ """Test multi-dimensional SDE"""
150
+ def drift(x, t):
151
+ return -0.5 * x
152
+
153
+ def diffusion(x, t):
154
+ return jnp.array([0.1, 0.2])
155
+
156
+ with brainstate.environ.context(dt=0.01):
157
+ brainstate.random.seed(456)
158
+ x0 = jnp.array([1.0, 1.0])
159
+ x1 = brainstate.nn.exp_euler_step(drift, diffusion, x0, None)
160
+
161
+ self.assertEqual(x1.shape, (2,))
162
+
163
+ def test_state_dependent_diffusion(self):
164
+ """Test SDE with state-dependent diffusion: dx = -x*dt + sqrt(x)*dW"""
165
+ def drift(x, t):
166
+ return -0.1 * x
167
+
168
+ def diffusion(x, t):
169
+ return jnp.sqrt(jnp.abs(x) + 1e-8)
170
+
171
+ with brainstate.environ.context(dt=0.01):
172
+ brainstate.random.seed(789)
173
+ x0 = jnp.array(1.0)
174
+ x1 = brainstate.nn.exp_euler_step(drift, diffusion, x0, None)
175
+
176
+ self.assertIsInstance(x1, (jnp.ndarray, float))
177
+
178
+
179
+ class TestExpEulerUnits(unittest.TestCase):
180
+ """Test cases for unit handling in exp_euler_step."""
181
+
182
+ def test_unit_compatibility_drift(self):
183
+ """Test that drift function units are validated correctly"""
184
+ def drift(x, tau):
185
+ return -x / tau
186
+
187
+ with brainstate.environ.context(dt=1.0 * u.ms):
188
+ x0 = 1.0 * u.mV
189
+ tau = 10.0 * u.ms
190
+ x1 = brainstate.nn.exp_euler_step(drift, x0, tau)
191
+
192
+ # Result should have same units as input
193
+ self.assertEqual(u.get_unit(x1), u.get_unit(x0))
194
+
195
+ def test_unit_mismatch_raises_error(self):
196
+ """Test that incompatible diffusion units raise an error"""
197
+ def drift(x, t):
198
+ return -x / (10.0 * u.ms)
199
+
200
+ def diffusion(x, t):
201
+ # Wrong units: should be mV/sqrt(ms) but returning mV
202
+ return 0.1 * u.mV
203
+
204
+ with brainstate.environ.context(dt=1.0 * u.ms):
205
+ x0 = 1.0 * u.mV
206
+ with self.assertRaises(ValueError):
207
+ brainstate.nn.exp_euler_step(drift, diffusion, x0, None)
208
+
209
+ def test_correct_diffusion_units(self):
210
+ """Test SDE with correct diffusion units"""
211
+ def drift(x, tau):
212
+ return -x / tau
213
+
214
+ def diffusion(x, t):
215
+ # Correct units: mV/sqrt(ms)
216
+ return 0.1 * u.mV / u.ms ** 0.5
217
+
218
+ with brainstate.environ.context(dt=1.0 * u.ms):
219
+ brainstate.random.seed(42)
220
+ x0 = 1.0 * u.mV
221
+ tau = 10.0 * u.ms
222
+ x1 = brainstate.nn.exp_euler_step(drift, diffusion, x0, tau)
223
+
224
+ self.assertEqual(u.get_unit(x1), u.get_unit(x0))
225
+
226
+ def test_dimensionless_with_time_units(self):
227
+ """Test dimensionless state with time units in dt"""
228
+ def drift(x, t):
229
+ return -2.0 * x / u.second
230
+
231
+ with brainstate.environ.context(dt=0.1 * u.second):
232
+ x0 = jnp.array(1.0)
233
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
234
+
235
+ expected = np.exp(-0.2)
236
+ np.testing.assert_allclose(x1, expected, rtol=1e-5)
237
+
238
+
239
+ class TestExpEulerInputValidation(unittest.TestCase):
240
+ """Test cases for input validation in exp_euler_step."""
241
+
242
+ def test_non_callable_drift_raises_error(self):
243
+ """Test that non-callable drift raises AssertionError"""
244
+ with brainstate.environ.context(dt=0.01):
245
+ x0 = jnp.array(1.0)
246
+ with self.assertRaises(AssertionError):
247
+ brainstate.nn.exp_euler_step("not a function", x0, None)
248
+
249
+ def test_no_state_variable_raises_error(self):
250
+ """Test that missing state variable raises AssertionError"""
251
+ def drift(x, t):
252
+ return -x
253
+
254
+ with brainstate.environ.context(dt=0.01):
255
+ with self.assertRaises(AssertionError):
256
+ brainstate.nn.exp_euler_step(drift)
257
+
258
+ def test_invalid_dtype_raises_error(self):
259
+ """Test that invalid dtype raises ValueError"""
260
+ def drift(x, t):
261
+ return -x
262
+
263
+ with brainstate.environ.context(dt=0.01):
264
+ x0 = jnp.array(1, dtype=jnp.int32)
265
+ with self.assertRaises(ValueError):
266
+ brainstate.nn.exp_euler_step(drift, x0, None)
267
+
268
+ def test_float16_dtype_accepted(self):
269
+ """Test that float16 dtype is accepted"""
270
+ def drift(x, t):
271
+ return -x
272
+
273
+ with brainstate.environ.context(dt=0.01):
274
+ x0 = jnp.array(1.0, dtype=jnp.float16)
275
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
276
+ self.assertEqual(x1.dtype, jnp.float16)
277
+
278
+ def test_bfloat16_dtype_accepted(self):
279
+ """Test that bfloat16 dtype is accepted"""
280
+ def drift(x, t):
281
+ return -x
282
+
283
+ with brainstate.environ.context(dt=0.01):
284
+ x0 = jnp.array(1.0, dtype=jnp.bfloat16)
285
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
286
+ self.assertEqual(x1.dtype, jnp.bfloat16)
287
+
288
+ def test_diffusion_without_state_raises_error(self):
289
+ """Test that diffusion function without state variable raises error"""
290
+ def drift(x, t):
291
+ return -x
292
+
293
+ def diffusion(x, t):
294
+ return 0.1
295
+
296
+ with brainstate.environ.context(dt=0.01):
30
297
  with self.assertRaises(AssertionError):
31
- r = brainstate.nn.exp_euler_step(fun, 1.0 * u.mV, 1. * u.ms)
298
+ brainstate.nn.exp_euler_step(drift, diffusion)
299
+
300
+
301
+ class TestExpEulerEdgeCases(unittest.TestCase):
302
+ """Test edge cases and special scenarios."""
303
+
304
+ def test_zero_initial_condition(self):
305
+ """Test with zero initial condition"""
306
+ def drift(x, t):
307
+ return -x + 1.0
308
+
309
+ with brainstate.environ.context(dt=0.01):
310
+ x0 = jnp.array(0.0)
311
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
312
+
313
+ # Should move toward equilibrium at x=1
314
+ self.assertGreater(x1, 0.0)
315
+
316
+ def test_very_small_timestep(self):
317
+ """Test with very small timestep"""
318
+ def drift(x, t):
319
+ return -x
320
+
321
+ with brainstate.environ.context(dt=1e-8):
322
+ x0 = jnp.array(1.0)
323
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
324
+
325
+ # Should barely change
326
+ np.testing.assert_allclose(x1, x0)
327
+
328
+ def test_large_timestep_stability(self):
329
+ """Test stability with large timestep (advantage of exponential Euler)"""
330
+ def drift(x, t):
331
+ return -10.0 * x
332
+
333
+ with brainstate.environ.context(dt=1.0):
334
+ x0 = jnp.array(1.0)
335
+ x1 = brainstate.nn.exp_euler_step(drift, x0, None)
336
+
337
+ # Should remain stable (not blow up or oscillate)
338
+ expected = np.exp(-10.0)
339
+ np.testing.assert_allclose(x1, expected, rtol=1e-1)
340
+ self.assertGreater(x1, 0.0)
341
+
342
+ def test_kwargs_passed_correctly(self):
343
+ """Test that kwargs are passed to drift and diffusion functions"""
344
+ def drift(x, scale=1., **kwargs):
345
+ return -scale * x
346
+
347
+ def diffusion(x, noise_level=0.1, **kwargs):
348
+ return noise_level
349
+
350
+ with brainstate.environ.context(dt=0.01):
351
+ brainstate.random.seed(42)
352
+ x0 = jnp.array(1.0)
353
+ x1 = brainstate.nn.exp_euler_step(
354
+ drift, diffusion, x0,
355
+ scale=2.0, noise_level=0.2
356
+ )
357
+
358
+ self.assertIsInstance(x1, (jnp.ndarray, float))
359
+
360
+ def test_reproducibility_with_seed(self):
361
+ """Test that results are reproducible with same random seed"""
362
+ def drift(x):
363
+ return -0.5 * x
364
+
365
+ def diffusion(x):
366
+ return 0.1
367
+
368
+ with brainstate.environ.context(dt=0.01):
369
+ x0 = jnp.array(1.0)
370
+
371
+ brainstate.random.seed(42)
372
+ x1_first = brainstate.nn.exp_euler_step(drift, diffusion, x0)
373
+
374
+ brainstate.random.seed(42)
375
+ x1_second = brainstate.nn.exp_euler_step(drift, diffusion, x0)
32
376
 
33
- with brainstate.environ.context(dt=1. * u.ms):
34
- r = brainstate.nn.exp_euler_step(fun, 1.0 * u.mV, 1. * u.ms)
35
- print(r)
377
+ np.testing.assert_array_equal(x1_first, x1_second)