flow-gym-suite 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 (87) hide show
  1. flow_gym_suite-0.1.0/LICENSE +21 -0
  2. flow_gym_suite-0.1.0/PKG-INFO +233 -0
  3. flow_gym_suite-0.1.0/README.md +178 -0
  4. flow_gym_suite-0.1.0/pyproject.toml +210 -0
  5. flow_gym_suite-0.1.0/setup.cfg +4 -0
  6. flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/PKG-INFO +233 -0
  7. flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/SOURCES.txt +85 -0
  8. flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/dependency_links.txt +1 -0
  9. flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/requires.txt +15 -0
  10. flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/top_level.txt +1 -0
  11. flow_gym_suite-0.1.0/src/flowgym/__init__.py +63 -0
  12. flow_gym_suite-0.1.0/src/flowgym/common/__init__.py +1 -0
  13. flow_gym_suite-0.1.0/src/flowgym/common/base/__init__.py +10 -0
  14. flow_gym_suite-0.1.0/src/flowgym/common/base/estimator.py +461 -0
  15. flow_gym_suite-0.1.0/src/flowgym/common/base/trainable_states.py +140 -0
  16. flow_gym_suite-0.1.0/src/flowgym/common/evaluation.py +200 -0
  17. flow_gym_suite-0.1.0/src/flowgym/common/filters.py +118 -0
  18. flow_gym_suite-0.1.0/src/flowgym/common/median.py +1286 -0
  19. flow_gym_suite-0.1.0/src/flowgym/common/preprocess.py +356 -0
  20. flow_gym_suite-0.1.0/src/flowgym/density/__init__.py +1 -0
  21. flow_gym_suite-0.1.0/src/flowgym/density/nn.py +129 -0
  22. flow_gym_suite-0.1.0/src/flowgym/density/simple.py +60 -0
  23. flow_gym_suite-0.1.0/src/flowgym/environment/__init__.py +1 -0
  24. flow_gym_suite-0.1.0/src/flowgym/environment/fluid_env.py +164 -0
  25. flow_gym_suite-0.1.0/src/flowgym/flow/__init__.py +1 -0
  26. flow_gym_suite-0.1.0/src/flowgym/flow/base.py +98 -0
  27. flow_gym_suite-0.1.0/src/flowgym/flow/consensus/__init__.py +7 -0
  28. flow_gym_suite-0.1.0/src/flowgym/flow/consensus/admm.py +271 -0
  29. flow_gym_suite-0.1.0/src/flowgym/flow/consensus/consensus.py +385 -0
  30. flow_gym_suite-0.1.0/src/flowgym/flow/consensus/consensus_algorithms.py +417 -0
  31. flow_gym_suite-0.1.0/src/flowgym/flow/consensus/objectives.py +426 -0
  32. flow_gym_suite-0.1.0/src/flowgym/flow/consensus/regularizers.py +106 -0
  33. flow_gym_suite-0.1.0/src/flowgym/flow/consensus/solvers.py +213 -0
  34. flow_gym_suite-0.1.0/src/flowgym/flow/deepflow.py +48 -0
  35. flow_gym_suite-0.1.0/src/flowgym/flow/dis/__init__.py +12 -0
  36. flow_gym_suite-0.1.0/src/flowgym/flow/dis/dis.py +174 -0
  37. flow_gym_suite-0.1.0/src/flowgym/flow/dis/dis_jax.py +210 -0
  38. flow_gym_suite-0.1.0/src/flowgym/flow/dis/process.py +1091 -0
  39. flow_gym_suite-0.1.0/src/flowgym/flow/dummy.py +42 -0
  40. flow_gym_suite-0.1.0/src/flowgym/flow/farneback.py +92 -0
  41. flow_gym_suite-0.1.0/src/flowgym/flow/hornschunck.py +135 -0
  42. flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/__init__.py +12 -0
  43. flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/openpiv.py +117 -0
  44. flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/openpiv_jax.py +215 -0
  45. flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/process.py +375 -0
  46. flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/__init__.py +270 -0
  47. flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/data_interpolation.py +114 -0
  48. flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/data_smoothing.py +155 -0
  49. flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/data_validation.py +283 -0
  50. flow_gym_suite-0.1.0/src/flowgym/flow/process.py +215 -0
  51. flow_gym_suite-0.1.0/src/flowgym/flow/raft/__init__.py +13 -0
  52. flow_gym_suite-0.1.0/src/flowgym/flow/raft/process.py +350 -0
  53. flow_gym_suite-0.1.0/src/flowgym/flow/raft/raft_jax.py +417 -0
  54. flow_gym_suite-0.1.0/src/flowgym/flow/raft/raft_piv_pytorch.py +243 -0
  55. flow_gym_suite-0.1.0/src/flowgym/flow/utils.py +182 -0
  56. flow_gym_suite-0.1.0/src/flowgym/make.py +234 -0
  57. flow_gym_suite-0.1.0/src/flowgym/nn/__init__.py +1 -0
  58. flow_gym_suite-0.1.0/src/flowgym/nn/blocks.py +437 -0
  59. flow_gym_suite-0.1.0/src/flowgym/nn/cnn.py +90 -0
  60. flow_gym_suite-0.1.0/src/flowgym/nn/raft_model.py +99 -0
  61. flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/__init__.py +6 -0
  62. flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/corr_block.py +55 -0
  63. flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/flowNetsRAFT.py +203 -0
  64. flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/submodules_RAFT_GRU.py +165 -0
  65. flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/submodules_RAFT_extractor.py +272 -0
  66. flow_gym_suite-0.1.0/src/flowgym/training/__init__.py +1 -0
  67. flow_gym_suite-0.1.0/src/flowgym/training/optimizer.py +99 -0
  68. flow_gym_suite-0.1.0/src/flowgym/training/schedules.py +84 -0
  69. flow_gym_suite-0.1.0/src/flowgym/training/utils.py +176 -0
  70. flow_gym_suite-0.1.0/src/flowgym/types.py +7 -0
  71. flow_gym_suite-0.1.0/src/flowgym/utils.py +441 -0
  72. flow_gym_suite-0.1.0/tests/test_data_interpolation.py +175 -0
  73. flow_gym_suite-0.1.0/tests/test_data_smoothing.py +191 -0
  74. flow_gym_suite-0.1.0/tests/test_data_validation.py +381 -0
  75. flow_gym_suite-0.1.0/tests/test_density_estimate_simple.py +97 -0
  76. flow_gym_suite-0.1.0/tests/test_dis_jax.py +1068 -0
  77. flow_gym_suite-0.1.0/tests/test_estimation_stats.py +99 -0
  78. flow_gym_suite-0.1.0/tests/test_filters.py +107 -0
  79. flow_gym_suite-0.1.0/tests/test_flow_estimate_base.py +49 -0
  80. flow_gym_suite-0.1.0/tests/test_flow_estimate_dis.py +117 -0
  81. flow_gym_suite-0.1.0/tests/test_median.py +158 -0
  82. flow_gym_suite-0.1.0/tests/test_openpiv_jax.py +275 -0
  83. flow_gym_suite-0.1.0/tests/test_postprocess_config.py +496 -0
  84. flow_gym_suite-0.1.0/tests/test_preprocess.py +156 -0
  85. flow_gym_suite-0.1.0/tests/test_preprocess_config.py +216 -0
  86. flow_gym_suite-0.1.0/tests/test_trainable_state.py +267 -0
  87. flow_gym_suite-0.1.0/tests/test_training_utils.py +137 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 Antonio Terpin
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,233 @@
1
+ Metadata-Version: 2.4
2
+ Name: flow-gym-suite
3
+ Version: 0.1.0
4
+ Summary: Flow Gym is a toolkit for research and deployment in flow quantification inspired by OpenAI Gym and Stable-Baselines3.
5
+ Author-email: Antonio Terpin <aterpin@ethz.ch>, Francesco Banelli <fbanelli@ethz.ch>, Alan Bonomi <abonomi@ethz.ch>
6
+ License: MIT License
7
+
8
+ Copyright (c) 2025 Antonio Terpin
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+
28
+ Project-URL: Homepage, https://github.com/antonioterpin/flowgym
29
+ Project-URL: Repository, https://github.com/antonioterpin/flowgym
30
+ Project-URL: Documentation, https://github.com/antonioterpin/flowgym#readme
31
+ Project-URL: Issues, https://github.com/antonioterpin/flowgym/issues
32
+ Classifier: Development Status :: 4 - Beta
33
+ Classifier: Intended Audience :: Science/Research
34
+ Classifier: Intended Audience :: Developers
35
+ Classifier: License :: OSI Approved :: MIT License
36
+ Classifier: Operating System :: OS Independent
37
+ Classifier: Programming Language :: Python :: 3
38
+ Requires-Python: >=3.10
39
+ Description-Content-Type: text/markdown
40
+ License-File: LICENSE
41
+ Requires-Dist: jax>=0.6.0
42
+ Requires-Dist: flax>=0.10.5
43
+ Requires-Dist: ipywidgets
44
+ Requires-Dist: robo-goggles>=0.1.6
45
+ Requires-Dist: pip>=25.3
46
+ Provides-Extra: other-methods
47
+ Requires-Dist: pyoptflow; extra == "other-methods"
48
+ Requires-Dist: openpiv; extra == "other-methods"
49
+ Requires-Dist: opencv-python-headless; extra == "other-methods"
50
+ Requires-Dist: torch>=1.13.1; extra == "other-methods"
51
+ Requires-Dist: torchvision>=0.14.1; extra == "other-methods"
52
+ Provides-Extra: cuda12
53
+ Requires-Dist: jax[cuda12]>=0.6.0; extra == "cuda12"
54
+ Dynamic: license-file
55
+
56
+ # 🌊 Flow Gym : A Research Toolkit for Fluid Flow Estimation
57
+
58
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
59
+ [![GitHub stars](https://img.shields.io/github/stars/antonioterpin/flowgym?style=social)](https://github.com/antonioterpin/flowgym/stargazers)
60
+ [![PyPI version](https://img.shields.io/pypi/v/flow-gym.svg)](https://pypi.org/project/flow-gym)
61
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
62
+ [![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
63
+
64
+ Flow Gym is a **toolkit for research and deployment in flow quantification**, inspired by **OpenAI Gym** and **Stable-Baselines3**.
65
+ It leverages **SynthPix** for synthetic image generation and provides a **unified interface** for testing, training, and deploying learning-based algorithms that estimate fluid flow from sequences of tracer-particle images.
66
+ The framework also includes a growing collection of **integrated algorithms** and **stable JAX re-implementations** of state-of-the-art methods.
67
+
68
+ ![The Flow Gym environment](https://raw.githubusercontent.com/antonioterpin/flowgym/main/docs/flowgym.jpg)
69
+
70
+ ## ✨ Features
71
+
72
+ - **Stable and reproducible environments**
73
+ Flow Gym provides **FluidEnv**, a unified environment for both supervised and reinforcement learning. It abstracts away infrastructure details and ensures consistent training and evaluation across experiments.
74
+
75
+ - **Unified estimator framework**
76
+ The **Estimator** API makes it easy to integrate, compare, and benchmark different flow estimation algorithms. It also supports extensions beyond velocity estimation — for example, pressure or tracer density — enabling ablation studies and cross-domain applications in fluid dynamics and control.
77
+
78
+ - **Ready for real-world deployment**
79
+ Flow Gym includes efficient **JAX implementations** of state-of-the-art flow estimation methods, optimized for real-time performance and suitable for **closed-loop control** or other **fluid-sensing** applications.
80
+
81
+ ## Getting Started 🚀
82
+
83
+ ```bash
84
+ # Basic installation
85
+ uv add flow-gym-suite # or pip install flow-gym-suite
86
+
87
+ # With CUDA12 support
88
+ uv add "flow-gym-suite[cuda12]"
89
+ ```
90
+
91
+ If you also want to use Flow Gym for training an evaluation, please clone the repository and install the development dependencies as outlined in the [CONTRIBUTING.md](./CONTRIBUTING.md) page.
92
+ There you can also find help if you have any issues installing Flow Gym.
93
+
94
+ ### Basic usage / Deployment 🧭
95
+
96
+ ```python
97
+ from flowgym.make import make_estimator
98
+ # Define the config (or load from YAML)
99
+ model_config = {
100
+ "estimator": "dis_jax",
101
+ "estimate_type": "flow",
102
+ "config": {"jit": True, ..., },
103
+ }
104
+
105
+ # Create the estimator and associated functions
106
+ trained_state, create_state_fn, compute_estimate_fn, model = make_estimator(
107
+ model_config=model_config,
108
+ image_shape=image_shape
109
+ )
110
+
111
+ # Compute an estimate for an image pair
112
+ est_state = create_state_fn(prev)
113
+ new_est_state, metrics = compute_estimate_fn(curr, est_state)
114
+ ```
115
+
116
+ ### Training a model 🧠
117
+
118
+ ```python
119
+ from flowgym.environment.fluid_env import FluidEnv
120
+
121
+ # Create the environment
122
+ env, env_state = FluidEnv.make(env_config)
123
+
124
+ for episode in range(num_episodes):
125
+ obs, env_state, done = env.reset(env_state)
126
+ est_state = create_state_fn(obs)
127
+ train_step_fn = model.create_train_step()
128
+
129
+ while not done.any():
130
+ # Estimator forward pass
131
+ est_state, metrics = compute_estimate_fn(obs, est_state, train_state)
132
+
133
+ # Extract the action (the last estimate)
134
+ action = estimation_state["estimates"][:, -1]
135
+ # Environment step
136
+ obs, env_state, reward, done = env.step(env_state, action)
137
+
138
+ # Optimization step
139
+ loss, train_state, _ = train_step_fn(
140
+ est_state, train_state, reward
141
+ )
142
+ ```
143
+
144
+ ### Implementing a custom estimator 🧩
145
+
146
+ ```python
147
+ from flowgym.base import Estimator
148
+ class MyEstimator(Estimator):
149
+ def init(self, rng, observation_space, action_space):
150
+ # Initialize model parameters and state
151
+ pass
152
+
153
+ def _estimate(self, images, state, trainable_state, extras):
154
+ # Define the forward pass of the model
155
+ return new_state, extras, metrics
156
+
157
+ def create_trainable_state(self, dummy_input, key):
158
+ # Create the trainable state if needed (e.g., network structure, optimizer, etc.)
159
+ pass
160
+
161
+ def create_train_step(self):
162
+ def train_step(self, estimation_state, trainable_state, rewards):
163
+ # Define the training step (loss computation and optimization)
164
+ return loss, new_trainable_state, metrics
165
+ ```
166
+
167
+ ## Custom Preprocessing 🧼
168
+
169
+ ```python
170
+ # flowgym/common/preprocess.py
171
+ def custom_preprocess_validate_params(params: Dict[str, Any]) -> None:
172
+ # Validate custom pre-processing parameters
173
+ pass
174
+
175
+ def custom_preprocess(
176
+ images: jnp.ndarray,
177
+ params: Dict[str, Any]
178
+ ) -> jnp.ndarray:
179
+ # Implement custom pre-processing logic
180
+ return processed_images
181
+
182
+ __all__ = [
183
+ "custom_preprocess_validate_params",
184
+ "custom_preprocess",
185
+ ]
186
+ ```
187
+
188
+ ## Custom Postprocessing 🎨
189
+
190
+ ```python
191
+ # flowgym/flow/postprocess/__init__.py
192
+ def custom_flow_postprocess_validate_params(params: Dict[str, Any]) -> None:
193
+ # Validate custom post-processing parameters
194
+ pass
195
+
196
+ def custom_flow_postprocess(
197
+ estimates: jnp.ndarray,
198
+ params: Dict[str, Any]
199
+ ) -> jnp.ndarray:
200
+ # Implement custom flow post-processing logic
201
+ return processed_estimates
202
+
203
+ __all__ = [
204
+ "custom_flow_postprocess_validate_params",
205
+ "custom_flow_postprocess",
206
+ ]
207
+ ```
208
+
209
+ ## Examples 📚
210
+
211
+ For more examples, please check our [training](https://raw.githubusercontent.com/antonioterpin/flowgym/main/src/flowgym/train.py) and [evaluation](https://raw.githubusercontent.com/antonioterpin/flowgym/main/src/flowgym/eval.py) scripts.
212
+
213
+ ## Contributing 🤗
214
+
215
+ Contributions are more than welcome! 🙏 Please check out our [how to contribute page](https://raw.githubusercontent.com/antonioterpin/flowgym/main/CONTRIBUTING.md), and feel free to open an issue for problems and feature requests⚠️.
216
+
217
+ If you have developed a new method integrating with Flow Gym, please do open a PR! New methods should come with reproducible experiments showcasing the features. We collect these experiments in the folder `experiments/your-method`.
218
+
219
+ ## Citation 📈
220
+
221
+ If you use this code in your research, please cite our paper:
222
+
223
+ ```bibtex
224
+ @article{banelli2025flowgym,
225
+ title={Flow Gym},
226
+ author={Banelli, Francesco and Terpin, Antonio and Bonomi, Alan and D'Andrea, Raffaello},
227
+ year={2025}
228
+ }
229
+ ```
230
+
231
+ ## License 📝
232
+
233
+ This project is licensed under the MIT License - see the [LICENSE](https://raw.githubusercontent.com/antonioterpin/flowgym/main/LICENSE) file for details.
@@ -0,0 +1,178 @@
1
+ # 🌊 Flow Gym : A Research Toolkit for Fluid Flow Estimation
2
+
3
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
4
+ [![GitHub stars](https://img.shields.io/github/stars/antonioterpin/flowgym?style=social)](https://github.com/antonioterpin/flowgym/stargazers)
5
+ [![PyPI version](https://img.shields.io/pypi/v/flow-gym.svg)](https://pypi.org/project/flow-gym)
6
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
7
+ [![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
8
+
9
+ Flow Gym is a **toolkit for research and deployment in flow quantification**, inspired by **OpenAI Gym** and **Stable-Baselines3**.
10
+ It leverages **SynthPix** for synthetic image generation and provides a **unified interface** for testing, training, and deploying learning-based algorithms that estimate fluid flow from sequences of tracer-particle images.
11
+ The framework also includes a growing collection of **integrated algorithms** and **stable JAX re-implementations** of state-of-the-art methods.
12
+
13
+ ![The Flow Gym environment](https://raw.githubusercontent.com/antonioterpin/flowgym/main/docs/flowgym.jpg)
14
+
15
+ ## ✨ Features
16
+
17
+ - **Stable and reproducible environments**
18
+ Flow Gym provides **FluidEnv**, a unified environment for both supervised and reinforcement learning. It abstracts away infrastructure details and ensures consistent training and evaluation across experiments.
19
+
20
+ - **Unified estimator framework**
21
+ The **Estimator** API makes it easy to integrate, compare, and benchmark different flow estimation algorithms. It also supports extensions beyond velocity estimation — for example, pressure or tracer density — enabling ablation studies and cross-domain applications in fluid dynamics and control.
22
+
23
+ - **Ready for real-world deployment**
24
+ Flow Gym includes efficient **JAX implementations** of state-of-the-art flow estimation methods, optimized for real-time performance and suitable for **closed-loop control** or other **fluid-sensing** applications.
25
+
26
+ ## Getting Started 🚀
27
+
28
+ ```bash
29
+ # Basic installation
30
+ uv add flow-gym-suite # or pip install flow-gym-suite
31
+
32
+ # With CUDA12 support
33
+ uv add "flow-gym-suite[cuda12]"
34
+ ```
35
+
36
+ If you also want to use Flow Gym for training an evaluation, please clone the repository and install the development dependencies as outlined in the [CONTRIBUTING.md](./CONTRIBUTING.md) page.
37
+ There you can also find help if you have any issues installing Flow Gym.
38
+
39
+ ### Basic usage / Deployment 🧭
40
+
41
+ ```python
42
+ from flowgym.make import make_estimator
43
+ # Define the config (or load from YAML)
44
+ model_config = {
45
+ "estimator": "dis_jax",
46
+ "estimate_type": "flow",
47
+ "config": {"jit": True, ..., },
48
+ }
49
+
50
+ # Create the estimator and associated functions
51
+ trained_state, create_state_fn, compute_estimate_fn, model = make_estimator(
52
+ model_config=model_config,
53
+ image_shape=image_shape
54
+ )
55
+
56
+ # Compute an estimate for an image pair
57
+ est_state = create_state_fn(prev)
58
+ new_est_state, metrics = compute_estimate_fn(curr, est_state)
59
+ ```
60
+
61
+ ### Training a model 🧠
62
+
63
+ ```python
64
+ from flowgym.environment.fluid_env import FluidEnv
65
+
66
+ # Create the environment
67
+ env, env_state = FluidEnv.make(env_config)
68
+
69
+ for episode in range(num_episodes):
70
+ obs, env_state, done = env.reset(env_state)
71
+ est_state = create_state_fn(obs)
72
+ train_step_fn = model.create_train_step()
73
+
74
+ while not done.any():
75
+ # Estimator forward pass
76
+ est_state, metrics = compute_estimate_fn(obs, est_state, train_state)
77
+
78
+ # Extract the action (the last estimate)
79
+ action = estimation_state["estimates"][:, -1]
80
+ # Environment step
81
+ obs, env_state, reward, done = env.step(env_state, action)
82
+
83
+ # Optimization step
84
+ loss, train_state, _ = train_step_fn(
85
+ est_state, train_state, reward
86
+ )
87
+ ```
88
+
89
+ ### Implementing a custom estimator 🧩
90
+
91
+ ```python
92
+ from flowgym.base import Estimator
93
+ class MyEstimator(Estimator):
94
+ def init(self, rng, observation_space, action_space):
95
+ # Initialize model parameters and state
96
+ pass
97
+
98
+ def _estimate(self, images, state, trainable_state, extras):
99
+ # Define the forward pass of the model
100
+ return new_state, extras, metrics
101
+
102
+ def create_trainable_state(self, dummy_input, key):
103
+ # Create the trainable state if needed (e.g., network structure, optimizer, etc.)
104
+ pass
105
+
106
+ def create_train_step(self):
107
+ def train_step(self, estimation_state, trainable_state, rewards):
108
+ # Define the training step (loss computation and optimization)
109
+ return loss, new_trainable_state, metrics
110
+ ```
111
+
112
+ ## Custom Preprocessing 🧼
113
+
114
+ ```python
115
+ # flowgym/common/preprocess.py
116
+ def custom_preprocess_validate_params(params: Dict[str, Any]) -> None:
117
+ # Validate custom pre-processing parameters
118
+ pass
119
+
120
+ def custom_preprocess(
121
+ images: jnp.ndarray,
122
+ params: Dict[str, Any]
123
+ ) -> jnp.ndarray:
124
+ # Implement custom pre-processing logic
125
+ return processed_images
126
+
127
+ __all__ = [
128
+ "custom_preprocess_validate_params",
129
+ "custom_preprocess",
130
+ ]
131
+ ```
132
+
133
+ ## Custom Postprocessing 🎨
134
+
135
+ ```python
136
+ # flowgym/flow/postprocess/__init__.py
137
+ def custom_flow_postprocess_validate_params(params: Dict[str, Any]) -> None:
138
+ # Validate custom post-processing parameters
139
+ pass
140
+
141
+ def custom_flow_postprocess(
142
+ estimates: jnp.ndarray,
143
+ params: Dict[str, Any]
144
+ ) -> jnp.ndarray:
145
+ # Implement custom flow post-processing logic
146
+ return processed_estimates
147
+
148
+ __all__ = [
149
+ "custom_flow_postprocess_validate_params",
150
+ "custom_flow_postprocess",
151
+ ]
152
+ ```
153
+
154
+ ## Examples 📚
155
+
156
+ For more examples, please check our [training](https://raw.githubusercontent.com/antonioterpin/flowgym/main/src/flowgym/train.py) and [evaluation](https://raw.githubusercontent.com/antonioterpin/flowgym/main/src/flowgym/eval.py) scripts.
157
+
158
+ ## Contributing 🤗
159
+
160
+ Contributions are more than welcome! 🙏 Please check out our [how to contribute page](https://raw.githubusercontent.com/antonioterpin/flowgym/main/CONTRIBUTING.md), and feel free to open an issue for problems and feature requests⚠️.
161
+
162
+ If you have developed a new method integrating with Flow Gym, please do open a PR! New methods should come with reproducible experiments showcasing the features. We collect these experiments in the folder `experiments/your-method`.
163
+
164
+ ## Citation 📈
165
+
166
+ If you use this code in your research, please cite our paper:
167
+
168
+ ```bibtex
169
+ @article{banelli2025flowgym,
170
+ title={Flow Gym},
171
+ author={Banelli, Francesco and Terpin, Antonio and Bonomi, Alan and D'Andrea, Raffaello},
172
+ year={2025}
173
+ }
174
+ ```
175
+
176
+ ## License 📝
177
+
178
+ This project is licensed under the MIT License - see the [LICENSE](https://raw.githubusercontent.com/antonioterpin/flowgym/main/LICENSE) file for details.
@@ -0,0 +1,210 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "flow-gym-suite"
7
+ version = "0.1.0"
8
+ description = "Flow Gym is a toolkit for research and deployment in flow quantification inspired by OpenAI Gym and Stable-Baselines3."
9
+ readme = "README.md"
10
+ license = { file = "LICENSE" }
11
+ authors = [
12
+ { name = "Antonio Terpin", email = "aterpin@ethz.ch" },
13
+ { name = "Francesco Banelli", email = "fbanelli@ethz.ch" },
14
+ { name = "Alan Bonomi", email = "abonomi@ethz.ch" }
15
+ ]
16
+ classifiers = [
17
+ "Development Status :: 4 - Beta",
18
+ "Intended Audience :: Science/Research",
19
+ "Intended Audience :: Developers",
20
+ "License :: OSI Approved :: MIT License",
21
+ "Operating System :: OS Independent",
22
+ "Programming Language :: Python :: 3",
23
+ ]
24
+ requires-python = ">=3.10"
25
+
26
+ # Core dependencies
27
+ dependencies = [
28
+ "jax>=0.6.0",
29
+ "flax>=0.10.5",
30
+ "ipywidgets",
31
+ "robo-goggles>=0.1.6",
32
+ "pip>=25.3",
33
+ ]
34
+
35
+ [project.urls]
36
+ Homepage = "https://github.com/antonioterpin/flowgym"
37
+ Repository = "https://github.com/antonioterpin/flowgym"
38
+ Documentation = "https://github.com/antonioterpin/flowgym#readme"
39
+ Issues = "https://github.com/antonioterpin/flowgym/issues"
40
+
41
+ [tool.setuptools.packages.find]
42
+ where = ["src"]
43
+ include = ["flowgym*"]
44
+ exclude = ["tests*", "docs*"]
45
+
46
+ # Dependency groups
47
+ [dependency-groups]
48
+ dev = [
49
+ "synthpix",
50
+ "pre_commit==4.0.1",
51
+ "snowballstemmer",
52
+ "pytest==7.4.4",
53
+ "wandb",
54
+ "scikit-image>=0.20.0",
55
+ ]
56
+
57
+ synthpix-cuda12 = [
58
+ "synthpix[cuda12]",
59
+ ]
60
+
61
+ dev-cuda12 = [
62
+ {include-group = "dev"},
63
+ {include-group = "synthpix-cuda12"},
64
+ ]
65
+
66
+ docs = [
67
+ "Sphinx==7.4.7",
68
+ "sphinx-copybutton==0.5.2",
69
+ "sphinx-rtd-theme==2.0.0",
70
+ "sphinx-tabs==3.4.7",
71
+ "sphinx-togglebutton==0.3.2",
72
+ "sphinxcontrib-applehelp==2.0.0",
73
+ "sphinxcontrib-bibtex==2.6.2",
74
+ "sphinxcontrib-devhelp==2.0.0",
75
+ "sphinxcontrib-htmlhelp==2.1.0",
76
+ "sphinxcontrib-jquery==4.1",
77
+ "sphinxcontrib-jsmath==1.0.1",
78
+ "sphinxcontrib-qthelp==2.0.0",
79
+ "sphinxcontrib-serializinghtml==2.0.0",
80
+ ]
81
+
82
+ # Optional dependencies
83
+ [project.optional-dependencies]
84
+ other_methods = [
85
+ "pyoptflow",
86
+ "openpiv",
87
+ "opencv-python-headless",
88
+ "torch>=1.13.1",
89
+ "torchvision>=0.14.1",
90
+ ]
91
+
92
+ cuda12 = [
93
+ "jax[cuda12]>=0.6.0",
94
+ ]
95
+
96
+ # autopep8 configuration for 4-space indentation
97
+ [tool.autopep8]
98
+ max_line_length = 80
99
+ indent_size = 4
100
+ aggressive = 2
101
+ in_place = true
102
+ exclude = ".eggs,.git,.hg,.mypy_cache,.tox,.venv,_build,buck-out,build,dist"
103
+
104
+ # isort configuration (updated for autopep8 compatibility)
105
+ [tool.isort]
106
+ line_length = 80
107
+ indent = 4
108
+ multi_line_output = 0
109
+ force_single_line = false
110
+ include_trailing_comma = false
111
+ force_grid_wrap = 0
112
+ use_parentheses = true
113
+ ensure_newline_before_comments = true
114
+
115
+ # MyPy configuration
116
+ [tool.mypy]
117
+ python_version = "3.10"
118
+ warn_return_any = true
119
+ warn_unused_configs = true
120
+ disallow_untyped_defs = true
121
+ disallow_incomplete_defs = true
122
+ check_untyped_defs = true
123
+ disallow_untyped_decorators = true
124
+ no_implicit_optional = true
125
+ warn_redundant_casts = true
126
+ warn_unused_ignores = true
127
+ warn_no_return = true
128
+ warn_unreachable = true
129
+ strict_equality = true
130
+ show_error_codes = true
131
+ explicit_package_bases = true
132
+
133
+ # Allow untyped decorators in test files for Hypothesis
134
+ [[tool.mypy.overrides]]
135
+ module = "tests.*"
136
+ disallow_untyped_decorators = false
137
+
138
+ # Pydocstyle configuration
139
+ [tool.pydocstyle]
140
+ convention = "google"
141
+ match-dir = "^(?!tests).*"
142
+
143
+ # Bandit security configuration
144
+ [tool.bandit]
145
+ exclude_dirs = ["tests", ".venv", "build", "dist"]
146
+
147
+ # Pytest configuration
148
+ [tool.pytest.ini_options]
149
+ minversion = "7.0"
150
+ addopts = [
151
+ "-ra",
152
+ "--strict-markers",
153
+ "--strict-config",
154
+ "--cov=goggles",
155
+ "--cov-branch",
156
+ "--cov-report=term-missing",
157
+ "--cov-report=html",
158
+ "--cov-report=xml",
159
+ ]
160
+ testpaths = ["tests"]
161
+ filterwarnings = [
162
+ "error",
163
+ "ignore::UserWarning",
164
+ "ignore::DeprecationWarning",
165
+ ]
166
+
167
+ # Ruff configuration for linting and modern typing enforcement
168
+ [tool.ruff]
169
+ line-length = 80
170
+ target-version = "py310"
171
+
172
+ [tool.ruff.lint]
173
+ # Enable comprehensive rule sets
174
+ select = [
175
+ "E", # pycodestyle errors
176
+ "W", # pycodestyle warnings
177
+ "F", # pyflakes (including unused imports F401)
178
+ "UP", # pyupgrade (modern typing patterns)
179
+ "I", # isort
180
+ "N", # pep8-naming
181
+ "D", # pydocstyle
182
+ "PL", # pylint
183
+ "RUF", # ruff-specific rules
184
+ ]
185
+
186
+ ignore = [
187
+ "D100", # Missing docstring in public module
188
+ "D104", # Missing docstring in public package
189
+ "D105", # Missing docstring in magic method
190
+ "PLR0913", # Too many arguments to function call
191
+ "PLR2004", # Magic value used in comparison
192
+ ]
193
+
194
+ # Specific rules for modern typing enforcement
195
+ [tool.ruff.lint.pyupgrade]
196
+ # Keep this empty to use all pyupgrade rules
197
+
198
+ [tool.ruff.lint.per-file-ignores]
199
+ "tests/*" = ["D", "PLR2004"] # Ignore docstring rules in tests
200
+ "examples/*" = ["D"] # Ignore docstring rules in examples
201
+
202
+ [tool.ruff.lint.pydocstyle]
203
+ convention = "google"
204
+
205
+ [tool.setuptools.dynamic]
206
+ version = {attr = "flowgym.__version__"}
207
+
208
+ [tool.uv]
209
+ default-groups = []
210
+
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+