SSTODE: Ocean-Atmosphere Physics-Informed Neural ODEs for Sea Surface Temperature Prediction

SSTODE is a physics-informed neural forecasting model for Sea Surface Temperature (SST). It explicitly incorporates advection–diffusion dynamics and energy flux terms into a neural ODE framework, enabling accurate and robust spatiotemporal forecasting across diverse oceanographic regions and timescales.

Installation

# Create and activate conda environment
conda env create -f environment.yml
conda activate SSTODE

# Install in development mode
python setup.py develop  # or pip install -e .

Dataset Preparation

1. OceanVP

   1. Download Dataset from Source Dataset：https://github.com/zhenglab/OceanVP
   2. Variable Used: water_temperature_0m

2. ERA5

   1. Download Data from ERA5 Dataset: ERA5 hourly data on single levels from 1940 to present (https://cds.climate.copernicus.eu/datasets/reanalysis-era5-single-levels?tab=overview)
   2. Select Variables:
      • Sea Surface Temperature (SST)
      • Surface net thermal radiation (LW)
      • Surface net solar radiation(SW)
      • Surface latent heat flux(SLF)
      • Surface sensible heat flux(SHF) variables
   3. Downsample with regrid.py script from https://github.com/pangeo-data/WeatherBench

3. Directory Structure

   data
   ├── era5
   │   ├── constants
   │   │   └── constants_6.625deg.nc
   │   ├── sea_surface_temperature
   │   │   ├── sea_surface_temperature_2000.nc
   │   ├── surface_sensible_heat_flux
   │   ├── surface_latent_heat_flux
   │   ├── surface_net_solar_radiation
   │   ├── surface_net_thermal_radiation
   
   ├── ocean
   │   ├── constants
   │   │   └── constants_6.625deg.nc
   │   ├── salinity_depth_0m
   │   │   ├── salinity_depth_0m_1994_32_64.nc
   │   │   ├── salinity_depth_0m_1995_32_64.nc
   │   │   ├── xxx.nc

Get Started

Run training script

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) \
python tools/train.py \
  -d ocean_t0_32_64 \
  -c configs/oceanvp/t0_32_64/QsourceODE.py \
  --ex_name oceanvp_t0_QsourceODE_6h_1_7 \
  --collect_type "qode" \
  --data_interval "6h" \
  --source_type "wQ" \
  --k "learn" \
  --pre_seq_length 1 \
  --aft_seq_length 7 \
  --total_length 8 \
  --auto_resume

Run evaluation script:

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) \
python tools/test.py \
  -d ocean_t0_32_64 \
  -c configs/oceanvp/t0_32_64/QsourceODE.py \
  --ex_name oceanvp_t0_QsourceODE_6h_1_7 \
  --collect_type "qode" \
  --data_interval "6h" \
  --source_type "wQ" \
  --k "learn" \
  --pre_seq_length 1 \
  --aft_seq_length 7 \
  --total_length 8 \
  --auto_resume