neuralrnn 0.3.5__tar.gz

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 (130) hide show
  1. neuralrnn-0.3.5/.gitignore +233 -0
  2. neuralrnn-0.3.5/LICENSE +21 -0
  3. neuralrnn-0.3.5/PKG-INFO +156 -0
  4. neuralrnn-0.3.5/README.md +128 -0
  5. neuralrnn-0.3.5/pyproject.toml +53 -0
  6. neuralrnn-0.3.5/src/neuralrnn/__init__.py +79 -0
  7. neuralrnn-0.3.5/src/neuralrnn/activations.py +122 -0
  8. neuralrnn-0.3.5/src/neuralrnn/analysis/__init__.py +171 -0
  9. neuralrnn-0.3.5/src/neuralrnn/analysis/connectivity.py +75 -0
  10. neuralrnn-0.3.5/src/neuralrnn/analysis/demixed.py +245 -0
  11. neuralrnn-0.3.5/src/neuralrnn/analysis/dimensionality.py +81 -0
  12. neuralrnn-0.3.5/src/neuralrnn/analysis/fixed_points.py +614 -0
  13. neuralrnn-0.3.5/src/neuralrnn/analysis/linalg_utils.py +247 -0
  14. neuralrnn-0.3.5/src/neuralrnn/analysis/line_attractor.py +246 -0
  15. neuralrnn-0.3.5/src/neuralrnn/analysis/linearization.py +58 -0
  16. neuralrnn-0.3.5/src/neuralrnn/analysis/lyapunov.py +53 -0
  17. neuralrnn-0.3.5/src/neuralrnn/analysis/manifolds.py +461 -0
  18. neuralrnn-0.3.5/src/neuralrnn/analysis/perturbation.py +163 -0
  19. neuralrnn-0.3.5/src/neuralrnn/analysis/population_structure.py +116 -0
  20. neuralrnn-0.3.5/src/neuralrnn/analysis/psychometric.py +146 -0
  21. neuralrnn-0.3.5/src/neuralrnn/analysis/sequentiality.py +170 -0
  22. neuralrnn-0.3.5/src/neuralrnn/analysis/stsp_metrics.py +116 -0
  23. neuralrnn-0.3.5/src/neuralrnn/analysis/vector_field.py +61 -0
  24. neuralrnn-0.3.5/src/neuralrnn/auto/__init__.py +9 -0
  25. neuralrnn-0.3.5/src/neuralrnn/auto/configuration_auto.py +57 -0
  26. neuralrnn-0.3.5/src/neuralrnn/auto/modeling_auto.py +84 -0
  27. neuralrnn-0.3.5/src/neuralrnn/configuration_utils.py +198 -0
  28. neuralrnn-0.3.5/src/neuralrnn/data/__init__.py +25 -0
  29. neuralrnn-0.3.5/src/neuralrnn/data/bartolo_monkey_dataset.py +395 -0
  30. neuralrnn-0.3.5/src/neuralrnn/data/base.py +111 -0
  31. neuralrnn-0.3.5/src/neuralrnn/data/cognitive_task_dataset.py +245 -0
  32. neuralrnn-0.3.5/src/neuralrnn/data/custom_dataset.py +488 -0
  33. neuralrnn-0.3.5/src/neuralrnn/data/download.py +111 -0
  34. neuralrnn-0.3.5/src/neuralrnn/data/neurogym_dataset.py +333 -0
  35. neuralrnn-0.3.5/src/neuralrnn/data/reconstruction_dataset.py +207 -0
  36. neuralrnn-0.3.5/src/neuralrnn/data/registry.py +200 -0
  37. neuralrnn-0.3.5/src/neuralrnn/data/tasks/__init__.py +129 -0
  38. neuralrnn-0.3.5/src/neuralrnn/data/tasks/checkerboard_task.py +196 -0
  39. neuralrnn-0.3.5/src/neuralrnn/data/tasks/dms_continuous_task.py +153 -0
  40. neuralrnn-0.3.5/src/neuralrnn/data/tasks/dms_task.py +151 -0
  41. neuralrnn-0.3.5/src/neuralrnn/data/tasks/go_nogo_task.py +99 -0
  42. neuralrnn-0.3.5/src/neuralrnn/data/tasks/lr_mante_task.py +14 -0
  43. neuralrnn-0.3.5/src/neuralrnn/data/tasks/mante2_task.py +167 -0
  44. neuralrnn-0.3.5/src/neuralrnn/data/tasks/mante_task.py +182 -0
  45. neuralrnn-0.3.5/src/neuralrnn/data/tasks/multitask_flexible_dataset.py +130 -0
  46. neuralrnn-0.3.5/src/neuralrnn/data/tasks/multitask_flexible_task.py +893 -0
  47. neuralrnn-0.3.5/src/neuralrnn/data/tasks/multitask_yang_dataset.py +94 -0
  48. neuralrnn-0.3.5/src/neuralrnn/data/tasks/multitask_yang_task.py +1180 -0
  49. neuralrnn-0.3.5/src/neuralrnn/data/tasks/raposo_task.py +143 -0
  50. neuralrnn-0.3.5/src/neuralrnn/data/tasks/rdm_task.py +107 -0
  51. neuralrnn-0.3.5/src/neuralrnn/data/tasks/task_base.py +78 -0
  52. neuralrnn-0.3.5/src/neuralrnn/data/tasks/wm_angle_task.py +95 -0
  53. neuralrnn-0.3.5/src/neuralrnn/data/tasks/wm_frequency_task.py +129 -0
  54. neuralrnn-0.3.5/src/neuralrnn/modeling_utils.py +462 -0
  55. neuralrnn-0.3.5/src/neuralrnn/models/__init__.py +11 -0
  56. neuralrnn-0.3.5/src/neuralrnn/models/constrained_rnn/__init__.py +24 -0
  57. neuralrnn-0.3.5/src/neuralrnn/models/constrained_rnn/configuration_constrained_rnn.py +202 -0
  58. neuralrnn-0.3.5/src/neuralrnn/models/constrained_rnn/modeling_constrained_rnn.py +364 -0
  59. neuralrnn-0.3.5/src/neuralrnn/models/ctrnn/__init__.py +8 -0
  60. neuralrnn-0.3.5/src/neuralrnn/models/ctrnn/configuration_ctrnn.py +113 -0
  61. neuralrnn-0.3.5/src/neuralrnn/models/ctrnn/modeling_ctrnn.py +174 -0
  62. neuralrnn-0.3.5/src/neuralrnn/models/gain_rnn/__init__.py +27 -0
  63. neuralrnn-0.3.5/src/neuralrnn/models/gain_rnn/configuration_gain_rnn.py +309 -0
  64. neuralrnn-0.3.5/src/neuralrnn/models/gain_rnn/modeling_gain_rnn.py +535 -0
  65. neuralrnn-0.3.5/src/neuralrnn/models/latent_circuit/__init__.py +9 -0
  66. neuralrnn-0.3.5/src/neuralrnn/models/latent_circuit/configuration_latent_circuit.py +82 -0
  67. neuralrnn-0.3.5/src/neuralrnn/models/latent_circuit/modeling_latent_circuit.py +210 -0
  68. neuralrnn-0.3.5/src/neuralrnn/models/lowrank/__init__.py +17 -0
  69. neuralrnn-0.3.5/src/neuralrnn/models/lowrank/configuration_lowrank.py +158 -0
  70. neuralrnn-0.3.5/src/neuralrnn/models/lowrank/modeling_lowrank.py +366 -0
  71. neuralrnn-0.3.5/src/neuralrnn/models/multiarea_rnn/__init__.py +22 -0
  72. neuralrnn-0.3.5/src/neuralrnn/models/multiarea_rnn/configuration_multiarea_rnn.py +112 -0
  73. neuralrnn-0.3.5/src/neuralrnn/models/multiarea_rnn/masks.py +126 -0
  74. neuralrnn-0.3.5/src/neuralrnn/models/multiarea_rnn/modeling_multiarea_rnn.py +96 -0
  75. neuralrnn-0.3.5/src/neuralrnn/models/plrnn/__init__.py +8 -0
  76. neuralrnn-0.3.5/src/neuralrnn/models/plrnn/configuration_plrnn.py +133 -0
  77. neuralrnn-0.3.5/src/neuralrnn/models/plrnn/modeling_plrnn.py +479 -0
  78. neuralrnn-0.3.5/src/neuralrnn/models/tiny_rnn/__init__.py +5 -0
  79. neuralrnn-0.3.5/src/neuralrnn/models/tiny_rnn/configuration_tiny_rnn.py +71 -0
  80. neuralrnn-0.3.5/src/neuralrnn/models/tiny_rnn/modeling_tiny_rnn.py +219 -0
  81. neuralrnn-0.3.5/src/neuralrnn/train/__init__.py +68 -0
  82. neuralrnn-0.3.5/src/neuralrnn/train/cv.py +103 -0
  83. neuralrnn-0.3.5/src/neuralrnn/train/losses/__init__.py +24 -0
  84. neuralrnn-0.3.5/src/neuralrnn/train/losses/loss_functions.py +90 -0
  85. neuralrnn-0.3.5/src/neuralrnn/train/losses/metrics.py +73 -0
  86. neuralrnn-0.3.5/src/neuralrnn/train/losses/regularizers.py +176 -0
  87. neuralrnn-0.3.5/src/neuralrnn/train/objectives/__init__.py +31 -0
  88. neuralrnn-0.3.5/src/neuralrnn/train/objectives/base.py +36 -0
  89. neuralrnn-0.3.5/src/neuralrnn/train/objectives/behavioral.py +56 -0
  90. neuralrnn-0.3.5/src/neuralrnn/train/objectives/constrained.py +38 -0
  91. neuralrnn-0.3.5/src/neuralrnn/train/objectives/reconstruction.py +182 -0
  92. neuralrnn-0.3.5/src/neuralrnn/train/objectives/registry.py +120 -0
  93. neuralrnn-0.3.5/src/neuralrnn/train/objectives/regularized_supervised.py +109 -0
  94. neuralrnn-0.3.5/src/neuralrnn/train/objectives/supervised.py +50 -0
  95. neuralrnn-0.3.5/src/neuralrnn/train/objectives/teacher_forcing.py +125 -0
  96. neuralrnn-0.3.5/src/neuralrnn/train/objectives/variational.py +60 -0
  97. neuralrnn-0.3.5/src/neuralrnn/train/trainer.py +413 -0
  98. neuralrnn-0.3.5/src/neuralrnn/train/training_args.py +66 -0
  99. neuralrnn-0.3.5/src/neuralrnn/visualization.py +1084 -0
  100. neuralrnn-0.3.5/test/test_activations.py +77 -0
  101. neuralrnn-0.3.5/test/test_analysis_utils.py +177 -0
  102. neuralrnn-0.3.5/test/test_constrained_rnn.py +178 -0
  103. neuralrnn-0.3.5/test/test_custom_dataset_reconstruction.py +180 -0
  104. neuralrnn-0.3.5/test/test_cv.py +73 -0
  105. neuralrnn-0.3.5/test/test_datasets.py +175 -0
  106. neuralrnn-0.3.5/test/test_dropout.py +147 -0
  107. neuralrnn-0.3.5/test/test_eirnn.py +249 -0
  108. neuralrnn-0.3.5/test/test_euler_alpha.py +139 -0
  109. neuralrnn-0.3.5/test/test_fixed_points.py +190 -0
  110. neuralrnn-0.3.5/test/test_freeze.py +331 -0
  111. neuralrnn-0.3.5/test/test_gain_rnn.py +434 -0
  112. neuralrnn-0.3.5/test/test_latent_circuit.py +387 -0
  113. neuralrnn-0.3.5/test/test_line_attractor.py +112 -0
  114. neuralrnn-0.3.5/test/test_losses.py +182 -0
  115. neuralrnn-0.3.5/test/test_lowrank_rnn.py +132 -0
  116. neuralrnn-0.3.5/test/test_lyapunov.py +73 -0
  117. neuralrnn-0.3.5/test/test_manifolds.py +120 -0
  118. neuralrnn-0.3.5/test/test_multiarea_rnn.py +226 -0
  119. neuralrnn-0.3.5/test/test_neurogym_dataset.py +197 -0
  120. neuralrnn-0.3.5/test/test_nonlinearity_mode.py +349 -0
  121. neuralrnn-0.3.5/test/test_objectives.py +357 -0
  122. neuralrnn-0.3.5/test/test_plrnn_dsr.py +305 -0
  123. neuralrnn-0.3.5/test/test_plrnn_variants.py +200 -0
  124. neuralrnn-0.3.5/test/test_reconstruction.py +371 -0
  125. neuralrnn-0.3.5/test/test_smoke.py +146 -0
  126. neuralrnn-0.3.5/test/test_stp_rnn.py +432 -0
  127. neuralrnn-0.3.5/test/test_task_refactor.py +279 -0
  128. neuralrnn-0.3.5/test/test_tiny_rnn.py +104 -0
  129. neuralrnn-0.3.5/test/test_trainer_logging.py +74 -0
  130. neuralrnn-0.3.5/test/test_visualization.py +253 -0
@@ -0,0 +1,233 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+ docs/
6
+ docs/ARCHITECTURE.md
7
+ docs/PORTING_GUIDE.md
8
+ docs/papers/*_cn.md
9
+ notebook/figs/
10
+ notebook/models/
11
+ notebook/test/
12
+ notebook/test_analysis_fixes.py
13
+ notebook/test.ipynb
14
+ notebook/*.py
15
+ notebook/**/*.py
16
+ **/__pycache__/
17
+ notebook/_*
18
+ **/*_executed.ipynb
19
+ temp/
20
+ test/_**
21
+ # C extensions
22
+ *.so
23
+ scripts/*
24
+ # Distribution / packaging
25
+ .Python
26
+ build/
27
+ develop-eggs/
28
+ dist/
29
+ downloads/
30
+ eggs/
31
+ .eggs/
32
+ lib/
33
+ lib64/
34
+ parts/
35
+ sdist/
36
+ var/
37
+ wheels/
38
+ share/python-wheels/
39
+ *.egg-info/
40
+ .installed.cfg
41
+ *.egg
42
+ MANIFEST
43
+
44
+ # PyInstaller
45
+ # Usually these files are written by a python script from a template
46
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
47
+ *.manifest
48
+ *.spec
49
+
50
+ # Installer logs
51
+ pip-log.txt
52
+ pip-delete-this-directory.txt
53
+
54
+ # Unit test / coverage reports
55
+ htmlcov/
56
+ .tox/
57
+ .nox/
58
+ .coverage
59
+ .coverage.*
60
+ .cache
61
+ nosetests.xml
62
+ coverage.xml
63
+ *.cover
64
+ *.py.cover
65
+ .hypothesis/
66
+ .pytest_cache/
67
+ cover/
68
+
69
+ # Translations
70
+ *.mo
71
+ *.pot
72
+
73
+ # Django stuff:
74
+ *.log
75
+ local_settings.py
76
+ db.sqlite3
77
+ db.sqlite3-journal
78
+
79
+ # Flask stuff:
80
+ instance/
81
+ .webassets-cache
82
+
83
+ # Scrapy stuff:
84
+ .scrapy
85
+
86
+ # Sphinx documentation
87
+ docs/_build/
88
+
89
+ # PyBuilder
90
+ .pybuilder/
91
+ target/
92
+
93
+ # Jupyter Notebook
94
+ .ipynb_checkpoints
95
+
96
+ # IPython
97
+ profile_default/
98
+ ipython_config.py
99
+
100
+ # pyenv
101
+ # For a library or package, you might want to ignore these files since the code is
102
+ # intended to run in multiple environments; otherwise, check them in:
103
+ # .python-version
104
+
105
+ # pipenv
106
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
107
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
108
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
109
+ # install all needed dependencies.
110
+ # Pipfile.lock
111
+
112
+ # UV
113
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
114
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
115
+ # commonly ignored for libraries.
116
+ # uv.lock
117
+
118
+ # poetry
119
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
120
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
121
+ # commonly ignored for libraries.
122
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
123
+ # poetry.lock
124
+ # poetry.toml
125
+
126
+ # pdm
127
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
128
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
129
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
130
+ # pdm.lock
131
+ # pdm.toml
132
+ .pdm-python
133
+ .pdm-build/
134
+
135
+ # pixi
136
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
137
+ # pixi.lock
138
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
139
+ # in the .venv directory. It is recommended not to include this directory in version control.
140
+ .pixi
141
+
142
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
143
+ __pypackages__/
144
+
145
+ # Celery stuff
146
+ celerybeat-schedule
147
+ celerybeat.pid
148
+
149
+ # Redis
150
+ *.rdb
151
+ *.aof
152
+ *.pid
153
+
154
+ # RabbitMQ
155
+ mnesia/
156
+ rabbitmq/
157
+ rabbitmq-data/
158
+
159
+ # ActiveMQ
160
+ activemq-data/
161
+
162
+ # SageMath parsed files
163
+ *.sage.py
164
+
165
+ # Environments
166
+ .env
167
+ .envrc
168
+ .venv
169
+ env/
170
+ venv/
171
+ ENV/
172
+ env.bak/
173
+ venv.bak/
174
+
175
+ # Spyder project settings
176
+ .spyderproject
177
+ .spyproject
178
+
179
+ # Rope project settings
180
+ .ropeproject
181
+
182
+ # mkdocs documentation
183
+ /site
184
+
185
+ # mypy
186
+ .mypy_cache/
187
+ .dmypy.json
188
+ dmypy.json
189
+
190
+ # Pyre type checker
191
+ .pyre/
192
+
193
+ # pytype static type analyzer
194
+ .pytype/
195
+
196
+ # Cython debug symbols
197
+ cython_debug/
198
+
199
+ # PyCharm
200
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
201
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
202
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
203
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
204
+ # .idea/
205
+
206
+ # Abstra
207
+ # Abstra is an AI-powered process automation framework.
208
+ # Ignore directories containing user credentials, local state, and settings.
209
+ # Learn more at https://abstra.io/docs
210
+ .abstra/
211
+
212
+ # Visual Studio Code
213
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
214
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
215
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
216
+ # you could uncomment the following to ignore the entire vscode folder
217
+ # .vscode/
218
+ # Temporary file for partial code execution
219
+ tempCodeRunnerFile.py
220
+
221
+ # Ruff stuff:
222
+ .ruff_cache/
223
+
224
+ # PyPI configuration file
225
+ .pypirc
226
+
227
+ # Marimo
228
+ marimo/_static/
229
+ marimo/_lsp/
230
+ __marimo__/
231
+
232
+ # Streamlit
233
+ .streamlit/secrets.toml
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 HH Chen
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,156 @@
1
+ Metadata-Version: 2.4
2
+ Name: neuralrnn
3
+ Version: 0.3.5
4
+ Summary: A unified framework for implementing RNN methods in cognitive neuroscience
5
+ Project-URL: Homepage, https://github.com/HHChen730/NeuralRNN
6
+ Project-URL: Documentation, https://neuralrnn.readthedocs.io/
7
+ Author-email: Honghua Chen <honghuachen69@gmail.com>
8
+ License: MIT
9
+ License-File: LICENSE
10
+ Keywords: PLRNN,RNN,computational-neuroscience,dynamical-systems,dynamics-reconstruction,interpretability
11
+ Classifier: License :: OSI Approved :: MIT License
12
+ Classifier: Programming Language :: Python :: 3.9
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Requires-Python: >=3.9
17
+ Requires-Dist: gym>=0.21
18
+ Requires-Dist: ipykernel
19
+ Requires-Dist: jupyter
20
+ Requires-Dist: matplotlib>=3.6
21
+ Requires-Dist: neurogym>=0.0.1
22
+ Requires-Dist: numpy>=1.23
23
+ Requires-Dist: safetensors>=0.4
24
+ Requires-Dist: scipy>=1.10
25
+ Requires-Dist: torch>=2.0
26
+ Requires-Dist: tqdm
27
+ Description-Content-Type: text/markdown
28
+
29
+ # NeuralRNN
30
+
31
+ ---
32
+
33
+ **NeuralRNN is A unified framework for implementing RNN methods in cognitive neuroscience** — bringing two major paradigms under a single [Transformers](https://github.com/huggingface/transformers)-style interface:
34
+
35
+ - **Paradigm A: Task Optimization**[^1]: Train RNNs on cognitive tasks, then reverse-engineer how they perform computation using analyses including fixed points, vector fields, dimensionality reduction, etc. The goal is to use RNNs as a proxy for cognitive computation.
36
+ - **Paradigm B: Dynamical System Reconstruction (DSR)**[^2][^3]: Fit generative RNNs directly from neural/behavioral time series that can reproduce attractors, power spectra, and Lyapunov spectra of the target system.
37
+
38
+ Both paradigms share a unified set of `model config`, `Trainer`, and `analysis` tools. **The only difference between the two paradigms is the `Objective`**: Paradigm A aims to optimize output for cognitive task performance, while Paradigm B aims to construct a dynamical system isomorphic to the target neural activity. Moreover, DSR can also be applied to reconstruct the dynamics of TBO-trained models for interpretability analysis[^4].
39
+
40
+ ---
41
+
42
+ ## Core Concept
43
+
44
+ All models are viewed as "discrete dynamical systems with downstream readout" $z_t=F_\theta(z_{t-1},x_t),\;y_t=G_\phi(z_t)$.
45
+
46
+ A model only needs to implement two methods:
47
+
48
+ ```python
49
+ def recurrence(self, x_t, z_prev, *, inputs=None): ... # single-step transition F
50
+ def readout(self, z_t): ... # readout G
51
+ ```
52
+
53
+ NeuralRNN provides an interface to automatically connect the model to the unified trainer and all analysis tools.
54
+
55
+ ## What NeuralRNN contains and not
56
+
57
+ NeuralRNN contains the pipeline to model RNN for neuroscience research, including (1) constructing dataset, (2) building and configuring RNN models, (3) training, and (4) model analysis (see the full pipeline in **[`custom pipeline`](notebook/03_custom_pipeline.ipynb)** and the documents in **[`docs`](docs/README.md)**). We also provide the guide to implement each built-in model through this framework (see **[`notebook`](notebook/README.md)**).
58
+
59
+ However, there are others methods using dynamical system methods as well, including [MARBLE](https://www.nature.com/articles/s41592-024-02582-2), [FINDR](https://www.nature.com/articles/s41586-025-09528-4), [neuralflow](https://www.nature.com/articles/s41586-025-09199-1), and [SSMLearn](https://arxiv.org/abs/2510.13519). These model-agnostic methods aim to inference interpretable representations of neural population dynamics exactly from the neural response, which are not included in NeuralRNN but can be suitably combined for the further analysis of RNN models.
60
+
61
+
62
+ ## Install
63
+
64
+ ```bash
65
+ $ git clone https://github.com/ChenHH730/NeuralRNN.git
66
+ $ cd NeuralRNN
67
+ $ pip install -e .
68
+ ```
69
+
70
+
71
+ ## Quickstart
72
+
73
+ ```python
74
+ from neuralrnn import AutoConfig, AutoModel, Trainer, TrainingArguments
75
+ from neuralrnn import TeacherForcingObjective, load_dataset
76
+
77
+ # 1) dataset
78
+ # use registered dataset or custom dataset
79
+ ds = load_dataset("lorenz63", sequence_length=200, batch_size=16, normalize=True)
80
+
81
+ # 2) model (config) + objective (based on the paradigm) + training
82
+ cfg = AutoConfig.for_model("shallow_plrnn", input_dim=0, latent_dim=3,
83
+ output_dim=3, hidden_dim=50, autonomous=True) # model config
84
+ model = AutoModel.from_config(cfg) # load model
85
+ Trainer(model, ds, TeacherForcingObjective(alpha=0.1),
86
+ TrainingArguments(max_steps=2000)).train() # train model
87
+
88
+ # 3) save and load (config.json + model.safetensors)
89
+ model.save_pretrained("ckpt/")
90
+ model = AutoModel.from_pretrained("ckpt/")
91
+
92
+ # 4) analysis (model agnostic)
93
+ from neuralrnn.analysis import find_fixed_points, max_lyapunov_exponent
94
+ fps = find_fixed_points(model)
95
+ ```
96
+
97
+ ## Content Structure
98
+
99
+ ```
100
+ src/neuralrnn/
101
+ configuration_utils.py modeling_utils.py # core contracts (Config / Model base classes)
102
+ auto/ # AutoConfig / AutoModel registration & dispatch
103
+ models/ # model zoo: ctrnn, ei_rnn, lowrank_rnn (Paradigm A), plrnn (Paradigm B), latent_circuit, tiny_rnn,
104
+ # constrained_rnn, multiarea_rnn, gain_rnn (gain_rnn + stp_rnn)
105
+ data/ # unified batching, datasets, open data registry + download cache
106
+ train/ # generic Trainer + paradigm Objectives + reusable loss terms /
107
+ # regularizers / metrics + nested cross-validation
108
+ analysis/ # fixed points / linearization / vector fields / dim reduction /
109
+ # Lyapunov / D_stsp, D_H / PLRNN invariant manifolds / sequentiality
110
+ docs/ # ARCHITECTURE.md · PORTING_GUIDE.md · theory/ · papers/
111
+ notebook/ # end-to-end tutorials for each paper
112
+ ```
113
+
114
+ ## Built-in Models
115
+
116
+ | Model | Paradigm | Status (mostly used) |
117
+ |---|---|---|
118
+ | continuous time RNN | A | ✅ |
119
+ | E-I RNN (Dale's principle) | A | ✅ |
120
+ | Latent Circuit Model | B | ✅ |
121
+ | piecewise linear RNN | B | ✅ |
122
+ | Tiny RNN | B | ✅ |
123
+ | low-rank RNN | AB | ✅ |
124
+ | constrained RNN | A | ✅ |
125
+ | seRNN | A | ✅ |
126
+ | multi-area RNN | A | ✅ |
127
+ | gain RNN | AB | ✅ |
128
+ | STP RNN | A | ✅ |
129
+
130
+ The second column shows the corresponding paradigm used in the original work.
131
+
132
+ ## Porting New Papers into the Framework
133
+
134
+ Core principle: **Porting = writing adapters (wrapping + verification), not rewriting mathematics**. Any model that implements `recurrence/readout` is plug-and-play; the analysis layer works only through the model's public contract and never imports specific model classes.
135
+
136
+ ## License
137
+
138
+ MIT, see [LICENSE](LICENSE). Original code of ported papers belongs to their respective authors; please follow their individual licenses when porting.
139
+
140
+ ## References
141
+
142
+ [^1]: [Training Excitatory-Inhibitory Recurrent Neural Networks for Cognitive Tasks](https://doi.org/10.1371/journal.pcbi.1004792).
143
+ Project: https://github.com/gyyang/nn-brain
144
+
145
+ [^2]: [Reconstructing computational dynamics from neural measurements with RNN](https://www.nature.com/articles/s41583-023-00740-7)
146
+ Project: https://github.com/DurstewitzLab/CNS-2023
147
+
148
+ [^3]: [Discovering cognitive strategies with tiny-RNN](https://www.nature.com/articles/s41586-025-09142-4)
149
+ Project: https://github.com/jil095/tinyRNN
150
+
151
+ [^4]: https://github.com/engellab/latentcircuit
152
+
153
+ [^5]: https://github.com/Dynamics-of-Neural-Systems-Lab/MARBLE
154
+
155
+ [^6]: https://github.com/NN4Neurosim/nn4n https://nn4n.org/
156
+ [^7]:
@@ -0,0 +1,128 @@
1
+ # NeuralRNN
2
+
3
+ ---
4
+
5
+ **NeuralRNN is A unified framework for implementing RNN methods in cognitive neuroscience** — bringing two major paradigms under a single [Transformers](https://github.com/huggingface/transformers)-style interface:
6
+
7
+ - **Paradigm A: Task Optimization**[^1]: Train RNNs on cognitive tasks, then reverse-engineer how they perform computation using analyses including fixed points, vector fields, dimensionality reduction, etc. The goal is to use RNNs as a proxy for cognitive computation.
8
+ - **Paradigm B: Dynamical System Reconstruction (DSR)**[^2][^3]: Fit generative RNNs directly from neural/behavioral time series that can reproduce attractors, power spectra, and Lyapunov spectra of the target system.
9
+
10
+ Both paradigms share a unified set of `model config`, `Trainer`, and `analysis` tools. **The only difference between the two paradigms is the `Objective`**: Paradigm A aims to optimize output for cognitive task performance, while Paradigm B aims to construct a dynamical system isomorphic to the target neural activity. Moreover, DSR can also be applied to reconstruct the dynamics of TBO-trained models for interpretability analysis[^4].
11
+
12
+ ---
13
+
14
+ ## Core Concept
15
+
16
+ All models are viewed as "discrete dynamical systems with downstream readout" $z_t=F_\theta(z_{t-1},x_t),\;y_t=G_\phi(z_t)$.
17
+
18
+ A model only needs to implement two methods:
19
+
20
+ ```python
21
+ def recurrence(self, x_t, z_prev, *, inputs=None): ... # single-step transition F
22
+ def readout(self, z_t): ... # readout G
23
+ ```
24
+
25
+ NeuralRNN provides an interface to automatically connect the model to the unified trainer and all analysis tools.
26
+
27
+ ## What NeuralRNN contains and not
28
+
29
+ NeuralRNN contains the pipeline to model RNN for neuroscience research, including (1) constructing dataset, (2) building and configuring RNN models, (3) training, and (4) model analysis (see the full pipeline in **[`custom pipeline`](notebook/03_custom_pipeline.ipynb)** and the documents in **[`docs`](docs/README.md)**). We also provide the guide to implement each built-in model through this framework (see **[`notebook`](notebook/README.md)**).
30
+
31
+ However, there are others methods using dynamical system methods as well, including [MARBLE](https://www.nature.com/articles/s41592-024-02582-2), [FINDR](https://www.nature.com/articles/s41586-025-09528-4), [neuralflow](https://www.nature.com/articles/s41586-025-09199-1), and [SSMLearn](https://arxiv.org/abs/2510.13519). These model-agnostic methods aim to inference interpretable representations of neural population dynamics exactly from the neural response, which are not included in NeuralRNN but can be suitably combined for the further analysis of RNN models.
32
+
33
+
34
+ ## Install
35
+
36
+ ```bash
37
+ $ git clone https://github.com/ChenHH730/NeuralRNN.git
38
+ $ cd NeuralRNN
39
+ $ pip install -e .
40
+ ```
41
+
42
+
43
+ ## Quickstart
44
+
45
+ ```python
46
+ from neuralrnn import AutoConfig, AutoModel, Trainer, TrainingArguments
47
+ from neuralrnn import TeacherForcingObjective, load_dataset
48
+
49
+ # 1) dataset
50
+ # use registered dataset or custom dataset
51
+ ds = load_dataset("lorenz63", sequence_length=200, batch_size=16, normalize=True)
52
+
53
+ # 2) model (config) + objective (based on the paradigm) + training
54
+ cfg = AutoConfig.for_model("shallow_plrnn", input_dim=0, latent_dim=3,
55
+ output_dim=3, hidden_dim=50, autonomous=True) # model config
56
+ model = AutoModel.from_config(cfg) # load model
57
+ Trainer(model, ds, TeacherForcingObjective(alpha=0.1),
58
+ TrainingArguments(max_steps=2000)).train() # train model
59
+
60
+ # 3) save and load (config.json + model.safetensors)
61
+ model.save_pretrained("ckpt/")
62
+ model = AutoModel.from_pretrained("ckpt/")
63
+
64
+ # 4) analysis (model agnostic)
65
+ from neuralrnn.analysis import find_fixed_points, max_lyapunov_exponent
66
+ fps = find_fixed_points(model)
67
+ ```
68
+
69
+ ## Content Structure
70
+
71
+ ```
72
+ src/neuralrnn/
73
+ configuration_utils.py modeling_utils.py # core contracts (Config / Model base classes)
74
+ auto/ # AutoConfig / AutoModel registration & dispatch
75
+ models/ # model zoo: ctrnn, ei_rnn, lowrank_rnn (Paradigm A), plrnn (Paradigm B), latent_circuit, tiny_rnn,
76
+ # constrained_rnn, multiarea_rnn, gain_rnn (gain_rnn + stp_rnn)
77
+ data/ # unified batching, datasets, open data registry + download cache
78
+ train/ # generic Trainer + paradigm Objectives + reusable loss terms /
79
+ # regularizers / metrics + nested cross-validation
80
+ analysis/ # fixed points / linearization / vector fields / dim reduction /
81
+ # Lyapunov / D_stsp, D_H / PLRNN invariant manifolds / sequentiality
82
+ docs/ # ARCHITECTURE.md · PORTING_GUIDE.md · theory/ · papers/
83
+ notebook/ # end-to-end tutorials for each paper
84
+ ```
85
+
86
+ ## Built-in Models
87
+
88
+ | Model | Paradigm | Status (mostly used) |
89
+ |---|---|---|
90
+ | continuous time RNN | A | ✅ |
91
+ | E-I RNN (Dale's principle) | A | ✅ |
92
+ | Latent Circuit Model | B | ✅ |
93
+ | piecewise linear RNN | B | ✅ |
94
+ | Tiny RNN | B | ✅ |
95
+ | low-rank RNN | AB | ✅ |
96
+ | constrained RNN | A | ✅ |
97
+ | seRNN | A | ✅ |
98
+ | multi-area RNN | A | ✅ |
99
+ | gain RNN | AB | ✅ |
100
+ | STP RNN | A | ✅ |
101
+
102
+ The second column shows the corresponding paradigm used in the original work.
103
+
104
+ ## Porting New Papers into the Framework
105
+
106
+ Core principle: **Porting = writing adapters (wrapping + verification), not rewriting mathematics**. Any model that implements `recurrence/readout` is plug-and-play; the analysis layer works only through the model's public contract and never imports specific model classes.
107
+
108
+ ## License
109
+
110
+ MIT, see [LICENSE](LICENSE). Original code of ported papers belongs to their respective authors; please follow their individual licenses when porting.
111
+
112
+ ## References
113
+
114
+ [^1]: [Training Excitatory-Inhibitory Recurrent Neural Networks for Cognitive Tasks](https://doi.org/10.1371/journal.pcbi.1004792).
115
+ Project: https://github.com/gyyang/nn-brain
116
+
117
+ [^2]: [Reconstructing computational dynamics from neural measurements with RNN](https://www.nature.com/articles/s41583-023-00740-7)
118
+ Project: https://github.com/DurstewitzLab/CNS-2023
119
+
120
+ [^3]: [Discovering cognitive strategies with tiny-RNN](https://www.nature.com/articles/s41586-025-09142-4)
121
+ Project: https://github.com/jil095/tinyRNN
122
+
123
+ [^4]: https://github.com/engellab/latentcircuit
124
+
125
+ [^5]: https://github.com/Dynamics-of-Neural-Systems-Lab/MARBLE
126
+
127
+ [^6]: https://github.com/NN4Neurosim/nn4n https://nn4n.org/
128
+ [^7]:
@@ -0,0 +1,53 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "neuralrnn"
7
+ version = "0.3.5"
8
+ description = "A unified framework for implementing RNN methods in cognitive neuroscience"
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = { text = "MIT" }
12
+ authors = [{ name = "Honghua Chen", email = "honghuachen69@gmail.com" }]
13
+ keywords = ["RNN", "computational-neuroscience", "dynamical-systems",
14
+ "dynamics-reconstruction", "interpretability", "PLRNN"]
15
+
16
+ dependencies = [
17
+ "torch>=2.0",
18
+ "numpy>=1.23",
19
+ "scipy>=1.10",
20
+ "safetensors>=0.4",
21
+ "neurogym>=0.0.1",
22
+ "gym>=0.21",
23
+ "matplotlib>=3.6",
24
+ "jupyter",
25
+ "ipykernel",
26
+ "tqdm",
27
+ ]
28
+
29
+ classifiers = [
30
+ "License :: OSI Approved :: MIT License",
31
+ "Programming Language :: Python :: 3.9",
32
+ "Programming Language :: Python :: 3.10",
33
+ "Programming Language :: Python :: 3.11",
34
+ "Programming Language :: Python :: 3.12",
35
+ ]
36
+
37
+
38
+ [project.urls]
39
+ Homepage = "https://github.com/HHChen730/NeuralRNN"
40
+ Documentation = "https://neuralrnn.readthedocs.io/"
41
+
42
+ [tool.hatch.build.targets.wheel]
43
+ packages = ["src/neuralrnn"]
44
+
45
+ [tool.ruff]
46
+ line-length = 100
47
+ target-version = "py39"
48
+
49
+ [tool.pytest.ini_options]
50
+ testpaths = ["test"]
51
+
52
+ [tool.setuptools.exclude-package-data]
53
+ "neuralrnn" = ["docs/*", "notebook/*", "*.ipynb"]
@@ -0,0 +1,79 @@
1
+ """NeuralRNN: a unified framework for RNN methods in cognitive neuroscience.
2
+
3
+ Brings two front-line paradigms under one transformers-style API:
4
+ Paradigm A -- optimize and train RNNs on cognitive tasks, with interpretability analysis
5
+ (fixed points / vector fields / dimensionality reduction);
6
+ Paradigm B -- reconstruct dynamics directly from neural / behavioral data
7
+ (PLRNN / LFADS / low-rank / Tiny RNN).
8
+
9
+ Core abstraction: every model is a "discrete dynamical system with readout". Implementing only
10
+ `recurrence` and `readout` is enough to plug into the unified Trainer (paradigm differences are
11
+ handled by Objective) and analysis modules.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ __version__ = "0.3.5"
17
+
18
+ # Core base classes / output container
19
+ from .configuration_utils import (
20
+ NeuralRNNConfig, resolve_euler_alpha,
21
+ SUPPORTED_NONLINEARITY_MODES, validate_nonlinearity_mode,
22
+ )
23
+ from .modeling_utils import NeuralDynamicsModel, DynamicsModelOutput
24
+
25
+ # Unified activation factory
26
+ from .activations import get_activation, SUPPORTED_ACTIVATIONS
27
+
28
+ # Auto factories
29
+ from .auto import (
30
+ AutoConfig, AutoModel,
31
+ register_config, register_model,
32
+ CONFIG_REGISTRY, MODEL_REGISTRY,
33
+ )
34
+
35
+ # Data
36
+ from .data import (
37
+ BaseDataset, StandardScaler, CustomDataset,
38
+ CognitiveTaskDataset, ReconstructionDataset,
39
+ DATASET_REGISTRY, DatasetSpec, load_dataset,
40
+ )
41
+
42
+ # Training
43
+ from .train import (
44
+ Trainer, TrainingArguments,
45
+ Objective, SupervisedObjective, RegularizedSupervisedObjective,
46
+ TeacherForcingObjective, BehavioralObjective, VariationalObjective,
47
+ ReconstructionObjective, ConstrainedSupervisedObjective,
48
+ build_objective, register_objective, OBJECTIVE_REGISTRY, AutoObjective,
49
+ masked_mse, masked_cross_entropy, masked_nll, loss_mse,
50
+ activity_l2, weight_l2, weight_l1,
51
+ orthogonality_penalty, model_orthogonality_penalty,
52
+ accuracy_classification, accuracy_general,
53
+ )
54
+
55
+ # Visualization
56
+ from . import visualization
57
+
58
+ __all__ = [
59
+ "__version__",
60
+ "NeuralRNNConfig", "NeuralDynamicsModel", "DynamicsModelOutput",
61
+ "resolve_euler_alpha",
62
+ "SUPPORTED_NONLINEARITY_MODES", "validate_nonlinearity_mode",
63
+ "get_activation", "SUPPORTED_ACTIVATIONS",
64
+ "AutoConfig", "AutoModel", "register_config", "register_model",
65
+ "CONFIG_REGISTRY", "MODEL_REGISTRY",
66
+ "BaseDataset", "StandardScaler", "CustomDataset",
67
+ "CognitiveTaskDataset", "ReconstructionDataset",
68
+ "DATASET_REGISTRY", "DatasetSpec", "load_dataset",
69
+ "Trainer", "TrainingArguments",
70
+ "Objective", "SupervisedObjective", "RegularizedSupervisedObjective",
71
+ "TeacherForcingObjective", "BehavioralObjective", "VariationalObjective",
72
+ "ReconstructionObjective", "ConstrainedSupervisedObjective",
73
+ "build_objective", "register_objective", "OBJECTIVE_REGISTRY", "AutoObjective",
74
+ "masked_mse", "masked_cross_entropy", "masked_nll", "loss_mse",
75
+ "activity_l2", "weight_l2", "weight_l1",
76
+ "orthogonality_penalty", "model_orthogonality_penalty",
77
+ "accuracy_classification", "accuracy_general",
78
+ ]
79
+