Flow Gym: a research toolkit for flow-field quantification

Flow Gym is a Python toolkit for developing, benchmarking, training, and deploying flow-field quantification methods. It provides a standardized estimator interface, JAX-first implementations, interoperable wrappers for external methods, and shared workflows for repeatable evaluation and training.

What Flow Gym provides

• Unified estimator construction through flowgym.make
• A shared Estimator interface for classical and learning-based methods
• Flow-field and density estimator implementations under flowgym.flow and flowgym.density
• JAX-native implementations together with wrappers for representative external libraries
• Training, evaluation, benchmarking, and comparison orchestration through src/main.py
• Cache-backed helpers for repeated dataset passes and expensive derived data

Install

uv add flow-gym-suite

With pip instead:

pip install flow-gym-suite

For local repository work:

uv sync --group dev

Minimal API example

import jax
import jax.numpy as jnp

from flowgym.make import make_estimator

estimator_config = {
    "estimator": "dis_jax",
    "estimate_type": "flow",
    "config": {
        "jit": False,
        "preset": 1,
    },
}

image_shape = (1, 64, 64)
estimate_shape = (1, 64, 64, 2)

trained_state, create_state_fn, compute_estimate_fn, estimator = make_estimator(
    estimator_config=estimator_config,
    image_shape=image_shape,
    estimate_shape=estimate_shape,
    rng=0,
)

prev = jnp.zeros(image_shape, dtype=jnp.float32)
curr = jnp.ones(image_shape, dtype=jnp.float32)
state = create_state_fn(prev, jax.random.PRNGKey(0))
new_state, metrics = compute_estimate_fn(curr, state, trained_state)

This is the smallest public API loop in Flow Gym. For repository-backed examples, benchmarking workflows, and CLI usage, use the docs pages linked below.

CLI example

uv run python src/main.py \
  --mode eval \
  --estimator src/flowgym/config/estimators/flow/real_dis.yaml \
  --dataset src/flowgym/config/piv_dataset_class1_eval.yaml

Documentation

The canonical human-facing docs start at docs/index.md.

• Getting started
• Example workflows
• API reference
• Architecture guide
• Contributing guide

If you are new to the project, start with docs/getting-started/index.md for the motivation, installation paths, and first estimator example.

Caching examples

The repository includes two documented caching walkthroughs:

• DIS caching example backed by examples/10_caching.py
• RAFT caching example backed by examples/11_caching.py

Contributing

Contributions are welcome. For current setup, quality gates, docs workflow, and repository navigation, use docs/contribute/contributing.md.

Citation

If you use this code in your research, please cite:

@article{banelli2026flowgym,
   title={Flow Gym: A framework for the development, benchmarking, training, and deployment of flow-field quantification methods},
   author={Banelli, Francesco and Terpin, Antonio and Bonomi, Alan and D'Andrea, Raffaello},
   journal={SoftwareX},
   volume={34},
   pages={102641},
   year={2026},
   doi={10.1016/j.softx.2026.102641}
}

License

This project is licensed under the MIT License. See LICENSE.