brainstate 0.1.9__py2.py3-none-any.whl → 0.2.0__py2.py3-none-any.whl

brainstate/{init/_random_inits_test.py → nn/init_test.py}

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX 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.
@@ -23,20 +23,20 @@ import brainstate
 class TestNormalInit(unittest.TestCase):
 
     def test_normal_init1(self):
-        init = brainstate.init.Normal()
+        init = brainstate.nn.init.Normal()
         for size in [(100,), (10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_normal_init2(self):
-        init = brainstate.init.Normal(scale=0.5)
+        init = brainstate.nn.init.Normal(scale=0.5)
         for size in [(100,), (10, 20)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_normal_init3(self):
-        init1 = brainstate.init.Normal(scale=0.5, seed=10)
-        init2 = brainstate.init.Normal(scale=0.5, seed=10)
+        init1 = brainstate.nn.init.Normal(scale=0.5, seed=10)
+        init2 = brainstate.nn.init.Normal(scale=0.5, seed=10)
         size = (10,)
         weights1 = init1(size)
         weights2 = init2(size)
@@ -46,13 +46,13 @@ class TestNormalInit(unittest.TestCase):
 
 class TestUniformInit(unittest.TestCase):
     def test_uniform_init1(self):
-        init = brainstate.init.Normal()
+        init = brainstate.nn.init.Normal()
         for size in [(100,), (10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_uniform_init2(self):
-        init = brainstate.init.Uniform(min_val=10, max_val=20)
+        init = brainstate.nn.init.Uniform(min_val=10, max_val=20)
         for size in [(100,), (10, 20)]:
             weights = init(size)
             assert weights.shape == size
@@ -60,20 +60,21 @@ class TestUniformInit(unittest.TestCase):
 
 class TestVarianceScaling(unittest.TestCase):
     def test_var_scaling1(self):
-        init = brainstate.init.VarianceScaling(scale=1., mode='fan_in', distribution='truncated_normal')
+        init = brainstate.nn.init.VarianceScaling(scale=1., mode='fan_in', distribution='truncated_normal')
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_var_scaling2(self):
-        init = brainstate.init.VarianceScaling(scale=2, mode='fan_out', distribution='normal')
+        init = brainstate.nn.init.VarianceScaling(scale=2, mode='fan_out', distribution='normal')
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_var_scaling3(self):
-        init = brainstate.init.VarianceScaling(scale=2 / 4, mode='fan_avg', in_axis=0, out_axis=1,
-                                               distribution='uniform')
+        init = brainstate.nn.init.VarianceScaling(
+            scale=2 / 4, mode='fan_avg', in_axis=0, out_axis=1, distribution='uniform'
+        )
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -81,7 +82,7 @@ class TestVarianceScaling(unittest.TestCase):
 
 class TestKaimingUniformUnit(unittest.TestCase):
     def test_kaiming_uniform_init(self):
-        init = brainstate.init.KaimingUniform()
+        init = brainstate.nn.init.KaimingUniform()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -89,7 +90,7 @@ class TestKaimingUniformUnit(unittest.TestCase):
 
 class TestKaimingNormalUnit(unittest.TestCase):
     def test_kaiming_normal_init(self):
-        init = brainstate.init.KaimingNormal()
+        init = brainstate.nn.init.KaimingNormal()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -97,7 +98,7 @@ class TestKaimingNormalUnit(unittest.TestCase):
 
 class TestXavierUniformUnit(unittest.TestCase):
     def test_xavier_uniform_init(self):
-        init = brainstate.init.XavierUniform()
+        init = brainstate.nn.init.XavierUniform()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -105,7 +106,7 @@ class TestXavierUniformUnit(unittest.TestCase):
 
 class TestXavierNormalUnit(unittest.TestCase):
     def test_xavier_normal_init(self):
-        init = brainstate.init.XavierNormal()
+        init = brainstate.nn.init.XavierNormal()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -113,7 +114,7 @@ class TestXavierNormalUnit(unittest.TestCase):
 
 class TestLecunUniformUnit(unittest.TestCase):
     def test_lecun_uniform_init(self):
-        init = brainstate.init.LecunUniform()
+        init = brainstate.nn.init.LecunUniform()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -121,7 +122,7 @@ class TestLecunUniformUnit(unittest.TestCase):
 
 class TestLecunNormalUnit(unittest.TestCase):
     def test_lecun_normal_init(self):
-        init = brainstate.init.LecunNormal()
+        init = brainstate.nn.init.LecunNormal()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -129,13 +130,13 @@ class TestLecunNormalUnit(unittest.TestCase):
 
 class TestOrthogonalUnit(unittest.TestCase):
     def test_orthogonal_init1(self):
-        init = brainstate.init.Orthogonal()
+        init = brainstate.nn.init.Orthogonal()
         for size in [(20, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_orthogonal_init2(self):
-        init = brainstate.init.Orthogonal(scale=2., axis=0)
+        init = brainstate.nn.init.Orthogonal(scale=2., axis=0)
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -143,7 +144,37 @@ class TestOrthogonalUnit(unittest.TestCase):
 
 class TestDeltaOrthogonalUnit(unittest.TestCase):
     def test_delta_orthogonal_init1(self):
-        init = brainstate.init.DeltaOrthogonal()
+        init = brainstate.nn.init.DeltaOrthogonal()
         for size in [(20, 20, 20), (10, 20, 30, 40), (50, 40, 30, 20, 20)]:
             weights = init(size)
             assert weights.shape == size
+
+
+class TestZeroInit(unittest.TestCase):
+    def test_zero_init(self):
+        init = brainstate.nn.init.ZeroInit()
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            weights = init(size)
+            assert weights.shape == size
+
+
+class TestOneInit(unittest.TestCase):
+    def test_one_init(self):
+        for size in [(100,), (10, 20), (10, 20, 30)]:
+            for value in [0., 1., -1.]:
+                init = brainstate.nn.init.Constant(value=value)
+                weights = init(size)
+                assert weights.shape == size
+                assert (weights == value).all()
+
+
+class TestIdentityInit(unittest.TestCase):
+    def test_identity_init(self):
+        for size in [(100,), (10, 20)]:
+            for value in [0., 1., -1.]:
+                init = brainstate.nn.init.Identity(value=value)
+                weights = init(size)
+                if len(size) == 1:
+                    assert weights.shape == (size[0], size[0])
+                else:
+                    assert weights.shape == size

brainstate/random/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX 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.
@@ -13,6 +13,252 @@
 # limitations under the License.
 # ==============================================================================
 
+"""
+Random number generation module for BrainState.
+
+This module provides a comprehensive set of random number generation functions and utilities
+for neural network simulations and scientific computing. It wraps JAX's random number
+generation capabilities with a stateful interface that simplifies usage while maintaining
+reproducibility and performance.
+
+The module includes:
+
+- Standard random distributions (uniform, normal, exponential, etc.)
+- Random state management with automatic key splitting
+- Seed management utilities for reproducible simulations
+- NumPy-compatible API for easy migration
+
+Key Features
+------------
+
+- **Stateful random generation**: Automatic management of JAX's PRNG keys
+- **NumPy compatibility**: Drop-in replacement for most NumPy random functions
+- **Reproducibility**: Robust seed management and state tracking
+- **Performance**: JIT-compiled random functions for efficient generation
+- **Thread-safe**: Proper handling of random state in parallel computations
+
+Random State Management
+-----------------------
+
+The module uses a global `DEFAULT` RandomState instance that automatically manages
+JAX's PRNG keys. This eliminates the need to manually track and split keys:
+
+.. code-block:: python
+
+    >>> import brainstate as bs
+    >>> import brainstate.random as bsr
+    >>>
+    >>> # Set a global seed for reproducibility
+    >>> bsr.seed(42)
+    >>>
+    >>> # Generate random numbers without manual key management
+    >>> x = bsr.normal(0, 1, size=(3, 3))
+    >>> y = bsr.uniform(0, 1, size=(100,))
+
+Custom Random States
+--------------------
+
+For more control, you can create custom RandomState instances:
+
+.. code-block:: python
+
+    >>> import brainstate.random as bsr
+    >>>
+    >>> # Create a custom random state
+    >>> rng = bsr.RandomState(seed=123)
+    >>>
+    >>> # Use it for generation
+    >>> data = rng.normal(0, 1, size=(10, 10))
+    >>>
+    >>> # Get the current key
+    >>> current_key = rng.value
+
+Available Distributions
+-----------------------
+
+The module provides a wide range of probability distributions:
+
+**Uniform Distributions:**
+
+- `rand`, `random`, `random_sample`, `ranf`, `sample` - Uniform [0, 1)
+- `randint`, `random_integers` - Uniform integers
+- `choice` - Random selection from array
+- `permutation`, `shuffle` - Random ordering
+
+**Normal Distributions:**
+
+- `randn`, `normal` - Normal (Gaussian) distribution
+- `standard_normal` - Standard normal distribution
+- `multivariate_normal` - Multivariate normal distribution
+- `truncated_normal` - Truncated normal distribution
+
+**Other Continuous Distributions:**
+
+- `beta` - Beta distribution
+- `exponential`, `standard_exponential` - Exponential distribution
+- `gamma`, `standard_gamma` - Gamma distribution
+- `gumbel` - Gumbel distribution
+- `laplace` - Laplace distribution
+- `logistic` - Logistic distribution
+- `pareto` - Pareto distribution
+- `rayleigh` - Rayleigh distribution
+- `standard_cauchy` - Cauchy distribution
+- `standard_t` - Student's t-distribution
+- `uniform` - Uniform distribution over [low, high)
+- `weibull` - Weibull distribution
+
+**Discrete Distributions:**
+
+- `bernoulli` - Bernoulli distribution
+- `binomial` - Binomial distribution
+- `poisson` - Poisson distribution
+
+Seed Management
+---------------
+
+The module provides utilities for managing random seeds:
+
+.. code-block:: python
+
+    >>> import brainstate.random as bsr
+    >>>
+    >>> # Set a global seed
+    >>> bsr.seed(42)
+    >>>
+    >>> # Get current seed/key
+    >>> key = bsr.get_key()
+    >>>
+    >>> # Split the key for parallel operations
+    >>> keys = bsr.split_key(n=4)
+    >>>
+    >>> # Use context manager for temporary seed
+    >>> with bsr.local_seed(123):
+    ...     x = bsr.normal(0, 1, (5,))  # Uses seed 123
+    >>> y = bsr.normal(0, 1, (5,))  # Uses original seed
+
+Examples
+--------
+
+**Basic random number generation:**
+
+.. code-block:: python
+
+    >>> import brainstate.random as bsr
+    >>> import jax.numpy as jnp
+    >>>
+    >>> # Set seed for reproducibility
+    >>> bsr.seed(0)
+    >>>
+    >>> # Generate uniform random numbers
+    >>> uniform_data = bsr.random((3, 3))
+    >>> print(uniform_data.shape)
+    (3, 3)
+    >>>
+    >>> # Generate normal random numbers
+    >>> normal_data = bsr.normal(loc=0, scale=1, size=(100,))
+    >>> print(f"Mean: {normal_data.mean():.3f}, Std: {normal_data.std():.3f}")
+    Mean: -0.045, Std: 0.972
+
+**Sampling and shuffling:**
+
+.. code-block:: python
+
+    >>> import brainstate.random as bsr
+    >>> import jax.numpy as jnp
+    >>>
+    >>> bsr.seed(42)
+    >>>
+    >>> # Random choice from array
+    >>> arr = jnp.array([1, 2, 3, 4, 5])
+    >>> samples = bsr.choice(arr, size=3, replace=False)
+    >>> print(samples)
+    [4 1 5]
+    >>>
+    >>> # Random permutation
+    >>> perm = bsr.permutation(10)
+    >>> print(perm)
+    [3 5 1 7 9 0 2 8 4 6]
+    >>>
+    >>> # In-place shuffle
+    >>> data = jnp.arange(5)
+    >>> bsr.shuffle(data)
+    >>> print(data)
+    [2 0 4 1 3]
+
+**Advanced distributions:**
+
+.. code-block:: python
+
+    >>> import brainstate.random as bsr
+    >>> import matplotlib.pyplot as plt
+    >>>
+    >>> bsr.seed(123)
+    >>>
+    >>> # Generate samples from different distributions
+    >>> normal_samples = bsr.normal(0, 1, 1000)
+    >>> exponential_samples = bsr.exponential(1.0, 1000)
+    >>> beta_samples = bsr.beta(2, 5, 1000)
+    >>>
+    >>> # Plot histograms
+    >>> fig, axes = plt.subplots(1, 3, figsize=(12, 4))
+    >>> axes[0].hist(normal_samples, bins=30, density=True)
+    >>> axes[0].set_title('Normal Distribution')
+    >>> axes[1].hist(exponential_samples, bins=30, density=True)
+    >>> axes[1].set_title('Exponential Distribution')
+    >>> axes[2].hist(beta_samples, bins=30, density=True)
+    >>> axes[2].set_title('Beta Distribution')
+    >>> plt.show()
+
+**Using with neural network simulations:**
+
+.. code-block:: python
+
+    >>> import brainstate as bs
+    >>> import brainstate.random as bsr
+    >>> import brainstate.nn as nn
+    >>>
+    >>> class NoisyNeuron(bs.Module):
+    ...     def __init__(self, n_neurons, noise_scale=0.1):
+    ...         super().__init__()
+    ...         self.n_neurons = n_neurons
+    ...         self.noise_scale = noise_scale
+    ...         self.membrane = bs.State(jnp.zeros(n_neurons))
+    ...
+    ...     def update(self, input_current):
+    ...         # Add noise to input current
+    ...         noise = bsr.normal(0, self.noise_scale, self.n_neurons)
+    ...         self.membrane.value += input_current + noise
+    ...         return self.membrane.value
+    >>>
+    >>> # Create and run noisy neuron model
+    >>> bsr.seed(42)
+    >>> neuron = NoisyNeuron(100)
+    >>> output = neuron.update(jnp.ones(100) * 0.5)
+
+Notes
+-----
+
+- This module is designed to work seamlessly with JAX's functional programming model
+- Random functions are JIT-compilable for optimal performance
+- The global DEFAULT state is thread-local to avoid race conditions
+- For deterministic results, always set a seed before random operations
+
+See Also
+--------
+
+jax.random : JAX's random number generation module
+numpy.random : NumPy's random number generation module
+RandomState : The stateful random number generator class
+
+References
+----------
+.. [1] JAX Random Number Generation:
+   https://jax.readthedocs.io/en/latest/jax.random.html
+.. [2] NumPy Random Sampling:
+   https://numpy.org/doc/stable/reference/random/index.html
+
+"""
+
 from ._rand_funs import *
 from ._rand_funs import __all__ as __all_random__
 from ._rand_seed import *