HydroGym

https://dynamicslab.github.io/hydrogym/
Admin
HydroGym is an open-source reinforcement learning platform for active flow control in fluid...
Tokens:102,001
Snippets:1,188
Trust Score:9.7
Update:1 month ago
Show doc for...
Tokens:
Raw
### Run Kolmogorov Environment Example

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jax

Quick start command to run the Kolmogorov environment example. No MPI or MPMD launch is required; JAX environments run in a single Python process.

```bash
cd examples/jax/getting_started/1_kolmogorov/
python test_kolmogorov_env.py
```

--------------------------------

### Full installation sequence

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jaxfluids

Complete installation sequence: install JAX, clone and install JAX-Fluids, then install HydroGym extras.

```bash
# 1. Install JAX with the appropriate hardware target  

pip install --upgrade "jax[cuda12]"   # or jax[cpu], jax[cuda13], etc.  

  

# 2. Install JAX-Fluids from source  

git clone https://github.com/tumaer/JAXFLUIDS.git  

cd JAXFLUIDS && pip install . && cd ..  

  

# 3. Install HydroGym with JAX-Fluids extras  

pip install hydrogym[jaxfluids]  
```

--------------------------------

### Verify NEK5000 build with examples

Source: https://dynamicslab.github.io/hydrogym/docs/installation/nek5000

Download examples, partition the mesh, compile, and run a test case on 2 MPI ranks.

```bash
cd ~  

git clone https://github.com/Nek5000/NekExamples.git  

cp -r ~/NekExamples/eddy_uv .  

cd eddy_uv  

  

genmap                  # partition the mesh (input: eddy_uv)  

makenek eddy_uv         # compile the case  

nekbmpi eddy_uv 2       # run on 2 MPI ranks  

tail logfile            # check solver output  
```

--------------------------------

### Install HydroGym JAX Extras

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jax

Installs all required packages for the JAX-based examples. If JAX is already installed with GPU support, the jax and jaxlib pins will not downgrade it.

```bash
pip install hydrogym[jax]
```

--------------------------------

### Install JAX for CPU

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jax

Installs the CPU-only version of JAX, sufficient for running the Kolmogorov and channel flow examples and training small agents on a workstation.

```bash
pip install -U jax
```

--------------------------------

### Run from_hf examples with mpirun

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/from_hf

Run the example scripts from examples/nek/getting_started/4_from_hf/ using MPI. The test command runs test_nek_DM.py with 100 steps, and the train command runs SB3 training with PPO for 100000 timesteps, both launching the nek5000 solver with 10 processes.

```bash
# Test  
mpirun -np 1 python test_nek_DM.py --steps 100 : -np 10 nek5000  

# Train  
mpirun -np 1 python train_sb3_from_hf.py \  
    --env TCFmini_3D_Re180 \  
    --algo PPO \  
    --total-timesteps 100000 \  
    : -np 10 nek5000
```

--------------------------------

### Install JAX-Fluids from source

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jaxfluids

JAX-Fluids is not on PyPI, so clone the repository and install it with pip.

```bash
git clone https://github.com/tumaer/JAXFLUIDS.git  

cd JAXFLUIDS  

pip install .  
```

--------------------------------

### Initialize and Prepare Environment

Source: https://dynamicslab.github.io/hydrogym/docs/api/data_manager

Example showing how to initialize the HFDataManager, retrieve an environment path, and prepare a working directory.

```python
# Use HF cache directly
dm = HFDataManager(use_clean_cache=False)
# Get environment path (downloads if needed)
env_path = dm.get_environment_path('Cylinder_2D_Re200')
# Prepare a solver-specific working directory
paths = dm.prepare_working_directory(env_path, './run_dir')
```

--------------------------------

### Quick start commands for HydroGym Firedrake

Source: https://dynamicslab.github.io/hydrogym/docs/examples/firedrake/getting_started

Three sequential commands to browse configuration options, verify the environment, and start training. The first command lists all configuration options. The second verifies the environment with a cylinder simulation for 10 steps. The third starts training with PPO for 100,000 timesteps.

```bash
# 1. Browse all configuration options  

python config_reference.py  

  

# 2. Verify the environment runs correctly  

python test_firedrake_env.py --environment cylinder --num-steps 10 --verbose  

  

# 3. Copy a configuration template from config_reference.py and start training  

python train_sb3_firedrake.py --env cylinder --algo PPO --total-timesteps 100000  

```

--------------------------------

### Install with specific solver backends

Source: https://dynamicslab.github.io/hydrogym/docs/developers/contributing

Install the environment with additional solver backends if you plan to work on environment code.

```bash
uv sync --locked --extra maia  

uv sync --locked --extra nek  

uv sync --locked --extra jax  

uv sync --locked --extra jaxfluids  
```

--------------------------------

### Running example scripts with MPI

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/hydrogym_control

Commands to run the example scripts using mpirun, with one process for the controller and 10 processes for the Nek5000 simulation. The first command tests controllers, the second runs the full train + evaluate workflow.

```bash
# Test controllers
mpirun -np 1 python test_nek_env_controller.py \
    --config test_config.yml \
    : -np 10 nek5000

# Full train + evaluate workflow
mpirun -np 1 python train_sb3_with_integrate.py \
    --env TCFmini_3D_Re180 \
    --algo PPO \
    --total-timesteps 50000 \
    --eval-steps 200 \
    : -np 10 nek5000
```

--------------------------------

### HFDataManager Usage Examples

Source: https://dynamicslab.github.io/hydrogym/docs/api/maia/hf_data_manager

Examples showing how to initialize the HFDataManager with different caching strategies.

```python
# Symlinks (recommended - best of both worlds!)
data_manager = HFDataManager(use_clean_cache=True)
# Copy files (if you need to modify them)
data_manager = HFDataManager(use_clean_cache='copy')
# Use HF cache directly
data_manager = HFDataManager(use_clean_cache=False)
```

--------------------------------

### Create Firedrake virtual environment and verify installation

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Create and activate a virtual environment for Firedrake, set PETSC_DIR and related variables, install Firedrake, and run firedrake-check.

```shell
python3 -m venv venv-firedrake  
. venv-firedrake/bin/activate  

# Set PETSC_DIR and related variables  
export $(python3 firedrake-configure --show-env)  

pip cache purge  

pip install --no-binary h5py 'firedrake[check]'  

# Verify the installation  
firedrake-check  

```

--------------------------------

### Run prepare_workspace.py for workspace setup

Source: https://dynamicslab.github.io/hydrogym/docs/examples/maia/getting_started

Downloads environment data and creates a workspace for HPC jobs. Specify the environment and working directory.

```bash
python prepare_workspace.py --env Cylinder_2D_Re200 --work-dir ./my_workspace
```

--------------------------------

### Import JAX and dependencies

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/ppo

Initial setup of required libraries including JAX, Flax, and Optax for the PPO training environment.

```python
import argparse  

import pickle  

from typing import NamedTuple, Sequence  

  

import distrax  

import flax.linen as nn  

import flax.serialization  

import jax  

import jax.numpy as jnp  

import matplotlib.pyplot as plt  

import numpy as np  

import optax  
```

--------------------------------

### Install system packages on Ubuntu/Debian

Source: https://dynamicslab.github.io/hydrogym/docs/installation/nek5000

Commands to install the necessary compilers, MPI, and optional tools for NEK5000.

```bash
sudo apt-get update  

sudo apt install build-essential       # gcc, make  

sudo apt install gfortran              # Fortran compiler  

sudo apt install libopenmpi-dev        # MPI  

sudo apt install cmake libx11-dev libxt-dev  # optional: NekTools + visnek  
```

--------------------------------

### Install PettingZoo and SuperSuit

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/pettingzoo

Installs the PettingZoo and SuperSuit packages required for the SB3 integration workflow.

```bash
pip install pettingzoo supersuit  
```

--------------------------------

### Run parallel environment examples with mpirun

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/parallel_environments

Commands to run the parallel environment examples. The first tests the environment with random actions for 100 steps, and the second trains with the DIY centralised wrapper using PPO for 100,000 timesteps. Both use mpirun to launch the Python script on one process and the NEK5000 executable on 10 processes.

```bash
# Test the parallel environment  
mpirun -np 1 python test_nek_parallel.py --steps 100 : -np 10 nek5000  

# Train with DIY centralised wrapper  
mpirun -np 1 python train_sb3_parallel.py \  
    --env TCFmini_3D_Re180 \  
    --algo PPO \  
    --total-timesteps 100000 \  
    : -np 10 nek5000
```

--------------------------------

### Download firedrake-configure helper script

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Download the firedrake-configure script that assists with native installation.

```shell
curl -O https://raw.githubusercontent.com/firedrakeproject/firedrake/release/scripts/firedrake-configure  

```

--------------------------------

### Start ParaView server in Apptainer

Source: https://dynamicslab.github.io/hydrogym/docs/installation/maia

Starts the ParaView server launcher inside the Apptainer environment. Prints a socket address such as node042:11111.

```bash
bash /home/easybuild/paraview-plugins/pvServerLaunch2024.sh 32  

```

--------------------------------

### Preview documentation locally

Source: https://dynamicslab.github.io/hydrogym/docs/developers/contributing

Navigate to the docs directory and start the Docusaurus development server to preview documentation changes.

```bash
cd docs  

npm ci  

npm start       # opens http://localhost:3000/hydrogym/  
```

--------------------------------

### Initialize and Step KolmogorovFlow Environment

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/getting_started

Example of importing the KolmogorovFlow environment, initializing it with default parameters, and performing a single step.

```python
import jax  

import jax.numpy as jnp  

from hydrogym.jax.envs.kolmogorov import KolmogorovFlow  

  

env = KolmogorovFlow(env_config={}, flow_config={})  

params = env.default_params  

  

key = jax.random.PRNGKey(0)  

obs, state = env.reset_env(key, params)  

  

action = jnp.zeros((params.action_dim,))  

obs, state, reward, done, info = env.step_env(key, state, action, params)
```

--------------------------------

### Install JAX for CPU or NVIDIA GPU

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jaxfluids

Install JAX with the appropriate hardware target. Choose the CPU-only or CUDA 12 version based on your hardware.

```bash
# CPU only  

pip install --upgrade "jax[cpu]"  

  

# NVIDIA GPU (CUDA 12)  

pip install --upgrade "jax[cuda12]"  
```

--------------------------------

### Basic cylinder environment setup

Source: https://dynamicslab.github.io/hydrogym/docs/examples/firedrake/getting_started

Initializes a cylinder flow environment and runs a simple simulation loop using the SemiImplicitBDF solver.

```python
from hydrogym import FlowEnv  

import hydrogym.firedrake as hgym  

  

env_config = {  

    'flow': hgym.Cylinder,  

    'flow_config': {'mesh': 'medium', 'Re': 100},  

    'solver': hgym.SemiImplicitBDF,  

    'solver_config': {'dt': 1e-2},  

}  

  

env = FlowEnv(env_config)  

obs, info = env.reset()  

  

for _ in range(100):  

    action = env.action_space.sample()  

    obs, reward, terminated, truncated, info = env.step(action)  
```

--------------------------------

### Run MAIA examples in Docker

Source: https://dynamicslab.github.io/hydrogym/docs/examples/maia/getting_started

Docker runner script for MAIA examples. Without arguments, it tests the environment; with 'train', it trains an SB3 agent.

```bash
# Test environment  
./run_example_docker.sh  

# Train SB3 agent  
./run_example_docker.sh train
```

--------------------------------

### Inspect training configuration

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/ppo

Prints the current training configuration and paths to the console for verification before starting the training process.

```python
print(f"=== PPO Training: {args.env} environment ===")  

print(f"  Total timesteps : {config['TOTAL_TIMESTEPS']}")  

print(f"  Num envs        : {config['NUM_ENVS']}")  

print(f"  Num steps       : {config['NUM_STEPS']}")  

print(f"  Learning rate   : {config['LR']}")  

print(f"  Model save path : {model_save_path}")  

print(f"  Plot save path  : {plot_path}")  

print("")  
```

--------------------------------

### Run HydroGym Quick Start Scripts

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/getting_started

Commands to activate the GPU environment and execute test scripts for Kolmogorov flow, channel flow, and PPO training.

```bash
# Activate the GPU environment  

source /home/easybuild/venvs/hydrogym_gpu/bin/activate  

  

# Test Kolmogorov flow  

cd getting_started/1_kolmogorov  

./run_kolmogorov_docker.sh  

  

# Test channel flow  

cd getting_started/2_channel  

./run_channel_docker.sh  

  

# Train PPO  

cd getting_started/3_ppo  

./run_ppo_docker.sh --env kolmogorov --total-timesteps 20000
```

--------------------------------

### Run config_reference.py

Source: https://dynamicslab.github.io/hydrogym/docs/examples/firedrake/getting_started

Runs the configuration reference script, which prints annotated output for ten progressively complex setups covering all common use cases.

```bash
python config_reference.py  
```

--------------------------------

### Quick start test

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jaxfluids

Run the JAX-Fluids environment test script. JAX-Fluids environments run in a single Python process, no MPI or MPMD launch required.

```bash
cd examples/jaxfluids/  

python test_jaxfluids_env.py  
```

--------------------------------

### Curriculum learning with multiple checkpoints

Source: https://dynamicslab.github.io/hydrogym/docs/examples/firedrake/getting_started

Passing a list of checkpoint paths in 'restart' causes each env.reset() to randomly select one, enabling a simple form of curriculum learning. The example prints the selected checkpoint index via info.get('checkpoint_index').

```python
env_config = {  

    'flow': hgym.Pinball,  

    'flow_config': {  

        'mesh': 'fine',  

        'Re': 30,  

        'restart': [  

            'checkpoint_early.h5',  

            'checkpoint_mid.h5',  

            'checkpoint_late.h5',  

        ],  

    },  

    'solver': hgym.SemiImplicitBDF,  

    'solver_config': {'dt': 1e-2},  

}  

  

env = FlowEnv(env_config)  

obs, info = env.reset()  

print(f"Started from checkpoint index: {info.get('checkpoint_index')}")  

```

--------------------------------

### Install HydroGym NEK extras

Source: https://dynamicslab.github.io/hydrogym/docs/installation/nek5000

Install the HydroGym Python package with NEK5000 support.

```bash
pip install hydrogym[nek]  
```

--------------------------------

### Import JAX and Hydrogym components

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/kolmogorov

Initializes the environment by importing JAX, plotting libraries, and the Pseudospectral Navier Stokes Solver.

```python
import os  

import time  

import jax.numpy as jnp  

import matplotlib.pyplot as plt  

import seaborn as sns  

from hydrogym.jax.solvers.base import RungeKuttaCrankNicolson  

from hydrogym.jax.utils import io as io  

from hydrogym.jax.envs.kolmogorov import FlowConfig, PseudoSpectralNavierStokes2D  
```

--------------------------------

### Editable install for JAX-Fluids development

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jaxfluids

Use an editable install if you plan to modify JAX-Fluids itself.

```bash
pip install -e .  
```

--------------------------------

### Run run_example_docker.sh (test or train)

Source: https://dynamicslab.github.io/hydrogym/docs/examples/firedrake/getting_started

Convenience wrapper for Docker image selection and volume mounting. Pass `train` to run the SB3 training script instead of the environment test.

```bash
./run_example_docker.sh          # test environment  
./run_example_docker.sh train    # train SB3 agent  
```

--------------------------------

### Install HydroGym MAIA extras via pip

Source: https://dynamicslab.github.io/hydrogym/docs/installation/maia

Installs the hydrogym[maia] extra. Adds mpi4py, omegaconf, einops, and toml alongside the core package. Requires a local m-AIA installation (Docker or native).

```bash
pip install hydrogym[maia]  

```

--------------------------------

### MaiaWorkspace.setup

Source: https://dynamicslab.github.io/hydrogym/docs/api/maia/workspace

Sets up the workspace by downloading environment data from Hugging Face Hub (or local fallback), creating the work directory structure, and creating symbolic links with proper file names and extensions.

```APIDOC
## MaiaWorkspace.setup

### Description
Sets up the workspace with all required files. It downloads/locates environment data from HF Hub, creates the work directory structure, and creates symbolic links with proper names and extensions.

### Method Signature
```
def setup(force_download: bool = False, verbose: bool = True) -> Dict[str, str]
```

### Parameters
- **force_download** (bool) - Optional - Force re-download from HF Hub. Defaults to False.
- **verbose** (bool) - Optional - Print setup progress. Defaults to True.

### Returns
Dictionary containing paths to key files:
- 'work_dir': Path to work directory
- 'properties_file': Path to properties file for MAIA
- 'config_file': Path to environment config file
- 'env_data_path': Path to source environment data
```

--------------------------------

### Install system packages with firedrake-configure

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Use firedrake-configure to install the required system packages on Ubuntu/Debian or macOS.

```shell
# Ubuntu / Debian  
sudo apt install $(python3 firedrake-configure --show-system-packages)  

# macOS (Homebrew)  
brew install $(python3 firedrake-configure --show-system-packages)  

```

--------------------------------

### PDEBase.initialize_state method

Source: https://dynamicslab.github.io/hydrogym/docs/api/core

Sets up the mesh, function spaces, and state vector.

```python
@abc.abstractmethod  

def initialize_state()
```

--------------------------------

### Prepare Working Directory

Source: https://dynamicslab.github.io/hydrogym/docs/api/data_manager

Creates a working directory and populates it with solver-specific symlinks.

```python
def prepare_working_directory(env_path: str,  

                              work_dir: str,  

                              profile: Optional[str] = None) -> Dict[str, str]
```

--------------------------------

### Install HydroGym in Firedrake environment

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Install the HydroGym package with the Firedrake extra while the Firedrake virtual environment is active.

```shell
pip install hydrogym[firedrake]  

```

--------------------------------

### Prepare, train, and evaluate with YAML config

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/pettingzoo

Configuration-driven workflow: prepare a workspace, train from a YAML config file, and evaluate. Run training from the workspace directory so that the compile_path: "." in the config resolves to the prepared case files, and ensure --nproc matches simulation.nproc in the config.

```bash
# 1. Prepare a workspace  

python ../prepare_workspace.py \  

  --local-dir ../../../packaged_envs \  

  --env TCFmini_3D_Re180 \  

  --work-dir ./train_run  

  

# 2. Train from a config file (run from the workspace directory)  

cd train_run  

mpirun -np 1 python ../train_sb3_pettingzoo.py \  

  --env TCFmini_3D_Re180 \  

  --nproc 10 \  

  --config-file ../configs/pettingzoo_tcfmini_re180.yml \  

  --algo TD3 \  

  --total-timesteps 5000000 \  

  : -np 10 nek5000  

  

# 3. Evaluate  

mpirun -np 1 python ../test_nek_pettingzoo.py \  

  --env TCFmini_3D_Re180 \  

  --nproc 10 \  

  --config-file ../configs/pettingzoo_tcfmini_re180.yml \  

  --steps 2500 \  

  : -np 10 nek5000  
```

--------------------------------

### Install JAX for AMD GPU (ROCm)

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jax

Installs JAX with ROCm support. ROCm support on Windows WSL2 is experimental.

```bash
pip install -U "jax[rocm7-local]"
```

--------------------------------

### Install HydroGym JAX-Fluids extras

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jaxfluids

Install the HydroGym package with the JAX-Fluids extras, which adds jax, jaxlib, flax, gitpython, h5py, optax, and pyvista.

```bash
pip install hydrogym[jaxfluids]  
```

--------------------------------

### CubeBase.__init__

Source: https://dynamicslab.github.io/hydrogym/docs/api/maia/envs/cube

Initializes the cube base environment with a configuration dictionary.

```APIDOC
## CubeBase.__init__

### Description
Initializes the cube base environment.

### Method
Constructor

### Signature
```python
def __init__(env_config: Dict)
```

### Parameters
- **env_config** (Dict) - Required - Environment configuration dictionary.

### Returns
None
```

--------------------------------

### Activate Firedrake environment and install HydroGym inside container

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Inside the container, activate the Firedrake virtual environment and install HydroGym with the firedrake extra.

```shell
. /home/firedrake/firedrake/bin/activate  

pip install hydrogym[firedrake]  

```

--------------------------------

### Load pre-packaged environments from HuggingFace

Source: https://dynamicslab.github.io/hydrogym/docs/installation/nek5000

Initialize a NEK5000 environment directly from the HuggingFace Hub.

```python
from hydrogym.nek import NekEnv  

  
env = NekEnv.from_hf('TCFmini_3D_Re180', nproc=10)  

obs, info = env.reset()  
```

--------------------------------

### Optional Firedrake extras installation

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Install optional Firedrake extras for SLEPc or VTK support. Always include --no-binary h5py to avoid conflicts.

```shell
pip install 'firedrake[check,slepc]'
```

```shell
pip install 'firedrake[check,vtk]'
```

--------------------------------

### Install JAX for NVIDIA GPU (CUDA 12 and 13)

Source: https://dynamicslab.github.io/hydrogym/docs/installation/jax

Installs JAX with CUDA support. The CUDA 12 wheel is built against CUDA 12.3 and requires CUDA ≥ 12.1; the CUDA 13 wheel requires CUDA ≥ 13.0. JAX installs its own CUDA and cuDNN libraries, so only the NVIDIA driver is required.

```bash
# CUDA 12 (driver ≥ 525 on Linux)
pip install -U "jax[cuda12]"
```

```bash
# CUDA 13 (driver ≥ 580 on Linux)
pip install -U "jax[cuda13]"
```

--------------------------------

### Cube.__init__

Source: https://dynamicslab.github.io/hydrogym/docs/api/maia/envs/cube

Initializes the cube environment with a configuration dictionary.

```APIDOC
## Cube.__init__

### Description
Initializes the cube environment.

### Method
Constructor

### Signature
```python
def __init__(env_config: Dict)
```

### Parameters
- **env_config** (Dict) - Required - Environment configuration dictionary.

### Returns
None
```

--------------------------------

### Prepare workspace with prepare_workspace.py

Source: https://dynamicslab.github.io/hydrogym/docs/examples/maia/getting_started

Downloads data from Hugging Face Hub (requires internet) and creates the working directory, configuration file, mesh symlink, and output directory. No MPI needed.

```bash
python prepare_workspace.py --env Cylinder_2D_Re200 --work-dir ./test_run_000
```

--------------------------------

### SimpleRewardLogger.new_episode

Source: https://dynamicslab.github.io/hydrogym/docs/api/nek/nek_lib/reward_logger

Marks the start of a new episode.

```APIDOC
## SimpleRewardLogger.new_episode

### Description
Mark the start of a new episode.

### Method
Method

### Signature
```python
def new_episode()
```
```

--------------------------------

### Run demo directly with mpirun

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/zeroshot_wing

Shows how to run the Python script `zeroshot_demo_pettingzoo.py` directly using mpirun, either with the default settings or with a custom policy template, policy root, and number of steps.

```bash
# Default template  
mpirun -np 1 python zeroshot_demo_pettingzoo.py : -np 12 nek5000  

  

# Custom template and policy root  
mpirun -np 1 python zeroshot_demo_pettingzoo.py \  
  --policy-template ./meta_policy_small_wing_template.py \  
  --policy-root /path/to/legacy_runs \  
  --steps 3000 \  
  : -np 12 nek5000  

```

--------------------------------

### JAXFlowEnv.get_environment_files_info

Source: https://dynamicslab.github.io/hydrogym/docs/api/jax/env_core

Get information about where environment files are stored.

```APIDOC
## JAXFlowEnv.get_environment_files_info

### Description
Get information about where environment files are stored. Useful for debugging and understanding file locations.

### Method
Function

### Signature
```python
def get_environment_files_info() -> Dict
```

### Returns
- **Dict** - Dictionary containing environment name, paths, and file information.
```

--------------------------------

### JAXFlowEnv.get_available_environments

Source: https://dynamicslab.github.io/hydrogym/docs/api/jax/env_core

Get list of all available environments from HF Hub.

```APIDOC
## JAXFlowEnv.get_available_environments

### Description
Get list of all available environments from HF Hub.

### Method
Function

### Signature
```python
def get_available_environments() -> List[str]
```

### Returns
- **List[str]** - List of environment names.
```

--------------------------------

### Run demo with Docker launcher

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/zeroshot_wing

Shows how to invoke the runner script `./run_pettingzoo_docker.sh`, both with the default settings and with a custom `--policy-root` to point to a directory containing pre-trained model runs.

```bash
# Default template  
./run_pettingzoo_docker.sh  

  
# Specify a root directory containing pre-trained model runs  
./run_pettingzoo_docker.sh --policy-root /workspace/legacy_runs  

```

--------------------------------

### Create NEK5000 environment with NekEnv.from_hf()

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/from_hf

Create a NEK5000 environment by name using NekEnv.from_hf(). Provide the environment name and nproc; the method auto-detects the config file and prepares the workspace. Setting use_clean_cache=False reuses an existing prepared workspace for iterative development.

```python
from hydrogym.nek import NekEnv  

env = NekEnv.from_hf(  
    'TCFmini_3D_Re180',  
    nproc=10,  
    use_clean_cache=False,         # reuse an existing prepared workspace  
    local_fallback_dir=None,       # optional: path to a local environment package  
)  

obs, info = env.reset()  
obs, reward, terminated, truncated, info = env.step(action)
```

--------------------------------

### SinWave load_node_info method

Source: https://dynamicslab.github.io/hydrogym/docs/api/nek/AFC

Get the node info to impose the Sinusoidal Wave.

```python
def load_node_info(Node_Info)
```

--------------------------------

### Configure Optimizer and Initialize TrainState

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/ppo

Sets up the Optax optimizer with gradient clipping and Adam, then initializes the Flax TrainState.

```python
        else:  

            tx = optax.chain(  

                optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),  

                optax.adam(config["LR"], eps=1e-5),  

            )  

        train_state = TrainState.create(  

            apply_fn=network.apply,  

            params=network_params,  

            tx=tx,  

        )
```

--------------------------------

### JAXFlowEnv get_available_environments Method

Source: https://dynamicslab.github.io/hydrogym/docs/api/jax/env_core

Get a list of all available environments from the Hugging Face Hub.

```python
def get_available_environments() -> List[str]
```

--------------------------------

### JAXFlowEnv get_environment_files_info Method

Source: https://dynamicslab.github.io/hydrogym/docs/api/jax/env_core

Get information about where environment files are stored for debugging purposes.

```python
def get_environment_files_info() -> Dict
```

--------------------------------

### SimpleRewardLogger.__init__

Source: https://dynamicslab.github.io/hydrogym/docs/api/nek/nek_lib/reward_logger

Initializes the SimpleRewardLogger with a log file path.

```APIDOC
## SimpleRewardLogger.__init__

### Description
Initialize simple reward logger.

### Method
Constructor

### Signature
```python
def __init__(log_file: Union[str, Path])
```

### Parameters
- **log_file** (Union[str, Path]) - Required - Path to the CSV log file
```

--------------------------------

### mpi_split

Source: https://dynamicslab.github.io/hydrogym/docs/api/nek/env

Split MPI world into master/worker inter-communicator. This is a utility function for MPI setup.

```APIDOC
## mpi_split

### Description
Split MPI world into master/worker inter-communicator.

### Method
mpi_split (function)

### Signature
```python
def mpi_split(comm_world: MPI.Comm, nproc: Optional[int] = None) -> MPI.Comm
```

### Parameters
- **comm_world** (MPI.Comm) - Required - MPI communicator
- **nproc** (Optional[int]) - Optional - Expected number of Nek workers (for validation)

### Returns
- MPI.Comm - Inter-communicator between controller and workers
```

--------------------------------

### Prepare workspace and train SB3 agent with MPMD

Source: https://dynamicslab.github.io/hydrogym/docs/examples/maia/getting_started

Workflow for training an SB3 agent: first prepare the workspace, then run training with MPMD execution, and finally monitor with TensorBoard. The training command uses PPO with 100000 timesteps.

```bash
# First, prepare workspace  
python prepare_workspace.py --env Cylinder_2D_Re200 --work-dir ./train_run  

# Then train with MPMD execution  
cd train_run  
mpirun -np 1 python ../train_sb3_maia.py --env Cylinder_2D_Re200 --algo PPO --total-timesteps 100000 : -np 1 maia properties.toml  

# Monitor training  
tensorboard --logdir logs/
```

--------------------------------

### Pull and run HydroGym Docker image

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Pull the pre-built HydroGym Docker image and start a container.

```shell
docker pull clagemann/hydrogym-nvhpc-26.1_cuda-12.9_hopper_blackwell  

docker run -it clagemann/hydrogym-nvhpc-26.1_cuda-12.9_hopper_blackwell  

```

--------------------------------

### NEK_INIT.init_restart

Source: https://dynamicslab.github.io/hydrogym/docs/api/nek/nek_lib/nek_utils

Copies the restart file to the target folder only if the restart file does not already exist.

```APIDOC
## NEK_INIT.init_restart

### Description
Copy the restart file to the target folder only if RSTART NOT EXIST.

### Method
Method

### Signature
```python
def init_restart()
```

```

--------------------------------

### Train NekEnv with Stable-Baselines3 wrappers

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/single_environment

Sets up a NekEnv environment with SB3's Monitor, DummyVecEnv, and VecNormalize wrappers, then trains a PPO model. The environment is saved as 'vec_normalize.pkl' after training. This demonstrates direct SB3 compatibility without additional wrapper code.

```python
from hydrogym.nek import NekEnv  
from stable_baselines3 import PPO  
from stable_baselines3.common.monitor import Monitor  
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize  

env_config = {  
    'environment_name': 'TCFmini_3D_Re180',  
    'nproc': 10,  
}  

env = NekEnv(env_config=env_config)  
env = Monitor(env)  
env = DummyVecEnv([lambda: env])  
env = VecNormalize(env, norm_obs=True, norm_reward=True)  

model = PPO("MlpPolicy", env, verbose=1)  
model.learn(total_timesteps=100_000)  
model.save("ppo_tcfmini")  
env.save("vec_normalize.pkl")  
```

--------------------------------

### Verify NEK5000 installation and configuration

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/getting_started

Commands to verify the availability of the nek5000 binary or define its path within a Python configuration dictionary.

```bash
# Verify that NEK5000 is available  

which nek5000
```

```python
# Or specify the path in the environment config  

env_config = {  

    'environment_name': 'TCFmini_3D_Re180',  

    'nek_path': '/opt/nek5000/bin/nek5000',  

    'nproc': 10,  

}
```

--------------------------------

### Run test and training scripts with mpirun

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/single_environment

Commands to test the environment and train with PPO using mpirun. The test command runs test_nek_direct.py with 100 steps, and the training command runs train_sb3_nek_direct.py with PPO for 100000 timesteps. Both use -np 1 for the Python controller and -np 10 for the NEK5000 solver.

```bash
# Test the environment  
mpirun -np 1 python test_nek_direct.py --steps 100 : -np 10 nek5000  

# Train with PPO  
mpirun -np 1 python train_sb3_nek_direct.py \  
    --env TCFmini_3D_Re180 \  
    --algo PPO \  
    --total-timesteps 100000 \  
    : -np 10 nek5000  
```

--------------------------------

### Sphere.__init__

Source: https://dynamicslab.github.io/hydrogym/docs/api/maia/envs/sphere

Initializes the sphere environment with a configuration dictionary.

```APIDOC
## Sphere.__init__

### Description
Initializes the sphere environment.

### Method
Constructor

### Signature
```python
def __init__(env_config: Dict)
```

### Parameters
- **env_config** (Dict) - Required - Environment configuration dictionary.

### Returns
None
```

--------------------------------

### HFDataManager.prepare_working_directory

Source: https://dynamicslab.github.io/hydrogym/docs/api/data_manager

Creates a working directory and populates it with solver-specific symlinks based on the detected or specified solver profile.

```APIDOC
## HFDataManager.prepare_working_directory

### Description
Create _work_dir_ and populate it with solver-specific symlinks. The mapping of source files/directories to their target locations inside _work_dir_ is defined by the `workspace_files` and `workspace_dirs` entries in `SOLVER_PROFILES`. If _profile_ is `None` the profile is auto-detected from sentinel files already present in _env_path_, then falls back to `self.fallback_profile`.

### Method
Method

### Signature
```python
def prepare_working_directory(env_path: str,
                              work_dir: str,
                              profile: Optional[str] = None) -> Dict[str, str]
```

### Parameters
- **env_path** (str) - Required - Path to the cached environment data (as returned by `get_environment_path`).
- **work_dir** (str) - Required - Target working directory (created if it does not exist).
- **profile** (Optional[str]) - Optional - Solver profile key (e.g. `'MAIA_LB'`). Auto-detected when `None`.

### Returns
Dictionary with at least `'work_dir'` and `'env_data_path'` keys. Additional keys may be present depending on the profile (e.g. `'properties_file'` for MAIA_LB).
```

--------------------------------

### Define action with jnp.zeros

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/channel

Creates a zero action vector with the correct dimension from the environment parameters. This is the default action to start interacting with the environment.

```python
action = jnp.zeros((params.action_dim,))  
```

--------------------------------

### ZPGTBLBase

Source: https://dynamicslab.github.io/hydrogym/docs/api/maia/envs/turbulent_boundary_layer

Base class for zero-pressure-gradient turbulent boundary layer environments. Provides methods for observation configuration, normalization setup, and reward computation.

```APIDOC
## ZPGTBLBase

### Description
Base class for zero-pressure-gradient turbulent boundary layer environments. Uses the structured-grid m-AIA solver (MAIA_STRCTRD). Both actuation variants (jet and surface wave) receive the solver force output [F_x, F_y, F_z, C_P, area] (nDim + 2 elements). Forces and power are pre-normalised by the solver using dynamic pressure q_inf = rho_inf * 0.5 * U_inf^2 (forces) and q_inf * U_inf (power), so dividing by the returned area gives the final dimensionless drag and power coefficients.

### Methods

#### configure_observations

```
def configure_observations() -> None
```

Configure the number of observations. Overrides the base class to account for the extended solver force output [F_x, F_y, F_z, C_P, area] (nDim + 2 elements per boundary instead of nDim).

#### setup_normalization

```
def setup_normalization() -> None
```

Set up observation normalization. Overrides the base class to handle the extended force vector [F_x, F_y, F_z, C_P, area] (nDim + 2 elements per boundary). For the 'U_inf' strategy, free-stream values computed from the isentropic relations in __init__ are used: velocities scaled by U_T (total free-stream speed), density by rho_inf, pressure by the dynamic pressure q_inf = 0.5 * rho_inf * U_T^2. Force-slot entries use loc = 0, scale = 1 since the solver already normalises them by q_inf (forces) / q_inf * U_T (power).

#### get_reward

```
def get_reward() -> Tuple[float, Dict]
```

Compute the unified reward: R = -C_D - omega * C_P. The solver returns [F_x, F_y, F_z, C_P, area]. Forces and power are pre-normalised by q_inf and q_inf * U_inf respectively; dividing by area yields the final dimensionless coefficients: C_D = forces[0] / area, C_P = forces[3] / area. For the jet environment C_P = 0 and the reward reduces to -C_D.

**Returns**: Tuple of (reward, info_dict). info_dict contains 'forces', 'C_D', and 'C_P'.
```

--------------------------------

### PDEBase.init_bcs

Source: https://dynamicslab.github.io/hydrogym/docs/api/core

Initialize any boundary conditions for the PDE environment.

```APIDOC
## PDEBase.init_bcs

### Description
Initialize any boundary conditions for the PDE. This is an abstract method.

### Method
`init_bcs`

### Signature
`def init_bcs()`
```

--------------------------------

### HFDataManager.get_environment_path

Source: https://dynamicslab.github.io/hydrogym/docs/api/data_manager

Gets the path to environment files, downloading from HF Hub if necessary. The solver profile is detected automatically from sentinel files before any download occurs.

```APIDOC
## HFDataManager.get_environment_path

### Description
Get path to environment files, downloading from HF Hub if necessary. The solver profile is detected automatically from sentinel files before any download occurs.

### Method
Method

### Signature
```python
def get_environment_path(env_name: str, force_download: bool = False) -> str
```

### Parameters
- **env_name** (str) - Required - Environment name (e.g. `'Cylinder_2D_Re200'`).
- **force_download** (bool) - Optional - Force re-download even if cached.

### Returns
Path to the local environment directory.
```

--------------------------------

### Implement basic RL training loop with SB3

Source: https://dynamicslab.github.io/hydrogym/docs/quickstart

A basic training loop using Stable Baselines3 PPO with a customized MAIA environment and probe locations.

```python
# See examples/maia/getting_started/train_sb3_maia.py for full implementation  

import numpy as np  

import hydrogym.maia as maia  

from stable_baselines3 import PPO  

from stable_baselines3.common.monitor import Monitor  

from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize  

  

# Define probes  

probe_locations = []  

for x in np.linspace(1.0, 8.0, 8):  

    for y in np.linspace(-1.0, 1.0, 5):  

        probe_locations.extend([x, y])  

  

  

def make_env():  

    env = maia.from_hf(  

        "Cylinder_2D_Re200",  

        use_clean_cache=False,  

        probe_locations=probe_locations,  

        obs_normalization_strategy="U_inf",  

    )  

    return Monitor(env)  

  

  

env = DummyVecEnv([make_env])  

env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10.0)  

  

model = PPO("MlpPolicy", env, verbose=1, tensorboard_log="./logs")  

model.learn(total_timesteps=100000)  

  

model.save("ppo_cylinder")  

env.save("vec_normalize.pkl")
```

--------------------------------

### Build NEK5000 tools

Source: https://dynamicslab.github.io/hydrogym/docs/installation/nek5000

Compile the mesh partitioner and other required tools.

```bash
cd ~/Nek5000/tools  

./maketools genmap     # mesh partitioner (required)  
```

--------------------------------

### Configure PETSc with SLEPc for stability analysis

Source: https://dynamicslab.github.io/hydrogym/docs/installation/firedrake

Adds SLEPc support to the PETSc configuration, required for stability analysis scripts in cylinder, cavity, and step examples.

```shell
python3 ../firedrake-configure --show-petsc-configure-options | \  
    xargs -L1 ./configure --download-slepc  

export SLEPC_DIR=$PETSC_DIR/$PETSC_ARCH  

```

--------------------------------

### init_restart Method

Source: https://dynamicslab.github.io/hydrogym/docs/api/nek/nek_lib/nek_utils

Copies the restart file to the target folder if it does not already exist.

```python
def init_restart()
```

--------------------------------

### Complete train + evaluate workflow

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/hydrogym_control

Full lifecycle from training with Monitor, VecNormalize, and PPO, to saving the model and normalization statistics, then loading them for evaluation with integrate().

```python
from hydrogym.nek import NekEnv, integrate
from stable_baselines3 import PPO
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize

# --- Phase 1: training ---
env = Monitor(NekEnv(env_config=config))
env = VecNormalize(DummyVecEnv([lambda: env]), norm_obs=True, norm_reward=True)

model = PPO("MlpPolicy", env)
model.learn(total_timesteps=50_000)
model.save("model.zip")
env.save("vec_normalize.pkl")

# --- Phase 2: evaluation ---
base_eval = NekEnv(env_config=config)
env_eval = VecNormalize.load("vec_normalize.pkl", DummyVecEnv([lambda: base_eval]))
model = PPO.load("model.zip")

integrate(env_eval, controller=model, num_steps=1000)
```

--------------------------------

### Basic usage of integrate() with classical and SB3 controllers

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/hydrogym_control

Shows how to create a NekEnv from Hugging Face, define a classical opposition control callable, and run integrate() with either the callable or a trained SB3 PPO model. integrate() handles the reset()/step() loop internally; for SB3 models it calls model.predict() automatically.

```python
from hydrogym.nek import NekEnv, integrate  

env = NekEnv.from_hf('TCFmini_3D_Re80', nproc=10)  

# Classical controller: opposition control
def opposition_control(t, obs, env):  
    return -obs[:env.action_space.shape[0]]  

integrate(env, controller=opposition_control, num_steps=1000)  

# Or pass a trained SB3 model directly
from stable_baselines3 import PPO
model = PPO.load("model.zip")
integrate(env, controller=model, num_steps=1000)
```

--------------------------------

### SSH tunnel to ParaView server

Source: https://dynamicslab.github.io/hydrogym/docs/installation/maia

Opens a tunnel from your local workstation to the compute node's ParaView port, allowing connection to localhost:11111 with ParaView 5.13 after starting the server.

```bash
ssh USER@REMOTE.WORKSTATION.ADDRESS -L 11111:node042:11111  

```

--------------------------------

### Initialize Environment

Source: https://dynamicslab.github.io/hydrogym/docs/examples/jax/ppo

Resets the environment using JAX random keys to prepare for the training loop.

```python
        # INIT ENV  

        rng, _rng = jax.random.split(rng)  

        reset_rng = jax.random.split(_rng, config["NUM_ENVS"])  

        obsv, env_state = env.reset(reset_rng, env_params)
```

--------------------------------

### Create PettingZoo environment from NekEnv

Source: https://dynamicslab.github.io/hydrogym/docs/examples/nek5000/pettingzoo

Creates a PettingZoo Parallel API environment by wrapping a NekEnv instance with make_pettingzoo_env. The example shows resetting the environment, generating actions per agent, and stepping with the resulting observations, rewards, termination flags, truncation flags, and infos.

```python
from hydrogym.nek import NekEnv  

from hydrogym.nek.pettingzoo_env import make_pettingzoo_env  

  

env_config = {'environment_name': 'TCFmini_3D_Re180', 'nproc': 10}  

base_env = NekEnv(env_config=env_config)  

  

env = make_pettingzoo_env(base_env)  

  

observations, infos = env.reset()  

actions = {agent: policy(obs) for agent, obs in observations.items()}  

observations, rewards, terminated, truncated, infos = env.step(actions)  
```

--------------------------------

### Load environment with local fallback

Source: https://dynamicslab.github.io/hydrogym/docs/installation/nek5000

Initialize a NEK5000 environment using a local directory for offline HPC usage.

```python
env = NekEnv.from_hf(  
    'TCFmini_3D_Re180',  
    nproc=10,  
    local_fallback_dir='/scratch/$USER/nek_envs',  
)  
```

--------------------------------

### Docker images for NVIDIA/AMD GPUs

Source: https://dynamicslab.github.io/hydrogym/docs/installation/maia

Pulls and runs pre-built one Docker images for m-AIA with HydroGym dependencies. Images are tagged by GPU hardware: CUDA 12.9 Hopper/Blackwell, CUDA 12.9 Turing/Ampere, and ROCm 6.3. Use --gpus all for NVIDIA; AMD requires no extra flag. Inside the container HydroGym and m-AIA are installed and on PATH.

```bash
# NVIDIA GPU (CUDA 12.9 Hopper & Blackwell)  
docker pull clagemann/hydrogym-nvhpc-26.1_cuda-12.9_hopper_blackwell  
docker run -it --gpus all clagemann/hydrogym-nvhpc-26.1_cuda-12.9_hopper_blackwell  
  
# NVIDIA GPU (CUDA 12.9 Turing & Ampere)  
docker pull clagemann/hydrogym-nvhpc-26.1_cuda-12.9_turing_ampere  
docker run -it --gpus all clagemann/hydrogym-nvhpc-26.1_cuda-12.9_turing_ampere  
  
# AMD GPU (ROCm 6.3)  
docker pull clagemann/hydrogym-rocm-6.3.3:latest  
docker run -it clagemann/hydrogym-rocm-6.3.3:latest  

```

--------------------------------

### CubeBase __init__ method

Source: https://dynamicslab.github.io/hydrogym/docs/api/maia/envs/cube

Initialize the cube base environment with an environment configuration dictionary.

```python
def __init__(env_config: Dict)
```
Tokens:
Raw