S3GM: Learning spatiotemporal dynamics with a pretrained generative model

Datasets of Kuramoto-Sivashinsky equation (KSE) and Kolmogorov flow.

Description of KSE data:

Each file for KSE datasets contains 4 dimensions in the following order: (B*V)*T*X*C. Details are listed in the following table:

B
number of varying initial conditions

V
number of varying parameters

T
number of temporal frames

X
spatial resolution

C
number of variables in solution (C = 1 for KSE)

values of parameter used to generate training dataset
1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.2, 2.4, 2.6, 2.8, 3.0, 3.2, 3.4, 3.6, 3.8, 4.0, 4.2, 4.4, 4.6, 4.8, 5.0

values of parameter used to generate test dataset
1.1, 2.5, 3.2

Description of Kolmogorov flow data:

Each file for Kolmogorov flow contains 5 dimensions inthe following order: (B*Re*K)*T*X*X*C. Details are listed in the following table:

B
number of varying initial conditions

Re
number of varying Reynolds numbers

K
number of varying source terms (controled by the value of k)

T
number of temporal frames

X
spatial resolution

C
number of variables in solution (C = 2 for Kolmogorov flow)

values of Reynolds number used to generate training dataset
100, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 800, 850, 900, 950, 1000, 1050

values of Reynolds number used to generate test dataset

50, 125, 575, 1100, 1500

values of k used to generate training dataset
2, 3, 4, 5, 6, 7, 8

values of k used to generate test dataset

2, 4, 6, 8

Description of ERA5 data:

Training and testing dataset for ERA5 contains 5 dimensions inthe following order: 1*T*X*X*C, which is manually collected from https://cds.climate.copernicus.eu/datasets/reanalysis-era5-single-levels?tab=download. Note that the quantities in the datasets are already rescaled (the scale factors are saved in the scalar_era5.npy file, which is a 4x2 array recording the means and stds for the 4 quantities we used). Details are listed in the following table:

T
number of temporal frames

X
spatial resolution

C
number of variables in solution (C = 4 for ERA5)

time span for training dataset
1979-2022

time span for test dataset

2023

Pretrained checkpoints:

The .zip file contains the pretrained checkpoints for KSE, Kolmogorov flow and ERA5. Within the .zip file, the folder'kse_v0' is the checkpoint for KSE, 'kol_v0' is the checkpoint for Kolmogorov flow, and 'era5_v0' is the checkpoint for ERA5.

Source code:

The source code is upload as Github repository in https://github.com/lzy12301/S3GM