|
456 | 456 | "source": [ |
457 | 457 | "post_size = 50 # Postsynaptic population size\n", |
458 | 458 | "\n", |
459 | | - "syn = brainpy.state.Expon.desc(post_size, tau=5*u.ms)" |
| 459 | + "syn = brainpy.state.Expon(post_size, tau=5*u.ms)" |
460 | 460 | ], |
461 | 461 | "outputs": [], |
462 | 462 | "execution_count": 9 |
|
479 | 479 | }, |
480 | 480 | "source": [ |
481 | 481 | "# Current-based output\n", |
482 | | - "out = brainpy.state.CUBA.desc() \n", |
| 482 | + "out = brainpy.state.CUBA() \n", |
483 | 483 | "\n", |
484 | 484 | "# Or conductance-based output\n", |
485 | | - "out = brainpy.state.COBA.desc(E=0*u.mV)" |
| 485 | + "out = brainpy.state.COBA(E=0*u.mV)" |
486 | 486 | ], |
487 | 487 | "outputs": [], |
488 | 488 | "execution_count": 10 |
|
889 | 889 | ], |
890 | 890 | "outputs": [], |
891 | 891 | "execution_count": 19 |
892 | | - }, |
893 | | - { |
894 | | - "cell_type": "markdown", |
895 | | - "metadata": {}, |
896 | | - "source": [ |
897 | | - "## Data Flow Example\n", |
898 | | - "\n", |
899 | | - "Here's how data flows through a typical ``brainpy.state`` simulation:" |
900 | | - ] |
901 | | - }, |
902 | | - { |
903 | | - "cell_type": "code", |
904 | | - "metadata": { |
905 | | - "ExecuteTime": { |
906 | | - "end_time": "2025-11-13T09:31:09.684116Z", |
907 | | - "start_time": "2025-11-13T09:31:09.550662Z" |
908 | | - } |
909 | | - }, |
910 | | - "source": [ |
911 | | - "# 1. Define network\n", |
912 | | - "class EINetwork(brainstate.nn.Module):\n", |
913 | | - " def __init__(self):\n", |
914 | | - " super().__init__()\n", |
915 | | - " self.E = brainpy.state.LIF(800, V_rest=-65*u.mV, V_th=-50*u.mV, tau=15*u.ms)\n", |
916 | | - " self.I = brainpy.state.LIF(200, V_rest=-65*u.mV, V_th=-50*u.mV, tau=10*u.ms)\n", |
917 | | - " \n", |
918 | | - " # Example projections (simplified - full setup requires more code)\n", |
919 | | - " # self.E2E = brainpy.state.AlignPostProj(...)\n", |
920 | | - " # self.E2I = brainpy.state.AlignPostProj(...)\n", |
921 | | - " # self.I2E = brainpy.state.AlignPostProj(...)\n", |
922 | | - " # self.I2I = brainpy.state.AlignPostProj(...)\n", |
923 | | - "\n", |
924 | | - " def update(self, input):\n", |
925 | | - " # Get spikes from last time step\n", |
926 | | - " e_spikes = self.E.get_spike()\n", |
927 | | - " i_spikes = self.I.get_spike()\n", |
928 | | - "\n", |
929 | | - " # Update projections (spikes → synaptic currents)\n", |
930 | | - " # self.E2E(e_spikes) # Updates E2E.syn.g\n", |
931 | | - " # self.E2I(e_spikes)\n", |
932 | | - " # self.I2E(i_spikes)\n", |
933 | | - " # self.I2I(i_spikes)\n", |
934 | | - "\n", |
935 | | - " # Update neurons (currents → new V and spikes)\n", |
936 | | - " self.E(input[:800] if len(input) >= 800 else input)\n", |
937 | | - " self.I(input[800:] if len(input) > 800 else jnp.zeros(200) * u.nA)\n", |
938 | | - "\n", |
939 | | - " return e_spikes, i_spikes\n", |
940 | | - "\n", |
941 | | - "# 2. Initialize\n", |
942 | | - "net = EINetwork()\n", |
943 | | - "brainstate.nn.init_all_states(net)\n", |
944 | | - "\n", |
945 | | - "# 3. Compile\n", |
946 | | - "@brainstate.transform.jit\n", |
947 | | - "def step(input):\n", |
948 | | - " return net.update(input)\n", |
949 | | - "\n", |
950 | | - "# 4. Simulate (commented out for quick execution)\n", |
951 | | - "# times = u.math.arange(0*u.ms, 1000*u.ms, 0.1*u.ms)\n", |
952 | | - "# results = brainstate.transform.for_loop(step, times)" |
953 | | - ], |
954 | | - "outputs": [], |
955 | | - "execution_count": 20 |
956 | | - }, |
957 | | - { |
958 | | - "cell_type": "markdown", |
959 | | - "metadata": {}, |
960 | | - "source": [ |
961 | | - "State Flow:\n", |
962 | | - "\n", |
963 | | - "```text\n", |
964 | | - "Time t:\n", |
965 | | - "┌──────────────────────────────────────────┐\n", |
966 | | - "│ States at t-1: │\n", |
967 | | - "│ E.V[t-1], E.spike[t-1] │\n", |
968 | | - "│ I.V[t-1], I.spike[t-1] │\n", |
969 | | - "│ E2E.syn.g[t-1], ... │\n", |
970 | | - "└──────────────────────────────────────────┘\n", |
971 | | - " ↓\n", |
972 | | - "┌──────────────────────────────────────────┐\n", |
973 | | - "│ Projection Updates: │\n", |
974 | | - "│ E2E.syn.g[t] = f(g[t-1], E.spike[t-1])│\n", |
975 | | - "│ ... (other projections) │\n", |
976 | | - "└──────────────────────────────────────────┘\n", |
977 | | - " ↓\n", |
978 | | - "┌──────────────────────────────────────────┐\n", |
979 | | - "│ Neuron Updates: │\n", |
980 | | - "│ E.V[t] = f(V[t-1], Σ currents[t]) │\n", |
981 | | - "│ E.spike[t] = E.V[t] >= V_th │\n", |
982 | | - "│ ... (other neurons) │\n", |
983 | | - "└──────────────────────────────────────────┘\n", |
984 | | - " ↓\n", |
985 | | - "Time t+1...\n", |
986 | | - "```" |
987 | | - ] |
988 | | - }, |
989 | | - { |
990 | | - "cell_type": "markdown", |
991 | | - "metadata": {}, |
992 | | - "source": [ |
993 | | - "## Performance Considerations\n", |
994 | | - "\n", |
995 | | - "### Memory Management\n", |
996 | | - "\n", |
997 | | - "- States are preallocated\n", |
998 | | - "- In-place updates when possible\n", |
999 | | - "- Efficient batching support\n", |
1000 | | - "- Automatic garbage collection\n", |
1001 | | - "\n", |
1002 | | - "### Compilation Strategy\n", |
1003 | | - "\n", |
1004 | | - "- Compile simulation loops\n", |
1005 | | - "- Batch operations when possible\n", |
1006 | | - "- Use ``for_loop`` for long sequences\n", |
1007 | | - "- Leverage JAX's XLA optimization\n", |
1008 | | - "\n", |
1009 | | - "### Hardware Acceleration\n", |
1010 | | - "\n", |
1011 | | - "- Automatic GPU dispatch for large arrays\n", |
1012 | | - "- TPU support for massive parallelism\n", |
1013 | | - "- Efficient CPU fallback for small problems" |
1014 | | - ] |
1015 | | - }, |
1016 | | - { |
1017 | | - "cell_type": "markdown", |
1018 | | - "metadata": {}, |
1019 | | - "source": [ |
1020 | | - "## Summary\n", |
1021 | | - "\n", |
1022 | | - "``brainpy.state`` 's architecture provides:\n", |
1023 | | - "\n", |
1024 | | - "✅ **Clear Abstractions**: Neurons, synapses, and projections with well-defined roles\n", |
1025 | | - "\n", |
1026 | | - "✅ **State Management**: Explicit, efficient handling of dynamical variables\n", |
1027 | | - "\n", |
1028 | | - "✅ **Modularity**: Compose complex models from simple components\n", |
1029 | | - "\n", |
1030 | | - "✅ **Performance**: JIT compilation and hardware acceleration\n", |
1031 | | - "\n", |
1032 | | - "✅ **Scientific Accuracy**: Integrated physical units\n", |
1033 | | - "\n", |
1034 | | - "✅ **Extensibility**: Easy to add custom components\n", |
1035 | | - "\n", |
1036 | | - "✅ **Modern Design**: Built on proven frameworks (JAX, brainstate)" |
1037 | | - ] |
1038 | | - }, |
1039 | | - { |
1040 | | - "cell_type": "markdown", |
1041 | | - "metadata": {}, |
1042 | | - "source": [ |
1043 | | - "## Next Steps\n", |
1044 | | - "\n", |
1045 | | - "- Learn about specific components: [neurons](neurons.ipynb), [synapses](synapses.ipynb), [projections](projections.ipynb)\n", |
1046 | | - "- Understand state management in depth: [state-management](state-management.ipynb)\n", |
1047 | | - "- See practical examples in the [tutorials](../tutorials/basic/01-lif-neuron.ipynb)\n", |
1048 | | - "- Explore the ecosystem: [brainstate docs](https://brainstate.readthedocs.io/)" |
1049 | | - ] |
1050 | 892 | } |
1051 | 893 | ], |
1052 | 894 | "metadata": { |
|
0 commit comments