# DynamisLab

**Machine Learning for Computational Fluid Dynamics**

DynamisLab is a research framework for applying reinforcement learning and machine learning techniques to computational fluid dynamics problems. Built on top of [CelerisLab](https://github.com/frank14f/CelerisLab), it provides standardized environments and training pipelines for active flow control tasks.

## Features

- 🌊 **CFD Environments**: Gymnasium-compatible environments for flow control
- 🤖 **RL Integration**: Ready-to-use with Stable-Baselines3 and other RL libraries
- 🚀 **GPU Acceleration**: Leverages CelerisLab's CUDA-accelerated LBM solver
- 📊 **Experiment Tracking**: Built-in TensorBoard integration
- 🔧 **Modular Design**: Clean separation of environments, configs, and training scripts
- 📦 **Standard Structure**: Follows Python packaging best practices (src layout)

## Project Structure

```
DynamisLabNew/
├── src/                            # Source code (src layout)
│   ├── __init__.py
│   ├── config.py                   # Configuration management
│   └── environments/               # Gymnasium environments
│       ├── __init__.py
│       └── cfd_env.py              # CFD flow control environment
├── scripts/                        # Training and evaluation scripts
│   └── train_ppo.py                # PPO training script
├── configs/                        # Configuration files
│   ├── config_cuda.json            # CUDA settings
│   ├── config_flowfield.json       # Flow field parameters
│   └── config_gym.json             # Environment settings
├── models/                         # Trained model checkpoints (gitignored)
├── output/                         # Training data and results (gitignored)
├── tensorboard/                    # TensorBoard logs (gitignored)
├── docs/                           # Documentation
├── requirements.txt                # Python dependencies
├── pyproject.toml                  # Package configuration
└── README.md                       # This file
```

## Installation

### Prerequisites

- Python 3.8+
- NVIDIA GPU with CUDA support
- CUDA Toolkit 11.0+

### Step 1: Clone the repository

```bash
git clone --recurse-submodules <your-repo-url> DynamisLab
cd DynamisLab
```

> **Note**: If CelerisLab is a submodule, use `--recurse-submodules` to clone it automatically.

### Step 2: Install CelerisLab

#### Option A: Install from submodule (recommended for development)

```bash
cd CelerisLab
pip install -e .
cd ..
```

#### Option B: Install from pip (if published)

```bash
pip install CelerisLab
```

### Step 3: Install DynamisLab dependencies

```bash
pip install -r requirements.txt
```

### Step 4: Install DynamisLab in development mode

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

## Quick Start

### Training a PPO Agent

Train a Proximal Policy Optimization agent for flow control:

```bash
python scripts/train_ppo.py \
    --run-name my_first_run \
    --device-id 0 \
    --total-timesteps 100 \
    --n-steps 3600 \
    --activation sin
```

**Arguments:**
- `--run-name`: Name for this training run (used for saving models and logs)
- `--device-id`: CUDA device ID for CFD simulation
- `--cuda-device`: CUDA device ID for PyTorch training (can be different from --device-id)
- `--total-timesteps`: Number of training iterations
- `--n-steps`: Environment steps per training iteration
- `--activation`: Activation function (`sin`, `tanh`, or `relu`)

### Monitoring Training

```bash
tensorboard --logdir tensorboard/
```

Then open http://localhost:6006 in your browser.

### Using the Environment Programmatically

```python
from environments import CFDFlowControlEnv
from config import load_celeris_configs

# Load configurations
config_cuda, config_field = load_celeris_configs()

# Create environment
env = CFDFlowControlEnv(
    device_id=0,
    config_cuda=config_cuda,
    config_field=config_field,
)

# Run episode
obs, info = env.reset()
for step in range(500):
    action = env.action_space.sample()  # Random action
    obs, reward, terminated, truncated, info = env.step(action)
    
    if terminated or truncated:
        break

env.close()
```

## Configuration

### CFD Configuration

Edit `configs/config_flowfield.json` to change flow parameters:

```json
{
  "viscosity": 0.01,          # Fluid viscosity
  "velocity": 0.1,            # Inlet velocity
  "field_dim_in_U": [400, 200, 1],  # Grid dimensions
  ...
}
```

### CUDA Configuration

Edit `configs/config_cuda.json` for GPU settings:

```json
{
  "threads_per_block": 256,
  "unit_dimensions": [16, 16, 1],
  ...
}
```

## Advanced Usage

### Resume Training

```bash
python scripts/train_ppo.py \
    --resume models/my_run_best.zip \
    --run-name my_run_continued
```

### Custom Hyperparameters

```bash
python scripts/train_ppo.py \
    --learning-rate 0.0003 \
    --gamma 0.99 \
    --batch-size 512 \
    --n-steps 7200
```

### Multi-GPU Setup

```bash
# CFD simulation on GPU 0, PyTorch training on GPU 1
python scripts/train_ppo.py \
    --device-id 0 \
    --cuda-device 1
```

## Environment Details

### CFDFlowControlEnv

The main environment for active flow control around a cylinder.

**Observation Space:**
- Dimensionality: `n_sensors × 2 × 2` (velocity components, current + derivative)
- Default: 12 dimensions (3 sensors × 2 velocities × 2)
- Normalized to zero mean and unit variance

**Action Space:**
- Dimensionality: `n_control_cylinders`
- Default: 3 (three controllable cylinders)
- Range: [-1, 1] (scaled internally to physical velocities)

**Reward:**
- Drag reduction: `-cd × 0.1`
- Lift minimization: `-|cl| × 0.05`
- Flow similarity: `-similarity_distance × 0.5`
- Total reward is sum of components

**Episode:**
- Max steps: 500 (configurable)
- Simulation runs at 800 LBM steps per environment step

## Development

### Project Guidelines

- Follow PEP 8 style guide
- Use type hints for function signatures
- Document classes and functions with docstrings
- Keep environments in `src/dynamis/environments/`
- Keep training scripts in `scripts/`
- Use `config.py` for all path and configuration management

### Adding a New Environment

1. Create new environment class in `src/dynamis/environments/`
2. Inherit from `gym.Env`
3. Register in `src/dynamis/environments/__init__.py`
4. Create corresponding training script in `scripts/`

### Running Tests

```bash
pytest tests/
```

## Citation

If you use DynamisLab in your research, please cite:

```bibtex
@software{dynamis2026,
  author = {Frank14f},
  title = {DynamisLab: Machine Learning for Computational Fluid Dynamics},
  year = {2026},
  url = {https://github.com/frank14f/DynamisLab}
}
```

Also cite CelerisLab:

```bibtex
@software{celerislab2026,
  author = {Frank14f},
  title = {CelerisLab: GPU-Accelerated Lattice Boltzmann Method Solver},
  year = {2026},
  url = {https://github.com/frank14f/CelerisLab}
}
```

## License

MIT License - see LICENSE file for details

## Acknowledgments

- Built on [CelerisLab](https://github.com/frank14f/CelerisLab) CFD solver
- Uses [Stable-Baselines3](https://github.com/DLR-RM/stable-baselines3) for RL
- Gymnasium API for standardized environments

## Contributing

Contributions are welcome! Please open an issue or pull request.

## Contact

For questions or issues, please open a GitHub issue or contact Frank14f.