hybrax 0.1.0__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 (186) hide show
  1. hybrax-0.1.0/.github/workflows/publish.yml +44 -0
  2. hybrax-0.1.0/.gitignore +242 -0
  3. hybrax-0.1.0/.python-version +1 -0
  4. hybrax-0.1.0/.readthedocs.yaml +12 -0
  5. hybrax-0.1.0/LICENSE +21 -0
  6. hybrax-0.1.0/MANIFEST.in +4 -0
  7. hybrax-0.1.0/PKG-INFO +132 -0
  8. hybrax-0.1.0/README.md +75 -0
  9. hybrax-0.1.0/pyproject.toml +118 -0
  10. hybrax-0.1.0/setup.cfg +4 -0
  11. hybrax-0.1.0/specs/format/01_design_rationale.md +188 -0
  12. hybrax-0.1.0/specs/format/02_data_model.md +531 -0
  13. hybrax-0.1.0/specs/format/03_serialization.md +298 -0
  14. hybrax-0.1.0/specs/format/04_validation.md +376 -0
  15. hybrax-0.1.0/specs/format/05_inspection.md +139 -0
  16. hybrax-0.1.0/specs/format/06_time_series.md +219 -0
  17. hybrax-0.1.0/specs/format/07_splines.md +329 -0
  18. hybrax-0.1.0/specs/format/08_mechanistic.md +264 -0
  19. hybrax-0.1.0/specs/format/09_simulation.md +212 -0
  20. hybrax-0.1.0/specs/format/README.md +103 -0
  21. hybrax-0.1.0/specs/train/01_design_rationale.md +148 -0
  22. hybrax-0.1.0/specs/train/02_cli_and_config.md +469 -0
  23. hybrax-0.1.0/specs/train/03_data_preparation.md +267 -0
  24. hybrax-0.1.0/specs/train/04_reaction_and_loss.md +628 -0
  25. hybrax-0.1.0/specs/train/05_train_forward_loo.md +302 -0
  26. hybrax-0.1.0/specs/train/06_serialization_inspect.md +194 -0
  27. hybrax-0.1.0/specs/train/README.md +115 -0
  28. hybrax-0.1.0/src/hybrax/__init__.py +5 -0
  29. hybrax-0.1.0/src/hybrax/format/__init__.py +160 -0
  30. hybrax-0.1.0/src/hybrax/format/_logging.py +18 -0
  31. hybrax-0.1.0/src/hybrax/format/dataclasses.py +599 -0
  32. hybrax-0.1.0/src/hybrax/format/inspect.py +1428 -0
  33. hybrax-0.1.0/src/hybrax/format/json_io.py +154 -0
  34. hybrax-0.1.0/src/hybrax/format/mechanistic.py +1051 -0
  35. hybrax-0.1.0/src/hybrax/format/monotone_spline.py +648 -0
  36. hybrax-0.1.0/src/hybrax/format/serialization.py +993 -0
  37. hybrax-0.1.0/src/hybrax/format/simulation.py +384 -0
  38. hybrax-0.1.0/src/hybrax/format/splines.py +1853 -0
  39. hybrax-0.1.0/src/hybrax/format/time_series/__init__.py +6 -0
  40. hybrax-0.1.0/src/hybrax/format/time_series/constants.py +11 -0
  41. hybrax-0.1.0/src/hybrax/format/time_series/grid_utils.py +74 -0
  42. hybrax-0.1.0/src/hybrax/format/time_series/io.py +49 -0
  43. hybrax-0.1.0/src/hybrax/format/time_series/ppoly.py +120 -0
  44. hybrax-0.1.0/src/hybrax/format/time_series/spline_ops.py +260 -0
  45. hybrax-0.1.0/src/hybrax/format/time_series/timeseries.py +660 -0
  46. hybrax-0.1.0/src/hybrax/format/validate.py +1383 -0
  47. hybrax-0.1.0/src/hybrax/train/__init__.py +277 -0
  48. hybrax-0.1.0/src/hybrax/train/augmentation.py +457 -0
  49. hybrax-0.1.0/src/hybrax/train/augmentation_plot.py +155 -0
  50. hybrax-0.1.0/src/hybrax/train/bounds_loss.py +244 -0
  51. hybrax-0.1.0/src/hybrax/train/checkpointing.py +245 -0
  52. hybrax-0.1.0/src/hybrax/train/cli.py +1312 -0
  53. hybrax-0.1.0/src/hybrax/train/constants.py +10 -0
  54. hybrax-0.1.0/src/hybrax/train/controls.py +490 -0
  55. hybrax-0.1.0/src/hybrax/train/controls_store.py +1431 -0
  56. hybrax-0.1.0/src/hybrax/train/defaults.py +684 -0
  57. hybrax-0.1.0/src/hybrax/train/dense.py +141 -0
  58. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/README.md +345 -0
  59. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/__init__.py +26 -0
  60. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/_callbacks.py +392 -0
  61. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/_solution.py +124 -0
  62. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/_solve.py +883 -0
  63. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/examples/01_basic_fed_batch.py +164 -0
  64. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/examples/02_multiple_callbacks.py +153 -0
  65. hybrax-0.1.0/src/hybrax/train/diffrax_callbacks/examples/03_hybrid_node_control.py +362 -0
  66. hybrax-0.1.0/src/hybrax/train/forward_plotting.py +367 -0
  67. hybrax-0.1.0/src/hybrax/train/harness.py +2731 -0
  68. hybrax-0.1.0/src/hybrax/train/inspect.py +424 -0
  69. hybrax-0.1.0/src/hybrax/train/logging.py +416 -0
  70. hybrax-0.1.0/src/hybrax/train/loo.py +1423 -0
  71. hybrax-0.1.0/src/hybrax/train/loo_metrics.py +1045 -0
  72. hybrax-0.1.0/src/hybrax/train/model_api.py +1157 -0
  73. hybrax-0.1.0/src/hybrax/train/physical_solve.py +421 -0
  74. hybrax-0.1.0/src/hybrax/train/postprocessing.py +743 -0
  75. hybrax-0.1.0/src/hybrax/train/preflight.py +171 -0
  76. hybrax-0.1.0/src/hybrax/train/prepare.py +599 -0
  77. hybrax-0.1.0/src/hybrax/train/run_config.py +802 -0
  78. hybrax-0.1.0/src/hybrax/train/runtime_artifact.py +1614 -0
  79. hybrax-0.1.0/src/hybrax/train/runtime_context.py +534 -0
  80. hybrax-0.1.0/src/hybrax/train/serialization.py +658 -0
  81. hybrax-0.1.0/src/hybrax/train/trainer.py +717 -0
  82. hybrax-0.1.0/src/hybrax/train/training_data.py +920 -0
  83. hybrax-0.1.0/src/hybrax/train/utils.py +97 -0
  84. hybrax-0.1.0/src/hybrax/train/validate.py +330 -0
  85. hybrax-0.1.0/src/hybrax/train/wrapper.py +645 -0
  86. hybrax-0.1.0/src/hybrax.egg-info/PKG-INFO +132 -0
  87. hybrax-0.1.0/src/hybrax.egg-info/SOURCES.txt +184 -0
  88. hybrax-0.1.0/src/hybrax.egg-info/dependency_links.txt +1 -0
  89. hybrax-0.1.0/src/hybrax.egg-info/entry_points.txt +2 -0
  90. hybrax-0.1.0/src/hybrax.egg-info/requires.txt +39 -0
  91. hybrax-0.1.0/src/hybrax.egg-info/scm_file_list.json +804 -0
  92. hybrax-0.1.0/src/hybrax.egg-info/scm_version.json +8 -0
  93. hybrax-0.1.0/src/hybrax.egg-info/top_level.txt +1 -0
  94. hybrax-0.1.0/tests/format/__init__.py +1 -0
  95. hybrax-0.1.0/tests/format/conftest.py +11 -0
  96. hybrax-0.1.0/tests/format/e2e_tests/__init__.py +0 -0
  97. hybrax-0.1.0/tests/format/e2e_tests/sim_1/README.md +177 -0
  98. hybrax-0.1.0/tests/format/e2e_tests/sim_1/__init__.py +0 -0
  99. hybrax-0.1.0/tests/format/e2e_tests/sim_1/load_utils.py +458 -0
  100. hybrax-0.1.0/tests/format/e2e_tests/sim_1/real_space_segments.py +168 -0
  101. hybrax-0.1.0/tests/format/e2e_tests/sim_1/sim_plots/process_variables.png +0 -0
  102. hybrax-0.1.0/tests/format/e2e_tests/sim_1/sim_plots/rates.png +0 -0
  103. hybrax-0.1.0/tests/format/e2e_tests/sim_1/sim_plots/reactor_states_and_volumes.png +0 -0
  104. hybrax-0.1.0/tests/format/e2e_tests/sim_1/sim_results/process_collection.json +24857 -0
  105. hybrax-0.1.0/tests/format/e2e_tests/sim_1/sim_results/simulation_dense_output.csv +2919 -0
  106. hybrax-0.1.0/tests/format/e2e_tests/sim_1/sim_results/simulation_events.csv +17 -0
  107. hybrax-0.1.0/tests/format/e2e_tests/sim_1/simulation.py +845 -0
  108. hybrax-0.1.0/tests/format/e2e_tests/sim_1/test_real_space_mechanistic_reintegration.py +281 -0
  109. hybrax-0.1.0/tests/format/e2e_tests/sim_1/test_real_space_segments.py +242 -0
  110. hybrax-0.1.0/tests/format/e2e_tests/sim_1/test_simulation_fixture_regeneration.py +122 -0
  111. hybrax-0.1.0/tests/format/unit_tests/__init__.py +0 -0
  112. hybrax-0.1.0/tests/format/unit_tests/test_dataclasses.py +760 -0
  113. hybrax-0.1.0/tests/format/unit_tests/test_discrete_jump_step.py +290 -0
  114. hybrax-0.1.0/tests/format/unit_tests/test_inspect.py +125 -0
  115. hybrax-0.1.0/tests/format/unit_tests/test_inspect_panels_pseudobatch.py +156 -0
  116. hybrax-0.1.0/tests/format/unit_tests/test_jax_compatibility.py +173 -0
  117. hybrax-0.1.0/tests/format/unit_tests/test_json_io.py +147 -0
  118. hybrax-0.1.0/tests/format/unit_tests/test_logging.py +35 -0
  119. hybrax-0.1.0/tests/format/unit_tests/test_mechanistic.py +1366 -0
  120. hybrax-0.1.0/tests/format/unit_tests/test_monotone_spline_pilot.py +201 -0
  121. hybrax-0.1.0/tests/format/unit_tests/test_monotone_spline_progress.py +106 -0
  122. hybrax-0.1.0/tests/format/unit_tests/test_monotone_splines.py +825 -0
  123. hybrax-0.1.0/tests/format/unit_tests/test_ppoly.py +114 -0
  124. hybrax-0.1.0/tests/format/unit_tests/test_pseudobatch_splines_sampling.py +415 -0
  125. hybrax-0.1.0/tests/format/unit_tests/test_public_api.py +55 -0
  126. hybrax-0.1.0/tests/format/unit_tests/test_serialization.py +1075 -0
  127. hybrax-0.1.0/tests/format/unit_tests/test_simulation.py +311 -0
  128. hybrax-0.1.0/tests/format/unit_tests/test_splines.py +1996 -0
  129. hybrax-0.1.0/tests/format/unit_tests/test_timeseries_api.py +137 -0
  130. hybrax-0.1.0/tests/format/unit_tests/test_timeseries_dtype.py +108 -0
  131. hybrax-0.1.0/tests/format/unit_tests/test_timeseries_runtime.py +315 -0
  132. hybrax-0.1.0/tests/format/unit_tests/test_timeseries_serialization.py +405 -0
  133. hybrax-0.1.0/tests/format/unit_tests/test_utils.py +895 -0
  134. hybrax-0.1.0/tests/format/unit_tests/test_validate.py +2117 -0
  135. hybrax-0.1.0/tests/train/conftest.py +11 -0
  136. hybrax-0.1.0/tests/train/diffrax_callback/test_callbacks.py +909 -0
  137. hybrax-0.1.0/tests/train/diffrax_callback/test_done_and_empty_presets.py +102 -0
  138. hybrax-0.1.0/tests/train/diffrax_callback/test_failed_segment_collapse.py +273 -0
  139. hybrax-0.1.0/tests/train/diffrax_callback/test_segment_dt_carry.py +123 -0
  140. hybrax-0.1.0/tests/train/diffrax_callback/test_total_step_budget.py +92 -0
  141. hybrax-0.1.0/tests/train/input.json +1 -0
  142. hybrax-0.1.0/tests/train/stateful_helpers.py +246 -0
  143. hybrax-0.1.0/tests/train/test_augmentation.py +1328 -0
  144. hybrax-0.1.0/tests/train/test_bounds_loss.py +517 -0
  145. hybrax-0.1.0/tests/train/test_checkpointing.py +700 -0
  146. hybrax-0.1.0/tests/train/test_cli.py +330 -0
  147. hybrax-0.1.0/tests/train/test_cli_run_dir.py +303 -0
  148. hybrax-0.1.0/tests/train/test_control_splines.py +494 -0
  149. hybrax-0.1.0/tests/train/test_controls_store.py +825 -0
  150. hybrax-0.1.0/tests/train/test_default_stateful.py +624 -0
  151. hybrax-0.1.0/tests/train/test_defaults.py +232 -0
  152. hybrax-0.1.0/tests/train/test_devices.py +133 -0
  153. hybrax-0.1.0/tests/train/test_e2e.py +100 -0
  154. hybrax-0.1.0/tests/train/test_event_eps.py +199 -0
  155. hybrax-0.1.0/tests/train/test_forward.py +2047 -0
  156. hybrax-0.1.0/tests/train/test_harness.py +2242 -0
  157. hybrax-0.1.0/tests/train/test_hek293_gallery_loss.py +188 -0
  158. hybrax-0.1.0/tests/train/test_initial_parameters.py +451 -0
  159. hybrax-0.1.0/tests/train/test_inspect.py +192 -0
  160. hybrax-0.1.0/tests/train/test_jump_ts.py +102 -0
  161. hybrax-0.1.0/tests/train/test_logging.py +257 -0
  162. hybrax-0.1.0/tests/train/test_loo.py +1762 -0
  163. hybrax-0.1.0/tests/train/test_loo_metrics.py +774 -0
  164. hybrax-0.1.0/tests/train/test_loo_metrics_divergence.py +88 -0
  165. hybrax-0.1.0/tests/train/test_model_api.py +973 -0
  166. hybrax-0.1.0/tests/train/test_model_load_predict.py +363 -0
  167. hybrax-0.1.0/tests/train/test_observe_latent_guard.py +83 -0
  168. hybrax-0.1.0/tests/train/test_output_grid.py +493 -0
  169. hybrax-0.1.0/tests/train/test_preflight.py +431 -0
  170. hybrax-0.1.0/tests/train/test_prepare.py +1041 -0
  171. hybrax-0.1.0/tests/train/test_run_config.py +494 -0
  172. hybrax-0.1.0/tests/train/test_runtime_artifact.py +1869 -0
  173. hybrax-0.1.0/tests/train/test_runtime_context.py +774 -0
  174. hybrax-0.1.0/tests/train/test_serialization.py +441 -0
  175. hybrax-0.1.0/tests/train/test_sharded_training.py +517 -0
  176. hybrax-0.1.0/tests/train/test_shared_reconstruction.py +773 -0
  177. hybrax-0.1.0/tests/train/test_solver_hook.py +151 -0
  178. hybrax-0.1.0/tests/train/test_stateful_convergence.py +56 -0
  179. hybrax-0.1.0/tests/train/test_stateful_grad.py +70 -0
  180. hybrax-0.1.0/tests/train/test_stateful_jump_invariance.py +24 -0
  181. hybrax-0.1.0/tests/train/test_stateful_loss_isolation.py +32 -0
  182. hybrax-0.1.0/tests/train/test_stateful_no_dilution.py +71 -0
  183. hybrax-0.1.0/tests/train/test_trainer.py +1559 -0
  184. hybrax-0.1.0/tests/train/test_training_data.py +825 -0
  185. hybrax-0.1.0/tests/train/test_wrapper.py +1710 -0
  186. hybrax-0.1.0/uv.lock +3720 -0
@@ -0,0 +1,44 @@
1
+ name: Publish to PyPI
2
+
3
+ on:
4
+ push:
5
+ tags:
6
+ - "v*"
7
+
8
+ jobs:
9
+ build:
10
+ runs-on: ubuntu-latest
11
+ permissions:
12
+ contents: read
13
+ steps:
14
+ - uses: actions/checkout@v4
15
+ with:
16
+ # setuptools-scm needs the tags to work out the version.
17
+ fetch-depth: 0
18
+
19
+ - uses: astral-sh/setup-uv@v6
20
+
21
+ - name: Build sdist and wheel
22
+ run: uv build
23
+
24
+ - uses: actions/upload-artifact@v4
25
+ with:
26
+ name: dist
27
+ path: dist/
28
+
29
+ publish:
30
+ needs: build
31
+ runs-on: ubuntu-latest
32
+ environment:
33
+ name: pypi
34
+ url: https://pypi.org/p/hybrax
35
+ permissions:
36
+ # Lets this job request an OIDC token for PyPI trusted publishing.
37
+ id-token: write
38
+ steps:
39
+ - uses: actions/download-artifact@v4
40
+ with:
41
+ name: dist
42
+ path: dist/
43
+
44
+ - uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,242 @@
1
+ # local dev files
2
+ *.override.*
3
+ dev/
4
+ .vscode/
5
+ reference/
6
+ AGENTS.md
7
+ CLAUDE.md
8
+
9
+ # docs build output (regenerated by docs/docs_rebuild.sh)
10
+ # docs/html/
11
+ docs/jupyter_execute/
12
+ docs/source/autoapi/
13
+ docs/source/_data/out/
14
+ docs/_scratch/
15
+
16
+ # hidden files
17
+ .*
18
+ # except for the below
19
+ !.github/
20
+ !.python-version
21
+ !.readthedocs.yaml
22
+
23
+ # Byte-compiled / optimized / DLL files
24
+ __pycache__/
25
+ *.py[codz]
26
+ *$py.class
27
+
28
+ # C extensions
29
+ *.so
30
+
31
+ # Distribution / packaging
32
+ .Python
33
+ build/
34
+ develop-eggs/
35
+ dist/
36
+ downloads/
37
+ eggs/
38
+ .eggs/
39
+ lib/
40
+ lib64/
41
+ parts/
42
+ sdist/
43
+ var/
44
+ wheels/
45
+ share/python-wheels/
46
+ *.egg-info/
47
+ .installed.cfg
48
+ *.egg
49
+ MANIFEST
50
+
51
+ # PyInstaller
52
+ # Usually these files are written by a python script from a template
53
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
54
+ *.manifest
55
+ *.spec
56
+
57
+ # Installer logs
58
+ pip-log.txt
59
+ pip-delete-this-directory.txt
60
+
61
+ # Unit test / coverage reports
62
+ htmlcov/
63
+ .tox/
64
+ .nox/
65
+ .coverage
66
+ .coverage.*
67
+ .cache
68
+ nosetests.xml
69
+ coverage.xml
70
+ *.cover
71
+ *.py.cover
72
+ .hypothesis/
73
+ .pytest_cache/
74
+ cover/
75
+
76
+ # Translations
77
+ *.mo
78
+ *.pot
79
+
80
+ # Django stuff:
81
+ *.log
82
+ local_settings.py
83
+ db.sqlite3
84
+ db.sqlite3-journal
85
+
86
+ # Flask stuff:
87
+ instance/
88
+ .webassets-cache
89
+
90
+ # Scrapy stuff:
91
+ .scrapy
92
+
93
+ # Sphinx documentation
94
+ docs/_build/
95
+
96
+ # PyBuilder
97
+ .pybuilder/
98
+ target/
99
+
100
+ # Jupyter Notebook
101
+ .ipynb_checkpoints
102
+
103
+ # IPython
104
+ profile_default/
105
+ ipython_config.py
106
+
107
+ # pyenv
108
+ # For a library or package, you might want to ignore these files since the code is
109
+ # intended to run in multiple environments; otherwise, check them in:
110
+ # .python-version
111
+
112
+ # pipenv
113
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
114
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
115
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
116
+ # install all needed dependencies.
117
+ #Pipfile.lock
118
+
119
+ # UV
120
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
121
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
122
+ # commonly ignored for libraries.
123
+ #uv.lock
124
+
125
+ # poetry
126
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
127
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
128
+ # commonly ignored for libraries.
129
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
130
+ #poetry.lock
131
+ #poetry.toml
132
+
133
+ # pdm
134
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
135
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
136
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
137
+ #pdm.lock
138
+ #pdm.toml
139
+ .pdm-python
140
+ .pdm-build/
141
+
142
+ # pixi
143
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
144
+ pixi.lock
145
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
146
+ # in the .venv directory. It is recommended not to include this directory in version control.
147
+ .pixi
148
+ # ignore .pixi.toml as well (since we're only using pixi for local dev)
149
+ pixi.toml
150
+
151
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
152
+ __pypackages__/
153
+
154
+ # Celery stuff
155
+ celerybeat-schedule
156
+ celerybeat.pid
157
+
158
+ # SageMath parsed files
159
+ *.sage.py
160
+
161
+ # Environments
162
+ .env
163
+ .envrc
164
+ .venv
165
+ env/
166
+ venv/
167
+ ENV/
168
+ env.bak/
169
+ venv.bak/
170
+
171
+ # Spyder project settings
172
+ .spyderproject
173
+ .spyproject
174
+
175
+ # Rope project settings
176
+ .ropeproject
177
+
178
+ # mkdocs documentation
179
+ /site
180
+
181
+ # mypy
182
+ .mypy_cache/
183
+ .dmypy.json
184
+ dmypy.json
185
+
186
+ # Pyre type checker
187
+ .pyre/
188
+
189
+ # pytype static type analyzer
190
+ .pytype/
191
+
192
+ # Cython debug symbols
193
+ cython_debug/
194
+
195
+ # PyCharm
196
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
197
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
198
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
199
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
200
+ #.idea/
201
+
202
+ # Abstra
203
+ # Abstra is an AI-powered process automation framework.
204
+ # Ignore directories containing user credentials, local state, and settings.
205
+ # Learn more at https://abstra.io/docs
206
+ .abstra/
207
+
208
+ # Visual Studio Code
209
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
210
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
211
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
212
+ # you could uncomment the following to ignore the entire vscode folder
213
+ # .vscode/
214
+
215
+ # Ruff stuff:
216
+ .ruff_cache/
217
+
218
+ # PyPI configuration file
219
+ .pypirc
220
+
221
+ # Cursor
222
+ # Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
223
+ # exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
224
+ # refer to https://docs.cursor.com/context/ignore-files
225
+ .cursorignore
226
+ .cursorindexingignore
227
+
228
+ # Marimo
229
+ marimo/_static/
230
+ marimo/_lsp/
231
+ __marimo__/
232
+
233
+ # Shitty Windows files
234
+ *:Zone.Identifier
235
+
236
+ # Local scratch directories
237
+ tmp/
238
+ **/tmp/
239
+
240
+ # MemPalace per-project files
241
+ mempalace.yaml
242
+ entities.json
@@ -0,0 +1 @@
1
+ 3.13.15
@@ -0,0 +1,12 @@
1
+ version: 2
2
+
3
+ build:
4
+ os: ubuntu-24.04
5
+ tools:
6
+ python: "3.13"
7
+ jobs:
8
+ build:
9
+ html:
10
+ - mkdir -p "$READTHEDOCS_OUTPUT/html"
11
+ - cp -a docs/html/. "$READTHEDOCS_OUTPUT/html/"
12
+
hybrax-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Mathias Gotsmy and Julian Libiseller-Egger
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,4 @@
1
+ # Keep the sdist small: built docs and example data are not needed to
2
+ # build or test the package.
3
+ prune docs
4
+ prune examples
hybrax-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,132 @@
1
+ Metadata-Version: 2.4
2
+ Name: hybrax
3
+ Version: 0.1.0
4
+ Summary: Bioprocess hybrid models and data management in JAX
5
+ Author: Mathias Gotsmy, Julian Libiseller-Egger
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/hybrax/hybrax
8
+ Project-URL: Repository, https://github.com/hybrax/hybrax
9
+ Project-URL: Documentation, https://hybrax.readthedocs.io
10
+ Keywords: bioprocess,fermentation,upstream modeling,hybrid modeling,jax,diffrax,neural ODE,UDE,universal differential equations,data structures
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Topic :: Scientific/Engineering
16
+ Requires-Python: >=3.12
17
+ Description-Content-Type: text/markdown
18
+ License-File: LICENSE
19
+ Requires-Dist: jax<0.10.1,>=0.4.30
20
+ Requires-Dist: jaxlib>=0.4.0
21
+ Requires-Dist: numpy>=1.26
22
+ Requires-Dist: pandas<3.1,>=3.0.5
23
+ Requires-Dist: equinox<0.14,>=0.13.8
24
+ Requires-Dist: scipy>=1.15.0
25
+ Requires-Dist: clarabel>=0.11
26
+ Requires-Dist: diffrax<0.8,>=0.7.2
27
+ Requires-Dist: sympy<1.15,>=1.14.0
28
+ Requires-Dist: ijson>=3.5.1
29
+ Requires-Dist: matplotlib>=3.8
30
+ Requires-Dist: optax<0.3,>=0.2.8
31
+ Requires-Dist: optimistix>=0.1.0
32
+ Requires-Dist: pydantic>=2.5.2
33
+ Provides-Extra: dev
34
+ Requires-Dist: pytest>=8.0; extra == "dev"
35
+ Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
36
+ Requires-Dist: pytest-xdist>=3.6; extra == "dev"
37
+ Requires-Dist: black>=22.0.0; extra == "dev"
38
+ Requires-Dist: flake8>=4.0.0; extra == "dev"
39
+ Requires-Dist: mypy>=0.950; extra == "dev"
40
+ Requires-Dist: ruff>=0.17.2; extra == "dev"
41
+ Requires-Dist: ipdb>=0.13.13; extra == "dev"
42
+ Requires-Dist: jupyter>=1.1.1; extra == "dev"
43
+ Requires-Dist: ipykernel>=7.2.0; extra == "dev"
44
+ Requires-Dist: openpyxl>=3.1.5; extra == "dev"
45
+ Provides-Extra: docs
46
+ Requires-Dist: sphinx>=7.2; extra == "docs"
47
+ Requires-Dist: furo>=2024.1; extra == "docs"
48
+ Requires-Dist: myst-nb>=1.1; extra == "docs"
49
+ Requires-Dist: linkify-it-py>=2.0; extra == "docs"
50
+ Requires-Dist: sphinx-autoapi>=3.1; extra == "docs"
51
+ Requires-Dist: sphinx-copybutton>=0.5; extra == "docs"
52
+ Requires-Dist: sphinx-design>=0.6; extra == "docs"
53
+ Requires-Dist: jupyter-cache>=1.0; extra == "docs"
54
+ Requires-Dist: ipykernel>=6.29; extra == "docs"
55
+ Requires-Dist: matplotlib==3.11.1; extra == "docs"
56
+ Dynamic: license-file
57
+
58
+ # Hybrax
59
+
60
+ Bioprocess hybrid models and data management in JAX.
61
+
62
+ Hybrid models, also known as universal differential equations (UDEs) or
63
+ grey-box models, embed learned components such as neural networks in a
64
+ mechanistic ODE.
65
+
66
+ Hybrax combines bioprocess data interchange and hybrid model training in one
67
+ package. `hybrax.format` provides the data model, validation, serialization,
68
+ and mechanistic right-hand side. `hybrax.train` provides preparation,
69
+ augmentation, training, prediction, and leave-one-out cross-validation.
70
+
71
+ Read the [documentation](https://hybrax.readthedocs.io) for tutorials, gallery
72
+ examples, and API reference.
73
+
74
+ ## Motivation
75
+
76
+ Bioprocess modeling research lacks standardized data formats, making it
77
+ difficult to compare modeling approaches across labs and publications.
78
+ `hybrax.format` provides a common hierarchical data structure with
79
+ built-in validation, serialization, and mechanistic modeling support.
80
+ Built on JAX and Equinox, it enables automatic differentiation through
81
+ ODE integration for gradient-based hybrid model training.
82
+
83
+ ## Installation
84
+
85
+ For local development, create the uv environment with the development extras:
86
+
87
+ ```bash
88
+ uv sync --extra dev
89
+ ```
90
+
91
+ Run tests with `uv run pytest`. Add `-n 12` to use 12 pytest-xdist workers.
92
+
93
+ ## Quick start
94
+
95
+ ```python
96
+ import hybrax.format as bp
97
+
98
+ # Load a process collection from JSON
99
+ collection = bp.serialization.load_process_collection("data.json")
100
+
101
+ # Explore the collection
102
+ bp.print_collection_structure(collection, verbosity=1)
103
+
104
+ # Access a specific process
105
+ process = collection.processes["run_1"]
106
+
107
+ # Validate data integrity
108
+ is_valid, messages = bp.validate_process(process)
109
+
110
+ # Inspect a single process
111
+ bp.print_process_structure(process, verbosity=3)
112
+
113
+ # Plot
114
+ fig = bp.plot_process(process)
115
+ ```
116
+
117
+ ## Modules
118
+
119
+ | Module | Description |
120
+ |--------|-------------|
121
+ | `hybrax.format.dataclasses` | Hierarchical data structures (`BioProcessCollection`, `BioProcess`, etc.) |
122
+ | `hybrax.format.time_series` | Time-series container with optional fitted spline coefficients (JAX pytree) |
123
+ | `hybrax.format.splines` | Pseudobatch transform, segmented spline fitting |
124
+ | `hybrax.format.mechanistic` | ODE right-hand side and control splines (integration lives in `hybrax.train`) |
125
+ | `hybrax.format.serialization` | JSON save/load for the full data hierarchy |
126
+ | `hybrax.format.validate` | Data integrity checks |
127
+ | `hybrax.format.inspect` | Text printing and matplotlib visualization |
128
+ | `hybrax.format.simulation` | Event bookkeeping for synthetic datasets |
129
+ | `hybrax.train` | Preparation, hybrid model training, prediction, and cross-validation |
130
+
131
+ See the [documentation](https://hybrax.readthedocs.io) for the data structure
132
+ and the training workflow.
hybrax-0.1.0/README.md ADDED
@@ -0,0 +1,75 @@
1
+ # Hybrax
2
+
3
+ Bioprocess hybrid models and data management in JAX.
4
+
5
+ Hybrid models, also known as universal differential equations (UDEs) or
6
+ grey-box models, embed learned components such as neural networks in a
7
+ mechanistic ODE.
8
+
9
+ Hybrax combines bioprocess data interchange and hybrid model training in one
10
+ package. `hybrax.format` provides the data model, validation, serialization,
11
+ and mechanistic right-hand side. `hybrax.train` provides preparation,
12
+ augmentation, training, prediction, and leave-one-out cross-validation.
13
+
14
+ Read the [documentation](https://hybrax.readthedocs.io) for tutorials, gallery
15
+ examples, and API reference.
16
+
17
+ ## Motivation
18
+
19
+ Bioprocess modeling research lacks standardized data formats, making it
20
+ difficult to compare modeling approaches across labs and publications.
21
+ `hybrax.format` provides a common hierarchical data structure with
22
+ built-in validation, serialization, and mechanistic modeling support.
23
+ Built on JAX and Equinox, it enables automatic differentiation through
24
+ ODE integration for gradient-based hybrid model training.
25
+
26
+ ## Installation
27
+
28
+ For local development, create the uv environment with the development extras:
29
+
30
+ ```bash
31
+ uv sync --extra dev
32
+ ```
33
+
34
+ Run tests with `uv run pytest`. Add `-n 12` to use 12 pytest-xdist workers.
35
+
36
+ ## Quick start
37
+
38
+ ```python
39
+ import hybrax.format as bp
40
+
41
+ # Load a process collection from JSON
42
+ collection = bp.serialization.load_process_collection("data.json")
43
+
44
+ # Explore the collection
45
+ bp.print_collection_structure(collection, verbosity=1)
46
+
47
+ # Access a specific process
48
+ process = collection.processes["run_1"]
49
+
50
+ # Validate data integrity
51
+ is_valid, messages = bp.validate_process(process)
52
+
53
+ # Inspect a single process
54
+ bp.print_process_structure(process, verbosity=3)
55
+
56
+ # Plot
57
+ fig = bp.plot_process(process)
58
+ ```
59
+
60
+ ## Modules
61
+
62
+ | Module | Description |
63
+ |--------|-------------|
64
+ | `hybrax.format.dataclasses` | Hierarchical data structures (`BioProcessCollection`, `BioProcess`, etc.) |
65
+ | `hybrax.format.time_series` | Time-series container with optional fitted spline coefficients (JAX pytree) |
66
+ | `hybrax.format.splines` | Pseudobatch transform, segmented spline fitting |
67
+ | `hybrax.format.mechanistic` | ODE right-hand side and control splines (integration lives in `hybrax.train`) |
68
+ | `hybrax.format.serialization` | JSON save/load for the full data hierarchy |
69
+ | `hybrax.format.validate` | Data integrity checks |
70
+ | `hybrax.format.inspect` | Text printing and matplotlib visualization |
71
+ | `hybrax.format.simulation` | Event bookkeeping for synthetic datasets |
72
+ | `hybrax.train` | Preparation, hybrid model training, prediction, and cross-validation |
73
+
74
+ See the [documentation](https://hybrax.readthedocs.io) for the data structure
75
+ and the training workflow.
@@ -0,0 +1,118 @@
1
+ [build-system]
2
+ requires = ["setuptools>=64", "setuptools-scm>=8", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "hybrax"
7
+ dynamic = ["version"]
8
+ description = "Bioprocess hybrid models and data management in JAX"
9
+ readme = "README.md"
10
+ requires-python = ">=3.12"
11
+ license = "MIT"
12
+ authors = [{ name = "Mathias Gotsmy" }, { name = "Julian Libiseller-Egger" }]
13
+ keywords = [
14
+ "bioprocess",
15
+ "fermentation",
16
+ "upstream modeling",
17
+ "hybrid modeling",
18
+ "jax",
19
+ "diffrax",
20
+ "neural ODE",
21
+ "UDE",
22
+ "universal differential equations",
23
+ "data structures",
24
+ ]
25
+ classifiers = [
26
+ "Development Status :: 3 - Alpha",
27
+ "Intended Audience :: Science/Research",
28
+ "Programming Language :: Python :: 3",
29
+ "Programming Language :: Python :: 3.12",
30
+ "Topic :: Scientific/Engineering",
31
+ ]
32
+
33
+ dependencies = [
34
+ "jax>=0.4.30,<0.10.1",
35
+ "jaxlib>=0.4.0",
36
+ "numpy>=1.26",
37
+ "pandas>=3.0.5,<3.1",
38
+ "equinox>=0.13.8,<0.14",
39
+ "scipy>=1.15.0",
40
+ "clarabel>=0.11",
41
+ "diffrax>=0.7.2,<0.8",
42
+ "sympy>=1.14.0,<1.15",
43
+ "ijson>=3.5.1",
44
+ "matplotlib>=3.8",
45
+ "optax>=0.2.8,<0.3",
46
+ "optimistix>=0.1.0",
47
+ "pydantic>=2.5.2",
48
+ ]
49
+
50
+ [project.scripts]
51
+ hybrax = "hybrax.train.cli:main"
52
+
53
+ [project.optional-dependencies]
54
+ dev = [
55
+ "pytest>=8.0",
56
+ "pytest-cov>=3.0.0",
57
+ "pytest-xdist>=3.6",
58
+ "black>=22.0.0",
59
+ "flake8>=4.0.0",
60
+ "mypy>=0.950",
61
+ "ruff>=0.17.2",
62
+ "ipdb>=0.13.13",
63
+ "jupyter>=1.1.1",
64
+ "ipykernel>=7.2.0",
65
+ "openpyxl>=3.1.5",
66
+ ]
67
+ docs = [
68
+ "sphinx>=7.2",
69
+ "furo>=2024.1",
70
+ "myst-nb>=1.1",
71
+ "linkify-it-py>=2.0",
72
+ "sphinx-autoapi>=3.1",
73
+ "sphinx-copybutton>=0.5",
74
+ "sphinx-design>=0.6",
75
+ "jupyter-cache>=1.0",
76
+ "ipykernel>=6.29",
77
+ # Pinned so the docs plots look the same on every rebuild.
78
+ "matplotlib==3.11.1",
79
+ ]
80
+
81
+ [project.urls]
82
+ Homepage = "https://github.com/hybrax/hybrax"
83
+ Repository = "https://github.com/hybrax/hybrax"
84
+ Documentation = "https://hybrax.readthedocs.io"
85
+
86
+ [tool.ruff.lint]
87
+ extend-select = ["E501"]
88
+ # Schema validation uses ValueError, including Pydantic field validators.
89
+ ignore = ["TRY004"]
90
+
91
+ [tool.ruff.format]
92
+ exclude = ["*.md"]
93
+
94
+ [tool.setuptools.packages.find]
95
+ where = ["src"]
96
+
97
+ # Version comes from the latest git tag (e.g. v0.2.0 -> 0.2.0).
98
+ [tool.setuptools_scm]
99
+ fallback_version = "0.0.0"
100
+
101
+ # Rebuild the editable install when the commit or tags change, so the
102
+ # reported version stays current.
103
+ [tool.uv]
104
+ cache-keys = [{ file = "pyproject.toml" }, { git = { commit = true, tags = true } }]
105
+
106
+ [tool.pytest.ini_options]
107
+ testpaths = ["tests"]
108
+ markers = ["integration: end-to-end tests that exercise the training pipeline"]
109
+ filterwarnings = [
110
+ "ignore:jax.core.mapped_aval is deprecated:DeprecationWarning",
111
+ "ignore:jax.core.unmapped_aval is deprecated:DeprecationWarning",
112
+ ]
113
+
114
+ [tool.mypy]
115
+ python_version = "3.12"
116
+ warn_return_any = true
117
+ warn_unused_configs = true
118
+ disallow_untyped_defs = false
hybrax-0.1.0/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+