WADEN (Wave Differentiable propagator Engine)

WADEN is a fully differentiable ocean wave dynamic propagation module implemented in PyTorch. It is designed to simulate the advection of wave action density in both geophysical and spectral spaces, mirroring the core physics of established models like WAVEWATCH III®. WADEN has been fully validated against multiple WW3 regression test cases including tp1.1, tp1.2, tp1.4, and tp2.1.

The dynamical core employs a third-order accurate QUICKEST scheme combined with the ULTIMATE Total Variance Diminishing (TVD) limiter, ensuring both accuracy and numerical stability. It supports both Cartesian and Spherical coordinate systems and integrates key physical factors (e.g., FACVX, FACVY, FACTH, FACK) as trainable torch.nn.Parameters.

This makes waden a powerful tool for researchers in physical oceanography and machine learning, enabling the development of hybrid wave models and facilitating gradient-based data assimilation or parameter estimation.

Test Case tp2.1 - 2D Cartesian Propagation: 2D Cartesian wave propagation showing excellent agreement between WW3 and WADEN in simulating wave energy advection across a uniform grid.

Test Case tp1.4 - Depth-induced Refraction: Depth-induced refraction demonstration where waves bend and focus due to varying bathymetry, showing excellent agreement between WW3 and WADEN.

Key Features

• Fully Differentiable: Built entirely in PyTorch, allowing for end-to-end backpropagation.
• High-Order Accuracy: Implements the 3rd-order QUICKEST/ULTIMATE scheme for robust advection.
• Physics-Informed: Based on the well-validated advection schemes from WAVEWATCH III®.
• Flexible: Supports Cartesian and Spherical coordinates, with configurable boundary conditions.
• Trainable Parameters: Key physical coefficients (FACVX, FACVY, FACTH, FACK, FDG, FACHFA) are exposed as nn.Parameters.
• Comprehensive Validation: Validated against multiple WW3 test cases (tp1.1, tp1.2, tp1.4, tp2.1).
• YAML Configuration: Easy-to-use configuration system for model parameters and simulation settings.

Project Structure

waden/
├── README.md                          # This file
├── setup.py                           # Package installation configuration
├── waden/                             # Main package directory
│   ├── __init__.py                   # Package initialization
│   ├── core.py                       # DynamicPropagator class
│   └── utils.py                      # Utility functions
├── configs/                           # Configuration files for different test cases
│   ├── nww_tp1.1.yaml               # Configuration for test case tp1.1
│   ├── nww_tp1.2.yaml               # Configuration for test case tp1.2
│   ├── nww_tp1.4.yaml               # Configuration for test case tp1.4
│   └── nww_tp2.1.yaml               # Configuration for test case tp2.1
├── hyparaFile_tp*.yaml               # Hyperparameter files for different cases
├── data/                             # Test case data from WW3
│   ├── ww3_tp1.1/                   # Test case 1.1 data
│   ├── ww3_tp1.2/                   # Test case 1.2 data
│   ├── ww3_tp1.4/                   # Test case 1.4 data
│   └── ww3_tp2.1/                   # Test case 2.1 data
├── outputs/                          # Simulation results
│   └── nww_tp*/                     # Results for each test case
├── validation/                       # Validation scripts and utilities
│   ├── validate.py                  # Main validation script
│   ├── NWW_visual*_WW3.py          # Visualization and comparison scripts
│   └── src/                         # Validation utilities
└── visual_results/                   # Generated visualization results
    ├── ww3_vs_nww_tp1_2.gif         # Comparison animation for tp1.2
    └── ww3_vs_nww_tp2_1.gif         # Comparison animation for tp2.1

Installation

git clone https://github.com/GaryYang77/waden.git
cd waden
pip install -e .

Dependencies

• PyTorch >= 1.9.0
• NumPy >= 1.19.0
• netCDF4 >= 1.5.0
• OmegaConf >= 2.1.0
• PyYAML >= 5.4.0
• matplotlib >= 3.3.0 (for visualization)
• xarray >= 0.19.0
• scipy >= 1.7.0

Configuration System

WADEN uses a flexible YAML-based configuration system with two types of configuration files:

1. Base Configuration (e.g., configs/nww_tp1.1.yaml)

case_name: "nww_tp1.1"

data:
  grid_file: "data/ww3_tp1.1/input/ww3.200001.nc"
  spec_file: "data/ww3_tp1.1/input/ww3.Center_200001_spec.nc"
  initial_condition_type: "point_center"

model_params:
  hypara_file: "hyparaFile_tp1.1.yaml"
  device: "cpu"
  precision: "half"

simulation:
  steps: 144

output:
  directory: "outputs"
  save_visualization: false
  results_nc_filename: "NWW_output.nc"

2. Hyperparameter Configuration (e.g., hyparaFile_tp1.1.yaml)

baseParams:
  coordinate: "Spherical"        # 'Spherical' or 'Cartesian'
  boundary: "simple"             # 'circle', 'close', 'absorb', 'open', 'simple'
  radius_earth: 6366197.7236758135
  dt_max: 3600                   # Maximum time step (s)
  dt_xy: 3600                    # Dynamic integration time step (s)
  dt_kth: 3600                   # Spectral integration time step (s)
  dt_stot: 3600                  # Source term integration time step (s)
  dynamic_intFrame: 'WW3'        # Integration framework: 'simple' or 'WW3'
  CFL_max: 0.75                  # CFL stability criterion
  FLCX: true                     # Enable X-direction advection
  FLCY: false                    # Enable Y-direction advection
  FLCTH: false                   # Enable directional advection
  FLCK: false                    # Enable frequency advection

hypersParams:
  USE_File: true                 # Use file-based parameters
  multi_params: false            # Use spatially varying parameters
  FACVX: 1.05                    # X-velocity factor
  FACVY: 1.0                     # Y-velocity factor
  FACTH: 1.0                     # Directional factor
  FACK: 1.0                      # Frequency factor
  FDG: 1.0                       # Diffusion factor
  FACHFA: 1.0                    # High-frequency factor

Quick Start

Method 1: Direct API Usage

Here is a minimal example of how to use waden to propagate a wave spectrum:

import torch
from waden import DynamicPropagator

# 1. Define the grid and physical parameters
device = 'cuda' if torch.cuda.is_available() else 'cpu'
lon = torch.linspace(0, 10, 64)      # Longitude points
lat = torch.linspace(0, 10, 64)      # Latitude points
theta = torch.linspace(0, 2 * torch.pi, 24)  # Direction bins (radians)
freq = torch.linspace(0.04, 1.0, 30)         # Frequency bins (Hz)

# 2. Create a hyperparameter configuration file
hypara_config = """
baseParams:
  coordinate: "Cartesian"
  radius_earth: 6366197.7236758135
  boundary: "absorb"
  dt_max: 600
  dt_xy: 600
  dt_kth: 600
  dt_stot: 600
  dynamic_intFrame: 'WW3'
  CFL_max: 0.75
  yfirst: true
  FLCX: true
  FLCY: true
  FLCTH: true
  FLCK: true
  learn_FieldStot: False

hypersParams:
  USE_File: true
  multi_params: false
  FACVX: 1.0
  FACVY: 1.0
  FACTH: 1.0
  FACK: 1.0
  FDG: 1.0
  FACHFA: 1.0
"""

with open("config.yaml", "w") as f:
    f.write(hypara_config)

# 3. Initialize the propagator
propagator = DynamicPropagator(
    lon=lon, lat=lat, theta=theta, freq=freq,
    config_file='config.yaml',
    device=device
)

# 4. Prepare initial conditions
batch_size = 1
initial_spectra = torch.rand(batch_size, len(lon), len(lat), len(theta), len(freq)).to(device)
depth = (torch.ones(batch_size, len(lon), len(lat)) * 50.0).to(device)  # 50m depth

# 5. Run one propagation step
propagated_spectra = propagator(
    spectra=initial_spectra,
    depth=depth
)

print("Propagation successful!")
print("Initial spectra shape:", initial_spectra.shape)
print("Propagated spectra shape:", propagated_spectra.shape)

Method 2: Using the Validation System

You can run pre-configured test cases using the validation system:

# Run all test cases
python validation/validate.py --config configs/nww_tp1.1.yaml
python validation/validate.py --config configs/nww_tp1.2.yaml
python validation/validate.py --config configs/nww_tp1.4.yaml
python validation/validate.py --config configs/nww_tp2.1.yaml

# Generate comparison visualizations
python validation/NWW_visual1_2_WW3.py  # Creates tp1.2 comparison
python validation/NWW_visual2_1_WW3.py  # Creates tp2.1 comparison

Advanced Usage

Trainable Parameters

All physical parameters can be made trainable for machine learning applications:

# Initialize with trainable parameters
propagator = DynamicPropagator(
    lon=lon, lat=lat, theta=theta, freq=freq,
    config_file='config.yaml',
    trainable=True,  # Make parameters trainable
    device=device
)

# Access trainable parameters
for name, param in propagator.named_parameters():
    print(f"{name}: {param.shape}, requires_grad={param.requires_grad}")

# Example gradient-based optimization
optimizer = torch.optim.Adam(propagator.parameters(), lr=0.01)
target_spectra = ...  # Your target data
loss = torch.nn.MSELoss()(propagated_spectra, target_spectra)
loss.backward()
optimizer.step()

Contributing

Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.

Development Setup

git clone https://github.com/GaryYang77/waden.git
cd waden
pip install -e ".[dev]"  # Install with development dependencies

License

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you use waden in your research, please consider citing:

@software{yang2025waden,
  title={WADEN: Wave Differentiable propagator Engine},
  author={Yang},
  year={2025},
  url={https://github.com/GaryYang77/waden},
  version={0.1.0}
}

Contact

• Author: Gary Yang
• Email: yanggy25@mail2.sysu.edu.cn
• GitHub: @GaryYang77

For questions, bug reports, or feature requests, please open an issue on GitHub.