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.
- flow_gym_suite-0.1.0/LICENSE +21 -0
- flow_gym_suite-0.1.0/PKG-INFO +233 -0
- flow_gym_suite-0.1.0/README.md +178 -0
- flow_gym_suite-0.1.0/pyproject.toml +210 -0
- flow_gym_suite-0.1.0/setup.cfg +4 -0
- flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/PKG-INFO +233 -0
- flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/SOURCES.txt +85 -0
- flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/dependency_links.txt +1 -0
- flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/requires.txt +15 -0
- flow_gym_suite-0.1.0/src/flow_gym_suite.egg-info/top_level.txt +1 -0
- flow_gym_suite-0.1.0/src/flowgym/__init__.py +63 -0
- flow_gym_suite-0.1.0/src/flowgym/common/__init__.py +1 -0
- flow_gym_suite-0.1.0/src/flowgym/common/base/__init__.py +10 -0
- flow_gym_suite-0.1.0/src/flowgym/common/base/estimator.py +461 -0
- flow_gym_suite-0.1.0/src/flowgym/common/base/trainable_states.py +140 -0
- flow_gym_suite-0.1.0/src/flowgym/common/evaluation.py +200 -0
- flow_gym_suite-0.1.0/src/flowgym/common/filters.py +118 -0
- flow_gym_suite-0.1.0/src/flowgym/common/median.py +1286 -0
- flow_gym_suite-0.1.0/src/flowgym/common/preprocess.py +356 -0
- flow_gym_suite-0.1.0/src/flowgym/density/__init__.py +1 -0
- flow_gym_suite-0.1.0/src/flowgym/density/nn.py +129 -0
- flow_gym_suite-0.1.0/src/flowgym/density/simple.py +60 -0
- flow_gym_suite-0.1.0/src/flowgym/environment/__init__.py +1 -0
- flow_gym_suite-0.1.0/src/flowgym/environment/fluid_env.py +164 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/__init__.py +1 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/base.py +98 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/consensus/__init__.py +7 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/consensus/admm.py +271 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/consensus/consensus.py +385 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/consensus/consensus_algorithms.py +417 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/consensus/objectives.py +426 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/consensus/regularizers.py +106 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/consensus/solvers.py +213 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/deepflow.py +48 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/dis/__init__.py +12 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/dis/dis.py +174 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/dis/dis_jax.py +210 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/dis/process.py +1091 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/dummy.py +42 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/farneback.py +92 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/hornschunck.py +135 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/__init__.py +12 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/openpiv.py +117 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/openpiv_jax.py +215 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/open_piv/process.py +375 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/__init__.py +270 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/data_interpolation.py +114 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/data_smoothing.py +155 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/postprocess/data_validation.py +283 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/process.py +215 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/raft/__init__.py +13 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/raft/process.py +350 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/raft/raft_jax.py +417 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/raft/raft_piv_pytorch.py +243 -0
- flow_gym_suite-0.1.0/src/flowgym/flow/utils.py +182 -0
- flow_gym_suite-0.1.0/src/flowgym/make.py +234 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/__init__.py +1 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/blocks.py +437 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/cnn.py +90 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/raft_model.py +99 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/__init__.py +6 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/corr_block.py +55 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/flowNetsRAFT.py +203 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/submodules_RAFT_GRU.py +165 -0
- flow_gym_suite-0.1.0/src/flowgym/nn/raft_torch_nn/submodules_RAFT_extractor.py +272 -0
- flow_gym_suite-0.1.0/src/flowgym/training/__init__.py +1 -0
- flow_gym_suite-0.1.0/src/flowgym/training/optimizer.py +99 -0
- flow_gym_suite-0.1.0/src/flowgym/training/schedules.py +84 -0
- flow_gym_suite-0.1.0/src/flowgym/training/utils.py +176 -0
- flow_gym_suite-0.1.0/src/flowgym/types.py +7 -0
- flow_gym_suite-0.1.0/src/flowgym/utils.py +441 -0
- flow_gym_suite-0.1.0/tests/test_data_interpolation.py +175 -0
- flow_gym_suite-0.1.0/tests/test_data_smoothing.py +191 -0
- flow_gym_suite-0.1.0/tests/test_data_validation.py +381 -0
- flow_gym_suite-0.1.0/tests/test_density_estimate_simple.py +97 -0
- flow_gym_suite-0.1.0/tests/test_dis_jax.py +1068 -0
- flow_gym_suite-0.1.0/tests/test_estimation_stats.py +99 -0
- flow_gym_suite-0.1.0/tests/test_filters.py +107 -0
- flow_gym_suite-0.1.0/tests/test_flow_estimate_base.py +49 -0
- flow_gym_suite-0.1.0/tests/test_flow_estimate_dis.py +117 -0
- flow_gym_suite-0.1.0/tests/test_median.py +158 -0
- flow_gym_suite-0.1.0/tests/test_openpiv_jax.py +275 -0
- flow_gym_suite-0.1.0/tests/test_postprocess_config.py +496 -0
- flow_gym_suite-0.1.0/tests/test_preprocess.py +156 -0
- flow_gym_suite-0.1.0/tests/test_preprocess_config.py +216 -0
- flow_gym_suite-0.1.0/tests/test_trainable_state.py +267 -0
- 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
|
+
[](https://www.python.org/downloads/)
|
|
59
|
+
[](https://github.com/antonioterpin/flowgym/stargazers)
|
|
60
|
+
[](https://pypi.org/project/flow-gym)
|
|
61
|
+
[](https://opensource.org/licenses/MIT)
|
|
62
|
+
[](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
|
+

|
|
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
|
+
[](https://www.python.org/downloads/)
|
|
4
|
+
[](https://github.com/antonioterpin/flowgym/stargazers)
|
|
5
|
+
[](https://pypi.org/project/flow-gym)
|
|
6
|
+
[](https://opensource.org/licenses/MIT)
|
|
7
|
+
[](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
|
+

|
|
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
|
+
|