Package: pdebench
PDEBenchThe code repository for the NeurIPS 2022 paperPDEBench: An Extensive Benchmark for Scientific Machine Learning:tada:SimTech Best Paper Award 2023:confetti_ball:PDEBench provides a diverse and comprehensive set of benchmarks for scientific machine learning, including challenging and realistic physical problems. This repository consists of the code used to generate the datasets, to upload and download the datasets from the data repository, as well as to train and evaluate different machine learning models as baselines. PDEBench features a much wider range of PDEs than existing benchmarks and includes realistic and difficult problems (both forward and inverse), larger ready-to-use datasets comprising various initial and boundary conditions, and PDE parameters. Moreover, PDEBench was created to make the source code extensible and we invite active participation from the SciML community to improve and extend the benchmark.Created and maintained by Makoto Takamoto<makoto.takamoto@neclab.eu, takamtmk@gmail.com>, Timothy Praditia<timothy.praditia@iws.uni-stuttgart.de>, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert.Datasets and Pretrained ModelsWe also provide datasets and pretrained machine learning models.PDEBench Datasets:https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2986PDEBench Pre-Trained Models:https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2987DOIsInstallationUsing pipLocally:pipinstall--upgradepipwheel
pipinstall.From PyPI:pipinstallpdebenchTo include dependencies for data generation:pipinstall"pdebench[datagen310]"pipinstall".[datagen310]"# locallyorpipinstall"pdebench[datagen39]"pipinstall".[datagen39]"# locallyGPU SupportFor GPU support there are additional platform-specific instructions:For PyTorch, the latest version we support is v1.13.1see previous-versions/#linux - CUDA 11.7.For JAX, which is approximately 6 times faster for simulations than PyTorch in our tests,see jax#pip-installation-gpu-cuda-installed-via-pipInstallation using conda:If you like you can also install dependencies using anaconda, we suggest to usemambaforgeas a distribution. Otherwise you may have toenable the conda-forgechannel for the following commands.Starting from a fresh environment:conda create -n myenv python=3.9
conda activate myenvInstall dependencies for model training:conda install deepxde hydra-core h5py -c conda-forgeAccording to your hardware availability, either install PyTorch with CUDA support:see previous-versions/#linux - CUDA 11.7.conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidiaor CPU only binaries.conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 cpuonly -c pytorchOptional dependencies for data generation:conda install clawpack jax jaxlib python-dotenvConfiguring DeepXDEIn our tests we used PyTorch as backend for DeepXDE. Pleasefollow the documentationto enable this.Data GenerationThe data generation codes are contained indata_gen:gen_diff_react.pyto generate the 2D diffusion-reaction data.gen_diff_sorp.pyto generate the 1D diffusion-sorption data.gen_radial_dam_break.pyto generate the 2D shallow-water data.gen_ns_incomp.pyto generate the 2D incompressible inhomogenous Navier-Stokes data.plot.pyto plot the generated data.uploader.pyto upload the generated data to the data repository..envis the environment data to store Dataverse URL and API token to upload the generated data. Note that the filename should be strictly.env(i.e. remove theexamplefrom the filename)configsdirectory contains the yaml files storing the configuration for the simulation. Arguments for the simulation are problem-specific and detailed explanation can be found in the simulation scripts.srcdirectory contains the simulation scripts for different problems:sim_diff_react-pyfor 2D diffusion-reaction,sim_diff_sorp.pyfor 1D diffusion-sorption, andswefor the shallow-water equation.Data Generation for 1D Advection/Burgers/Reaction-Diffusion/2D DarcyFlow/Compressible Navier-Stokes EquationsThe data generation codes are contained indata_gen_NLE:utils.pyutil file for data generation, mainly boundary conditions and initial conditions.AdvectionEqdirectory with the source codes to generate 1D Advection equation training samplesBurgersEqdirectory with the source codes to generate 1D Burgers equation training samplesCompressibleFluiddirectory with the source codes to generate compressible Navier-Stokes equations training samplesReactionDiffusionEqdirectory with the source codes to generate 1D Reaction-Diffusion equation training samples (Note:DarcyFlow data can be generated by run_DarcyFlow2D.shin this folder.)savedirectory saving the generated training samplesA typical example to generate training samples (1D Advection Equation):
(indata_gen/data_gen_NLE/AdvectionEq/)python3advection_multi_solution_Hydra.py+multi=beta1e0.yamlwhich is assumed to be performed in each directory.Examples for generating other PDEs are provided inrun_trainset.shin each PDE's directories.
The config files for Hydra are stored inconfigdirectory in each PDE's directory.Data Transformaion and Merge into HDF5 format1D Advection/Burgers/Reaction-Diffusion/2D DarcyFlow/Compressible Navier-Stokes Equations save data as a numpy array.
So, to read those data via our dataloaders, the data transformation/merge should be performed.
This can be done usingdata_gen_NLE/Data_Merge.pywhose config file is located at:data_gen/data_gen_NLE/config/config.yaml.
After properly setting the parameters in the config file (type: name of PDEs, dim: number of spatial-dimension, bd: boundary condition),
the corresponding HDF5 file could be obtained as:python3Data_Merge.pyConfigurationYou can set the default values for data locations for this project by putting config vars like this in the.envfile:WORKING_DIR=~/Data/Working
ARCHIVE_DATA_DIR=~/Data/ArchiveThere is an example inexample.env.Data DownloadThe download scripts are provided indata_download. There are two options to download data.Usingdownload_direct.py(recommended)Retrieves data shards directly using URLs. Sample command for each PDE is given in the README file in thedata_downloaddirectory.Usingdownload_easydataverse.py(might be slow and you could encounter errors/issues; hence, not recommended!)Use the config files from theconfigdirectory that contains the yaml files storing the configuration. Any files in the dataset matchingargs.filenamewill be downloaded intoargs.data_folder.Baseline ModelsIn this work, we provide three different ML models to be trained and evaluated against the benchmark datasets, namelyFNO,U-Net, andPINN.
The codes for the baseline model implementations are contained inmodels:train_models_forward.pyis the main script to train and evaluate the model. It will call on model-specific script based on the input argument.train_models_inverse.pyis the main script to train and evaluate the model for inverse problems. It will call on model-specific script based on the input argument.metrics.pyis the script to evaluate the trained models based on various evaluation metrics described in our paper. Additionally, it also plots the prediction and target data.analyse_result_forward.pyis the script to convert the saved pickle file from the metrics calculation script into pandas dataframe format and save it as a CSV file. Additionally it also plots a bar chart to compare the results between different models.analyse_result_inverse.pyis the script to convert the saved pickle file from the metrics calculation script into pandas dataframe format and save it as a CSV file. This script is used for the inverse problems. Additionally it also plots a bar chart to compare the results between different models.fnocontains the scripts of FNO implementation. These are partly adapted from theFNO repository.unetcontains the scripts of U-Net implementation. These are partly adapted from theU-Net repository.pinncontains the scripts of PINN implementation. These utilize theDeepXDE library.inversecontains the model for inverse model based on gradient.configcontains the yaml files for the model training input. The default templates for different equations are provided in theargsdirectory. User just needs to copy and paste them to the args keyword in theconfig.yamlfile.An example to run the forward model training can be found inrun_forward_1D.sh, and an example to run the inverse model training can be found inrun_inverse.sh.Short explanations on the config argsmodel_name: string, containing the baseline model name, either 'FNO', 'Unet', or 'PINN'.if_training: bool, set True for training, or False for evaluation.continue_training: bool, set True to continute training from a checkpoint.num_workers: int, number of workers for the PyTorch dataloader.batch_size: int, training batch size.initial_step: int, number of time steps used as input for FNO and U-Net.t_train: int, number of the last time step used for training (for extrapolation testing, set this to be < Nt).model_update: int, number of epochs to save model.filename: str, has to match the dataset filename.single_file: bool, set False for 2D diffusion-reaction, 1D diffusion-sorption, and the radial dam break scenarios, and set True otherwise.reduced_resolution: int, factor to downsample spatial resolution.reduced_resolution_t: int, factor to downsample temporal resolution.reduced_batch: int, factor to downsample sample size used for training.epochs: int, total epochs used for training.learning_rate: float, learning rate of the optimizer.scheduler_step: int, number of epochs to update the learning rate scheduler.scheduler_gamma: float, decay rate of the learning rate.U-Net specific args:in_channels: int, number of input channelsout_channels: int, number of output channelsar_mode: bool, set True for fully autoregressive or pushforward training.pushforward: bool, set True for pushforward training, False otherwise (ar_mode also has to be set True).unroll_step: int, number of time steps to backpropagate in the pushforward training.FNO specific args:num_channels: int, number of channels (variables).modes: int, number of Fourier modes to multiply.width: int, number of channels for the Fourier layer.INVERSE specific args:base_path: string, location of the data directorytraining_type: string, type of training,  autoregressive, singlemcmc_num_samples: int, number of generated samplesmcmc_warmup_steps: 10mcmc_num_chains: 1num_samples_max: 1000in_channels_hid: 64inverse_model_type: string, type of inverse inference model, ProbRasterLatent, InitialConditionInterpinverse_epochs: int, number of epochs for the  gradient based methodinverse_learning_rate: float, learning rate for the gradient based methodinverse_verbose_flag: bool, some printingPlotting specific args:plot: bool, set True to activate plotting.channel_plot: int, determines which channel/variable to plot.x_min: float, left spatial domain.x_max: float, right spatial domain.y_min: float, lower spatial domain.y_max: float, upper spatial domain.t_min: float, start of temporal domain.t_max: float, end of temporal domain.Datasets and pretrained modelsWe provide the benchmark datasets we used in the paper through ourDaRUS data repository.
The data generation configuration can be found in the paper.
Additionally, the pretrained models are also available to be downloaded fromPDEBench Pretrained ModelsDaRus repository. To use the pretrained models, users can specify the argumentcontinue_training: Truein theconfig file.Directory TourBelow is an illustration of the directory structure of PDEBench.📂 pdebench
|_📁 models
  |_📁 pinn    # Model: Physics-Informed Neural Network 
    |_📄 train.py  
    |_📄 utils.py
    |_📄 pde_definitions.py
  |_📁 fno     # Model: Fourier Neural Operator
    |_📄 train.py
    |_📄 utils.py
    |_📄 fno.py
  |_📁 unet    # Model: U-Net
    |_📄 train.py
    |_📄 utils.py
    |_📄 unet.py
  |_📁 inverse # Model: Gradient-Based Inverse Method
    |_📄 train.py
    |_📄 utils.py
    |_📄 inverse.py
  |_📁 config  # Config: All config files reside here
  |_📄 train_models_inverse.py
  |_📄 run_forward_1D.sh
  |_📄 analyse_result_inverse.py
  |_📄 train_models_forward.py
  |_📄 run_inverse.sh
  |_📄 metrics.py
  |_📄 analyse_result_forward.py
|_📁 data_download  # Data: Scripts to download data from DaRUS
  |_📁 config
  |_📄 download_direct.py
  |_📄 download_easydataverse.py
  |_📄 visualize_pdes.py
  |_📄 README.md
  |_📄 download_metadata.csv
|_📁 data_gen   # Data: Scripts to generate data
  |_📁 configs
  |_📁 data_gen_NLE
  |_📁 src
  |_📁 notebooks
  |_📄 gen_diff_sorp.py
  |_📄 plot.py
  |_📄 example.env
  |_📄 gen_ns_incomp.py
  |_📄 gen_diff_react.py
  |_📄 uploader.py
  |_📄 gen_radial_dam_break.py
|_📄 __init__.pyCitationsPlease cite the following papers if you use PDEBench datasets and/or source code in your research.PDEBench: An Extensive Benchmark for Scientific Machine Learning - NeurIPS'2022@inproceedings{PDEBench2022,
author = {Takamoto, Makoto and Praditia, Timothy and Leiteritz, Raphael and MacKinlay, Dan and Alesiani, Francesco and Pflüger, Dirk and Niepert, Mathias},
title = {{PDEBench: An Extensive Benchmark for Scientific Machine Learning}},
year = {2022},
booktitle = {36th Conference on Neural Information Processing Systems (NeurIPS 2022) Track on Datasets and Benchmarks},
url = {https://arxiv.org/abs/2210.07182}
}PDEBench Datasets - NeurIPS'2022@data{darus-2986_2022,
author = {Takamoto, Makoto and Praditia, Timothy and Leiteritz, Raphael and MacKinlay, Dan and Alesiani, Francesco and Pflüger, Dirk and Niepert, Mathias},
publisher = {DaRUS},
title = {{PDEBench Datasets}},
year = {2022},
doi = {10.18419/darus-2986},
url = {https://doi.org/10.18419/darus-2986}
}Learning Neural PDE Solvers with Parameter-Guided Channel Attention - ICML'2023@article{cape-takamoto:2023,
     author   = {Makoto Takamoto and
                 Francesco Alesiani and
                 Mathias Niepert},
 title        = {Learning Neural {PDE} Solvers with Parameter-Guided Channel Attention},
 journal      = {CoRR},
 volume       = {abs/2304.14118},
 year         = {2023},
 url          = {https://doi.org/10.48550/arXiv.2304.14118},
 doi          = {10.48550/arXiv.2304.14118},
 eprinttype    = {arXiv},
 eprint       = {2304.14118},
 }Code contributorsMakato Takamoto(NEC laboratories Europe)Timothy Praditia(Stuttgart Center for Simulation Science | University of Stuttgart)Raphael Leiteritz(Stuttgart Center for Simulation Science | University of Stuttgart)Francesco Alesiani(NEC laboratories Europe)Dan MacKinlay(CSIRO’s Data61)Mario Kalimuthu(Stuttgart Center for Simulation Science | University of Stuttgart)John Kim(ANU TechLauncher/CSIRO’s Data61)Gefei Shan(ANU TechLauncher/CSIRO’s Data61)Yizhou Yang(ANU TechLauncher/CSIRO’s Data61)Ran Zhang(ANU TechLauncher/CSIRO’s Data61)Simon Brown(ANU TechLauncher/CSIRO’s Data61)LicenseMIT licensed, except where otherwise stated.
SeeLICENSE.txtfile.