diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 0000000..b20527d --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,24 @@ +// For format details, see https://aka.ms/devcontainer.json. For config options, see the +// README at: https://github.com/devcontainers/templates/tree/main/src/python +{ + "name": "Python 3", + // Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile + "image": "mcr.microsoft.com/devcontainers/python:1-3.10-bookworm", + + // Features to add to the dev container. More info: https://containers.dev/features. + // "features": {}, + + // Use 'forwardPorts' to make a list of ports inside the container available locally. + // "forwardPorts": [], + + "runArgs": ["--gpus", "all"], + + // Use 'postCreateCommand' to run commands after the container is created. + "postCreateCommand": "sudo apt update -y && sudo apt upgrade -y && sudo apt install gfortran -y && pip3 install --user -e .[dev,test,datagen310]" + + // Configure tool-specific properties. + // "customizations": {}, + + // Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root. + // "remoteUser": "root" +} diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..63fa9ed --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,12 @@ +# To get started with Dependabot version updates, you'll need to specify which +# package ecosystems to update and where the package manifests are located. +# Please see the documentation for more information: +# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates +# https://containers.dev/guide/dependabot + +version: 2 +updates: + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: weekly diff --git a/.github/workflows/precommit.yml b/.github/workflows/precommit.yml new file mode 100644 index 0000000..d0707c1 --- /dev/null +++ b/.github/workflows/precommit.yml @@ -0,0 +1,14 @@ +name: pre-commit + +on: + pull_request: + push: + branches: [main] + +jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + - uses: pre-commit/action@v3.0.1 diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml new file mode 100644 index 0000000..9c714d1 --- /dev/null +++ b/.github/workflows/python-package.yml @@ -0,0 +1,33 @@ +# This workflow will install Python dependencies, run tests and lint with a variety of Python versions +# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python + +name: Python package + +on: + push: + branches: ["main"] + pull_request: + branches: ["main"] + +jobs: + build: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9", "3.10"] + + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install --user -e .[test] + - name: Lint and test with nox + run: | + # stop the build if there are Python syntax errors or undefined names + nox -s pylint tests diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml new file mode 100644 index 0000000..c60b460 --- /dev/null +++ b/.github/workflows/python-publish.yml @@ -0,0 +1,38 @@ +# This workflow will upload a Python Package using Twine when a release is created +# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries + +# This workflow uses actions that are not certified by GitHub. +# They are provided by a third-party and are governed by +# separate terms of service, privacy policy, and support +# documentation. + +name: Upload Python Package + +on: + release: + types: [published] + +permissions: + contents: read + +jobs: + deploy: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.x" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install build + - name: Build package + run: python -m build + - name: Publish package + uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc + with: + user: __token__ + password: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.gitignore b/.gitignore index 109c47f..4780cfa 100644 --- a/.gitignore +++ b/.gitignore @@ -158,7 +158,7 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +.idea/ #linters and ... .vscode/ @@ -175,3 +175,7 @@ debug/ # macos .DS_Store + +# neural net, sciml artefacts +.pt +.hdf5 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..5155e1c --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,88 @@ +ci: + autoupdate_commit_msg: "chore: update pre-commit hooks" + autofix_commit_msg: "style: pre-commit fixes" + +repos: + - repo: https://github.com/adamchainz/blacken-docs + rev: "1.16.0" + hooks: + - id: blacken-docs + additional_dependencies: [black==23.*] + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: "v4.5.0" + hooks: + - id: check-added-large-files + - id: check-case-conflict + - id: check-merge-conflict + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: end-of-file-fixer + - id: mixed-line-ending + - id: name-tests-test + args: ["--pytest-test-first"] + - id: requirements-txt-fixer + - id: trailing-whitespace + + - repo: https://github.com/pre-commit/pygrep-hooks + rev: "v1.10.0" + hooks: + - id: rst-backticks + - id: rst-directive-colons + - id: rst-inline-touching-normal + + - repo: https://github.com/pre-commit/mirrors-prettier + rev: "v3.1.0" + hooks: + - id: prettier + types_or: [yaml, markdown, html, css, scss, javascript, json] + args: [--prose-wrap=always] + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: "v0.8.4" + hooks: + - id: ruff + args: ["--fix", "--show-fixes"] + - id: ruff-format + + # - repo: https://github.com/pre-commit/mirrors-mypy + # rev: "v1.8.0" + # hooks: + # - id: mypy + # files: pdebench|tests + # args: [] + # additional_dependencies: + # - pytest + + - repo: https://github.com/codespell-project/codespell + rev: "v2.2.6" + hooks: + - id: codespell + exclude_types: [jupyter] + + - repo: https://github.com/shellcheck-py/shellcheck-py + rev: "v0.10.0.1" + hooks: + - id: shellcheck + + - repo: local + hooks: + - id: disallow-caps + name: Disallow improper capitalization + language: pygrep + entry: PyBind|Numpy|Cmake|CCache|Github|PyTest + exclude: .pre-commit-config.yaml + + - repo: https://github.com/abravalheri/validate-pyproject + rev: "v0.16" + hooks: + - id: validate-pyproject + additional_dependencies: ["validate-pyproject-schema-store[all]"] + + - repo: https://github.com/python-jsonschema/check-jsonschema + rev: "0.27.3" + hooks: + - id: check-dependabot + - id: check-github-workflows + - id: check-readthedocs diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 index 0000000..49e3e47 --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,9 @@ +Except where otherwise stated this code is released under the MIT license. + + Copyright 2022 NEC Labs Europe GmbH, Stuttgart University, CSIRO and PDEbench contributors + + Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md index 1ed849b..ea37558 100644 --- a/README.md +++ b/README.md @@ -1,57 +1,210 @@ # PDEBench -This repository contains the code for the paper: -PDEBench: An Extensive Benchmark for Scientific Machine Learning +The code repository for the NeurIPS 2022 paper +[PDEBench: An Extensive Benchmark for Scientific Machine Learning](https://arxiv.org/abs/2210.07182) -In this work, we provide a diverse and comprehensive benchmark for scientific machine learning, including a variety of challenging and representative range of physical problems. -This repository consists of the codes used to generate the datasets, upload and download the datasets from the data repository, train and evaluate different machine learning models as baseline. -PDEBench features a much wider range of PDEs than existing approaches including realistic and difficult problems (both forward and inverse), larger ready-to-use datasets comprising various initial or boundary conditions and model parameters, and extensible source codes. +:tada: +[**SimTech Best Paper Award 2023**](https://www.simtech.uni-stuttgart.de/press/SimTech-Best-Paper-Award-2023-Benchmark-for-ML-for-scientific-simulations) +: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. + +![Visualizations of some PDE problems covered by the benchmark.](https://github.com/pdebench/PDEBench/blob/main/pdebench_examples.PNG) + +Created and maintained by Makoto Takamoto +``, Timothy Praditia +``, Raphael Leiteritz, Dan MacKinlay, +Francesco Alesiani, Dirk Pflรผger, and Mathias Niepert. + +--- + +## Datasets and Pretrained Models + +We also provide datasets and pretrained machine learning models. + +PDEBench Datasets: +https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2986 + +PDEBench Pre-Trained Models: +https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2987 + +DOIs + +[![DOI:10.18419/darus-2986](https://img.shields.io/badge/DOI-doi%3A10.18419%2Fdarus--2986-red)](https://doi.org/10.18419/darus-2986) +[![DOI:10.18419/darus-2987](https://img.shields.io/badge/DOI-doi%3A10.18419%2Fdarus--2987-red)](https://doi.org/10.18419/darus-2987) ## Installation - pip install pdebench +### Using pip -## Requirements +Locally: ```bash -python3 -m venv ./venv --prompt pde_benchmark --system-site-packages -. ./venv/bin/activate pip install --upgrade pip wheel -pip install -r requirements.txt +pip install . ``` -The required packages are listed in [requirements.txt](./requirements.txt) +From PyPI: -For GPU support there are additional platform-specific instructions: +```bash +pip install pdebench +``` -For pytorch, +To include dependencies for data generation: ```bash -pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113 +pip install "pdebench[datagen310]" +pip install ".[datagen310]" # locally ``` -For jax, which is approximately 6 times faster for simulations than pytorch in our tests, +or ```bash -pip install "jax[cuda11_cudnn82]>=0.3.0" -f https://storage.googleapis.com/jax-releases/jax_releases.html +pip install "pdebench[datagen39]" +pip install ".[datagen39]" # locally +``` + +### GPU Support + +For GPU support there are additional platform-specific instructions: + +For PyTorch, the latest version we support is v1.13.1 +[see previous-versions/#linux - CUDA 11.7](https://pytorch.org/get-started/previous-versions/#linux-and-windows-2). + +For JAX, which is approximately 6 times faster for simulations than PyTorch in +our tests, +[see jax#pip-installation-gpu-cuda-installed-via-pip](https://github.com/google/jax#pip-installation-gpu-cuda-installed-via-pip-easier) + +## Installation using conda: + +If you like you can also install dependencies using anaconda, we suggest to use +[mambaforge](https://github.com/conda-forge/miniforge#mambaforge) as a +distribution. Otherwise you may have to **enable the conda-forge** channel for +the following commands. + +Starting from a fresh environment: + +``` +conda create -n myenv python=3.9 +conda activate myenv ``` +Install dependencies for model training: + +``` +conda install deepxde hydra-core h5py -c conda-forge +``` + +According to your hardware availability, either install PyTorch with CUDA +support: + +- [see previous-versions/#linux - CUDA 11.7](https://pytorch.org/get-started/previous-versions/#linux-and-windows-2). + +``` +conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia +``` + +- [or CPU only binaries](https://pytorch.org/get-started/previous-versions/#linux-and-windows-2). + +``` +conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 cpuonly -c pytorch +``` + +Optional dependencies for data generation: + +``` +conda install clawpack jax jaxlib python-dotenv +``` + +## Configuring DeepXDE + +In our tests we used PyTorch as backend for DeepXDE. Please +[follow the documentation](https://deepxde.readthedocs.io/en/latest/user/installation.html#working-with-different-backends) +to enable this. + ## Data Generation + The data generation codes are contained in [data_gen](./pdebench/data_gen): + - `gen_diff_react.py` to generate the 2D diffusion-reaction data. - `gen_diff_sorp.py` to generate the 1D diffusion-sorption data. - `gen_radial_dam_break.py` to generate the 2D shallow-water data. -- `gen_radial_dam_break.py` to generate the 2D shallow-water data. -- `gen_ns_incomp.py` to generate the 2D incompressible inhomogenous Navier-Stokes data. +- `gen_ns_incomp.py` to generate the 2D incompressible inhomogeneous + Navier-Stokes data. - `plot.py` to plot the generated data. - `uploader.py` to upload the generated data to the data repository. -- `.env` is 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 the `example` from the filename) -- `configs` directory 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. -- `src` directory contains the simulation scripts for different problems: `sim_diff_react-py` for 2D diffusion-reaction, `sim_diff_sorp.py` for 1D diffusion-sorption, and `swe` for the shallow-water equation. +- `.env` is 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 the `example` from the filename) +- `configs` directory 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. +- `src` directory contains the simulation scripts for different problems: + `sim_diff_react-py` for 2D diffusion-reaction, `sim_diff_sorp.py` for 1D + diffusion-sorption, and `swe` for the shallow-water equation. + +### Data Generation for 1D Advection/Burgers/Reaction-Diffusion/2D DarcyFlow/Compressible Navier-Stokes Equations + +The data generation codes are contained in +[data_gen_NLE](./pdebench/data_gen/data_gen_NLE/): + +- `utils.py` util file for data generation, mainly boundary conditions and + initial conditions. +- `AdvectionEq` directory with the source codes to generate 1D Advection + equation training samples +- `BurgersEq` directory with the source codes to generate 1D Burgers equation + training samples +- `CompressibleFluid` directory with the source codes to generate compressible + Navier-Stokes equations training samples + + - `ReactionDiffusionEq` directory with the source codes to generate 1D + Reaction-Diffusion equation training samples (**Note: + [DarcyFlow data can be generated by run_DarcyFlow2D.sh](pdebench/data_gen/data_gen_NLE/README.md) + in this folder.**) + +- `save` directory saving the generated training samples + +A typical example to generate training samples (1D Advection Equation): (in +`data_gen/data_gen_NLE/AdvectionEq/`) + +```bash +python3 advection_multi_solution_Hydra.py +multi=beta1e0.yaml +``` + +which is assumed to be performed in each directory. + +Examples for generating other PDEs are provided in `run_trainset.sh` in each +PDE's directories. The config files for Hydra are stored in `config` directory +in each PDE's directory. + +#### Data Transformaion and Merge into HDF5 format + +1D 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 +using `data_gen_NLE/Data_Merge.py` whose 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: + +```bash +python3 Data_Merge.py +``` ## Configuration -You can sent the default values for data locations for this project by putting config vars like this in the `.env` file: +You can set the default values for data locations for this project by putting +config vars like this in the `.env` file: ``` WORKING_DIR=~/Data/Working @@ -60,42 +213,82 @@ ARCHIVE_DATA_DIR=~/Data/Archive There is an example in `example.env`. - ## Data Download -The data download codes are contained in [data_download](./pdebench/data_download): -- `download.py` to download the data. -- `.env` is the environment data to store Dataverse URL and API token to download the generated data. Note that the filename should be strictly `.env` (i.e. remove the `example` from the filename) -- `config` directory contains the yaml files storing the configuration for the data downloader. The argument 'filename' should match the filename in the data repository. +The download scripts are provided in [data_download](./pdebench/data_download). +There are two options to download data. + +1. Using `download_direct.py` (**recommended**) + - Retrieves data shards directly using URLs. Sample command for each PDE is + given in the README file in the [data_download](./pdebench/data_download) + directory. +2. Using `download_easydataverse.py` (might be slow and you could encounter + errors/issues; hence, not recommended!) + - Use the config files from the `config` directory that contains the yaml + files storing the configuration. Any files in the dataset matching + `args.filename` will be downloaded into `args.data_folder`. ## Baseline Models -In this work, we provide three different ML models to be trained and evaluated against the benchmark datasets, namely [FNO](https://arxiv.org/pdf/2010.08895.pdf), [U-Net](https://www.sciencedirect.com/science/article/abs/pii/S0010482519301520?via%3Dihub), and [PINN](https://www.sciencedirect.com/science/article/pii/S0021999118307125). -The codes for the baseline model implementations are contained in [models](./pdebench/models): -- `train_models_forward.py` is the main script to train and evaluate the model. It will call on model-specific script based on the input argument. -- `train_models_inverse.py` is 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.py` is 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.py` is 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.py` is 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. -- `fno` contains the scripts of FNO implementation. These are partly adapted from the [FNO repository](https://github.com/zongyi-li/fourier_neural_operator). -- `unet` contains the scripts of U-Net implementation. These are partly adapted from the [U-Net repository](https://github.com/mateuszbuda/brain-segmentation-pytorch). -- `pinn` contains the scripts of PINN implementation. These utilize the [DeepXDE library](https://github.com/lululxvi/deepxde). -- `inverse` contains the model for inverse model based on gradient. -- `config` contains the yaml files for the model training input. The default templates for different equations are provided in the [args](./pdebench/models/config/args) directory. User just needs to copy and paste them to the args keyword in the [config.yaml](./pdebench/models/config/config.yaml) file. -An example to run the forward model training can be found in [run_forward_1D.sh](./pdebench/models/run_forward_1D.sh), and an example to run the inverse model training can be found in [run_inverse.sh](./pdebench/models/run_inverse.sh). +In this work, we provide three different ML models to be trained and evaluated +against the benchmark datasets, namely +[FNO](https://arxiv.org/pdf/2010.08895.pdf), +[U-Net](https://www.sciencedirect.com/science/article/abs/pii/S0010482519301520?via%3Dihub), +and [PINN](https://www.sciencedirect.com/science/article/pii/S0021999118307125). +The codes for the baseline model implementations are contained in +[models](./pdebench/models): + +- `train_models_forward.py` is the main script to train and evaluate the model. + It will call on model-specific script based on the input argument. +- `train_models_inverse.py` is 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.py` is 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.py` is 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.py` is 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. +- `fno` contains the scripts of FNO implementation. These are partly adapted + from the + [FNO repository](https://github.com/zongyi-li/fourier_neural_operator). +- `unet` contains the scripts of U-Net implementation. These are partly adapted + from the + [U-Net repository](https://github.com/mateuszbuda/brain-segmentation-pytorch). +- `pinn` contains the scripts of PINN implementation. These utilize the + [DeepXDE library](https://github.com/lululxvi/deepxde). +- `inverse` contains the model for inverse model based on gradient. +- `config` contains the yaml files for the model training input. The default + templates for different equations are provided in the + [args](./pdebench/models/config/args) directory. User just needs to copy and + paste them to the args keyword in the + [config.yaml](./pdebench/models/config/config.yaml) file. +An example to run the forward model training can be found in +[run_forward_1D.sh](./pdebench/models/run_forward_1D.sh), and an example to run +the inverse model training can be found in +[run_inverse.sh](./pdebench/models/run_inverse.sh). ### Short explanations on the config args -- model_name: string, containing the baseline model name, either 'FNO', 'Unet', or 'PINN'. + +- model_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. +- continue_training: bool, set True to continue 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). +- 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. +- 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. @@ -105,31 +298,38 @@ An example to run the forward model training can be found in [run_forward_1D.sh] - scheduler_gamma: float, decay rate of the learning rate. #### U-Net specific args: + - in_channels: int, number of input channels - out_channels: int, number of output channels - ar_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. +- 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 directory - - training_type: string, type of training, autoregressive, single - - mcmc_num_samples: int, number of generated samples - - mcmc_warmup_steps: 10 - - mcmc_num_chains: 1 - - num_samples_max: 1000 - - in_channels_hid: 64 - - inverse_model_type: striung, type of inverse inference model, ProbRasterLatent, InitialConditionInterp - - inverse_epochs: int, number of epochs for the gradint based method - - inverse_learning_rate: float, learning rate for the gradint based method - - inverse_verbose_flag: bool, some printing + +- base_path: string, location of the data directory +- training_type: string, type of training, autoregressive, single +- mcmc_num_samples: int, number of generated samples +- mcmc_warmup_steps: 10 +- mcmc_num_chains: 1 +- num_samples_max: 1000 +- in_channels_hid: 64 +- inverse_model_type: string, type of inverse inference model, ProbRasterLatent, + InitialConditionInterp +- inverse_epochs: int, number of epochs for the gradient based method +- inverse_learning_rate: float, learning rate for the gradient based method +- inverse_verbose_flag: bool, some printing #### Plotting specific args: + - plot: bool, set True to activate plotting. - channel_plot: int, determines which channel/variable to plot. - x_min: float, left spatial domain. @@ -140,64 +340,209 @@ An example to run the forward model training can be found in [run_forward_1D.sh] - t_max: float, end of temporal domain. ## Datasets and pretrained models -We provide the benchmark datasets we used in the paper through our [data repository](https://darus.uni-stuttgart.de/privateurl.xhtml?token=1be27526-348a-40ed-9fd0-c62f588efc01). -The data generation configuration can be found in the paper. -Additionally, the pretrained models are also available to be downloaded [here](https://darus.uni-stuttgart.de/privateurl.xhtml?token=cd862f8c-8e1b-49d2-b4da-b35f8df5ac85). +We provide the benchmark datasets we used in the paper through our +[DaRUS data repository](https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2986). +The data generation configuration can be found in the paper. Additionally, the +pretrained models are also available to be downloaded from +[PDEBench Pretrained Models](https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2987) +DaRus repository. To use the pretrained models, users can specify the argument +`continue_training: True` in the +[config file](./pdebench/models/config/config.yaml). + +--- + +## Directory Tour + +Below 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__.py +``` + +--- + +## Publications & Citations -## Citations +Please 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 + +
``` -@online{PDEBenchDataset, - author = {Makoto Takamoto and Timothy Pradita and Raphael Leiteritz and Dan MacKinlay and Francesco Alesiani and Dirk Pflรผger and Mathias Niepert}, - title = {{PDEBench}: A Diverse and Comprehensive Benchmark for Scientific Machine Learning}, - year = 2022, - doi = {doi:10.18419/darus-2986}, - urldoi = {http://dx.doi.org/10.18419/darus-2986}, - url = {https://darus.uni-stuttgart.de/privateurl.xhtml?token=1be27526-348a-40ed-9fd0-c62f588efc01}, - urldate = {2022-06-06} +@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} } +``` + +
-@misc{li2020fourier, - title={Fourier Neural Operator for Parametric Partial Differential Equations}, - author={Zongyi Li and Nikola Kovachki and Kamyar Azizzadenesheli and Burigede Liu and Kaushik Bhattacharya and Andrew Stuart and Anima Anandkumar}, - year={2020}, - eprint={2010.08895}, - archivePrefix={arXiv}, - primaryClass={cs.LG} +
+ + 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} } +``` -@article{buda2019association, - title={Association of genomic subtypes of lower-grade gliomas with shape features automatically extracted by a deep learning algorithm}, - author={Buda, Mateusz and Saha, Ashirbani and Mazurowski, Maciej A}, - journal={Computers in Biology and Medicine}, - volume={109}, - year={2019}, - publisher={Elsevier}, - doi={10.1016/j.compbiomed.2019.05.002} +
+ +
+ + 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}, + } +``` + +
+ +
+ + Vectorized Conditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations - ICLR-W'2024 & ICML'2024 + +
+ +``` +@inproceedings{vcnef-vectorized-conditional-neural-fields-hagnberger:2024, +author = {Hagnberger, Jan and Kalimuthu, Marimuthu and Musekamp, Daniel and Niepert, Mathias}, +title = {{Vectorized Conditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations}}, +year = {2024}, +booktitle = {Proceedings of the 41st International Conference on Machine Learning (ICML 2024)} } +``` -@article{raissi2019pinn, -title = {Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations}, -journal = {Journal of Computational Physics}, -volume = {378}, -pages = {686-707}, -year = {2019}, -issn = {0021-9991}, -doi = {https://doi.org/10.1016/j.jcp.2018.10.045}, -author = {M. Raissi and P. Perdikaris and G.E. Karniadakis} +
+ +
+ + Active Learning for Neural PDE Solvers - NeurIPS-W'2024 + +
+ +``` +@article{active-learn-neuralpde-benchmark-musekamp:2024, + author = {Daniel Musekamp and + Marimuthu Kalimuthu and + David Holzm{\"{u}}ller and + Makoto Takamoto and + Mathias Niepert}, + title = {Active Learning for Neural {PDE} Solvers}, + journal = {CoRR}, + volume = {abs/2408.01536}, + year = {2024}, + url = {https://doi.org/10.48550/arXiv.2408.01536}, + doi = {10.48550/ARXIV.2408.01536}, + eprinttype = {arXiv}, + eprint = {2408.01536}, } ``` +
+ +--- + ## Code contributors -* [Makato Takamoto](https://github.com/mtakamoto-D) -* [Timothy Praditia](https://github.com/timothypraditia) -* [Raphael Leiteritz](https://github.com/leiterrl) -* [Francesco Alesiani](https://github.com/falesiani) -* [Dan MacKinlay](https://danmackinlay.name/) -* [John Kim](https://github.com/johnmjkim) -* [Gefei Shan](https://github.com/davecatmeow) -* [Yizhou Yang](https://github.com/verdantwynnd) -* [Ran Zhang](https://github.com/maphyca) -* [Simon Brown](https://github.com/SimonSyBrown) \ No newline at end of file +- [Makato Takamoto](https://github.com/mtakamoto-D) + ([NEC laboratories Europe](https://www.neclab.eu/)) +- [Timothy Praditia](https://github.com/timothypraditia) + ([Stuttgart Center for Simulation Science | University of Stuttgart](https://www.simtech.uni-stuttgart.de/)) +- [Raphael Leiteritz](https://github.com/leiterrl) + ([Stuttgart Center for Simulation Science | University of Stuttgart](https://www.simtech.uni-stuttgart.de/)) +- [Francesco Alesiani](https://github.com/falesiani) + ([NEC laboratories Europe](https://www.neclab.eu/)) +- [Dan MacKinlay](https://danmackinlay.name/) + ([CSIROโ€™s Data61](https://data61.csiro.au/)) +- [Marimuthu Kalimuthu](https://github.com/kmario23) + ([Stuttgart Center for Simulation Science | University of Stuttgart](https://www.simtech.uni-stuttgart.de/)) +- [John Kim](https://github.com/johnmjkim) + ([ANU TechLauncher](https://comp.anu.edu.au/TechLauncher/)/[CSIROโ€™s Data61](https://data61.csiro.au/)) +- [Gefei Shan](https://github.com/davecatmeow) + ([ANU TechLauncher](https://comp.anu.edu.au/TechLauncher/)/[CSIROโ€™s Data61](https://data61.csiro.au/)) +- [Yizhou Yang](https://github.com/verdantwynnd) + ([ANU TechLauncher](https://comp.anu.edu.au/TechLauncher/)/[CSIROโ€™s Data61](https://data61.csiro.au/)) +- [Ran Zhang](https://github.com/maphyca) + ([ANU TechLauncher](https://comp.anu.edu.au/TechLauncher/)/[CSIROโ€™s Data61](https://data61.csiro.au/)) +- [Simon Brown](https://github.com/SimonSyBrown) + ([ANU TechLauncher](https://comp.anu.edu.au/TechLauncher/)/[CSIROโ€™s Data61](https://data61.csiro.au/)) + +## License + +MIT licensed, except where otherwise stated. See `LICENSE.txt` file. diff --git a/noxfile.py b/noxfile.py new file mode 100644 index 0000000..9536e44 --- /dev/null +++ b/noxfile.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import shutil +from pathlib import Path + +import nox + +DIR = Path(__file__).parent.resolve() + +nox.needs_version = ">=2024.3.2" +nox.options.sessions = ["precommit", "pylint", "tests", "build"] +nox.options.default_venv_backend = "uv|mamba|virtualenv" + + +@nox.session(python=["3.10"]) +def precommit(session: nox.Session) -> None: + """ + Run the linter. + """ + session.install("pre-commit") + session.run( + "pre-commit", "run", "--all-files", "--show-diff-on-failure", *session.posargs + ) + + +@nox.session(python=["3.10"]) +def pylint(session: nox.Session) -> None: + """ + Run PyLint. + """ + # This needs to be installed into the package environment, and is slower + # than a pre-commit check + session.install(".", "pylint>=3.2") + session.run("pylint", "pdebench", *session.posargs) + + +@nox.session(python=["3.10"]) +def tests(session: nox.Session) -> None: + """ + Run the unit and regular tests. + """ + session.install(".[test]") + session.run("pytest", *session.posargs) + + +@nox.session(python=["3.10"]) +def build(session: nox.Session) -> None: + """ + Build an SDist and wheel. + """ + + build_path = DIR.joinpath("build") + if build_path.exists(): + shutil.rmtree(build_path) + + session.install("build") + session.run("python", "-m", "build") diff --git a/pdebench/__init__.py b/pdebench/__init__.py new file mode 100644 index 0000000..78651a5 --- /dev/null +++ b/pdebench/__init__.py @@ -0,0 +1,32 @@ +""" +PDEBench: An Extensive Benchmark for Scientific Machine Learning + +In this work, we provide a diverse and comprehensive benchmark for scientific machine learning, including a variety of challenging and representative range of physical problems. +This repository consists of the codes used to generate the datasets, upload and download the datasets from the data repository, train and evaluate different machine learning models as baseline. +PDEBench features a much wider range of PDEs than existing approaches including realistic and difficult problems (both forward and inverse), larger ready-to-use datasets comprising various initial or boundary conditions and model parameters, and extensible source codes. + +Citations + +@online{PDEBenchDataset, + author = {Makoto Takamoto and Timothy Pradita and Raphael Leiteritz and Dan MacKinlay and Francesco Alesiani and Dirk Pflรผger and Mathias Niepert}, + title = {{PDEBench}: A Diverse and Comprehensive Benchmark for Scientific Machine Learning}, + year = 2022, + doi = {doi:10.18419/darus-2986}, + urldoi = {http://dx.doi.org/10.18419/darus-2986}, + url = {https://darus.uni-stuttgart.de/privateurl.xhtml?token=1be27526-348a-40ed-9fd0-c62f588efc01}, + urldate = {2022-06-06} +} + + +""" + +from __future__ import annotations + +import logging + +_logger = logging.getLogger(__name__) +_logger.propagate = False + +__version__ = "0.0.1" +__author__ = "Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflรผger, Mathias Niepert" +__credits__ = "NEC labs Europe, University of Stuttgart, CSIRO's Data61" diff --git a/pdebench/_version.pyi b/pdebench/_version.pyi new file mode 100644 index 0000000..91744f9 --- /dev/null +++ b/pdebench/_version.pyi @@ -0,0 +1,4 @@ +from __future__ import annotations + +version: str +version_tuple: tuple[int, int, int] | tuple[int, int, int, str, str] diff --git a/pdebench/data_download/README.md b/pdebench/data_download/README.md new file mode 100644 index 0000000..5be2a5d --- /dev/null +++ b/pdebench/data_download/README.md @@ -0,0 +1,121 @@ +# Downloading PDEBench Datasets :earth_asia: + +Here we enumerate the list of all available PDEs in PDEBench and the commands to +download them. + +| PDEs | Dataset Download | Dataset Size | +| ----------- | :------------------------------------------------------------------------------- | ------------ | +| advection | `python download_direct.py --root_folder $proj_home/data --pde_name advection` | 47 GB | +| burgers | `python download_direct.py --root_folder $proj_home/data --pde_name burgers` | 93 GB | +| 1d_cfd | `python download_direct.py --root_folder $proj_home/data --pde_name 1d_cfd` | 88 GB | +| diff_sorp | `python download_direct.py --root_folder $proj_home/data --pde_name diff_sorp` | 4 GB | +| 1d_reacdiff | `python download_direct.py --root_folder $proj_home/data --pde_name 1d_reacdiff` | 62 GB | +| 2d_reacdiff | `python download_direct.py --root_folder $proj_home/data --pde_name 2d_reacdiff` | 13 GB | +| 2d_cfd | `python download_direct.py --root_folder $proj_home/data --pde_name 2d_cfd` | 551 GB | +| 3d_cfd | `python download_direct.py --root_folder $proj_home/data --pde_name 3d_cfd` | 285 GB | +| darcy | `python download_direct.py --root_folder $proj_home/data --pde_name darcy` | 6.2 GB | +| ns_incom | `python download_direct.py --root_folder $proj_home/data --pde_name ns_incom` | 2.3 TB | +| swe | `python download_direct.py --root_folder $proj_home/data --pde_name swe` | 6.2 GB | + +--- + +# Visualizing PDEs :ocean: + +Below are some illustrations for how to visualize a certain PDE. It is assumed +that you first download the data shard you'd like to visualize for a desired +PDE. Then you can use the `visualize_pde.py` script to generate an animation +(i.e., `.gif`). + +###### 1D Diffusion Sorption Eqn + +``` +# get data: 1D_diff-sorp_NA_NA.h5 +https://darus.uni-stuttgart.de/api/access/datafile/133020 + +# visualize +python visualize_pdes.py --pde_name "diff_sorp" --data_path "./" +``` + +--- + +###### 1D Diffusion Reaction Eqn + +``` +# get data: ReacDiff_Nu1.0_Rho1.0.hdf5 +https://darus.uni-stuttgart.de/api/access/datafile/133181 + +# visualize +python visualize_pdes.py --pde_name "1d_reacdiff" +``` + +--- + +###### 1D Advection Eqn + +``` +# get data: 1D_Advection_Sols_beta0.4.hdf5 +https://darus.uni-stuttgart.de/api/access/datafile/133110 + +# visualize +python visualize_pdes.py --pde_name "advection" --param 0.4 +``` + +--- + +###### 1D Burgers Eqn + +``` +# get data: 1D_Burgers_Sols_Nu0.01.hdf5 +https://darus.uni-stuttgart.de/api/access/datafile/133136 + +# visualize +python visualize_pdes.py --pde_name "burgers" --param 0.01 +``` + +--- + +###### 1D CFD Eqn + +``` +# get data: 1D_CFD_Rand_Eta1.e-8_Zeta1.e-8_periodic_Train.hdf5 +https://darus.uni-stuttgart.de/api/access/datafile/135485 + +# visualize +python visualize_pdes.py --pde_name "1d_cfd" +``` + +--- + +###### 2D Diffusion Reaction Eqn + +``` +# get data: 2D_diff-react_NA_NA.h5 +https://darus.uni-stuttgart.de/api/access/datafile/133017 + +# visualize +python visualize_pdes.py --pde_name "2d_reacdiff" +``` + +--- + +###### 2D Darcy Flow Eqn + +``` +# get data: 2D_DarcyFlow_beta1.0_Train.hdf5 +https://darus.uni-stuttgart.de/api/access/datafile/133219 + +# visualize +python visualize_pdes.py --pde_name "darcy" +``` + +--- + +###### 2D Shallow Water Eqn + +``` +# get data: 2D_rdb_NA_NA.h5 +https://darus.uni-stuttgart.de/api/access/datafile/133021 + +# visualize +python visualize_pdes.py --pde_name "swe" --data_path "./" +``` diff --git a/pdebench/data_download/config/config.yaml b/pdebench/data_download/config/config.yaml index a2897f3..13d7b8b 100644 --- a/pdebench/data_download/config/config.yaml +++ b/pdebench/data_download/config/config.yaml @@ -7,6 +7,9 @@ hydra: output_subdir: null run: dir: . - + args: - filename: '1D_diff-sorp_NA_NA' \ No newline at end of file + filename: "Advection_beta" + dataverse_url: "https://darus.uni-stuttgart.de" + dataset_id: "doi:10.18419/darus-2986" + data_folder: "data" diff --git a/pdebench/data_download/download_direct.py b/pdebench/data_download/download_direct.py new file mode 100644 index 0000000..5a87249 --- /dev/null +++ b/pdebench/data_download/download_direct.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import pandas as pd +from torchvision.datasets.utils import download_url +from tqdm import tqdm + + +def parse_metadata(pde_names): + """ + This function parses the argument to filter the metadata of files that need to be downloaded. + + Args: + pde_names: List containing the name of the PDE to be downloaded + df : The provided dataframe loaded from the csv file + + Options for pde_names: + - Advection + - Burgers + - 1D_CFD + - Diff-Sorp + - 1D_ReacDiff + - 2D_CFD + - Darcy + - 2D_ReacDiff + - NS_Incom + - SWE + - 3D_CFD + + Returns: + pde_df : Filtered dataframe containing metadata of files to be downloaded + """ + + meta_df = pd.read_csv("pdebench_data_urls.csv") + + # Ensure the pde_name is defined + pde_list = [ + "advection", + "burgers", + "1d_cfd", + "diff_sorp", + "1d_reacdiff", + "2d_cfd", + "darcy", + "2d_reacdiff", + "ns_incom", + "swe", + "3d_cfd", + ] + + assert all(name.lower() in pde_list for name in pde_names), "PDE name not defined." + + # Filter the files to be downloaded + meta_df["PDE"] = meta_df["PDE"].str.lower() + + return meta_df[meta_df["PDE"].isin(pde_names)] + + +def download_data(root_folder, pde_name): + """ " + Download data splits specific to a given PDE. + + Args: + root_folder: The root folder where the data will be downloaded + pde_name : The name of the PDE for which the data to be downloaded + """ + + # print(f"Downloading data for {pde_name} ...") + + # Load and parse metadata csv file + pde_df = parse_metadata(pde_name) + + # Iterate filtered dataframe and download the files + for _, row in tqdm(pde_df.iterrows(), total=pde_df.shape[0]): + file_path = Path(root_folder) / row["Path"] + download_url(row["URL"], file_path, row["Filename"], md5=row["MD5"]) + + +if __name__ == "__main__": + arg_parser = argparse.ArgumentParser( + prog="Download Script", + description="Helper script to download the PDEBench datasets", + epilog="", + ) + + arg_parser.add_argument( + "--root_folder", + type=str, + required=True, + help="Root folder where the data will be downloaded", + ) + arg_parser.add_argument( + "--pde_name", + action="append", + help="Name of the PDE dataset to download. You can use this flag multiple times to download multiple datasets", + ) + + args = arg_parser.parse_args() + + download_data(args.root_folder, args.pde_name) diff --git a/pdebench/data_download/download.py b/pdebench/data_download/download_easydataverse.py similarity index 51% rename from pdebench/data_download/download.py rename to pdebench/data_download/download_easydataverse.py index 4622202..64c87d9 100644 --- a/pdebench/data_download/download.py +++ b/pdebench/data_download/download_easydataverse.py @@ -1,17 +1,13 @@ -import os -import dotenv +from __future__ import annotations -dotenv.load_dotenv() +import logging +import os import hydra +from easyDataverse import Dataset from hydra.utils import get_original_cwd from omegaconf import DictConfig -import logging - -import glob -from pyDataverse.api import NativeApi, DataAccessApi -from pyDaRUS import Dataset -from easyDataverse.core.downloader import download_files +from pyDataverse.api import NativeApi log = logging.getLogger(__name__) @@ -19,7 +15,7 @@ @hydra.main(config_path="config/", config_name="config") def main(config: DictConfig): """ - use config specifications to generate dataset + use config specifications to download files from a dataset """ # Imports should be nested inside @hydra.main to optimize tab completion @@ -28,37 +24,32 @@ def main(config: DictConfig): # Change to original working directory os.chdir(get_original_cwd()) - - path_to_data = os.path.join(os.path.abspath("../data"), config.args.filename) - - file_list = sorted(glob.glob(path_to_data + "/*.h5")) - doi = "doi:10.18419/darus-2986" + os.environ["DATAVERSE_URL"] = config.args.dataverse_url # Extract dataset from the given DOI dataset = Dataset() - setattr(dataset, "p_id", doi) + dataset.p_id = config.args.dataset_id # Extract file list contained in the dataset - api = NativeApi(os.getenv("DATAVERSE_URL"), os.getenv("DATAVERSE_API_TOKEN")) - data_api = DataAccessApi( - os.getenv("DATAVERSE_URL"), os.getenv("DATAVERSE_API_TOKEN") - ) - dv_dataset = api.get_dataset(doi) + api = NativeApi(config.args.dataverse_url) + dv_dataset = api.get_dataset(config.args.dataset_id) files_list = dv_dataset.json()["data"]["latestVersion"]["files"] # Compile list of files that matches the desired filename files = [] - for i, file in enumerate(files_list): + for _, file in enumerate(files_list): if config.args.filename in file["dataFile"]["filename"]: - files.append(file) + files.append(file["dataFile"]["filename"]) # Download the files - download_files(data_api, dataset, files, os.path.abspath("../data/")) - - return + dataset = Dataset.from_dataverse_doi( + doi=config.args.dataset_id, + dataverse_url=config.args.dataverse_url, + filenames=files, + filedir=config.args.data_folder, + ) -import os if __name__ == "__main__": main() diff --git a/pdebench/data_download/example.env b/pdebench/data_download/example.env deleted file mode 100644 index 399e3dc..0000000 --- a/pdebench/data_download/example.env +++ /dev/null @@ -1,3 +0,0 @@ -DATAVERSE_URL=https://darus.uni-stuttgart.de -DATAVERSE_API_TOKEN=abc_0123 -DATAVERSE_ID=doi:10.18419/darus-2986 \ No newline at end of file diff --git a/pdebench/data_download/pdebench_data_urls.csv b/pdebench/data_download/pdebench_data_urls.csv new file mode 100644 index 0000000..3332b67 --- /dev/null +++ b/pdebench/data_download/pdebench_data_urls.csv @@ -0,0 +1,376 @@ +PDE,Filename,URL,Path,MD5 +Advection,1D_Advection_Sols_beta0.1.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255672,1D/Advection/Train/,b4be2fc3383f737c76033073e6d2ccfb +Advection,1D_Advection_Sols_beta0.2.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255671,1D/Advection/Train/,d47b52c1be4b6bcaa8748a8df2f5fbb8 +Advection,1D_Advection_Sols_beta0.4.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255674,1D/Advection/Train/,d595bbfd2c659df995a93cd40d6ea568 +Advection,1D_Advection_Sols_beta0.7.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255666,1D/Advection/Train/,f89030c240a1e90b55649776854a84c2 +Advection,1D_Advection_Sols_beta1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255675,1D/Advection/Train/,1fe41923a4123db55bf4e89bea32e142 +Advection,1D_Advection_Sols_beta2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255677,1D/Advection/Train/,0e9beff98693089c38e213a1f2b9ac43 +Advection,1D_Advection_Sols_beta4.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255676,1D/Advection/Train/,a87b03ff425496360e2732921805f47d +Advection,1D_Advection_Sols_beta7.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/255664,1D/Advection/Train/,72a286ad2fcba574bd20bc045ba82d58 +Burgers,1D_Burgers_Sols_Nu0.001.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268190,1D/Burgers/Train/,44cb784d5a07aa2b1c864cabdcf625f9 +Burgers,1D_Burgers_Sols_Nu0.002.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268193,1D/Burgers/Train/,edf1cd13622d151dfde3e4e5af7a95b4 +Burgers,1D_Burgers_Sols_Nu0.004.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268191,1D/Burgers/Train/,435e1fecb8a64a0b4563bb8d09f81c33 +Burgers,1D_Burgers_Sols_Nu0.01.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/281363,1D/Burgers/Train/,e6d9a4f62baf9a29121a816b919e2770 +Burgers,1D_Burgers_Sols_Nu0.02.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268189,1D/Burgers/Train/,7c8c717a3a7818145877baa57106b090 +Burgers,1D_Burgers_Sols_Nu0.04.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/281362,1D/Burgers/Train/,16f4c7afaf8c16238be157e54c9297c7 +Burgers,1D_Burgers_Sols_Nu0.1.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268185,1D/Burgers/Train/,660ba1008d3843bf4e28d2895eb607ce +Burgers,1D_Burgers_Sols_Nu0.2.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268187,1D/Burgers/Train/,01d9333254dff2f3cad090b61f5cf695 +Burgers,1D_Burgers_Sols_Nu0.4.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268192,1D/Burgers/Train/,58327a6107e8f9defb5075677cc42533 +Burgers,1D_Burgers_Sols_Nu1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/281365,1D/Burgers/Train/,9021eba35332d127306f11ef84c1a60f +Burgers,1D_Burgers_Sols_Nu2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/281364,1D/Burgers/Train/,70fe0c24d9313e70f6059e5212bdb3d7 +Burgers,1D_Burgers_Sols_Nu4.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/268188,1D/Burgers/Train/,8af7f70166c00ffad377c80fa477d81f +1D_CFD,1D_CFD_Rand_Eta0.01_Zeta0.01_periodic_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164672,1D/CFD/Train/,2fdd00138f1d7abc794fb953021a9f43 +1D_CFD,1D_CFD_Rand_Eta0.1_Zeta0.1_periodic_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164668,1D/CFD/Train/,45655bd77d006ab539c52b7fbcf099b9 +1D_CFD,1D_CFD_Rand_Eta1.e-8_Zeta1.e-8_periodic_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/135485,1D/CFD/Train/,ad22fe08b8abc721179dbb34f2cc8b2a +1D_CFD,1D_CFD_Rand_Eta1.e-8_Zeta1.e-8_trans_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133155,1D/CFD/Train/,79857a11ad9a69d77b8cf249a1c72d06 +1D_CFD,1D_CFD_Shock_Eta1.e-8_Zeta1.e-8_trans_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133156,1D/CFD/Train/,08dffc3d84033c7574a70614703dd753 +1D_CFD,Sod1.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133145,1D/CFD/Test/ShockTube/,ef8271ce332b986e961a7aaedcae24e7 +1D_CFD,Sod2.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133146,1D/CFD/Test/ShockTube/,10e47c5c6ee0d1028f912a52bbf613d0 +1D_CFD,Sod3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133147,1D/CFD/Test/ShockTube/,495f2d2e0cee9b4ae28d7fef29bb7c68 +1D_CFD,Sod4.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133148,1D/CFD/Test/ShockTube/,544399b3d05ae514d66fc96982ee9db8 +1D_CFD,Sod5.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133149,1D/CFD/Test/ShockTube/,ca7fa6631201e8bd506914d3b0b9fe50 +1D_CFD,Sod6.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133150,1D/CFD/Train/ShockTube/,adb2d95bf0d48e03bc0d8f4a2cbcd1c6 +1D_CFD,Sod7.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133151,1D/CFD/Test/ShockTube/,b8d73e1e3ef862ef732a5fa43c64612e +Diff_Sorp,1D_diff-sorp_NA_NA.h5,https://darus.uni-stuttgart.de/api/access/datafile/133020,1D/diffusion-sorption/,9d466d1213065619d087319e16d9a938 +1D_ReacDiff,ReacDiff_Nu0.5_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133177,1D/ReactionDiffusion/Train/,69a429239778d529cd419ed5888ea835 +1D_ReacDiff,ReacDiff_Nu0.5_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133178,1D/ReactionDiffusion/Train/,ff7c724b18e7ebe02e19c179852f48ee +1D_ReacDiff,ReacDiff_Nu0.5_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133179,1D/ReactionDiffusion/Train/,ac907daa7e483d203a5c77567cdea561 +1D_ReacDiff,ReacDiff_Nu0.5_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133180,1D/ReactionDiffusion/Train/,fb149e7540d8977af158bb8fec1048a3 +1D_ReacDiff,ReacDiff_Nu1.0_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133181,1D/ReactionDiffusion/Train/,bd73c2f3448d03e95e98c3831fc8fa70 +1D_ReacDiff,ReacDiff_Nu1.0_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133182,1D/ReactionDiffusion/Train/,a94e65631881a27ddae3ef74caf53093 +1D_ReacDiff,ReacDiff_Nu1.0_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133183,1D/ReactionDiffusion/Train/,112c01a76447162bd67c8c1073f58ca2 +1D_ReacDiff,ReacDiff_Nu1.0_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133184,1D/ReactionDiffusion/Train/,fa224c9d143de37ac6914d391e70f425 +1D_ReacDiff,ReacDiff_Nu2.0_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133185,1D/ReactionDiffusion/Train/,925a7e2c9b9e40ad2dd44b7002ec882b +1D_ReacDiff,ReacDiff_Nu2.0_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133186,1D/ReactionDiffusion/Train/,a3e25a9fb8a99f010352ad3dd1afd596 +1D_ReacDiff,ReacDiff_Nu2.0_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133187,1D/ReactionDiffusion/Train/,426353b241acfbc64067acb1bdc80ade +1D_ReacDiff,ReacDiff_Nu2.0_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133188,1D/ReactionDiffusion/Train/,f2890bdac5103a3b78fac5f80e57b760 +1D_ReacDiff,ReacDiff_Nu5.0_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133189,1D/ReactionDiffusion/Train/,0ed75b55f61bec11c47d379e34959e54 +1D_ReacDiff,ReacDiff_Nu5.0_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133190,1D/ReactionDiffusion/Train/,264c80af0e2a6cc1f70e87275e0f6ac4 +1D_ReacDiff,ReacDiff_Nu5.0_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133191,1D/ReactionDiffusion/Train/,f927f5d85f2e9bd97dff31b4dab051ae +1D_ReacDiff,ReacDiff_Nu5.0_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133192,1D/ReactionDiffusion/Train/,c9c26e2b5f2d4bf5bcbbd4ba34fefd5e +1D_ReacDiff,ReacDiff_react_Nu0.5_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133157,1D/ReactionDiffusion/Test/,6dc95d36e7126c104c5316485670f2c0 +1D_ReacDiff,ReacDiff_react_Nu0.5_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133158,1D/ReactionDiffusion/Test/,5bf21fbc428ac4346c0e7a04a5899f9b +1D_ReacDiff,ReacDiff_react_Nu0.5_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133159,1D/ReactionDiffusion/Test/,5198a1572c201657dc19ad8592e24ac1 +1D_ReacDiff,ReacDiff_react_Nu0.5_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133160,1D/ReactionDiffusion/Test/,9968701e89a8b406f18ffbc2b74045d8 +1D_ReacDiff,ReacDiff_react_Nu1.0_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133161,1D/ReactionDiffusion/Test/,a258ff5be23d17d89cb10ac94f512165 +1D_ReacDiff,ReacDiff_react_Nu1.0_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133162,1D/ReactionDiffusion/Test/,09279a597860dcf48ef6f6539c46920f +1D_ReacDiff,ReacDiff_react_Nu1.0_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133163,1D/ReactionDiffusion/Test/,27113922509f6336f5c8be8d99109828 +1D_ReacDiff,ReacDiff_react_Nu1.0_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133164,1D/ReactionDiffusion/Test/,24ab563c27bf07d653b881d3c652bfa5 +1D_ReacDiff,ReacDiff_react_Nu10.0_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133165,1D/ReactionDiffusion/Test/,5a0f4d1ff304b11b9b22c3b00b1f63d4 +1D_ReacDiff,ReacDiff_react_Nu10.0_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133166,1D/ReactionDiffusion/Test/,4cf67490b01e54f80d45c8e7936b804d +1D_ReacDiff,ReacDiff_react_Nu10.0_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133167,1D/ReactionDiffusion/Test/,3c54c1cd3710b1f5b02a3de3782201ed +1D_ReacDiff,ReacDiff_react_Nu10.0_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133168,1D/ReactionDiffusion/Test/,580f08d5755ee393e7092eafedb70e46 +1D_ReacDiff,ReacDiff_react_Nu2.0_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133169,1D/ReactionDiffusion/Test/,208882dcf594d52f8e5017904efd84c9 +1D_ReacDiff,ReacDiff_react_Nu2.0_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133170,1D/ReactionDiffusion/Test/,891f8b47e49348b5411882ecebbc1bfd +1D_ReacDiff,ReacDiff_react_Nu2.0_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133171,1D/ReactionDiffusion/Test/,56e8e1e55ac92acf21b253511c370642 +1D_ReacDiff,ReacDiff_react_Nu2.0_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133172,1D/ReactionDiffusion/Test/,4498954ef4cd5c63ac530ee46c351790 +1D_ReacDiff,ReacDiff_react_Nu5.0_Rho1.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133173,1D/ReactionDiffusion/Test/,6b2c45748711bea8779dea3b48a53a7e +1D_ReacDiff,ReacDiff_react_Nu5.0_Rho10.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133174,1D/ReactionDiffusion/Test/,b04c22629bf86d85929386aaf4a5f95d +1D_ReacDiff,ReacDiff_react_Nu5.0_Rho2.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133175,1D/ReactionDiffusion/Test/,1738e4374456af974e4d0058d3426a69 +1D_ReacDiff,ReacDiff_react_Nu5.0_Rho5.0.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133176,1D/ReactionDiffusion/Test/,4e3cf0ac66e85dbb77c79861076db8e8 +2D_CFD,2D_CFD_Rand_M0.1_Eta0.01_Zeta0.01_periodic_128_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164687,2D/CFD/2D_Train_Rand/,5b21dcccaef4d2145ca579a71153c580 +2D_CFD,2D_CFD_Rand_M0.1_Eta0.1_Zeta0.1_periodic_128_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164688,2D/CFD/2D_Train_Rand/,b2733888745a1d64e36df813e979910b +2D_CFD,2D_CFD_Rand_M0.1_Eta1e-08_Zeta1e-08_periodic_512_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164689,2D/CFD/2D_Train_Rand/,ae2552864d0856f58dc74132a2f95d20 +2D_CFD,2D_CFD_Rand_M1.0_Eta0.01_Zeta0.01_periodic_128_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164690,2D/CFD/2D_Train_Rand/,21969082d0e9524bcc4708e216148e60 +2D_CFD,2D_CFD_Rand_M1.0_Eta0.1_Zeta0.1_periodic_128_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164691,2D/CFD/2D_Train_Rand/,75a34f34a3dcd9bc8ae4b7adf8bf372c +2D_CFD,2D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_512_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164692,2D/CFD/2D_Train_Rand/,04fcc44c39acea69e2869b84cd653d4a +2D_CFD,2D_CFD_Turb_M0.1_Eta1e-08_Zeta1e-08_periodic_512_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164685,2D/CFD/2D_Train_Turb/,844555000d342d2947162c6cf46798e7 +2D_CFD,2D_CFD_Turb_M1.0_Eta1e-08_Zeta1e-08_periodic_512_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164686,2D/CFD/2D_Train_Turb/,3f2c7376cde5fb072db0f9814f1c6992 +2D_CFD,2D_shock.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133193,2D/CFD/Test/2DShock/,b4098c20e2e0a25cd8ee229c9294b899 +2D_CFD,KH_M01_dk10_Re1e3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133209,2D/CFD/Test/KH/,54e278b3a7419107a5cd6a8a019225d5 +2D_CFD,KH_M01_dk1_Re1e3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133210,2D/CFD/Test/KH/,29ce31d3ce61de3b90ac0b42716c061c +2D_CFD,KH_M01_dk2_Re1e3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133211,2D/CFD/Test/KH/,de0b2a3343de028e4c162723ec571713 +2D_CFD,KH_M01_dk5_Re1e3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133212,2D/CFD/Test/KH/,5bcaed9eb7db5fd55d51f7347a938413 +2D_CFD,KH_M02_dk1_Re1e3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133213,2D/CFD/Test/KH/,fbc392f83affc2224a5703dad7bda28a +2D_CFD,KH_M04_dk1_Re1e3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133214,2D/CFD/Test/KH/,66626c935ab070a1e5df52bcf6aeccc7 +2D_CFD,KH_M1_dk1_Re1e3.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133215,2D/CFD/Test/KH/,338393c118196048a38d1df731a33b0a +2D_CFD,OTVortex.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133216,2D/CFD/Test/TOV/,3b68cb29e3d1906a19ce2cb4ce71faa4 +Darcy,2D_DarcyFlow_beta0.01_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133217,2D/DarcyFlow/,d05c287d4c0b7d3178b0097084238251 +Darcy,2D_DarcyFlow_beta0.1_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133218,2D/DarcyFlow/,294f9a03a4aa16b0e386469ca8b471be +Darcy,2D_DarcyFlow_beta1.0_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133219,2D/DarcyFlow/,81694ed31306ff2e5f6b76349b0b4389 +Darcy,2D_DarcyFlow_beta10.0_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133220,2D/DarcyFlow/,a7f23cf8011fc211b180828af39b7d1a +Darcy,2D_DarcyFlow_beta100.0_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133221,2D/DarcyFlow/,7c09f3b1bb097737d3fd52ffb0c6f1c8 +2D_ReacDiff,2D_diff-react_NA_NA.h5,https://darus.uni-stuttgart.de/api/access/datafile/133017,2D/diffusion-reaction/,b8d0b86064193195ddc30c33be5dc949 +NS_Incom,ns_incom_inhom_2d_512-0.h5,https://darus.uni-stuttgart.de/api/access/datafile/133280,2D/NS_incom/,54109d46f9c957317bd670ddb2068ac0 +NS_Incom,ns_incom_inhom_2d_512-1.h5,https://darus.uni-stuttgart.de/api/access/datafile/136439,2D/NS_incom/,e280fd3208fccb8ad5c1ff46c4796864 +NS_Incom,ns_incom_inhom_2d_512-10.h5,https://darus.uni-stuttgart.de/api/access/datafile/133309,2D/NS_incom/,f497297295ce3f2ebfefd81d00de31da +NS_Incom,ns_incom_inhom_2d_512-100.h5,https://darus.uni-stuttgart.de/api/access/datafile/133267,2D/NS_incom/,fcf3b743fc82ac8706cfa85a7b259dcd +NS_Incom,ns_incom_inhom_2d_512-101.h5,https://darus.uni-stuttgart.de/api/access/datafile/133289,2D/NS_incom/,9d1cd827c95f80975d09eae0b4e6c3da +NS_Incom,ns_incom_inhom_2d_512-102.h5,https://darus.uni-stuttgart.de/api/access/datafile/133291,2D/NS_incom/,ce1f0a7a5e52bf49db43084eef367639 +NS_Incom,ns_incom_inhom_2d_512-103.h5,https://darus.uni-stuttgart.de/api/access/datafile/133294,2D/NS_incom/,4e0a5f39d14745adf96655f220329550 +NS_Incom,ns_incom_inhom_2d_512-104.h5,https://darus.uni-stuttgart.de/api/access/datafile/133298,2D/NS_incom/,52f59ffdef04338f2630e98e32448086 +NS_Incom,ns_incom_inhom_2d_512-105.h5,https://darus.uni-stuttgart.de/api/access/datafile/133290,2D/NS_incom/,1010c2d0c807e5e355e336fe8d905067 +NS_Incom,ns_incom_inhom_2d_512-106.h5,https://darus.uni-stuttgart.de/api/access/datafile/133374,2D/NS_incom/,1127b6e69bd83b128a9420ae58187729 +NS_Incom,ns_incom_inhom_2d_512-107.h5,https://darus.uni-stuttgart.de/api/access/datafile/133375,2D/NS_incom/,38e676bfd97cc946cd6db2af48c487bf +NS_Incom,ns_incom_inhom_2d_512-108.h5,https://darus.uni-stuttgart.de/api/access/datafile/133376,2D/NS_incom/,0355430cc9cee8e921378eda7c1c6b84 +NS_Incom,ns_incom_inhom_2d_512-109.h5,https://darus.uni-stuttgart.de/api/access/datafile/133305,2D/NS_incom/,eff2dd6e71649bb3f9f37f951bd54d63 +NS_Incom,ns_incom_inhom_2d_512-11.h5,https://darus.uni-stuttgart.de/api/access/datafile/133336,2D/NS_incom/,3342616df260d3ec60fbf24dfaabf9ae +NS_Incom,ns_incom_inhom_2d_512-110.h5,https://darus.uni-stuttgart.de/api/access/datafile/133313,2D/NS_incom/,d9739419a2100b390742acf88e4548b3 +NS_Incom,ns_incom_inhom_2d_512-111.h5,https://darus.uni-stuttgart.de/api/access/datafile/133318,2D/NS_incom/,ea5239ab7204d53c476f48739a4e0d78 +NS_Incom,ns_incom_inhom_2d_512-112.h5,https://darus.uni-stuttgart.de/api/access/datafile/133324,2D/NS_incom/,cdaf99ff4ef69ee79ad60648a06259d7 +NS_Incom,ns_incom_inhom_2d_512-113.h5,https://darus.uni-stuttgart.de/api/access/datafile/133377,2D/NS_incom/,aee8074da1cf6e19ee421a7b86caeaa3 +NS_Incom,ns_incom_inhom_2d_512-114.h5,https://darus.uni-stuttgart.de/api/access/datafile/133378,2D/NS_incom/,f1b4550b00a887771e79a81c7336a4b0 +NS_Incom,ns_incom_inhom_2d_512-115.h5,https://darus.uni-stuttgart.de/api/access/datafile/133379,2D/NS_incom/,0756d685a754b84337d1de7280572c85 +NS_Incom,ns_incom_inhom_2d_512-116.h5,https://darus.uni-stuttgart.de/api/access/datafile/133380,2D/NS_incom/,f7c55c110dcf0aeaed03699bc0942f0d +NS_Incom,ns_incom_inhom_2d_512-117.h5,https://darus.uni-stuttgart.de/api/access/datafile/133381,2D/NS_incom/,375331e30111457028de38b89e03c404 +NS_Incom,ns_incom_inhom_2d_512-118.h5,https://darus.uni-stuttgart.de/api/access/datafile/133383,2D/NS_incom/,6d2fbe4a0e32e1934d349616ca6b85d8 +NS_Incom,ns_incom_inhom_2d_512-119.h5,https://darus.uni-stuttgart.de/api/access/datafile/133385,2D/NS_incom/,993112f4fe3910230005cc35099c7d59 +NS_Incom,ns_incom_inhom_2d_512-12.h5,https://darus.uni-stuttgart.de/api/access/datafile/133574,2D/NS_incom/,71f407084513d7c0e93377c1b46e4bb4 +NS_Incom,ns_incom_inhom_2d_512-120.h5,https://darus.uni-stuttgart.de/api/access/datafile/133386,2D/NS_incom/,f7c0b147e5b3f217a1d534f303ab52f4 +NS_Incom,ns_incom_inhom_2d_512-121.h5,https://darus.uni-stuttgart.de/api/access/datafile/133441,2D/NS_incom/,785d733406eeec2694eb89769f72971b +NS_Incom,ns_incom_inhom_2d_512-122.h5,https://darus.uni-stuttgart.de/api/access/datafile/133346,2D/NS_incom/,eeb2488aa0a5d1f14069dcfaaa0da71d +NS_Incom,ns_incom_inhom_2d_512-123.h5,https://darus.uni-stuttgart.de/api/access/datafile/133442,2D/NS_incom/,57e341bea82fd33ec095e7070fd1ae38 +NS_Incom,ns_incom_inhom_2d_512-124.h5,https://darus.uni-stuttgart.de/api/access/datafile/133443,2D/NS_incom/,6aba12029ab0ddb165c1c71f3b1c8f27 +NS_Incom,ns_incom_inhom_2d_512-125.h5,https://darus.uni-stuttgart.de/api/access/datafile/133444,2D/NS_incom/,2b15572d1dda8048ab9713c2318d9322 +NS_Incom,ns_incom_inhom_2d_512-126.h5,https://darus.uni-stuttgart.de/api/access/datafile/133303,2D/NS_incom/,56bcc085c7c0d303dedf147e6c608730 +NS_Incom,ns_incom_inhom_2d_512-127.h5,https://darus.uni-stuttgart.de/api/access/datafile/133292,2D/NS_incom/,71b5191ed81e86ec68f0a7bcd6a07759 +NS_Incom,ns_incom_inhom_2d_512-128.h5,https://darus.uni-stuttgart.de/api/access/datafile/133454,2D/NS_incom/,72b8110ec99098530494bbf0b64f51f7 +NS_Incom,ns_incom_inhom_2d_512-129.h5,https://darus.uni-stuttgart.de/api/access/datafile/133293,2D/NS_incom/,0e2cda441dec2c19ffaefd8ad035220e +NS_Incom,ns_incom_inhom_2d_512-13.h5,https://darus.uni-stuttgart.de/api/access/datafile/133304,2D/NS_incom/,424f6d2ab607612599645442a45234cf +NS_Incom,ns_incom_inhom_2d_512-130.h5,https://darus.uni-stuttgart.de/api/access/datafile/133295,2D/NS_incom/,8a6f6b785ce64aecc9f46a31650662d0 +NS_Incom,ns_incom_inhom_2d_512-131.h5,https://darus.uni-stuttgart.de/api/access/datafile/133275,2D/NS_incom/,3c4541c0ed411fdc9641b27978318ad6 +NS_Incom,ns_incom_inhom_2d_512-132.h5,https://darus.uni-stuttgart.de/api/access/datafile/133296,2D/NS_incom/,3c4541c0ed411fdc9641b27978318ad6 +NS_Incom,ns_incom_inhom_2d_512-133.h5,https://darus.uni-stuttgart.de/api/access/datafile/133297,2D/NS_incom/,ba2bb63092d5d0cb8fae29a0c18ff851 +NS_Incom,ns_incom_inhom_2d_512-134.h5,https://darus.uni-stuttgart.de/api/access/datafile/133269,2D/NS_incom/,4b0b6e0c3d22772ea20dc4c3e07c5e7d +NS_Incom,ns_incom_inhom_2d_512-135.h5,https://darus.uni-stuttgart.de/api/access/datafile/133274,2D/NS_incom/,e48583df141230425ea36d093b0b46f6 +NS_Incom,ns_incom_inhom_2d_512-136.h5,https://darus.uni-stuttgart.de/api/access/datafile/133299,2D/NS_incom/,4b0b6e0c3d22772ea20dc4c3e07c5e7d +NS_Incom,ns_incom_inhom_2d_512-137.h5,https://darus.uni-stuttgart.de/api/access/datafile/133300,2D/NS_incom/,e48583df141230425ea36d093b0b46f6 +NS_Incom,ns_incom_inhom_2d_512-138.h5,https://darus.uni-stuttgart.de/api/access/datafile/133302,2D/NS_incom/,294b49a65b2fc861e23dfc40098d6e91 +NS_Incom,ns_incom_inhom_2d_512-139.h5,https://darus.uni-stuttgart.de/api/access/datafile/133575,2D/NS_incom/,14c9705c20d630e08f4387cc4880a646 +NS_Incom,ns_incom_inhom_2d_512-14.h5,https://darus.uni-stuttgart.de/api/access/datafile/133281,2D/NS_incom/,f8099cba29f6374b6454424ef521b86d +NS_Incom,ns_incom_inhom_2d_512-140.h5,https://darus.uni-stuttgart.de/api/access/datafile/133698,2D/NS_incom/,f8c12c7ee3973eab55a531a24689ce95 +NS_Incom,ns_incom_inhom_2d_512-141.h5,https://darus.uni-stuttgart.de/api/access/datafile/133307,2D/NS_incom/,95a897a07713b654ce497b513ada21df +NS_Incom,ns_incom_inhom_2d_512-142.h5,https://darus.uni-stuttgart.de/api/access/datafile/133306,2D/NS_incom/,b0c964793163e49d46ddfa783b7f2a87 +NS_Incom,ns_incom_inhom_2d_512-143.h5,https://darus.uni-stuttgart.de/api/access/datafile/133590,2D/NS_incom/,b35dd28ee302e41dd9500e65e681bd93 +NS_Incom,ns_incom_inhom_2d_512-144.h5,https://darus.uni-stuttgart.de/api/access/datafile/133591,2D/NS_incom/,d7b000376bcb6cf4949e2dec4232a926 +NS_Incom,ns_incom_inhom_2d_512-145.h5,https://darus.uni-stuttgart.de/api/access/datafile/133592,2D/NS_incom/,e90fad8268cf369136cebba09f3b06e8 +NS_Incom,ns_incom_inhom_2d_512-146.h5,https://darus.uni-stuttgart.de/api/access/datafile/133593,2D/NS_incom/,568fac9aab2f599c20da783e9c7f3f0a +NS_Incom,ns_incom_inhom_2d_512-147.h5,https://darus.uni-stuttgart.de/api/access/datafile/133599,2D/NS_incom/,43194bdefada7b59cac975cc08e6d12d +NS_Incom,ns_incom_inhom_2d_512-148.h5,https://darus.uni-stuttgart.de/api/access/datafile/133600,2D/NS_incom/,41a78815f65ff7b607d69dd508b3ffea +NS_Incom,ns_incom_inhom_2d_512-149.h5,https://darus.uni-stuttgart.de/api/access/datafile/133601,2D/NS_incom/,3ebb3a014d7c33787cda458a689f7728 +NS_Incom,ns_incom_inhom_2d_512-15.h5,https://darus.uni-stuttgart.de/api/access/datafile/133609,2D/NS_incom/,74dd42aaf8b809925ff6dce345be11fd +NS_Incom,ns_incom_inhom_2d_512-150.h5,https://darus.uni-stuttgart.de/api/access/datafile/133602,2D/NS_incom/,562bd9a0d25ce189a9635d385e3d6a3c +NS_Incom,ns_incom_inhom_2d_512-151.h5,https://darus.uni-stuttgart.de/api/access/datafile/133271,2D/NS_incom/,1acb06d08a4ae3a17ddbe04a2e3fd5cf +NS_Incom,ns_incom_inhom_2d_512-152.h5,https://darus.uni-stuttgart.de/api/access/datafile/133603,2D/NS_incom/,0023490cb7c00a4a75f9a252af9be71c +NS_Incom,ns_incom_inhom_2d_512-153.h5,https://darus.uni-stuttgart.de/api/access/datafile/133594,2D/NS_incom/,669788ec079a3e695b68821c17d79a26 +NS_Incom,ns_incom_inhom_2d_512-154.h5,https://darus.uni-stuttgart.de/api/access/datafile/133604,2D/NS_incom/,547cfa5abca50bcd927630b8264d6876 +NS_Incom,ns_incom_inhom_2d_512-155.h5,https://darus.uni-stuttgart.de/api/access/datafile/133605,2D/NS_incom/,f28fcf6d99c0d1f1bb08a82be94114eb +NS_Incom,ns_incom_inhom_2d_512-156.h5,https://darus.uni-stuttgart.de/api/access/datafile/133606,2D/NS_incom/,8c52c6258bdbbc7c4fe1cb7892b5bd23 +NS_Incom,ns_incom_inhom_2d_512-157.h5,https://darus.uni-stuttgart.de/api/access/datafile/133607,2D/NS_incom/,91f4ab6bd4eef91ab43908290539bdab +NS_Incom,ns_incom_inhom_2d_512-158.h5,https://darus.uni-stuttgart.de/api/access/datafile/133321,2D/NS_incom/,299521efb9f871919b003e86d0e8d628 +NS_Incom,ns_incom_inhom_2d_512-159.h5,https://darus.uni-stuttgart.de/api/access/datafile/133608,2D/NS_incom/,54b54c80708366c74d4aee82688be26e +NS_Incom,ns_incom_inhom_2d_512-16.h5,https://darus.uni-stuttgart.de/api/access/datafile/133311,2D/NS_incom/,b83a1465fee7f5dcb9826d1bbb4b6f71 +NS_Incom,ns_incom_inhom_2d_512-160.h5,https://darus.uni-stuttgart.de/api/access/datafile/133610,2D/NS_incom/,e90d31aff2172f508846473126c78947 +NS_Incom,ns_incom_inhom_2d_512-161.h5,https://darus.uni-stuttgart.de/api/access/datafile/133611,2D/NS_incom/,24b4b52f0a03dc52f7159ef10652c40a +NS_Incom,ns_incom_inhom_2d_512-162.h5,https://darus.uni-stuttgart.de/api/access/datafile/133683,2D/NS_incom/,87d592f2be6a76751b2b36e8bf0d0a88 +NS_Incom,ns_incom_inhom_2d_512-163.h5,https://darus.uni-stuttgart.de/api/access/datafile/133612,2D/NS_incom/,2acf733c2833fb815487271bbb5ce66f +NS_Incom,ns_incom_inhom_2d_512-164.h5,https://darus.uni-stuttgart.de/api/access/datafile/136435,2D/NS_incom/,8470816b05ecbabd2ec6036f346d4d79 +NS_Incom,ns_incom_inhom_2d_512-165.h5,https://darus.uni-stuttgart.de/api/access/datafile/133685,2D/NS_incom/,61d677b1edd03987c26648d7d136672d +NS_Incom,ns_incom_inhom_2d_512-166.h5,https://darus.uni-stuttgart.de/api/access/datafile/133308,2D/NS_incom/,4f1dd0e87f4ad925dc82444da4df7d3a +NS_Incom,ns_incom_inhom_2d_512-167.h5,https://darus.uni-stuttgart.de/api/access/datafile/133684,2D/NS_incom/,8e90f2ab307a6c5f8a2e1a3befa02d8f +NS_Incom,ns_incom_inhom_2d_512-168.h5,https://darus.uni-stuttgart.de/api/access/datafile/133701,2D/NS_incom/,1f8ef01efadb38eb638c383e242d9511 +NS_Incom,ns_incom_inhom_2d_512-169.h5,https://darus.uni-stuttgart.de/api/access/datafile/133310,2D/NS_incom/,d2ae69de21cb38a24dbc174f34afeb70 +NS_Incom,ns_incom_inhom_2d_512-17.h5,https://darus.uni-stuttgart.de/api/access/datafile/133325,2D/NS_incom/,9c206312c353aa6dcaa3144d53753075 +NS_Incom,ns_incom_inhom_2d_512-170.h5,https://darus.uni-stuttgart.de/api/access/datafile/133312,2D/NS_incom/,d7f9e313387dfb3efc356c1d5606ebc3 +NS_Incom,ns_incom_inhom_2d_512-171.h5,https://darus.uni-stuttgart.de/api/access/datafile/133679,2D/NS_incom/,f6ade3872e3e0db0d6b8af4cc9fd7ed6 +NS_Incom,ns_incom_inhom_2d_512-172.h5,https://darus.uni-stuttgart.de/api/access/datafile/133314,2D/NS_incom/,986851271af43e55264a8db4229b2049 +NS_Incom,ns_incom_inhom_2d_512-173.h5,https://darus.uni-stuttgart.de/api/access/datafile/133315,2D/NS_incom/,e064a3b043136572de413744f470e795 +NS_Incom,ns_incom_inhom_2d_512-174.h5,https://darus.uni-stuttgart.de/api/access/datafile/133316,2D/NS_incom/,d4906a07e5e94b6e147570613205d19d +NS_Incom,ns_incom_inhom_2d_512-175.h5,https://darus.uni-stuttgart.de/api/access/datafile/133317,2D/NS_incom/,7d7230752c444c2570396927e9f9b4bf +NS_Incom,ns_incom_inhom_2d_512-176.h5,https://darus.uni-stuttgart.de/api/access/datafile/133319,2D/NS_incom/,47a59095761d69cecc886ddb3027b556 +NS_Incom,ns_incom_inhom_2d_512-177.h5,https://darus.uni-stuttgart.de/api/access/datafile/133320,2D/NS_incom/,979de8be5c402f6dc50060e92ca6245a +NS_Incom,ns_incom_inhom_2d_512-178.h5,https://darus.uni-stuttgart.de/api/access/datafile/133342,2D/NS_incom/,a72fc652b38b2dc0b0aa406bb35cd676 +NS_Incom,ns_incom_inhom_2d_512-179.h5,https://darus.uni-stuttgart.de/api/access/datafile/133697,2D/NS_incom/,6dd8a78a85ad7c0e03daaabcb3826fdb +NS_Incom,ns_incom_inhom_2d_512-18.h5,https://darus.uni-stuttgart.de/api/access/datafile/136437,2D/NS_incom/,272c1cdf65a90dc3e4308df8a15ef7dc +NS_Incom,ns_incom_inhom_2d_512-180.h5,https://darus.uni-stuttgart.de/api/access/datafile/133326,2D/NS_incom/,f9ffc612a7438dc17524818b48acb7ac +NS_Incom,ns_incom_inhom_2d_512-181.h5,https://darus.uni-stuttgart.de/api/access/datafile/133327,2D/NS_incom/,78ee160778aab06d0e88507944f661b4 +NS_Incom,ns_incom_inhom_2d_512-182.h5,https://darus.uni-stuttgart.de/api/access/datafile/133328,2D/NS_incom/,e9a36056e3e9fe23937fbf65e0dfc16a +NS_Incom,ns_incom_inhom_2d_512-183.h5,https://darus.uni-stuttgart.de/api/access/datafile/133329,2D/NS_incom/,5107343596d93f8739b4f09926daeef7 +NS_Incom,ns_incom_inhom_2d_512-184.h5,https://darus.uni-stuttgart.de/api/access/datafile/133330,2D/NS_incom/,f4a7aff84a03fff6e1958617abf373c4 +NS_Incom,ns_incom_inhom_2d_512-185.h5,https://darus.uni-stuttgart.de/api/access/datafile/133331,2D/NS_incom/,f59cb88424414a28aad308499935619b +NS_Incom,ns_incom_inhom_2d_512-186.h5,https://darus.uni-stuttgart.de/api/access/datafile/133666,2D/NS_incom/,d9bc339d2d2047f8f6a369ad2e4e9357 +NS_Incom,ns_incom_inhom_2d_512-187.h5,https://darus.uni-stuttgart.de/api/access/datafile/136436,2D/NS_incom/,4d0461fc34683816d5906d4809c84432 +NS_Incom,ns_incom_inhom_2d_512-188.h5,https://darus.uni-stuttgart.de/api/access/datafile/133332,2D/NS_incom/,3565b7f3ab620c3716a25865670b98fb +NS_Incom,ns_incom_inhom_2d_512-189.h5,https://darus.uni-stuttgart.de/api/access/datafile/133668,2D/NS_incom/,d385ef198ae29ef5eff9aad6d129621a +NS_Incom,ns_incom_inhom_2d_512-19.h5,https://darus.uni-stuttgart.de/api/access/datafile/133335,2D/NS_incom/,f67fec4521fc9a113cd16de7f7c283ff +NS_Incom,ns_incom_inhom_2d_512-190.h5,https://darus.uni-stuttgart.de/api/access/datafile/136438,2D/NS_incom/,2dce11f971c3561083d580f7299f3e02 +NS_Incom,ns_incom_inhom_2d_512-191.h5,https://darus.uni-stuttgart.de/api/access/datafile/133663,2D/NS_incom/,fda967e4ca21881b00d5c67b984d4c78 +NS_Incom,ns_incom_inhom_2d_512-192.h5,https://darus.uni-stuttgart.de/api/access/datafile/133664,2D/NS_incom/,85542754064fcab3675930ac607b68e0 +NS_Incom,ns_incom_inhom_2d_512-193.h5,https://darus.uni-stuttgart.de/api/access/datafile/133371,2D/NS_incom/,2f67289963f5d9d3041acf67f45a38b7 +NS_Incom,ns_incom_inhom_2d_512-194.h5,https://darus.uni-stuttgart.de/api/access/datafile/133669,2D/NS_incom/,609d2d5dbd19ccdb6cee8ab67df66c38 +NS_Incom,ns_incom_inhom_2d_512-195.h5,https://darus.uni-stuttgart.de/api/access/datafile/133699,2D/NS_incom/,b39dbe9518dbc19da27345ce569489e4 +NS_Incom,ns_incom_inhom_2d_512-196.h5,https://darus.uni-stuttgart.de/api/access/datafile/133665,2D/NS_incom/,fc636092ce4175b3d908d156db66deb8 +NS_Incom,ns_incom_inhom_2d_512-197.h5,https://darus.uni-stuttgart.de/api/access/datafile/133702,2D/NS_incom/,0b3bd457d75ccb3120f649bcdb0089f2 +NS_Incom,ns_incom_inhom_2d_512-198.h5,https://darus.uni-stuttgart.de/api/access/datafile/133368,2D/NS_incom/,526cce686110114720bc9d84eb8d16ee +NS_Incom,ns_incom_inhom_2d_512-199.h5,https://darus.uni-stuttgart.de/api/access/datafile/133667,2D/NS_incom/,123272aaaf52747c39e8c32b4955390b +NS_Incom,ns_incom_inhom_2d_512-2.h5,https://darus.uni-stuttgart.de/api/access/datafile/133721,2D/NS_incom/,1d7a2aac41a410bea6c887f624273d20 +NS_Incom,ns_incom_inhom_2d_512-20.h5,https://darus.uni-stuttgart.de/api/access/datafile/133283,2D/NS_incom/,859c8c8cede4fd611548af291acfa362 +NS_Incom,ns_incom_inhom_2d_512-200.h5,https://darus.uni-stuttgart.de/api/access/datafile/136440,2D/NS_incom/,621c5bdf211ab8aa6d586ede4b51c5a2 +NS_Incom,ns_incom_inhom_2d_512-201.h5,https://darus.uni-stuttgart.de/api/access/datafile/133367,2D/NS_incom/,89b6667026a2f289b81e48b6d75ef4d6 +NS_Incom,ns_incom_inhom_2d_512-202.h5,https://darus.uni-stuttgart.de/api/access/datafile/133661,2D/NS_incom/,a70f3e38089caec832075b7c3ce86744 +NS_Incom,ns_incom_inhom_2d_512-203.h5,https://darus.uni-stuttgart.de/api/access/datafile/133660,2D/NS_incom/,f0bc2659a6f671813e88409bdb08531a +NS_Incom,ns_incom_inhom_2d_512-204.h5,https://darus.uni-stuttgart.de/api/access/datafile/133671,2D/NS_incom/,c109d2aa5cf0688c5f0565b03d1e3ce0 +NS_Incom,ns_incom_inhom_2d_512-205.h5,https://darus.uni-stuttgart.de/api/access/datafile/133688,2D/NS_incom/,2fa87303d667e4d50935fe26f04437b4 +NS_Incom,ns_incom_inhom_2d_512-206.h5,https://darus.uni-stuttgart.de/api/access/datafile/136441,2D/NS_incom/,28a557b57ffce43cb9de4cc48642fc05 +NS_Incom,ns_incom_inhom_2d_512-207.h5,https://darus.uni-stuttgart.de/api/access/datafile/133282,2D/NS_incom/,01a63c9ad30ac4d140f86bb2fa929274 +NS_Incom,ns_incom_inhom_2d_512-208.h5,https://darus.uni-stuttgart.de/api/access/datafile/133340,2D/NS_incom/,d2aa05820caf6a38be4711ef797d3bf4 +NS_Incom,ns_incom_inhom_2d_512-209.h5,https://darus.uni-stuttgart.de/api/access/datafile/133333,2D/NS_incom/,9e17d624a2eea451f0956f9184ce63d8 +NS_Incom,ns_incom_inhom_2d_512-21.h5,https://darus.uni-stuttgart.de/api/access/datafile/133334,2D/NS_incom/,859c8c8cede4fd611548af291acfa362 +NS_Incom,ns_incom_inhom_2d_512-210.h5,https://darus.uni-stuttgart.de/api/access/datafile/133703,2D/NS_incom/,fb5de2b943891ed1ec1bbb682520bc78 +NS_Incom,ns_incom_inhom_2d_512-211.h5,https://darus.uni-stuttgart.de/api/access/datafile/133286,2D/NS_incom/,24f8fe56ffe24baaa03f6bf12762ef95 +NS_Incom,ns_incom_inhom_2d_512-212.h5,https://darus.uni-stuttgart.de/api/access/datafile/133284,2D/NS_incom/,8e4561fe4ae69d62000df136f980fae3 +NS_Incom,ns_incom_inhom_2d_512-213.h5,https://darus.uni-stuttgart.de/api/access/datafile/133285,2D/NS_incom/,d1101482e0d766334b92c4d0e11d73f3 +NS_Incom,ns_incom_inhom_2d_512-214.h5,https://darus.uni-stuttgart.de/api/access/datafile/133337,2D/NS_incom/,d1101482e0d766334b92c4d0e11d73f3 +NS_Incom,ns_incom_inhom_2d_512-215.h5,https://darus.uni-stuttgart.de/api/access/datafile/133287,2D/NS_incom/,f615b96dd71864673bba274240acfb12 +NS_Incom,ns_incom_inhom_2d_512-216.h5,https://darus.uni-stuttgart.de/api/access/datafile/133288,2D/NS_incom/,892986d49b49b431f3f172d7baa9977b +NS_Incom,ns_incom_inhom_2d_512-217.h5,https://darus.uni-stuttgart.de/api/access/datafile/133338,2D/NS_incom/,a471e758505de149dba27c9fa2d29198 +NS_Incom,ns_incom_inhom_2d_512-218.h5,https://darus.uni-stuttgart.de/api/access/datafile/133339,2D/NS_incom/,f615b96dd71864673bba274240acfb12 +NS_Incom,ns_incom_inhom_2d_512-219.h5,https://darus.uni-stuttgart.de/api/access/datafile/133341,2D/NS_incom/,1965889c070e802965ebe19b01cdf2db +NS_Incom,ns_incom_inhom_2d_512-22.h5,https://darus.uni-stuttgart.de/api/access/datafile/133344,2D/NS_incom/,2235dd8e60d0d43b49aba99819ddb8d3 +NS_Incom,ns_incom_inhom_2d_512-220.h5,https://darus.uni-stuttgart.de/api/access/datafile/133343,2D/NS_incom/,588ce18cf51ca3517527f3dd0b0ee443 +NS_Incom,ns_incom_inhom_2d_512-221.h5,https://darus.uni-stuttgart.de/api/access/datafile/133345,2D/NS_incom/,e57ad44f1ae4fa9f31ddff57578507ce +NS_Incom,ns_incom_inhom_2d_512-222.h5,https://darus.uni-stuttgart.de/api/access/datafile/133347,2D/NS_incom/,892986d49b49b431f3f172d7baa9977b +NS_Incom,ns_incom_inhom_2d_512-223.h5,https://darus.uni-stuttgart.de/api/access/datafile/133369,2D/NS_incom/,9da97766e3b311af6d79a946f8db5184 +NS_Incom,ns_incom_inhom_2d_512-224.h5,https://darus.uni-stuttgart.de/api/access/datafile/133370,2D/NS_incom/,44a864a879590427273a3dbf1aa2c190 +NS_Incom,ns_incom_inhom_2d_512-225.h5,https://darus.uni-stuttgart.de/api/access/datafile/133372,2D/NS_incom/,8e4561fe4ae69d62000df136f980fae3 +NS_Incom,ns_incom_inhom_2d_512-226.h5,https://darus.uni-stuttgart.de/api/access/datafile/133373,2D/NS_incom/,adf37d1dd066c2dc8fb3886369b5434f +NS_Incom,ns_incom_inhom_2d_512-227.h5,https://darus.uni-stuttgart.de/api/access/datafile/136442,2D/NS_incom/,de46866be488bb59d0d1c8ddc13a4dcb +NS_Incom,ns_incom_inhom_2d_512-228.h5,https://darus.uni-stuttgart.de/api/access/datafile/133723,2D/NS_incom/,43975b57db7c922f94703d1c63fdaebe +NS_Incom,ns_incom_inhom_2d_512-229.h5,https://darus.uni-stuttgart.de/api/access/datafile/136443,2D/NS_incom/,805f1a7336301b2bcc4a9f0fcce43864 +NS_Incom,ns_incom_inhom_2d_512-23.h5,https://darus.uni-stuttgart.de/api/access/datafile/133617,2D/NS_incom/,d026b809986795e73f13368a763e56d2 +NS_Incom,ns_incom_inhom_2d_512-230.h5,https://darus.uni-stuttgart.de/api/access/datafile/136445,2D/NS_incom/,e43acdfb560a45181111256e4ce45cc1 +NS_Incom,ns_incom_inhom_2d_512-231.h5,https://darus.uni-stuttgart.de/api/access/datafile/133613,2D/NS_incom/,f56497f0ba427fdf9e215c885c6c1d3a +NS_Incom,ns_incom_inhom_2d_512-232.h5,https://darus.uni-stuttgart.de/api/access/datafile/133692,2D/NS_incom/,599ff6b1ea10572a829a13e5fc9c2523 +NS_Incom,ns_incom_inhom_2d_512-233.h5,https://darus.uni-stuttgart.de/api/access/datafile/133614,2D/NS_incom/,b0cbda1f7435ccbe04bc6a09840624df +NS_Incom,ns_incom_inhom_2d_512-234.h5,https://darus.uni-stuttgart.de/api/access/datafile/133700,2D/NS_incom/,6331197878327890af4da21d5c4065bf +NS_Incom,ns_incom_inhom_2d_512-235.h5,https://darus.uni-stuttgart.de/api/access/datafile/133693,2D/NS_incom/,dd937dace08671e5b61a8f960d845561 +NS_Incom,ns_incom_inhom_2d_512-236.h5,https://darus.uni-stuttgart.de/api/access/datafile/133615,2D/NS_incom/,07a4596a86c3099cf7f6e7e5c47b71b6 +NS_Incom,ns_incom_inhom_2d_512-237.h5,https://darus.uni-stuttgart.de/api/access/datafile/133616,2D/NS_incom/,4c977f28e59fd8cc877f6d95d96011e5 +NS_Incom,ns_incom_inhom_2d_512-238.h5,https://darus.uni-stuttgart.de/api/access/datafile/133695,2D/NS_incom/,9d3bad375edf84cd86a648caca49b6b5 +NS_Incom,ns_incom_inhom_2d_512-239.h5,https://darus.uni-stuttgart.de/api/access/datafile/136446,2D/NS_incom/,bae42dea9f17c4ca5ec33997e3eb966f +NS_Incom,ns_incom_inhom_2d_512-24.h5,https://darus.uni-stuttgart.de/api/access/datafile/133618,2D/NS_incom/,c7e0bdab4f18748914263534cf21422d +NS_Incom,ns_incom_inhom_2d_512-240.h5,https://darus.uni-stuttgart.de/api/access/datafile/133268,2D/NS_incom/,10b14e2a9eb6fcf5a88a7a5447481aa7 +NS_Incom,ns_incom_inhom_2d_512-241.h5,https://darus.uni-stuttgart.de/api/access/datafile/136447,2D/NS_incom/,0aa745b2377c49379f91f7f000d8f23e +NS_Incom,ns_incom_inhom_2d_512-242.h5,https://darus.uni-stuttgart.de/api/access/datafile/136448,2D/NS_incom/,79033da1fbd3167b560769c8292c15f4 +NS_Incom,ns_incom_inhom_2d_512-243.h5,https://darus.uni-stuttgart.de/api/access/datafile/136450,2D/NS_incom/,896aa796fbf2052da573c30694744794 +NS_Incom,ns_incom_inhom_2d_512-244.h5,https://darus.uni-stuttgart.de/api/access/datafile/136451,2D/NS_incom/,3063dee5e25740331b4c606939c3960e +NS_Incom,ns_incom_inhom_2d_512-245.h5,https://darus.uni-stuttgart.de/api/access/datafile/136457,2D/NS_incom/,daf18bcc864d5db453f7321de3c199bf +NS_Incom,ns_incom_inhom_2d_512-246.h5,https://darus.uni-stuttgart.de/api/access/datafile/136458,2D/NS_incom/,8f97135545fdea8eb504c6f76878b77c +NS_Incom,ns_incom_inhom_2d_512-247.h5,https://darus.uni-stuttgart.de/api/access/datafile/136459,2D/NS_incom/,2e1f4cf38135b2f07a72462e0f995c2d +NS_Incom,ns_incom_inhom_2d_512-248.h5,https://darus.uni-stuttgart.de/api/access/datafile/136460,2D/NS_incom/,edaec61956c5af7be9f8996b8498c8fd +NS_Incom,ns_incom_inhom_2d_512-249.h5,https://darus.uni-stuttgart.de/api/access/datafile/136461,2D/NS_incom/,b05738fc7cdf18fecca88361ba58875f +NS_Incom,ns_incom_inhom_2d_512-25.h5,https://darus.uni-stuttgart.de/api/access/datafile/133619,2D/NS_incom/,d5e37c12c61353533ab82c6c9e48e93f +NS_Incom,ns_incom_inhom_2d_512-250.h5,https://darus.uni-stuttgart.de/api/access/datafile/166273,2D/NS_incom/,c9731467c80b809ce0f3f337ce3ed71e +NS_Incom,ns_incom_inhom_2d_512-251.h5,https://darus.uni-stuttgart.de/api/access/datafile/166274,2D/NS_incom/,3a515e95b641360734702248b3bfcb6d +NS_Incom,ns_incom_inhom_2d_512-252.h5,https://darus.uni-stuttgart.de/api/access/datafile/166275,2D/NS_incom/,283244af2adc86953dd6021bf6722ba4 +NS_Incom,ns_incom_inhom_2d_512-253.h5,https://darus.uni-stuttgart.de/api/access/datafile/166276,2D/NS_incom/,a471e758505de149dba27c9fa2d29198 +NS_Incom,ns_incom_inhom_2d_512-254.h5,https://darus.uni-stuttgart.de/api/access/datafile/166277,2D/NS_incom/,9da97766e3b311af6d79a946f8db5184 +NS_Incom,ns_incom_inhom_2d_512-255.h5,https://darus.uni-stuttgart.de/api/access/datafile/166278,2D/NS_incom/,1965889c070e802965ebe19b01cdf2db +NS_Incom,ns_incom_inhom_2d_512-256.h5,https://darus.uni-stuttgart.de/api/access/datafile/166279,2D/NS_incom/,588ce18cf51ca3517527f3dd0b0ee443 +NS_Incom,ns_incom_inhom_2d_512-257.h5,https://darus.uni-stuttgart.de/api/access/datafile/166281,2D/NS_incom/,44a864a879590427273a3dbf1aa2c190 +NS_Incom,ns_incom_inhom_2d_512-258.h5,https://darus.uni-stuttgart.de/api/access/datafile/166282,2D/NS_incom/,adf37d1dd066c2dc8fb3886369b5434f +NS_Incom,ns_incom_inhom_2d_512-259.h5,https://darus.uni-stuttgart.de/api/access/datafile/166283,2D/NS_incom/,588200460d5168ce1ba582101aa3547e +NS_Incom,ns_incom_inhom_2d_512-26.h5,https://darus.uni-stuttgart.de/api/access/datafile/133689,2D/NS_incom/,563b6261312a32460d848e5b9ca82f79 +NS_Incom,ns_incom_inhom_2d_512-260.h5,https://darus.uni-stuttgart.de/api/access/datafile/166284,2D/NS_incom/,2e1e5883bec2e8e50048738266785bae +NS_Incom,ns_incom_inhom_2d_512-261.h5,https://darus.uni-stuttgart.de/api/access/datafile/166285,2D/NS_incom/,b5f6b09ad77bbcafbc7d52710388669c +NS_Incom,ns_incom_inhom_2d_512-262.h5,https://darus.uni-stuttgart.de/api/access/datafile/166286,2D/NS_incom/,8a8026b2cd2e59a7b1d796dca4404e3a +NS_Incom,ns_incom_inhom_2d_512-263.h5,https://darus.uni-stuttgart.de/api/access/datafile/166287,2D/NS_incom/,8958869c63e7a496816d5c551d4d8115 +NS_Incom,ns_incom_inhom_2d_512-264.h5,https://darus.uni-stuttgart.de/api/access/datafile/166288,2D/NS_incom/,69e4911decfab00ff945884915084104 +NS_Incom,ns_incom_inhom_2d_512-265.h5,https://darus.uni-stuttgart.de/api/access/datafile/166290,2D/NS_incom/,3516776e1d47d573ee91815ca70360b6 +NS_Incom,ns_incom_inhom_2d_512-266.h5,https://darus.uni-stuttgart.de/api/access/datafile/166291,2D/NS_incom/,f50c0e5f4ad659c1b5d6f7343e072bdf +NS_Incom,ns_incom_inhom_2d_512-267.h5,https://darus.uni-stuttgart.de/api/access/datafile/166292,2D/NS_incom/,85a8bd3bd539acb7685393393294e916 +NS_Incom,ns_incom_inhom_2d_512-268.h5,https://darus.uni-stuttgart.de/api/access/datafile/166293,2D/NS_incom/,3b43956b04c4e0e619ba671cdfbbc834 +NS_Incom,ns_incom_inhom_2d_512-269.h5,https://darus.uni-stuttgart.de/api/access/datafile/166294,2D/NS_incom/,1ca4a6528b2d73b55042eb6dbe487ff1 +NS_Incom,ns_incom_inhom_2d_512-27.h5,https://darus.uni-stuttgart.de/api/access/datafile/133620,2D/NS_incom/,3536aace9614abf8868d9c03bbccf9d4 +NS_Incom,ns_incom_inhom_2d_512-270.h5,https://darus.uni-stuttgart.de/api/access/datafile/166295,2D/NS_incom/,8bf9ac2f7dc47e3617ba9de2ac10d582 +NS_Incom,ns_incom_inhom_2d_512-271.h5,https://darus.uni-stuttgart.de/api/access/datafile/166296,2D/NS_incom/,907d7cd2c84312f52f94b27197128d00 +NS_Incom,ns_incom_inhom_2d_512-272.h5,https://darus.uni-stuttgart.de/api/access/datafile/166297,2D/NS_incom/,7d1dc19e1b08c09d86bed9a0c7448ceb +NS_Incom,ns_incom_inhom_2d_512-273.h5,https://darus.uni-stuttgart.de/api/access/datafile/166298,2D/NS_incom/,e5e51712cc8921fcc4e89b9fe4690c7f +NS_Incom,ns_incom_inhom_2d_512-274.h5,https://darus.uni-stuttgart.de/api/access/datafile/166299,2D/NS_incom/,5a3ea4a7cdeade672038ffdbf3a7886a +NS_Incom,ns_incom_inhom_2d_512-28.h5,https://darus.uni-stuttgart.de/api/access/datafile/136462,2D/NS_incom/,aef3e1e349adc36e3d738d7d5534548b +NS_Incom,ns_incom_inhom_2d_512-29.h5,https://darus.uni-stuttgart.de/api/access/datafile/136463,2D/NS_incom/,b5f7a9f28aab2baf06d6129a436dca79 +NS_Incom,ns_incom_inhom_2d_512-3.h5,https://darus.uni-stuttgart.de/api/access/datafile/136466,2D/NS_incom/,e9c92a19854e96c918d3a46d39994be4 +NS_Incom,ns_incom_inhom_2d_512-30.h5,https://darus.uni-stuttgart.de/api/access/datafile/136464,2D/NS_incom/,fa2275aaf761334ec2d80affb10d324f +NS_Incom,ns_incom_inhom_2d_512-31.h5,https://darus.uni-stuttgart.de/api/access/datafile/133621,2D/NS_incom/,a30bc0bb2120dcb42a8335c6d264ac7d +NS_Incom,ns_incom_inhom_2d_512-32.h5,https://darus.uni-stuttgart.de/api/access/datafile/133626,2D/NS_incom/,6579fe51ea46095aa1cd336be4701b34 +NS_Incom,ns_incom_inhom_2d_512-33.h5,https://darus.uni-stuttgart.de/api/access/datafile/133622,2D/NS_incom/,63ee81664764659f3b6436c055981c95 +NS_Incom,ns_incom_inhom_2d_512-34.h5,https://darus.uni-stuttgart.de/api/access/datafile/133623,2D/NS_incom/,e7c32d512bf95b97699c36381fea07ca +NS_Incom,ns_incom_inhom_2d_512-35.h5,https://darus.uni-stuttgart.de/api/access/datafile/133624,2D/NS_incom/,c2715bd899018326e9679775d57e1dfd +NS_Incom,ns_incom_inhom_2d_512-36.h5,https://darus.uni-stuttgart.de/api/access/datafile/136465,2D/NS_incom/,b085ae82d91f9baa2a7f9c8720f5ca48 +NS_Incom,ns_incom_inhom_2d_512-37.h5,https://darus.uni-stuttgart.de/api/access/datafile/133625,2D/NS_incom/,80c39f8aaf4168974f8bde94df177ab6 +NS_Incom,ns_incom_inhom_2d_512-38.h5,https://darus.uni-stuttgart.de/api/access/datafile/133645,2D/NS_incom/,d01ee492fcdb64cd52b7d0ae7b606370 +NS_Incom,ns_incom_inhom_2d_512-39.h5,https://darus.uni-stuttgart.de/api/access/datafile/133630,2D/NS_incom/,3f4f083589ab6d314584fa7ec795669e +NS_Incom,ns_incom_inhom_2d_512-4.h5,https://darus.uni-stuttgart.de/api/access/datafile/166289,2D/NS_incom/,306a56ac14ea5686921cbb3f09c7dcb6 +NS_Incom,ns_incom_inhom_2d_512-40.h5,https://darus.uni-stuttgart.de/api/access/datafile/133627,2D/NS_incom/,c906a76f920bb63d575e347e76e47caf +NS_Incom,ns_incom_inhom_2d_512-41.h5,https://darus.uni-stuttgart.de/api/access/datafile/133696,2D/NS_incom/,b55a275050a22b9c515d9d896dba6da0 +NS_Incom,ns_incom_inhom_2d_512-42.h5,https://darus.uni-stuttgart.de/api/access/datafile/133642,2D/NS_incom/,5aa8b626b64ba979905797b9e56dd34c +NS_Incom,ns_incom_inhom_2d_512-43.h5,https://darus.uni-stuttgart.de/api/access/datafile/133650,2D/NS_incom/,d6ef4cf80d79f696dfc34d4947b14a77 +NS_Incom,ns_incom_inhom_2d_512-44.h5,https://darus.uni-stuttgart.de/api/access/datafile/133686,2D/NS_incom/,b8527daec6934026f5128a6d1f5db1e5 +NS_Incom,ns_incom_inhom_2d_512-45.h5,https://darus.uni-stuttgart.de/api/access/datafile/136467,2D/NS_incom/,7899ab1897fbd4665e4d7daea7c90bc2 +NS_Incom,ns_incom_inhom_2d_512-46.h5,https://darus.uni-stuttgart.de/api/access/datafile/136468,2D/NS_incom/,6853e9a78b6a2b2b756c4e6f8bcdac3e +NS_Incom,ns_incom_inhom_2d_512-47.h5,https://darus.uni-stuttgart.de/api/access/datafile/136469,2D/NS_incom/,4737e1e1283748baae9617ceff1e777a +NS_Incom,ns_incom_inhom_2d_512-48.h5,https://darus.uni-stuttgart.de/api/access/datafile/136470,2D/NS_incom/,5d3abc3ecbddeacb5270030954cc8064 +NS_Incom,ns_incom_inhom_2d_512-5.h5,https://darus.uni-stuttgart.de/api/access/datafile/136471,2D/NS_incom/,6a939d7418593d70faf7f90a191f841d +NS_Incom,ns_incom_inhom_2d_512-50.h5,https://darus.uni-stuttgart.de/api/access/datafile/133678,2D/NS_incom/,f30303f48d036e9aeac07fd70e61c97b +NS_Incom,ns_incom_inhom_2d_512-51.h5,https://darus.uni-stuttgart.de/api/access/datafile/133659,2D/NS_incom/,e82c2adbbc22b59168b7b587012d296d +NS_Incom,ns_incom_inhom_2d_512-52.h5,https://darus.uni-stuttgart.de/api/access/datafile/133657,2D/NS_incom/,1a072500b92f728f24526182a8e22d71 +NS_Incom,ns_incom_inhom_2d_512-53.h5,https://darus.uni-stuttgart.de/api/access/datafile/133670,2D/NS_incom/,e82c2adbbc22b59168b7b587012d296d +NS_Incom,ns_incom_inhom_2d_512-54.h5,https://darus.uni-stuttgart.de/api/access/datafile/133672,2D/NS_incom/,3516776e1d47d573ee91815ca70360b6 +NS_Incom,ns_incom_inhom_2d_512-55.h5,https://darus.uni-stuttgart.de/api/access/datafile/133662,2D/NS_incom/,49732762c231a410904cf8bddb0444cd +NS_Incom,ns_incom_inhom_2d_512-56.h5,https://darus.uni-stuttgart.de/api/access/datafile/133673,2D/NS_incom/,f50c0e5f4ad659c1b5d6f7343e072bdf +NS_Incom,ns_incom_inhom_2d_512-57.h5,https://darus.uni-stuttgart.de/api/access/datafile/133674,2D/NS_incom/,1a072500b92f728f24526182a8e22d71 +NS_Incom,ns_incom_inhom_2d_512-58.h5,https://darus.uni-stuttgart.de/api/access/datafile/133658,2D/NS_incom/,ae1e262d7f1a56369b13b6d5de80e03d +NS_Incom,ns_incom_inhom_2d_512-59.h5,https://darus.uni-stuttgart.de/api/access/datafile/133656,2D/NS_incom/,b99c5b112d4a7dfca053cb2c78ebe581 +NS_Incom,ns_incom_inhom_2d_512-6.h5,https://darus.uni-stuttgart.de/api/access/datafile/136472,2D/NS_incom/,36156ac1f8cecb47d076a1ee0707dec5 +NS_Incom,ns_incom_inhom_2d_512-60.h5,https://darus.uni-stuttgart.de/api/access/datafile/133675,2D/NS_incom/,85a8bd3bd539acb7685393393294e916 +NS_Incom,ns_incom_inhom_2d_512-61.h5,https://darus.uni-stuttgart.de/api/access/datafile/133676,2D/NS_incom/,3b43956b04c4e0e619ba671cdfbbc834 +NS_Incom,ns_incom_inhom_2d_512-62.h5,https://darus.uni-stuttgart.de/api/access/datafile/133677,2D/NS_incom/,8bf9ac2f7dc47e3617ba9de2ac10d582 +NS_Incom,ns_incom_inhom_2d_512-63.h5,https://darus.uni-stuttgart.de/api/access/datafile/133680,2D/NS_incom/,49732762c231a410904cf8bddb0444cd +NS_Incom,ns_incom_inhom_2d_512-64.h5,https://darus.uni-stuttgart.de/api/access/datafile/133628,2D/NS_incom/,c0e54ed459c84470986818d6d89be0c8 +NS_Incom,ns_incom_inhom_2d_512-65.h5,https://darus.uni-stuttgart.de/api/access/datafile/133635,2D/NS_incom/,69f1b9013bb76fba1006272951e12845 +NS_Incom,ns_incom_inhom_2d_512-66.h5,https://darus.uni-stuttgart.de/api/access/datafile/133629,2D/NS_incom/,c79245b195ec850a40c0957056bffad7 +NS_Incom,ns_incom_inhom_2d_512-67.h5,https://darus.uni-stuttgart.de/api/access/datafile/133277,2D/NS_incom/,f8465b71d9afe577ba30cbe5d0c490e6 +NS_Incom,ns_incom_inhom_2d_512-68.h5,https://darus.uni-stuttgart.de/api/access/datafile/133681,2D/NS_incom/,ae1e262d7f1a56369b13b6d5de80e03d +NS_Incom,ns_incom_inhom_2d_512-69.h5,https://darus.uni-stuttgart.de/api/access/datafile/133631,2D/NS_incom/,b5159c70d9b14e480bb520961bafcae4 +NS_Incom,ns_incom_inhom_2d_512-7.h5,https://darus.uni-stuttgart.de/api/access/datafile/136473,2D/NS_incom/,e1b06fd07d09227a0767baa260f361e2 +NS_Incom,ns_incom_inhom_2d_512-70.h5,https://darus.uni-stuttgart.de/api/access/datafile/133632,2D/NS_incom/,3aee2aa9b1a14c7d04ce2770e95a354a +NS_Incom,ns_incom_inhom_2d_512-71.h5,https://darus.uni-stuttgart.de/api/access/datafile/133634,2D/NS_incom/,2f5b2bf56e72d61c41d4b056ac540fb0 +NS_Incom,ns_incom_inhom_2d_512-72.h5,https://darus.uni-stuttgart.de/api/access/datafile/133682,2D/NS_incom/,b99c5b112d4a7dfca053cb2c78ebe581 +NS_Incom,ns_incom_inhom_2d_512-73.h5,https://darus.uni-stuttgart.de/api/access/datafile/133273,2D/NS_incom/,9d1c5e98ee4fe97c7498d2f574c165ae +NS_Incom,ns_incom_inhom_2d_512-74.h5,https://darus.uni-stuttgart.de/api/access/datafile/133633,2D/NS_incom/,475427b7337f1fd7114b7e6c32a2f709 +NS_Incom,ns_incom_inhom_2d_512-75.h5,https://darus.uni-stuttgart.de/api/access/datafile/133694,2D/NS_incom/,170b32d051ed72716aab0f40b13f359e +NS_Incom,ns_incom_inhom_2d_512-76.h5,https://darus.uni-stuttgart.de/api/access/datafile/133636,2D/NS_incom/,2666826a3944f5999a1053242e9588a9 +NS_Incom,ns_incom_inhom_2d_512-77.h5,https://darus.uni-stuttgart.de/api/access/datafile/133638,2D/NS_incom/,c58cc6dd09fb89c1b5fb95081f1220c9 +NS_Incom,ns_incom_inhom_2d_512-78.h5,https://darus.uni-stuttgart.de/api/access/datafile/133637,2D/NS_incom/,b183a32a3e2230968d5b7f20d9dbf818 +NS_Incom,ns_incom_inhom_2d_512-79.h5,https://darus.uni-stuttgart.de/api/access/datafile/133651,2D/NS_incom/,def3145be356b7d020ea09e5232bb60f +NS_Incom,ns_incom_inhom_2d_512-8.h5,https://darus.uni-stuttgart.de/api/access/datafile/136449,2D/NS_incom/,2e6a62550680a255d4dfec24a761d5cf +NS_Incom,ns_incom_inhom_2d_512-80.h5,https://darus.uni-stuttgart.de/api/access/datafile/133639,2D/NS_incom/,07a1edc16fb7d981155ad4173c42abf1 +NS_Incom,ns_incom_inhom_2d_512-81.h5,https://darus.uni-stuttgart.de/api/access/datafile/133276,2D/NS_incom/,3a58395f316b2158ea61844eadd4109f +NS_Incom,ns_incom_inhom_2d_512-82.h5,https://darus.uni-stuttgart.de/api/access/datafile/133640,2D/NS_incom/,a0f4505418beb4da8b7f84a83efc3904 +NS_Incom,ns_incom_inhom_2d_512-83.h5,https://darus.uni-stuttgart.de/api/access/datafile/133272,2D/NS_incom/,89993b4f3ac2653f6377941e640e2233 +NS_Incom,ns_incom_inhom_2d_512-84.h5,https://darus.uni-stuttgart.de/api/access/datafile/133643,2D/NS_incom/,c5269398ec4edb0f5891d4efdb2fc89d +NS_Incom,ns_incom_inhom_2d_512-85.h5,https://darus.uni-stuttgart.de/api/access/datafile/133641,2D/NS_incom/,090862c29ecbc2c4232ddbacca3fb230 +NS_Incom,ns_incom_inhom_2d_512-86.h5,https://darus.uni-stuttgart.de/api/access/datafile/133647,2D/NS_incom/,a712449bef87318a620d1663c4ffde27 +NS_Incom,ns_incom_inhom_2d_512-87.h5,https://darus.uni-stuttgart.de/api/access/datafile/133644,2D/NS_incom/,920a3dcb3b216f0d9a4ab0b9b4c8745a +NS_Incom,ns_incom_inhom_2d_512-88.h5,https://darus.uni-stuttgart.de/api/access/datafile/133655,2D/NS_incom/,cc96b1e357f264f12dcfbb41736a53bf +NS_Incom,ns_incom_inhom_2d_512-89.h5,https://darus.uni-stuttgart.de/api/access/datafile/133646,2D/NS_incom/,7fd484851410b7139bfea4166074fe00 +NS_Incom,ns_incom_inhom_2d_512-9.h5,https://darus.uni-stuttgart.de/api/access/datafile/136475,2D/NS_incom/,254bb44710b625f716b61d51d5d1e58e +NS_Incom,ns_incom_inhom_2d_512-90.h5,https://darus.uni-stuttgart.de/api/access/datafile/133652,2D/NS_incom/,c6c4e27d6ba19b4ccb5f27e67dd9cdf1 +NS_Incom,ns_incom_inhom_2d_512-91.h5,https://darus.uni-stuttgart.de/api/access/datafile/136474,2D/NS_incom/,527b23b7e4ca4d18c5c036c41e15e4fe +NS_Incom,ns_incom_inhom_2d_512-92.h5,https://darus.uni-stuttgart.de/api/access/datafile/133648,2D/NS_incom/,84fd5f64db474f8e20b670d01929d57a +NS_Incom,ns_incom_inhom_2d_512-93.h5,https://darus.uni-stuttgart.de/api/access/datafile/133649,2D/NS_incom/,20583d7edf7da9a80f16a26c8388a65e +NS_Incom,ns_incom_inhom_2d_512-94.h5,https://darus.uni-stuttgart.de/api/access/datafile/133691,2D/NS_incom/,8540e7f1cc6a9e2bc64726a325ddacfb +NS_Incom,ns_incom_inhom_2d_512-95.h5,https://darus.uni-stuttgart.de/api/access/datafile/133690,2D/NS_incom/,25d4c08d5534d0f10c877220e258640c +NS_Incom,ns_incom_inhom_2d_512-96.h5,https://darus.uni-stuttgart.de/api/access/datafile/133270,2D/NS_incom/,93e772cbfd36d434e97a079d701d5382 +NS_Incom,ns_incom_inhom_2d_512-97.h5,https://darus.uni-stuttgart.de/api/access/datafile/133653,2D/NS_incom/,ffa1d46e58f67aa1932fc4ffd2e7716a +NS_Incom,ns_incom_inhom_2d_512-98.h5,https://darus.uni-stuttgart.de/api/access/datafile/133687,2D/NS_incom/,c28112e1339cda43f67dce03536e3455 +NS_Incom,ns_incom_inhom_2d_512-99.h5,https://darus.uni-stuttgart.de/api/access/datafile/133654,2D/NS_incom/,abfe923a224cb4515dd545d5858b2124 +SWE,2D_rdb_NA_NA.h5,https://darus.uni-stuttgart.de/api/access/datafile/133021,2D/shallow-water/,75d838c47aa410694bdc912ea7f22282 +3D_CFD,3D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164693,3D/Train/,45892a12d1066d54af74badae55c438e +3D_CFD,3D_CFD_Turb_M1.0_Eta1e-08_Zeta1e-08_periodic_Train.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/164694,3D/Train/,4f0eacaec5ff000bbd9a31026ea207b8 +3D_CFD,BlastWave.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133264,3D/Test/BlastWave/,f901a1ec75925c1d861ac99a825878f3 +3D_CFD,Turb_M01.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133224,3D/Test/Turbulence/,b27495031f434d01762bb0b819ac71e2 +3D_CFD,Turb_M05.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133225,3D/Test/Turbulence/,f9407dff1a75a1d14d93e7dd570af728 +3D_CFD,Turb_M1.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/135833,3D/Test/Turbulence/,3758f23f71684ac666e0b1e91da0a1c4 +3D_CFD,Turb_M2.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133227,3D/Test/Turbulence/,12e528dc8ab800f69474600ec58b24d3 +3D_CFD,Turb_M4.hdf5,https://darus.uni-stuttgart.de/api/access/datafile/133228,3D/Test/Turbulence/,8db384feba75903a8c5b21ebeba40083 diff --git a/pdebench/data_download/visualize_pdes.py b/pdebench/data_download/visualize_pdes.py new file mode 100644 index 0000000..1c75218 --- /dev/null +++ b/pdebench/data_download/visualize_pdes.py @@ -0,0 +1,547 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import h5py +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import animation +from tqdm import tqdm + +pdes = ( + "advection", + "burgers", + "1d_cfd", + "diff_sorp", + "1d_reacdiff", + "2d_cfd", + "darcy", + "2d_reacdiff", + "ns_incom", + "swe", + "3d_cfd", +) + + +def visualize_diff_sorp(path, seed=None): + """ + This function animates the diffusion sorption data. + + Args: + path : path to the desired file + seed : seed to select a specific sample/batch, ranging from 0-9999 + """ + + # Read the h5 file and store the data + h5_file = h5py.File(Path(path) / "1D_diff-sorp_NA_NA.h5", "r") + num_samples = len(h5_file.keys()) + + # randomly choose a seed for picking a sample that will subsequently be visualized + rng = np.random.default_rng() + if not seed: + seed = rng.integers(low=0, high=num_samples, size=1).item() + + # Ensure the seed number is defined + assert seed < num_samples, "Seed number too high!" + + seed = str(seed).zfill(4) + data = np.array(h5_file[f"{seed}/data"], dtype="f") # dim = [101, 1024, 1] + + h5_file.close() + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + for i in tqdm(range(data.shape[0])): + if i == 0: + im = ax.plot( + data[0].squeeze(), animated=True, color="blue" + ) # show an initial one first + else: + im = ax.plot(data[i].squeeze(), animated=True, color="blue") + ims.append([im[0]]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie_diff_sorp.gif", writer=writer) + + +def visualize_2d_reacdiff(path, seed=None): + """ + This function animates the 2D reaction-diffusion data. + + Args: + path : path to the desired file + seed : seed to select a specific sample/batch, ranging from 0-9999 + """ + + # Read the h5 file and store the data + h5_file = h5py.File(Path(path) / "2D_diff-react_NA_NA.h5", "r") + num_samples = len(h5_file.keys()) + + # randomly choose a seed for picking a sample that will subsequently be visualized + rng = np.random.default_rng() + if not seed: + seed = rng.integers(low=0, high=num_samples, size=1).item() + + # Ensure the seed number is defined + assert seed < num_samples, "Seed number too high!" + + seed = str(seed).zfill(4) + # dim = [101, 128, 128, 2] + data = np.array(h5_file[f"{seed}/data"], dtype="f") + + h5_file.close() + + # Initialize plot + fig, ax = plt.subplots(1, 2) + + # Store the plot handle at each time step in the 'ims' list + ims = [] + for i in tqdm(range(data.shape[0])): + im1 = ax[0].imshow(data[i, ..., 0].squeeze(), animated=True) + im2 = ax[1].imshow(data[i, ..., 1].squeeze(), animated=True) + if i == 0: + # show an initial one first + ax[0].imshow(data[0, ..., 0].squeeze()) + # show an initial one first + ax[1].imshow(data[0, ..., 1].squeeze()) + ims.append([im1, im2]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie_2d_reacdiff.gif", writer=writer) + + +def visualize_swe(path, seed=None): + """ + This function animates the shallow water equation data. + + Args: + path : path to the desired file + seed : seed to select a specific sample/batch, ranging from 0-9999 + """ + + # Read the h5 file and store the data + h5_file = h5py.File(Path(path) / "2D_rdb_NA_NA.h5", "r") + num_samples = len(h5_file.keys()) + + # randomly choose a seed for picking a sample that will subsequently be visualized + rng = np.random.default_rng() + if not seed: + seed = rng.integers(low=0, high=num_samples, size=1).item() + + # Ensure the seed number is defined + assert seed < num_samples, "Seed number too high!" + + seed = str(seed).zfill(4) + # dim = [101, 128, 128, 1] + data = np.array(h5_file[f"{seed}/data"], dtype="f") + + h5_file.close() + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + for i in tqdm(range(data.shape[0])): + im = ax.imshow(data[i].squeeze(), animated=True) + if i == 0: + ax.imshow(data[0].squeeze()) # show an initial one first + ims.append([im]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie_swe.gif", writer=writer) + + +def visualize_burgers(path, param=None): + """ + This function animates the Burgers equation + + Args: + path : path to the desired file + param: PDE parameter of the data shard to be visualized + """ + + # Read the h5 file and store the data + if param is not None: + flnm = "1D_Burgers_Sols_Nu" + str(param) + ".hdf5" + assert Path(path + flnm).is_file(), "no such file! " + path + flnm + else: + flnm = "1D_Burgers_Sols_Nu0.01.hdf5" + + nb = 0 + with h5py.File(Path(path) / flnm, "r") as h5_file: + xcrd = np.array(h5_file["x-coordinate"], dtype=np.float32) + data = np.array(h5_file["tensor"], dtype=np.float32)[ + nb + ] # (batch, t, x, channel) --> (t, x, channel) + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + + for i in tqdm(range(data.shape[0])): + if i == 0: + im = ax.plot(xcrd, data[i].squeeze(), animated=True, color="blue") + else: + im = ax.plot( + xcrd, data[i].squeeze(), animated=True, color="blue" + ) # show an initial one first + ims.append([im[0]]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie_burgers.gif", writer=writer) + + +def visualize_advection(path, param=None): + """ + This function animates the Advection equation + + Args: + path : path to the desired file + param: PDE parameter of the data shard to be visualized + """ + + # Read the h5 file and store the data + if param is not None: + flnm = "1D_Advection_Sols_beta" + str(param) + ".hdf5" + assert Path(path + flnm).is_file(), "no such file! " + path + flnm + else: + flnm = "1D_Advection_Sols_beta0.4.hdf5" + + nb = 0 + with h5py.File(Path(path) / flnm, "r") as h5_file: + xcrd = np.array(h5_file["x-coordinate"], dtype=np.float32) + data = np.array(h5_file["tensor"], dtype=np.float32)[ + nb + ] # (batch, t, x, channel) --> (t, x, channel) + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + for i in tqdm(range(data.shape[0])): + im = ax.plot(xcrd, data[i].squeeze(), animated=True) + if i == 0: + ax.plot(xcrd, data[i].squeeze()) # show an initial one first + ims.append([im[0]]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie_advection.gif", writer=writer) + + +def visualize_1d_cfd(path, param=None): + """ + This function animates the Burgers equation + + Args: + path : path to the desired file + param: PDE parameter of the data shard to be visualized + """ + + # Read the h5 file and store the data + if param is not None: + assert len(param) == 4, "param should include type,eta,zeta,boundary as list" + flnm = ( + "1D_CFD_" + + str(param[0]) + + "_Eta" + + str(param[1]) + + "_Zeta" + + str(param[2]) + + "_" + + str(param[3]) + + "_Train.hdf5" + ) + assert Path(path + flnm).is_file(), "no such file! " + path + flnm + else: + flnm = "1D_CFD_Rand_Eta1.e-8_Zeta1.e-8_periodic_Train.hdf5" + + nb = 0 + with h5py.File(Path(path) / flnm, "r") as h5_file: + xcrd = np.array(h5_file["x-coordinate"], dtype=np.float32) + dd = np.array(h5_file["density"], dtype=np.float32)[ + nb + ] # (batch, t, x, channel) --> (t, x, channel) + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + ax.set_title("density") + for i in tqdm(range(dd.shape[0])): + im = ax.plot(xcrd, dd[i].squeeze(), animated=True) + if i == 0: + ax.plot(xcrd, dd[i].squeeze()) # show an initial one first + ims.append([im[0]]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie_1d_cfd.gif", writer=writer) + + +def visualize_2d_cfd(path, param=None): + """ + This function animates the Burgers equation + + Args: + path : path to the desired file + param: PDE parameter of the data shard to be visualized + """ + + # Read the h5 file and store the data + if param is not None: + assert ( + len(param) == 6 + ), "param should include type,M,eta,zeta,boundary, resolution as list" + flnm = ( + "2D_CFD_" + + str(param[0]) + + "_M" + + str(param[1]) + + "_Eta" + + str(param[2]) + + "_Zeta" + + str(param[3]) + + "_" + + str(param[4]) + + "_" + + str(param[5]) + + "_Train.hdf5" + ) + assert Path(path + flnm).is_file(), "no such file! " + path + flnm + else: + flnm = "2D_CFD_Rand_M0.1_Eta1e-8_Zeta1e-8_periodic_512_Train.hdf5" + + nb = 0 + with h5py.File(Path(path) / flnm, "r") as h5_file: + dd = np.array(h5_file["density"], dtype=np.float32)[ + nb + ] # (batch, t, x, y, channel) --> (t, x, y, channel) + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + ax.set_title("density") + for i in range(dd.shape[0]): + im = ax.imshow(dd[i].squeeze(), animated=True) + ims.append([im]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie.gif", writer=writer) + + +def visualize_3d_cfd(path, param=None): + # Read the h5 file and store the data + if param is not None: + assert len(param) == 5, "param should include type,M,eta,zeta,boundary as list" + flnm = ( + "3D_CFD_" + + str(param[0]) + + "_M" + + str(param[1]) + + "_Eta" + + str(param[2]) + + "_Zeta" + + str(param[3]) + + "_" + + str(param[4]) + + "_Train.hdf5" + ) + assert Path(path + flnm).is_file(), "no such file! " + path + flnm + else: + flnm = "3D_CFD_Rand_M1.0_Eta1e-8_Zeta1e-8_periodic_Train.hdf5" + + nb = 0 + with h5py.File(Path(path) / flnm, "r") as h5_file: + dd = np.array(h5_file["density"], dtype=np.float32)[ + nb + ] # (batch, t, x, y, channel) --> (t, x, y, channel) + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + ax.set_title("density") + for i in range(dd.shape[0]): + im = ax.imshow(dd[i, :, :, 32].squeeze(), animated=True) + ims.append([im]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie.gif", writer=writer) + + +def visualize_ns_incom() -> None: + pass + + +def visualize_darcy(path, param=None): + """ + This function animates Darcy Flow equation + + Args: + path : path to the desired file + param: PDE parameter of the data shard to be visualized + """ + + # Read the h5 file and store the data + if param is not None: + flnm = "2D_DarcyFlow_beta" + str(param) + "_Train.hdf5" + assert Path(path + flnm).is_file(), "no such file! " + path + flnm + else: + flnm = "2D_DarcyFlow_beta1.0_Train.hdf5" + + nb = 0 + with h5py.File(Path(path) / flnm, "r") as h5_file: + data = np.array(h5_file["tensor"], dtype=np.float32)[ + nb + ] # (batch, t, x, y, channel) --> (t, x, y, channel) + nu = np.array(h5_file["nu"], dtype=np.float32)[ + nb + ] # (batch, t, x, y, channel) --> (t, x, y, channel) + + # Initialize plot + fig, ax = plt.subplots(1, 2, figsize=(16, 8)) + + ax[0].imshow(data.squeeze()) + ax[1].imshow(nu.squeeze()) + ax[0].set_title("Data u") + ax[1].set_title("diffusion coefficient nu") + plt.savefig("2D_DarcyFlow.pdf") + + +def visualize_1d_reacdiff(path, param=None): + """ + This function animates 1D Reaction Diffusion equation + + Args: + path : path to the desired file + param: PDE parameter of the data shard to be visualized + """ + + # Read the h5 file and store the data + if param is not None: + assert len(param) == 2, "param should include Nu and Rho as list" + flnm = "ReacDiff_Nu" + str(param[0]) + "_Rho" + str(param[1]) + ".hdf5" + assert Path(path + flnm).is_file(), "no such file! " + path + flnm + else: + flnm = "ReacDiff_Nu1.0_Rho1.0.hdf5" + + nb = 0 + with h5py.File(Path(path) / flnm, "r") as h5_file: + xcrd = np.array(h5_file["x-coordinate"], dtype=np.float32) + data = np.array(h5_file["tensor"], dtype=np.float32)[ + nb + ] # (batch, t, x, channel) --> (t, x, channel) + + # Initialize plot + fig, ax = plt.subplots() + + # Store the plot handle at each time step in the 'ims' list + ims = [] + for i in tqdm(range(data.shape[0])): + im = ax.plot(xcrd, data[i].squeeze(), animated=True) + if i == 0: + ax.plot(xcrd, data[i].squeeze()) # show an initial one first + ims.append([im[0]]) + + # Animate the plot + ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000) + + writer = animation.PillowWriter(fps=15, bitrate=1800) + ani.save("movie_1d_reacdiff.gif", writer=writer) + + +if __name__ == "__main__": + arg_parser = argparse.ArgumentParser( + prog="Visualize PDEs", + description="Helper script to visualize PDEs in the PDEBench dataset", + epilog="", + ) + + arg_parser.add_argument( + "--data_path", + type=str, + default="./", + help="Path to the hdf5 data where the downloaded data reside", + ) + arg_parser.add_argument( + "--pde_name", type=str, help="Name of the PDE dataset to download" + ) + arg_parser.add_argument( + "--seed_number", + type=int, + default=None, + help="Seed number to define which sample/batch to be plotted", + ) + arg_parser.add_argument( + "--param", + type=float, + default=None, + help="PDE parameter to be plotted", + ) + arg_parser.add_argument( + "--params", + nargs="+", + default=None, + help="PDE parameters to be plotted", + ) + + args = arg_parser.parse_args() + + if args.pde_name == "diff_sorp": + visualize_diff_sorp(args.data_path, args.seed_number) + elif args.pde_name == "2d_reacdiff": + visualize_2d_reacdiff(args.data_path, args.seed_number) + elif args.pde_name == "swe": + visualize_swe(args.data_path, args.seed_number) + elif args.pde_name == "burgers": + visualize_burgers(args.data_path, args.param) + elif args.pde_name == "advection": + visualize_advection(args.data_path, args.param) + elif args.pde_name == "1d_cfd": + visualize_1d_cfd(args.data_path, args.params) + elif args.pde_name == "2d_cfd": + visualize_2d_cfd(args.data_path, args.params) + elif args.pde_name == "3d_cfd": + visualize_3d_cfd(args.data_path, args.params) + elif args.pde_name == "darcy": + visualize_darcy(args.data_path, args.param) + elif args.pde_name == "1d_reacdiff": + visualize_1d_reacdiff(args.data_path, args.params) + else: + errmsg = "PDE name not recognized!" + raise ValueError(errmsg) diff --git a/pdebench/data_gen/configs/diff-react.yaml b/pdebench/data_gen/configs/diff-react.yaml index 39372cc..92ab753 100644 --- a/pdebench/data_gen/configs/diff-react.yaml +++ b/pdebench/data_gen/configs/diff-react.yaml @@ -34,13 +34,12 @@ sim: y_top: 1.0 ydim: 128 n: 1 - seed: '???' - -plot: - t_idx: 1.0 # Fraction of the final time step idx to be plotted - dim: 2 # Spatial dimension - channel_idx: 0 # Index of the variable to be plotted + seed: "???" +plot: + t_idx: 1.0 # Fraction of the final time step idx to be plotted + dim: 2 # Spatial dimension + channel_idx: 0 # Index of the variable to be plotted dataverse: lib_name: pyDaRUS @@ -53,8 +52,9 @@ dataverse: identifier_scheme: ORCID identifier: 0000-0003-3619-9122 description: - - text: 2D diffusion-reaction dataset generated for the PDE benchmark paper - date: '2022' + - text: + 2D diffusion-reaction dataset generated for the PDE benchmark paper + date: "2022" contact: - name: Timothy Praditia affiliation: Universitรคt Stuttgart @@ -66,10 +66,18 @@ dataverse: process: processing_methods: - name: FVM - description: Finite Volume Method is a spatial discretization method to calculate spatial derivative in a Partial Differential Equation. It integrates the fluxes at all discrete cell boundaries so that it ensures conservation. + description: + Finite Volume Method is a spatial discretization method to calculate + spatial derivative in a Partial Differential Equation. It integrates + the fluxes at all discrete cell boundaries so that it ensures + conservation. parameters: cell length, cell width - name: RK45 - description: Explicit Runge-Kutta method of order 5(4) is a time integration method to solve the temporal derivative in a Partial Differential Equation. It is an adaptive time integration scheme to ensure better accuracy and computation efficiency. + description: + Explicit Runge-Kutta method of order 5(4) is a time integration + method to solve the temporal derivative in a Partial Differential + Equation. It is an adaptive time integration scheme to ensure better + accuracy and computation efficiency. parameters: time step size, total time, error tolerance method_parameters: - name: cell length diff --git a/pdebench/data_gen/configs/diff-sorp.yaml b/pdebench/data_gen/configs/diff-sorp.yaml index e45a24e..2ad7e1f 100644 --- a/pdebench/data_gen/configs/diff-sorp.yaml +++ b/pdebench/data_gen/configs/diff-sorp.yaml @@ -34,13 +34,12 @@ sim: x_right: 1.0 xdim: 1024 n: 1 - seed: '???' - + seed: "???" + plot: - t_idx: 1.0 # Fraction of the final time step idx to be plotted - dim: 1 # Spatial dimension - channel_idx: 0 # Index of the variable to be plotted - + t_idx: 1.0 # Fraction of the final time step idx to be plotted + dim: 1 # Spatial dimension + channel_idx: 0 # Index of the variable to be plotted dataverse: lib_name: pyDaRUS @@ -53,8 +52,9 @@ dataverse: identifier_scheme: ORCID identifier: 0000-0003-3619-9122 description: - - text: 1D diffusion-sorption dataset generated for the PDE benchmark paper - date: '2022' + - text: + 1D diffusion-sorption dataset generated for the PDE benchmark paper + date: "2022" contact: - name: Timothy Praditia affiliation: Universitรคt Stuttgart @@ -66,10 +66,18 @@ dataverse: process: processing_methods: - name: FVM - description: Finite Volume Method is a spatial discretization method to calculate spatial derivative in a Partial Differential Equation. It integrates the fluxes at all discrete cell boundaries so that it ensures conservation. + description: + Finite Volume Method is a spatial discretization method to calculate + spatial derivative in a Partial Differential Equation. It integrates + the fluxes at all discrete cell boundaries so that it ensures + conservation. parameters: cell length - name: RK45 - description: Explicit Runge-Kutta method of order 5(4) is a time integration method to solve the temporal derivative in a Partial Differential Equation. It is an adaptive time integration scheme to ensure better accuracy and computation efficiency. + description: + Explicit Runge-Kutta method of order 5(4) is a time integration + method to solve the temporal derivative in a Partial Differential + Equation. It is an adaptive time integration scheme to ensure better + accuracy and computation efficiency. parameters: time step size, total time, error tolerance method_parameters: - name: cell length diff --git a/pdebench/data_gen/configs/mode/debug.yaml b/pdebench/data_gen/configs/mode/debug.yaml index e17c0d7..2225b85 100644 --- a/pdebench/data_gen/configs/mode/debug.yaml +++ b/pdebench/data_gen/configs/mode/debug.yaml @@ -14,4 +14,3 @@ hydra: subdir: ${hydra.job.num} launcher: n_jobs: 1 - diff --git a/pdebench/data_gen/configs/mode/default.yaml b/pdebench/data_gen/configs/mode/default.yaml index ecb1b8f..bd34f5a 100644 --- a/pdebench/data_gen/configs/mode/default.yaml +++ b/pdebench/data_gen/configs/mode/default.yaml @@ -8,4 +8,3 @@ hydra: sweep: dir: ${oc.env:WORKING_DIR,multirun}/${hydra.job.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname} subdir: ${hydra.job.num} - diff --git a/pdebench/data_gen/configs/mode/slurm.yaml b/pdebench/data_gen/configs/mode/slurm.yaml index c3c619e..eaff80a 100644 --- a/pdebench/data_gen/configs/mode/slurm.yaml +++ b/pdebench/data_gen/configs/mode/slurm.yaml @@ -15,4 +15,3 @@ hydra: tasks_per_node: 1 mem_gb: 16 timeout_min: 719 # just under 12 hours is good for many clusters - diff --git a/pdebench/data_gen/configs/ns_incomp.yaml b/pdebench/data_gen/configs/ns_incomp.yaml index ec76433..086b771 100644 --- a/pdebench/data_gen/configs/ns_incomp.yaml +++ b/pdebench/data_gen/configs/ns_incomp.yaml @@ -7,25 +7,25 @@ artefact_dir: ${oc.env:ARTEFACT_DIR,artefacts} dataverse: dataset_id: doi:10.18419/darus-2984 -sim_name: 'ns_sim_2d' +sim_name: "ns_sim_2d" label: null -# Solver Parameters +# Solver Parameters domain_size: [1, 1] -grid_size: [256,256] +grid_size: [256, 256] #['scalar_grid', extrapolation_x:(type or bound), extrapolation_y:(type or bound)] -particle_extrapolation: 'BOUNDARY' - +particle_extrapolation: "BOUNDARY" + #['staggered_grid', extrapolation_x:(type or bound), extrapolation_y:(type or bound)] -velocity_extrapolation: 'ZERO' +velocity_extrapolation: "ZERO" # Fluid characteristics NU: 0.01 #(kinematic viscosity) # External force # enable_gravity: false -force_extrapolation: 'ZERO' +force_extrapolation: "ZERO" # Fluctuation Generator Parameters (Noise) seed: 1 @@ -36,7 +36,7 @@ force_scale: 0.15 #params for IncompressibleFlow(Physics) n_steps: 100000 -DT : 0.00005 +DT: 0.00005 frame_int: 100 n_batch: 1 @@ -45,10 +45,9 @@ n_batch: 1 # save_images: false # save_gif: false save_h5: true -profile: false # Run performance profiling -upload: false # upload to DARUS - requires key +profile: false # Run performance profiling +upload: false # upload to DARUS - requires key -backend: 'jax' -device: 'GPU' +backend: "jax" +device: "GPU" jit: true - diff --git a/pdebench/data_gen/configs/radial_dam_break.yaml b/pdebench/data_gen/configs/radial_dam_break.yaml index 3071940..fb5e13c 100644 --- a/pdebench/data_gen/configs/radial_dam_break.yaml +++ b/pdebench/data_gen/configs/radial_dam_break.yaml @@ -33,13 +33,12 @@ sim: x_right: 2.5 y_bottom: -2.5 y_top: 2.5 - seed: '???' + seed: "???" - plot: - t_idx: 1.0 # Fraction of the final time step idx to be plotted - dim: 2 # Spatial dimension - channel_idx: 0 # Index of the variable to be plotted + t_idx: 1.0 # Fraction of the final time step idx to be plotted + dim: 2 # Spatial dimension + channel_idx: 0 # Index of the variable to be plotted dataverse: lib_name: pyDaRUS @@ -52,8 +51,10 @@ dataverse: identifier_scheme: ORCID identifier: 0000-0001-8070-2384 description: - - text: 2D shallow-water equation dataset generated for the PDE benchmark paper - date: '2022' + - text: + 2D shallow-water equation dataset generated for the PDE benchmark + paper + date: "2022" contact: - name: Raphael Leiteritz affiliation: Universitรคt Stuttgart @@ -65,7 +66,11 @@ dataverse: process: processing_methods: - name: FVM - description: Finite Volume Method is a spatial discretization method to calculate spatial derivative in a Partial Differential Equation. It integrates the fluxes at all discrete cell boundaries so that it ensures conservation. + description: + Finite Volume Method is a spatial discretization method to calculate + spatial derivative in a Partial Differential Equation. It integrates + the fluxes at all discrete cell boundaries so that it ensures + conservation. parameters: cell length, cell width method_parameters: - name: cell length @@ -85,4 +90,4 @@ dataverse: unit: s value: 1 engMeta: {} - codeMeta: {} \ No newline at end of file + codeMeta: {} diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py new file mode 100644 index 0000000..cdd336b --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py @@ -0,0 +1,215 @@ +""" + + + File: advection_exact_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import logging +import time +from math import ceil + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put + +# Hydra +from omegaconf import DictConfig + +logger = logging.getLogger(__name__) + + +# Init arguments with Hydra +@hydra.main(config_path="config") +def main(cfg: DictConfig) -> None: + logger.info("advection velocity: %f", cfg.args.beta) + + # cell edge coordinate + xe = jnp.linspace(cfg.args.xL, cfg.args.xR, cfg.args.nx + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * (xe[1] - xe[0]) + # t-coordinate + it_tot = ceil((cfg.args.fin_time - cfg.args.ini_time) / cfg.args.dt_save) + 1 + tc = jnp.arange(it_tot + 1) * cfg.args.dt_save + + def evolve(u): + t = cfg.args.ini_time + i_save = 0 + tm_ini = time.time() + + it_tot = ceil((cfg.args.fin_time - cfg.args.ini_time) / cfg.args.dt_save) + 1 + uu = jnp.zeros([it_tot, u.shape[0]]) + uu = uu.at[0].set(u) + + while t < cfg.args.fin_time: + logger.info("save data at t = %f", t) + u = set_function(xc, t, cfg.args.beta) + uu = uu.at[i_save].set(u) + t += cfg.args.dt_save + i_save += 1 + + tm_fin = time.time() + logger.info("total elapsed time is %f sec", tm_fin - tm_ini) + uu = uu.at[-1].set(u) + return uu, t + + @jax.jit + def set_function(x, t, beta): + return jnp.sin(2.0 * jnp.pi * (x - beta * t)) + + u = set_function(xc, t=0, beta=cfg.args.beta) + u = device_put(u) # putting variables in GPU (not necessary??) + uu, t = evolve(u) + logger.info("final time is: %f", t) + + logger.info("data saving...") + cwd = hydra.utils.get_original_cwd() + "/" + jnp.save(cwd + cfg.args.save + "/Advection_beta" + str(cfg.args.beta), uu) + jnp.save(cwd + cfg.args.save + "/x_coordinate", xe) + jnp.save(cwd + cfg.args.save + "/t_coordinate", tc) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py new file mode 100644 index 0000000..d736f43 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py @@ -0,0 +1,301 @@ +""" + + + File: advection_multi_solution_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import random +import sys + +# Hydra +from math import ceil, exp, log +from pathlib import Path + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put, lax +from omegaconf import DictConfig + +sys.path.append("..") +import logging + +from utils import Courant, bc, init_multi, limiting + +logger = logging.getLogger(__name__) + + +def _pass(carry): + return carry + + +# Init arguments with Hydra +@hydra.main(config_path="config") +def main(cfg: DictConfig) -> None: + # basic parameters + dx = (cfg.multi.xR - cfg.multi.xL) / cfg.multi.nx + dx_inv = 1.0 / dx + + # cell edge coordinate + xe = jnp.linspace(cfg.multi.xL, cfg.multi.xR, cfg.multi.nx + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + # t-coordinate + it_tot = ceil((cfg.multi.fin_time - cfg.multi.ini_time) / cfg.multi.dt_save) + 1 + tc = jnp.arange(it_tot + 1) * cfg.multi.dt_save + + show_steps = cfg.multi.show_steps + ini_time = cfg.multi.ini_time + fin_time = cfg.multi.fin_time + dt_save = cfg.multi.dt_save + CFL = cfg.multi.CFL + if cfg.multi.if_rand_param: + beta = exp( + random.uniform(log(0.01), log(100)) + ) # uniform number between 0.01 to 100 + else: + beta = cfg.multi.beta + + logger.info("beta: %f", beta) + + @jax.jit + def evolve(u): + t = ini_time + tsave = t + steps = 0 + i_save = 0 + dt = 0.0 + uu = jnp.zeros([it_tot, u.shape[0]]) + uu = uu.at[0].set(u) + + def cond_fun(x): + return x[0] < fin_time + + def _body_fun(carry): + def _show(_carry): + u, tsave, i_save, uu = _carry + uu = uu.at[i_save].set(u) + tsave += dt_save + i_save += 1 + return (u, tsave, i_save, uu) + + t, tsave, steps, i_save, dt, u, uu = carry + + carry = (u, tsave, i_save, uu) + u, tsave, i_save, uu = lax.cond(t >= tsave, _show, _pass, carry) + + carry = (u, t, dt, steps, tsave) + u, t, dt, steps, tsave = lax.fori_loop(0, show_steps, simulation_fn, carry) + + return (t, tsave, steps, i_save, dt, u, uu) + + carry = t, tsave, steps, i_save, dt, u, uu + t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) + return uu.at[-1].set(u) + + @jax.jit + def simulation_fn(i, carry): + u, t, dt, steps, tsave = carry + dt = Courant(jnp.array([beta]), dx) * CFL + dt = jnp.min(jnp.array([dt, fin_time - t, tsave - t])) + + def _update(carry): + u, dt = carry + # preditor step for calculating t+dt/2-th time step + u_tmp = update(u, u, dt * 0.5) + # update using flux at t+dt/2-th time step + u = update(u, u_tmp, dt) + return u, dt + + carry = u, dt + u, dt = lax.cond(dt > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return u, t, dt, steps, tsave + + @jax.jit + def update(u, u_tmp, dt): + f = flux(u_tmp) + u -= dt * dx_inv * (f[1 : cfg.multi.nx + 1] - f[0 : cfg.multi.nx]) + return u + + def flux(u): + _u = bc( + u, dx, Ncell=cfg.multi.nx + ) # index 2 for _U is equivalent with index 0 for u + uL, uR = limiting(_u, cfg.multi.nx, if_second_order=cfg.multi.if_second_order) + fL = uL * beta + fR = uR * beta + # upwind advection scheme + return 0.5 * ( + fR[1 : cfg.multi.nx + 2] + + fL[2 : cfg.multi.nx + 3] + - jnp.abs(beta) * (uL[2 : cfg.multi.nx + 3] - uR[1 : cfg.multi.nx + 2]) + ) + + u = init_multi(xc, numbers=cfg.multi.numbers, k_tot=4, init_key=cfg.multi.init_key) + u = device_put(u) # putting variables in GPU (not necessary??) + + # vm_evolve = vmap(evolve, 0, 0) + # uu = vm_evolve(u) + vm_evolve = jax.pmap(jax.vmap(evolve, axis_name="j"), axis_name="i") + local_devices = jax.local_device_count() + + uu = vm_evolve(u.reshape([local_devices, cfg.multi.numbers // local_devices, -1])) + + # reshape before saving + uu = uu.reshape((-1, *uu.shape[2:])) + + logger.info("data saving...") + cwd = hydra.utils.get_original_cwd() + "/" + Path(cwd + cfg.multi.save).mkdir(parents=True, exist_ok=True) + jnp.save(cwd + cfg.multi.save + "1D_Advection_Sols_beta" + str(beta)[:5], uu) + jnp.save(cwd + cfg.multi.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.multi.save + "/t_coordinate", tc) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e-1.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e-1.yaml new file mode 100644 index 0000000..8c484b7 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e-1.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 1.e-1 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e0.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e0.yaml new file mode 100644 index 0000000..bebc582 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e0.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 1.e0 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e1.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e1.yaml new file mode 100644 index 0000000..db14c50 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta1e1.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 1.e1 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta2e-1.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta2e-1.yaml new file mode 100644 index 0000000..aa5539e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta2e-1.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 2.e-1 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta2e0.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta2e0.yaml new file mode 100644 index 0000000..73832a2 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta2e0.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 2.e0 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta4e-1.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta4e-1.yaml new file mode 100644 index 0000000..d8114c8 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta4e-1.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 4.e-1 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta4e0.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta4e0.yaml new file mode 100644 index 0000000..30ab03a --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/beta4e0.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 4.e0 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/config.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/config.yaml new file mode 100644 index 0000000..bebc582 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/args/config.yaml @@ -0,0 +1,10 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 1.e0 +if_show: 1 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta1e-1.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta1e-1.yaml new file mode 100644 index 0000000..0bdcd23 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta1e-1.yaml @@ -0,0 +1,15 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 1.e-1 +if_show: 1 +numbers: 10000 +CFL: 4.e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta1e0.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta1e0.yaml new file mode 100644 index 0000000..5d8f5fd --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta1e0.yaml @@ -0,0 +1,15 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 1.e0 +if_show: 1 +numbers: 10000 +CFL: 4.e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta2e-1.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta2e-1.yaml new file mode 100644 index 0000000..cb46b32 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta2e-1.yaml @@ -0,0 +1,15 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 2.e-1 +if_show: 1 +numbers: 10000 +CFL: 4.e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta2e0.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta2e0.yaml new file mode 100644 index 0000000..3185fc7 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta2e0.yaml @@ -0,0 +1,15 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 2.e0 +if_show: 1 +numbers: 10000 +CFL: 4.e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta4e-1.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta4e-1.yaml new file mode 100644 index 0000000..0125f48 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta4e-1.yaml @@ -0,0 +1,15 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 4.e-1 +if_show: 1 +numbers: 10000 +CFL: 4.e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta4e0.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta4e0.yaml new file mode 100644 index 0000000..d6d8772 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/beta4e0.yaml @@ -0,0 +1,15 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 4.e0 +if_show: 1 +numbers: 10000 +CFL: 4.e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/config.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/config.yaml new file mode 100644 index 0000000..ce8be8b --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/config.yaml @@ -0,0 +1,15 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +beta: 1.e0 +if_show: 1 +numbers: 100 +CFL: 3.e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: 1 diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/config2D.yaml b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/config2D.yaml new file mode 100644 index 0000000..21ce8e9 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/config/multi/config2D.yaml @@ -0,0 +1,19 @@ +save: "../save/advection/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 32 +ny: 32 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +betaX: 1.e0 +betaY: 1.e0 +if_show: 1 +numbers: 4 +CFL: 2.5e-1 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/run_testset.sh b/pdebench/data_gen/data_gen_NLE/AdvectionEq/run_testset.sh new file mode 100644 index 0000000..08e2ebd --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/run_testset.sh @@ -0,0 +1,8 @@ +#! /bin/bash +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta1e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta1e1.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta1e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta2e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta2e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta4e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta4e-1.yaml diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/run_trainset.sh b/pdebench/data_gen/data_gen_NLE/AdvectionEq/run_trainset.sh new file mode 100644 index 0000000..48f5d93 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/run_trainset.sh @@ -0,0 +1,7 @@ +#! /bin/bash +CUDA_VISIBLE_DEVICES='2,3' python3 advection_multi_solution_Hydra.py +multi=beta1e0.yaml +CUDA_VISIBLE_DEVICES='2,3' python3 advection_multi_solution_Hydra.py +multi=beta1e-1.yaml +CUDA_VISIBLE_DEVICES='2,3' python3 advection_multi_solution_Hydra.py +multi=beta2e0.yaml +CUDA_VISIBLE_DEVICES='2,3' python3 advection_multi_solution_Hydra.py +multi=beta2e-1.yaml +CUDA_VISIBLE_DEVICES='2,3' python3 advection_multi_solution_Hydra.py +multi=beta4e0.yaml +CUDA_VISIBLE_DEVICES='2,3' python3 advection_multi_solution_Hydra.py +multi=beta4e-1.yaml diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py new file mode 100644 index 0000000..b3bd376 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py @@ -0,0 +1,327 @@ +""" + + + File: burgers_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import logging +import sys +import time +from math import ceil + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put, lax + +# Hydra +from omegaconf import DictConfig + +sys.path.append("..") +from utils import Courant, Courant_diff, bc, init, limiting + +logger = logging.getLogger(__name__) + + +def _pass(carry): + return carry + + +# Init arguments with Hydra +@hydra.main(config_path="config") +def main(cfg: DictConfig) -> None: + # basic parameters + pi_inv = 1.0 / jnp.pi + dx = (cfg.args.xR - cfg.args.xL) / cfg.args.nx + dx_inv = 1.0 / dx + + # cell edge coordinate + xe = jnp.linspace(cfg.args.xL, cfg.args.xR, cfg.args.nx + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + + show_steps = cfg.args.show_steps + ini_time = cfg.args.ini_time + fin_time = cfg.args.fin_time + dt_save = cfg.args.dt_save + CFL = cfg.args.CFL + + # t-coordinate + it_tot = ceil((fin_time - ini_time) / dt_save) + 1 + tc = jnp.arange(it_tot + 1) * dt_save + + @jax.jit + def evolve(u): + t = ini_time + tsave = t + steps = 0 + i_save = 0 + dt = 0.0 + uu = jnp.zeros([it_tot, u.shape[0]]) + uu = uu.at[0].set(u) + + tm_ini = time.time() + + def cond_fun(x): + return x[0] < fin_time + + def _body_fun(carry): + def _save(_carry): + u, tsave, i_save, uu = _carry + uu = uu.at[i_save].set(u) + tsave += dt_save + i_save += 1 + return (u, tsave, i_save, uu) + + t, tsave, steps, i_save, dt, u, uu = carry + + # if save data + carry = (u, tsave, i_save, uu) + u, tsave, i_save, uu = lax.cond(t >= tsave, _save, _pass, carry) + + carry = (u, t, dt, steps, tsave) + u, t, dt, steps, tsave = lax.fori_loop(0, show_steps, simulation_fn, carry) + + return (t, tsave, steps, i_save, dt, u, uu) + + carry = t, tsave, steps, i_save, dt, u, uu + t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) + uu = uu.at[-1].set(u) + + tm_fin = time.time() + logger.info("total elapsed time is %f sec", tm_fin - tm_ini) + return uu, t + + @jax.jit + def simulation_fn(i, carry): + u, t, dt, steps, tsave = carry + dt_adv = Courant(u, dx) * CFL + dt_dif = Courant_diff(dx, cfg.args.epsilon * pi_inv) * CFL + dt = jnp.min(jnp.array([dt_adv, dt_dif, fin_time - t, tsave - t])) + + def _update(carry): + u, dt = carry + # preditor step for calculating t+dt/2-th time step + u_tmp = update(u, u, dt * 0.5) + # update using flux at t+dt/2-th time step + u = update(u, u_tmp, dt) + return u, dt + + carry = u, dt + u, dt = lax.cond(dt > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return u, t, dt, steps, tsave + + @jax.jit + def update(u, u_tmp, dt): + f = flux(u_tmp) + u -= dt * dx_inv * (f[1 : cfg.args.nx + 1] - f[0 : cfg.args.nx]) + return u + + def flux(u): + _u = bc( + u, dx, Ncell=cfg.args.nx + ) # index 2 for _U is equivalent with index 0 for u + uL, uR = limiting(_u, cfg.args.nx, if_second_order=1.0) + fL = 0.5 * uL**2 + fR = 0.5 * uR**2 + # upwind advection scheme + f_upwd = 0.5 * ( + fR[1 : cfg.args.nx + 2] + + fL[2 : cfg.args.nx + 3] + - 0.5 + * jnp.abs(uL[2 : cfg.args.nx + 3] + uR[1 : cfg.args.nx + 2]) + * (uL[2 : cfg.args.nx + 3] - uR[1 : cfg.args.nx + 2]) + ) + # source term + f_upwd += ( + -cfg.args.epsilon + * pi_inv + * (_u[2 : cfg.args.nx + 3] - _u[1 : cfg.args.nx + 2]) + * dx_inv + ) + return f_upwd + + u = init(xc=xc, mode=cfg.args.init_mode, u0=cfg.args.u0, du=cfg.args.du) + u = device_put(u) # putting variables in GPU (not necessary??) + uu, t = evolve(u) + logger.info("final time is: %.3f", t) + + logger.info("data saving...") + cwd = hydra.utils.get_original_cwd() + "/" + if cfg.args.init_mode == "sinsin": + jnp.save( + cwd + + cfg.args.save + + "/Burgers_" + + cfg.args.init_mode + + "_u" + + str(cfg.args.u0) + + "_du" + + str(cfg.args.du) + + "_Nu" + + str(cfg.args.epsilon), + uu, + ) + else: + jnp.save( + cwd + + cfg.args.save + + "/Burgers_" + + cfg.args.init_mode + + "_u" + + str(cfg.args.u0) + + "_Nu" + + str(cfg.args.epsilon), + uu, + ) + jnp.save(cwd + cfg.args.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.args.save + "/t_coordinate", tc) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py new file mode 100644 index 0000000..8d84253 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py @@ -0,0 +1,315 @@ +""" + + + File: burgers_multi_solution_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import logging +import random +import sys +from math import ceil, exp, log +from pathlib import Path + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put, lax + +# Hydra +from omegaconf import DictConfig + +sys.path.append("..") +from utils import Courant, Courant_diff, bc, init_multi, limiting + +logger = logging.getLogger(__name__) + + +def _pass(carry): + return carry + + +# Init arguments with Hydra +@hydra.main(config_path="config") +def main(cfg: DictConfig) -> None: + # basic parameters + pi_inv = 1.0 / jnp.pi + dx = (cfg.multi.xR - cfg.multi.xL) / cfg.multi.nx + dx_inv = 1.0 / dx + + # cell edge coordinate + xe = jnp.linspace(cfg.multi.xL, cfg.multi.xR, cfg.multi.nx + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + + show_steps = cfg.multi.show_steps + ini_time = cfg.multi.ini_time + fin_time = cfg.multi.fin_time + dt_save = cfg.multi.dt_save + CFL = cfg.multi.CFL + if cfg.multi.if_rand_param: + epsilon = exp( + random.uniform(log(0.001), log(10)) + ) # uniform number between 0.01 to 100 + else: + epsilon = cfg.multi.epsilon + logger.info("epsilon: %f", epsilon) + # t-coordinate + it_tot = ceil((fin_time - ini_time) / dt_save) + 1 + tc = jnp.arange(it_tot + 1) * dt_save + + @jax.jit + def evolve(u): + t = ini_time + tsave = t + steps = 0 + i_save = 0 + dt = 0.0 + uu = jnp.zeros([it_tot, u.shape[0]]) + uu = uu.at[0].set(u) + + def cond_fun(x): + return x[0] < fin_time + + def _body_fun(carry): + def _show(_carry): + u, tsave, i_save, uu = _carry + uu = uu.at[i_save].set(u) + tsave += dt_save + i_save += 1 + return (u, tsave, i_save, uu) + + t, tsave, steps, i_save, dt, u, uu = carry + + carry = (u, tsave, i_save, uu) + u, tsave, i_save, uu = lax.cond(t >= tsave, _show, _pass, carry) + + carry = (u, t, dt, steps, tsave) + u, t, dt, steps, tsave = lax.fori_loop(0, show_steps, simulation_fn, carry) + + return (t, tsave, steps, i_save, dt, u, uu) + + carry = t, tsave, steps, i_save, dt, u, uu + t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) + return uu.at[-1].set(u) + + @jax.jit + def simulation_fn(i, carry): + u, t, dt, steps, tsave = carry + dt_adv = Courant(u, dx) * CFL + dt_dif = Courant_diff(dx, epsilon * pi_inv) * CFL + dt = jnp.min(jnp.array([dt_adv, dt_dif, fin_time - t, tsave - t])) + + def _update(carry): + u, dt = carry + # preditor step for calculating t+dt/2-th time step + u_tmp = update(u, u, dt * 0.5) + # update using flux at t+dt/2-th time step + u = update(u, u_tmp, dt) + return u, dt + + carry = u, dt + u, dt = lax.cond(dt > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return u, t, dt, steps, tsave + + @jax.jit + def update(u, u_tmp, dt): + f = flux(u_tmp) + u -= dt * dx_inv * (f[1 : cfg.multi.nx + 1] - f[0 : cfg.multi.nx]) + return u + + def flux(u): + _u = bc( + u, dx, Ncell=cfg.multi.nx + ) # index 2 for _U is equivalent with index 0 for u + uL, uR = limiting(_u, cfg.multi.nx, if_second_order=1.0) + fL = 0.5 * uL**2 + fR = 0.5 * uR**2 + # upwind advection scheme + f_upwd = 0.5 * ( + fR[1 : cfg.multi.nx + 2] + + fL[2 : cfg.multi.nx + 3] + - 0.5 + * jnp.abs(uL[2 : cfg.multi.nx + 3] + uR[1 : cfg.multi.nx + 2]) + * (uL[2 : cfg.multi.nx + 3] - uR[1 : cfg.multi.nx + 2]) + ) + # source term + f_upwd += ( + -epsilon + * pi_inv + * (_u[2 : cfg.multi.nx + 3] - _u[1 : cfg.multi.nx + 2]) + * dx_inv + ) + return f_upwd + + u = init_multi(xc, numbers=cfg.multi.numbers, k_tot=4, init_key=cfg.multi.init_key) + u = device_put(u) # putting variables in GPU (not necessary??) + + # vm_evolve = vmap(evolve, 0, 0) + # uu = vm_evolve(u) + vm_evolve = jax.pmap(jax.vmap(evolve, axis_name="j"), axis_name="i") + local_devices = jax.local_device_count() + uu = vm_evolve(u.reshape([local_devices, cfg.multi.numbers // local_devices, -1])) + + cwd = hydra.utils.get_original_cwd() + "/" + jnp.save(cwd + cfg.multi.save + "1D_Burgers_Sols_Nu" + str(epsilon)[:5], uu) + jnp.save(cwd + cfg.multi.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.multi.save + "/t_coordinate", tc) + + # reshape before saving + uu = uu.reshape((-1, *uu.shape[2:])) + + logger.info("data saving...") + cwd = hydra.utils.get_original_cwd() + "/" + Path(cwd + cfg.multi.save).mkdir(parents=True, exist_ok=True) + jnp.save(cwd + cfg.multi.save + "1D_Burgers_Sols_Nu" + str(epsilon)[:5], uu) + jnp.save(cwd + cfg.multi.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.multi.save + "/t_coordinate", tc) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/config.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/config.yaml new file mode 100644 index 0000000..bdb4851 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/config.yaml @@ -0,0 +1,17 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" +init_key: 2022 +if_rand_param: None diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-1.yaml new file mode 100644 index 0000000..2e2faaa --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-1.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-1 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "possin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-2.yaml new file mode 100644 index 0000000..a069740 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-2.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "possin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-3.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-3.yaml new file mode 100644 index 0000000..88907d2 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e-3.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-3 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "possin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e0.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e0.yaml new file mode 100644 index 0000000..0563653 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e0.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e0 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "possin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e1.yaml new file mode 100644 index 0000000..214752e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e1.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e1 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "possin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e2.yaml new file mode 100644 index 0000000..6161d32 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/possin_eps1e2.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e2 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "possin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-1.yaml new file mode 100644 index 0000000..1a6d55a --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-1.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-1 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2.yaml new file mode 100644 index 0000000..137e61c --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e-1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e-1.yaml new file mode 100644 index 0000000..67f0974 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e-1.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1.e-1 +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e-2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e-2.yaml new file mode 100644 index 0000000..9078bcd --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e-2.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1.e-2 +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e1.yaml new file mode 100644 index 0000000..fef36a8 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e1.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1.e1 +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e2.yaml new file mode 100644 index 0000000..efe85dd --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-2_u01e2.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1.e2 +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-3.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-3.yaml new file mode 100644 index 0000000..59dbfe7 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e-3.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-3 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e0.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e0.yaml new file mode 100644 index 0000000..4aaea8b --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e0.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e0 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e1.yaml new file mode 100644 index 0000000..4e03682 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e1.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e1 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e2.yaml new file mode 100644 index 0000000..8526ee6 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sin_eps1e2.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e2 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du01.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du01.yaml new file mode 100644 index 0000000..1bb49e3 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du01.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 0.1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sinsin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du025.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du025.yaml new file mode 100644 index 0000000..051f70a --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du025.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 0.25 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sinsin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du05.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du05.yaml new file mode 100644 index 0000000..fa55f69 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du05.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 0.5 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sinsin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du1.yaml new file mode 100644 index 0000000..d1d512a --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du1.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 1 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sinsin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du2.yaml new file mode 100644 index 0000000..e0d48c8 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du2.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 2 +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sinsin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du5.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du5.yaml new file mode 100644 index 0000000..6f0ea11 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/args/sinsin_eps1e-2_du5.yaml @@ -0,0 +1,15 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +u0: 1. +du: 5. +CFL: 4.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "sinsin" diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-1.yaml new file mode 100644 index 0000000..76d5873 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-1.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 1.e-1 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-2.yaml new file mode 100644 index 0000000..f8fccd3 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-2.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 1.e-2 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-3.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-3.yaml new file mode 100644 index 0000000..798bc67 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e-3.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 1.e-3 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e0.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e0.yaml new file mode 100644 index 0000000..6562ba4 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/1e0.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 1.e0 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-1.yaml new file mode 100644 index 0000000..93d252b --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-1.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 2.e-1 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-2.yaml new file mode 100644 index 0000000..638746a --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-2.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 2.e-2 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-3.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-3.yaml new file mode 100644 index 0000000..afe96db --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e-3.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 2.e-3 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e0.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e0.yaml new file mode 100644 index 0000000..a867e22 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/2e0.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 2.e0 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-1.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-1.yaml new file mode 100644 index 0000000..9b2f9e6 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-1.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 4.e-1 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-2.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-2.yaml new file mode 100644 index 0000000..c3b6706 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-2.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 4.e-2 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-3.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-3.yaml new file mode 100644 index 0000000..15c940f --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e-3.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 4.e-3 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e0.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e0.yaml new file mode 100644 index 0000000..59c9465 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/4e0.yaml @@ -0,0 +1,14 @@ +save: "../save/burgers/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: 0. +xR: 1. +epsilon: 4.e0 +CFL: 2.5e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/config.yaml b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/config.yaml new file mode 100644 index 0000000..3bf3e39 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/config/multi/config.yaml @@ -0,0 +1,12 @@ +save: "../save/burgers/" +dt_save: 0.05 +ini_time: 0. +fin_time: 2. +nx: 1024 +xL: -1. +xR: 1. +epsilon: 1.e-2 +CFL: 4.e-1 +if_second_order: 1. +numbers: 10000 +show_steps: 100 diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/run_testset.sh b/pdebench/data_gen/data_gen_NLE/BurgersEq/run_testset.sh new file mode 100644 index 0000000..95db658 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/run_testset.sh @@ -0,0 +1,26 @@ +#!/bin/bash +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=possin_eps1e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=possin_eps1e1.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=possin_eps1e2.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=possin_eps1e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=possin_eps1e-2.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=possin_eps1e-3.yaml +# +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e1.yaml +#CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e2.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e-2.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e-3.yaml +# +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e-2_u01e1.yaml +#CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e-2_u01e2.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e-2_u01e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sin_eps1e-2_u01e-2.yaml +# +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sinsin_eps1e-2_du1.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sinsin_eps1e-2_du01.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sinsin_eps1e-2_du2.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sinsin_eps1e-2_du5.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sinsin_eps1e-2_du05.yaml +CUDA_VISIBLE_DEVICES='3' python3 burgers_Hydra.py +args=sinsin_eps1e-2_du025.yaml diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/run_trainset.sh b/pdebench/data_gen/data_gen_NLE/BurgersEq/run_trainset.sh new file mode 100644 index 0000000..25994e9 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/run_trainset.sh @@ -0,0 +1,13 @@ +#!/bin/sh +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=1e0.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=1e-1.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=1e-2.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=1e-3.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=2e0.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=2e-1.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=2e-2.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=2e-3.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=4e0.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=4e-1.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=4e-2.yaml +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=4e-3.yaml diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py new file mode 100644 index 0000000..fe7b86f --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py @@ -0,0 +1,705 @@ +""" + + + File: CFD_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import logging +import sys +import time +from functools import partial + +import hydra +import jax.numpy as jnp +from jax import device_put, jit, lax + +# Hydra +from omegaconf import DictConfig + +# if double precision +# from jax.config import config +# config.update("jax_enable_x64", True) + +sys.path.append("..") +from utils import Courant_HD, Courant_vis_HD, bc_HD, init_HD, limiting_HD, save_data_HD + +logger = logging.getLogger(__name__) + + +def _pass(carry): + return carry + + +# Init arguments with Hydra +@hydra.main(config_path="config", config_name="config") +def main(cfg: DictConfig) -> None: + # physical constants + gamma = cfg.args.gamma # 3D non-relativistic gas + gammi1 = gamma - 1.0 + gamminv1 = 1.0 / gammi1 + gamgamm1inv = gamma * gamminv1 + gammi1 = gamma - 1.0 + + visc = cfg.args.zeta + cfg.args.eta / 3.0 + + BCs = ["trans", "periodic", "KHI"] # reflect + assert cfg.args.bc in BCs, "bc should be in 'trans, reflect, periodic'" + + dx = (cfg.args.xR - cfg.args.xL) / cfg.args.nx + dx_inv = 1.0 / dx + dy = (cfg.args.yR - cfg.args.yL) / cfg.args.ny + dy_inv = 1.0 / dy + dz = (cfg.args.zR - cfg.args.zL) / cfg.args.nz + dz_inv = 1.0 / dz + + # cell edge coordinate + xe = jnp.linspace(cfg.args.xL, cfg.args.xR, cfg.args.nx + 1) + ye = jnp.linspace(cfg.args.yL, cfg.args.yR, cfg.args.ny + 1) + ze = jnp.linspace(cfg.args.zL, cfg.args.zR, cfg.args.nz + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + yc = ye[:-1] + 0.5 * dy + zc = ze[:-1] + 0.5 * dz + + def evolve(Q): + t = cfg.args.ini_time + tsave = t + steps = 0 + i_save = 0 + tm_ini = time.time() + dt = 0.0 + + while t < cfg.args.fin_time: + if t >= tsave: + logger.info(f"save data at t = {t:.3f}") + save_data_HD(Q[:, 2:-2, 2:-2, 2:-2], xc, yc, zc, i_save, cfg.args.save) + tsave += cfg.args.dt_save + i_save += 1 + + if steps % cfg.args.show_steps == 0 and cfg.args.if_show: + logger.info(f"now {steps:d}-steps, t = {t:.3f}, dt = {dt:.3f}") + + carry = (Q, t, dt, steps, tsave) + Q, t, dt, steps, tsave = lax.fori_loop( + 0, cfg.args.show_steps, simulation_fn, carry + ) + + tm_fin = time.time() + logger.info(f"total elapsed time is {tm_fin - tm_ini} sec") + save_data_HD( + Q[:, 2:-2, 2:-2, 2:-2], + xc, + yc, + zc, + i_save, + cfg.args.save, + cfg.args.dt_save, + if_final=True, + ) + return t + + @jit + def simulation_fn(i, carry): + Q, t, dt, steps, tsave = carry + dt = ( + Courant_HD(Q[:, 2:-2, 2:-2, 2:-2], dx, dy, dz, cfg.args.gamma) + * cfg.args.CFL + ) + dt = jnp.min(jnp.array([dt, cfg.args.fin_time - t, tsave - t])) + + def _update(carry): + Q, dt = carry + + # preditor step for calculating t+dt/2-th time step + Q_tmp = bc_HD( + Q, mode=cfg.args.bc + ) # index 2 for _U is equivalent with index 0 for u + Q_tmp = update(Q, Q_tmp, dt * 0.5) + # update using flux at t+dt/2-th time step + Q_tmp = bc_HD( + Q_tmp, mode=cfg.args.bc + ) # index 2 for _U is equivalent with index 0 for u + Q = update(Q, Q_tmp, dt) + + # update via viscosity + # d_min = jnp.min(Q[0]) + # dt_vis = Courant_vis_HD(dx, dy, dz, eta/d_min, zeta/d_min) * cfg.args.CFL # for realistic viscosity + dt_vis = ( + Courant_vis_HD(dx, dy, dz, cfg.args.eta, cfg.args.zeta) * cfg.args.CFL + ) + dt_vis = jnp.min(jnp.array([dt_vis, dt])) + t_vis = 0.0 + + carry = Q, dt, dt_vis, t_vis + Q, dt, dt_vis, t_vis = lax.while_loop( + lambda x: x[1] - x[3] > 1.0e-8, update_vis, carry + ) + return Q, dt + + carry = Q, dt + Q, dt = lax.cond(dt > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return Q, t, dt, steps, tsave + + @jit + def update(Q, Q_tmp, dt): + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * Q[0] + My = Q[2] * Q[0] + Mz = Q[3] * Q[0] + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + # calculate flux + fx = flux_x(Q_tmp) + fy = flux_y(Q_tmp) + fz = flux_z(Q_tmp) + + # update conservative variables + dtdx, dtdy, dtdz = dt * dx_inv, dt * dy_inv, dt * dz_inv + D0 -= ( + dtdx * (fx[0, 1:, 2:-2, 2:-2] - fx[0, :-1, 2:-2, 2:-2]) + + dtdy * (fy[0, 2:-2, 1:, 2:-2] - fy[0, 2:-2, :-1, 2:-2]) + + dtdz * (fz[0, 2:-2, 2:-2, 1:] - fz[0, 2:-2, 2:-2, :-1]) + ) + + Mx -= ( + dtdx * (fx[1, 1:, 2:-2, 2:-2] - fx[1, :-1, 2:-2, 2:-2]) + + dtdy * (fy[1, 2:-2, 1:, 2:-2] - fy[1, 2:-2, :-1, 2:-2]) + + dtdz * (fz[1, 2:-2, 2:-2, 1:] - fz[1, 2:-2, 2:-2, :-1]) + ) + + My -= ( + dtdx * (fx[2, 1:, 2:-2, 2:-2] - fx[2, :-1, 2:-2, 2:-2]) + + dtdy * (fy[2, 2:-2, 1:, 2:-2] - fy[2, 2:-2, :-1, 2:-2]) + + dtdz * (fz[2, 2:-2, 2:-2, 1:] - fz[2, 2:-2, 2:-2, :-1]) + ) + + Mz -= ( + dtdx * (fx[3, 1:, 2:-2, 2:-2] - fx[3, :-1, 2:-2, 2:-2]) + + dtdy * (fy[3, 2:-2, 1:, 2:-2] - fy[3, 2:-2, :-1, 2:-2]) + + dtdz * (fz[3, 2:-2, 2:-2, 1:] - fz[3, 2:-2, 2:-2, :-1]) + ) + + E0 -= ( + dtdx * (fx[4, 1:, 2:-2, 2:-2] - fx[4, :-1, 2:-2, 2:-2]) + + dtdy * (fy[4, 2:-2, 1:, 2:-2] - fy[4, 2:-2, :-1, 2:-2]) + + dtdz * (fz[4, 2:-2, 2:-2, 1:] - fz[4, 2:-2, 2:-2, :-1]) + ) + + # reverse primitive variables + Q = Q.at[0, 2:-2, 2:-2, 2:-2].set(D0) # d + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + return Q.at[4].set(jnp.where(Q[4] > 1.0e-8, Q[4], cfg.args.p_floor)) + + @jit + def update_vis(carry): + eta = cfg.args.eta + + def _update_vis_x(carry): + Q, dt = carry + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * D0 + My = Q[2] * D0 + Mz = Q[3] * D0 + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + # calculate flux + dtdx = dt * dx_inv + # here the viscosity is eta*D0, so that dv/dt = eta*d^2v/dx^2 (not realistic viscosity but fast to calculate) + Dm = 0.5 * (D0[2:-1, 2:-2, 2:-2] + D0[1:-2, 2:-2, 2:-2]) + + fMx = ( + (eta + visc) + * Dm + * dx_inv + * (Q[1, 2:-1, 2:-2, 2:-2] - Q[1, 1:-2, 2:-2, 2:-2]) + ) + fMy = eta * Dm * dx_inv * (Q[2, 2:-1, 2:-2, 2:-2] - Q[2, 1:-2, 2:-2, 2:-2]) + fMz = eta * Dm * dx_inv * (Q[3, 2:-1, 2:-2, 2:-2] - Q[3, 1:-2, 2:-2, 2:-2]) + fE = 0.5 * (eta + visc) * Dm * dx_inv * ( + Q[1, 2:-1, 2:-2, 2:-2] ** 2 - Q[1, 1:-2, 2:-2, 2:-2] ** 2 + ) + 0.5 * eta * Dm * dx_inv * ( + (Q[2, 2:-1, 2:-2, 2:-2] ** 2 - Q[2, 1:-2, 2:-2, 2:-2] ** 2) + + (Q[3, 2:-1, 2:-2, 2:-2] ** 2 - Q[3, 1:-2, 2:-2, 2:-2] ** 2) + ) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + Mx += dtdx * (fMx[1:, :, :] - fMx[:-1, :, :]) + My += dtdx * (fMy[1:, :, :] - fMy[:-1, :, :]) + Mz += dtdx * (fMz[1:, :, :] - fMz[:-1, :, :]) + E0 += dtdx * (fE[1:, :, :] - fE[:-1, :, :]) + + # reverse primitive variables + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + + return Q, dt + + def _update_vis_y(carry): + Q, dt = carry + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * D0 + My = Q[2] * D0 + Mz = Q[3] * D0 + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + # calculate flux + dtdy = dt * dy_inv + # here the viscosity is eta*D0, so that dv/dt = eta*d^2v/dx^2 (not realistic viscosity but fast to calculate) + Dm = 0.5 * (D0[2:-2, 2:-1, 2:-2] + D0[2:-2, 1:-2, 2:-2]) + + fMx = eta * Dm * dy_inv * (Q[1, 2:-2, 2:-1, 2:-2] - Q[1, 2:-2, 1:-2, 2:-2]) + fMy = ( + (eta + visc) + * Dm + * dy_inv + * (Q[2, 2:-2, 2:-1, 2:-2] - Q[2, 2:-2, 1:-2, 2:-2]) + ) + fMz = eta * Dm * dy_inv * (Q[3, 2:-2, 2:-1, 2:-2] - Q[3, 2:-2, 1:-2, 2:-2]) + fE = 0.5 * (eta + visc) * Dm * dy_inv * ( + Q[2, 2:-2, 2:-1, 2:-2] ** 2 - Q[2, 2:-2, 1:-2, 2:-2] ** 2 + ) + 0.5 * eta * Dm * dy_inv * ( + (Q[3, 2:-2, 2:-1, 2:-2] ** 2 - Q[3, 2:-2, 1:-2, 2:-2] ** 2) + + (Q[1, 2:-2, 2:-1, 2:-2] ** 2 - Q[1, 2:-2, 1:-2, 2:-2] ** 2) + ) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + Mx += dtdy * (fMx[:, 1:, :] - fMx[:, :-1, :]) + My += dtdy * (fMy[:, 1:, :] - fMy[:, :-1, :]) + Mz += dtdy * (fMz[:, 1:, :] - fMz[:, :-1, :]) + E0 += dtdy * (fE[:, 1:, :] - fE[:, :-1, :]) + + # reverse primitive variables + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + + return Q, dt + + def _update_vis_z(carry): + Q, dt = carry + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * D0 + My = Q[2] * D0 + Mz = Q[3] * D0 + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + # calculate flux + dtdz = dt * dz_inv + # here the viscosity is eta*D0, so that dv/dt = eta*d^2v/dx^2 (not realistic viscosity but fast to calculate) + Dm = 0.5 * (D0[2:-2, 2:-2, 2:-1] + D0[2:-2, 2:-2, 1:-2]) + + fMx = eta * Dm * dz_inv * (Q[1, 2:-2, 2:-2, 2:-1] - Q[1, 2:-2, 2:-2, 1:-2]) + fMy = eta * Dm * dz_inv * (Q[2, 2:-2, 2:-2, 2:-1] - Q[2, 2:-2, 2:-2, 1:-2]) + fMz = ( + (eta + visc) + * Dm + * dz_inv + * (Q[3, 2:-2, 2:-2, 2:-1] - Q[3, 2:-2, 2:-2, 1:-2]) + ) + fE = 0.5 * (eta + visc) * Dm * dz_inv * ( + Q[3, 2:-2, 2:-2, 2:-1] ** 2 - Q[3, 2:-2, 2:-2, 1:-2] ** 2 + ) + 0.5 * eta * Dm * dz_inv * ( + (Q[1, 2:-2, 2:-2, 2:-1] ** 2 - Q[1, 2:-2, 2:-2, 1:-2] ** 2) + + (Q[2, 2:-2, 2:-2, 2:-1] ** 2 - Q[2, 2:-2, 2:-2, 1:-2] ** 2) + ) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + Mx += dtdz * (fMx[:, :, 1:] - fMx[:, :, :-1]) + My += dtdz * (fMy[:, :, 1:] - fMy[:, :, :-1]) + Mz += dtdz * (fMz[:, :, 1:] - fMz[:, :, :-1]) + E0 += dtdz * (fE[:, :, 1:] - fE[:, :, :-1]) + + # reverse primitive variables + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + + return Q, dt + + Q, dt, dt_vis, t_vis = carry + Q = bc_HD( + Q, mode=cfg.args.bc + ) # index 2 for _U is equivalent with index 0 for u + dt_ev = jnp.min(jnp.array([dt, dt_vis, dt - t_vis])) + + carry = Q, dt_ev + # directional split + carry = _update_vis_x(carry) # x + carry = _update_vis_y(carry) # y + Q, d_ev = _update_vis_z(carry) # z + + t_vis += dt_ev + + return Q, dt, dt_vis, t_vis + + @jit + def flux_x(Q): + QL, QR = limiting_HD(Q, if_second_order=cfg.args.if_second_order) + # f_Riemann = HLL(QL, QR, direc=0) + return HLLC(QL, QR, direc=0) + + @jit + def flux_y(Q): + _Q = jnp.transpose(Q, (0, 2, 3, 1)) # (y, z, x) + QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) + # f_Riemann = jnp.transpose(HLL(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) + return jnp.transpose(HLLC(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) + + @jit + def flux_z(Q): + _Q = jnp.transpose(Q, (0, 3, 1, 2)) # (z, x, y) + QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) + # f_Riemann = jnp.transpose(HLL(QL, QR, direc=2), (0, 2, 3, 1)) + return jnp.transpose(HLLC(QL, QR, direc=2), (0, 2, 3, 1)) + + @partial(jit, static_argnums=(2,)) + def HLL(QL, QR, direc): + # direc = 0, 1, 2: (X, Y, Z) + iX, iY, iZ = direc + 1, (direc + 1) % 3 + 1, (direc + 2) % 3 + 1 + cfL = jnp.sqrt(gamma * QL[4] / QL[0]) + cfR = jnp.sqrt(gamma * QR[4] / QR[0]) + Sfl = jnp.minimum(QL[iX, 2:-1], QR[iX, 1:-2]) - jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # left-going wave + Sfr = jnp.maximum(QL[iX, 2:-1], QR[iX, 1:-2]) + jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # right-going wave + dcfi = 1.0 / (Sfr - Sfl + 1.0e-8) + + UL, UR = jnp.zeros_like(QL), jnp.zeros_like(QR) + UL = UL.at[0].set(QL[0]) + UL = UL.at[iX].set(QL[0] * QL[iX]) + UL = UL.at[iY].set(QL[0] * QL[iY]) + UL = UL.at[iZ].set(QL[0] * QL[iZ]) + UL = UL.at[4].set( + gamminv1 * QL[4] + + 0.5 * (UL[iX] * QL[iX] + UL[iY] * QL[iY] + UL[iZ] * QL[iZ]) + ) + UR = UR.at[0].set(QR[0]) + UR = UR.at[iX].set(QR[0] * QR[iX]) + UR = UR.at[iY].set(QR[0] * QR[iY]) + UR = UR.at[iZ].set(QR[0] * QR[iZ]) + UR = UR.at[4].set( + gamminv1 * QR[4] + + 0.5 * (UR[iX] * QR[iX] + UR[iY] * QR[iY] + UR[iZ] * QR[iZ]) + ) + + fL, fR = jnp.zeros_like(QL), jnp.zeros_like(QR) + fL = fL.at[0].set(UL[iX]) + fL = fL.at[iX].set(UL[iX] * QL[iX] + QL[4]) + fL = fL.at[iY].set(UL[iX] * QL[iY]) + fL = fL.at[iZ].set(UL[iX] * QL[iZ]) + fL = fL.at[4].set((UL[4] + QL[4]) * QL[iX]) + fR = fR.at[0].set(UR[iX]) + fR = fR.at[iX].set(UR[iX] * QR[iX] + QR[4]) + fR = fR.at[iY].set(UR[iX] * QR[iY]) + fR = fR.at[iZ].set(UR[iX] * QR[iZ]) + fR = fR.at[4].set((UR[4] + QR[4]) * QR[iX]) + # upwind advection scheme + fHLL = dcfi * ( + Sfr * fR[:, 1:-2] + - Sfl * fL[:, 2:-1] + + Sfl * Sfr * (UL[:, 2:-1] - UR[:, 1:-2]) + ) + + # L: left of cell = right-going, R: right of cell: left-going + f_Riemann = jnp.where(Sfl > 0.0, fR[:, 1:-2], fHLL) + return jnp.where(Sfr < 0.0, fL[:, 2:-1], f_Riemann) + + @partial(jit, static_argnums=(2,)) + def HLLC(QL, QR, direc): + """full-Godunov method -- exact shock solution""" + + iX, iY, iZ = direc + 1, (direc + 1) % 3 + 1, (direc + 2) % 3 + 1 + cfL = jnp.sqrt(gamma * QL[4] / QL[0]) + cfR = jnp.sqrt(gamma * QR[4] / QR[0]) + Sfl = jnp.minimum(QL[iX, 2:-1], QR[iX, 1:-2]) - jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # left-going wave + Sfr = jnp.maximum(QL[iX, 2:-1], QR[iX, 1:-2]) + jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # right-going wave + + UL, UR = jnp.zeros_like(QL), jnp.zeros_like(QR) + UL = UL.at[0].set(QL[0]) + UL = UL.at[iX].set(QL[0] * QL[iX]) + UL = UL.at[iY].set(QL[0] * QL[iY]) + UL = UL.at[iZ].set(QL[0] * QL[iZ]) + UL = UL.at[4].set( + gamminv1 * QL[4] + + 0.5 * (UL[iX] * QL[iX] + UL[iY] * QL[iY] + UL[iZ] * QL[iZ]) + ) + UR = UR.at[0].set(QR[0]) + UR = UR.at[iX].set(QR[0] * QR[iX]) + UR = UR.at[iY].set(QR[0] * QR[iY]) + UR = UR.at[iZ].set(QR[0] * QR[iZ]) + UR = UR.at[4].set( + gamminv1 * QR[4] + + 0.5 * (UR[iX] * QR[iX] + UR[iY] * QR[iY] + UR[iZ] * QR[iZ]) + ) + + Va = ( + (Sfr - QL[iX, 2:-1]) * UL[iX, 2:-1] + - (Sfl - QR[iX, 1:-2]) * UR[iX, 1:-2] + - QL[4, 2:-1] + + QR[4, 1:-2] + ) + Va /= (Sfr - QL[iX, 2:-1]) * QL[0, 2:-1] - (Sfl - QR[iX, 1:-2]) * QR[0, 1:-2] + Pa = QR[4, 1:-2] + QR[0, 1:-2] * (Sfl - QR[iX, 1:-2]) * (Va - QR[iX, 1:-2]) + + # shock jump condition + Dal = QR[0, 1:-2] * (Sfl - QR[iX, 1:-2]) / (Sfl - Va) # right-hand density + Dar = QL[0, 2:-1] * (Sfr - QL[iX, 2:-1]) / (Sfr - Va) # left-hand density + + fL, fR = jnp.zeros_like(QL), jnp.zeros_like(QR) + fL = fL.at[0].set(UL[iX]) + fL = fL.at[iX].set(UL[iX] * QL[iX] + QL[4]) + fL = fL.at[iY].set(UL[iX] * QL[iY]) + fL = fL.at[iZ].set(UL[iX] * QL[iZ]) + fL = fL.at[4].set((UL[4] + QL[4]) * QL[iX]) + fR = fR.at[0].set(UR[iX]) + fR = fR.at[iX].set(UR[iX] * QR[iX] + QR[4]) + fR = fR.at[iY].set(UR[iX] * QR[iY]) + fR = fR.at[iZ].set(UR[iX] * QR[iZ]) + fR = fR.at[4].set((UR[4] + QR[4]) * QR[iX]) + # upwind advection scheme + far, fal = jnp.zeros_like(QL[:, 2:-1]), jnp.zeros_like(QR[:, 1:-2]) + far = far.at[0].set(Dar * Va) + far = far.at[iX].set(Dar * Va**2 + Pa) + far = far.at[iY].set(Dar * Va * QL[iY, 2:-1]) + far = far.at[iZ].set(Dar * Va * QL[iZ, 2:-1]) + far = far.at[4].set( + ( + gamgamm1inv * Pa + + 0.5 * Dar * (Va**2 + QL[iY, 2:-1] ** 2 + QL[iZ, 2:-1] ** 2) + ) + * Va + ) + fal = fal.at[0].set(Dal * Va) + fal = fal.at[iX].set(Dal * Va**2 + Pa) + fal = fal.at[iY].set(Dal * Va * QR[iY, 1:-2]) + fal = fal.at[iZ].set(Dal * Va * QR[iZ, 1:-2]) + fal = fal.at[4].set( + ( + gamgamm1inv * Pa + + 0.5 * Dal * (Va**2 + QR[iY, 1:-2] ** 2 + QR[iZ, 1:-2] ** 2) + ) + * Va + ) + + f_Riemann = jnp.where( + Sfl > 0.0, fR[:, 1:-2], fL[:, 2:-1] + ) # Sf2 > 0 : supersonic + f_Riemann = jnp.where( + Sfl * Va < 0.0, fal, f_Riemann + ) # SL < 0 and Va > 0 : sub-sonic + return jnp.where( + Sfr * Va < 0.0, far, f_Riemann + ) # Va < 0 and SR > 0 : sub-sonic + # f_Riemann = jnp.where(Sfr < 0., fL[:, 2:-1], f_Riemann) # SR < 0 : supersonic + + Q = jnp.zeros([5, cfg.args.nx + 4, cfg.args.ny + 4, cfg.args.nz + 4]) + Q = init_HD( + Q, + xc, + yc, + zc, + mode=cfg.args.init_mode, + direc="x", + init_key=cfg.args.init_key, + M0=cfg.args.M0, + dk=cfg.args.dk, + gamma=cfg.args.gamma, + ) + Q = device_put(Q) # putting variables in GPU (not necessary??) + t = evolve(Q) + logger.info(f"final time is: {t:.3f}") + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py new file mode 100644 index 0000000..ecc4865 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py @@ -0,0 +1,1004 @@ +""" + + + File: CFD_multi_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import logging +import os +import random +import sys +import time +from functools import partial +from math import ceil, exp, log + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put, jit, lax +from omegaconf import DictConfig +from utils import ( + Courant_HD, + Courant_vis_HD, + bc_HD, + init_multi_HD, + init_multi_HD_2DRand, + init_multi_HD_2DTurb, + init_multi_HD_3DRand, + init_multi_HD_3DTurb, + init_multi_HD_KH, + init_multi_HD_shock, + limiting_HD, +) + +logger = logging.getLogger(__name__) + +# Hydra + +os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" +os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".9" + +sys.path.append("..") + +# if double precision +# from jax.config import config +# config.update("jax_enable_x64", True) + + +def _pass(carry): + return carry + + +# Init arguments with Hydra +@hydra.main(config_path="config", config_name="config") +def main(cfg: DictConfig) -> None: + # physical constants + gamma = cfg.args.gamma # 3D non-relativistic gas + gammi1 = gamma - 1.0 + gamminv1 = 1.0 / gammi1 + gamgamm1inv = gamma * gamminv1 + gammi1 = gamma - 1.0 + + BCs = ["trans", "periodic", "KHI"] # reflect + assert cfg.args.bc in BCs, "bc should be in 'trans, reflect, periodic'" + + dx = (cfg.args.xR - cfg.args.xL) / cfg.args.nx + dx_inv = 1.0 / dx + dy = (cfg.args.yR - cfg.args.yL) / cfg.args.ny + dy_inv = 1.0 / dy + dz = (cfg.args.zR - cfg.args.zL) / cfg.args.nz + dz_inv = 1.0 / dz + + # cell edge coordinate + xe = jnp.linspace(cfg.args.xL, cfg.args.xR, cfg.args.nx + 1) + ye = jnp.linspace(cfg.args.yL, cfg.args.yR, cfg.args.ny + 1) + ze = jnp.linspace(cfg.args.zL, cfg.args.zR, cfg.args.nz + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + yc = ye[:-1] + 0.5 * dy + zc = ze[:-1] + 0.5 * dz + + show_steps = cfg.args.show_steps + ini_time = cfg.args.ini_time + fin_time = cfg.args.fin_time + dt_save = cfg.args.dt_save + + # t-coordinate + it_tot = ceil((fin_time - ini_time) / dt_save) + 1 + tc = jnp.arange(it_tot + 1) * dt_save + + # set viscosity + if cfg.args.if_rand_param: + zeta = exp( + random.uniform(log(0.001), log(10)) + ) # uniform number between 0.01 to 100 + eta = exp( + random.uniform(log(0.001), log(10)) + ) # uniform number between 0.01 to 100 + else: + zeta = cfg.args.zeta + eta = cfg.args.eta + logger.info(f"zeta: {zeta:>5f}, eta: {eta:>5f}") + visc = zeta + eta / 3.0 + + def evolve(Q): + t = ini_time + tsave = t + steps = 0 + i_save = 0 + dt = 0.0 + + tm_ini = time.time() + + DDD = jnp.zeros([it_tot, cfg.args.nx, cfg.args.ny, cfg.args.nz]) + VVx = jnp.zeros([it_tot, cfg.args.nx, cfg.args.ny, cfg.args.nz]) + VVy = jnp.zeros([it_tot, cfg.args.nx, cfg.args.ny, cfg.args.nz]) + VVz = jnp.zeros([it_tot, cfg.args.nx, cfg.args.ny, cfg.args.nz]) + PPP = jnp.zeros([it_tot, cfg.args.nx, cfg.args.ny, cfg.args.nz]) + # initial time-step + DDD = DDD.at[0].set(Q[0, 2:-2, 2:-2, 2:-2]) + VVx = VVx.at[0].set(Q[1, 2:-2, 2:-2, 2:-2]) + VVy = VVy.at[0].set(Q[2, 2:-2, 2:-2, 2:-2]) + VVz = VVz.at[0].set(Q[3, 2:-2, 2:-2, 2:-2]) + PPP = PPP.at[0].set(Q[4, 2:-2, 2:-2, 2:-2]) + + cond_fun = lambda x: x[0] < fin_time + + def _body_fun(carry): + def _save(_carry): + Q, tsave, i_save, DDD, VVx, VVy, VVz, PPP = _carry + + DDD = DDD.at[i_save].set(Q[0, 2:-2, 2:-2, 2:-2]) + VVx = VVx.at[i_save].set(Q[1, 2:-2, 2:-2, 2:-2]) + VVy = VVy.at[i_save].set(Q[2, 2:-2, 2:-2, 2:-2]) + VVz = VVz.at[i_save].set(Q[3, 2:-2, 2:-2, 2:-2]) + PPP = PPP.at[i_save].set(Q[4, 2:-2, 2:-2, 2:-2]) + + tsave += dt_save + i_save += 1 + return (Q, tsave, i_save, DDD, VVx, VVy, VVz, PPP) + + t, tsave, steps, i_save, dt, Q, DDD, VVx, VVy, VVz, PPP = carry + + # if save data + carry = (Q, tsave, i_save, DDD, VVx, VVy, VVz, PPP) + Q, tsave, i_save, DDD, VVx, VVy, VVz, PPP = lax.cond( + t >= tsave, _save, _pass, carry + ) + + carry = (Q, t, dt, steps, tsave) + Q, t, dt, steps, tsave = lax.fori_loop(0, show_steps, simulation_fn, carry) + + return (t, tsave, steps, i_save, dt, Q, DDD, VVx, VVy, VVz, PPP) + + carry = t, tsave, steps, i_save, dt, Q, DDD, VVx, VVy, VVz, PPP + t, tsave, steps, i_save, dt, Q, DDD, VVx, VVy, VVz, PPP = lax.while_loop( + cond_fun, _body_fun, carry + ) + + tm_fin = time.time() + logger.info(f"total elapsed time is {tm_fin - tm_ini} sec") + DDD = DDD.at[-1].set(Q[0, 2:-2, 2:-2, 2:-2]) + VVx = VVx.at[-1].set(Q[1, 2:-2, 2:-2, 2:-2]) + VVy = VVy.at[-1].set(Q[2, 2:-2, 2:-2, 2:-2]) + VVz = VVz.at[-1].set(Q[3, 2:-2, 2:-2, 2:-2]) + PPP = PPP.at[-1].set(Q[4, 2:-2, 2:-2, 2:-2]) + return t, DDD, VVx, VVy, VVz, PPP + + @jit + def simulation_fn(i, carry): + Q, t, dt, steps, tsave = carry + dt = ( + Courant_HD(Q[:, 2:-2, 2:-2, 2:-2], dx, dy, dz, cfg.args.gamma) + * cfg.args.CFL + ) + dt = jnp.min(jnp.array([dt, cfg.args.fin_time - t, tsave - t])) + + def _update(carry): + Q, dt = carry + + # preditor step for calculating t+dt/2-th time step + Q_tmp = bc_HD( + Q, mode=cfg.args.bc + ) # index 2 for _U is equivalent with index 0 for u + Q_tmp = update(Q, Q_tmp, dt * 0.5) + # update using flux at t+dt/2-th time step + Q_tmp = bc_HD( + Q_tmp, mode=cfg.args.bc + ) # index 2 for _U is equivalent with index 0 for u + Q = update(Q, Q_tmp, dt) + + # update via viscosity + # d_min = jnp.min(Q[0]) + # dt_vis = Courant_vis_HD(dx, dy, dz, eta/d_min, zeta/d_min) * cfg.args.CFL # for realistic viscosity + + dt_vis = Courant_vis_HD(dx, dy, dz, eta, zeta) * cfg.args.CFL + dt_vis = jnp.min(jnp.array([dt_vis, dt])) + t_vis = 0.0 + + carry = Q, dt, dt_vis, t_vis + Q, dt, dt_vis, t_vis = lax.while_loop( + lambda x: x[1] - x[3] > 1.0e-8, update_vis, carry + ) + return Q, dt + + carry = Q, dt + Q, dt = lax.cond(dt > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return Q, t, dt, steps, tsave + + @jit + def update(Q, Q_tmp, dt): + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * Q[0] + My = Q[2] * Q[0] + Mz = Q[3] * Q[0] + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + # calculate flux + fx = flux_x(Q_tmp) + fy = flux_y(Q_tmp) + fz = flux_z(Q_tmp) + + # update conservative variables + dtdx, dtdy, dtdz = dt * dx_inv, dt * dy_inv, dt * dz_inv + D0 -= ( + dtdx * (fx[0, 1:, 2:-2, 2:-2] - fx[0, :-1, 2:-2, 2:-2]) + + dtdy * (fy[0, 2:-2, 1:, 2:-2] - fy[0, 2:-2, :-1, 2:-2]) + + dtdz * (fz[0, 2:-2, 2:-2, 1:] - fz[0, 2:-2, 2:-2, :-1]) + ) + + Mx -= ( + dtdx * (fx[1, 1:, 2:-2, 2:-2] - fx[1, :-1, 2:-2, 2:-2]) + + dtdy * (fy[1, 2:-2, 1:, 2:-2] - fy[1, 2:-2, :-1, 2:-2]) + + dtdz * (fz[1, 2:-2, 2:-2, 1:] - fz[1, 2:-2, 2:-2, :-1]) + ) + + My -= ( + dtdx * (fx[2, 1:, 2:-2, 2:-2] - fx[2, :-1, 2:-2, 2:-2]) + + dtdy * (fy[2, 2:-2, 1:, 2:-2] - fy[2, 2:-2, :-1, 2:-2]) + + dtdz * (fz[2, 2:-2, 2:-2, 1:] - fz[2, 2:-2, 2:-2, :-1]) + ) + + Mz -= ( + dtdx * (fx[3, 1:, 2:-2, 2:-2] - fx[3, :-1, 2:-2, 2:-2]) + + dtdy * (fy[3, 2:-2, 1:, 2:-2] - fy[3, 2:-2, :-1, 2:-2]) + + dtdz * (fz[3, 2:-2, 2:-2, 1:] - fz[3, 2:-2, 2:-2, :-1]) + ) + + E0 -= ( + dtdx * (fx[4, 1:, 2:-2, 2:-2] - fx[4, :-1, 2:-2, 2:-2]) + + dtdy * (fy[4, 2:-2, 1:, 2:-2] - fy[4, 2:-2, :-1, 2:-2]) + + dtdz * (fz[4, 2:-2, 2:-2, 1:] - fz[4, 2:-2, 2:-2, :-1]) + ) + + # reverse primitive variables + Q = Q.at[0, 2:-2, 2:-2, 2:-2].set(D0) # d + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + return Q.at[4].set(jnp.where(Q[4] > 1.0e-8, Q[4], cfg.args.p_floor)) + + @jit + def update_vis(carry): + def _update_vis_x(carry): + Q, dt = carry + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * D0 + My = Q[2] * D0 + Mz = Q[3] * D0 + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + # calculate flux + dtdx = dt * dx_inv + # here the viscosity is eta*D0, so that dv/dt = eta*d^2v/dx^2 (not realistic viscosity but fast to calculate) + Dm = 0.5 * (D0[2:-1, 2:-2, 2:-2] + D0[1:-2, 2:-2, 2:-2]) + + fMx = ( + (eta + visc) + * Dm + * dx_inv + * (Q[1, 2:-1, 2:-2, 2:-2] - Q[1, 1:-2, 2:-2, 2:-2]) + ) + fMy = eta * Dm * dx_inv * (Q[2, 2:-1, 2:-2, 2:-2] - Q[2, 1:-2, 2:-2, 2:-2]) + fMz = eta * Dm * dx_inv * (Q[3, 2:-1, 2:-2, 2:-2] - Q[3, 1:-2, 2:-2, 2:-2]) + fE = 0.5 * (eta + visc) * Dm * dx_inv * ( + Q[1, 2:-1, 2:-2, 2:-2] ** 2 - Q[1, 1:-2, 2:-2, 2:-2] ** 2 + ) + 0.5 * eta * Dm * dx_inv * ( + (Q[2, 2:-1, 2:-2, 2:-2] ** 2 - Q[2, 1:-2, 2:-2, 2:-2] ** 2) + + (Q[3, 2:-1, 2:-2, 2:-2] ** 2 - Q[3, 1:-2, 2:-2, 2:-2] ** 2) + ) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + Mx += dtdx * (fMx[1:, :, :] - fMx[:-1, :, :]) + My += dtdx * (fMy[1:, :, :] - fMy[:-1, :, :]) + Mz += dtdx * (fMz[1:, :, :] - fMz[:-1, :, :]) + E0 += dtdx * (fE[1:, :, :] - fE[:-1, :, :]) + + # reverse primitive variables + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + + return Q, dt + + def _update_vis_y(carry): + Q, dt = carry + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * D0 + My = Q[2] * D0 + Mz = Q[3] * D0 + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + # calculate flux + dtdy = dt * dy_inv + # here the viscosity is eta*D0, so that dv/dt = eta*d^2v/dx^2 (not realistic viscosity but fast to calculate) + Dm = 0.5 * (D0[2:-2, 2:-1, 2:-2] + D0[2:-2, 1:-2, 2:-2]) + + fMx = eta * Dm * dy_inv * (Q[1, 2:-2, 2:-1, 2:-2] - Q[1, 2:-2, 1:-2, 2:-2]) + fMy = ( + (eta + visc) + * Dm + * dy_inv + * (Q[2, 2:-2, 2:-1, 2:-2] - Q[2, 2:-2, 1:-2, 2:-2]) + ) + fMz = eta * Dm * dy_inv * (Q[3, 2:-2, 2:-1, 2:-2] - Q[3, 2:-2, 1:-2, 2:-2]) + fE = 0.5 * (eta + visc) * Dm * dy_inv * ( + Q[2, 2:-2, 2:-1, 2:-2] ** 2 - Q[2, 2:-2, 1:-2, 2:-2] ** 2 + ) + 0.5 * eta * Dm * dy_inv * ( + (Q[3, 2:-2, 2:-1, 2:-2] ** 2 - Q[3, 2:-2, 1:-2, 2:-2] ** 2) + + (Q[1, 2:-2, 2:-1, 2:-2] ** 2 - Q[1, 2:-2, 1:-2, 2:-2] ** 2) + ) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + Mx += dtdy * (fMx[:, 1:, :] - fMx[:, :-1, :]) + My += dtdy * (fMy[:, 1:, :] - fMy[:, :-1, :]) + Mz += dtdy * (fMz[:, 1:, :] - fMz[:, :-1, :]) + E0 += dtdy * (fE[:, 1:, :] - fE[:, :-1, :]) + + # reverse primitive variables + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + + return Q, dt + + def _update_vis_z(carry): + Q, dt = carry + # calculate conservative variables + D0 = Q[0] + Mx = Q[1] * D0 + My = Q[2] * D0 + Mz = Q[3] * D0 + E0 = Q[4] * gamminv1 + 0.5 * (Mx * Q[1] + My * Q[2] + Mz * Q[3]) + + # calculate flux + dtdz = dt * dz_inv + # here the viscosity is eta*D0, so that dv/dt = eta*d^2v/dx^2 (not realistic viscosity but fast to calculate) + Dm = 0.5 * (D0[2:-2, 2:-2, 2:-1] + D0[2:-2, 2:-2, 1:-2]) + + fMx = eta * Dm * dz_inv * (Q[1, 2:-2, 2:-2, 2:-1] - Q[1, 2:-2, 2:-2, 1:-2]) + fMy = eta * Dm * dz_inv * (Q[2, 2:-2, 2:-2, 2:-1] - Q[2, 2:-2, 2:-2, 1:-2]) + fMz = ( + (eta + visc) + * Dm + * dz_inv + * (Q[3, 2:-2, 2:-2, 2:-1] - Q[3, 2:-2, 2:-2, 1:-2]) + ) + fE = 0.5 * (eta + visc) * Dm * dz_inv * ( + Q[3, 2:-2, 2:-2, 2:-1] ** 2 - Q[3, 2:-2, 2:-2, 1:-2] ** 2 + ) + 0.5 * eta * Dm * dz_inv * ( + (Q[1, 2:-2, 2:-2, 2:-1] ** 2 - Q[1, 2:-2, 2:-2, 1:-2] ** 2) + + (Q[2, 2:-2, 2:-2, 2:-1] ** 2 - Q[2, 2:-2, 2:-2, 1:-2] ** 2) + ) + + D0 = D0[2:-2, 2:-2, 2:-2] + Mx = Mx[2:-2, 2:-2, 2:-2] + My = My[2:-2, 2:-2, 2:-2] + Mz = Mz[2:-2, 2:-2, 2:-2] + E0 = E0[2:-2, 2:-2, 2:-2] + + Mx += dtdz * (fMx[:, :, 1:] - fMx[:, :, :-1]) + My += dtdz * (fMy[:, :, 1:] - fMy[:, :, :-1]) + Mz += dtdz * (fMz[:, :, 1:] - fMz[:, :, :-1]) + E0 += dtdz * (fE[:, :, 1:] - fE[:, :, :-1]) + + # reverse primitive variables + Q = Q.at[1, 2:-2, 2:-2, 2:-2].set(Mx / D0) # vx + Q = Q.at[2, 2:-2, 2:-2, 2:-2].set(My / D0) # vy + Q = Q.at[3, 2:-2, 2:-2, 2:-2].set(Mz / D0) # vz + Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( + gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) + ) # p + + return Q, dt + + Q, dt, dt_vis, t_vis = carry + Q = bc_HD( + Q, mode=cfg.args.bc + ) # index 2 for _U is equivalent with index 0 for u + dt_ev = jnp.min(jnp.array([dt, dt_vis, dt - t_vis])) + + carry = Q, dt_ev + # directional split + carry = _update_vis_x(carry) # x + carry = _update_vis_y(carry) # y + Q, d_ev = _update_vis_z(carry) # z + + t_vis += dt_ev + + return Q, dt, dt_vis, t_vis + + @jit + def flux_x(Q): + QL, QR = limiting_HD(Q, if_second_order=cfg.args.if_second_order) + # f_Riemann = HLL(QL, QR, direc=0) + return HLLC(QL, QR, direc=0) + + @jit + def flux_y(Q): + _Q = jnp.transpose(Q, (0, 2, 3, 1)) # (y, z, x) + QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) + # f_Riemann = jnp.transpose(HLL(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) + return jnp.transpose(HLLC(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) + + @jit + def flux_z(Q): + _Q = jnp.transpose(Q, (0, 3, 1, 2)) # (z, x, y) + QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) + # f_Riemann = jnp.transpose(HLL(QL, QR, direc=2), (0, 2, 3, 1)) + return jnp.transpose(HLLC(QL, QR, direc=2), (0, 2, 3, 1)) + + @partial(jit, static_argnums=(2,)) + def HLL(QL, QR, direc): + # direc = 0, 1, 2: (X, Y, Z) + iX, iY, iZ = direc + 1, (direc + 1) % 3 + 1, (direc + 2) % 3 + 1 + cfL = jnp.sqrt(gamma * QL[4] / QL[0]) + cfR = jnp.sqrt(gamma * QR[4] / QR[0]) + Sfl = jnp.minimum(QL[iX, 2:-1], QR[iX, 1:-2]) - jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # left-going wave + Sfr = jnp.maximum(QL[iX, 2:-1], QR[iX, 1:-2]) + jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # right-going wave + dcfi = 1.0 / (Sfr - Sfl + 1.0e-8) + + UL, UR = jnp.zeros_like(QL), jnp.zeros_like(QR) + UL = UL.at[0].set(QL[0]) + UL = UL.at[iX].set(QL[0] * QL[iX]) + UL = UL.at[iY].set(QL[0] * QL[iY]) + UL = UL.at[iZ].set(QL[0] * QL[iZ]) + UL = UL.at[4].set( + gamminv1 * QL[4] + + 0.5 * (UL[iX] * QL[iX] + UL[iY] * QL[iY] + UL[iZ] * QL[iZ]) + ) + UR = UR.at[0].set(QR[0]) + UR = UR.at[iX].set(QR[0] * QR[iX]) + UR = UR.at[iY].set(QR[0] * QR[iY]) + UR = UR.at[iZ].set(QR[0] * QR[iZ]) + UR = UR.at[4].set( + gamminv1 * QR[4] + + 0.5 * (UR[iX] * QR[iX] + UR[iY] * QR[iY] + UR[iZ] * QR[iZ]) + ) + + fL, fR = jnp.zeros_like(QL), jnp.zeros_like(QR) + fL = fL.at[0].set(UL[iX]) + fL = fL.at[iX].set(UL[iX] * QL[iX] + QL[4]) + fL = fL.at[iY].set(UL[iX] * QL[iY]) + fL = fL.at[iZ].set(UL[iX] * QL[iZ]) + fL = fL.at[4].set((UL[4] + QL[4]) * QL[iX]) + fR = fR.at[0].set(UR[iX]) + fR = fR.at[iX].set(UR[iX] * QR[iX] + QR[4]) + fR = fR.at[iY].set(UR[iX] * QR[iY]) + fR = fR.at[iZ].set(UR[iX] * QR[iZ]) + fR = fR.at[4].set((UR[4] + QR[4]) * QR[iX]) + # upwind advection scheme + fHLL = dcfi * ( + Sfr * fR[:, 1:-2] + - Sfl * fL[:, 2:-1] + + Sfl * Sfr * (UL[:, 2:-1] - UR[:, 1:-2]) + ) + + # L: left of cell = right-going, R: right of cell: left-going + f_Riemann = jnp.where(Sfl > 0.0, fR[:, 1:-2], fHLL) + return jnp.where(Sfr < 0.0, fL[:, 2:-1], f_Riemann) + + @partial(jit, static_argnums=(2,)) + def HLLC(QL, QR, direc): + """full-Godunov method -- exact shock solution""" + + iX, iY, iZ = direc + 1, (direc + 1) % 3 + 1, (direc + 2) % 3 + 1 + cfL = jnp.sqrt(gamma * QL[4] / QL[0]) + cfR = jnp.sqrt(gamma * QR[4] / QR[0]) + Sfl = jnp.minimum(QL[iX, 2:-1], QR[iX, 1:-2]) - jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # left-going wave + Sfr = jnp.maximum(QL[iX, 2:-1], QR[iX, 1:-2]) + jnp.maximum( + cfL[2:-1], cfR[1:-2] + ) # right-going wave + + UL, UR = jnp.zeros_like(QL), jnp.zeros_like(QR) + UL = UL.at[0].set(QL[0]) + UL = UL.at[iX].set(QL[0] * QL[iX]) + UL = UL.at[iY].set(QL[0] * QL[iY]) + UL = UL.at[iZ].set(QL[0] * QL[iZ]) + UL = UL.at[4].set( + gamminv1 * QL[4] + + 0.5 * (UL[iX] * QL[iX] + UL[iY] * QL[iY] + UL[iZ] * QL[iZ]) + ) + UR = UR.at[0].set(QR[0]) + UR = UR.at[iX].set(QR[0] * QR[iX]) + UR = UR.at[iY].set(QR[0] * QR[iY]) + UR = UR.at[iZ].set(QR[0] * QR[iZ]) + UR = UR.at[4].set( + gamminv1 * QR[4] + + 0.5 * (UR[iX] * QR[iX] + UR[iY] * QR[iY] + UR[iZ] * QR[iZ]) + ) + + Va = ( + (Sfr - QL[iX, 2:-1]) * UL[iX, 2:-1] + - (Sfl - QR[iX, 1:-2]) * UR[iX, 1:-2] + - QL[4, 2:-1] + + QR[4, 1:-2] + ) + Va /= (Sfr - QL[iX, 2:-1]) * QL[0, 2:-1] - (Sfl - QR[iX, 1:-2]) * QR[0, 1:-2] + Pa = QR[4, 1:-2] + QR[0, 1:-2] * (Sfl - QR[iX, 1:-2]) * (Va - QR[iX, 1:-2]) + + # shock jump condition + Dal = QR[0, 1:-2] * (Sfl - QR[iX, 1:-2]) / (Sfl - Va) # right-hand density + Dar = QL[0, 2:-1] * (Sfr - QL[iX, 2:-1]) / (Sfr - Va) # left-hand density + + fL, fR = jnp.zeros_like(QL), jnp.zeros_like(QR) + fL = fL.at[0].set(UL[iX]) + fL = fL.at[iX].set(UL[iX] * QL[iX] + QL[4]) + fL = fL.at[iY].set(UL[iX] * QL[iY]) + fL = fL.at[iZ].set(UL[iX] * QL[iZ]) + fL = fL.at[4].set((UL[4] + QL[4]) * QL[iX]) + fR = fR.at[0].set(UR[iX]) + fR = fR.at[iX].set(UR[iX] * QR[iX] + QR[4]) + fR = fR.at[iY].set(UR[iX] * QR[iY]) + fR = fR.at[iZ].set(UR[iX] * QR[iZ]) + fR = fR.at[4].set((UR[4] + QR[4]) * QR[iX]) + # upwind advection scheme + far, fal = jnp.zeros_like(QL[:, 2:-1]), jnp.zeros_like(QR[:, 1:-2]) + far = far.at[0].set(Dar * Va) + far = far.at[iX].set(Dar * Va**2 + Pa) + far = far.at[iY].set(Dar * Va * QL[iY, 2:-1]) + far = far.at[iZ].set(Dar * Va * QL[iZ, 2:-1]) + far = far.at[4].set( + ( + gamgamm1inv * Pa + + 0.5 * Dar * (Va**2 + QL[iY, 2:-1] ** 2 + QL[iZ, 2:-1] ** 2) + ) + * Va + ) + fal = fal.at[0].set(Dal * Va) + fal = fal.at[iX].set(Dal * Va**2 + Pa) + fal = fal.at[iY].set(Dal * Va * QR[iY, 1:-2]) + fal = fal.at[iZ].set(Dal * Va * QR[iZ, 1:-2]) + fal = fal.at[4].set( + ( + gamgamm1inv * Pa + + 0.5 * Dal * (Va**2 + QR[iY, 1:-2] ** 2 + QR[iZ, 1:-2] ** 2) + ) + * Va + ) + + f_Riemann = jnp.where( + Sfl > 0.0, fR[:, 1:-2], fL[:, 2:-1] + ) # Sf2 > 0 : supersonic + f_Riemann = jnp.where( + Sfl * Va < 0.0, fal, f_Riemann + ) # SL < 0 and Va > 0 : sub-sonic + return jnp.where( + Sfr * Va < 0.0, far, f_Riemann + ) # Va < 0 and SR > 0 : sub-sonic + # f_Riemann = jnp.where(Sfr < 0., fL[:, 2:-1], f_Riemann) # SR < 0 : supersonic + + Q = jnp.zeros( + [cfg.args.numbers, 5, cfg.args.nx + 4, cfg.args.ny + 4, cfg.args.nz + 4] + ) + if cfg.args.init_mode_Multi == "1D_rand": + Q = Q.at[:, 0, 2:-2, 2:-2, 2:-2].set( + init_multi_HD( + xc, + yc, + zc, + numbers=cfg.args.numbers, + k_tot=3, + init_key=cfg.args.init_key, + num_choise_k=2, + umin=1.0e0, + umax=1.0e1, + if_renorm=True, + ) + ) + Q = Q.at[:, 4, 2:-2, 2:-2, 2:-2].set( + init_multi_HD( + xc, + yc, + zc, + numbers=cfg.args.numbers, + k_tot=3, + init_key=cfg.args.init_key + 1, + num_choise_k=2, + umin=1.0e1, + umax=1.0e2, + if_renorm=True, + ) + ) + Q = Q.at[:, 1, 2:-2, 2:-2, 2:-2].set( + init_multi_HD( + xc, + yc, + zc, + numbers=cfg.args.numbers, + k_tot=3, + init_key=cfg.args.init_key + 2, + num_choise_k=2, + if_renorm=False, + ) + ) + elif cfg.args.init_mode_Multi == "1D_shocks": + Q = Q.at[:, 0, 2:-2, 2:-2, 2:-2].set( + init_multi_HD_shock( + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key, + umin=1.0e0, + umax=1.0e1, + ) + ) + Q = Q.at[:, 4, 2:-2, 2:-2, 2:-2].set( + init_multi_HD_shock( + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key + 1, + umin=1.0e1, + umax=1.0e2, + ) + ) + Q = Q.at[:, 1, 2:-2, 2:-2, 2:-2].set( + init_multi_HD_shock( + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key + 2, + umin=-0.5e0, + umax=0.5e0, + ) + ) + elif cfg.args.init_mode_Multi == "KHs": + assert 2.0 * yc[0] - (yc[1] - yc[0]) == 0.0, "yL is assumed 0!" + logger.info("now we are coming into KHs...") + Q = init_multi_HD_KH( + Q, + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key, + M0=cfg.args.M0, + dkMx=cfg.args.dkMx, + gamma=cfg.args.gamma, + ) + elif cfg.args.init_mode_Multi == "2D_Turbs": + logger.info("now we are coming into 2DTurbs......") + Q = init_multi_HD_2DTurb( + Q, + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key, + M0=cfg.args.M0, + k_tot=cfg.args.k_tot, + gamma=cfg.args.gamma, + ) + elif cfg.args.init_mode_Multi == "2D_rand": + assert ( # noqa: PT018 + xe[0] == 0.0 and ye[0] == 0.0 and xe[-1] == 1.0 and ye[-1] == 1.0 + ), "xc, yc should be between 0 and 1!" + logger.info("now we are coming into 2Drand......") + Q = init_multi_HD_2DRand( + Q, + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key, + M0=cfg.args.M0, + k_tot=cfg.args.k_tot, + gamma=cfg.args.gamma, + ) + elif cfg.args.init_mode_Multi == "3D_Turbs": + logger.info("now we are coming into 3DTurbs......") + Q = init_multi_HD_3DTurb( + Q, + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key, + M0=cfg.args.M0, + k_tot=cfg.args.k_tot, + gamma=cfg.args.gamma, + ) + elif cfg.args.init_mode_Multi == "3D_rand": + logger.info("now we are coming into 3Drand......") + Q = init_multi_HD_3DRand( + Q, + xc, + yc, + zc, + numbers=cfg.args.numbers, + init_key=cfg.args.init_key, + M0=cfg.args.M0, + k_tot=cfg.args.k_tot, + gamma=cfg.args.gamma, + ) + logger.info("initial conditions were prepared!!") + Q = device_put(Q) # putting variables in GPU (not necessary??) + + local_device_count = jax.local_device_count() + pm_evolve = jax.pmap(jax.vmap(evolve, axis_name="j"), axis_name="i") + t, DDD, VVx, VVy, VVz, PPP = pm_evolve( + Q.reshape( + [ + local_device_count, + cfg.args.numbers // local_device_count, + 5, + cfg.args.nx + 4, + cfg.args.ny + 4, + cfg.args.nz + 4, + ] + ) + ) + + itot = DDD.shape[2] + DDD = DDD.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) + VVx = VVx.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) + VVy = VVy.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) + VVz = VVz.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) + PPP = PPP.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) + logger.info("now data saving...") + jnp.save( + cfg.args.save + + "HD_Sols_" + + cfg.args.init_mode_Multi + + "_Eta" + + str(eta)[:5] + + "_Zeta" + + str(zeta)[:5] + + "_M" + + str(cfg.args.M0) + + "_key" + + str(cfg.args.init_key) + + "_D", + DDD, + ) + jnp.save( + cfg.args.save + + "HD_Sols_" + + cfg.args.init_mode_Multi + + "_Eta" + + str(eta)[:5] + + "_Zeta" + + str(zeta)[:5] + + "_M" + + str(cfg.args.M0) + + "_key" + + str(cfg.args.init_key) + + "_Vx", + VVx, + ) + jnp.save( + cfg.args.save + + "HD_Sols_" + + cfg.args.init_mode_Multi + + "_Eta" + + str(eta)[:5] + + "_Zeta" + + str(zeta)[:5] + + "_M" + + str(cfg.args.M0) + + "_key" + + str(cfg.args.init_key) + + "_Vy", + VVy, + ) + jnp.save( + cfg.args.save + + "HD_Sols_" + + cfg.args.init_mode_Multi + + "_Eta" + + str(eta)[:5] + + "_Zeta" + + str(zeta)[:5] + + "_M" + + str(cfg.args.M0) + + "_key" + + str(cfg.args.init_key) + + "_Vz", + VVz, + ) + jnp.save( + cfg.args.save + + "HD_Sols_" + + cfg.args.init_mode_Multi + + "_Eta" + + str(eta)[:5] + + "_Zeta" + + str(zeta)[:5] + + "_M" + + str(cfg.args.M0) + + "_key" + + str(cfg.args.init_key) + + "_P", + PPP, + ) + jnp.save(cfg.args.save + "/x_coordinate", xc) + jnp.save(cfg.args.save + "/y_coordinate", yc) + jnp.save(cfg.args.save + "/z_coordinate", zc) + jnp.save(cfg.args.save + "/t_coordinate", tc) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi.yaml new file mode 100644 index 0000000..b291196 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi.yaml @@ -0,0 +1,27 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-2 +zeta: 1.e-2 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +p_floor: 1.e-4 +numbers: 1000 +init_mode_Multi: "1D_rand" +init_key: 2022 +if_rand_param: False +M0: 0. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi_shock.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi_shock.yaml new file mode 100644 index 0000000..8140b8f --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi_shock.yaml @@ -0,0 +1,27 @@ +save: "../save/CFD/" +dt_save: 0.005 +ini_time: 0. +fin_time: 0.4 +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +p_floor: 1.e-4 +numbers: 1000 +init_mode_Multi: "1D_shocks" +init_key: 2020 +if_rand_param: False +M0: 0 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi_trans.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi_trans.yaml new file mode 100644 index 0000000..6c172a3 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_Multi_trans.yaml @@ -0,0 +1,27 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +p_floor: 1.e-4 +numbers: 1000 +init_mode_Multi: "1D_rand" +init_key: 2022 +if_rand_param: False +M0: 0 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube.yaml new file mode 100644 index 0000000..3c8a1f5 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 0.4 +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube1" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 +M0: 1. +dk: 1. +dkMx: 1. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube2.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube2.yaml new file mode 100644 index 0000000..9bd4d9e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube2.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 0.15 +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube2" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 +M0: 1. +dk: 1. +dkMx: 1. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube3.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube3.yaml new file mode 100644 index 0000000..7f07b3e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube3.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.001 +ini_time: 0. +fin_time: 0.012 +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube3" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 +M0: 1. +dk: 1. +dkMx: 1. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube4.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube4.yaml new file mode 100644 index 0000000..469a755 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube4.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.001 +ini_time: 0. +fin_time: 0.035 +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube4" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 +M0: 1. +dk: 1. +dkMx: 1. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube5.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube5.yaml new file mode 100644 index 0000000..23569c7 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube5.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.001 +ini_time: 0. +fin_time: 0.012 +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube5" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 +M0: 1. +dk: 1. +dkMx: 1. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube6.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube6.yaml new file mode 100644 index 0000000..c5e3579 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube6.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube6" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 +M0: 1. +dk: 1. +dkMx: 1. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube7.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube7.yaml new file mode 100644 index 0000000..682e7fb --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/1D_ShockTube7.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 2. +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube7" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 +M0: 1. +dk: 1. +dkMx: 1. diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk1.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk1.yaml new file mode 100644 index 0000000..07b2df6 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk1.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/KH/KH_M01_dk1_Re1e3/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 0.1 +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk10.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk10.yaml new file mode 100644 index 0000000..6520922 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk10.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 0.1 +dk: 10. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk2.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk2.yaml new file mode 100644 index 0000000..f5fcd7c --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk2.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/KH/KH_M01_dk2_Re1e3/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 0.1 +dk: 2. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk5.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk5.yaml new file mode 100644 index 0000000..0893cfc --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M01_dk5.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 0.1 +dk: 5. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M02_dk1.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M02_dk1.yaml new file mode 100644 index 0000000..a81f895 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M02_dk1.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 0.2 +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M04_dk1.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M04_dk1.yaml new file mode 100644 index 0000000..4c46603 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M04_dk1.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 0.4 +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M1_dk1.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M1_dk1.yaml new file mode 100644 index 0000000..9d36c74 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M1_dk1.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/KH/KH_M1_dk1_Re1e3/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 1. +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M2_dk1.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M2_dk1.yaml new file mode 100644 index 0000000..2c176a1 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_KH_M2_dk1.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.1 +ini_time: 0. +fin_time: 5. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-3 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "KHI" +M0: 2. +dk: 1 +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_KH.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_KH.yaml new file mode 100644 index 0000000..4d80a12 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_KH.yaml @@ -0,0 +1,27 @@ +save: "../save/CFD/" +dt_save: 0.1 +ini_time: 0. +fin_time: 2. +nx: 32 +ny: 32 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "KHI" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +p_floor: 1.e-4 +numbers: 4 +init_mode_Multi: "KHs" +M0: 0.1 +dkMx: 2. +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Rand.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Rand.yaml new file mode 100644 index 0000000..a0a24e4 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Rand.yaml @@ -0,0 +1,28 @@ +save: "../save/CFD/" +dt_save: 0.05 +ini_time: 0. +fin_time: 1. +nx: 128 +ny: 128 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +p_floor: 1.e-4 +numbers: 100 +init_mode_Multi: "2D_rand" +M0: 0.1 +k_tot: 4 +init_key: 2020 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Rand_HR.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Rand_HR.yaml new file mode 100644 index 0000000..42055c1 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Rand_HR.yaml @@ -0,0 +1,28 @@ +save: "../save/CFD/" +dt_save: 0.05 +ini_time: 0. +fin_time: 1. +nx: 512 +ny: 512 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +p_floor: 1.e-4 +numbers: 20 +init_mode_Multi: "2D_rand" +M0: 0.5 +k_tot: 4 +init_key: 2020 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Turb.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Turb.yaml new file mode 100644 index 0000000..acaf6a8 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_Multi_Turb.yaml @@ -0,0 +1,28 @@ +save: "../save/CFD/" +dt_save: 0.05 +ini_time: 0. +fin_time: 1. +nx: 512 +ny: 512 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +p_floor: 1.e-4 +numbers: 20 +init_mode_Multi: "2DTurbs" +M0: 0.1 +k_tot: 4 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_ShockTube.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_ShockTube.yaml new file mode 100644 index 0000000..257764c --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_ShockTube.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 4000 +p_floor: 1.e-4 +init_mode: "2D-shock" +M0: 1. +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_TOV.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_TOV.yaml new file mode 100644 index 0000000..4fc92bc --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/2D_TOV.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +ny: 1024 +nz: 1 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 5.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +init_mode: "OTVortex" +M0: 1. +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_BlastWave.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_BlastWave.yaml new file mode 100644 index 0000000..7d753a6 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_BlastWave.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.025 +ini_time: 0. +fin_time: 0.5 +nx: 256 +ny: 256 +nz: 256 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +init_mode: "BlastWave" +M0: 1. +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_Multi_Rand.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_Multi_Rand.yaml new file mode 100644 index 0000000..f36ce4b --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_Multi_Rand.yaml @@ -0,0 +1,28 @@ +save: "../save/CFD/" +dt_save: 0.05 +ini_time: 0. +fin_time: 1. +nx: 128 +ny: 128 +nz: 128 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "3D_rand" +M0: 1. +k_tot: 4 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_Multi_TurbM1.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_Multi_TurbM1.yaml new file mode 100644 index 0000000..75bb721 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_Multi_TurbM1.yaml @@ -0,0 +1,28 @@ +save: "../save/CFD/" +dt_save: 0.05 +ini_time: 0. +fin_time: 1. +nx: 64 +ny: 64 +nz: 64 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +numbers: 20 +init_mode_Multi: "3D_Turbs" +M0: 1. +k_tot: 4 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM01.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM01.yaml new file mode 100644 index 0000000..3978a73 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM01.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.025 +ini_time: 0. +fin_time: 1. +nx: 256 +ny: 256 +nz: 256 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +init_mode: "turbulence" +M0: 1.e-1 +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM05.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM05.yaml new file mode 100644 index 0000000..59d5cc5 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM05.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.025 +ini_time: 0. +fin_time: 1. +nx: 256 +ny: 256 +nz: 256 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +init_mode: "turbulence" +M0: 5.e-1 +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM1.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM1.yaml new file mode 100644 index 0000000..3978a73 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM1.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.025 +ini_time: 0. +fin_time: 1. +nx: 256 +ny: 256 +nz: 256 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +init_mode: "turbulence" +M0: 1.e-1 +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM2.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM2.yaml new file mode 100644 index 0000000..a47492e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM2.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.025 +ini_time: 0. +fin_time: 1. +nx: 256 +ny: 256 +nz: 256 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +init_mode: "turbulence" +M0: 2. +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM4.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM4.yaml new file mode 100644 index 0000000..fbe6aa0 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/3D_TurbM4.yaml @@ -0,0 +1,29 @@ +save: "../save/CFD/" +dt_save: 0.025 +ini_time: 0. +fin_time: 1. +nx: 256 +ny: 256 +nz: 256 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +zL: 0. +zR: 1. +bc: "periodic" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 400 +p_floor: 1.e-4 +init_mode: "turbulence" +M0: 4. +dk: 1. +dkMx: 1. +numbers: 4 +init_mode_Multi: "KHs" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/default.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/default.yaml new file mode 100644 index 0000000..185a013 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/args/default.yaml @@ -0,0 +1,26 @@ +save: "../save/CFD/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +ny: 1 +nz: 1 +xL: -1. +xR: 1. +yL: -1. +yR: 1. +zL: -1. +zR: 1. +bc: "trans" +gamma: 1.6666666666666667 +eta: 1.e-8 +zeta: 1.e-8 +CFL: 3.e-1 +if_second_order: 1. +if_show: 1 +show_steps: 100 +init_mode: "shocktube1" +p_floor: 1.e-4 +numbers: 10 +init_mode_Multi: "1D_shocks" +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/config.yaml b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/config.yaml new file mode 100644 index 0000000..b3b9da6 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/config/config.yaml @@ -0,0 +1,9 @@ +defaults: + - _self_ + - override hydra/hydra_logging: disabled + - override hydra/job_logging: disabled + +hydra: + output_subdir: null + run: + dir: . diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset.sh new file mode 100644 index 0000000..1b02061 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta1e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta1e1.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta1e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta2e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta2e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta4e0.yaml +CUDA_VISIBLE_DEVICES='3' python3 advection_exact_Hydra.py +args=beta4e-1.yaml diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset_3DTurb.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset_3DTurb.sh new file mode 100644 index 0000000..b29f590 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset_3DTurb.sh @@ -0,0 +1,6 @@ +#!/bin/bash +CUDA_VISIBLE_DEVICES='0' python3 CFD_Hydra.py +args=3D_TurbM01.yaml +CUDA_VISIBLE_DEVICES='0' python3 CFD_Hydra.py +args=3D_TurbM05.yaml +CUDA_VISIBLE_DEVICES='0' python3 CFD_Hydra.py +args=3D_TurbM1.yaml +CUDA_VISIBLE_DEVICES='0' python3 CFD_Hydra.py +args=3D_TurbM2.yaml +CUDA_VISIBLE_DEVICES='0' python3 CFD_Hydra.py +args=3D_TurbM4.yaml diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset_KHI.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset_KHI.sh new file mode 100644 index 0000000..03a6151 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_testset_KHI.sh @@ -0,0 +1,9 @@ +#!/bin/bash +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M01_dk1.yaml +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M02_dk1.yaml +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M04_dk1.yaml +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M1_dk1.yaml +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M2_dk1.yaml +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M01_dk2.yaml +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M01_dk5.yaml +CUDA_VISIBLE_DEVICES='1' python3 CFD_Hydra.py +args=2D_KH_M01_dk10.yaml diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1D.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1D.sh new file mode 100644 index 0000000..b1e6886 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1D.sh @@ -0,0 +1,10 @@ +#!/bin/bash +nn=1 +key=2020 +while [ "$nn" -le 10 ]; do + CUDA_VISIBLE_DEVICES='0,1' python3 CFD_multi_Hydra.py +args=1D_Multi.yaml ++args.init_key="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1DShock.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1DShock.sh new file mode 100644 index 0000000..0324a72 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1DShock.sh @@ -0,0 +1,10 @@ +#!/bin/bash +nn=1 +key=2031 +while [ "$nn" -le 10 ]; do + CUDA_VISIBLE_DEVICES='0,1' python3 CFD_multi_Hydra.py +args=1D_Multi_shock.yaml ++args.init_key="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1D_trans.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1D_trans.sh new file mode 100644 index 0000000..2ce7891 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_1D_trans.sh @@ -0,0 +1,10 @@ +#!/bin/bash +nn=1 +key=2020 +while [ "$nn" -le 10 ]; do + CUDA_VISIBLE_DEVICES='0,1' python3 CFD_multi_Hydra.py +args=1D_Multi_trans.yaml ++args.init_key="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_2D.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_2D.sh new file mode 100644 index 0000000..d7752c6 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_2D.sh @@ -0,0 +1,11 @@ +#!/bin/bash +nn=1 +key=2031 +while [ "$nn" -le 100 ]; do + CUDA_VISIBLE_DEVICES='0,1,2,3' python3 CFD_multi_Hydra.py +args=2D_Multi_Rand.yaml ++args.init_key="$key" + #CUDA_VISIBLE_DEVICES='0,1,2,3' python3 CFD_multi_Hydra.py +args=2D_Multi_Rand_HR.yaml ++args.init_key="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_2DTurb.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_2DTurb.sh new file mode 100644 index 0000000..bdd281e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_2DTurb.sh @@ -0,0 +1,11 @@ +#!/bin/bash +nn=1 +key=2031 +#while [ "$nn" -le 100 ]; do +while [ "$nn" -le 55 ]; do + CUDA_VISIBLE_DEVICES='0,1,2,3' python3 CFD_multi_Hydra.py +args=2D_Multi_Turb.yaml ++args.init_key="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_3D.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_3D.sh new file mode 100644 index 0000000..ed0028f --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_3D.sh @@ -0,0 +1,10 @@ +#!/bin/bash +nn=1 +key=2031 +while [ "$nn" -le 10 ]; do + CUDA_VISIBLE_DEVICES='0,1' python3 CFD_multi_Hydra.py +args=3D_Multi_Rand.yaml ++args.init_key="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_3DTurb.sh b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_3DTurb.sh new file mode 100644 index 0000000..9ce9b1f --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/run_trainset_3DTurb.sh @@ -0,0 +1,10 @@ +#!/bin/bash +nn=1 +key=2031 +while [ "$nn" -le 6 ]; do + CUDA_VISIBLE_DEVICES='0,1,2,3' python3 CFD_multi_Hydra.py +args=3D_Multi_TurbM1.yaml ++args.init_key="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/Data_Merge.py b/pdebench/data_gen/data_gen_NLE/Data_Merge.py new file mode 100644 index 0000000..16e930b --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/Data_Merge.py @@ -0,0 +1,408 @@ +""" + + + File: Data_Merge.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +from pathlib import Path + +import h5py +import hydra +import numpy as np +from omegaconf import DictConfig + +""" +Data_Merge.py +This is a script creating HDF5 from the generated data (numpy array) by our data generation scripts. +A more detailed explanation how to use this script is provided in the README. +""" + + +def _mergeRD(var, DataND, savedir): + _vars = ["2D", "nu"] + if var not in _vars: + return None + + idx = 0 + data_list = Path(savedir).glob(var + "*key*.npy") + data_list.sort() + for data_file in data_list: + test = np.load(data_file).squeeze() + batch = min(test.shape[0], DataND.shape[0] - idx) + if var == "2D": + DataND[idx : idx + batch] = test[:batch, -2] + else: + DataND[idx : idx + batch] = test[:batch] + idx += batch + + return DataND[:idx] + + +def _merge(var, DataND, dim, savedir): + if dim == 1: + _vars = ["D", "P", "Vx"] + elif dim == 2: + _vars = ["D", "P", "Vx", "Vy"] + elif dim == 3: + _vars = ["D", "P", "Vx", "Vy", "Vz"] + if var not in _vars: + return None + + idx = 0 + data = Path(savedir).glob("HD*" + var + ".npy") + data.sort() + for data_file in data: + test = np.load(data_file).squeeze() + batch = min(test.shape[0], DataND.shape[0] - idx) + DataND[idx : idx + batch] = test[:batch] + idx += batch + + return DataND[:idx] + + +def nan_check(data): + data = np.abs(data).reshape([data.shape[0], data.shape[1], -1]).sum(axis=-1) + return np.where(data[:, -2] < 1.0e-6)[0], np.where(data[:, -2] > 1.0e-6)[0] + + +def merge(type, dim, bd, nbatch, savedir): + if type == "CFD": + data = Path(savedir).glob("HD*D.npy") + data.sort() + test = np.load(data[0]) + __nbatch, nt, nx, ny, nz = test.shape + _nbatch = __nbatch * len(data) + assert ( + nbatch <= _nbatch + ), "nbatch should be equal or less than the number of generated samples" + assert ( + 2 * nbatch > _nbatch + ), "2*nbatch should be larger than the number of generated samples" + + if dim == 1: + DataND = np.zeros([2 * nbatch, nt, nx], dtype=np.float32) + vars = ["D", "P", "Vx"] + elif dim == 2: + DataND = np.zeros([2 * nbatch, nt, nx, ny], dtype=np.float32) + vars = ["D", "P", "Vx", "Vy"] + elif dim == 3: + DataND = np.zeros([2 * nbatch, nt, nx, ny, nz], dtype=np.float32) + vars = ["D", "P", "Vx", "Vy", "Vz"] + + elif type == "ReacDiff": + data = Path(savedir).glob("nu*.npy") + data.sort() + test = np.load(data[0]) + __nbatch, nx, ny = test.shape + _nbatch = __nbatch * len(data) + assert ( + nbatch == _nbatch + ), "nbatch should be equal or less than the number of generated samples" + DataND = np.zeros([nbatch, nx, ny], dtype=np.float32) + vars = ["2D", "nu"] + + for var in vars: + if type == "CFD": + _DataND = _merge(var, DataND, dim, savedir) + if var == "D": + idx_neg, idx_pos = nan_check(_DataND) + if len(idx_pos) < nbatch: + pass + _DataND = _DataND[idx_pos] + _DataND = _DataND[:nbatch] + np.save(savedir + "/" + var + ".npy", _DataND) + elif type == "ReacDiff": + DataND = _mergeRD(var, DataND, savedir) + np.save(savedir + "/" + var + ".npy", DataND) + + data = Path(savedir).glob("*npy") + data.sort() + + if type == "CFD": + zcrd = np.load(data[-1]) + del data[-1] + ycrd = np.load(data[-1]) + del data[-1] + xcrd = np.load(data[-1]) + del data[-1] + tcrd = np.load(data[-1]) + del data[-1] + if type == "ReacDiff": + # data = glob.glob('save/' + type + '/nu*key*npy') + data = Path(savedir).glob("nu*key*npy") + data.sort() + _beta = data[0].split("/")[-1].split("_")[3] + flnm = savedir + "/2D_DecayFlow_" + _beta + "_Train.hdf5" + with h5py.File(flnm, "w") as f: + f.create_dataet("tensor", data=np.load(savedir + "/2D.npy")[:, None, :, :]) + f.create_dataet("nu", data=np.load(savedir + "/nu.npy")) + f.create_dataet("x-coordinate", data=xcrd) + f.create_dataet("y-coordinate", data=ycrd) + f.attrs["beta"] = float(_beta[4:]) + return 0 + + mode = data[1].split("/")[-1].split("_")[3] + _eta = data[1].split("/")[-1].split("_")[4] + _zeta = data[1].split("/")[-1].split("_")[5] + _M = data[1].split("/")[-1].split("_")[6] + if dim == 1: + flnm = ( + savedir + + "/1D_CFD_" + + mode + + "_" + + _eta + + "_" + + _zeta + + "_" + + bd + + "_Train.hdf5" + ) + elif dim == 2: + flnm = ( + savedir + + "/2D_CFD_" + + mode + + "_" + + _eta + + "_" + + _zeta + + "_" + + _M + + "_" + + bd + + "_Train.hdf5" + ) + elif dim == 3: + flnm = ( + savedir + + "/3D_CFD_" + + mode + + "_" + + _eta + + "_" + + _zeta + + "_" + + _M + + "_" + + bd + + "_Train.hdf5" + ) + + del DataND + + with h5py.File(flnm, "w") as f: + f.create_dataet("density", data=np.load(savedir + "/D.npy")) + f.create_dataet("pressure", data=np.load(savedir + "/P.npy")) + f.create_dataet("Vx", data=np.load(savedir + "/Vx.npy")) + if dim > 1: + f.create_dataet("Vy", data=np.load(savedir + "/Vy.npy")) + f.create_dataet("y-coordinate", data=ycrd) + if dim == 3: + f.create_dataet("Vz", data=np.load(savedir + "/Vz.npy")) + f.create_dataet("z-coordinate", data=zcrd) + f.create_dataet("x-coordinate", data=xcrd) + f.create_dataet("t-coordinate", data=tcrd) + eta = float(_eta[3:]) + zeta = float(_zeta[4:]) + f.attrs["eta"] = eta + f.attrs["zeta"] = zeta + if dim > 1: + M = float(_M[1:]) + f.attrs["M"] = M + return None + return None + + +def transform(type, savedir): + data = Path(savedir).glob("*npy") + data.sort() + xcrd = np.load(data[-1]) + del data[-1] + tcrd = np.load(data[-1]) + del data[-1] + + flnm = data[0] + with h5py.File(flnm[:-3] + "hdf5", "w") as f: + _data = np.load(flnm) + + f.create_dataset("tensor", data=_data.astype(np.float32)) + f.create_dataset("x-coordinate", data=xcrd) + f.create_dataset("t-coordinate", data=tcrd) + if type == "advection": + beta = float(flnm.split("/")[-1].split("_")[3][4:-4]) # advection train + f.attrs["beta"] = beta + + elif type == "burgers": + Nu = float(flnm.split("/")[-1].split("_")[-1][2:-4]) # Burgers test/train + f.attrs["Nu"] = Nu + + elif type == "ReacDiff": + Rho = float(flnm.split("/")[-1].split("_")[-1][3:-4]) # reac-diff test + Nu = float(flnm.split("/")[-1].split("_")[-2][2:]) # reac-diff test + f.attrs["Nu"] = Nu + f.attrs["rho"] = Rho + + +# Init arguments with Hydra +@hydra.main(config_path="config", config_name="config") +def main(cfg: DictConfig) -> None: + pde1ds = ["advection", "burgers", "ReacDiff"] + if cfg.args.type in pde1ds and cfg.args.dim == 1: + transform(type=cfg.args.type, savedir=cfg.args.savedir) + else: + bds = ["periodic", "trans"] + assert cfg.args.bd in bds, "bd should be either periodic or trans" + merge( + type=cfg.args.type, + dim=cfg.args.dim, + bd=cfg.args.bd, + nbatch=cfg.args.nbatch, + savedir=cfg.args.savedir, + ) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/README.md b/pdebench/data_gen/data_gen_NLE/README.md new file mode 100644 index 0000000..695c5cb --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/README.md @@ -0,0 +1,112 @@ +## Data Generation + +#### Data generation for DarcyFlow Equation: + +- Run the shell script: + +```bash +bash data_gen/data_gen_NLE/ReactionDiffusionEq/run_DarcyFlow2D.sh +``` + +which will in turn run the python script +`data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_soluion_Hydra.py` + +- Update `data_gen/data_gen_NLE/config/config.yaml` to: + +```yaml +type: "ReacDiff" # 'advection'/'ReacDiff'/'burgers'/'CFD' +dim: 2 +``` + +- Finally, run the data merge script: + +```bash +python data_gen/data_gen_NLE/Data_Merge.py +``` + +--- + +#### Data generation for 1D Advection Equation: + +``` +# generate data and save as .npy array +cd PDEBench/pdebench/data_gen/data_gen_NLE/AdvectionEq + +# Either generate a single file +CUDA_VISIBLE_DEVICES='2,3' python3 advection_multi_solution_Hydra.py +multi=beta1e0.yaml + +# Or generate all files +bash run_trainset.sh +``` + +- Update `data_gen/data_gen_NLE/config/config.yaml` to: + +```yaml +type: "advection" # 'advection'/'ReacDiff'/'burgers'/'CFD' +dim: 1 +savedir: "./save/advection" +``` + +``` +# serialize to hdf5 by transforming npy file +cd .. +python Data_Merge.py +``` + +--- + +#### Data generation for 1D Burgers' Equation: + +``` +# generate data and save as .npy array +cd PDEBench/pdebench/data_gen/data_gen_NLE/BurgersEq/ + +# Either generate a single file +CUDA_VISIBLE_DEVICES='0,2' python3 burgers_multi_solution_Hydra.py +multi=1e-1.yaml + +# Or generate all files +bash run_trainset.sh +``` + +- Update `data_gen/data_gen_NLE/config/config.yaml` to: + +```yaml +type: "burgers" # 'advection'/'ReacDiff'/'burgers'/'CFD' +dim: 1 +savedir: "./save/burgers" +``` + +``` +# serialize to hdf5 by transforming npy file +cd .. +python Data_Merge.py +``` + +--- + +#### Data generation for 1D Reaction Diffusion Equation: + +``` +# generate data and save as .npy array +cd PDEBench/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/ + +# Either generate a single file +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho2e0_Nu5e0.yaml + +# Or generate all files +bash run_trainset.sh +``` + +- Update `data_gen/data_gen_NLE/config/config.yaml` to: + +```yaml +type: "ReacDiff" # 'advection'/'ReacDiff'/'burgers'/'CFD' +dim: 1 +savedir: "./save/ReacDiff" +``` + +``` +# serialize to hdf5 by transforming npy file +cd .. +python Data_Merge.py +``` diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu1e0.yaml new file mode 100644 index 0000000..82892c5 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu1e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 1.e0 +rho: 1.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu2e0.yaml new file mode 100644 index 0000000..b3a41a4 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu2e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 2.e0 +rho: 1.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu5e-1.yaml new file mode 100644 index 0000000..3c4721f --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu5e-1.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e-1 +rho: 1.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu5e0.yaml new file mode 100644 index 0000000..986fcf1 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e0_Nu5e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e0 +rho: 1.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu1e0.yaml new file mode 100644 index 0000000..21aef17 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu1e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 1.e0 +rho: 1.e1 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu2e0.yaml new file mode 100644 index 0000000..2df6cd9 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu2e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 2.e0 +rho: 1.e1 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu5e-1.yaml new file mode 100644 index 0000000..114a987 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu5e-1.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e-1 +rho: 1.e1 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu5e0.yaml new file mode 100644 index 0000000..a2c6bea --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho1e1_Nu5e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e0 +rho: 1.e1 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu1e0.yaml new file mode 100644 index 0000000..087cd63 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu1e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 1.e0 +rho: 2.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu2e0.yaml new file mode 100644 index 0000000..4121f05 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu2e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 2.e0 +rho: 2.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu5e-1.yaml new file mode 100644 index 0000000..c90ac52 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu5e-1.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e-1 +rho: 2.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu5e0.yaml new file mode 100644 index 0000000..5861512 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho2e0_Nu5e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e0 +rho: 2.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu1e0.yaml new file mode 100644 index 0000000..eea3bbe --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu1e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 1.e0 +rho: 5.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu2e0.yaml new file mode 100644 index 0000000..97781ac --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu2e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 2.e0 +rho: 5.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu5e-1.yaml new file mode 100644 index 0000000..f584839 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu5e-1.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e-1 +rho: 5.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu5e0.yaml new file mode 100644 index 0000000..d2ce48e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/Rho5e0_Nu5e0.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 5.e0 +rho: 5.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config.yaml new file mode 100644 index 0000000..82892c5 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config.yaml @@ -0,0 +1,13 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 6.28318530718 +nu: 1.e0 +rho: 1.e0 +CFL: 4.e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config_2D.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config_2D.yaml new file mode 100644 index 0000000..8909062 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/args/config_2D.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 64 +ny: 64 +xL: 0. +xR: 6.28318530718 +yL: 0. +yR: 6.28318530718 +nu: 1.e0 +rho: 1.e0 +CFL: 2.5e-1 +if_show: 1 +show_steps: 100 +init_mode: "react" diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu1e0.yaml new file mode 100644 index 0000000..180ed0a --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu1e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 1.e0 +rho: 1.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu1e1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu1e1.yaml new file mode 100644 index 0000000..f313da7 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu1e1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 1.e1 +rho: 1.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu2e0.yaml new file mode 100644 index 0000000..3aaa868 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu2e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 2.e0 +rho: 1.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu5e-1.yaml new file mode 100644 index 0000000..63445a0 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu5e-1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 5.e-1 +rho: 1.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu5e0.yaml new file mode 100644 index 0000000..d483cf8 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e0_Nu5e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 256 +xL: 0. +xR: 1. +nu: 5.e0 +rho: 1.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu1e0.yaml new file mode 100644 index 0000000..f12dea6 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu1e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 1.e0 +rho: 1.e1 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu1e1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu1e1.yaml new file mode 100644 index 0000000..da255bc --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu1e1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 1.e1 +rho: 1.e1 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu2e0.yaml new file mode 100644 index 0000000..3b12386 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu2e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 2.e0 +rho: 1.e1 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu5e-1.yaml new file mode 100644 index 0000000..17d81ce --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu5e-1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 5.e-1 +rho: 1.e1 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu5e0.yaml new file mode 100644 index 0000000..1400217 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho1e1_Nu5e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 256 +xL: 0. +xR: 1. +nu: 5.e0 +rho: 1.e1 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu1e0.yaml new file mode 100644 index 0000000..35594b5 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu1e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 1.e0 +rho: 2.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu1e1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu1e1.yaml new file mode 100644 index 0000000..b12a5af --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu1e1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 1.e1 +rho: 2.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu2e0.yaml new file mode 100644 index 0000000..f3d9c47 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu2e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 2.e0 +rho: 2.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu5e-1.yaml new file mode 100644 index 0000000..86db6f7 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu5e-1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 5.e-1 +rho: 2.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu5e0.yaml new file mode 100644 index 0000000..b8ed164 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho2e0_Nu5e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 256 +xL: 0. +xR: 1. +nu: 5.e0 +rho: 2.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu1e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu1e0.yaml new file mode 100644 index 0000000..ff4cc51 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu1e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 1.e0 +rho: 5.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu1e1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu1e1.yaml new file mode 100644 index 0000000..472c8fb --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu1e1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 1.e1 +rho: 5.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 100 +if_second_order: 1. +show_steps: 10000 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu2e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu2e0.yaml new file mode 100644 index 0000000..999ef74 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu2e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 512 +xL: 0. +xR: 1. +nu: 2.e0 +rho: 5.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu5e-1.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu5e-1.yaml new file mode 100644 index 0000000..5092fab --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu5e-1.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 5.e-1 +rho: 5.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 10000 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu5e0.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu5e0.yaml new file mode 100644 index 0000000..261a646 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/Rho5e0_Nu5e0.yaml @@ -0,0 +1,16 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 256 +xL: 0. +xR: 1. +nu: 5.e0 +rho: 5.e0 +CFL: 2.5e-1 +if_show: 1 +numbers: 10000 +if_second_order: 1. +show_steps: 100 +init_key: 2022 +if_rand_param: False diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/config.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/config.yaml new file mode 100644 index 0000000..c26af21 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/config.yaml @@ -0,0 +1,14 @@ +save: "../save/ReacDiff/" +dt_save: 0.01 +ini_time: 0. +fin_time: 1. +nx: 1024 +xL: 0. +xR: 1. +nu: 5.e0 +rho: 1.e1 +CFL: 4.e-1 +if_show: 1 +numbers: 100 +if_second_order: 1. +show_steps: 100 diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/config_2D.yaml b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/config_2D.yaml new file mode 100644 index 0000000..5c1891e --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/config/multi/config_2D.yaml @@ -0,0 +1,17 @@ +save: "../save/ReacDiff/" +dt_save: 0.25 +ini_time: 0. +fin_time: 2. +nx: 128 +ny: 128 +xL: 0. +xR: 1. +yL: 0. +yR: 1. +beta: 1. +CFL: 2.5e-1 +if_show: 1 +numbers: 200 +if_second_order: 1. +show_steps: 100 +init_key: 2022 diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_solution_Hydra.py new file mode 100644 index 0000000..a0a7243 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_solution_Hydra.py @@ -0,0 +1,351 @@ +""" + + + File: reaction_diffusion_2D_multi_solution_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import sys +from math import ceil + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put, lax, vmap + +# Hydra +from omegaconf import DictConfig + +sys.path.append("..") +from utils import Courant_diff_2D, bc_2D, init_multi_2DRand + + +def _pass(carry): + return carry + + +# Init arguments with Hydra +@hydra.main(config_path="config") +def main(cfg: DictConfig) -> None: + # basic parameters + dx = (cfg.multi.xR - cfg.multi.xL) / cfg.multi.nx + dx_inv = 1.0 / dx + dy = (cfg.multi.yR - cfg.multi.yL) / cfg.multi.ny + dy_inv = 1.0 / dy + + # cell edge coordinate + xe = jnp.linspace(cfg.multi.xL, cfg.multi.xR, cfg.multi.nx + 1) + ye = jnp.linspace(cfg.multi.yL, cfg.multi.yR, cfg.multi.ny + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + yc = ye[:-1] + 0.5 * dy + + show_steps = cfg.multi.show_steps + ini_time = cfg.multi.ini_time + fin_time = cfg.multi.fin_time + dt_save = cfg.multi.dt_save + CFL = cfg.multi.CFL + beta = cfg.multi.beta + + # t-coordinate + it_tot = ceil((fin_time - ini_time) / dt_save) + 1 + tc = jnp.arange(it_tot + 1) * dt_save + + @jax.jit + def evolve(u, nu): + t = ini_time + tsave = t + steps = 0 + i_save = 0 + dt = 0.0 + uu = jnp.zeros([it_tot, u.shape[0], u.shape[1]]) + uu = uu.at[0].set(u) + + cond_fun = lambda x: x[0] < fin_time + + def _body_fun(carry): + def _show(_carry): + u, tsave, i_save, uu = _carry + uu = uu.at[i_save].set(u) + tsave += dt_save + i_save += 1 + return (u, tsave, i_save, uu) + + t, tsave, steps, i_save, dt, u, uu, nu = carry + + carry = (u, tsave, i_save, uu) + u, tsave, i_save, uu = lax.cond(t >= tsave, _show, _pass, carry) + + carry = (u, t, dt, steps, tsave, nu) + u, t, dt, steps, tsave, nu = lax.fori_loop( + 0, show_steps, simulation_fn, carry + ) + + return (t, tsave, steps, i_save, dt, u, uu, nu) + + carry = t, tsave, steps, i_save, dt, u, uu, nu + t, tsave, steps, i_save, dt, u, uu, nu = lax.while_loop( + cond_fun, _body_fun, carry + ) + return uu.at[-1].set(u) + + @jax.jit + def simulation_fn(i, carry): + u, t, dt, steps, tsave, nu = carry + dt = Courant_diff_2D(dx, dy, nu) * CFL + dt = jnp.min(jnp.array([dt, fin_time - t, tsave - t])) + + def _update(carry): + u, dt, nu = carry + # preditor step for calculating t+dt/2-th time step + u_tmp = update(u, u, dt * 0.5, nu) + # update using flux at t+dt/2-th time step + u = update(u, u_tmp, dt, nu) + return u, dt, nu + + carry = u, dt, nu + u, dt, nu = lax.cond(dt > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return u, t, dt, steps, tsave, nu + + @jax.jit + def update(u, u_tmp, dt, nu): + # boundary condition + _u = bc_2D(u_tmp, mode="Neumann") + # diffusion + dtdx = dt * dx_inv + dtdy = dt * dy_inv + fx = ( + -0.5 + * (nu[2:-1, 2:-2] + nu[1:-2, 2:-2]) + * dx_inv + * (_u[2:-1, 2:-2] - _u[1:-2, 2:-2]) + ) + fy = ( + -0.5 + * (nu[2:-2, 2:-1] + nu[2:-2, 1:-2]) + * dy_inv + * (_u[2:-2, 2:-1] - _u[2:-2, 1:-2]) + ) + u -= dtdx * (fx[1:, :] - fx[:-1, :]) + dtdy * (fy[:, 1:] - fy[:, :-1]) + # source term: f = 1 * beta + u += dt * beta + return u + + u = init_multi_2DRand( + xc, yc, numbers=cfg.multi.numbers, k_tot=4, init_key=cfg.multi.init_key + ) + u = device_put(u) # putting variables in GPU (not necessary??) + + # generate random diffusion coefficient + key = jax.random.PRNGKey(cfg.multi.init_key) + xms = jax.random.uniform( + key, shape=[cfg.multi.numbers, 5], minval=cfg.multi.xL, maxval=cfg.multi.xR + ) + key, subkey = jax.random.split(key) + yms = jax.random.uniform( + key, shape=[cfg.multi.numbers, 5], minval=cfg.multi.yL, maxval=cfg.multi.yR + ) + + key, subkey = jax.random.split(key) + stds = ( + 0.5 + * (cfg.multi.xR - cfg.multi.xL) + * jax.random.uniform(key, shape=[cfg.multi.numbers, 5]) + ) + nu = jnp.zeros_like(u) + for i in range(5): + nu += jnp.exp( + -( + (xc[None, :, None] - xms[:, None, None, i]) ** 2 + + (yc[None, None, :] - yms[:, None, None, i]) ** 2 + ) + / stds[:, None, None, i] + ) + nu = jnp.where(nu > nu.mean(), 1, 0.1) + nu = vmap(bc_2D, axis_name="i")(nu) + + local_devices = jax.local_device_count() + if local_devices > 1: + nb, nx, ny = u.shape + vm_evolve = jax.pmap(jax.vmap(evolve, axis_name="j"), axis_name="i") + uu = vm_evolve( + u.reshape([local_devices, cfg.multi.numbers // local_devices, nx, ny]), + nu.reshape( + [local_devices, cfg.multi.numbers // local_devices, nx + 4, ny + 4] + ), + ) + uu = uu.reshape([nb, -1, nx, ny]) + else: + vm_evolve = vmap(evolve, 0, 0) + uu = vm_evolve(u, nu) + + cwd = hydra.utils.get_original_cwd() + "/" + jnp.save( + cwd + + cfg.multi.save + + "/2D_ReacDiff_Multi_beta" + + str(beta)[:5] + + "_key" + + str(cfg.multi.init_key), + uu, + ) + jnp.save(cwd + cfg.multi.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.multi.save + "/y_coordinate", yc) + jnp.save(cwd + cfg.multi.save + "/t_coordinate", tc) + jnp.save( + cwd + + cfg.multi.save + + "/nu_diff_coef_beta" + + str(beta)[:5] + + "_key" + + str(cfg.multi.init_key), + nu[:, 2:-2, 2:-2], + ) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_Hydra.py b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_Hydra.py new file mode 100644 index 0000000..cfcfe26 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_Hydra.py @@ -0,0 +1,282 @@ +""" + + + File: reaction_diffusion_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import logging +import sys +import time +from math import ceil + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put, lax + +# Hydra +from omegaconf import DictConfig + +sys.path.append("..") +from utils import Courant_diff, bc, init + +logger = logging.getLogger(__name__) + + +# Init arguments with Hydra +@hydra.main(config_path="config") +def main(cfg: DictConfig) -> None: + # basic parameters + dx = (cfg.args.xR - cfg.args.xL) / cfg.args.nx + dx_inv = 1.0 / dx + + # cell edge coordinate + xe = jnp.linspace(cfg.args.xL, cfg.args.xR, cfg.args.nx + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + # t-coordinate + it_tot = ceil((cfg.args.fin_time - cfg.args.ini_time) / cfg.args.dt_save) + 1 + tc = jnp.arange(it_tot + 1) * cfg.args.dt_save + + def evolve(u): + t = cfg.args.ini_time + tsave = t + steps = 0 + i_save = 0 + tm_ini = time.time() + dt = 0.0 + + uu = jnp.zeros([it_tot, u.shape[0]]) + uu = uu.at[0].set(u) + + while t < cfg.args.fin_time: + if t >= tsave: + uu = uu.at[i_save].set(u) + tsave += cfg.args.dt_save + i_save += 1 + + if steps % cfg.args.show_steps == 0 and cfg.args.if_show: + logger.info(f"now {steps:d}-steps, t = {t:.3f}, dt = {dt:.3f}") + + carry = (u, t, dt, steps, tsave) + u, t, dt, steps, tsave = lax.fori_loop( + 0, cfg.args.show_steps, simulation_fn, carry + ) + + tm_fin = time.time() + logger.info(f"total elapsed time is {tm_fin - tm_ini} sec") + return uu, t + + @jax.jit + def simulation_fn(i, carry): + u, t, dt, steps, tsave = carry + dt = Courant_diff(dx, cfg.args.nu) * cfg.args.CFL + dt = jnp.min(jnp.array([dt, cfg.args.fin_time - t, tsave - t])) + + def _update(carry): + u, dt = carry + # preditor step for calculating t+dt/2-th time step + u_tmp = update(u, u, dt * 0.5) + # update using flux at t+dt/2-th time step + u = update(u, u_tmp, dt) + return u, dt + + def _pass(carry): + return carry + + carry = u, dt + u, dt = lax.cond(t > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return u, t, dt, steps, tsave + + @jax.jit + def update(u, u_tmp, dt): + # stiff part + u = Piecewise_Exact_Solution(u, dt) + # diffusion + f = flux(u_tmp) + u -= dt * dx_inv * (f[1 : cfg.args.nx + 1] - f[0 : cfg.args.nx]) + return u + + @jax.jit + def flux(u): + _u = bc( + u, dx, Ncell=cfg.args.nx + ) # index 2 for _U is equivalent with index 0 for u + # source term + return ( + -cfg.args.nu * (_u[2 : cfg.args.nx + 3] - _u[1 : cfg.args.nx + 2]) * dx_inv + ) + + @jax.jit + def Piecewise_Exact_Solution(u, dt): # Piecewise_Exact_Solution method + # stiff equation + return 1.0 / (1.0 + jnp.exp(-cfg.args.rho * dt) * (1.0 - u) / u) + + u = init(xc=xc, mode=cfg.args.init_mode) + u = device_put(u) # putting variables in GPU (not necessary??) + uu, t = evolve(u) + logger.info(f"final time is: {t:.3f}") + + cwd = hydra.utils.get_original_cwd() + "/" + jnp.save( + cwd + + cfg.args.save + + "/ReacDiff_" + + cfg.args.init_mode + + "_Nu" + + str(cfg.args.nu) + + "_Rho" + + str(cfg.args.rho), + uu, + ) + jnp.save(cwd + cfg.args.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.args.save + "/t_coordinate", tc) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_multi_solution_Hydra.py new file mode 100644 index 0000000..f29e3b4 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_multi_solution_Hydra.py @@ -0,0 +1,314 @@ +""" + + + File: reaction_diffusion_multi_solution_Hydra.py + Authors: Makoto Takamoto (makoto.takamoto@neclab.eu) + +NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. + + PROPRIETARY INFORMATION --- + +SOFTWARE LICENSE AGREEMENT + +ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY + +BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS +LICENSE AGREEMENT. IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR +DOWNLOAD THE SOFTWARE. + +This is a license agreement ("Agreement") between your academic institution +or non-profit organization or self (called "Licensee" or "You" in this +Agreement) and NEC Laboratories Europe GmbH (called "Licensor" in this +Agreement). All rights not specifically granted to you in this Agreement +are reserved for Licensor. + +RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: Licensor retains exclusive +ownership of any copy of the Software (as defined below) licensed under this +Agreement and hereby grants to Licensee a personal, non-exclusive, +non-transferable license to use the Software for noncommercial research +purposes, without the right to sublicense, pursuant to the terms and +conditions of this Agreement. NO EXPRESS OR IMPLIED LICENSES TO ANY OF +LICENSOR'S PATENT RIGHTS ARE GRANTED BY THIS LICENSE. As used in this +Agreement, the term "Software" means (i) the actual copy of all or any +portion of code for program routines made accessible to Licensee by Licensor +pursuant to this Agreement, inclusive of backups, updates, and/or merged +copies permitted hereunder or subsequently supplied by Licensor, including +all or any file structures, programming instructions, user interfaces and +screen formats and sequences as well as any and all documentation and +instructions related to it, and (ii) all or any derivatives and/or +modifications created or made by You to any of the items specified in (i). + +CONFIDENTIALITY/PUBLICATIONS: Licensee acknowledges that the Software is +proprietary to Licensor, and as such, Licensee agrees to receive all such +materials and to use the Software only in accordance with the terms of this +Agreement. Licensee agrees to use reasonable effort to protect the Software +from unauthorized use, reproduction, distribution, or publication. All +publication materials mentioning features or use of this software must +explicitly include an acknowledgement the software was developed by NEC +Laboratories Europe GmbH. + +COPYRIGHT: The Software is owned by Licensor. + +PERMITTED USES: The Software may be used for your own noncommercial +internal research purposes. You understand and agree that Licensor is not +obligated to implement any suggestions and/or feedback you might provide +regarding the Software, but to the extent Licensor does so, you are not +entitled to any compensation related thereto. + +DERIVATIVES: You may create derivatives of or make modifications to the +Software, however, You agree that all and any such derivatives and +modifications will be owned by Licensor and become a part of the Software +licensed to You under this Agreement. You may only use such derivatives and +modifications for your own noncommercial internal research purposes, and you +may not otherwise use, distribute or copy such derivatives and modifications +in violation of this Agreement. + +BACKUPS: If Licensee is an organization, it may make that number of copies +of the Software necessary for internal noncommercial use at a single site +within its organization provided that all information appearing in or on the +original labels, including the copyright and trademark notices are copied +onto the labels of the copies. + +USES NOT PERMITTED: You may not distribute, copy or use the Software except +as explicitly permitted herein. Licensee has not been granted any trademark +license as part of this Agreement. Neither the name of NEC Laboratories +Europe GmbH nor the names of its contributors may be used to endorse or +promote products derived from this Software without specific prior written +permission. + +You may not sell, rent, lease, sublicense, lend, time-share or transfer, in +whole or in part, or provide third parties access to prior or present +versions (or any parts thereof) of the Software. + +ASSIGNMENT: You may not assign this Agreement or your rights hereunder +without the prior written consent of Licensor. Any attempted assignment +without such consent shall be null and void. + +TERM: The term of the license granted by this Agreement is from Licensee's +acceptance of this Agreement by downloading the Software or by using the +Software until terminated as provided below. + +The Agreement automatically terminates without notice if you fail to comply +with any provision of this Agreement. Licensee may terminate this Agreement +by ceasing using the Software. Upon any termination of this Agreement, +Licensee will delete any and all copies of the Software. You agree that all +provisions which operate to protect the proprietary rights of Licensor shall +remain in force should breach occur and that the obligation of +confidentiality described in this Agreement is binding in perpetuity and, as +such, survives the term of the Agreement. + +FEE: Provided Licensee abides completely by the terms and conditions of this +Agreement, there is no fee due to Licensor for Licensee's use of the +Software in accordance with this Agreement. + +DISCLAIMER OF WARRANTIES: THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY +OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR +FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON- INFRINGEMENT. LICENSEE +BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND +RELATED MATERIALS. + +SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is +provided as part of this Agreement. + +EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent +permitted under applicable law, Licensor shall not be liable for direct, +indirect, special, incidental, or consequential damages or lost profits +related to Licensee's use of and/or inability to use the Software, even if +Licensor is advised of the possibility of such damage. + +EXPORT REGULATION: Licensee agrees to comply with any and all applicable +export control laws, regulations, and/or other laws related to embargoes and +sanction programs administered by law. + +SEVERABILITY: If any provision(s) of this Agreement shall be held to be +invalid, illegal, or unenforceable by a court or other tribunal of competent +jurisdiction, the validity, legality and enforceability of the remaining +provisions shall not in any way be affected or impaired thereby. + +NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right +or remedy under this Agreement shall be construed as a waiver of any future +or other exercise of such right or remedy by Licensor. + +GOVERNING LAW: This Agreement shall be construed and enforced in accordance +with the laws of Germany without reference to conflict of laws principles. +You consent to the personal jurisdiction of the courts of this country and +waive their rights to venue outside of Germany. + +ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and +entire agreement between Licensee and Licensor as to the matter set forth +herein and supersedes any previous agreements, understandings, and +arrangements between the parties relating hereto. + + THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. +""" + +from __future__ import annotations + +import random +import sys +from math import ceil, exp, log +from pathlib import Path + +import hydra +import jax +import jax.numpy as jnp +from jax import device_put, lax + +# Hydra +from omegaconf import DictConfig + +sys.path.append("..") +from utils import Courant_diff, bc, init_multi + + +def _pass(carry): + return carry + + +# Init arguments with Hydra +@hydra.main(config_path="config") +def main(cfg: DictConfig) -> None: + # basic parameters + dx = (cfg.multi.xR - cfg.multi.xL) / cfg.multi.nx + dx_inv = 1.0 / dx + + # cell edge coordinate + xe = jnp.linspace(cfg.multi.xL, cfg.multi.xR, cfg.multi.nx + 1) + # cell center coordinate + xc = xe[:-1] + 0.5 * dx + + show_steps = cfg.multi.show_steps + ini_time = cfg.multi.ini_time + fin_time = cfg.multi.fin_time + dt_save = cfg.multi.dt_save + CFL = cfg.multi.CFL + if cfg.multi.if_rand_param: + rho = exp( + random.uniform(log(0.001), log(10)) + ) # uniform number between 0.01 to 100 + nu = exp( + random.uniform(log(0.001), log(10)) + ) # uniform number between 0.01 to 100 + else: + rho = cfg.multi.rho + nu = cfg.multi.nu + + # t-coordinate + it_tot = ceil((fin_time - ini_time) / dt_save) + 1 + tc = jnp.arange(it_tot + 1) * dt_save + + @jax.jit + def evolve(u): + t = ini_time + tsave = t + steps = 0 + i_save = 0 + dt = 0.0 + uu = jnp.zeros([it_tot, u.shape[0]]) + uu = uu.at[0].set(u) + + cond_fun = lambda x: x[0] < fin_time + + def _body_fun(carry): + def _show(_carry): + u, tsave, i_save, uu = _carry + uu = uu.at[i_save].set(u) + tsave += dt_save + i_save += 1 + return (u, tsave, i_save, uu) + + t, tsave, steps, i_save, dt, u, uu = carry + + carry = (u, tsave, i_save, uu) + u, tsave, i_save, uu = lax.cond(t >= tsave, _show, _pass, carry) + + carry = (u, t, dt, steps, tsave) + u, t, dt, steps, tsave = lax.fori_loop(0, show_steps, simulation_fn, carry) + + return (t, tsave, steps, i_save, dt, u, uu) + + carry = t, tsave, steps, i_save, dt, u, uu + t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) + return uu.at[-1].set(u) + + @jax.jit + def simulation_fn(i, carry): + u, t, dt, steps, tsave = carry + dt = Courant_diff(dx, nu) * CFL + dt = jnp.min(jnp.array([dt, fin_time - t, tsave - t])) + + def _update(carry): + u, dt = carry + # preditor step for calculating t+dt/2-th time step + u_tmp = update(u, u, dt * 0.5) + # update using flux at t+dt/2-th time step + u = update(u, u_tmp, dt) + return u, dt + + carry = u, dt + u, dt = lax.cond(dt > 1.0e-8, _update, _pass, carry) + + t += dt + steps += 1 + return u, t, dt, steps, tsave + + @jax.jit + def update(u, u_tmp, dt): + # stiff part + u = Piecewise_Exact_Solution(u, dt) + # diffusion + f = flux(u_tmp) + u -= dt * dx_inv * (f[1 : cfg.multi.nx + 1] - f[0 : cfg.multi.nx]) + return u + + @jax.jit + def flux(u): + _u = bc( + u, dx, Ncell=cfg.multi.nx + ) # index 2 for _U is equivalent with index 0 for u + # 2nd-order diffusion flux + return -nu * (_u[2 : cfg.multi.nx + 3] - _u[1 : cfg.multi.nx + 2]) * dx_inv + + @jax.jit + def Piecewise_Exact_Solution(u, dt): # Piecewise_Exact_Solution method + # stiff equation + return 1.0 / (1.0 + jnp.exp(-rho * dt) * (1.0 - u) / u) + + u = init_multi( + xc, + numbers=cfg.multi.numbers, + k_tot=4, + init_key=cfg.multi.init_key, + if_norm=True, + ) + u = device_put(u) # putting variables in GPU (not necessary??) + + # vm_evolve = vmap(evolve, 0, 0) + # uu = vm_evolve(u) + vm_evolve = jax.pmap(jax.vmap(evolve, axis_name="j"), axis_name="i") + local_devices = jax.local_device_count() + uu = vm_evolve(u.reshape([local_devices, cfg.multi.numbers // local_devices, -1])) + + cwd = hydra.utils.get_original_cwd() + "/" + jnp.save( + cwd + cfg.multi.save + "/ReacDiff_Nu" + str(nu)[:5] + "_Rho" + str(rho)[:5], uu + ) + jnp.save(cwd + cfg.multi.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.multi.save + "/t_coordinate", tc) + + # reshape based on device count + uu = uu.reshape((-1, *uu.shape[2:])) + + cwd = hydra.utils.get_original_cwd() + "/" + Path(cwd + cfg.multi.save).mkdir(parents=True, exist_ok=True) + jnp.save( + cwd + cfg.multi.save + "ReacDiff_Nu" + str(nu)[:5] + "_Rho" + str(rho)[:5], uu + ) + jnp.save(cwd + cfg.multi.save + "/x_coordinate", xc) + jnp.save(cwd + cfg.multi.save + "/t_coordinate", tc) + + +if __name__ == "__main__": + main() diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_DarcyFlow2D.sh b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_DarcyFlow2D.sh new file mode 100644 index 0000000..c948681 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_DarcyFlow2D.sh @@ -0,0 +1,11 @@ +#! /bin/bash +nn=1 +key=2020 +while [ "$nn" -le 50 ]; do + CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_2D_multi_solution_Hydra.py +multi=config_2D.yaml ++multi.init_k\ +ey="$key" + nn=$(${nn} + 1) + key=$(${key} + 1) + echo "$nn" + echo "$key" +done diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_testset.sh b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_testset.sh new file mode 100644 index 0000000..513aadc --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_testset.sh @@ -0,0 +1,21 @@ +#! /bin/bash +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e0_Nu1e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e0_Nu2e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e0_Nu5e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e0_Nu1e1.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho2e0_Nu1e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho2e0_Nu2e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho2e0_Nu5e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho2e0_Nu1e1.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho5e0_Nu1e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho5e0_Nu2e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho5e0_Nu5e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho5e0_Nu1e1.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e1_Nu1e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e1_Nu2e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e1_Nu5e0.yaml +#CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e1_Nu1e1.yaml +CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e1_Nu5e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho5e0_Nu5e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho2e0_Nu5e-1.yaml +CUDA_VISIBLE_DEVICES='3' python3 reaction_diffusion_Hydra.py +args=Rho1e0_Nu5e-1.yaml diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_trainset.sh b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_trainset.sh new file mode 100644 index 0000000..21b8a51 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/run_trainset.sh @@ -0,0 +1,17 @@ +#! /bin/bash +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e0_Nu1e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e0_Nu2e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e0_Nu5e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho2e0_Nu1e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho2e0_Nu2e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho2e0_Nu5e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho5e0_Nu1e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho5e0_Nu2e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho5e0_Nu5e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e1_Nu1e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e1_Nu2e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e1_Nu5e0.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho2e0_Nu5e-1.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho5e0_Nu5e-1.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e0_Nu5e-1.yaml +CUDA_VISIBLE_DEVICES='0,1' python3 reaction_diffusion_multi_solution_Hydra.py +multi=Rho1e1_Nu5e-1.yaml diff --git a/pdebench/data_gen/data_gen_NLE/config/config.yaml b/pdebench/data_gen/data_gen_NLE/config/config.yaml new file mode 100644 index 0000000..89a88d5 --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/config/config.yaml @@ -0,0 +1,16 @@ +defaults: + - _self_ + - override hydra/hydra_logging: disabled + - override hydra/job_logging: disabled + +hydra: + output_subdir: null + run: + dir: . + +args: + type: "ReacDiff" # "advection"/"ReacDiff"/"burgers"/"CFD" + dim: 1 + bd: "periodic" + nbatch: 1000 + savedir: "./save/CFD/" diff --git a/pdebench/data_gen/data_gen_NLE/utils.py b/pdebench/data_gen/data_gen_NLE/utils.py new file mode 100644 index 0000000..c9f63ce --- /dev/null +++ b/pdebench/data_gen/data_gen_NLE/utils.py @@ -0,0 +1,1608 @@ +from __future__ import annotations + +import math as mt +from functools import partial + +import jax +import jax.numpy as jnp +import numpy as np +from jax import jit, lax, nn, random, scipy, vmap + +# if double precision +# from jax.config import config +# config.update("jax_enable_x64", True) + + +def init(xc, mode="sin", u0=1.0, du=0.1): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + modes = ["sin", "sinsin", "Gaussian", "react", "possin"] + assert mode in modes, "mode is not defined!!" + if mode == "sin": # sinusoidal wave + u = u0 * jnp.sin((xc + 1.0) * jnp.pi) + elif mode == "sinsin": # sinusoidal wave + u = jnp.sin((xc + 1.0) * jnp.pi) + du * jnp.sin((xc + 1.0) * jnp.pi * 8.0) + elif mode == "Gaussian": # for diffusion check + t0 = 0.01 + u = jnp.exp(-(xc**2) * jnp.pi / (4.0 * t0)) / jnp.sqrt(2.0 * t0) + elif mode == "react": # for reaction-diffusion eq. + logu = -0.5 * (xc - jnp.pi) ** 2 / (0.25 * jnp.pi) ** 2 + u = jnp.exp(logu) + elif mode == "possin": # sinusoidal wave + u = u0 * jnp.abs(jnp.sin((xc + 1.0) * jnp.pi)) + return u + + +@partial(jit, static_argnums=(1, 2, 3, 4)) +def init_multi( + xc, numbers=10000, k_tot=8, init_key=2022, num_choise_k=2, if_norm=False +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + + def _pass(carry): + return carry + + def select_A(carry): + def _func(carry): + return jnp.abs(carry) + + cond, value = carry + value = lax.cond(cond == 1, _func, _pass, value) + return cond, value + + def select_W(carry): + def _window(carry): + xx, val, xL, xR, trns = carry + val = 0.5 * (jnp.tanh((xx - xL) / trns) - jnp.tanh((xx - xR) / trns)) + return xx, val, xL, xR, trns + + cond, value, xx, xL, xR, trns = carry + + carry = xx, value, xL, xR, trns + xx, value, xL, xR, trns = lax.cond(cond == 1, _window, _pass, carry) + return cond, value, xx, xL, xR, trns + + def normalize(carry): + def _norm(carry): + u = carry + u -= jnp.min(u, axis=1, keepdims=True) # positive value + u /= jnp.max(u, axis=1, keepdims=True) # normalize + return u + + cond, u = carry + u = lax.cond(cond is True, _norm, _pass, u) + return cond, u + + key = random.PRNGKey(init_key) + + selected = random.randint( + key, shape=[numbers, num_choise_k], minval=0, maxval=k_tot + ) + selected = nn.one_hot(selected, k_tot, dtype=int).sum(axis=1) + kk = jnp.pi * 2.0 * jnp.arange(1, k_tot + 1) * selected / (xc[-1] - xc[0]) + amp = random.uniform(key, shape=[numbers, k_tot, 1]) + + key, subkey = random.split(key) + + phs = 2.0 * jnp.pi * random.uniform(key, shape=[numbers, k_tot, 1]) + _u = amp * jnp.sin(kk[:, :, jnp.newaxis] * xc[jnp.newaxis, jnp.newaxis, :] + phs) + _u = jnp.sum(_u, axis=1) + + # perform absolute value function + cond = random.choice(key, 2, p=jnp.array([0.9, 0.1]), shape=([numbers])) + carry = (cond, _u) + + cond, _u = vmap(select_A, 0, 0)(carry) + sgn = random.choice(key, a=jnp.array([1, -1]), shape=([numbers, 1])) + _u *= sgn # random flip of signature + + # perform window function + key, subkey = random.split(key) + cond = random.choice(key, 2, p=jnp.array([0.9, 0.1]), shape=([numbers])) + _xc = jnp.repeat(xc[None, :], numbers, axis=0) + mask = jnp.ones_like(_xc) + xL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + xR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + trns = 0.01 * jnp.ones_like(cond) + carry = cond, mask, _xc, xL, xR, trns + cond, mask, _xc, xL, xR, trns = vmap(select_W, 0, 0)(carry) + + _u *= mask + + carry = if_norm, _u + _, _u = normalize( + carry + ) # normalize value between [0, 1] for reaction-diffusion eq. + + return _u + + +def init_multi_2DRand(xc, yc, numbers=10000, init_key=2022, k_tot=4, duMx=1.0e1): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + def _pass(carry): + return carry + + def select_W(carry): + def _window(carry): + xx, yy, val, xL, xR, yL, yR, trns = carry + x_win = 0.5 * (jnp.tanh((xx - xL) / trns) - jnp.tanh((xx - xR) / trns)) + y_win = 0.5 * (jnp.tanh((yy - yL) / trns) - jnp.tanh((yy - yR) / trns)) + val = x_win[:, None] * y_win[None, :] + return xx, yy, val, xL, xR, yL, yR, trns + + cond, value, xx, yy, xL, xR, yL, yR, trns = carry + + carry = xx, yy, value, xL, xR, yL, yR, trns + xx, yy, value, xL, xR, yL, yR, trns = lax.cond(cond == 1, _window, _pass, carry) + return cond, value, xx, xL, xR, trns + + def __create_2DRand_init(u0, delu): + nx, ny = xc.shape[0], yc.shape[0] + + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + + qLx = dx * nx + qLy = dy * ny + + # random field + u = jnp.zeros([nx, ny]) + + key = random.PRNGKey(init_key) + kx0 = jnp.pi * 2.0 / qLx + ky0 = jnp.pi * 2.0 / qLy + + for j in range(-k_tot, k_tot + 1): + ky = ky0 * j # from 1 to k_tot + for i in range(-k_tot, k_tot + 1): + kx = kx0 * i # from 1 to k_tot + if i * j == 0: # avoiding uniform velocity + continue + # random phase + key, subkey = random.split(key) + phs = 2.0 * jnp.pi * random.uniform(key, shape=[1]) # (vi, k) + + uk = 1.0 / jnp.sqrt(jnp.sqrt(kx**2 + ky**2)) + kdx = kx * xc[:, None] + ky * yc[None, :] + u += uk * jnp.sin(kdx + phs) + + # renormalize total velocity + return u0 + delu * u / jnp.abs(u).mean() + + key = random.PRNGKey(init_key) + u0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=duMx) + key, subkey = random.split(key) + delu = random.uniform(key, shape=([numbers, 1]), minval=1.0e-2, maxval=0.5) + u = jax.vmap(__create_2DRand_init, axis_name="i")(u0, delu) + + # perform window function + key, subkey = random.split(key) + cond = random.choice(key, 2, p=jnp.array([0.5, 0.5]), shape=([numbers])) + mask = jnp.ones([numbers, xc.shape[0], yc.shape[0]]) + xL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + xR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + key, subkey = random.split(key) + yL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + yR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + _xc = jnp.repeat(xc[None, :], numbers, axis=0) + _yc = jnp.repeat(xc[None, :], numbers, axis=0) + trns = 0.01 * jnp.ones_like(cond) + carry = cond, mask, _xc, _yc, xL, xR, yL, yR, trns + cond, mask, _xc, xL, xR, trns = vmap(select_W, 0, 0)(carry) + + u = u * mask + return u + u0[:, :, None] * (1.0 - mask) + + +def init_HD( + u, + xc, + yc, + zc, + mode="shocktube1", + direc="x", + init_key=2022, + M0=0.1, + dk=1, + gamma=0.1666666667, +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + modes = [ + "shocktube0", + "shocktube1", + "shocktube2", + "shocktube3", + "shocktube4", + "shocktube5", + "shocktube6", + "shocktube7", + "2D-shock", + "OTVortex", + "KHI", + "turbulence", + "sound_wave", + "c_discon", + "BlastWave", + ] + assert mode in modes, "mode is not defined!!" + + _, nx, ny, nz = u.shape + + if mode[:-1] == "shocktube": # shock tube + if direc == "x": + iX, iY, iZ = 1, 2, 3 + Ncell = nx + _u = jnp.zeros_like(u) + elif direc == "y": + iX, iY, iZ = 2, 3, 1 + Ncell = ny + _u = jnp.transpose(u, (0, 2, 3, 1)) + if direc == "z": + iX, iY, iZ = 3, 1, 2 + Ncell = nz + _u = jnp.transpose(u, (0, 3, 1, 2)) + + if mode[-1] == "0": # test 0 for viscosity + nx0 = int(0.5 * Ncell) + uL = [1.0, 0.75, 0.2, -0.3, 1.0] + uR = [0.125, 0.0, 0.1, 0.9, 0.1] + elif mode[-1] == "1": # test 1 + nx0 = int(0.3 * Ncell) + uL = [1.0, 0.75, 0.0, 0.0, 1.0] + uR = [0.125, 0.0, 0.0, 0.0, 0.1] + elif mode[-1] == "2": # test 2 + nx0 = int(0.5 * Ncell) + uL = [1.0, -2.0, 0.0, 0.0, 0.4] + uR = [1.0, 2.0, 0.0, 0.0, 0.4] + elif mode[-1] == "3": # test 3 + nx0 = int(0.5 * Ncell) + uL = [1.0, 0.0, 0.0, 0.0, 1.0e3] + uR = [1.0, 0.0, 0.0, 0.0, 0.01] + elif mode[-1] == "4": # test 4 + nx0 = int(0.4 * Ncell) + uL = [5.99924, 19.5975, 0.0, 0.0, 460.894] + uR = [5.99242, -6.19633, 0.0, 0.0, 46.095] + elif mode[-1] == "5": # test 5 + nx0 = int(0.8 * Ncell) + uL = [1.0, -19.59745, 0.0, 0.0, 1.0e3] + uR = [1.0, -19.59745, 0.0, 0.0, 0.01] + elif mode[-1] == "6": # test 6 + nx0 = int(0.5 * Ncell) + uL = [1.4, 0.0, 0.0, 0.0, 1.0] + uR = [1.0, 0.0, 0.0, 0.0, 1.0] + elif mode[-1] == "7": # test 7 + nx0 = int(0.5 * Ncell) + uL = [1.4, 0.1, 0.0, 0.0, 1.0] + uR = [1.0, 0.1, 0.0, 0.0, 1.0] + + # left + _u = _u.loc[0, :nx0].set(uL[0]) + _u = _u.loc[iX, :nx0].set(uL[1]) + _u = _u.loc[iY, :nx0].set(uL[2]) + _u = _u.loc[iZ, :nx0].set(uL[3]) + _u = _u.loc[4, :nx0].set(uL[4]) + # right + _u = _u.loc[0, nx0:].set(uR[0]) + _u = _u.loc[iX, nx0:].set(uR[1]) + _u = _u.loc[iY, nx0:].set(uR[2]) + _u = _u.loc[iZ, nx0:].set(uR[3]) + _u = _u.loc[4, nx0:].set(uR[4]) + + if direc == "x": + u = _u + elif direc == "y": + u = jnp.transpose(_u, (0, 3, 1, 2)) + elif direc == "z": + u = jnp.transpose(_u, (0, 2, 3, 1)) + elif mode == "2D-shock": # shock tube + u1 = [0.5, 0.0, 0.0, 0.0, 0.1] + u2 = [0.1, 0.0, 1.0, 0.0, 1.0] + u3 = [0.1, 1.0, 0.0, 0.0, 1.0] + u4 = [0.1, 0.0, 0.0, 0.0, 0.01] + + # left-bottom + u = u.loc[0, : nx // 2, : ny // 2].set(u1[0]) + u = u.loc[1, : nx // 2, : ny // 2].set(u1[1]) + u = u.loc[2, : nx // 2, : ny // 2].set(u1[2]) + u = u.loc[3, : nx // 2, : ny // 2].set(u1[3]) + u = u.loc[4, : nx // 2, : ny // 2].set(u1[4]) + # right-bottom + u = u.loc[0, nx // 2 :, : ny // 2].set(u2[0]) + u = u.loc[1, nx // 2 :, : ny // 2].set(u2[1]) + u = u.loc[2, nx // 2 :, : ny // 2].set(u2[2]) + u = u.loc[3, nx // 2 :, : ny // 2].set(u2[3]) + u = u.loc[4, nx // 2 :, : ny // 2].set(u2[4]) + # left-top + u = u.loc[0, : nx // 2, ny // 2 :].set(u3[0]) + u = u.loc[1, : nx // 2, ny // 2 :].set(u3[1]) + u = u.loc[2, : nx // 2, ny // 2 :].set(u3[2]) + u = u.loc[3, : nx // 2, ny // 2 :].set(u3[3]) + u = u.loc[4, : nx // 2, ny // 2 :].set(u3[4]) + # right-top + u = u.loc[0, nx // 2 :, ny // 2 :].set(u4[0]) + u = u.loc[1, nx // 2 :, ny // 2 :].set(u4[1]) + u = u.loc[2, nx // 2 :, ny // 2 :].set(u4[2]) + u = u.loc[3, nx // 2 :, ny // 2 :].set(u4[3]) + u = u.loc[4, nx // 2 :, ny // 2 :].set(u4[4]) + + elif mode == "OTVortex": # shock tube + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + qLx = dx * xc.shape[0] + qLy = dy * yc.shape[0] + _xc = jnp.zeros([xc.shape[0] + 4]) + _yc = jnp.zeros([yc.shape[0] + 4]) + _xc = _xc.at[2:-2].set(xc) + _yc = _yc.at[2:-2].set(yc) + _xc = _xc.at[:2].set(jnp.array([-2 * dx, -dx])) + _yc = _yc.at[:2].set(jnp.array([-2 * dy, -dy])) + _xc = _xc.at[-2:].set(jnp.array([xc[-1] + dx, xc[-1] + 2.0 * dx])) + _yc = _yc.at[-2:].set(jnp.array([yc[-1] + dy, yc[-1] + 2.0 * dy])) + + u = u.loc[0].add(gamma**2) + u = u.loc[1].set(-jnp.sin(2.0 * jnp.pi * _yc[None, :, None] / qLy)) + u = u.loc[2].set(jnp.sin(2.0 * jnp.pi * _xc[:, None, None] / qLx)) + u = u.loc[3].add(0.0) + u = u.loc[4].add(gamma) + + elif mode == "KHI": # Kelvin-Helmholtz instability + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + # gamma = 1.666666666666667 + # k = 1. # moved to the external input + d0_u = 2.0 / (dk + 1.0) + d0_d = dk * d0_u + d0 = 0.5 * (d0_u + d0_d) + # M0 = 0.1 # Mach number # moved to external input + ux = 1.0 + cs = ux / M0 + # ux = 0.1 * cs # << cs + p0 = cs**2 * d0 / gamma + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + qLx = dx * nx + qLy = dy * ny + kk = 4.0 # wave number + kx = kk * 2.0 * jnp.pi / qLx + dl = 5.0e-3 * qLy + + bound = 0.5 * qLy + dl * jnp.sin(kx * xc) # assuming yL = 0 + + vx = jnp.zeros([nx, ny, nz]) + dd = jnp.zeros([nx, ny, nz]) + for i in range(nx): + _vx = jnp.where(yc > bound[i], ux, -ux) + _dd = jnp.where(yc > bound[i], d0_u, d0_d) + vx = vx.at[i, :, :].set(_vx[:, None]) + dd = dd.at[i, :, :].set(_dd[:, None]) + + u = u.loc[0, 2:-2, 2:-2, 2:-2].set(dd) + u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) + u = u.loc[2].set(0.0) + u = u.loc[3].add(0.0) + u = u.loc[4].add(p0) + + elif mode == "turbulence": # 3D decaying turbulence + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + d0 = 1.0 + cs = 1.0 / M0 + u0 = 1.0 # fixed + p0 = cs**2 * d0 / gamma + + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + dz = zc[1] - zc[0] + qLx = dx * nx + qLy = dy * ny + qLz = dz * nz + + # random velocity field + k_tot = 3 + vx, vy, vz = ( + np.zeros([nx, ny, nz]), + np.zeros([nx, ny, nz]), + np.zeros([nx, ny, nz]), + ) + + key = random.PRNGKey(init_key) + + kx0 = jnp.pi * 2.0 / qLx + ky0 = jnp.pi * 2.0 / qLy + kz0 = jnp.pi * 2.0 / qLz + + for k in range(-k_tot, k_tot + 1): + kz = kz0 * k # from 1 to k_tot + for j in range(-k_tot, k_tot + 1): + ky = ky0 * j # from 1 to k_tot + for i in range(-k_tot, k_tot + 1): + kx = kx0 * i # from 1 to k_tot + if i * j * k == 0: # avoiding uniform velocity + continue + # random phase + key, subkey = random.split(key) + phs = 2.0 * jnp.pi * random.uniform(key, shape=[3]) # (vi, k) + + uk = 1.0 / jnp.sqrt(kx**2 + ky**2 + kz**2) + kdx = ( + kx * xc[:, None, None] + + ky * yc[None, :, None] + + kz * zc[None, None, :] + ) + vx += uk * jnp.sin(kdx + phs[0]) + vy += uk * jnp.sin(kdx + phs[1]) + vz += uk * jnp.sin(kdx + phs[2]) + + del (kdx, uk, phs) + + # Helmholtz decomposition to subtract expansion: k.vk + dfx, dfy, dfz = 1.0 / qLx, 1.0 / qLy, 1.0 / qLz + fx = dfx * (np.arange(nx) - 1.0 - nx // 2) + fy = dfy * (np.arange(ny) - 1.0 - ny // 2) + fz = dfz * (np.arange(nz) - 1.0 - nz // 2) + + vkx = np.fft.fftn(vx) * dx * dy * dz + vky = np.fft.fftn(vy) * dx * dy * dz + vkz = np.fft.fftn(vz) * dx * dy * dz + + # shift to kxi=0 is at the center + vkx = np.fft.fftshift(vkx) + vky = np.fft.fftshift(vky) + vkz = np.fft.fftshift(vkz) + + # for k in range(nz): + # for j in range(ny): + # for i in range(nx): + # ff = (fx[i]**2 + fy[j]**2 + fz[k]**2) + # fi = np.where(ff > 1.e-8, 1./ff, 0.) + # # subtract expansion k.vk + # fdv = fx[i] * vkx[i, j, k] + fy[j] * vky[i, j, k] + fz[k] * vkz[i, j, k] + # vkx -= fdv * fx[i] * fi + # vky -= fdv * fy[j] * fi + # vkz -= fdv * fz[k] * fi + + fi = fx[:, None, None] ** 2 + fy[None, :, None] ** 2 + fz[None, None, :] ** 2 + fi = np.where(fi > 1.0e-8, 1.0 / fi, 0.0) + + fdv = ( + fx[:, None, None] * vkx + fy[None, :, None] * vky + fz[None, None, :] * vkz + ) * fi + vkx -= fdv * fx[:, None, None] + vky -= fdv * fy[None, :, None] + vkz -= fdv * fz[None, None, :] + del (fi, fdv) + + # shift back to original order + vkx = np.fft.ifftshift(vkx) + vky = np.fft.ifftshift(vky) + vkz = np.fft.ifftshift(vkz) + + # inverse FFT + vx = np.fft.ifftn(vkx).real * dfx * dfy * dfz + vy = np.fft.ifftn(vky).real * dfx * dfy * dfz + vz = np.fft.ifftn(vkz).real * dfx * dfy * dfz + + # renormalize total velocity + vtot = np.sqrt(vx**2 + vy**2 + vz**2).mean() + vx *= u0 / vtot + vy *= u0 / vtot + vz *= u0 / vtot + + u = u.loc[0].set(d0) + u = u.loc[1, 2:-2, 2:-2, 2:-2].set(jnp.array(vx)) + u = u.loc[2, 2:-2, 2:-2, 2:-2].set(jnp.array(vy)) + u = u.loc[3, 2:-2, 2:-2, 2:-2].add(jnp.array(vz)) + u = u.loc[4].add(p0) + + elif mode == "BlastWave": # Kelvin-Helmholtz instability + """ Stone Gardiner 2009 without B """ + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + db = 1.0 + pb = 0.1 + + pc = 1.0e2 # central region + + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + dz = zc[1] - zc[0] + qLx = dx * nx + qLy = dy * ny + qLz = dz * nz + qL = (qLx + qLy + qLz) / 3.0 + + # p0 = jnp.ones([nx, ny, nz]) * pb + RR = jnp.sqrt( + (xc[:, None, None] - xc[nx // 2]) ** 2 + + (yc[None, :, None] - yc[ny // 2]) ** 2 + + (zc[None, None, :] - zc[nz // 2]) ** 2 + ) + p0 = jnp.where(0.05 * qL < RR, pb, pc) + # for k in range(nz): + # for j in range(ny): + # for i in range(nx): + # RR = jnp.sqrt((xc[i] - 0.5 * qLx)**2 + (yc[j] - 0.5 * qLy)**2 + (zc[k] - 0.5 * qLz)**2) + # if RR < 0.1 * qL: + # p0 = p0.at[i,j,k].set(pc) + + u = u.loc[0].set(db) + u = u.loc[1].set(0.0) + u = u.loc[2].set(0.0) + u = u.loc[3].set(0.0) + u = u.loc[4, 2:-2, 2:-2, 2:-2].set(p0) + + elif mode == "sound_wave": # sound wave + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + gamma = 1.666666666666667 + d0 = 1.0 + cs = 2.0 + p0 = cs**2 * d0 / gamma + if direc == "x": + iX, iY, iZ = 1, 2, 3 + XC = xc + qL = (xc[1] - xc[0]) * nx + _u = jnp.zeros_like(u) + elif direc == "y": + iX, iY, iZ = 2, 3, 1 + XC = yc + qL = (yc[1] - yc[0]) * ny + _u = jnp.transpose(u, (0, 2, 3, 1)) + if direc == "z": + iX, iY, iZ = 3, 1, 2 + XC = zc + qL = (zc[1] - zc[0]) * nz + _u = jnp.transpose(u, (0, 3, 1, 2)) + + kk = 2.0 * jnp.pi / qL + _u = _u.loc[0, 2:-2].set(d0 * (1.0 + 1.0e-3 * jnp.sin(kk * XC[:, None, None]))) + _u = _u.loc[iX].set((_u[0] - d0) * cs / d0) + _u = _u.loc[4].set(p0 + cs**2 * (_u[0] - d0)) + + if direc == "x": + u = _u + elif direc == "y": + u = jnp.transpose(_u, (0, 3, 1, 2)) + elif direc == "z": + u = jnp.transpose(_u, (0, 2, 3, 1)) + + elif mode == "c_discon": # tangent discontinuity + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + d0 = 1.0 + p0 = 1.0 + vy0 = 0.1 + if direc == "x": + iX, iY, iZ = 1, 2, 3 + XC = xc + qL = (xc[1] - xc[0]) * nx + _u = jnp.zeros_like(u) + elif direc == "y": + iX, iY, iZ = 2, 3, 1 + XC = yc + qL = (yc[1] - yc[0]) * ny + _u = jnp.transpose(u, (0, 2, 3, 1)) + if direc == "z": + iX, iY, iZ = 3, 1, 2 + XC = zc + qL = (zc[1] - zc[0]) * nz + _u = jnp.transpose(u, (0, 3, 1, 2)) + + _u = _u.loc[0].set(d0) + _u = _u.loc[iY, 2:-2].set( + vy0 * scipy.special.erf(0.5 * XC[:, None, None] / jnp.sqrt(0.1)) + ) + _u = _u.loc[4].set(p0) + + if direc == "x": + u = _u + elif direc == "y": + u = jnp.transpose(_u, (0, 3, 1, 2)) + elif direc == "z": + u = jnp.transpose(_u, (0, 2, 3, 1)) + + return u + + +@partial(jit, static_argnums=(3, 4, 5, 6, 7, 8, 9)) +def init_multi_HD( + xc, + yc, + zc, + numbers=10000, + k_tot=10, + init_key=2022, + num_choise_k=2, + if_renorm=False, + umax=1.0e4, + umin=1.0e-8, +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + + def _pass(carry): + return carry + + def select_A(carry): + def _func(carry): + return jnp.abs(carry) + + cond, value = carry + value = lax.cond(cond == 1, _func, _pass, value) + return cond, value + + def select_W(carry): + def _window(carry): + xx, val, xL, xR, trns = carry + val = 0.5 * (jnp.tanh((xx - xL) / trns) - jnp.tanh((xx - xR) / trns)) + return xx, val, xL, xR, trns + + cond, value, xx, xL, xR, trns = carry + + carry = xx, value, xL, xR, trns + xx, value, xL, xR, trns = lax.cond(cond == 1, _window, _pass, carry) + return cond, value, xx, xL, xR, trns + + def renormalize(carry): + def _norm(carry): + u, key = carry + u -= jnp.min(u, axis=1, keepdims=True) # positive value + u /= jnp.max(u, axis=1, keepdims=True) # normalize + + key, subkey = random.split(key) + m_val = random.uniform( + key, shape=[numbers], minval=mt.log(umin), maxval=mt.log(umax) + ) + m_val = jnp.exp(m_val) + key, subkey = random.split(key) + b_val = random.uniform( + key, shape=[numbers], minval=mt.log(umin), maxval=mt.log(umax) + ) + b_val = jnp.exp(b_val) + return u * m_val[:, None] + b_val[:, None], key + + cond, u, key = carry + carry = u, key + u, key = lax.cond(cond is True, _norm, _pass, carry) + return cond, u, key + + assert yc.shape[0] == 1 and zc.shape[0] == 1, "ny and nz is assumed to be 1!!" # noqa: PT018 + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + key = random.PRNGKey(init_key) + + selected = random.randint( + key, shape=[numbers, num_choise_k], minval=0, maxval=k_tot + ) + selected = nn.one_hot(selected, k_tot, dtype=int).sum(axis=1) + kk = jnp.pi * 2.0 * jnp.arange(1, k_tot + 1) * selected / (xc[-1] - xc[0]) + amp = random.uniform(key, shape=[numbers, k_tot, 1]) + + key, subkey = random.split(key) + + phs = 2.0 * jnp.pi * random.uniform(key, shape=[numbers, k_tot, 1]) + _u = amp * jnp.sin(kk[:, :, jnp.newaxis] * xc[jnp.newaxis, jnp.newaxis, :] + phs) + _u = jnp.sum(_u, axis=1) + + # perform absolute value function + cond = random.choice(key, 2, p=jnp.array([0.9, 0.1]), shape=([numbers])) + carry = (cond, _u) + + cond, _u = vmap(select_A, 0, 0)(carry) + sgn = random.choice(key, a=jnp.array([1, -1]), shape=([numbers, 1])) + _u *= sgn # random flip of signature + + # perform window function + key, subkey = random.split(key) + cond = random.choice(key, 2, p=jnp.array([0.5, 0.5]), shape=([numbers])) + _xc = jnp.repeat(xc[None, :], numbers, axis=0) + mask = jnp.ones_like(_xc) + xL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + xR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + trns = 0.01 * jnp.ones_like(cond) + carry = cond, mask, _xc, xL, xR, trns + cond, mask, _xc, xL, xR, trns = vmap(select_W, 0, 0)(carry) + + _u *= mask + + carry = if_renorm, _u, key + _, _u, _ = renormalize(carry) # renormalize value between a given values + + return _u[..., None, None] + + +# @partial(jit, static_argnums=(3, 4, 5, 6)) +def init_multi_HD_shock( + xc, yc, zc, numbers=10000, init_key=2022, umax=1.0e4, umin=1.0e-8 +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + assert yc.shape[0] == 1 and zc.shape[0] == 1, "ny and nz is assumed to be 1!!" # noqa: PT018 + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + def select_var(carry): + def _func(carry): + vmin, vmax = carry + return jnp.log(vmin), jnp.log(vmax) + + def _pass(carry): + return carry + + vmin, vmax = carry + vmin, vmax = lax.cond(vmin > 0.0, _func, _pass, carry) + return vmin, vmax + + nx = xc.shape[0] + + carry = umin, umax + u_min, u_max = select_var(carry) + + key = random.PRNGKey(init_key) + QLs = random.uniform(key, shape=([numbers, 1]), minval=u_min, maxval=u_max) + QLs = jnp.exp(QLs) + key, subkey = random.split(key) + QRs = random.uniform(key, shape=([numbers, 1]), minval=u_min, maxval=u_max) + QRs = jnp.exp(QRs) + + nx0s = nx * random.uniform(key, shape=([numbers, 1]), minval=0.25, maxval=0.75) + nx0s = nx0s.astype(int) + + u = jnp.arange(xc.shape[0]) + u = jnp.tile(u, (numbers, 1)) + + u = jax.vmap(jnp.where, axis_name="i")(u < nx0s, QLs, QRs) + return u[..., None, None] + + +# @partial(jit, static_argnums=(4, 5, 6, 7, 8, 9)) +def init_multi_HD_KH( + u, + xc, + yc, + zc, + numbers=10000, + init_key=2022, + M0=0.1, + dkMx=2.0, + kmax=4.0, + gamma=1.666666667, +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + assert zc.shape[0] == 1, "nz is assumed to be 1!!" + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + def __create_KH_init(u, dk, kk): + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + d0_u = 2.0 / (dk + 1.0) + d0_d = dk * d0_u + d0 = 0.5 * (d0_u + d0_d) + ux = 1.0 + cs = ux / M0 + p0 = cs**2 * d0 / gamma + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + qLx = dx * nx + qLy = dy * ny + kx = kk * 2.0 * jnp.pi / qLx + dl = 5.0e-3 * qLy + # (numbers, nx) + bound = 0.5 * qLy + dl * jnp.sin(kx * xc) # assuming yL = 0 + + vx = jnp.zeros([nx, ny, nz]) + dd = jnp.zeros([nx, ny, nz]) + for i in range(nx): + _vx = jnp.where(yc > bound[i], ux, -ux) + _dd = jnp.where(yc > bound[i], d0_u, d0_d) + vx = vx.at[i, :, :].set(_vx[:, None]) + dd = dd.at[i, :, :].set(_dd[:, None]) + + u = u.loc[0, 2:-2, 2:-2, 2:-2].set(dd) + u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) + u = u.loc[2].set(0.0) + u = u.loc[3].add(0.0) + return u.loc[4].add(p0) + + # create random density ratio + key = random.PRNGKey(init_key) + dk = random.uniform(key, shape=([numbers, 1]), minval=1.0 / dkMx, maxval=dkMx) + # create random wave-numbers + key, subkey = random.split(key) + kk = random.randint(key, shape=([numbers, 1]), minval=1, maxval=kmax) + return jax.vmap(__create_KH_init, axis_name="i")(u, dk, kk) + + +# @partial(jit, static_argnums=(4, 5, 6, 7, 8)) +def init_multi_HD_2DTurb( + u, xc, yc, zc, numbers=10000, init_key=2022, M0=0.1, k_tot=4.0, gamma=1.666666667 +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + assert zc.shape[0] == 1, "nz is assumed to be 1!!" + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + def __create_2DTurb_init(u, keys): + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + d0 = 1.0 + cs = 1.0 / M0 + u0 = 1.0 # fixed + p0 = cs**2 * d0 / gamma + + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + + qLx = dx * nx + qLy = dy * ny + + # random velocity field + vx, vy = jnp.zeros([nx, ny, nz]), jnp.zeros([nx, ny, nz]) + + key = random.PRNGKey(keys) + + kx0 = jnp.pi * 2.0 / qLx + ky0 = jnp.pi * 2.0 / qLy + + for j in range(-k_tot, k_tot + 1): + ky = ky0 * j # from 1 to k_tot + for i in range(-k_tot, k_tot + 1): + kx = kx0 * i # from 1 to k_tot + if i * j == 0: # avoiding uniform velocity + continue + # random phase + key, subkey = random.split(key) + phs = 2.0 * jnp.pi * random.uniform(key, shape=[2]) # (vi, k) + + uk = 1.0 / jnp.sqrt(jnp.sqrt(kx**2 + ky**2)) + kdx = kx * xc[:, None, None] + ky * yc[None, :, None] + vx += uk * jnp.sin(kdx + phs[0]) + vy += uk * jnp.sin(kdx + phs[1]) + + del (kdx, uk, phs) + + # Helmholtz decomposition to subtract expansion: k.vk + dfx, dfy = 1.0 / qLx, 1.0 / qLy + fx = dfx * (jnp.arange(nx) - 1.0 - nx // 2) + fy = dfy * (jnp.arange(ny) - 1.0 - ny // 2) + + vkx = jnp.fft.fftn(vx) * dx * dy + vky = jnp.fft.fftn(vy) * dx * dy + + # shift to kxi=0 is at the center + vkx = jnp.fft.fftshift(vkx) + vky = jnp.fft.fftshift(vky) + + fi = fx[:, None, None] ** 2 + fy[None, :, None] ** 2 + fi = jnp.where(fi > 1.0e-8, 1.0 / fi, 0.0) + + fdv = (fx[:, None, None] * vkx + fy[None, :, None] * vky) * fi + vkx -= fdv * fx[:, None, None] + vky -= fdv * fy[None, :, None] + del (fi, fdv) + + # shift back to original order + vkx = jnp.fft.ifftshift(vkx) + vky = jnp.fft.ifftshift(vky) + + # inverse FFT + vx = jnp.fft.ifftn(vkx).real * dfx * dfy + vy = jnp.fft.ifftn(vky).real * dfx * dfy + + # renormalize total velocity + vtot = jnp.sqrt(vx**2 + vy**2).mean() + vx *= u0 / vtot + vy *= u0 / vtot + + u = u.loc[0].set(d0) + u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) + u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) + return u.loc[4].add(p0) + + key = random.PRNGKey(init_key) + keys = random.randint( + key, + [ + numbers, + ], + minval=0, + maxval=10000000, + ) + return jax.vmap(__create_2DTurb_init, axis_name="i")(u, keys) + + +def init_multi_HD_2DRand( + u, + xc, + yc, + zc, + numbers=10000, + init_key=2022, + M0=0.1, + k_tot=4.0, + gamma=1.666666667, + dMx=1.0e1, + TMx=1.0e1, +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + assert zc.shape[0] == 1, "nz is assumed to be 1!!" + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + def _pass(carry): + return carry + + def select_W(carry): + def _window(carry): + xx, yy, val, xL, xR, yL, yR, trns = carry + x_win = 0.5 * (jnp.tanh((xx - xL) / trns) - jnp.tanh((xx - xR) / trns)) + y_win = 0.5 * (jnp.tanh((yy - yL) / trns) - jnp.tanh((yy - yR) / trns)) + val = x_win[:, None] * y_win[None, :] + return xx, yy, val, xL, xR, yL, yR, trns + + cond, value, xx, yy, xL, xR, yL, yR, trns = carry + + carry = xx, yy, value, xL, xR, yL, yR, trns + xx, yy, value, xL, xR, yL, yR, trns = lax.cond(cond == 1, _window, _pass, carry) + return cond, value, xx, yy, xL, xR, yL, yR, trns + + def __create_2DRand_init(u, d0, T0, delD, delP, keys): + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + + p0 = d0 * T0 + cs = jnp.sqrt(T0 * gamma) + u0 = M0 * cs + + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + + qLx = dx * nx + qLy = dy * ny + + # random velocity field + d, p, vx, vy = ( + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + ) + + key = random.PRNGKey(keys) + kx0 = jnp.pi * 2.0 / qLx + ky0 = jnp.pi * 2.0 / qLy + + for j in range(-k_tot, k_tot + 1): + ky = ky0 * j # from 1 to k_tot + for i in range(-k_tot, k_tot + 1): + kx = kx0 * i # from 1 to k_tot + if i * j == 0: # avoiding uniform velocity + continue + # random phase + key, subkey = random.split(key) + phs = 2.0 * jnp.pi * random.uniform(key, shape=[4]) # (vi, k) + + uk = 1.0 / jnp.sqrt(jnp.sqrt(kx**2 + ky**2)) + kdx = kx * xc[:, None, None] + ky * yc[None, :, None] + vx += uk * jnp.sin(kdx + phs[0]) + vy += uk * jnp.sin(kdx + phs[1]) + p += uk * jnp.sin(kdx + phs[2]) + d += uk * jnp.sin(kdx + phs[3]) + + del (kdx, uk, phs) + + # renormalize total velocity + vtot = jnp.sqrt(vx**2 + vy**2).mean() + vx *= u0 / vtot + vy *= u0 / vtot + # d = d0 + delD * d / jnp.abs(d).mean() + # p = p0 + delP * p / jnp.abs(p).mean() + d = d0 * (1.0 + delD * d / jnp.abs(d).mean()) + p = p0 * (1.0 + delP * p / jnp.abs(p).mean()) + + u = u.loc[0, 2:-2, 2:-2, 2:-2].set(d) + u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) + u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) + return u.loc[4, 2:-2, 2:-2, 2:-2].set(p) + + key = random.PRNGKey(init_key) + d0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=dMx) + key, subkey = random.split(key) + delD = random.uniform(key, shape=([numbers, 1]), minval=1.0e-2, maxval=0.2) + key, subkey = random.split(key) + T0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=TMx) + key, subkey = random.split(key) + delP = random.uniform(key, shape=([numbers, 1]), minval=1.0e-2, maxval=0.2) + + key, subkey = random.split(key) + keys = random.randint( + key, + shape=( + [ + numbers, + ] + ), + minval=0, + maxval=10000000, + ) + u = jax.vmap(__create_2DRand_init, axis_name="i")(u, d0, T0, delD, delP, keys) + + # perform window function + key, subkey = random.split(key) + cond = random.choice(key, 2, p=jnp.array([0.5, 0.5]), shape=([numbers])) + mask = jnp.ones([numbers, xc.shape[0], yc.shape[0]]) + xL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + xR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + key, subkey = random.split(key) + yL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + yR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + _xc = jnp.repeat(xc[None, :], numbers, axis=0) # add batch + _yc = jnp.repeat(yc[None, :], numbers, axis=0) # add batch + trns = 0.01 * jnp.ones_like(cond) + carry = cond, mask, _xc, _yc, xL, xR, yL, yR, trns + cond, mask, _xc, _yc, xL, xR, yL, yR, trns = vmap(select_W, 0, 0)(carry) + + u = u.loc[:, :, 2:-2, 2:-2, 2:-2].set( + u[:, :, 2:-2, 2:-2, 2:-2] * mask[:, None, :, :, None] + ) + u = u.loc[:, 0, 2:-2, 2:-2, 2:-2].add( + d0[:, :, None, None] * (1.0 - mask[:, :, :, None]) + ) + return u.loc[:, 4, 2:-2, 2:-2, 2:-2].add( + d0[:, :, None, None] * T0[:, :, None, None] * (1.0 - mask[:, :, :, None]) + ) + + +def init_multi_HD_3DTurb( + u, xc, yc, zc, numbers=100, init_key=2022, M0=0.1, k_tot=4.0, gamma=1.666666667 +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + def __create_3DTurb_init(u, keys): + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + d0 = 1.0 + cs = 1.0 / M0 + u0 = 1.0 # fixed + p0 = cs**2 * d0 / gamma + + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + dz = zc[1] - zc[0] + + qLx = dx * nx + qLy = dy * ny + qLz = dz * nz + + # random velocity field + vx, vy, vz = ( + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + ) + + key = random.PRNGKey(keys) + + kx0 = jnp.pi * 2.0 / qLx + ky0 = jnp.pi * 2.0 / qLy + kz0 = jnp.pi * 2.0 / qLz + + for k in range(-k_tot, k_tot + 1): + kz = kz0 * k # from 1 to k_tot + for j in range(-k_tot, k_tot + 1): + ky = ky0 * j # from 1 to k_tot + for i in range(-k_tot, k_tot + 1): + kx = kx0 * i # from 1 to k_tot + if i * j * k == 0: # avoiding uniform velocity + continue + # random phase + key, subkey = random.split(key) + phs = 2.0 * jnp.pi * random.uniform(key, shape=[3]) # (vi, k) + + uk = 1.0 / jnp.sqrt(kx**2 + ky**2 + kz**2) + kdx = ( + kx * xc[:, None, None] + + ky * yc[None, :, None] + + kz * zc[None, None, :] + ) + vx += uk * jnp.sin(kdx + phs[0]) + vy += uk * jnp.sin(kdx + phs[1]) + vz += uk * jnp.sin(kdx + phs[2]) + + del (kdx, uk, phs) + + # Helmholtz decomposition to subtract expansion: k.vk + dfx, dfy, dfz = 1.0 / qLx, 1.0 / qLy, 1.0 / qLz + fx = dfx * (jnp.arange(nx) - 1.0 - nx // 2) + fy = dfy * (jnp.arange(ny) - 1.0 - ny // 2) + fz = dfz * (jnp.arange(nz) - 1.0 - nz // 2) + + vkx = jnp.fft.fftn(vx) * dx * dy * dz + vky = jnp.fft.fftn(vy) * dx * dy * dz + vkz = jnp.fft.fftn(vz) * dx * dy * dz + + # shift to kxi=0 is at the center + vkx = jnp.fft.fftshift(vkx) + vky = jnp.fft.fftshift(vky) + vkz = jnp.fft.fftshift(vkz) + + fi = fx[:, None, None] ** 2 + fy[None, :, None] ** 2 + fz[None, None, :] ** 2 + fi = jnp.where(fi > 1.0e-8, 1.0 / fi, 0.0) + + fdv = ( + fx[:, None, None] * vkx + fy[None, :, None] * vky + fz[None, None, :] * vkz + ) * fi + vkx -= fdv * fx[:, None, None] + vky -= fdv * fy[None, :, None] + vkz -= fdv * fz[None, None, :] + del (fi, fdv) + + # shift back to original order + vkx = jnp.fft.ifftshift(vkx) + vky = jnp.fft.ifftshift(vky) + vkz = jnp.fft.ifftshift(vkz) + + # inverse FFT + vx = jnp.fft.ifftn(vkx).real * dfx * dfy * dfz + vy = jnp.fft.ifftn(vky).real * dfx * dfy * dfz + vz = jnp.fft.ifftn(vkz).real * dfx * dfy * dfz + + # renormalize total velocity + vtot = jnp.sqrt(vx**2 + vy**2 + vz**2).mean() + vx *= u0 / vtot + vy *= u0 / vtot + vz *= u0 / vtot + + u = u.loc[0].set(d0) + u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) + u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) + u = u.loc[3, 2:-2, 2:-2, 2:-2].set(vz) + return u.loc[4].add(p0) + + key = random.PRNGKey(init_key) + keys = random.randint( + key, + [ + numbers, + ], + minval=0, + maxval=10000000, + ) + return jax.vmap(__create_3DTurb_init, axis_name="i")(u, keys) + + +def init_multi_HD_3DRand( + u, + xc, + yc, + zc, + numbers=10000, + init_key=2022, + M0=0.1, + k_tot=4.0, + gamma=1.666666667, + dMx=1.0e1, + TMx=1.0e1, +): + """ + :param xc: cell center coordinate + :param mode: initial condition + :return: 1D scalar function u at cell center + """ + assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" + + def _pass(carry): + return carry + + def select_W(carry): + def _window(carry): + xx, yy, zz, val, xL, xR, yL, yR, zL, zR, trns = carry + x_win = 0.5 * (jnp.tanh((xx - xL) / trns) - jnp.tanh((xx - xR) / trns)) + y_win = 0.5 * (jnp.tanh((yy - yL) / trns) - jnp.tanh((yy - yR) / trns)) + z_win = 0.5 * (jnp.tanh((zz - zL) / trns) - jnp.tanh((zz - zR) / trns)) + val = x_win[:, None, None] * y_win[None, :, None] * z_win[None, None, :] + return xx, yy, zz, val, xL, xR, yL, yR, zL, zR, trns + + cond, value, xx, yy, zz, xL, xR, yL, yR, zL, zR, trns = carry + + carry = xx, yy, zz, value, xL, xR, yL, yR, zL, zR, trns + xx, yy, zz, value, xL, xR, yL, yR, zL, zR, trns = lax.cond( + cond == 1, _window, _pass, carry + ) + return cond, value, xx, yy, zz, xL, xR, yL, yR, zL, zR, trns + + def __create_3DRand_init(u, d0, T0, delD, delP, keys): + nx, ny, nz = xc.shape[0], yc.shape[0], zc.shape[0] + + p0 = d0 * T0 + cs = jnp.sqrt(T0 * gamma) + u0 = M0 * cs + + dx = xc[1] - xc[0] + dy = yc[1] - yc[0] + dz = zc[1] - zc[0] + + qLx = dx * nx + qLy = dy * ny + qLz = dz * nz + + # random velocity field + d, p, vx, vy, vz = ( + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + jnp.zeros([nx, ny, nz]), + ) + + key = random.PRNGKey(keys) + kx0 = jnp.pi * 2.0 / qLx + ky0 = jnp.pi * 2.0 / qLy + kz0 = jnp.pi * 2.0 / qLz + + for k in range(-k_tot, k_tot + 1): + kz = kz0 * k # from 1 to k_tot + for j in range(-k_tot, k_tot + 1): + ky = ky0 * j # from 1 to k_tot + for i in range(-k_tot, k_tot + 1): + kx = kx0 * i # from 1 to k_tot + if i * j * k == 0: # avoiding uniform velocity + continue + # random phase + key, subkey = random.split(key) + phs = 2.0 * jnp.pi * random.uniform(key, shape=[5]) # (vi, k) + + uk = 1.0 / jnp.sqrt(kx**2 + ky**2 + kz**2) + kdx = ( + kx * xc[:, None, None] + + ky * yc[None, :, None] + + kz * zc[None, None, :] + ) + vx += uk * jnp.sin(kdx + phs[0]) + vy += uk * jnp.sin(kdx + phs[1]) + vz += uk * jnp.sin(kdx + phs[2]) + p += uk * jnp.sin(kdx + phs[3]) + d += uk * jnp.sin(kdx + phs[4]) + + del (kdx, uk, phs) + + # renormalize total velocity + vtot = jnp.sqrt(vx**2 + vy**2 + vz**2).mean() + vx *= u0 / vtot + vy *= u0 / vtot + vz *= u0 / vtot + # d = d0 + delD * d / jnp.abs(d).mean() + # p = p0 + delP * p / jnp.abs(p).mean() + d = d0 * (1.0 + delD * d / jnp.abs(d).mean()) + p = p0 * (1.0 + delP * p / jnp.abs(p).mean()) + + u = u.loc[0, 2:-2, 2:-2, 2:-2].set(d) + u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) + u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) + u = u.loc[3, 2:-2, 2:-2, 2:-2].set(vz) + return u.loc[4, 2:-2, 2:-2, 2:-2].set(p) + + key = random.PRNGKey(init_key) + d0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=dMx) + key, subkey = random.split(key) + delD = random.uniform(key, shape=([numbers, 1]), minval=1.0e-2, maxval=0.2) + key, subkey = random.split(key) + T0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=TMx) + key, subkey = random.split(key) + delP = random.uniform(key, shape=([numbers, 1]), minval=1.0e-2, maxval=0.2) + + key, subkey = random.split(key) + keys = random.randint( + key, + [ + numbers, + ], + minval=0, + maxval=10000000, + ) + u = jax.vmap(__create_3DRand_init, axis_name="i")(u, d0, T0, delD, delP, keys) + + # perform window function + key, subkey = random.split(key) + cond = random.choice(key, 2, p=jnp.array([0.5, 0.5]), shape=([numbers])) + mask = jnp.ones([numbers, xc.shape[0], yc.shape[0], zc.shape[0]]) + xL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + xR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + key, subkey = random.split(key) + yL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + yR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + key, subkey = random.split(key) + zL = random.uniform(key, shape=([numbers]), minval=0.1, maxval=0.45) + zR = random.uniform(key, shape=([numbers]), minval=0.55, maxval=0.9) + _xc = jnp.repeat(xc[None, :], numbers, axis=0) + _yc = jnp.repeat(yc[None, :], numbers, axis=0) + _zc = jnp.repeat(zc[None, :], numbers, axis=0) + trns = 0.01 * jnp.ones_like(cond) + carry = cond, mask, _xc, _yc, _zc, xL, xR, yL, yR, zL, zR, trns + cond, mask, _xc, _yc, _zc, xL, xR, yL, yR, zL, zR, trns = vmap(select_W, 0, 0)( + carry + ) + + u = u.loc[:, :, 2:-2, 2:-2, 2:-2].set( + u[:, :, 2:-2, 2:-2, 2:-2] * mask[:, None, :, :, :] + ) + u = u.loc[:, 0, 2:-2, 2:-2, 2:-2].add( + d0[:, :, None, None] * (1.0 - mask[:, :, :, :]) + ) + return u.loc[:, 4, 2:-2, 2:-2, 2:-2].add( + d0[:, :, None, None] * T0[:, :, None, None] * (1.0 - mask[:, :, :, :]) + ) + + +def bc(u, dx, Ncell, mode="periodic"): + _u = jnp.zeros(Ncell + 4) # because of 2nd-order precision in space + _u = _u.loc[2 : Ncell + 2].set(u) + if mode == "periodic": # periodic boundary condition + _u = _u.loc[0:2].set(u[-2:]) # left hand side + _u = _u.loc[Ncell + 2 : Ncell + 4].set(u[0:2]) # right hand side + elif mode == "reflection": + _u = _u.loc[0].set(-u[3]) # left hand side + _u = _u.loc[1].set(-u[2]) # left hand side + _u = _u.loc[-2].set(-u[-3]) # right hand side + _u = _u.loc[-1].set(-u[-4]) # right hand side + elif mode == "copy": + _u = _u.loc[0].set(u[3]) # left hand side + _u = _u.loc[1].set(u[2]) # left hand side + _u = _u.loc[-2].set(u[-3]) # right hand side + _u = _u.loc[-1].set(u[-4]) # right hand side + + return _u + + +def bc_2D(_u, mode="trans"): + Nx, Ny = _u.shape + u = jnp.zeros([Nx + 4, Ny + 4]) # because of 2nd-order precision in space + u = u.loc[2:-2, 2:-2].set(_u) + Nx += 2 + Ny += 2 + + if mode == "periodic": # periodic boundary condition + # left hand side + u = u.loc[0:2, 2:-2].set(u[Nx - 2 : Nx, 2:-2]) # x + u = u.loc[2:-2, 0:2].set(u[2:-2, Ny - 2 : Ny]) # y + # right hand side + u = u.loc[Nx : Nx + 2, 2:-2].set(u[2:4, 2:-2]) + u = u.loc[2:-2, Ny : Ny + 2].set(u[2:-2, 2:4]) + elif mode == "trans": # periodic boundary condition + # left hand side + u = u.loc[0, 2:-2].set(u[3, 2:-2]) # x + u = u.loc[2:-2, 0].set(u[2:-2, 3]) # y + u = u.loc[1, 2:-2].set(u[2, 2:-2]) # x + u = u.loc[2:-2, 1].set(u[2:-2, 2]) # y + # right hand side + u = u.loc[-2, 2:-2].set(u[-3, 2:-2]) + u = u.loc[2:-2, -2].set(u[2:-2, -3]) + u = u.loc[-1, 2:-2].set(u[-4, 2:-2]) + u = u.loc[2:-2, -1].set(u[2:-2, -4]) + elif mode == "Neumann": # periodic boundary condition + # left hand side + u = u.loc[0, 2:-2].set(0.0) # x + u = u.loc[2:-2, 0].set(0.0) # y + u = u.loc[1, 2:-2].set(0.0) # x + u = u.loc[2:-2, 1].set(0.0) # y + # right hand side + u = u.loc[-2, 2:-2].set(0.0) + u = u.loc[2:-2, -2].set(0.0) + u = u.loc[-1, 2:-2].set(0.0) + u = u.loc[2:-2, -1].set(0.0) + return u + + +def bc_HD(u, mode): + _, Nx, Ny, Nz = u.shape + Nx -= 2 + Ny -= 2 + Nz -= 2 + if mode == "periodic": # periodic boundary condition + # left hand side + u = u.loc[:, 0:2, 2:-2, 2:-2].set(u[:, Nx - 2 : Nx, 2:-2, 2:-2]) # x + u = u.loc[:, 2:-2, 0:2, 2:-2].set(u[:, 2:-2, Ny - 2 : Ny, 2:-2]) # y + u = u.loc[:, 2:-2, 2:-2, 0:2].set(u[:, 2:-2, 2:-2, Nz - 2 : Nz]) # z + # right hand side + u = u.loc[:, Nx : Nx + 2, 2:-2, 2:-2].set(u[:, 2:4, 2:-2, 2:-2]) + u = u.loc[:, 2:-2, Ny : Ny + 2, 2:-2].set(u[:, 2:-2, 2:4, 2:-2]) + u = u.loc[:, 2:-2, 2:-2, Nz : Nz + 2].set(u[:, 2:-2, 2:-2, 2:4]) + elif mode == "trans": # periodic boundary condition + # left hand side + u = u.loc[:, 0, 2:-2, 2:-2].set(u[:, 3, 2:-2, 2:-2]) # x + u = u.loc[:, 2:-2, 0, 2:-2].set(u[:, 2:-2, 3, 2:-2]) # y + u = u.loc[:, 2:-2, 2:-2, 0].set(u[:, 2:-2, 2:-2, 3]) # z + u = u.loc[:, 1, 2:-2, 2:-2].set(u[:, 2, 2:-2, 2:-2]) # x + u = u.loc[:, 2:-2, 1, 2:-2].set(u[:, 2:-2, 2, 2:-2]) # y + u = u.loc[:, 2:-2, 2:-2, 1].set(u[:, 2:-2, 2:-2, 2]) # z + # right hand side + u = u.loc[:, -2, 2:-2, 2:-2].set(u[:, -3, 2:-2, 2:-2]) + u = u.loc[:, 2:-2, -2, 2:-2].set(u[:, 2:-2, -3, 2:-2]) + u = u.loc[:, 2:-2, 2:-2, -2].set(u[:, 2:-2, 2:-2, -3]) + u = u.loc[:, -1, 2:-2, 2:-2].set(u[:, -4, 2:-2, 2:-2]) + u = u.loc[:, 2:-2, -1, 2:-2].set(u[:, 2:-2, -4, 2:-2]) + u = u.loc[:, 2:-2, 2:-2, -1].set(u[:, 2:-2, 2:-2, -4]) + elif mode == "KHI": # x: periodic, y, z : trans + # left hand side + u = u.loc[:, 0:2, 2:-2, 2:-2].set(u[:, Nx - 2 : Nx, 2:-2, 2:-2]) # x + u = u.loc[:, 2:-2, 0, 2:-2].set(u[:, 2:-2, 3, 2:-2]) # y + u = u.loc[:, 2:-2, 2:-2, 0].set(u[:, 2:-2, 2:-2, 3]) # z + u = u.loc[:, 2:-2, 1, 2:-2].set(u[:, 2:-2, 2, 2:-2]) # y + u = u.loc[:, 2:-2, 2:-2, 1].set(u[:, 2:-2, 2:-2, 2]) # z + # right hand side + u = u.loc[:, Nx : Nx + 2, 2:-2, 2:-2].set(u[:, 2:4, 2:-2, 2:-2]) + u = u.loc[:, 2:-2, -2, 2:-2].set(u[:, 2:-2, -3, 2:-2]) + u = u.loc[:, 2:-2, 2:-2, -2].set(u[:, 2:-2, 2:-2, -3]) + u = u.loc[:, 2:-2, -1, 2:-2].set(u[:, 2:-2, -4, 2:-2]) + u = u.loc[:, 2:-2, 2:-2, -1].set(u[:, 2:-2, 2:-2, -4]) + return u + + +def bc_HD_vis(u, if_periodic=True): # for viscosity + """ + for the moment, assuming periodic/copy boundary + seemingly, copy boundary does not work well... + """ + _, Nx, Ny, Nz = u.shape + Nx -= 2 + Ny -= 2 + Nz -= 2 + + if if_periodic: + u = u.loc[:, 0:2, 0:2, 2:-2].set(u[:, Nx - 2 : Nx, Ny - 2 : Ny, 2:-2]) # xByB + u = u.loc[:, 0:2, 2:-2, 0:2].set(u[:, Nx - 2 : Nx, 2:-2, Nz - 2 : Nz]) # xBzB + u = u.loc[:, 0:2, Ny : Ny + 2, 2:-2].set(u[:, Nx - 2 : Nx, 2:4, 2:-2]) # xByT + u = u.loc[:, 0:2, 2:-2, Nz : Nz + 2].set(u[:, Nx - 2 : Nx, 2:-2, 2:4]) # xBzT + u = u.loc[:, Nx : Nx + 2, 0:2, 2:-2].set(u[:, 2:4, Ny - 2 : Ny, 2:-2]) # xTyB + u = u.loc[:, Nx : Nx + 2, 2:-2, 0:2].set(u[:, 2:4, 2:-2, Nz - 2 : Nz]) # xTzB + u = u.loc[:, Nx : Nx + 2, Ny : Ny + 2, 2:-2].set(u[:, 2:4, 2:4, 2:-2]) # xTyT + u = u.loc[:, Nx : Nx + 2, 2:-2, Nz : Nz + 2].set(u[:, 2:4, 2:-2, 2:4]) # xTzT + else: # trans + u = u.loc[:, 0:2, 0:2, 2:-2].set(u[:, 4:2, 4:2, 2:-2]) # xByT + u = u.loc[:, 0:2, 2:-2, 0:2].set(u[:, 4:2, 2:-2, 4:2]) # xBzB + u = u.loc[:, 0:2, Ny : Ny + 2, 2:-2].set(u[:, 4:2, Ny : Ny - 2, 2:-2]) # xByB + u = u.loc[:, 0:2, 2:-2, Nz : Nz + 2].set(u[:, 4:2, 2:-2, Nz : Nz - 2]) # xBzT + u = u.loc[:, Nx : Nx + 2, 0:2, 2:-2].set(u[:, Nx : Nx - 2, 4:2, 2:-2]) # xTyB + u = u.loc[:, Nx : Nx + 2, 2:-2, 0:2].set(u[:, Nx : Nx - 2, 2:-2, 4:2]) # xTzB + u = u.loc[:, Nx : Nx + 2, Ny : Ny + 2, 2:-2].set( + u[:, Nx : Nx - 2, Ny : Ny - 2, 2:-2] + ) # xTyT + u = u.loc[:, Nx : Nx + 2, 2:-2, Nz : Nz + 2].set( + u[:, Nx : Nx - 2, 2:-2, Nz : Nz - 2] + ) # xTzT + + return u + + +def VLlimiter(a, b, c, alpha=2.0): + return ( + jnp.sign(c) + * (0.5 + 0.5 * jnp.sign(a * b)) + * jnp.minimum(alpha * jnp.minimum(jnp.abs(a), jnp.abs(b)), jnp.abs(c)) + ) + + +def limiting(u, Ncell, if_second_order): + # under construction + du_L = u[1 : Ncell + 3] - u[0 : Ncell + 2] + du_R = u[2 : Ncell + 4] - u[1 : Ncell + 3] + du_M = (u[2 : Ncell + 4] - u[0 : Ncell + 2]) * 0.5 + gradu = VLlimiter(du_L, du_R, du_M) * if_second_order + # -1:Ncell + # uL, uR = jnp.zeros(Ncell+4), jnp.zeros(Ncell+4) + uL, uR = jnp.zeros_like(u), jnp.zeros_like(u) + # left of cell + uL = uL.loc[1 : Ncell + 3].set(u[1 : Ncell + 3] - 0.5 * gradu) + # right of cell + uR = uR.loc[1 : Ncell + 3].set(u[1 : Ncell + 3] + 0.5 * gradu) + return uL, uR + + +def limiting_HD(u, if_second_order): + _, nx, _, _ = u.shape + uL, uR = u, u + nx -= 4 + + du_L = u[:, 1 : nx + 3, :, :] - u[:, 0 : nx + 2, :, :] + du_R = u[:, 2 : nx + 4, :, :] - u[:, 1 : nx + 3, :, :] + du_M = (u[:, 2 : nx + 4, :, :] - u[:, 0 : nx + 2, :, :]) * 0.5 + gradu = VLlimiter(du_L, du_R, du_M) * if_second_order + # -1:Ncell + uL = uL.loc[:, 1 : nx + 3, :, :].set( + u[:, 1 : nx + 3, :, :] - 0.5 * gradu + ) # left of cell + uR = uR.loc[:, 1 : nx + 3, :, :].set( + u[:, 1 : nx + 3, :, :] + 0.5 * gradu + ) # right of cell + + uL = jnp.where(uL[0] > 0.0, uL, u) + uL = jnp.where(uL[4] > 0.0, uL, u) + uR = jnp.where(uR[0] > 0.0, uR, u) + uR = jnp.where(uR[4] > 0.0, uR, u) + + return uL, uR + + +def save_data(u, xc, i_save, save_dir, dt_save=None, if_final=False): + if if_final: + jnp.save(save_dir + "/x_coordinate", xc) + tc = jnp.arange(i_save + 1) * dt_save + jnp.save(save_dir + "/t_coordinate", tc) + flnm = save_dir + "/Data_" + str(i_save).zfill(4) + jnp.save(flnm, u) + else: + flnm = save_dir + "/Data_" + str(i_save).zfill(4) + jnp.save(flnm, u) + + +def save_data_HD(u, xc, yc, zc, i_save, save_dir, dt_save=None, if_final=False): + if if_final: + jnp.save(save_dir + "/x_coordinate", xc) + jnp.save(save_dir + "/y_coordinate", yc) + jnp.save(save_dir + "/z_coordinate", zc) + tc = jnp.arange(i_save + 1) * dt_save + jnp.save(save_dir + "/t_coordinate", tc) + flnm = save_dir + "/Data_" + str(i_save).zfill(4) + jnp.save(flnm, u) + else: + flnm = save_dir + "/Data_" + str(i_save).zfill(4) + jnp.save(flnm, u) + + +def Courant(u, dx): + return dx / (jnp.max(jnp.abs(u)) + 1.0e-8) + + +def Courant_diff(dx, epsilon=1.0e-3): + return 0.5 * dx**2 / (epsilon + 1.0e-8) + + +def Courant_diff_2D(dx, dy, epsilon=1.0e-3): + stability_dif_x = 0.5 * dx**2 / (epsilon + 1.0e-8) + stability_dif_y = 0.5 * dy**2 / (epsilon + 1.0e-8) + return jnp.min(jnp.array([stability_dif_x, stability_dif_y])) + + +def Courant_HD(u, dx, dy, dz, gamma): + cs = jnp.sqrt(gamma * u[4] / u[0]) # sound velocity + stability_adv_x = dx / (jnp.max(cs + jnp.abs(u[1])) + 1.0e-8) + stability_adv_y = dy / (jnp.max(cs + jnp.abs(u[2])) + 1.0e-8) + stability_adv_z = dz / (jnp.max(cs + jnp.abs(u[3])) + 1.0e-8) + return jnp.min(jnp.array([stability_adv_x, stability_adv_y, stability_adv_z])) + + +def Courant_vis_HD(dx, dy, dz, eta, zeta): + # visc = jnp.max(jnp.array([eta, zeta])) + visc = 4.0 / 3.0 * eta + zeta # maximum + stability_dif_x = 0.5 * dx**2 / (visc + 1.0e-8) + stability_dif_y = 0.5 * dy**2 / (visc + 1.0e-8) + stability_dif_z = 0.5 * dz**2 / (visc + 1.0e-8) + return jnp.min(jnp.array([stability_dif_x, stability_dif_y, stability_dif_z])) diff --git a/pdebench/data_gen/gen_diff_react.py b/pdebench/data_gen/gen_diff_react.py index 5696b85..49d1a26 100644 --- a/pdebench/data_gen/gen_diff_react.py +++ b/pdebench/data_gen/gen_diff_react.py @@ -1,7 +1,21 @@ -#!/usr/bin/env python +from __future__ import annotations +import logging +import multiprocessing as mp import os +import time +from itertools import repeat +from pathlib import Path + import dotenv +import h5py +import hydra +import numpy as np +from hydra.utils import get_original_cwd +from omegaconf import DictConfig, OmegaConf +from pdebench.data_gen.src import utils +from pdebench.data_gen.uploader import dataverse_upload + # load environment variables from `.env` file if it exists # recursively searches for `.env` in all folders starting from work dir # this allows us to keep defaults local to the machine @@ -17,52 +31,61 @@ os.environ["VECLIB_MAXIMUM_THREADS"] = num_threads os.environ["NUMEXPR_NUM_THREADS"] = num_threads -import dotenv -import hydra -from hydra.utils import get_original_cwd -from omegaconf import DictConfig, OmegaConf -import logging -import multiprocessing as mp -from itertools import repeat -import numpy as np - -from src import utils -import h5py -from uploader import dataverse_upload log = logging.getLogger(__name__) -def simulator(config, i, data_f): - - from src import sim_diff_react - + +def simulator(config, i): + from pdebench.data_gen.src import sim_diff_react + config.sim.seed = i log.info(f"Starting seed {i}") + start_time = time.time() sim_obj = sim_diff_react.Simulator(**config.sim) data_sample = sim_obj.generate_sample() + duration = time.time() - start_time + log.info(f"Seed {config.sim.seed} took {duration} to finish") seed_str = str(i).zfill(4) - - ## Chunking for compression and data access - ## https://docs.h5py.org/en/stable/high/dataset.html#chunked-storage - ## should be by batch and less than 1MB - ## lzf compression for float32 is kind of pointless though. - data_f.create_dataset( - f"{seed_str}/data", - data=data_sample, - dtype="float32", - compression="lzf" - ) - data_f.create_dataset( - f"{seed_str}/grid/x", data = sim_obj.x, dtype="float32", compression="lzf" - ) - data_f.create_dataset( - f"{seed_str}/grid/y", data=sim_obj.y, dtype="float32", compression="lzf" - ) - data_f.create_dataset( - f"{seed_str}/grid/t", data=sim_obj.t, dtype="float32", compression="lzf" - ) - data_f.attrs[f"{seed_str}/config"] = OmegaConf.to_yaml(config) + + while True: + try: + with h5py.File(utils.expand_path(config.output_path), "a") as data_f: + ## Chunking for compression and data access + ## https://docs.h5py.org/en/stable/high/dataset.html#chunked-storage + ## should be by batch and less than 1MB + ## lzf compression for float32 is kind of pointless though. + data_f.create_dataset( + f"{seed_str}/data", + data=data_sample, + dtype="float32", + compression="lzf", + ) + data_f.create_dataset( + f"{seed_str}/grid/x", + data=sim_obj.x, + dtype="float32", + compression="lzf", + ) + data_f.create_dataset( + f"{seed_str}/grid/y", + data=sim_obj.y, + dtype="float32", + compression="lzf", + ) + data_f.create_dataset( + f"{seed_str}/grid/t", + data=sim_obj.t, + dtype="float32", + compression="lzf", + ) + seed_group = data_f[seed_str] + seed_group.attrs["config"] = OmegaConf.to_yaml(config) + except OSError: + time.sleep(0.1) + continue + else: + break @hydra.main(config_path="configs/", config_name="diff-react") @@ -76,47 +99,35 @@ def main(config: DictConfig): # Change to original working directory to import modules - temp_path = os.getcwd() + temp_path = Path.cwd() os.chdir(get_original_cwd()) - from src import utils - import h5py - # Change back to the hydra working directory os.chdir(temp_path) - - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) + '.h5' + + work_path = Path(config.work_dir) + output_path: Path = work_path / config.data_dir / config.output_path + output_path.mkdir(output_path, exist_ok=True, parents=True) + config.output_path = (output_path / config.output_path).with_suffix(".h5") num_samples_init = 0 num_samples_final = 1000 - with h5py.File(utils.expand_path(config.output_path), "w") as data_f: - - pool = mp.Pool(mp.cpu_count()) - seed = np.arange(num_samples_init, num_samples_final) - seed = seed.tolist() - pool.starmap(simulator, zip(repeat(config), seed, data_f)) + pool = mp.Pool(mp.cpu_count()) + seed = np.arange(num_samples_init, num_samples_final) + seed = seed.tolist() + pool.starmap(simulator, zip(repeat(config), seed)) if config.upload: dataverse_upload( file_path=config.output_path, - dataverse_url=os.getenv( - 'DATAVERSE_URL', 'https://darus.uni-stuttgart.de'), - dataverse_token=os.getenv( - 'DATAVERSE_API_TOKEN', ''), + dataverse_url=os.getenv("DATAVERSE_URL", "https://darus.uni-stuttgart.de"), + dataverse_token=os.getenv("DATAVERSE_API_TOKEN", ""), dataverse_dir=config.name, - dataverse_id=os.getenv( - 'DATAVERSE_ID', ''), - log=log) + dataverse_id=os.getenv("DATAVERSE_ID", ""), + log=log, + ) - return - - -import os if __name__ == "__main__": test = main() diff --git a/pdebench/data_gen/gen_diff_sorp.py b/pdebench/data_gen/gen_diff_sorp.py index 09a3759..2f1462b 100644 --- a/pdebench/data_gen/gen_diff_sorp.py +++ b/pdebench/data_gen/gen_diff_sorp.py @@ -1,13 +1,28 @@ -#!/usr/bin/env python +from __future__ import annotations +import logging +import multiprocessing as mp import os +import time +from itertools import repeat +from pathlib import Path + import dotenv +import h5py +import hydra +import numpy as np +from hydra.utils import get_original_cwd +from omegaconf import DictConfig, OmegaConf +from pdebench.data_gen.src import utils +from pdebench.data_gen.uploader import dataverse_upload + # load environment variables from `.env` file if it exists # recursively searches for `.env` in all folders starting from work dir # this allows us to keep defaults local to the machine # e.g. HPC versus local laptop dotenv.load_dotenv() + # or if the environment variables will be fixed for all executions, we can hard-code the environment variables like this: num_threads = "4" @@ -17,50 +32,55 @@ os.environ["VECLIB_MAXIMUM_THREADS"] = num_threads os.environ["NUMEXPR_NUM_THREADS"] = num_threads -import dotenv -import hydra -from hydra.utils import get_original_cwd -from omegaconf import DictConfig, OmegaConf -import logging -import multiprocessing as mp -from itertools import repeat -import numpy as np - -from src import utils -import h5py -from uploader import dataverse_upload log = logging.getLogger(__name__) -def simulator(config, i, data_f): - - from src import sim_diff_sorp - +def simulator(config, i): + from pdebench.data_gen.src import sim_diff_sorp + config.sim.seed = i log.info(f"Starting seed {i}") + start_time = time.time() sim_obj = sim_diff_sorp.Simulator(**config.sim) data_sample = sim_obj.generate_sample() - + duration = time.time() - start_time + log.info(f"Seed {config.sim.seed} took {duration} to finish") + seed_str = str(i).zfill(4) - - ## Chunking for compression and data access - ## https://docs.h5py.org/en/stable/high/dataset.html#chunked-storage - ## should be by batch and less than 1MB - ## lzf compression for float32 is kind of pointless though. - data_f.create_dataset( - f"{seed_str}/data", - data=data_sample, - dtype="float32", - compression="lzf" - ) - data_f.create_dataset( - f"{seed_str}/grid/x", data = sim_obj.x, dtype="float32", compression="lzf" - ) - data_f.create_dataset( - f"{seed_str}/grid/t", data=sim_obj.t, dtype="float32", compression="lzf" - ) - data_f.attrs[f"{seed_str}/config"] = OmegaConf.to_yaml(config) + + while True: + try: + with h5py.File(utils.expand_path(config.output_path), "a") as data_f: + ## Chunking for compression and data access + ## https://docs.h5py.org/en/stable/high/dataset.html#chunked-storage + ## should be by batch and less than 1MB + ## lzf compression for float32 is kind of pointless though. + data_f.create_dataset( + f"{seed_str}/data", + data=data_sample, + dtype="float32", + compression="lzf", + ) + data_f.create_dataset( + f"{seed_str}/grid/x", + data=sim_obj.x, + dtype="float32", + compression="lzf", + ) + data_f.create_dataset( + f"{seed_str}/grid/t", + data=sim_obj.t, + dtype="float32", + compression="lzf", + ) + seed_group = data_f[seed_str] + seed_group.attrs["config"] = OmegaConf.to_yaml(config) + except OSError: + time.sleep(0.1) + continue + else: + break @hydra.main(config_path="configs/", config_name="diff-sorp") @@ -71,51 +91,38 @@ def main(config: DictConfig): # Imports should be nested inside @hydra.main to optimize tab completion # Read more here: https://github.com/facebookresearch/hydra/issues/934 - + # Change to original working directory to import modules - - temp_path = os.getcwd() + + temp_path = Path.cwd() os.chdir(get_original_cwd()) - - from src import utils - import h5py - - # Change back to the hydra working directory + + # Change back to the hydra working directory os.chdir(temp_path) - - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) + '.h5' - + + work_path = Path(config.work_dir).parent + output_path: Path = work_path / config.data_dir / config.output_path + output_path.mkdir(parents=True, exist_ok=True) + config.output_path = (output_path / config.output_path).with_suffix(".h5") + num_samples_init = 0 num_samples_final = 10000 - - with h5py.File(utils.expand_path(config.output_path), "w") as data_f: - - pool = mp.Pool(mp.cpu_count()) - seed = np.arange(num_samples_init, num_samples_final) - seed = seed.tolist() - pool.starmap(simulator, zip(repeat(config), seed, data_f)) - + + pool = mp.Pool(mp.cpu_count()) + seed = np.arange(num_samples_init, num_samples_final) + seed = seed.tolist() + pool.starmap(simulator, zip(repeat(config), seed)) + if config.upload: dataverse_upload( file_path=config.output_path, - dataverse_url=os.getenv( - 'DATAVERSE_URL', 'https://darus.uni-stuttgart.de'), - dataverse_token=os.getenv( - 'DATAVERSE_API_TOKEN', ''), + dataverse_url=os.getenv("DATAVERSE_URL", "https://darus.uni-stuttgart.de"), + dataverse_token=os.getenv("DATAVERSE_API_TOKEN", ""), dataverse_dir=config.name, - dataverse_id=os.getenv( - 'DATAVERSE_ID', ''), - log=log) - - return + dataverse_id=os.getenv("DATAVERSE_ID", ""), + log=log, + ) -import os - if __name__ == "__main__": test = main() - \ No newline at end of file diff --git a/pdebench/data_gen/gen_ns_incomp.py b/pdebench/data_gen/gen_ns_incomp.py index 2542965..82eaa2a 100644 --- a/pdebench/data_gen/gen_ns_incomp.py +++ b/pdebench/data_gen/gen_ns_incomp.py @@ -1,8 +1,8 @@ -#!/usr/bin/env python +from __future__ import annotations -import hydra import dotenv -from omegaconf import DictConfig, OmegaConf +import hydra +from omegaconf import DictConfig # load environment variables from `.env` file if it exists # recursively searches for `.env` in all folders starting from work dir @@ -17,8 +17,9 @@ def main(config: DictConfig): Args: config: This function uses hydra configuration for all parameters. """ - + from src import sim_ns_incomp_2d + sim_ns_incomp_2d.ns_sim(config=config, **config) diff --git a/pdebench/data_gen/gen_radial_dam_break.py b/pdebench/data_gen/gen_radial_dam_break.py index 8586e3e..798f3ff 100644 --- a/pdebench/data_gen/gen_radial_dam_break.py +++ b/pdebench/data_gen/gen_radial_dam_break.py @@ -1,13 +1,25 @@ -#!/usr/bin/env python +from __future__ import annotations -from copy import deepcopy +import logging +import multiprocessing as mp import os +import time +from copy import deepcopy +from itertools import repeat # load environment variables from `.env` file if it exists # recursively searches for `.env` in all folders starting from work dir # this allows us to keep defaults local to the machine # e.g. HPC versus local laptop import dotenv +import h5py +import hydra +import numpy as np +from hydra.utils import get_original_cwd +from omegaconf import DictConfig, OmegaConf +from pdebench.data_gen.src import utils +from pdebench.data_gen.src.sim_radial_dam_break import RadialDamBreak2D +from pdebench.data_gen.uploader import dataverse_upload dotenv.load_dotenv() @@ -21,31 +33,18 @@ os.environ["NUMEXPR_NUM_THREADS"] = num_threads os.environ["NUMEXPR_MAX_THREADS"] = num_threads -import hydra -from hydra.utils import get_original_cwd -from omegaconf import DictConfig, OmegaConf -import h5py -import logging -import multiprocessing as mp -from itertools import repeat -from src import utils -import numpy as np -from uploader import dataverse_upload log = logging.getLogger(__name__) -def simulator(base_config, i, data_f): - - from src.sim_radial_dam_break.py import RadialDamBreak2D - +def simulator(base_config, i): config = deepcopy(base_config) config.sim.seed = i - log.info(f"Starting seed {i}") + log.info("Starting seed %d", i) - np.random.seed(config.sim.seed) + rng = np.random.default_rng(config.sim.seed) # config.sim.inner_height = np.random.uniform(1.5, 2.5) - config.sim.dam_radius = np.random.uniform(0.3, 0.7) + config.sim.dam_radius = rng.uniform(0.3, 0.7) scenario = RadialDamBreak2D( grav=config.sim.gravity, @@ -54,14 +53,33 @@ def simulator(base_config, i, data_f): ydim=config.sim.ydim, ) + start_time = time.time() scenario.run(T=config.sim.T_end, tsteps=config.sim.n_time_steps, plot=False) - scenario.save_state_to_disk(data_f, i) - + duration = time.time() - start_time seed_str = str(i).zfill(4) - - data_f.attrs[f"{seed_str}/config"] = OmegaConf.to_yaml(config) - - + log.info("Seed %s took %s to finish", seed_str, duration) + + while True: + try: + with h5py.File(utils.expand_path(config.output_path), "a") as h5_file: + scenario.save_state_to_disk(h5_file, seed_str) + seed_group = h5_file[seed_str] + seed_group.attrs["config"] = OmegaConf.to_yaml(config) + except OSError: + time.sleep(0.1) + continue + else: + break + + if config.upload: + dataverse_upload( + file_path=config.output_path, + dataverse_url=os.getenv("DATAVERSE_URL", "https://darus.uni-stuttgart.de"), + dataverse_token=os.getenv("DATAVERSE_API_TOKEN", ""), + dataverse_dir=config.name, + dataverse_id=os.getenv("DATAVERSE_ID", ""), + log=log, + ) @hydra.main(config_path="configs/", config_name="radial_dam_break") @@ -75,42 +93,27 @@ def main(config: DictConfig): # Change to original working directory to import modules import os + from pathlib import Path - temp_path = os.getcwd() + temp_path = Path.cwd() os.chdir(get_original_cwd()) - - # Change back to the hydra working directory + + # Change back to the hydra working directory os.chdir(temp_path) - - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) + '.h5' + + work_path = Path(config.work_dir) + output_path = work_path / config.data_dir / config.output_path + if not output_path.is_dir(): + output_path.mkdir(parents=True) + config.output_path = output_path / config.output_path.with_suffix(".h5") num_samples_init = 0 - num_samples_final = 1000 - - with h5py.File(utils.expand_path(config.output_path), "w") as data_f: - - pool = mp.Pool(mp.cpu_count()) - seed = np.arange(num_samples_init, num_samples_final) - seed = seed.tolist() - pool.starmap(simulator, zip(repeat(config), seed, data_f)) - - if config.upload: - dataverse_upload( - file_path=config.output_path, - dataverse_url=os.getenv( - 'DATAVERSE_URL', 'https://darus.uni-stuttgart.de'), - dataverse_token=os.getenv( - 'DATAVERSE_API_TOKEN', ''), - dataverse_dir=config.name, - dataverse_id=os.getenv( - 'DATAVERSE_ID', ''), - log=log) + num_samples_final = 10000 - return + pool = mp.Pool(mp.cpu_count()) + seed = np.arange(num_samples_init, num_samples_final) + seed = seed.tolist() + pool.starmap(simulator, zip(repeat(config), seed)) if __name__ == "__main__": diff --git a/pdebench/data_gen/notebooks/Analysis.ipynb b/pdebench/data_gen/notebooks/Analysis.ipynb index 045cb58..1d1c2e0 100644 --- a/pdebench/data_gen/notebooks/Analysis.ipynb +++ b/pdebench/data_gen/notebooks/Analysis.ipynb @@ -2,29 +2,32 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ + "from __future__ import annotations\n", + "\n", "%load_ext autoreload\n", "%autoreload 2" ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sys\n", - "import os\n", - "parentdir = os.path.dirname(os.getcwd())\n", - "sys.path.append(parentdir)" + "from pathlib import Path\n", + "\n", + "parentdir = Path.cwd()\n", + "sys.path.append(str(parentdir))" ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [ { @@ -40,6 +43,7 @@ ], "source": [ "import dotenv\n", + "\n", "# load environment variables from `.env` file if it exists\n", "# recursively searches for `.env` in all folders starting from work dir\n", "dotenv.load_dotenv()" @@ -47,7 +51,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -55,10 +59,6 @@ "import h5py\n", "import matplotlib.pyplot as plt\n", "from einops import rearrange\n", - "from phi.flow import *\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from phi.vis import *\n", "from src.utils import resolve_path" ] }, @@ -76,8 +76,8 @@ } ], "source": [ - "data_path = resolve_path('${WORKING_DIR}/*/*/*/*/*/*.h5', idx=-1, unique=False)\n", - "print(data_path)" + "data_path = resolve_path(\"${WORKING_DIR}/*/*/*/*/*/*.h5\", idx=-1, unique=False)\n", + "print(data_path) # noqa: T201" ] }, { @@ -97,7 +97,7 @@ } ], "source": [ - "data_f = h5py.File(data_path, 'r')\n", + "data_f = h5py.File(data_path, \"r\")\n", "data_f.keys()" ] }, @@ -118,7 +118,7 @@ } ], "source": [ - "data_f['particles'].shape" + "data_f[\"particles\"].shape" ] }, { @@ -140,16 +140,15 @@ } ], "source": [ - "\n", "columns = 5\n", "fsize = 12\n", "\n", - "arr = data_f['particles'][0, :columns, :, :, 0]\n", + "arr = data_f[\"particles\"][0, :columns, :, :, 0]\n", "\n", - "fig = plt.figure(figsize=(fsize, fsize/columns))\n", - "plt.imshow(rearrange(arr, 't x y -> x (t y)'))\n", + "fig = plt.figure(figsize=(fsize, fsize / columns))\n", + "plt.imshow(rearrange(arr, \"t x y -> x (t y)\"))\n", "plt.gca().set_axis_off()\n", - "plt.tight_layout(pad = 1)\n", + "plt.tight_layout(pad=1)\n", "plt.show()" ] }, @@ -172,16 +171,15 @@ } ], "source": [ - "\n", "columns = 5\n", "fsize = 12\n", "\n", - "arr = data_f['particles'][0, 1000:1000+columns, :, :, 0]\n", + "arr = data_f[\"particles\"][0, 1000 : 1000 + columns, :, :, 0]\n", "\n", - "fig = plt.figure(figsize=(fsize, fsize/columns))\n", - "plt.imshow(rearrange(arr, 't x y -> x (t y)'))\n", + "fig = plt.figure(figsize=(fsize, fsize / columns))\n", + "plt.imshow(rearrange(arr, \"t x y -> x (t y)\"))\n", "plt.gca().set_axis_off()\n", - "plt.tight_layout(pad = 1)\n", + "plt.tight_layout(pad=1)\n", "plt.show()" ] }, @@ -204,16 +202,15 @@ } ], "source": [ - "\n", "columns = 5\n", "fsize = 12\n", "\n", - "arr = data_f['force'][:columns, :, :, 0]\n", + "arr = data_f[\"force\"][:columns, :, :, 0]\n", "\n", - "fig = plt.figure(figsize=(fsize, fsize/columns))\n", - "plt.imshow(rearrange(arr, 'b x y -> x (b y)'))\n", + "fig = plt.figure(figsize=(fsize, fsize / columns))\n", + "plt.imshow(rearrange(arr, \"b x y -> x (b y)\"))\n", "plt.gca().set_axis_off()\n", - "plt.tight_layout(pad = 1)\n", + "plt.tight_layout(pad=1)\n", "plt.show()" ] }, @@ -226,11 +223,8 @@ } ], "metadata": { - "interpreter": { - "hash": "38db623c9a8c22d26705f5e96929a68995fd63f2ffb05d6c4e502eb3b9fc3e8d" - }, "kernelspec": { - "display_name": "Python 3.9.4 ('venv2': venv)", + "display_name": "pde_bench_tmp", "language": "python", "name": "python3" }, @@ -244,14 +238,9 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.4" + "version": "3.10.12" }, - "orig_nbformat": 4, - "vscode": { - "interpreter": { - "hash": "7a0d8eacbb66928d1b3dec5c8bc3662d91007e449cc8cc497b49c1d19e71d254" - } - } + "orig_nbformat": 4 }, "nbformat": 4, "nbformat_minor": 2 diff --git a/pdebench/data_gen/plot.py b/pdebench/data_gen/plot.py index 5453f0a..99abcdc 100644 --- a/pdebench/data_gen/plot.py +++ b/pdebench/data_gen/plot.py @@ -1,52 +1,55 @@ -# -*- coding: utf-8 -*- """ Created on Wed May 4 09:53:18 2022 @author: timot """ +from __future__ import annotations +import os +from pathlib import Path + +import h5py import hydra -from omegaconf import DictConfig import numpy as np -import h5py -import matplotlib.pyplot as plt from hydra.utils import get_original_cwd -from src.plots import plot_data +from omegaconf import DictConfig +from pdebench.data_gen.src.plots import plot_data + @hydra.main(config_path="configs/", config_name="diff-sorp") def main(config: DictConfig): """ use config specifications to generate dataset """ + work_path = Path(config.work_dir) + output_path = work_path / config.data_dir / config.output_path + if not output_path.is_dir(): + output_path.mkdir(parents=True) + config.output_path = output_path / config.output_path - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) - # Open and load file - data_path = config.output_path + '.h5' + data_path = config.output_path + ".h5" h5_file = h5py.File(data_path, "r") - - if "seed" in config.sim.keys(): + rng = np.random.default_rng() + + if "seed" in config.sim: # Choose random sample number idx_max = 10000 if config.plot.dim == 1 else 1000 - config.sim.seed = np.random.randint(0, idx_max) + config.sim.seed = rng.integers(0, idx_max) postfix = str(config.sim.seed).zfill(4) data = np.array(h5_file[f"{postfix}/data"], dtype="f") t = np.array(h5_file[f"{postfix}/grid/t"], dtype="f") # data dim = [t, x1, ..., xd, v] else: idx_max = 10000 if config.plot.dim == 1 else 1000 - postfix = np.random.randint(0, idx_max) + postfix = rng.randint(0, idx_max) data = np.array(h5_file["data"], dtype="f") data = data[postfix] t = np.array(h5_file["grid/t"], dtype="f") t = t[postfix] # data dim = [t, x1, ..., xd, v] - + h5_file.close() os.chdir(get_original_cwd()) @@ -60,10 +63,6 @@ def main(config: DictConfig): config.name + "_" + postfix + ".png", ) - return - - -import os if __name__ == "__main__": main() diff --git a/pdebench/data_gen/src/_attic/grf.py b/pdebench/data_gen/src/_attic/grf.py index 5004f65..09ce407 100644 --- a/pdebench/data_gen/src/_attic/grf.py +++ b/pdebench/data_gen/src/_attic/grf.py @@ -1,17 +1,19 @@ +from __future__ import annotations + import jax import jax.numpy as jnp -def grf( - seed: int = 1234, - xdim: int = 256, - ydim: int = 256, - sigma: float = 0.1, - rho: float = 0.1, - n: int = 1, - ): +def grf( + seed: int = 1234, + xdim: int = 256, + ydim: int = 256, + sigma: float = 0.1, + rho: float = 0.1, + n: int = 1, +): """ - Variables seeting for random + Variables seeding for random seed : random seed sigma : scale(?) @@ -25,21 +27,21 @@ def grf( """ rng = jax.random.PRNGKey(seed) fx, fy = jnp.meshgrid( - jnp.fft.fftfreq(xdim) * xdim, - jnp.fft.rfftfreq(ydim) * ydim, - indexing = 'ij') + jnp.fft.fftfreq(xdim) * xdim, jnp.fft.rfftfreq(ydim) * ydim, indexing="ij" + ) nfx, nfy = fx.shape fnorm = jnp.sqrt(fx**2 + fy**2) - power = jnp.exp(-(fnorm**2/rho)) - gain = jnp.sqrt(sigma**2 * power/power.sum()) # Lazy not calculating normalisation + power = jnp.exp(-(fnorm**2 / rho)) + gain = jnp.sqrt( + sigma**2 * power / power.sum() + ) # Lazy not calculating normalisation noise = ( jax.random.normal(rng, (n, nfx, nfy)) + jax.random.normal(rng, (n, nfx, nfy)) * 1j ) - noise = noise.at[...,0].set(jnp.abs(noise[..., 0])) - ## TODO: This is the rbf kernel; Matern kernel has more plausible smoothness. - ## Matern 3/2 PSD is - #(18 * jnp.sqrt(3)* jnp.pi * sigma**2)/((4 * k^2 * jnp.pi**2 + 3/(rho**2))^(5/2) rho^3) - field = jnp.fft.irfft2(noise * gain, (xdim, ydim), norm="forward") - return field + noise = noise.at[..., 0].set(jnp.abs(noise[..., 0])) + # TODO: This is the rbf kernel; Matern kernel has more plausible smoothness. + # Matern 3/2 PSD is + # (18 * jnp.sqrt(3)* jnp.pi * sigma**2)/((4 * k^2 * jnp.pi**2 + 3/(rho**2))^(5/2) rho^3) + return jnp.fft.irfft2(noise * gain, (xdim, ydim), norm="forward") diff --git a/pdebench/data_gen/src/data_io.py b/pdebench/data_gen/src/data_io.py index 64ea494..8004258 100644 --- a/pdebench/data_gen/src/data_io.py +++ b/pdebench/data_gen/src/data_io.py @@ -1,128 +1,123 @@ -import os, os.path -import subprocess +from __future__ import annotations + import json import logging +import subprocess -from phi.flow import * -from phi.field import Field -from phi.math import Shape - -import numpy as np import h5py - -from omegaconf import DictConfig, OmegaConf - +import numpy as np +from omegaconf import OmegaConf +from phi.field import Field +from phi.flow import CenteredGrid +from phi.math import Shape log = logging.getLogger(__name__) -def dims_for( - n_steps=1000, - grid_size=(100, 100), - frame_int = 1, - n_batch = 1, - **kwargs): +def dims_for(n_steps=1000, grid_size=(100, 100), frame_int=1, n_batch=1, **kwargs): """ return a dict of fields and their shapes """ - n_frames = ((n_steps-1)//frame_int) + 1 - return dict( - velocity = (n_batch, n_frames, *grid_size, len(grid_size)), - particles= (n_batch, n_frames, *grid_size, 1), - force= (n_batch, *grid_size, len(grid_size)), - t= (n_batch, n_frames), - ) + n_frames = ((n_steps - 1) // frame_int) + 1 + return { + "velocity": (n_batch, n_frames, *grid_size, len(grid_size)), + "particles": (n_batch, n_frames, *grid_size, 1), + "force": (n_batch, *grid_size, len(grid_size)), + "t": (n_batch, n_frames), + } def dict_for(config): spec = dims_for(**config) - data_store = dict( - latest_index = -1, - config = config - ) + data_store = {"latest_index": -1, "config": config} for field_name, full_shape in spec.items(): - data_store[field_name] = np.ndarray(full_shape, dtype='float32') + data_store[field_name] = np.ndarray(full_shape, dtype="float32") return data_store def h5_for(config): - log.info(f"config: {config}") + log.info("config: %s", config) spec = dims_for(**config) - log.info(f"spec: {spec}") + log.info("spec: %s", spec) fname = f"{config['sim_name']}-{config['seed']}.h5" - data_store = h5py.File(fname, 'a') - data_store.attrs['config'] = OmegaConf.to_yaml(config) - data_store.attrs['latestIndex'] = -1 + data_store = h5py.File(fname, "a") + data_store.attrs["config"] = OmegaConf.to_yaml(config) + data_store.attrs["latestIndex"] = -1 for field_name, full_shape in spec.items(): # dataset shape is (batch, t_length, x1, ..., xd, v) - chunk_shape = (1, 1, *full_shape[2:]) # chunk shape in (1, 1, x1, ..., xd, v) - # Open a dataset, creating it if it doesnโ€™t exist. - data_store.require_dataset( + chunk_shape = (1, 1, *full_shape[2:]) # chunk shape in (1, 1, x1, ..., xd, v) + # Open a dataset, creating it if it doesn't exist. + data_store.require_dataset( field_name, full_shape, - 'float32', + "float32", compression="lzf", chunks=chunk_shape, - shuffle=True) + shuffle=True, + ) return data_store def to_centre_grid(field: Field) -> CenteredGrid: - ''' + """ resample the input `Field` and return a corresponding `CenterGrid` used because the `StaggeredGrid`, which is usually the Object for velocity, does pack into nice tensors for typical neural nets - ''' + """ if isinstance(field, CenteredGrid): return field return CenteredGrid(field, resolution=field.shape.spatial, bounds=field.bounds) def _get_dim_order(shape: Shape): - ''' + """ Return a tuple of string, represents the order of dimensions e.g. ('batch','x','y','vector') If the current Shape does not have channel dims, fill in "vector" as 1. - ''' - batchNames = shape.batch.names if (shape.batch_rank > 0) else ('batch',) - channelNames = shape.channel.names if (shape.channel_rank > 0) else ('vector',) + """ + batchNames = shape.batch.names if (shape.batch_rank > 0) else ("batch",) + channelNames = shape.channel.names if (shape.channel_rank > 0) else ("vector",) return batchNames + shape.spatial.names + channelNames def to_ndarray(field: Field) -> np.ndarray: - ''' + """ Turn the current Field into ndarray, with shape (batch, x1, ..., xd, v) - ''' + """ centered = to_centre_grid(field) order = _get_dim_order(centered.shape) - ndarray = centered.values.numpy(order=order) - return ndarray + return centered.values.numpy(order=order) def dataverse_upload( - file_path, - dataverse_url, - dataverse_token, - dataverse_id, - dataverse_dir=None, - retry=10): - ''' + file_path, + dataverse_url, + dataverse_token, + dataverse_id, + dataverse_dir=None, + retry=10, +): + """ Upload a file to dataverse - ''' - darus_struct = { - "description":"", - "categories":["Data"], - "restrict": "false" - } + """ + darus_struct = {"description": "", "categories": ["Data"], "restrict": "false"} if dataverse_dir is not None: darus_struct["directoryLabel"] = f"{dataverse_dir}/" cmd = [ "curl", - "-X", "POST", - "-H", f"X-Dataverse-key:{dataverse_token}", - "-F", f"file=@{file_path}", - "-F", 'jsonData='+json.dumps(darus_struct), + "-X", + "POST", + "-H", + f"X-Dataverse-key:{dataverse_token}", + "-F", + f"file=@{file_path}", + "-F", + "jsonData=" + json.dumps(darus_struct), f"{dataverse_url}/api/datasets/:persistentId/add?persistentId={dataverse_id}", - "--retry", str(retry)] - log.info(f"upload cmd {cmd}") + "--retry", + str(retry), + ] + log.info("upload cmd %s", cmd) subprocess.Popen(cmd) - log.info(f"upload cmd {os.getcwd()}$ {' '.join(cmd)}") + from pathlib import Path + + log.info("upload cmd %s$ %s", Path.cwd(), " ".join(cmd)) diff --git a/pdebench/data_gen/src/plots.py b/pdebench/data_gen/src/plots.py index 29b95c5..ec94db2 100644 --- a/pdebench/data_gen/src/plots.py +++ b/pdebench/data_gen/src/plots.py @@ -2,11 +2,14 @@ Author : John Kim, Simon Brown, Timothy Praditia PDE Simulation packages """ + +from __future__ import annotations + +import h5py +import imageio import matplotlib.pyplot as plt import numpy as np -import imageio import phi.vis as phivis -import os def plot_data(data, t, dim, channel, t_fraction, config, filename): @@ -16,7 +19,8 @@ def plot_data(data, t, dim, channel, t_fraction, config, filename): plt.figure() plt.title(f"$t={t[t_idx]}$") if dim == 1: - x = np.array(h5_file["grid"]["x"], dtype="f") + with h5py.File(config.data_path, "r") as h5_file: + x = np.array(h5_file["grid"]["x"], dtype="f") plt.plot(x.squeeze(), data[t_idx, ..., channel]) plt.xlabel("$x$") else: @@ -37,7 +41,7 @@ def plot_data(data, t, dim, channel, t_fraction, config, filename): plt.savefig(filename) -def save_phi_plot(result, title, filepath, bbox_inches='tight', pad_inches=0): +def save_phi_plot(result, title, filepath, bbox_inches="tight", pad_inches=0): """ save one custom figure from an array """ @@ -47,40 +51,50 @@ def save_phi_plot(result, title, filepath, bbox_inches='tight', pad_inches=0): plt.close() -def phi_plots(results, T_results, title, filepath, scale = 1, bbox_inches='tight', pad_inches=0): +def phi_plots( + results, T_results, title, filepath, scale=1, bbox_inches="tight", pad_inches=0 +): """ Save simulation custom figures, get images list """ images = [] upperfilepath = filepath - for i, arr in enumerate(T_results): - filename = '{}.png'.format(title) - if upperfilepath == '': - filepath = filename - else: - filepath = upperfilepath + '/{}'.format(filename) + for i, _ in enumerate(T_results): + filename = f"{title}.png" + filepath = filename if upperfilepath == "" else upperfilepath + f"/{filename}" save_phi_plot( - scale * results[i], title, filepath, bbox_inches=bbox_inches, pad_inches=pad_inches) + scale * results[i], + title, + filepath, + bbox_inches=bbox_inches, + pad_inches=pad_inches, + ) images.append(imageio.imread(filepath)) return images -def save_sim_figures(results, T_results, simulation_name, kinematic_value, filepath, scale = 1, bbox_inches='tight', pad_inches=0): +def save_sim_figures( + T_results, + simulation_name, + kinematic_value, + filepath, + scale=1, + bbox_inches="tight", + pad_inches=0, +): """ save figures, get images list """ images = [] upperfilepath = filepath - for i, arr in enumerate(T_results): + for _, arr in enumerate(T_results): res = arr[0] - title = '{}_{}_t={}'.format(simulation_name, kinematic_value, round(T_results[i], 2)) - filename = '{}.png'.format(title) - if upperfilepath == '': - filepath = filename - else: - filepath = upperfilepath + '/{}'.format(filename) + title = f"{simulation_name}_{kinematic_value}_t={round(arr, 2)}" + filename = f"{title}.png" + filepath = filename if upperfilepath == "" else upperfilepath + f"/{filename}" save_phi_plot( - scale * res, title, filepath, bbox_inches=bbox_inches, pad_inches=pad_inches) + scale * res, title, filepath, bbox_inches=bbox_inches, pad_inches=pad_inches + ) images.append(imageio.imread(filepath)) return images diff --git a/pdebench/data_gen/src/pytorch_dataset.py b/pdebench/data_gen/src/pytorch_dataset.py index 4cd2c6b..4c84da4 100644 --- a/pdebench/data_gen/src/pytorch_dataset.py +++ b/pdebench/data_gen/src/pytorch_dataset.py @@ -1,7 +1,13 @@ -import h5py -from torch.utils.data import Dataset, DataLoader +from __future__ import annotations + +import logging from pathlib import Path + +import h5py from pytorch_lightning import LightningDataModule +from torch.utils.data import DataLoader, Dataset + +logger = logging.getLogger(__name__) class HDF5Dataset(Dataset): @@ -14,7 +20,7 @@ def __init__(self, dir_path, transform=None): super().__init__() path = Path(dir_path) assert path.is_dir() - files_path = list(path.glob('*.h5')) # all .h5 files' path + files_path = list(path.glob("*.h5")) # all .h5 files' path assert len(files_path) > 0 self.data_info = {} @@ -23,9 +29,9 @@ def __init__(self, dir_path, transform=None): self.config = [] self.names = [] - for files_path in files_path: - with h5py.File(str(files_path.resolve())) as f: - config = f.attrs.get('config') + for file_path in files_path: + with h5py.File(str(file_path.resolve())) as f: + config = f.attrs.get("config") for ds_name, ds in f.items(): self.names.append(ds_name) b = ds.shape[0] @@ -60,6 +66,7 @@ def _load_data(self, idx): # PATH_DATASETS = 'dummy_dataset' + class HDF5DatasetLightning(LightningDataModule): def __init__(self, data_dir: str, batch_size: int = 64, transforms=None): super().__init__() @@ -73,12 +80,11 @@ def setup(self, stage=None): self.train = HDF5Dataset(self.data_dir, transform=self.transforms) def train_dataloader(self): - print(self.train is None) return DataLoader(self.train, batch_size=self.batch_size) if __name__ == "__main__": - dir_path = 'download_dataset' # random_force_field--ns_sim--10.h5 in this directory + dir_path = "download_dataset" # random_force_field--ns_sim--10.h5 in this directory # test pytorch dataset dataset = HDF5Dataset(dir_path=dir_path, transform=None) @@ -86,8 +92,8 @@ def train_dataloader(self): dataloader = DataLoader(dataset, batch_size=64, shuffle=True) data, config = next(iter(dataloader)) for i, d in enumerate(data): - print(f'{names[i].upper()} batched data shape: ', d.size()) - print('number of config files: ', len(config)) + logger.info("%s batched data shape: %s", names[i].upper(), d.size()) + logger.info("number of config files: %s", len(config)) # test pytorch lightning dataset lightning_dataset = HDF5DatasetLightning(dir_path, batch_size=64, transforms=None) @@ -95,5 +101,5 @@ def train_dataloader(self): lightning_dataloader = lightning_dataset.train_dataloader() data, config = next(iter(lightning_dataloader)) for i, d in enumerate(data): - print(f'{names[i].upper()} batched data shape: ', d.size()) - print('number of config files: ', len(config)) + logger.info("%s batched data shape: %s", names[i].upper(), d.size()) + logger.info("number of config files: %s", len(config)) diff --git a/pdebench/data_gen/src/sim_diff_react.py b/pdebench/data_gen/src/sim_diff_react.py index 2d16185..7c175e6 100644 --- a/pdebench/data_gen/src/sim_diff_react.py +++ b/pdebench/data_gen/src/sim_diff_react.py @@ -1,25 +1,29 @@ +from __future__ import annotations + +import logging + import numpy as np from scipy.integrate import solve_ivp from scipy.sparse import diags -import logging + class Simulator: - - def __init__(self, - Du: float = 1E-3, - Dv: float = 5E-3, - k: float = 5E-3, - t: float = 50, - tdim: int = 501, - x_left: float = -1.0, - x_right: float = 1.0, - xdim: int = 50, - y_bottom: float = -1.0, - y_top: float = 1.0, - ydim: int = 50, - n: int = 1, - seed: int = 0): - + def __init__( + self, + Du: float = 1e-3, + Dv: float = 5e-3, + k: float = 5e-3, + t: float = 50, + tdim: int = 501, + x_left: float = -1.0, + x_right: float = 1.0, + xdim: int = 50, + y_bottom: float = -1.0, + y_top: float = 1.0, + ydim: int = 50, + n: int = 1, # noqa: ARG002 + seed: int = 0, + ): """ Constructor method initializing the parameters for the diffusion sorption problem. @@ -47,70 +51,72 @@ def __init__(self, self.X1 = x_right self.Y0 = y_bottom self.Y1 = y_top - + self.Nx = xdim self.Ny = ydim self.Nt = tdim - - # Calculate grid size and generate grid - self.dx = (self.X1 - self.X0)/(self.Nx) - self.dy = (self.Y1 - self.Y0)/(self.Ny) - - self.x = np.linspace(self.X0 + self.dx/2, self.X1 - self.dx/2, self.Nx) - self.y = np.linspace(self.Y0 + self.dy/2, self.Y1 - self.dy/2, self.Ny) - + + # Calculate grid size and generate grid + self.dx = (self.X1 - self.X0) / (self.Nx) + self.dy = (self.Y1 - self.Y0) / (self.Ny) + + self.x = np.linspace(self.X0 + self.dx / 2, self.X1 - self.dx / 2, self.Nx) + self.y = np.linspace(self.Y0 + self.dy / 2, self.Y1 - self.dy / 2, self.Ny) + # Time steps to store the simulation results self.t = np.linspace(0, self.T, self.Nt) - + # Initialize the logger self.log = logging.getLogger(__name__) - + self.seed = seed - + def generate_sample(self): """ Single sample generation using the parameters of this simulator. :return: The generated sample as numpy array(t, x, y, num_features) """ - - np.random.seed(self.seed) - - u0 = np.random.randn(self.Nx*self.Ny) - v0 = np.random.randn(self.Nx*self.Ny) - - u0 = u0.reshape(self.Nx*self.Ny) - v0 = v0.reshape(self.Nx*self.Ny) - u0 = np.concatenate((u0,v0)) - + + rng = np.random.default_rng(self.seed) + + u0 = rng.standard_normal(self.Nx * self.Ny) + v0 = rng.standard_normal(self.Nx * self.Ny) + + u0 = u0.reshape(self.Nx * self.Ny) + v0 = v0.reshape(self.Nx * self.Ny) + u0 = np.concatenate((u0, v0)) + # # Normalize u0 # u0 = 2 * (u0 - u0.min()) / (u0.max() - u0.min()) - 1 # Generate arrays as diagonal inputs to the Laplacian matrix - main_diag = -2*np.ones(self.Nx)/self.dx**2 -2*np.ones(self.Nx)/self.dy**2 - main_diag[0] = -1/self.dx**2 -2/self.dy**2 - main_diag[-1] = -1/self.dx**2 -2/self.dy**2 + main_diag = ( + -2 * np.ones(self.Nx) / self.dx**2 - 2 * np.ones(self.Nx) / self.dy**2 + ) + main_diag[0] = -1 / self.dx**2 - 2 / self.dy**2 + main_diag[-1] = -1 / self.dx**2 - 2 / self.dy**2 main_diag = np.tile(main_diag, self.Ny) - main_diag[:self.Nx] = -2/self.dx**2 -1/self.dy**2 - main_diag[self.Nx*(self.Ny-1):] = -2/self.dx**2 -1/self.dy**2 - main_diag[0] = -1/self.dx**2 -1/self.dy**2 - main_diag[self.Nx-1] = -1/self.dx**2 -1/self.dy**2 - main_diag[self.Nx*(self.Ny-1)] = -1/self.dx**2 -1/self.dy**2 - main_diag[-1] = -1/self.dx**2 -1/self.dy**2 - + main_diag[: self.Nx] = -2 / self.dx**2 - 1 / self.dy**2 + main_diag[self.Nx * (self.Ny - 1) :] = -2 / self.dx**2 - 1 / self.dy**2 + main_diag[0] = -1 / self.dx**2 - 1 / self.dy**2 + main_diag[self.Nx - 1] = -1 / self.dx**2 - 1 / self.dy**2 + main_diag[self.Nx * (self.Ny - 1)] = -1 / self.dx**2 - 1 / self.dy**2 + main_diag[-1] = -1 / self.dx**2 - 1 / self.dy**2 + left_diag = np.ones(self.Nx) left_diag[0] = 0 left_diag = np.tile(left_diag, self.Ny) - left_diag = left_diag[1:]/self.dx**2 - + left_diag = left_diag[1:] / self.dx**2 + right_diag = np.ones(self.Nx) right_diag[-1] = 0 right_diag = np.tile(right_diag, self.Ny) - right_diag = right_diag[:-1]/self.dx**2 - - bottom_diag = np.ones(self.Nx*(self.Ny-1))/self.dy**2 - - top_diag = np.ones(self.Nx*(self.Ny-1))/self.dy**2 - + right_diag = right_diag[:-1] / self.dx**2 + + bottom_diag = np.ones(self.Nx * (self.Ny - 1)) / self.dy**2 + + top_diag = np.ones(self.Nx * (self.Ny - 1)) / self.dy**2 + # Generate the sparse Laplacian matrix diagonals = [main_diag, left_diag, right_diag, bottom_diag, top_diag] offsets = [0, -1, 1, -self.Nx, self.Nx] @@ -120,35 +126,34 @@ def generate_sample(self): prob = solve_ivp(self.rc_ode, (0, self.T), u0, t_eval=self.t) ode_data = prob.y - sample_u = np.transpose(ode_data[:self.Nx*self.Ny]).reshape(-1,self.Ny,self.Nx) - sample_v = np.transpose(ode_data[self.Nx*self.Ny:]).reshape(-1,self.Ny,self.Nx) + sample_u = np.transpose(ode_data[: self.Nx * self.Ny]).reshape( + -1, self.Ny, self.Nx + ) + sample_v = np.transpose(ode_data[self.Nx * self.Ny :]).reshape( + -1, self.Ny, self.Nx + ) - return np.stack((sample_u, sample_v),axis=-1) + return np.stack((sample_u, sample_v), axis=-1) - def rc_ode(self, t, y): + def rc_ode(self, t, y): # noqa: ARG002 """ Solves a given equation for a particular time step. :param t: The current time step :param y: The equation values to solve :return: A finite volume solution """ - + # Separate y into u and v - u = y[:self.Nx*self.Ny] - v = y[self.Nx*self.Ny:] - + u = y[: self.Nx * self.Ny] + v = y[self.Nx * self.Ny :] + # Calculate reaction function for each unknown react_u = u - u**3 - self.k - v react_v = u - v - + # Calculate time derivative for each unknown u_t = react_u + self.Du * (self.lap @ u) v_t = react_v + self.Dv * (self.lap @ v) - + # Stack the time derivative into a single array y_t - y_t = np.concatenate((u_t,v_t)) - - # Log the simulation progress - self.log.info('t = ' + str(t)) - - return y_t + return np.concatenate((u_t, v_t)) diff --git a/pdebench/data_gen/src/sim_diff_sorp.py b/pdebench/data_gen/src/sim_diff_sorp.py index 5a5f0eb..b4c2e1c 100644 --- a/pdebench/data_gen/src/sim_diff_sorp.py +++ b/pdebench/data_gen/src/sim_diff_sorp.py @@ -1,25 +1,28 @@ +from __future__ import annotations + +import logging + import numpy as np from scipy.integrate import solve_ivp from scipy.sparse import diags -import logging + class Simulator: - - def __init__(self, - D: float = 5E-4, - por: float = 0.29, - rho_s: float = 2880, - k_f: float = 3.5E-4, - n_f: float = 0.874, - sol: float = 1.0, - t: float = 2500, - tdim: int = 501, - x_left: float = 0.0, - x_right: float = 1.0, - xdim: int = 50, - n: int = 1, - seed: int = 0): - + def __init__( + self, + D: float = 5e-4, + por: float = 0.29, + rho_s: float = 2880, + k_f: float = 3.5e-4, + n_f: float = 0.874, + sol: float = 1.0, + t: float = 2500, + tdim: int = 501, + x_left: float = 0.0, + x_right: float = 1.0, + xdim: int = 50, + seed: int = 0, + ): """ Constructor method initializing the parameters for the diffusion sorption problem. @@ -34,7 +37,6 @@ def __init__(self, :param x_left: Left end of the 2D simulation field :param x_right: Right end of the 2D simulation field :param xdim: Number of spatial steps between x_left and x_right - :param n: Number of batches """ # Set class parameters @@ -48,82 +50,73 @@ def __init__(self, self.T = t self.X0 = x_left self.X1 = x_right - + self.Nx = xdim self.Nt = tdim - - # Calculate grid size and generate grid - self.dx = (self.X1 - self.X0)/(self.Nx) - self.x = np.linspace(self.X0 + self.dx/2, self.X1 - self.dx/2, self.Nx) - + + # Calculate grid size and generate grid + self.dx = (self.X1 - self.X0) / (self.Nx) + self.x = np.linspace(self.X0 + self.dx / 2, self.X1 - self.dx / 2, self.Nx) + # Time steps to store the simulation results self.t = np.linspace(0, self.T, self.Nt) - + # Initialize the logger self.log = logging.getLogger(__name__) - + self.seed = seed - - def generate_sample(self): + + def generate_sample(self) -> np.ndarray: """ Single sample generation using the parameters of this simulator. :return: The generated sample as numpy array(t, x, y, num_features) """ - np.random.seed(self.seed) - + generator = np.random.default_rng(self.seed) # Generate initial condition - u0 = np.ones(self.Nx) * np.random.uniform(0,0.2) + u0 = np.ones(self.Nx) * generator.uniform(0, 0.2) # Generate arrays as diagonal inputs to the Laplacian matrix - main_diag = -2*np.ones(self.Nx)/self.dx**2 - - left_diag = np.ones(self.Nx-1)/self.dx**2 - - right_diag = np.ones(self.Nx-1)/self.dx**2 - + main_diag = -2 * np.ones(self.Nx) / self.dx**2 + + left_diag = np.ones(self.Nx - 1) / self.dx**2 + + right_diag = np.ones(self.Nx - 1) / self.dx**2 + # Generate the sparse Laplacian matrix diagonals = [main_diag, left_diag, right_diag] offsets = [0, -1, 1] self.lap = diags(diagonals, offsets) - + # Initialize the right hand side to account for the boundary condition self.rhs = np.zeros(self.Nx) # Solve the diffusion reaction problem prob = solve_ivp(self.rc_ode, (0, self.T), u0, t_eval=self.t) ode_data = prob.y - + sample_c = np.transpose(ode_data) - + return np.expand_dims(sample_c, axis=-1) - def rc_ode(self, t, y): + def rc_ode(self, t: float, y): # noqa: ARG002 """ Solves a given equation for a particular time step. :param t: The current time step :param y: The equation values to solve :return: A finite volume solution """ - - c = y - + # Define left and right boundary conditions left_BC = self.sol - right_BC = (c[-2]-c[-1])/self.dx * self.D - + right_BC = (y[-2] - y[-1]) / self.dx * self.D + # Calculate the Freundlich retardation factor - retardation = 1 + ((1 - self.por)/self.por)*self.rho_s\ - *self.k_f*self.n_f*(c + 1e-6)**(self.n_f-1) - + retardation = 1 + ( + (1 - self.por) / self.por + ) * self.rho_s * self.k_f * self.n_f * (y + 1e-6) ** (self.n_f - 1) + # Calculate the right hand side - self.rhs[0] = self.D/retardation[0]/(self.dx**2)*left_BC - self.rhs[-1] = self.D/retardation[-1]/(self.dx**2)*right_BC - - # Calculate time derivative - c_t = self.D/retardation * (self.lap @ c) + self.rhs - y_t = c_t - - # Log the simulation progress - self.log.info('t = ' + str(t)) - - return y_t + self.rhs[0] = self.D / retardation[0] / (self.dx**2) * left_BC + self.rhs[-1] = self.D / retardation[-1] / (self.dx**2) * right_BC + + return self.D / retardation * (self.lap @ y) + self.rhs diff --git a/pdebench/data_gen/src/sim_ns_incomp_2d.py b/pdebench/data_gen/src/sim_ns_incomp_2d.py index 12b8c58..092243f 100644 --- a/pdebench/data_gen/src/sim_ns_incomp_2d.py +++ b/pdebench/data_gen/src/sim_ns_incomp_2d.py @@ -2,83 +2,119 @@ Author : John Kim, Ran Zhang, Dan MacKinlay PDE Simulation packages """ -from contextlib import nullcontext -from typing import Optional + +from __future__ import annotations + import logging import os +from contextlib import nullcontext +from typing import Optional import imageio import numpy as np -from tqdm import tqdm -import hydra -from src import data_io -from src import image_processor - -log = logging.getLogger(__name__) +from pdebench.data_gen.src import data_io +from tqdm import tqdm +logger = logging.getLogger(__name__) # import wandb -def call_many(fn_list, *args, **kwargs): + +def call_many(fn_list, *args, **kwargs) -> list[callable]: """ Surely there is already a helper function for this somewhere. inverse map? """ return [fn(*args, **kwargs) for fn in fn_list] - + def ns_sim( - seed: int, - label: Optional[str]=None, - sim_name: str='ns_sim_2d', - particle_extrapolation:str='BOUNDARY', - velocity_extrapolation:str='ZERO', - NU: float=0.01, - scale: float= 10.0, - smoothness: float=3.0, - grid_size=(100,100), - enable_gravity: bool=False, - enable_obstacles: bool=False, - force_extrapolation: str='ZERO', - save_images: bool=True, - save_gif: bool=True, - save_h5:bool=True, - n_steps: int=10, - DT: float=0.01, - frame_int: int=1, - n_batch=1, - backend='jax', - device='GPU', - jit=True, - profile: bool=False, - upload: bool=False, - exec_dir: Optional[str]=None, - artefact_dir: Optional[str] = None, #hackish way of writing artefacts to a good location without fighting hydra's conf interpolation - dataverse: Optional[dict] = None, - config={}, - ): + seed: int, + label: Optional[str] = None, + sim_name: str = "ns_sim_2d", + particle_extrapolation: str = "BOUNDARY", + velocity_extrapolation: str = "ZERO", + NU: float = 0.01, + scale: float = 10.0, + smoothness: float = 3.0, + grid_size=(100, 100), + enable_gravity: bool = False, + enable_obstacles: bool = False, + force_extrapolation: str = "ZERO", + save_images: bool = True, + save_gif: bool = True, + save_h5: bool = True, + n_steps: int = 10, + DT: float = 0.01, + frame_int: int = 1, + n_batch: int = 1, + backend: str = "jax", + device: str = "GPU", + jit: bool = True, + profile: bool = False, + upload: bool = False, + dataverse: Optional[dict] = None, + config: dict | None = None, +): """ Run the actual simulation. """ - # log.info(f"exec_dir {exec_dir}") - # log.info(f"orig_dir {hydra.utils.get_original_cwd()}") - # log.info(f"artefact_dir {artefact_dir}") - # log.info(f"WORKING_DIR {os.getenv('WORKING_DIR')}") - # log.info(f"cwd {os.getcwd()}") - - if backend == 'jax': - from phi.jax.flow import extrapolation, Box, Obstacle, advect, diffuse, fluid, math, Solve, CenteredGrid, StaggeredGrid, Noise, batch - elif backend == 'pytorch': - from phi.torch.flow import extrapolation, Box, Obstacle, advect, diffuse, fluid, math, Solve, CenteredGrid, StaggeredGrid, Noise, batch + if config is None: + config = {} + + if backend == "jax": + from phi.jax.flow import ( + Box, + CenteredGrid, + Noise, + Obstacle, + Solve, + StaggeredGrid, + advect, + batch, + diffuse, + extrapolation, + fluid, + math, + ) + elif backend == "pytorch": from phi.torch import TORCH + from phi.torch.flow import ( + Box, + CenteredGrid, + Noise, + Obstacle, + Solve, + StaggeredGrid, + advect, + batch, + diffuse, + extrapolation, + fluid, + math, + ) + TORCH.set_default_device(device) else: - from phi.flow import extrapolation, Box, Obstacle, advect, diffuse, fluid, math, Solve, CenteredGrid, StaggeredGrid, Noise, batch + from phi.flow import ( + Box, + CenteredGrid, + Noise, + Obstacle, + Solve, + StaggeredGrid, + advect, + batch, + diffuse, + extrapolation, + fluid, + math, + ) # from torch.profiler import profile, record_function, ProfilerActivity - def bounds_select(x, y): + def bounds_select(x: int | None = None, y: int | None = None) -> Box: """ This function generates a 2D Phiflow Box with scale from zero to the number indicated by the bounds or infinite. @@ -89,16 +125,15 @@ def bounds_select(x, y): Returns: Box: A Box type of Phiflow """ - if x == None: + if x is None: return Box[:, 0:y] - elif y == None: + if y is None: return Box[0:x, :] - else: - return Box[0:x, 0:y] - + return Box[0:x, 0:y] - - def cauchy_momentum_step(velocity, particles, body_acceleration, NU, DT, obstacles = None): + def cauchy_momentum_step( + velocity, particles, body_acceleration, NU, DT, obstacles: tuple | None = None + ) -> tuple[fluid.Field, fluid.Field]: """ Navier-Stokes Simulation cauchy_momentum_step returns velocity and particles by solving cauchy momentum equation for one step @@ -112,23 +147,25 @@ def cauchy_momentum_step(velocity, particles, body_acceleration, NU, DT, obstacl **kwargs : Other obstacles (Simulation constraints etc) """ # Set empty obstacle as empty tuple - if obstacles == None: + if obstacles is None: obstacles = () # Computing velocity term first # Cauchy-momentum equation - velocity = advect.semi_lagrangian(velocity, velocity, dt=DT) # advection - velocity = diffuse.explicit(velocity, NU, dt=DT) # diffusion - + velocity = advect.semi_lagrangian(velocity, velocity, dt=DT) # advection + velocity = diffuse.explicit(velocity, NU, dt=DT) # diffusion + # Add external body_acceleration, constraints - velocity += DT * particles * body_acceleration # external body_acceleration - velocity = fluid.apply_boundary_conditions(velocity, obstacles) # obstacles - + velocity += DT * particles * body_acceleration # external body_acceleration + velocity = fluid.apply_boundary_conditions(velocity, obstacles) # obstacles + # Make incompressible - velocity, _ = fluid.make_incompressible(velocity, obstacles, solve=Solve('CG-adaptive', 1e-3, 0, x0=None)) # particles - - # Computing particles term next + velocity, _ = fluid.make_incompressible( + velocity, obstacles, solve=Solve("CG-adaptive", 1e-3, 0, x0=None) + ) # particles + + # Computing particles term next particles = advect.semi_lagrangian(particles, velocity, dt=DT) - + return velocity, particles # Setting the random seed for simulation. This is a global seed for Phiflow. @@ -140,18 +177,16 @@ def cauchy_momentum_step(velocity, particles, body_acceleration, NU, DT, obstacl if save_h5: data_store = data_io.h5_for(config) h5_path = data_store.filename + def _store(frame_i, t, particles, velocity, **kwargs): - data_store['particles'][:, frame_i, ...] = data_io.to_ndarray(particles) - data_store['velocity'][:, frame_i, ...] = data_io.to_ndarray(velocity) - data_store['t'][:, frame_i] = t - data_store.attrs['latestIndex'] = frame_i - callbacks.append( - _store - ) - cleanups.append( - lambda *args, **kwargs: data_store.close() - ) - ## Move output to artefacts dir here + data_store["particles"][:, frame_i, ...] = data_io.to_ndarray(particles) + data_store["velocity"][:, frame_i, ...] = data_io.to_ndarray(velocity) + data_store["t"][:, frame_i] = t + data_store.attrs["latestIndex"] = frame_i + + callbacks.append(_store) + cleanups.append(lambda *args, **kwargs: data_store.close()) + # Move output to artefacts dir here # if artefact_dir is not None: # cleanups.append( # lambda *args, **kwargs: data_store.close() @@ -162,30 +197,30 @@ def _store(frame_i, t, particles, velocity, **kwargs): lambda *args, **kwargs: data_io.dataverse_upload( file_path=h5_path, dataverse_url=os.getenv( - 'DATAVERSE_URL', 'https://darus.uni-stuttgart.de'), - dataverse_token=os.getenv( - 'DATAVERSE_API_TOKEN', ''), + "DATAVERSE_URL", "https://darus.uni-stuttgart.de" + ), + dataverse_token=os.getenv("DATAVERSE_API_TOKEN", ""), dataverse_dir=label, - dataverse_id=dataverse['dataset_id'], + dataverse_id=dataverse["dataset_id"], ) ) else: data_store = data_io.dict_for(config) + def _store(frame_i, t, particles, velocity, **kwargs): - data_store['particles'][:, frame_i, ...] = data_io.to_ndarray(particles) - data_store['velocity'][:, frame_i, ...] = data_io.to_ndarray(velocity) - data_store['t'][:, frame_i] = t - callbacks.append( - _store - ) + data_store["particles"][:, frame_i, ...] = data_io.to_ndarray(particles) + data_store["velocity"][:, frame_i, ...] = data_io.to_ndarray(velocity) + data_store["t"][:, frame_i] = t + + callbacks.append(_store) if save_images: + def _save_img(frame_i, t, particles, velocity, **kwargs): particles_images.append() velocity_images.append() - callbacks.append( - _save_img - ) + + callbacks.append(_save_img) profile_ctx = nullcontext() if profile: @@ -194,8 +229,7 @@ def _save_img(frame_i, t, particles, velocity, **kwargs): with profile_ctx as prof: # Initialization of the particles (i.e. density of the flow) grid with a Phiflow Noise() method particles = CenteredGrid( - Noise(batch(batch=n_batch), - scale=scale, smoothness=smoothness), + Noise(batch(batch=n_batch), scale=scale, smoothness=smoothness), extrapolation=getattr(extrapolation, particle_extrapolation), x=grid_size[0], y=grid_size[1], @@ -213,17 +247,17 @@ def _save_img(frame_i, t, particles, velocity, **kwargs): # Initialization of the force grid. The force is also a staggered grid with a Phiflow Noise() method or using gravity if enable_gravity: - force = math.tensor(batch(batch=n_batch),[0, -9.81]) + force = math.tensor(batch(batch=n_batch), [0, -9.81]) else: force = StaggeredGrid( - Noise(batch(batch=n_batch),vector=2), + Noise(batch(batch=n_batch), vector=2), extrapolation=getattr(extrapolation, force_extrapolation), x=grid_size[0], y=grid_size[1], bounds=bounds_select(*grid_size), ) - data_store['force'][:,...] = data_io.to_ndarray(force) + data_store["force"][:, ...] = data_io.to_ndarray(force) # Set the obstacles. Obstacles are not enabled by default. obstacles = [] @@ -235,44 +269,53 @@ def _save_img(frame_i, t, particles, velocity, **kwargs): # Use "python gen_ns_incomp.py save_gif=false" for not saving .gif animation and only save to pictures call_many( - callbacks, frame_i=0, t=0.0, - velocity=velocity, particles=particles, prof=prof) + callbacks, + frame_i=0, + t=0.0, + velocity=velocity, + particles=particles, + prof=prof, + ) - def sim_step(velocity, particles): - return cauchy_momentum_step( - velocity, particles, force, NU, DT, obstacles) + def sim_step(velocity, particles) -> tuple[fluid.Field, fluid.Field]: + return cauchy_momentum_step(velocity, particles, force, NU, DT, obstacles) if jit: sim_step = math.jit_compile(sim_step) - ts = np.linspace(0, n_steps*DT, n_steps, endpoint=False) - n_steps_actual = ((n_steps-1)//frame_int) * frame_int + 1 + ts = np.linspace(0, n_steps * DT, n_steps, endpoint=False) + n_steps_actual = ((n_steps - 1) // frame_int) * frame_int + 1 ts = ts[1:n_steps_actual] - # log.info("ts: {}".format(ts)) for step, t in enumerate(tqdm(ts), start=1): velocity, particles = sim_step( - velocity, particles,) - + velocity, + particles, + ) + if step % frame_int == 0: - frame_i = step//frame_int - log.info(f"step {step} frame_i {frame_i}") + frame_i = step // frame_int + msg = f"step {step} frame_i {frame_i}" + logger.info(msg) call_many( - callbacks, frame_i=frame_i, t=t, - velocity=velocity, particles=particles, prof=prof) + callbacks, + frame_i=frame_i, + t=t, + velocity=velocity, + particles=particles, + prof=prof, + ) if save_images and save_gif: # Saving images into .gif animation imageio.mimsave( - "{}_velocity.gif".format(sim_name), + f"{sim_name}_velocity.gif", velocity_images, duration=DT, ) imageio.mimsave( - "{}_particles.gif".format(sim_name), + f"{sim_name}_particles.gif", particles_images, duration=DT, ) call_many(cleanups) - - diff --git a/pdebench/data_gen/src/sim_radial_dam_break.py b/pdebench/data_gen/src/sim_radial_dam_break.py index d0127d8..34408e6 100644 --- a/pdebench/data_gen/src/sim_radial_dam_break.py +++ b/pdebench/data_gen/src/sim_radial_dam_break.py @@ -1,18 +1,14 @@ -from abc import abstractmethod -from abc import ABC +from __future__ import annotations -import os -import sys -import time +import logging +from abc import ABC, abstractmethod +from pathlib import Path -import h5py import numpy as np import torch -from clawpack import riemann -from clawpack import pyclaw +from clawpack import pyclaw, riemann - -sys.path.append(os.path.join(sys.path[0], "..", "utils")) +logger = logging.getLogger(__name__) class Basic2DScenario(ABC): @@ -32,7 +28,7 @@ def __init__(self): self.set_boundary_conditions() self.set_initial_conditions() self.register_state_getters() - self.outdir = os.sep.join(["./", self.name.replace(" ", "") + "2D"]) + self.outdir = Path.cwd() / (self.name.replace(" ", "") + "2D") @abstractmethod def setup_solver(self): @@ -69,7 +65,7 @@ def __get_hu(self): def __get_hv(self): return self.claw_state.q[self.momentumId_y, :].tolist() - def register_state_getters(self): + def register_state_getters(self) -> None: self.state_getters = { "h": self.__get_h, "u": self.__get_u, @@ -78,11 +74,11 @@ def register_state_getters(self): "hv": self.__get_hv, } - def add_save_state(self): + def add_save_state(self) -> None: for key, getter in self.state_getters.items(): self.save_state[key].append(getter()) - def init_save_state(self, T, tsteps): + def init_save_state(self, T: float, tsteps: int) -> None: self.save_state = {} self.save_state["x"] = self.domain.grid.x.centers.tolist() self.save_state["y"] = self.domain.grid.y.centers.tolist() @@ -90,42 +86,46 @@ def init_save_state(self, T, tsteps): for key, getter in self.state_getters.items(): self.save_state[key] = [getter()] - def save_state_to_disk(self, data_f, seed): + def save_state_to_disk(self, data_f, seed_str) -> None: T = np.asarray(self.save_state["t"]) X = np.asarray(self.save_state["x"]) Y = np.asarray(self.save_state["y"]) H = np.expand_dims(np.asarray(self.save_state["h"]), -1) - - seed_str = str(seed).zfill(4) data_f.create_dataset(f"{seed_str}/data", data=H, dtype="f") data_f.create_dataset(f"{seed_str}/grid/x", data=X, dtype="f") data_f.create_dataset(f"{seed_str}/grid/y", data=Y, dtype="f") data_f.create_dataset(f"{seed_str}/grid/t", data=T, dtype="f") - def simulate(self, t): + def simulate(self, t) -> None: if all(v is not None for v in [self.domain, self.claw_state, self.solver]): self.solver.evolve_to_time(self.solution, t) else: - print("Simulate failed: No scenario defined.") + msg = "Simulate failed: No scenario defined." + logger.info(msg) - def run(self, T=1.0, tsteps=20, plot=False): + def run(self, T: float = 1.0, tsteps: int = 20) -> None: self.init_save_state(T, tsteps) self.solution = pyclaw.Solution(self.claw_state, self.domain) - dt = T / tsteps - start = time.time() + dt = T / float(tsteps) + for tstep in range(1, tsteps + 1): t = tstep * dt - print("Simulating timestep {}/{} at t={:f}".format(tstep, tsteps, t)) self.simulate(t) self.add_save_state() - print("Simulation took: {}".format(time.time() - start)) class RadialDamBreak2D(Basic2DScenario): name = "RadialDamBreak" - def __init__(self, xdim, ydim, grav=1.0, dam_radius=0.5, inner_height=2.0): + def __init__( + self, + xdim, + ydim, + grav: float = 1.0, + dam_radius: float = 0.5, + inner_height: float = 2.0, + ): self.depthId = 0 self.momentumId_x = 1 self.momentumId_y = 2 @@ -137,7 +137,7 @@ def __init__(self, xdim, ydim, grav=1.0, dam_radius=0.5, inner_height=2.0): super().__init__() # self.state_getters['bathymetry'] = self.__get_bathymetry - def setup_solver(self): + def setup_solver(self) -> None: rs = riemann.shallow_roe_with_efix_2D self.solver = pyclaw.ClawSolver2D(rs) self.solver.limiters = pyclaw.limiters.tvd.MC @@ -148,7 +148,7 @@ def setup_solver(self): self.momentumId_x = 1 self.momentumId_y = 2 - def create_domain(self): + def create_domain(self) -> None: self.xlower = -2.5 self.xupper = 2.5 self.ylower = -2.5 @@ -160,7 +160,7 @@ def create_domain(self): self.domain = pyclaw.Domain([x, y]) self.claw_state = pyclaw.State(self.domain, self.solver.num_eqn) - def set_boundary_conditions(self): + def set_boundary_conditions(self) -> None: """ Sets homogeneous Neumann boundary conditions at each end for q=(u, h*u) and for the bathymetry (auxiliary variable). @@ -170,8 +170,7 @@ def set_boundary_conditions(self): self.solver.bc_lower[1] = pyclaw.BC.extrap self.solver.bc_upper[1] = pyclaw.BC.extrap - @staticmethod - def initial_h(coords): + def initial_h(self, coords): x0 = 0.0 y0 = 0.0 x = coords[:, 0] @@ -182,17 +181,17 @@ def initial_h(coords): return h_in * (r <= self.dam_radius) + h_out * (r > self.dam_radius) @staticmethod - def initial_momentum_x(coords): + def initial_momentum_x() -> torch.Tensor: return torch.tensor(0.0) @staticmethod - def initial_momentum_y(coords): + def initial_momentum_y() -> torch.Tensor: return torch.tensor(0.0) def __get_bathymetry(self): return self.claw_state.aux[0, :].tolist() - def set_initial_conditions(self): + def set_initial_conditions(self) -> None: self.claw_state.problem_data["grav"] = self.grav xc = self.claw_state.grid.x.centers diff --git a/pdebench/data_gen/src/utils.py b/pdebench/data_gen/src/utils.py index 1193d91..bcde22a 100644 --- a/pdebench/data_gen/src/utils.py +++ b/pdebench/data_gen/src/utils.py @@ -1,25 +1,25 @@ -import logging -import warnings -from typing import List, Sequence +from __future__ import annotations + import os -import glob +from pathlib import Path from pprint import pprint from omegaconf import DictConfig, OmegaConf -def expand_path(path, unique=True): + +def expand_path(path): """ Resolve a path that may contain variables and user home directory references. """ - return os.path.expandvars(os.path.expanduser(path)) + return os.path.expandvars(Path(path).expanduser()) def matching_paths(glob_exp): """ return a list of paths matching a glob expression """ - path = os.path.expandvars(os.path.expanduser(glob_exp)) - return glob.glob(path) + path = os.path.expandvars(Path(glob_exp).expanduser()) + return list(Path(path).glob("*")) def resolve_path(path, idx=None, unique=True): @@ -28,23 +28,27 @@ def resolve_path(path, idx=None, unique=True): if "unique" is True, and there are many matches, panic. Otherwise return the result at index "idx", which could reasonably be 0 or -1; if it is, we sort the list of files """ - matches = matching_paths(path) + matches = matching_paths(path) if idx is None: idx = 0 else: matches = sorted(matches) - + if unique and len(matches) > 1: - raise ValueError("Too many matches for glob: {}".format(path)) - else: - try: - return matches[idx] - except IndexError: - raise FileNotFoundError("No matches for glob: {}".format(path)) + valerrmsg = f"Too many matches for glob: {path}" + raise ValueError(valerrmsg) + try: + return matches[idx] + except IndexError: + idxerrmsg = f"No matches for glob: {path}" + raise FileNotFoundError(idxerrmsg) from None -def print_config(config: DictConfig, resolve: bool = True,): +def print_config( + config: DictConfig, + resolve: bool = True, +): """ basic pretty-printer for omegaconf configs """ - pprint(OmegaConf.to_yaml(config, resolve=resolve)) + pprint(OmegaConf.to_yaml(config, resolve=resolve)) # noqa: T203 diff --git a/pdebench/data_gen/src/vorticity.py b/pdebench/data_gen/src/vorticity.py new file mode 100644 index 0000000..26cac35 --- /dev/null +++ b/pdebench/data_gen/src/vorticity.py @@ -0,0 +1,150 @@ +r"""Generate vorticity field :math:`\boldsymbol{\omega} = \nabla \times \boldsymbol{v}` given + velocity field :math:`\boldsymbol{v}` using numerical approximation. + +Assuming the velocitiy field of shape [n, sx, sy, sz, 3] (5D) consists of a trajectory of equidistant cells, +the vorticity field is calculated by using spectral derivatives and the Fast Fourier Transform +such that :math:`\mathcal{F}\{\frac{\partial f}{\partial x}\ = i \omega \mathcal{F}\{f\}}`. + +The code is inspired by +Brunton, S. L.,; Kutz, J. N. (2022). Data-Driven Science and Engineering: Machine Learning, Dynamical Systems, and Control (2nd ed.). +and adapted to operate on 5D data. + +This module provides the functions + - compute_spectral_vorticity_np (numpy) + - compute_spectral_vorticity_jnp (jax.numpy) + +for approximating the vorticity field. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np + + +def compute_spectral_vorticity_np( + velocities: np.ndarray, dx: float, dy: float, dz: float +) -> np.ndarray: + r"""Compute vorcitity field of a [n, sx, sy, sz, 3] field, + where n denotes the number of timesteps, sx, sy, sz are the number + of bins in x-, y-, z-direction. + + :param velocities: 5D Velocity field of shape [n, sx, sy, sz, 3]. + :param lx: length x-direction + :type lx: float + :param ly: length y-direction + :type ly: float + :param lz: length z-direction + :type lz: float + :type velocities: np.ndarray + :raises ValueError: If `velocities` is not a 5D array or + of shape [n, sx, sy, sz, 3]. + :return: Vorticity field :math:`\boldsymbol{\omega} \in \mathbb{R}^{n \times s_x \times s_y \times s_z \times 3}` + :rtype: np.ndarray + """ + + if velocities.ndim != 5 or velocities.shape[-1] != 3: + msg = "Expected 5D array of shape [n, sx, sy, sz, 3]!" + raise ValueError(msg) + + dx = abs(dx) + dy = abs(dy) + dz = abs(dz) + + vx = velocities[..., 0] + vy = velocities[..., 1] + vz = velocities[..., 2] + + fxy = np.fft.fft(vx, axis=2) + fyx = np.fft.fft(vy, axis=1) + fxz = np.fft.fft(vx, axis=3) + fzx = np.fft.fft(vz, axis=1) + fyz = np.fft.fft(vy, axis=3) + fzy = np.fft.fft(vz, axis=2) + + kappa_xy = 2.0 * np.pi * np.fft.fftfreq(n=fxy.shape[2], d=dy) + kappa_yx = 2.0 * np.pi * np.fft.fftfreq(n=fyx.shape[1], d=dx) + kappa_xz = 2.0 * np.pi * np.fft.fftfreq(n=fxz.shape[3], d=dz) + kappa_zx = 2.0 * np.pi * np.fft.fftfreq(n=fzx.shape[1], d=dx) + kappa_yz = 2.0 * np.pi * np.fft.fftfreq(n=fyz.shape[3], d=dz) + kappa_zy = 2.0 * np.pi * np.fft.fftfreq(n=fzy.shape[2], d=dy) + + vxy = np.fft.ifft(1j * kappa_xy[None, None, :, None] * fxy, axis=2).real + vyx = np.fft.ifft(1j * kappa_yx[None, :, None, None] * fyx, axis=1).real + vxz = np.fft.ifft(1j * kappa_xz[None, None, None, :] * fxz, axis=3).real + vzx = np.fft.ifft(1j * kappa_zx[None, :, None, None] * fzx, axis=1).real + vyz = np.fft.ifft(1j * kappa_yz[None, None, None, :] * fyz, axis=3).real + vzy = np.fft.ifft(1j * kappa_zy[None, None, :, None] * fzy, axis=2).real + + omega_x = vzy - vyz + omega_y = vxz - vzx + omega_z = vyx - vxy + + return np.concatenate( + [omega_x[..., None], omega_y[..., None], omega_z[..., None]], axis=-1 + ) + + +@jax.jit # type: ignore[misc] +def compute_spectral_vorticity_jnp( + velocities: jnp.ndarray, dx: float, dy: float, dz: float +) -> jnp.ndarray: + r"""Compute vorcitity field of a [n, sx, sy, sz, 3] field, + where n denotes the number of timesteps, sx, sy, sz are the number + of bins in x-, y-, z-direction. In this case computations are performed on GPU + + :param velocities: 5D Velocity field of shape [n, sx, sy, sz, 3]. + :param lx: length x-direction + :type lx: float + :param ly: length y-direction + :type ly: float + :param lz: length z-direction + :type lz: float + :type velocities: np.ndarray + :raises ValueError: If `velocities` is not a 5D array or + of shape [n, sx, sy, sz, 3]. + :return: Vorticity field :math:`\boldsymbol{\omega} \in \mathbb{R}^{n \times s_x \times s_y \times s_z \times 3}` + :rtype: np.ndarray + """ + + if velocities.ndim != 5 or velocities.shape[-1] != 3: + msg = "Expected 5D array of shape [n, sx, sy, sz, 3]!" + raise ValueError(msg) + + dx = abs(dx) + dy = abs(dy) + dz = abs(dz) + + vx = velocities[..., 0] + vy = velocities[..., 1] + vz = velocities[..., 2] + + fxy = jnp.fft.fft(vx, axis=2) + fyx = jnp.fft.fft(vy, axis=1) + fxz = jnp.fft.fft(vx, axis=3) + fzx = jnp.fft.fft(vz, axis=1) + fyz = jnp.fft.fft(vy, axis=3) + fzy = jnp.fft.fft(vz, axis=2) + + kappa_xy = 2.0 * jnp.pi * jnp.fft.fftfreq(n=fxy.shape[2], d=dy) + kappa_yx = 2.0 * jnp.pi * jnp.fft.fftfreq(n=fyx.shape[1], d=dx) + kappa_xz = 2.0 * jnp.pi * jnp.fft.fftfreq(n=fxz.shape[3], d=dz) + kappa_zx = 2.0 * jnp.pi * jnp.fft.fftfreq(n=fzx.shape[1], d=dx) + kappa_yz = 2.0 * jnp.pi * jnp.fft.fftfreq(n=fyz.shape[3], d=dz) + kappa_zy = 2.0 * jnp.pi * jnp.fft.fftfreq(n=fzy.shape[2], d=dy) + + vxy = jnp.fft.ifft(1j * kappa_xy[None, None, :, None] * fxy, axis=2).real + vyx = jnp.fft.ifft(1j * kappa_yx[None, :, None, None] * fyx, axis=1).real + vxz = jnp.fft.ifft(1j * kappa_xz[None, None, None, :] * fxz, axis=3).real + vzx = jnp.fft.ifft(1j * kappa_zx[None, :, None, None] * fzx, axis=1).real + vyz = jnp.fft.ifft(1j * kappa_yz[None, None, None, :] * fyz, axis=3).real + vzy = jnp.fft.ifft(1j * kappa_zy[None, None, :, None] * fzy, axis=2).real + + omega_x = vzy - vyz + omega_y = vxz - vzx + omega_z = vyx - vxy + + return jnp.concatenate( + [omega_x[..., None], omega_y[..., None], omega_z[..., None]], axis=-1 + ) diff --git a/pdebench/data_gen/uploader.py b/pdebench/data_gen/uploader.py index c111903..02eb262 100644 --- a/pdebench/data_gen/uploader.py +++ b/pdebench/data_gen/uploader.py @@ -1,29 +1,42 @@ -import subprocess +from __future__ import annotations + import json +import subprocess + def dataverse_upload( - file_path, - dataverse_url, - dataverse_token, - dataverse_dir, - dataverse_id, - log, - retry=10): - ''' + file_path, + dataverse_url, + dataverse_token, + dataverse_dir, + dataverse_id, + log, + retry=10, +): + """ Upload a file to dataverse - ''' + """ cmd = [ "curl", - "-X", "POST", - "-H", f"X-Dataverse-key:{dataverse_token}", - "-F", f"file=@{file_path}", - "-F", 'jsonData='+json.dumps({ - "description":"", - "directoryLabel":f"{dataverse_dir}/", - "categories":["Data"], - "restrict": "false" - }), + "-X", + "POST", + "-H", + f"X-Dataverse-key:{dataverse_token}", + "-F", + f"file=@{file_path}", + "-F", + "jsonData=" + + json.dumps( + { + "description": "", + "directoryLabel": f"{dataverse_dir}/", + "categories": ["Data"], + "restrict": "false", + } + ), f"{dataverse_url}/api/datasets/:persistentId/add?persistentId={dataverse_id}", - "--retry", str(retry)] + "--retry", + str(retry), + ] log.info(cmd) - subprocess.Popen(cmd) \ No newline at end of file + subprocess.Popen(cmd) diff --git a/pdebench/data_gen/velocity2vorticity.py b/pdebench/data_gen/velocity2vorticity.py new file mode 100644 index 0000000..807b3be --- /dev/null +++ b/pdebench/data_gen/velocity2vorticity.py @@ -0,0 +1,104 @@ +"""Convert velocity- to vorticity field assuming 3D CFD exampe was downloaded. +The resulting .hdf5 file does not store pressure and density. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import h5py as h5 +import jax.numpy as jnp +import numpy as np +from tqdm import tqdm + +from .src.vorticity import ( + compute_spectral_vorticity_jnp, +) + + +def convert_velocity() -> None: + parser = argparse.ArgumentParser("Convert velocity field to vorticity!") + parser.add_argument( + "-d", + "--data", + type=str, + required=True, + dest="data", + help=" Specify path to .hdf5 data file", + ) + + args = parser.parse_args() + + if not Path(args.data).exists(): + msg = f"{args.data} does not exist!" + raise FileNotFoundError(msg) + + h5file = h5.File(args.data, "r") + fname = args.data.split("/")[-1] + fname = fname.split(".hdf5")[0] + outpath = str(Path(args.data).parent / fname) + "_vorticity.hdf5" + dx = h5file["x-coordinate"][1] - h5file["x-coordinate"][0] + dy = h5file["y-coordinate"][1] - h5file["y-coordinate"][0] + dz = h5file["z-coordinate"][1] - h5file["z-coordinate"][0] + + if not Path(str(outpath)).exists(): + outfile = h5.File(str(outpath), "a") + outfile.create_dataset( + "omega_x", shape=h5file["Vx"].shape, dtype=h5file["Vx"].dtype + ) + outfile.create_dataset( + "omega_y", shape=h5file["Vy"].shape, dtype=h5file["Vy"].dtype + ) + outfile.create_dataset( + "omega_z", shape=h5file["Vz"].shape, dtype=h5file["Vz"].dtype + ) + outfile.create_dataset( + "t-coordinate", + shape=h5file["t-coordinate"].shape, + dtype=h5file["t-coordinate"].dtype, + ) + outfile.create_dataset( + "x-coordinate", + shape=h5file["x-coordinate"].shape, + dtype=h5file["x-coordinate"].dtype, + ) + outfile.create_dataset( + "y-coordinate", + shape=h5file["y-coordinate"].shape, + dtype=h5file["y-coordinate"].dtype, + ) + outfile.create_dataset( + "z-coordinate", + shape=h5file["z-coordinate"].shape, + dtype=h5file["z-coordinate"].dtype, + ) + + xcoords = h5file["x-coordinate"][:] + ycoords = h5file["y-coordinate"][:] + zcoords = h5file["z-coordinate"][:] + + outfile["t-coordinate"][:] = h5file["t-coordinate"][:] + outfile["x-coordinate"][:] = xcoords + outfile["y-coordinate"][:] = ycoords + outfile["z-coordinate"][:] = zcoords + + trials = h5file["Vx"].shape[0] + + for i in tqdm(range(trials), total=trials): + vx = h5file["Vx"][i][:] + vy = h5file["Vy"][i][:] + vz = h5file["Vz"][i][:] + velocity = np.concatenate( + [vx[..., None], vy[..., None], vz[..., None]], axis=-1 + ) + + vorticity = compute_spectral_vorticity_jnp(jnp.array(velocity), dx, dy, dz) + + outfile["omega_x"][i] = np.array(vorticity[..., 0]) + outfile["omega_y"][i] = np.array(vorticity[..., 1]) + outfile["omega_z"][i] = np.array(vorticity[..., 2]) + + +if __name__ == "__main__": + convert_velocity() diff --git a/pdebench/models/analyse_result_forward.py b/pdebench/models/analyse_result_forward.py index 4d4373b..484635c 100644 --- a/pdebench/models/analyse_result_forward.py +++ b/pdebench/models/analyse_result_forward.py @@ -144,79 +144,109 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import pandas as pd -import numpy as np -import glob + +from __future__ import annotations + import _pickle as cPickle +from pathlib import Path + import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + def main(): # get results - files = glob.glob('./*pickle') + files = list(Path().glob("*.pickle")) files.sort() - + # metric names - var_names = ['MSE', 'normalized MSE', 'Conservation MSE', 'Maximum Error', 'MSE at boundary', - 'MSE FT low', 'MSE FT mid', 'MSE FT high'] - + var_names = [ + "MSE", + "normalized MSE", + "Conservation MSE", + "Maximum Error", + "MSE at boundary", + "MSE FT low", + "MSE FT mid", + "MSE FT high", + ] + # define index index1, index2, index3 = [], [], [] - for j, fl in enumerate(files): - with open(fl, 'rb') as f: - title = fl.split('\\')[-1][:-7].split('_') - print(title) - if title[0] == '1D': - index1.append(title[1]) - index2.append(title[3]) - index3.append(title[4]) - elif title[0] == '2D': - index1.append(title[1]) - index2.append(title[2]) - index3.append(title[4]) + for _j, fl in enumerate(files): + with Path(fl).open("rb") as f: + title = fl.split("\\")[-1][:-7].split("_") + if title[0] == "1D": + if title[1] == "CFD": + index1.append(title[0] + title[1]) + index2.append(title[3] + title[4] + "_" + title[2] + "_" + title[5]) + index3.append(title[7]) + else: + index1.append(title[1]) + index2.append(title[3]) + index3.append(title[4]) + elif title[0] == "2D": + if title[1] == "CFD": + index1.append(title[0] + title[1]) + index2.append( + title[3] + title[3] + title[4] + "_" + title[2] + "_" + title[6] + ) + index3.append(title[9]) + else: + index1.append(title[1]) + index2.append(title[2]) + index3.append(title[4]) + elif title[0] == "3D": + index1.append(title[0] + title[1]) + index2.append( + title[3] + title[4] + title[5] + "_" + title[2] + "_" + title[6] + ) + index3.append(title[8]) else: index1.append(title[0]) index2.append(title[1] + title[2]) index3.append(title[3]) indexes = [index1, index2, index3] - + # create dataframe data = np.zeros([len(files), 8]) for j, fl in enumerate(files): - with open(fl, 'rb') as f: + with Path(fl).open("rb") as f: test = cPickle.load(f) for i, var in enumerate(test): - if i==5: + if i == 5: data[j, i:] = var else: data[j, i] = var - - index = pd.MultiIndex.from_arrays(indexes, names=('PDE', 'param', 'model')) + + index = pd.MultiIndex.from_arrays(indexes, names=("PDE", "param", "model")) data = pd.DataFrame(data, columns=var_names, index=index) - data.to_csv('Results.csv') - + data.to_csv("Results.csv") + pdes = index.get_level_values(0).drop_duplicates() num_pdes = len(pdes) models = index.get_level_values(2).drop_duplicates() num_models = len(models) x = np.arange(num_pdes) - width = 0.5/(num_models-1) - - fig, ax = plt.subplots(figsize=(8,6)) + + width = 0.5 if num_models == 1 else 0.5 / (num_models - 1) + + fig, ax = plt.subplots(figsize=(8, 6)) for i in range(num_models): - pos = x-0.3 + 0.5/(num_models-1)*i - ax.bar(pos, data[data.index.isin([models[i]],level=2)]['MSE'], width) - + pos = x - 0.3 + 0.5 / (num_models - 1) * i + ax.bar(pos, data[data.index.isin([models[i]], level=2)]["MSE"], width) + ax.set_xticks(x) - ax.set_xticklabels(pdes,fontsize=16) - ax.tick_params(axis='y',labelsize=16) - ax.set_yscale('log') - ax.set_xlabel('PDEs',fontsize=16) - ax.set_ylabel('MSE',fontsize=16) - fig.legend(models,loc=8,ncol=num_models,fontsize=16) - plt.tight_layout(rect=[0,0.1,1,1]) - plt.savefig('Results.pdf') - + ax.set_xticklabels(pdes, fontsize=30) + ax.tick_params(axis="y", labelsize=30) + ax.set_yscale("log") + ax.set_xlabel("PDEs", fontsize=30) + ax.set_ylabel("MSE", fontsize=30) + fig.legend(models, loc=8, ncol=num_models, fontsize=20) + plt.tight_layout(rect=[0, 0.1, 1, 1]) + plt.savefig("Results.pdf") + if __name__ == "__main__": main() - print("Done.") \ No newline at end of file diff --git a/pdebench/models/analyse_result_inverse.py b/pdebench/models/analyse_result_inverse.py index 5420232..0a624c2 100644 --- a/pdebench/models/analyse_result_inverse.py +++ b/pdebench/models/analyse_result_inverse.py @@ -144,40 +144,44 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import pandas as pd -import numpy as np -import glob -import _pickle as cPickle + +from __future__ import annotations + import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + def main(): - filename = 'inverse.csv' + filename = "inverse.csv" data = pd.read_csv(filename) - pdes = data['pde'].drop_duplicates() + pdes = data["pde"].drop_duplicates() num_pdes = len(pdes) - models = list(data.columns.values[-2:]) + models = list(data.columns.to_numpy()[-2:]) num_models = len(models) x = np.arange(num_pdes) - width = 0.5/(num_models) - - fig, ax = plt.subplots(figsize=(8,6)) + width = 0.5 / (num_models) + + fig, ax = plt.subplots(figsize=(8, 6)) for i in range(num_models): - pos = x - 0.125 + 0.5/(num_models)*i - ax.bar(pos, data[data.iloc[:,1] == 'mean'][models[i]], - yerr = data[data.iloc[:,1] == 'std'][models[i]], width=width) - print(width, pos) - + pos = x - 0.125 + 0.5 / (num_models) * i + ax.bar( + pos, + data[data.iloc[:, 1] == "mean"][models[i]], + yerr=data[data.iloc[:, 1] == "std"][models[i]], + width=width, + ) + ax.set_xticks(x) - ax.set_xticklabels(pdes,rotation=45,fontsize=16) - ax.tick_params(axis='y',labelsize=16) - ax.set_yscale('log') - ax.set_xlabel('PDEs',fontsize=16) - ax.set_ylabel('MSE',fontsize=16) - fig.legend(models,loc=1,ncol=num_models,fontsize=16) + ax.set_xticklabels(pdes, rotation=45, fontsize=30) + ax.tick_params(axis="y", labelsize=30) + ax.set_yscale("log") + ax.set_xlabel("PDEs", fontsize=30) + ax.set_ylabel("MSE", fontsize=30) + fig.legend(models, loc=1, ncol=num_models, fontsize=20) plt.tight_layout() - plt.savefig('ResultsInverse.pdf') - + plt.savefig("ResultsInverse.pdf") + if __name__ == "__main__": main() - print("Done.") \ No newline at end of file diff --git a/pdebench/models/config/README.md b/pdebench/models/config/README.md new file mode 100644 index 0000000..41a1a85 --- /dev/null +++ b/pdebench/models/config/README.md @@ -0,0 +1,40 @@ +# Config Documentation + +This is the documentation of the config files that were used to generate the +provided +[pre-trained models](https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2987). +Since the default config files for all problems are already provided, this file +only provides the values for the arguments that need to be changed. N/A values +mean that the default values can be used. The complete explanation of the +arguments can be found in the [README file](/README.md) + +| Pre-trained model | Config filename | model_name | filename (data) | ar_mode | pushforward | unroll_step | modes | width | +| :--------------------------------------------------------------- | :-------------------------- | :--------- | :------------------------------------------------------ | :------ | :---------- | ----------: | ----: | ----: | +| 1D_diff-sorp_NA_NA_FNO.pt | config_diff-sorp.yaml | FNO | 1D_diff-sorp_NA_NA | N/A | N/A | N/A | 16 | 64 | +| 1D_diff-sorp_NA_NA_Unet-1-step.pt | config_diff-sorp.yaml | Unet | 1D_diff-sorp_NA_NA | False | False | N/A | N/A | N/A | +| 1D_diff-sorp_NA_NA_Unet-AR.pt | config_diff-sorp.yaml | Unet | 1D_diff-sorp_NA_NA | True | False | N/A | N/A | N/A | +| 1D_diff-sorp_NA_NA_Unet-PF-20.pt | config_diff-sorp.yaml | Unet | 1D_diff-sorp_NA_NA | True | True | 20 | N/A | N/A | +| 1D_diff-sorp_NA_NA_0001.h5_PINN.pt-15000.pt | config_pinn_diff-sorp.yaml | PINN | 1D_diff-sorp_NA_NA.h5 | N/A | N/A | N/A | N/A | N/A | +| 1D_CFD_Shock_trans_Train_FNO.pt | config_1DCFD.yaml | FNO | 1D_CFD_Shock_Eta1.e-8_Zeta1.e-8_trans_Train.hdf5 | N/A | N/A | N/A | 12 | 20 | +| 1D_CFD_Shock_trans_Train_Unet.pt | config_1DCFD.yaml | Unet | 1D_CFD_Shock_Eta1.e-8_Zeta1.e-8_trans_Train.hdf5 | True | True | 20 | N/A | N/A | +| ReacDiff_Nu1.0_Rho2.0_FNO.pt | config_ReacDiff.yaml | FNO | ReacDiff_Nu1.0_Rho2.0.hdf5 | N/A | N/A | N/A | 12 | 20 | +| ReacDiff_Nu1.0_Rho2.0_Unet.pt | config_ReacDiff.yaml | Unet | ReacDiff_Nu1.0_Rho2.0.hdf5 | True | True | 10 | N/A | N/A | +| 1D_Advection_Sols_beta4.0_FNO.pt | config_Adv.yaml | FNO | 1D_Advection_Sols_beta4.0.hdf5 | N/A | N/A | N/A | 12 | 20 | +| 1D_Advection_Sols_beta4.0_Unet.pt | config_Adv.yaml | Unet | 1D_Advection_Sols_beta4.0.hdf5 | True | True | 20 | N/A | N/A | +| 1D_Advection_Sols_beta4.0_PINN.pt-15000.pt | config_pinn_pde1d.yaml | PINN | 1D_Advection_Sols_beta4.0.hdf5 | N/A | N/A | N/A | N/A | N/A | +| 1D_Burgers_Sols_Nu1.0_FNO.pt | config_Bgs.yaml | FNO | 1D_Burgers_Sols_Nu1.0.hdf5 | N/A | N/A | N/A | 12 | 20 | +| 1D_Burgers_Sols_Nu1.0_Unet-PF-20.pt | config_Bgs.yaml | Unet | 1D_Burgers_Sols_Nu1.0.hdf5 | True | True | 20 | N/A | N/A | +| 2D_diff-react_NA_NA_FNO.pt | config_diff-react.yaml | FNO | 2D_diff-react_NA_NA | N/A | N/A | N/A | 12 | 20 | +| 2D_diff-react_NA_NA_Unet-1-step.pt | config_diff-react.yaml | Unet | 2D_diff-react_NA_NA | False | False | N/A | N/A | N/A | +| 2D_diff-react_NA_NA_Unet-AR.pt | config_diff-react.yaml | Unet | 2D_diff-react_NA_NA | True | False | N/A | N/A | N/A | +| 2D_diff-react_NA_NA_Unet-PF-20.pt | config_diff-react.yaml | Unet | 2D_diff-react_NA_NA | True | True | 20 | N/A | N/A | +| 2D_diff-react_NA_NA_0000.h5_PINN.pt-15000.pt | config_pinn_diff-react.yaml | PINN | 2D_diff-react_NA_NA.h5 | N/A | N/A | N/A | N/A | N/A | +| 2D_rdb_NA_NA_FNO.pt | config_rdb.yaml | FNO | 2D_rdb_NA_NA | N/A | N/A | N/A | 12 | 20 | +| 2D_rdb_NA_NA_Unet-1-step.pt | config_rdb.yaml | Unet | 2D_rdb_NA_NA | False | False | N/A | N/A | N/A | +| 2D_rdb_NA_NA_Unet-AR.pt | config_rdb.yaml | Unet | 2D_rdb_NA_NA | True | False | N/A | N/A | N/A | +| 2D_rdb_NA_NA_Unet-PF-20.pt | config_rdb.yaml | Unet | 2D_rdb_NA_NA | True | True | 20 | N/A | N/A | +| 2D_rdb_NA_NA_0000.h5_PINN.pt-15000.pt | config_pinn_swe2d.yaml | PINN | 2D_rdb_NA_NA.h5 | N/A | N/A | N/A | N/A | N/A | +| 2D_DarcyFlow_beta0.01_Train_FNO.pt | config_Darcy.yaml | FNO | 2D_DarcyFlow_beta0.01_Train.hdf5 | N/A | N/A | N/A | 12 | 20 | +| 2D_DarcyFlow_beta0.01_Train_Unet_PF_1.pt | config_Darcy.yaml | Unet | 2D_DarcyFlow_beta0.01_Train.hdf5 | False | False | N/A | N/A | N/A | +| 3D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_Train_FNO.pt | config_3DCFD.yaml | FNO | 3D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_Train.hdf5 | N/A | N/A | N/A | 12 | 20 | +| 3D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_Train_Unet-PF-20.pt | config_3DCFD.yaml | Unet | 3D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_Train.hdf5 | True | True | 20 | N/A | N/A | diff --git a/pdebench/models/config/args/config_1DCFD.yaml b/pdebench/models/config/args/config_1DCFD.yaml index 87863a1..d8d6524 100644 --- a/pdebench/models/config/args/config_1DCFD.yaml +++ b/pdebench/models/config/args/config_1DCFD.yaml @@ -1,4 +1,4 @@ -model_name: 'FNO' +model_name: "FNO" if_training: True continue_training: False num_workers: 2 @@ -6,7 +6,7 @@ batch_size: 100 initial_step: 10 t_train: 100 model_update: 2 -filename: '1D_CFD_periodic_Train.hdf5' +filename: "1D_CFD_periodic_Train.hdf5" single_file: True reduced_resolution: 8 reduced_resolution_t: 5 @@ -34,8 +34,8 @@ mcmc_warmup_steps: 10 mcmc_num_chains: 1 num_samples_max: 1000 in_channels_hid: 64 -inverse_model_type: InitialConditionInterp +inverse_model_type: InitialConditionInterp #Inverse grad inverse_epochs: 100 inverse_learning_rate: 0.2 -inverse_verbose_flag: False \ No newline at end of file +inverse_verbose_flag: False diff --git a/pdebench/models/config/args/config_2DCFD.yaml b/pdebench/models/config/args/config_2DCFD.yaml new file mode 100644 index 0000000..c276c0f --- /dev/null +++ b/pdebench/models/config/args/config_2DCFD.yaml @@ -0,0 +1,25 @@ +model_name: "FNO" +if_training: True +continue_training: False +batch_size: 20 +unroll_step: 20 +num_workers: 2 +t_train: 21 +model_update: 2 +filename: "2D_CFD_M0.1_Eta0.01_Zeta0.01_periodic_128_Train.hdf5" +single_file: True +reduced_resolution: 2 +reduced_resolution_t: 1 +reduced_batch: 1 +epochs: 500 +learning_rate: 1.e-3 +#Unet +in_channels: 4 +out_channels: 4 +#FNO +num_channels: 4 +modes: 12 +width: 20 +scheduler_step: 100 +scheduler_gamma: 0.5 +initial_step: 10 diff --git a/pdebench/models/config/args/config_3DCFD.yaml b/pdebench/models/config/args/config_3DCFD.yaml new file mode 100644 index 0000000..25df7aa --- /dev/null +++ b/pdebench/models/config/args/config_3DCFD.yaml @@ -0,0 +1,24 @@ +model_name: "Unet" +if_training: True +t_train: 21 +continue_training: False +batch_size: 5 +unroll_step: 20 +model_update: 1 +filename: "3D_CFD_Rand_M1.0_Eta1e-08_Zeta1e-08_periodic_Train.hdf5" +single_file: True +reduced_resolution: 2 +reduced_resolution_t: 1 +reduced_batch: 1 +epochs: 500 +learning_rate: 1.e-3 +#Unet +in_channels: 5 +out_channels: 5 +#FNO +num_channels: 5 +modes: 12 +width: 20 +scheduler_step: 100 +scheduler_gamma: 0.5 +initial_step: 10 # should be the same value to unroll_step ?? diff --git a/pdebench/models/config/args/config_Adv.yaml b/pdebench/models/config/args/config_Adv.yaml new file mode 100644 index 0000000..54129d6 --- /dev/null +++ b/pdebench/models/config/args/config_Adv.yaml @@ -0,0 +1,27 @@ +model_name: "Unet" +if_training: True +continue_training: False +batch_size: 50 +unroll_step: 20 +t_train: 200 +model_update: 1 +filename: "1D_Advection_Sols_beta4.0.hdf5" +single_file: True +reduced_resolution: 4 +reduced_resolution_t: 5 +reduced_batch: 1 +epochs: 500 +learning_rate: 1.e-3 +num_workers: 0 +#Unet +in_channels: 1 +out_channels: 1 +ar_mode: True +pushforward: True +#FNO +num_channels: 1 +modes: 12 +width: 20 +scheduler_step: 100 +scheduler_gamma: 0.5 +initial_step: 10 diff --git a/pdebench/models/config/args/config_Bgs.yaml b/pdebench/models/config/args/config_Bgs.yaml new file mode 100644 index 0000000..3702b3b --- /dev/null +++ b/pdebench/models/config/args/config_Bgs.yaml @@ -0,0 +1,27 @@ +model_name: "Unet" +if_training: True +continue_training: False +batch_size: 50 +t_train: 200 +model_update: 1 +filename: "1D_Burgers_Sols_Nu1.0.hdf5" +single_file: True +reduced_resolution: 4 +reduced_resolution_t: 5 +reduced_batch: 1 +epochs: 500 +learning_rate: 1.e-3 +num_workers: 0 +#Unet +in_channels: 1 +out_channels: 1 +ar_mode: True +pushforward: True +unroll_step: 20 +#FNO +num_channels: 1 +modes: 12 +width: 20 +scheduler_step: 100 +scheduler_gamma: 0.5 +initial_step: 10 diff --git a/pdebench/models/config/args/config_Darcy.yaml b/pdebench/models/config/args/config_Darcy.yaml index bca0bfa..bb9b0d8 100644 --- a/pdebench/models/config/args/config_Darcy.yaml +++ b/pdebench/models/config/args/config_Darcy.yaml @@ -1,12 +1,12 @@ -model_name: 'FNO' +model_name: "FNO" if_training: True continue_training: False num_workers: 2 batch_size: 50 initial_step: 10 -t_train: 1 +t_train: 2 model_update: 2 -filename: '2D_DarcyFlow_beta0.01_Train.hdf5' +filename: "2D_DarcyFlow_beta0.01_Train.hdf5" single_file: True reduced_resolution: 2 reduced_resolution_t: 1 @@ -34,8 +34,8 @@ mcmc_warmup_steps: 10 mcmc_num_chains: 1 num_samples_max: 1000 in_channels_hid: 64 -inverse_model_type: InitialConditionInterp +inverse_model_type: InitialConditionInterp #Inverse grad inverse_epochs: 100 inverse_learning_rate: 0.2 -inverse_verbose_flag: False \ No newline at end of file +inverse_verbose_flag: False diff --git a/pdebench/models/config/args/config_ReacDiff.yaml b/pdebench/models/config/args/config_ReacDiff.yaml new file mode 100644 index 0000000..5ebfe91 --- /dev/null +++ b/pdebench/models/config/args/config_ReacDiff.yaml @@ -0,0 +1,27 @@ +model_name: "Unet" +if_training: True +continue_training: False +batch_size: 50 +t_train: 30 +model_update: 1 +filename: "ReacDiff_Nu0.5_Rho1.0.hdf5" +single_file: True +reduced_resolution: 4 +reduced_resolution_t: 1 +reduced_batch: 1 +epochs: 500 +learning_rate: 1.e-3 +num_workers: 0 +#Unet +in_channels: 1 +out_channels: 1 +ar_mode: True +pushforward: True +unroll_step: 10 +#FNO +num_channels: 1 +modes: 12 +width: 20 +scheduler_step: 100 +scheduler_gamma: 0.5 +initial_step: 5 diff --git a/pdebench/models/config/args/config_diff-react.yaml b/pdebench/models/config/args/config_diff-react.yaml index d7a8abe..a650a07 100644 --- a/pdebench/models/config/args/config_diff-react.yaml +++ b/pdebench/models/config/args/config_diff-react.yaml @@ -1,4 +1,4 @@ -model_name: 'FNO' +model_name: "FNO" if_training: False continue_training: False num_workers: 2 @@ -6,7 +6,7 @@ batch_size: 5 initial_step: 10 t_train: 101 model_update: 10 -filename: '2D_diff-react_NA_NA' +filename: "2D_diff-react_NA_NA" single_file: False reduced_resolution: 1 reduced_resolution_t: 1 @@ -34,7 +34,7 @@ mcmc_warmup_steps: 10 mcmc_num_chains: 1 num_samples_max: 1000 in_channels_hid: 64 -inverse_model_type: InitialConditionInterp +inverse_model_type: InitialConditionInterp #Inverse grad inverse_epochs: 100 inverse_learning_rate: 0.2 @@ -47,4 +47,4 @@ x_max: 1 y_min: -1 y_max: 1 t_min: 0 -t_max: 5 \ No newline at end of file +t_max: 5 diff --git a/pdebench/models/config/args/config_diff-sorp.yaml b/pdebench/models/config/args/config_diff-sorp.yaml index c352fc5..7ed4371 100644 --- a/pdebench/models/config/args/config_diff-sorp.yaml +++ b/pdebench/models/config/args/config_diff-sorp.yaml @@ -1,4 +1,4 @@ -model_name: 'FNO' +model_name: "FNO" if_training: False continue_training: False num_workers: 2 @@ -6,7 +6,7 @@ batch_size: 50 initial_step: 10 t_train: 101 model_update: 10 -filename: '1D_diff-sorp_NA_NA' +filename: "1D_diff-sorp_NA_NA" single_file: False reduced_resolution: 1 reduced_resolution_t: 1 @@ -34,7 +34,7 @@ mcmc_warmup_steps: 10 mcmc_num_chains: 1 num_samples_max: 1000 in_channels_hid: 64 -inverse_model_type: InitialConditionInterp +inverse_model_type: InitialConditionInterp #Inverse grad inverse_epochs: 100 inverse_learning_rate: 0.2 @@ -47,4 +47,4 @@ x_max: 1 y_min: 0 y_max: 1 t_min: 0 -t_max: 500 \ No newline at end of file +t_max: 500 diff --git a/pdebench/models/config/args/config_pinn_CFD1d.yaml b/pdebench/models/config/args/config_pinn_CFD1d.yaml new file mode 100644 index 0000000..3b688a5 --- /dev/null +++ b/pdebench/models/config/args/config_pinn_CFD1d.yaml @@ -0,0 +1,15 @@ +model_name: "PINN" +scenario: "pde1D" +model_update: 500 +filename: "1D_CFD_Rand_Eta1.e-8_Zeta1.e-8_periodic_Train.hdf5" +epochs: 15000 +input_ch: 2 +output_ch: 3 +learning_rate: 1.e-3 +root_path: "../data" +val_num: 10 +if_periodic_bc: True +period: 5000 +val_time: 1.0 +val_batch_idx: 10 +aux_params: [1.6666666667] diff --git a/pdebench/models/config/args/config_pinn_diff-react.yaml b/pdebench/models/config/args/config_pinn_diff-react.yaml index 28ffa96..f251aac 100644 --- a/pdebench/models/config/args/config_pinn_diff-react.yaml +++ b/pdebench/models/config/args/config_pinn_diff-react.yaml @@ -1,8 +1,15 @@ -model_name: 'PINN' +model_name: "PINN" -scenario: 'diff-react' +scenario: "diff-react" model_update: 500 -filename: '2D_diff-react_NA_NA_0000.h5' +filename: "2D_diff-react_NA_NA/2D_diff-react_NA_NA.h5" epochs: 100 learning_rate: 1.e-3 -seed: '0000' \ No newline at end of file +seed: "0000" +# unused arguments +input_ch: 0 +output_ch: 1 +root_path: "." +val_num: 1 +if_periodic_bc: False +aux_params: 0 diff --git a/pdebench/models/config/args/config_pinn_diff-sorp.yaml b/pdebench/models/config/args/config_pinn_diff-sorp.yaml index e1dab70..d3e6512 100644 --- a/pdebench/models/config/args/config_pinn_diff-sorp.yaml +++ b/pdebench/models/config/args/config_pinn_diff-sorp.yaml @@ -1,8 +1,15 @@ -model_name: 'PINN' +model_name: "PINN" -scenario: 'diff-sorp' +scenario: "diff-sorp" model_update: 500 -filename: '2D_diff-react_NA_NA_0000.h5' +filename: "1D_diff-sorp_NA_NA/1D_diff-sorp_NA_NA.h5" epochs: 15000 learning_rate: 1.e-3 -seed: '0000' \ No newline at end of file +seed: "0000" +# unused arguments +input_ch: 0 +output_ch: 1 +root_path: "." +val_num: 1 +if_periodic_bc: False +aux_params: 0 diff --git a/pdebench/models/config/args/config_pinn_pde1d.yaml b/pdebench/models/config/args/config_pinn_pde1d.yaml new file mode 100644 index 0000000..b76d3cc --- /dev/null +++ b/pdebench/models/config/args/config_pinn_pde1d.yaml @@ -0,0 +1,15 @@ +model_name: "PINN" +scenario: "pde1D" +model_update: 500 +filename: "1D_Advection_Sols_beta0.1.hdf5" +epochs: 15000 +input_ch: 2 +output_ch: 1 +learning_rate: 1.e-3 +root_path: "../data" +val_num: 10 +if_periodic_bc: True +period: 5000 +val_time: 2.0 +val_batch_idx: 10 +aux_params: [0.1] diff --git a/pdebench/models/config/args/config_pinn_swe2d.yaml b/pdebench/models/config/args/config_pinn_swe2d.yaml index d739f9a..bff678a 100644 --- a/pdebench/models/config/args/config_pinn_swe2d.yaml +++ b/pdebench/models/config/args/config_pinn_swe2d.yaml @@ -1,8 +1,15 @@ -model_name: 'PINN' +model_name: "PINN" -scenario: 'swe2d' +scenario: "swe2d" model_update: 500 -filename: '2D_rdb_NA_NA_0000.h5' +filename: "2D_rdb_NA_NA/2D_rdb_NA_NA.h5" epochs: 15000 learning_rate: 1.e-3 -seed: '0000' \ No newline at end of file +seed: "0000" +# unused arguments +input_ch: 0 +output_ch: 1 +root_path: "." +val_num: 1 +if_periodic_bc: False +aux_params: 0 diff --git a/pdebench/models/config/args/config_rdb.yaml b/pdebench/models/config/args/config_rdb.yaml index 1048098..1690d9d 100644 --- a/pdebench/models/config/args/config_rdb.yaml +++ b/pdebench/models/config/args/config_rdb.yaml @@ -1,12 +1,12 @@ -model_name: 'FNO' +model_name: "FNO" if_training: False continue_training: False num_workers: 2 batch_size: 5 initial_step: 10 -t_train: 100 +t_train: 101 model_update: 10 -filename: '2D_rdb_NA_NA' +filename: "2D_rdb_NA_NA" single_file: False reduced_resolution: 1 reduced_resolution_t: 1 @@ -34,7 +34,7 @@ mcmc_warmup_steps: 10 mcmc_num_chains: 1 num_samples_max: 1000 in_channels_hid: 64 -inverse_model_type: InitialConditionInterp +inverse_model_type: InitialConditionInterp #Inverse grad inverse_epochs: 100 inverse_learning_rate: 0.2 @@ -47,4 +47,4 @@ x_max: 2.5 y_min: -2.5 y_max: 2.5 t_min: 0 -t_max: 1 \ No newline at end of file +t_max: 1 diff --git a/pdebench/models/config/config.yaml b/pdebench/models/config/config.yaml index 198ad5c..63ca75a 100644 --- a/pdebench/models/config/config.yaml +++ b/pdebench/models/config/config.yaml @@ -7,55 +7,55 @@ hydra: output_subdir: null run: dir: . - + args: - model_name: 'FNO' - if_training: False - continue_training: False - num_workers: 2 - batch_size: 5 - initial_step: 10 - t_train: 101 - model_update: 10 - filename: '2D_diff-react_NA_NA' - single_file: False - reduced_resolution: 1 - reduced_resolution_t: 1 - reduced_batch: 1 - epochs: 500 - learning_rate: 1.e-3 - scheduler_step: 100 - scheduler_gamma: 0.5 - #Unet - in_channels: 2 - out_channels: 2 - ar_mode: True - pushforward: True - unroll_step: 20 - #FNO - num_channels: 2 - modes: 12 - width: 20 - #Inverse - base_path: ../data/ - training_type: autoregressive - #Inverse MCMC - mcmc_num_samples: 20 - mcmc_warmup_steps: 10 - mcmc_num_chains: 1 - num_samples_max: 1000 - in_channels_hid: 64 - inverse_model_type: InitialConditionInterp - #Inverse grad - inverse_epochs: 100 - inverse_learning_rate: 0.2 - inverse_verbose_flag: False - #Plotting - plot: False - channel_plot: 0 # Which channel/variable to be plotted - x_min: -1 - x_max: 1 - y_min: -1 - y_max: 1 - t_min: 0 - t_max: 5 \ No newline at end of file + model_name: "FNO" + if_training: False + continue_training: False + num_workers: 2 + batch_size: 5 + initial_step: 10 + t_train: 101 + model_update: 10 + filename: "2D_diff-react_NA_NA" + single_file: False + reduced_resolution: 1 + reduced_resolution_t: 1 + reduced_batch: 1 + epochs: 500 + learning_rate: 1.e-3 + scheduler_step: 100 + scheduler_gamma: 0.5 + #Unet + in_channels: 2 + out_channels: 2 + ar_mode: True + pushforward: True + unroll_step: 20 + #FNO + num_channels: 2 + modes: 12 + width: 20 + #Inverse + base_path: ../data/ + training_type: autoregressive + #Inverse MCMC + mcmc_num_samples: 20 + mcmc_warmup_steps: 10 + mcmc_num_chains: 1 + num_samples_max: 1000 + in_channels_hid: 64 + inverse_model_type: InitialConditionInterp + #Inverse grad + inverse_epochs: 100 + inverse_learning_rate: 0.2 + inverse_verbose_flag: False + #Plotting + plot: False + channel_plot: 0 # Which channel/variable to be plotted + x_min: -1 + x_max: 1 + y_min: -1 + y_max: 1 + t_min: 0 + t_max: 5 diff --git a/pdebench/models/config/config_darcy.yaml b/pdebench/models/config/config_darcy.yaml new file mode 100644 index 0000000..2949cb0 --- /dev/null +++ b/pdebench/models/config/config_darcy.yaml @@ -0,0 +1,61 @@ +defaults: + - _self_ + - override hydra/hydra_logging: disabled + - override hydra/job_logging: disabled + +hydra: + output_subdir: null + run: + dir: . + +args: + model_name: "FNO" + if_training: True + continue_training: False + num_workers: 2 + batch_size: 50 + initial_step: 1 + t_train: 2 # steady-state + model_update: 2 + filename: "2D_DarcyFlow_beta0.01_Train.hdf5" + single_file: True + reduced_resolution: 1 + reduced_resolution_t: 1 + reduced_batch: 1 + epochs: 500 + learning_rate: 1.e-3 + scheduler_step: 100 + scheduler_gamma: 0.5 + #Unet + in_channels: 1 + out_channels: 1 + ar_mode: False + pushforward: False + unroll_step: 1 + #FNO + num_channels: 1 + modes: 12 + width: 20 + #Inverse + data_path: "../data/2D/DarcyFlow/Train/" + training_type: "single" #autoregressive + #Inverse MCMC + mcmc_num_samples: 20 + mcmc_warmup_steps: 10 + mcmc_num_chains: 1 + num_samples_max: 1000 + in_channels_hid: 64 + inverse_model_type: InitialConditionInterp + #Inverse grad + inverse_epochs: 100 + inverse_learning_rate: 0.2 + inverse_verbose_flag: False + #plotting + plot: False + channel_plot: 0 # Which channel/variable to be plotted + x_min: -1 + x_max: 1 # spatial dimension x: [-1, 1] + y_min: -1 + y_max: 1 # spatial dimension y: [-1, 1] + t_min: 0 + t_max: 5 # time dimension t: [0, 5] diff --git a/pdebench/models/config/config_rdb.yaml b/pdebench/models/config/config_rdb.yaml new file mode 100644 index 0000000..e28a163 --- /dev/null +++ b/pdebench/models/config/config_rdb.yaml @@ -0,0 +1,61 @@ +defaults: + - _self_ + - override hydra/hydra_logging: disabled + - override hydra/job_logging: disabled + +hydra: + output_subdir: null + run: + dir: . + +args: + model_name: "FNO" + if_training: False + continue_training: False + num_workers: 2 + batch_size: 5 + initial_step: 10 + t_train: 101 + model_update: 2 + filename: "2D_rdb_NA_NA.h5" + single_file: True + data_path: "/path/to/swe2d/h5" + reduced_resolution: 1 + reduced_resolution_t: 1 + reduced_batch: 1 + epochs: 500 + learning_rate: 1.e-3 + scheduler_step: 100 + scheduler_gamma: 0.5 + #Unet + in_channels: 1 + out_channels: 1 + ar_mode: True + pushforward: True + unroll_step: 20 + #FNO + num_channels: 1 + modes: 12 + width: 20 + #Inverse + training_type: "autoregressive" + #Inverse MCMC + mcmc_num_samples: 20 + mcmc_warmup_steps: 10 + mcmc_num_chains: 1 + num_samples_max: 1000 + in_channels_hid: 64 + inverse_model_type: InitialConditionInterp + #Inverse grad + inverse_epochs: 100 + inverse_learning_rate: 0.2 + inverse_verbose_flag: False + #Plotting + plot: False + channel_plot: 0 # Which channel/variable to be plotted + x_min: -2.5 + x_max: 2.5 + y_min: -2.5 + y_max: 2.5 + t_min: 0 + t_max: 1 diff --git a/pdebench/models/config/results.yaml b/pdebench/models/config/results.yaml index 12e1728..67195b9 100644 --- a/pdebench/models/config/results.yaml +++ b/pdebench/models/config/results.yaml @@ -7,13 +7,18 @@ hydra: output_subdir: null run: dir: . - -args: - model_names: [FNO, Unet] - base_path : /home/alesiani/python/pde_benchmark/pdebench/data/ - inverse_model_type : InitialConditionInterp - filenames : [/1D/Advection/Train/1D_Advection_Sols_beta4.0.hdf5,/1D/Burgers/Train/1D_Burgers_Sols_Nu1.0.hdf5, /1D/ReactionDiffusion/Train/ReacDiff_Nu1.0_Rho2.0.hdf5, /1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5] - shortfilenames : [Advection, Burgers, ReacDiff, CFD] - results_values : [ mseloss_pred_u0 ] - result_filename : csv/results_inverse.csv +args: + model_names: [FNO, Unet] + base_path: /home/alesiani/python/pde_benchmark/pdebench/data/ + inverse_model_type: InitialConditionInterp + filenames: + [ + /1D/Advection/Train/1D_Advection_Sols_beta4.0.hdf5, + /1D/Burgers/Train/1D_Burgers_Sols_Nu1.0.hdf5, + /1D/ReactionDiffusion/Train/ReacDiff_Nu1.0_Rho2.0.hdf5, + /1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5, + ] + shortfilenames: [Advection, Burgers, ReacDiff, CFD] + results_values: [mseloss_pred_u0] + result_filename: csv/results_inverse.csv diff --git a/pdebench/models/fno/fno.py b/pdebench/models/fno/fno.py index 6371228..19e21bd 100644 --- a/pdebench/models/fno/fno.py +++ b/pdebench/models/fno/fno.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python3 - """ FNO. Implementation taken and modified from https://github.com/zongyi-li/fourier_neural_operator @@ -27,26 +25,33 @@ SOFTWARE. """ +from __future__ import annotations + import torch -import torch.nn as nn -import numpy as np import torch.nn.functional as F +from torch import nn class SpectralConv1d(nn.Module): def __init__(self, in_channels, out_channels, modes1): - super(SpectralConv1d, self).__init__() + super().__init__() """ - 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. + 1D Fourier layer. It does FFT, linear transform, and Inverse FFT. """ self.in_channels = in_channels self.out_channels = out_channels - self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 + self.modes1 = ( + # Number of Fourier modes to multiply, at most floor(N/2) + 1 + modes1 + ) - self.scale = (1 / (in_channels*out_channels)) - self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat)) + self.scale = 1 / (in_channels * out_channels) + self.weights1 = nn.Parameter( + self.scale + * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat) + ) # Complex multiplication def compl_mul1d(self, input, weights): @@ -55,20 +60,28 @@ def compl_mul1d(self, input, weights): def forward(self, x): batchsize = x.shape[0] - #Compute Fourier coeffcients up to factor of e^(- something constant) + # Compute Fourier coefficients up to factor of e^(- something constant) x_ft = torch.fft.rfft(x) - # Multiply relevant Fourier modes - out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat) - out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1) + # Multiply relevant Fourier modes + out_ft = torch.zeros( + batchsize, + self.out_channels, + x.size(-1) // 2 + 1, + device=x.device, + dtype=torch.cfloat, + ) + out_ft[:, :, : self.modes1] = self.compl_mul1d( + x_ft[:, :, : self.modes1], self.weights1 + ) + + # Return to physical space + return torch.fft.irfft(out_ft, n=x.size(-1)) - #Return to physical space - x = torch.fft.irfft(out_ft, n=x.size(-1)) - return x class FNO1d(nn.Module): def __init__(self, num_channels, modes=16, width=64, initial_step=10): - super(FNO1d, self).__init__() + super().__init__() """ The overall network. It contains 4 layers of the Fourier layer. @@ -76,7 +89,7 @@ def __init__(self, num_channels, modes=16, width=64, initial_step=10): 2. 4 layers of the integral operators u' = (W + K)(u). W defined by self.w; K defined by self.conv . 3. Project from the channel space to the output space by self.fc1 and self.fc2 . - + input: the solution of the initial condition and location (a(x), x) input shape: (batchsize, x=s, c=2) output: the solution of a later timestep @@ -85,8 +98,10 @@ def __init__(self, num_channels, modes=16, width=64, initial_step=10): self.modes1 = modes self.width = width - self.padding = 2 # pad the domain if input is non-periodic - self.fc0 = nn.Linear(initial_step*num_channels+1, self.width) # input channel is 2: (a(x), x) + self.padding = 2 # pad the domain if input is non-periodic + self.fc0 = nn.Linear( + initial_step * num_channels + 1, self.width + ) # input channel is 2: (a(x), x) self.conv0 = SpectralConv1d(self.width, self.width, self.modes1) self.conv1 = SpectralConv1d(self.width, self.width, self.modes1) @@ -105,8 +120,9 @@ def forward(self, x, grid): x = torch.cat((x, grid), dim=-1) x = self.fc0(x) x = x.permute(0, 2, 1) - - x = F.pad(x, [0, self.padding]) # pad the domain if input is non-periodic + + # pad the domain if input is non-periodic + x = F.pad(x, [0, self.padding]) x1 = self.conv0(x) x2 = self.w0(x) @@ -127,7 +143,7 @@ def forward(self, x, grid): x2 = self.w3(x) x = x1 + x2 - x = x[..., :-self.padding] + x = x[..., : -self.padding] x = x.permute(0, 2, 1) x = self.fc1(x) x = F.gelu(x) @@ -137,20 +153,33 @@ def forward(self, x, grid): class SpectralConv2d_fast(nn.Module): def __init__(self, in_channels, out_channels, modes1, modes2): - super(SpectralConv2d_fast, self).__init__() + super().__init__() """ - 2D Fourier layer. It does FFT, linear transform, and Inverse FFT. + 2D Fourier layer. It does FFT, linear transform, and Inverse FFT. """ self.in_channels = in_channels self.out_channels = out_channels - self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 + self.modes1 = ( + # Number of Fourier modes to multiply, at most floor(N/2) + 1 + modes1 + ) self.modes2 = modes2 - self.scale = (1 / (in_channels * out_channels)) - self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)) - self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)) + self.scale = 1 / (in_channels * out_channels) + self.weights1 = nn.Parameter( + self.scale + * torch.rand( + in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat + ) + ) + self.weights2 = nn.Parameter( + self.scale + * torch.rand( + in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat + ) + ) # Complex multiplication def compl_mul2d(self, input, weights): @@ -159,23 +188,32 @@ def compl_mul2d(self, input, weights): def forward(self, x): batchsize = x.shape[0] - #Compute Fourier coeffcients up to factor of e^(- something constant) + # Compute Fourier coefficients up to factor of e^(- something constant) x_ft = torch.fft.rfft2(x) # Multiply relevant Fourier modes - out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1)//2 + 1, dtype=torch.cfloat, device=x.device) - out_ft[:, :, :self.modes1, :self.modes2] = \ - self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1) - out_ft[:, :, -self.modes1:, :self.modes2] = \ - self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2) + out_ft = torch.zeros( + batchsize, + self.out_channels, + x.size(-2), + x.size(-1) // 2 + 1, + dtype=torch.cfloat, + device=x.device, + ) + out_ft[:, :, : self.modes1, : self.modes2] = self.compl_mul2d( + x_ft[:, :, : self.modes1, : self.modes2], self.weights1 + ) + out_ft[:, :, -self.modes1 :, : self.modes2] = self.compl_mul2d( + x_ft[:, :, -self.modes1 :, : self.modes2], self.weights2 + ) + + # Return to physical space + return torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1))) - #Return to physical space - x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1))) - return x class FNO2d(nn.Module): def __init__(self, num_channels, modes1=12, modes2=12, width=20, initial_step=10): - super(FNO2d, self).__init__() + super().__init__() """ The overall network. It contains 4 layers of the Fourier layer. @@ -183,7 +221,7 @@ def __init__(self, num_channels, modes1=12, modes2=12, width=20, initial_step=10 2. 4 layers of the integral operators u' = (W + K)(u). W defined by self.w; K defined by self.conv . 3. Project from the channel space to the output space by self.fc1 and self.fc2 . - + input: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y), ..., u(t-1, x, y), x, y) input shape: (batchsize, x, y, c) output: the solution of the next timestep @@ -193,14 +231,22 @@ def __init__(self, num_channels, modes1=12, modes2=12, width=20, initial_step=10 self.modes1 = modes1 self.modes2 = modes2 self.width = width - self.padding = 2 # pad the domain if input is non-periodic - self.fc0 = nn.Linear(initial_step*num_channels+2, self.width) + self.padding = 2 # pad the domain if input is non-periodic + self.fc0 = nn.Linear(initial_step * num_channels + 2, self.width) # input channel is 12: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y), ..., u(t-1, x, y), x, y) - self.conv0 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) - self.conv1 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) - self.conv2 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) - self.conv3 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2) + self.conv0 = SpectralConv2d_fast( + self.width, self.width, self.modes1, self.modes2 + ) + self.conv1 = SpectralConv2d_fast( + self.width, self.width, self.modes1, self.modes2 + ) + self.conv2 = SpectralConv2d_fast( + self.width, self.width, self.modes1, self.modes2 + ) + self.conv3 = SpectralConv2d_fast( + self.width, self.width, self.modes1, self.modes2 + ) self.w0 = nn.Conv2d(self.width, self.width, 1) self.w1 = nn.Conv2d(self.width, self.width, 1) self.w2 = nn.Conv2d(self.width, self.width, 1) @@ -214,7 +260,7 @@ def forward(self, x, grid): x = torch.cat((x, grid), dim=-1) x = self.fc0(x) x = x.permute(0, 3, 1, 2) - + # Pad tensor with boundary condition x = F.pad(x, [0, self.padding, 0, self.padding]) @@ -237,34 +283,77 @@ def forward(self, x, grid): x2 = self.w3(x) x = x1 + x2 - x = x[..., :-self.padding, :-self.padding] # Unpad the tensor + x = x[..., : -self.padding, : -self.padding] # Unpad the tensor x = x.permute(0, 2, 3, 1) x = self.fc1(x) x = F.gelu(x) x = self.fc2(x) - + return x.unsqueeze(-2) - + class SpectralConv3d(nn.Module): def __init__(self, in_channels, out_channels, modes1, modes2, modes3): - super(SpectralConv3d, self).__init__() + super().__init__() """ - 3D Fourier layer. It does FFT, linear transform, and Inverse FFT. + 3D Fourier layer. It does FFT, linear transform, and Inverse FFT. """ self.in_channels = in_channels self.out_channels = out_channels - self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1 + self.modes1 = ( + # Number of Fourier modes to multiply, at most floor(N/2) + 1 + modes1 + ) self.modes2 = modes2 self.modes3 = modes3 - self.scale = (1 / (in_channels * out_channels)) - self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat)) - self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat)) - self.weights3 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat)) - self.weights4 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat)) + self.scale = 1 / (in_channels * out_channels) + self.weights1 = nn.Parameter( + self.scale + * torch.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=torch.cfloat, + ) + ) + self.weights2 = nn.Parameter( + self.scale + * torch.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=torch.cfloat, + ) + ) + self.weights3 = nn.Parameter( + self.scale + * torch.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=torch.cfloat, + ) + ) + self.weights4 = nn.Parameter( + self.scale + * torch.rand( + in_channels, + out_channels, + self.modes1, + self.modes2, + self.modes3, + dtype=torch.cfloat, + ) + ) # Complex multiplication def compl_mul3d(self, input, weights): @@ -273,27 +362,41 @@ def compl_mul3d(self, input, weights): def forward(self, x): batchsize = x.shape[0] - #Compute Fourier coeffcients up to factor of e^(- something constant) - x_ft = torch.fft.rfftn(x, dim=[-3,-2,-1]) + # Compute Fourier coefficients up to factor of e^(- something constant) + x_ft = torch.fft.rfftn(x, dim=[-3, -2, -1]) # Multiply relevant Fourier modes - out_ft = torch.zeros(batchsize, self.out_channels, x.size(-3), x.size(-2), x.size(-1)//2 + 1, dtype=torch.cfloat, device=x.device) - out_ft[:, :, :self.modes1, :self.modes2, :self.modes3] = \ - self.compl_mul3d(x_ft[:, :, :self.modes1, :self.modes2, :self.modes3], self.weights1) - out_ft[:, :, -self.modes1:, :self.modes2, :self.modes3] = \ - self.compl_mul3d(x_ft[:, :, -self.modes1:, :self.modes2, :self.modes3], self.weights2) - out_ft[:, :, :self.modes1, -self.modes2:, :self.modes3] = \ - self.compl_mul3d(x_ft[:, :, :self.modes1, -self.modes2:, :self.modes3], self.weights3) - out_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3] = \ - self.compl_mul3d(x_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.weights4) - - #Return to physical space - x = torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1))) - return x + out_ft = torch.zeros( + batchsize, + self.out_channels, + x.size(-3), + x.size(-2), + x.size(-1) // 2 + 1, + dtype=torch.cfloat, + device=x.device, + ) + out_ft[:, :, : self.modes1, : self.modes2, : self.modes3] = self.compl_mul3d( + x_ft[:, :, : self.modes1, : self.modes2, : self.modes3], self.weights1 + ) + out_ft[:, :, -self.modes1 :, : self.modes2, : self.modes3] = self.compl_mul3d( + x_ft[:, :, -self.modes1 :, : self.modes2, : self.modes3], self.weights2 + ) + out_ft[:, :, : self.modes1, -self.modes2 :, : self.modes3] = self.compl_mul3d( + x_ft[:, :, : self.modes1, -self.modes2 :, : self.modes3], self.weights3 + ) + out_ft[:, :, -self.modes1 :, -self.modes2 :, : self.modes3] = self.compl_mul3d( + x_ft[:, :, -self.modes1 :, -self.modes2 :, : self.modes3], self.weights4 + ) + + # Return to physical space + return torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1))) + class FNO3d(nn.Module): - def __init__(self, num_channels, modes1=8, modes2=8, modes3=8, width=20, initial_step=10): - super(FNO3d, self).__init__() + def __init__( + self, num_channels, modes1=8, modes2=8, modes3=8, width=20, initial_step=10 + ): + super().__init__() """ The overall network. It contains 4 layers of the Fourier layer. @@ -301,7 +404,7 @@ def __init__(self, num_channels, modes1=8, modes2=8, modes3=8, width=20, initial 2. 4 layers of the integral operators u' = (W + K)(u). W defined by self.w; K defined by self.conv . 3. Project from the channel space to the output space by self.fc1 and self.fc2 . - + input: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y), x, y, t). It's a constant function in time, except for the last index. input shape: (batchsize, x=64, y=64, t=40, c=13) output: the solution of the next 40 timesteps @@ -312,14 +415,22 @@ def __init__(self, num_channels, modes1=8, modes2=8, modes3=8, width=20, initial self.modes2 = modes2 self.modes3 = modes3 self.width = width - self.padding = 6 # pad the domain if input is non-periodic - self.fc0 = nn.Linear(initial_step*num_channels+3, self.width) + self.padding = 6 # pad the domain if input is non-periodic + self.fc0 = nn.Linear(initial_step * num_channels + 3, self.width) # input channel is 12: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y), x, y, t) - self.conv0 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3) - self.conv1 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3) - self.conv2 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3) - self.conv3 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3) + self.conv0 = SpectralConv3d( + self.width, self.width, self.modes1, self.modes2, self.modes3 + ) + self.conv1 = SpectralConv3d( + self.width, self.width, self.modes1, self.modes2, self.modes3 + ) + self.conv2 = SpectralConv3d( + self.width, self.width, self.modes1, self.modes2, self.modes3 + ) + self.conv3 = SpectralConv3d( + self.width, self.width, self.modes1, self.modes2, self.modes3 + ) self.w0 = nn.Conv3d(self.width, self.width, 1) self.w1 = nn.Conv3d(self.width, self.width, 1) self.w2 = nn.Conv3d(self.width, self.width, 1) @@ -337,8 +448,9 @@ def forward(self, x, grid): x = torch.cat((x, grid), dim=-1) x = self.fc0(x) x = x.permute(0, 4, 1, 2, 3) - - x = F.pad(x, [0, self.padding]) # pad the domain if input is non-periodic + + # pad the domain if input is non-periodic + x = F.pad(x, [0, self.padding]) x1 = self.conv0(x) x2 = self.w0(x) @@ -359,9 +471,9 @@ def forward(self, x, grid): x2 = self.w3(x) x = x1 + x2 - x = x[..., :-self.padding] - x = x.permute(0, 2, 3, 4, 1) # pad the domain if input is non-periodic + x = x[..., : -self.padding] + x = x.permute(0, 2, 3, 4, 1) # pad the domain if input is non-periodic x = self.fc1(x) x = F.gelu(x) x = self.fc2(x) - return x.unsqueeze(-2) \ No newline at end of file + return x.unsqueeze(-2) diff --git a/pdebench/models/fno/train.py b/pdebench/models/fno/train.py index 65a2846..0a51062 100644 --- a/pdebench/models/fno/train.py +++ b/pdebench/models/fno/train.py @@ -1,198 +1,219 @@ -import sys -import torch -import numpy as np -import pickle -import torch.nn as nn -import torch.nn.functional as F +from __future__ import annotations -import operator -from functools import reduce -from functools import partial +import pickle +from pathlib import Path -from timeit import default_timer +import numpy as np +import torch +from pdebench.models.fno.fno import FNO1d, FNO2d, FNO3d +from pdebench.models.fno.utils import FNODatasetMult, FNODatasetSingle +from pdebench.models.metrics import metrics +from torch import nn # torch.manual_seed(0) # np.random.seed(0) +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def run_training( + if_training, + continue_training, + num_workers, + modes, + width, + initial_step, + t_train, + num_channels, + batch_size, + epochs, + learning_rate, + scheduler_step, + scheduler_gamma, + model_update, + flnm, + single_file, + reduced_resolution, + reduced_resolution_t, + reduced_batch, + plot, + channel_plot, + x_min, + x_max, + y_min, + y_max, + t_min, + t_max, + base_path="../data/", + training_type="autoregressive", +): + # print( + # f"Epochs = {epochs}, learning rate = {learning_rate}, scheduler step = {scheduler_step}, scheduler gamma = {scheduler_gamma}" + # ) -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - -sys.path.append('.') -from .fno import FNO1d, FNO2d, FNO3d -from .utils import FNODatasetSingle, FNODatasetMult -from metrics import metrics - -def run_training(if_training, - continue_training, - num_workers, - modes, - width, - initial_step, - t_train, - num_channels, - batch_size, - epochs, - learning_rate, - scheduler_step, - scheduler_gamma, - model_update, - flnm, - single_file, - reduced_resolution, - reduced_resolution_t, - reduced_batch, - plot, - channel_plot, - x_min, - x_max, - y_min, - y_max, - t_min, - t_max, - base_path='../data/', - training_type='autoregressive' - ): - - print(f'Epochs = {epochs}, learning rate = {learning_rate}, scheduler step = {scheduler_step}, scheduler gamma = {scheduler_gamma}') - ################################################################ # load data ################################################################ - + if single_file: # filename - model_name = flnm[:-5] + '_FNO' - print("FNODatasetSingle") + model_name = flnm[:-5] + "_FNO" + # print("FNODatasetSingle") # Initialize the dataset and dataloader - train_data = FNODatasetSingle(flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - initial_step=initial_step, - saved_folder = base_path - ) - val_data = FNODatasetSingle(flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - initial_step=initial_step, - if_test=True, - saved_folder = base_path - ) - + train_data = FNODatasetSingle( + flnm, + reduced_resolution=reduced_resolution, + reduced_resolution_t=reduced_resolution_t, + reduced_batch=reduced_batch, + initial_step=initial_step, + saved_folder=base_path, + ) + val_data = FNODatasetSingle( + flnm, + reduced_resolution=reduced_resolution, + reduced_resolution_t=reduced_resolution_t, + reduced_batch=reduced_batch, + initial_step=initial_step, + if_test=True, + saved_folder=base_path, + ) + else: # filename - model_name = flnm + '_FNO' - - print("FNODatasetMult") - train_data = FNODatasetMult(flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - saved_folder = base_path - ) - val_data = FNODatasetMult(flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - if_test=True, - saved_folder = base_path) - - train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, - num_workers=num_workers, shuffle=True) - val_loader = torch.utils.data.DataLoader(val_data, batch_size=batch_size, - num_workers=num_workers, shuffle=False) - + model_name = flnm + "_FNO" + + # print("FNODatasetMult") + train_data = FNODatasetMult( + flnm, + saved_folder=base_path, + ) + val_data = FNODatasetMult( + flnm, + if_test=True, + saved_folder=base_path, + ) + + train_loader = torch.utils.data.DataLoader( + train_data, batch_size=batch_size, num_workers=num_workers, shuffle=True + ) + val_loader = torch.utils.data.DataLoader( + val_data, batch_size=batch_size, num_workers=num_workers, shuffle=False + ) + ################################################################ # training and evaluation ################################################################ - + _, _data, _ = next(iter(val_loader)) dimensions = len(_data.shape) - print('Spatial Dimension', dimensions - 3) + # print("Spatial Dimension", dimensions - 3) if dimensions == 4: - model = FNO1d(num_channels=num_channels, - width=width, - modes=modes, - initial_step=initial_step).to(device) + model = FNO1d( + num_channels=num_channels, + width=width, + modes=modes, + initial_step=initial_step, + ).to(device) elif dimensions == 5: - model = FNO2d(num_channels=num_channels, - width=width, - modes1=modes, - modes2=modes, - initial_step=initial_step).to(device) + model = FNO2d( + num_channels=num_channels, + width=width, + modes1=modes, + modes2=modes, + initial_step=initial_step, + ).to(device) elif dimensions == 6: - model = FNO3d(num_channels=num_channels, - width=width, - modes1=modes, - modes2=modes, - modes3=modes, - initial_step=initial_step).to(device) - + model = FNO3d( + num_channels=num_channels, + width=width, + modes1=modes, + modes2=modes, + modes3=modes, + initial_step=initial_step, + ).to(device) + # Set maximum time step of the data to train - if t_train > _data.shape[-2]: - t_train = _data.shape[-2] + t_train = min(t_train, _data.shape[-2]) model_path = model_name + ".pt" - - total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) - print(f'Total parameters = {total_params}') - - optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4) - scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma) - + + # total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + # print(f"Total parameters = {total_params}") + + optimizer = torch.optim.Adam( + model.parameters(), lr=learning_rate, weight_decay=1e-4 + ) + scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, step_size=scheduler_step, gamma=scheduler_gamma + ) + loss_fn = nn.MSELoss(reduction="mean") - loss_val_min = np.infty - + loss_val_min = np.inf + start_epoch = 0 if not if_training: checkpoint = torch.load(model_path, map_location=device) - model.load_state_dict(checkpoint['model_state_dict']) + model.load_state_dict(checkpoint["model_state_dict"]) model.to(device) model.eval() - Lx, Ly, Lz = 1., 1., 1. - errs = metrics(val_loader, model, Lx, Ly, Lz, plot, channel_plot, - model_name, x_min, x_max, y_min, y_max, - t_min, t_max, initial_step=initial_step) - pickle.dump(errs, open(model_name+'.pickle', "wb")) - + Lx, Ly, Lz = 1.0, 1.0, 1.0 + errs = metrics( + val_loader, + model, + Lx, + Ly, + Lz, + plot, + channel_plot, + model_name, + x_min, + x_max, + y_min, + y_max, + t_min, + t_max, + initial_step=initial_step, + ) + with Path(model_name + ".pickle").open("wb") as pb: + pickle.dump(errs, pb) + return # If desired, restore the network by loading the weights saved in the .pt # file if continue_training: - print('Restoring model (that is the network\'s weights) from file...') + # print("Restoring model (that is the network's weights) from file...") checkpoint = torch.load(model_path, map_location=device) - model.load_state_dict(checkpoint['model_state_dict']) + model.load_state_dict(checkpoint["model_state_dict"]) model.to(device) model.train() - + # Load optimizer state dict - optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.to(device) - - start_epoch = checkpoint['epoch'] - loss_val_min = checkpoint['loss'] - + + start_epoch = checkpoint["epoch"] + loss_val_min = checkpoint["loss"] + for ep in range(start_epoch, epochs): model.train() - t1 = default_timer() + # t1 = default_timer() train_l2_step = 0 train_l2_full = 0 for xx, yy, grid in train_loader: loss = 0 - + # xx: input tensor (first few time steps) [b, x1, ..., xd, t_init, v] # yy: target tensor [b, x1, ..., xd, t, v] # grid: meshgrid [b, x1, ..., xd, dims] - xx = xx.to(device) - yy = yy.to(device) - grid = grid.to(device) - + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 + grid = grid.to(device) # noqa: PLW2901 + # Initialize the prediction tensor pred = yy[..., :initial_step, :] # Extract shape of the input tensor for reshaping (i.e. stacking the @@ -200,16 +221,15 @@ def run_training(if_training, inp_shape = list(xx.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - - if training_type in ['autoregressive']: + + if training_type in ["autoregressive"]: # Autoregressive loop for t in range(initial_step, t_train): - # Reshape input tensor into [b, x1, ..., xd, t_init*v] inp = xx.reshape(inp_shape) - + # Extract target at current time step - y = yy[..., t:t+1, :] + y = yy[..., t : t + 1, :] # Model run im = model(inp, grid) @@ -217,34 +237,35 @@ def run_training(if_training, # Loss calculation _batch = im.size(0) loss += loss_fn(im.reshape(_batch, -1), y.reshape(_batch, -1)) - + # Concatenate the prediction at current time step into the # prediction tensor pred = torch.cat((pred, im), -2) - + # Concatenate the prediction at the current time step to be used # as input for the next time step - xx = torch.cat((xx[..., 1:, :], im), dim=-2) + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 train_l2_step += loss.item() _batch = yy.size(0) - l2_full = loss_fn(pred.reshape(_batch, -1), yy.reshape(_batch, -1)) + _yy = yy[..., :t_train, :] # if t_train is not -1 + l2_full = loss_fn(pred.reshape(_batch, -1), _yy.reshape(_batch, -1)) train_l2_full += l2_full.item() - + optimizer.zero_grad() loss.backward() optimizer.step() - if training_type in ['single']: - x = xx[..., 0 , :] - y = yy[..., t_train:t_train+1 , :] + if training_type in ["single"]: + x = xx[..., 0, :] + y = yy[..., t_train - 1 : t_train, :] pred = model(x, grid) _batch = yy.size(0) loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1)) - + train_l2_step += loss.item() train_l2_full += loss.item() - + optimizer.zero_grad() loss.backward() optimizer.step() @@ -255,57 +276,68 @@ def run_training(if_training, with torch.no_grad(): for xx, yy, grid in val_loader: loss = 0 - xx = xx.to(device) - yy = yy.to(device) - grid = grid.to(device) - - if training_type in ['autoregressive']: + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 + grid = grid.to(device) # noqa: PLW2901 + + if training_type in ["autoregressive"]: pred = yy[..., :initial_step, :] inp_shape = list(xx.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - + for t in range(initial_step, yy.shape[-2]): inp = xx.reshape(inp_shape) - y = yy[..., t:t+1, :] + y = yy[..., t : t + 1, :] im = model(inp, grid) _batch = im.size(0) - loss += loss_fn(im.reshape(_batch, -1), y.reshape(_batch, -1)) + loss += loss_fn( + im.reshape(_batch, -1), y.reshape(_batch, -1) + ) pred = torch.cat((pred, im), -2) - - xx = torch.cat((xx[..., 1:, :], im), dim=-2) - + + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 + val_l2_step += loss.item() _batch = yy.size(0) - val_l2_full += loss_fn(pred.reshape(_batch, -1), yy.reshape(_batch, -1)).item() + _pred = pred[..., initial_step:t_train, :] + _yy = yy[..., initial_step:t_train, :] + val_l2_full += loss_fn( + _pred.reshape(_batch, -1), _yy.reshape(_batch, -1) + ).item() - if training_type in ['single']: - x = xx[..., 0 , :] - y = yy[..., t_train:t_train+1 , :] + if training_type in ["single"]: + x = xx[..., 0, :] + y = yy[..., t_train - 1 : t_train, :] pred = model(x, grid) _batch = yy.size(0) loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1)) - + val_l2_step += loss.item() val_l2_full += loss.item() - - if val_l2_full < loss_val_min: + + if val_l2_full < loss_val_min: loss_val_min = val_l2_full - torch.save({ - 'epoch': ep, - 'model_state_dict': model.state_dict(), - 'optimizer_state_dict': optimizer.state_dict(), - 'loss': loss_val_min - }, model_path) - - - t2 = default_timer() + torch.save( + { + "epoch": ep, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "loss": loss_val_min, + }, + model_path, + ) + + # t2 = default_timer() scheduler.step() - print('epoch: {0}, loss: {1:.5f}, t2-t1: {2:.5f}, trainL2: {3:.5f}, testL2: {4:.5f}'\ - .format(ep, loss.item(), t2 - t1, train_l2_full, val_l2_full)) - + # print( + # "epoch: {0}, loss: {1:.5f}, t2-t1: {2:.5f}, trainL2: {3:.5f}, testL2: {4:.5f}".format( + # ep, loss.item(), t2 - t1, train_l2_full, val_l2_full + # ) + # ) + + if __name__ == "__main__": - run_training() - print("Done.") \ No newline at end of file + # print("Done.") diff --git a/pdebench/models/fno/utils.py b/pdebench/models/fno/utils.py index 6b52982..aae1232 100644 --- a/pdebench/models/fno/utils.py +++ b/pdebench/models/fno/utils.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- """ @@ -148,28 +147,32 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import torch -from torch.utils.data import Dataset, IterableDataset -from torch.utils.data import DataLoader -import os -import glob +from __future__ import annotations + +import math as mt +from pathlib import Path + import h5py import numpy as np -import math as mt +import torch +from torch.utils.data import Dataset + class FNODatasetSingle(Dataset): - def __init__(self, filename, - initial_step=10, - saved_folder='../data/', - reduced_resolution=1, - reduced_resolution_t=1, - reduced_batch=1, - if_test=False, - test_ratio=0.1, - num_samples_max = -1 - ): + def __init__( + self, + filename, + initial_step=10, + saved_folder="../data/", + reduced_resolution=1, + reduced_resolution_t=1, + reduced_batch=1, + if_test=False, + test_ratio=0.1, + num_samples_max=-1, + ): """ - + :param filename: filename that contains the dataset :type filename: STR :param filenum: array containing indices of filename included in the dataset @@ -178,118 +181,329 @@ def __init__(self, filename, :type initial_step: INT, optional """ - + # Define path to files - root_path = os.path.abspath(saved_folder + filename) - assert filename[-2:] != 'h5', 'HDF5 data is assumed!!' - - with h5py.File(root_path, 'r') as f: - keys = list(f.keys()) - keys.sort() - if 'tensor' not in keys: - _data = np.array(f['density'], dtype=np.float32) # batch, time, x,... - idx_cfd = _data.shape - if len(idx_cfd)==3: # 1D - self.data = np.zeros([idx_cfd[0]//reduced_batch, - idx_cfd[2]//reduced_resolution, - mt.ceil(idx_cfd[1]/reduced_resolution_t), - 3], - dtype=np.float32) - #density - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data[:, :, :], (0, 2, 1)) - self.data[...,0] = _data # batch, x, t, ch - # pressure - _data = np.array(f['pressure'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data[:, :, :], (0, 2, 1)) - self.data[...,1] = _data # batch, x, t, ch - # Vx - _data = np.array(f['Vx'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data[:, :, :], (0, 2, 1)) - self.data[...,2] = _data # batch, x, t, ch - - self.grid = np.array(f["x-coordinate"], dtype=np.float32) - self.grid = torch.tensor(self.grid[::reduced_resolution], dtype=torch.float).unsqueeze(-1) - print(self.data.shape) - if len(idx_cfd)==4: # 2D - self.data = np.zeros([idx_cfd[0]//reduced_batch, - idx_cfd[2]//reduced_resolution, - idx_cfd[3]//reduced_resolution, - mt.ceil(idx_cfd[1]/reduced_resolution_t), - 4], - dtype=np.float32) - # density - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,0] = _data # batch, x, t, ch - # pressure - _data = np.array(f['pressure'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,1] = _data # batch, x, t, ch - # Vx - _data = np.array(f['Vx'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,2] = _data # batch, x, t, ch - # Vy - _data = np.array(f['Vy'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,3] = _data # batch, x, t, ch - - x = np.array(f["x-coordinate"], dtype=np.float32) - y = np.array(f["y-coordinate"], dtype=np.float32) - x = torch.tensor(x, dtype=torch.float) - y = torch.tensor(y, dtype=torch.float) - X, Y = torch.meshgrid(x, y) - self.grid = torch.stack((X, Y), axis=-1)[::reduced_resolution, ::reduced_resolution] - - else: # scalar equations - ## data dim = [t, x1, ..., xd, v] - _data = np.array(f['tensor'], dtype=np.float32) # batch, time, x,... - if len(_data.shape) == 3: # 1D - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data[:, :, :], (0, 2, 1)) - self.data = _data[:, :, :, None] # batch, x, t, ch - - self.grid = np.array(f["x-coordinate"], dtype=np.float32) - self.grid = torch.tensor(self.grid[::reduced_resolution], dtype=torch.float).unsqueeze(-1) - if len(_data.shape) == 4: # 2D Darcy flow - # u: label - _data = _data[::reduced_batch,:,::reduced_resolution,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1)) - #if _data.shape[-1]==1: # if nt==1 - # _data = np.tile(_data, (1, 1, 1, 2)) - self.data = _data - # nu: input - _data = np.array(f['nu'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch, None,::reduced_resolution,::reduced_resolution] - ## convert to [x1, ..., xd, t, v] - _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1)) - self.data = np.concatenate([_data, self.data], axis=-1) - self.data = self.data[:, :, :, :, None] # batch, x, y, t, ch - - x = np.array(f["x-coordinate"], dtype=np.float32) - y = np.array(f["y-coordinate"], dtype=np.float32) - x = torch.tensor(x, dtype=torch.float) - y = torch.tensor(y, dtype=torch.float) - X, Y = torch.meshgrid(x, y) - self.grid = torch.stack((X, Y), axis=-1)[::reduced_resolution, ::reduced_resolution] - - if num_samples_max>0: - num_samples_max = min(num_samples_max,self.data.shape[0]) + root_path = Path(Path(saved_folder).resolve()) / filename + if filename[-2:] != "h5": + # print(".HDF5 file extension is assumed hereafter") + + with h5py.File(root_path, "r") as f: + keys = list(f.keys()) + keys.sort() + if "tensor" not in keys: + _data = np.array( + f["density"], dtype=np.float32 + ) # batch, time, x,... + idx_cfd = _data.shape + if len(idx_cfd) == 3: # 1D + self.data = np.zeros( + [ + idx_cfd[0] // reduced_batch, + idx_cfd[2] // reduced_resolution, + mt.ceil(idx_cfd[1] / reduced_resolution_t), + 3, + ], + dtype=np.float32, + ) + # density + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data[:, :, :], (0, 2, 1)) + self.data[..., 0] = _data # batch, x, t, ch + # pressure + _data = np.array( + f["pressure"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data[:, :, :], (0, 2, 1)) + self.data[..., 1] = _data # batch, x, t, ch + # Vx + _data = np.array( + f["Vx"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data[:, :, :], (0, 2, 1)) + self.data[..., 2] = _data # batch, x, t, ch + + self.grid = np.array(f["x-coordinate"], dtype=np.float32) + self.grid = torch.tensor( + self.grid[::reduced_resolution], dtype=torch.float + ).unsqueeze(-1) + # print(self.data.shape) + if len(idx_cfd) == 4: # 2D + self.data = np.zeros( + [ + idx_cfd[0] // reduced_batch, + idx_cfd[2] // reduced_resolution, + idx_cfd[3] // reduced_resolution, + mt.ceil(idx_cfd[1] / reduced_resolution_t), + 4, + ], + dtype=np.float32, + ) + # density + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 1)) + self.data[..., 0] = _data # batch, x, t, ch + # pressure + _data = np.array( + f["pressure"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 1)) + self.data[..., 1] = _data # batch, x, t, ch + # Vx + _data = np.array( + f["Vx"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 1)) + self.data[..., 2] = _data # batch, x, t, ch + # Vy + _data = np.array( + f["Vy"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 1)) + self.data[..., 3] = _data # batch, x, t, ch + + x = np.array(f["x-coordinate"], dtype=np.float32) + y = np.array(f["y-coordinate"], dtype=np.float32) + x = torch.tensor(x, dtype=torch.float) + y = torch.tensor(y, dtype=torch.float) + X, Y = torch.meshgrid(x, y, indexing="ij") + self.grid = torch.stack((X, Y), axis=-1)[ + ::reduced_resolution, ::reduced_resolution + ] + + if len(idx_cfd) == 5: # 3D + self.data = np.zeros( + [ + idx_cfd[0] // reduced_batch, + idx_cfd[2] // reduced_resolution, + idx_cfd[3] // reduced_resolution, + idx_cfd[4] // reduced_resolution, + mt.ceil(idx_cfd[1] / reduced_resolution_t), + 5, + ], + dtype=np.float32, + ) + # density + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 0] = _data # batch, x, t, ch + # pressure + _data = np.array( + f["pressure"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 1] = _data # batch, x, t, ch + # Vx + _data = np.array( + f["Vx"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 2] = _data # batch, x, t, ch + # Vy + _data = np.array( + f["Vy"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 3] = _data # batch, x, t, ch + # Vz + _data = np.array( + f["Vz"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 4] = _data # batch, x, t, ch + + x = np.array(f["x-coordinate"], dtype=np.float32) + y = np.array(f["y-coordinate"], dtype=np.float32) + z = np.array(f["z-coordinate"], dtype=np.float32) + x = torch.tensor(x, dtype=torch.float) + y = torch.tensor(y, dtype=torch.float) + z = torch.tensor(z, dtype=torch.float) + X, Y, Z = torch.meshgrid(x, y, z, indexing="ij") + self.grid = torch.stack((X, Y, Z), axis=-1)[ + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + + else: # scalar equations + ## data dim = [t, x1, ..., xd, v] + _data = np.array( + f["tensor"], dtype=np.float32 + ) # batch, time, x,... + if len(_data.shape) == 3: # 1D + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data[:, :, :], (0, 2, 1)) + self.data = _data[:, :, :, None] # batch, x, t, ch + + self.grid = np.array(f["x-coordinate"], dtype=np.float32) + self.grid = torch.tensor( + self.grid[::reduced_resolution], dtype=torch.float + ).unsqueeze(-1) + if len(_data.shape) == 4: # 2D Darcy flow + # u: label + _data = _data[ + ::reduced_batch, + :, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1)) + # if _data.shape[-1]==1: # if nt==1 + # _data = np.tile(_data, (1, 1, 1, 2)) + self.data = _data + # nu: input + _data = np.array( + f["nu"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + None, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1)) + self.data = np.concatenate([_data, self.data], axis=-1) + self.data = self.data[:, :, :, :, None] # batch, x, y, t, ch + + x = np.array(f["x-coordinate"], dtype=np.float32) + y = np.array(f["y-coordinate"], dtype=np.float32) + x = torch.tensor(x, dtype=torch.float) + y = torch.tensor(y, dtype=torch.float) + X, Y = torch.meshgrid(x, y, indexing="ij") + self.grid = torch.stack((X, Y), axis=-1)[ + ::reduced_resolution, ::reduced_resolution + ] + + elif filename[-2:] == "h5": # SWE-2D (RDB) + # print(".H5 file extension is assumed hereafter") + + with h5py.File(root_path, "r") as f: + keys = list(f.keys()) + keys.sort() + + data_arrays = [ + np.array(f[key]["data"], dtype=np.float32) for key in keys + ] + _data = torch.from_numpy( + np.stack(data_arrays, axis=0) + ) # [batch, nt, nx, ny, nc] + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ..., + ] + _data = torch.permute(_data, (0, 2, 3, 1, 4)) # [batch, nx, ny, nt, nc] + gridx, gridy = ( + np.array(f["0023"]["grid"]["x"], dtype=np.float32), + np.array(f["0023"]["grid"]["y"], dtype=np.float32), + ) + mgridX, mgridY = np.meshgrid(gridx, gridy, indexing="ij") + _grid = torch.stack( + (torch.from_numpy(mgridX), torch.from_numpy(mgridY)), axis=-1 + ) + _grid = _grid[::reduced_resolution, ::reduced_resolution, ...] + _tsteps_t = torch.from_numpy( + np.array(f["0023"]["grid"]["t"], dtype=np.float32) + ) + + tsteps_t = _tsteps_t[::reduced_resolution_t] + self.data = _data + self.grid = _grid + self.tsteps_t = tsteps_t + + if num_samples_max > 0: + num_samples_max = min(num_samples_max, self.data.shape[0]) else: num_samples_max = self.data.shape[0] @@ -302,27 +516,26 @@ def __init__(self, filename, # Time steps used as initial conditions self.initial_step = initial_step - self.data = torch.tensor(self.data) + self.data = self.data if torch.is_tensor(self.data) else torch.tensor(self.data) def __len__(self): return len(self.data) - + def __getitem__(self, idx): - - return self.data[idx,...,:self.initial_step,:], self.data[idx], self.grid + return self.data[idx, ..., : self.initial_step, :], self.data[idx], self.grid class FNODatasetMult(Dataset): - def __init__(self, filename, - initial_step=10, - saved_folder='../data/', - reduced_resolution=1, - reduced_resolution_t=1, - reduced_batch=1, - if_test=False, test_ratio=0.1 - ): + def __init__( + self, + filename, + initial_step=10, + saved_folder="../data/", + if_test=False, + test_ratio=0.1, + ): """ - + :param filename: filename that contains the dataset :type filename: STR :param filenum: array containing indices of filename included in the dataset @@ -331,64 +544,63 @@ def __init__(self, filename, :type initial_step: INT, optional """ - + # Define path to files - self.file_path = os.path.abspath(saved_folder + filename + ".h5") - + self.file_path = Path(saved_folder + filename + ".h5").resolve() + # Extract list of seeds - with h5py.File(self.file_path, 'r') as h5_file: + with h5py.File(self.file_path, "r") as h5_file: data_list = sorted(h5_file.keys()) - test_idx = int(len(data_list) * (1-test_ratio)) + test_idx = int(len(data_list) * (1 - test_ratio)) if if_test: self.data_list = np.array(data_list[test_idx:]) else: self.data_list = np.array(data_list[:test_idx]) - + # Time steps used as initial conditions self.initial_step = initial_step def __len__(self): return len(self.data_list) - + def __getitem__(self, idx): - # Open file and read data - with h5py.File(self.file_path, 'r') as h5_file: + with h5py.File(self.file_path, "r") as h5_file: seed_group = h5_file[self.data_list[idx]] - + # data dim = [t, x1, ..., xd, v] - data = np.array(seed_group["data"], dtype='f') + data = np.array(seed_group["data"], dtype="f") data = torch.tensor(data, dtype=torch.float) - + # convert to [x1, ..., xd, t, v] - permute_idx = list(range(1,len(data.shape)-1)) - permute_idx.extend(list([0, -1])) + permute_idx = list(range(1, len(data.shape) - 1)) + permute_idx.extend([0, -1]) data = data.permute(permute_idx) - + # Extract spatial dimension of data - dim = len(data.shape) - 2 - + dim = len(data.shape) - 2 + # x, y and z are 1-D arrays # Convert the spatial coordinates to meshgrid if dim == 1: - grid = np.array(seed_group["grid"]["x"], dtype='f') + grid = np.array(seed_group["grid"]["x"], dtype="f") grid = torch.tensor(grid, dtype=torch.float).unsqueeze(-1) elif dim == 2: - x = np.array(seed_group["grid"]["x"], dtype='f') - y = np.array(seed_group["grid"]["y"], dtype='f') + x = np.array(seed_group["grid"]["x"], dtype="f") + y = np.array(seed_group["grid"]["y"], dtype="f") x = torch.tensor(x, dtype=torch.float) y = torch.tensor(y, dtype=torch.float) - X, Y = torch.meshgrid(x, y) - grid = torch.stack((X,Y),axis=-1) + X, Y = torch.meshgrid(x, y, indexing="ij") + grid = torch.stack((X, Y), axis=-1) elif dim == 3: - x = np.array(seed_group["grid"]["x"], dtype='f') - y = np.array(seed_group["grid"]["y"], dtype='f') - z = np.array(seed_group["grid"]["z"], dtype='f') + x = np.array(seed_group["grid"]["x"], dtype="f") + y = np.array(seed_group["grid"]["y"], dtype="f") + z = np.array(seed_group["grid"]["z"], dtype="f") x = torch.tensor(x, dtype=torch.float) y = torch.tensor(y, dtype=torch.float) z = torch.tensor(z, dtype=torch.float) - X, Y, Z = torch.meshgrid(x, y, z) - grid = torch.stack((X,Y,Z),axis=-1) - - return data[...,:self.initial_step,:], data, grid \ No newline at end of file + X, Y, Z = torch.meshgrid(x, y, z, indexing="ij") + grid = torch.stack((X, Y, Z), axis=-1) + + return data[..., : self.initial_step, :], data, grid diff --git a/pdebench/models/inverse/inverse.py b/pdebench/models/inverse/inverse.py index 44c4cff..54f3e82 100644 --- a/pdebench/models/inverse/inverse.py +++ b/pdebench/models/inverse/inverse.py @@ -1,10 +1,9 @@ -# -*- coding: utf-8 -*- """ File: inverse.py Authors: Francesco Alesiani (makoto.takamoto@neclab.eu) - Dan MacKinlay (Dan.MacKinlay@data61.csiro.au) + Dan MacKinlay (Dan.MacKinlay@data61.csiro.au) NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. @@ -146,48 +145,44 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import sys -import torch -import numpy as np -import pickle -import torch.nn as nn -import torch.nn.functional as F - -import operator -from functools import reduce -from functools import partial +from __future__ import annotations import pyro -from pyro.nn import PyroModule, PyroSample import pyro.distributions as dist -from pyro.infer import MCMC, NUTS -from pyro import poutine +import torch +import torch.nn.functional as F +from numpy import prod +from pyro.nn import PyroModule, PyroSample +from torch import nn class ElementStandardScaler: - def fit(self, x): - self.mean = x.mean() - self.std = x.std(unbiased=False) - def transform(self, x): - eps = 1e-20 - x = x - self.mean - x = x/(self.std + eps) - return x - def fit_transform(self, x): - self.fit(x) - return self.transform(x) + def fit(self, x): + self.mean = x.mean() + self.std = x.std(unbiased=False) + + def transform(self, x): + eps = 1e-20 + x = x - self.mean + return x / (self.std + eps) + + def fit_transform(self, x): + self.fit(x) + return self.transform(x) + class ProbRasterLatent(PyroModule): def __init__( - self, - process_predictor: "nn.Module", - dims = (256,256), - latent_dims = (16,16), - interpolation = "bilinear", - prior_scale = 0.01, - obs_scale = 0.01, - prior_std = 0.01, - device=None): + self, + process_predictor: nn.Module, + dims=(256, 256), + latent_dims=(16, 16), + interpolation="bilinear", + prior_scale=0.01, + obs_scale=0.01, + prior_std=0.01, + device=None, + ): super().__init__() self.dims = dims self.device = device @@ -200,64 +195,46 @@ def __init__( self.obs_scale = torch.tensor(obs_scale, device=self.device, dtype=torch.float) self.process_predictor = process_predictor process_predictor.train(False) - ## Do not fit the process predictor weights + # Do not fit the process predictor weights for param in self.process_predictor.parameters(): param.requires_grad = False - _m,_s = torch.tensor([0], device=self.device, dtype=torch.float), torch.tensor([self.prior_std], device=self.device, dtype=torch.float) - self.latent = PyroSample(dist.Normal(_m,_s).expand(latent_dims).to_event(2)) - print(self.latent_dims,self.dims) + _m, _s = ( + torch.tensor([0], device=self.device, dtype=torch.float), + torch.tensor([self.prior_std], device=self.device, dtype=torch.float), + ) + self.latent = PyroSample(dist.Normal(_m, _s).expand(latent_dims).to_event(2)) def get_latent(self): - if self.latent_dims==self.dims: + if self.latent_dims == self.dims: return self.latent.unsqueeze(0) # `mini-batch x channels x [optional depth] x [optional height] x width`. - l = F.interpolate( + return F.interpolate( self.latent.unsqueeze(1), self.dims, mode=self.interpolation, - align_corners=False - ).squeeze(0) #squeeze/unsqueeze is because of weird interpolate semantics - return l + align_corners=False, + ).squeeze(0) # squeeze/unsqueeze is because of weird interpolate semantics - def latent2source(self,latent): - if latent.shape==self.dims: + def latent2source(self, latent): + if latent.shape == self.dims: return latent.unsqueeze(0) # `mini-batch x channels x [optional depth] x [optional height] x width`. - l = F.interpolate( - latent.unsqueeze(1), - self.dims, - mode=self.interpolation, - align_corners=False - ).squeeze(0) #squeeze/unsqueeze is because of weird interpolate semantics - return l + return F.interpolate( + latent.unsqueeze(1), self.dims, mode=self.interpolation, align_corners=False + ).squeeze(0) # squeeze/unsqueeze is because of weird interpolate semantics def forward(self, grid, y=None): - #overwrite process predictor batch with my own latent + # overwrite process predictor batch with my own latent x = self.get_latent() # print("forward:x.shape,grid.shape=",x.shape,grid.shape) - mean = self.process_predictor(x.to(self.device),grid.to(self.device)) - o = pyro.sample( - "obs", dist.Normal(mean, self.obs_scale).to_event(2), - obs=y) - return o - - -import sys -import torch -import numpy as np -import pickle -import torch.nn as nn -import torch.nn.functional as F + mean = self.process_predictor(x.to(self.device), grid.to(self.device)) + return pyro.sample("obs", dist.Normal(mean, self.obs_scale).to_event(2), obs=y) -import operator -from functools import reduce -from functools import partial -from numpy import prod class InitialConditionInterp(nn.Module): """ InitialConditionInterp - Class for the inital conditions using interpoliation. Works for 1d,2d and 3d + Class for the initial conditions using interpoliation. Works for 1d,2d and 3d model_ic = InitialConditionInterp([16],[8]) model_ic = InitialConditionInterp([16,16],[8,8]) @@ -265,30 +242,31 @@ class InitialConditionInterp(nn.Module): June 2022, F.Alesiani """ + def __init__(self, dims, hidden_dim): - super(InitialConditionInterp, self).__init__() + super().__init__() self.spatial_dim = len(hidden_dim) - self.dims = [1]+dims if len(dims)==1 else dims + self.dims = [1, *dims] if len(dims) == 1 else dims # self.dims = [1,1,1]+dims - self.hidden_dim = [1]+hidden_dim if len(hidden_dim)==1 else hidden_dim - self.interpolation = "bilinear" if len(hidden_dim)<3 else "trilinear" - self.scale = (1 / prod(hidden_dim)) - self.latent = nn.Parameter(self.scale * torch.rand(1, 1, *self.hidden_dim, dtype=torch.float)) + self.hidden_dim = [1, *hidden_dim] if len(hidden_dim) == 1 else hidden_dim + self.interpolation = "bilinear" if len(hidden_dim) < 3 else "trilinear" + self.scale = 1 / prod(hidden_dim) + self.latent = nn.Parameter( + self.scale * torch.rand(1, 1, *self.hidden_dim, dtype=torch.float) + ) # print(self.latent.shape) - def latent2source(self,latent): - if latent.shape[2:]==self.dims: + def latent2source(self, latent): + if latent.shape[2:] == self.dims: return latent # `mini-batch x channels x [optional depth] x [optional height] x width`. - l = F.interpolate( - latent, - self.dims, - mode=self.interpolation, - align_corners=False - ) - return l.view(self.dims) + latent = F.interpolate( + latent, self.dims, mode=self.interpolation, align_corners=False + ) + return latent.view(self.dims) + def forward(self): x = self.latent2source(self.latent) if self.spatial_dim == 1: - x = x.squeeze(0) - return x \ No newline at end of file + x = x.squeeze(0) + return x diff --git a/pdebench/models/inverse/train.py b/pdebench/models/inverse/train.py index d8cfcf1..d0e5d92 100644 --- a/pdebench/models/inverse/train.py +++ b/pdebench/models/inverse/train.py @@ -1,10 +1,9 @@ -# -*- coding: utf-8 -*- """ File: inverse.py Authors: Francesco Alesiani (makoto.takamoto@neclab.eu) - Dan MacKinlay (Dan.MacKinlay@data61.csiro.au) + Dan MacKinlay (Dan.MacKinlay@data61.csiro.au) NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. @@ -146,265 +145,308 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import sys -import torch -import numpy as np -import pickle -import torch.nn as nn -import torch.nn.functional as F - -import operator -from functools import reduce -from functools import partial - -import pyro -from pyro.nn import PyroModule, PyroSample -import pyro.distributions as dist -from pyro.infer import MCMC, NUTS -from pyro import poutine +from __future__ import annotations +import logging from timeit import default_timer - -import sys, os import hydra +import pandas as pd +import torch from omegaconf import DictConfig -from omegaconf import OmegaConf -from omegaconf import open_dict - - -import pdebench as pde -from pdebench.models.fno.fno import FNO1d,FNO2d,FNO3d -from pdebench.models.fno.utils import FNODatasetSingle, FNODatasetMult - +from pdebench.models.fno.fno import FNO1d, FNO2d, FNO3d +from pdebench.models.fno.utils import FNODatasetSingle +from pdebench.models.inverse.inverse import ( + ElementStandardScaler, + InitialConditionInterp, + ProbRasterLatent, +) +from pdebench.models.metrics import inverse_metrics from pdebench.models.unet.unet import UNet1d, UNet2d, UNet3d -from pdebench.models.unet.utils import UNetDatasetSingle,UNetDatasetMult - -from pdebench.models import metrics -from pdebench.models.metrics import LpLoss,FftLpLoss,FftMseLoss,inverse_metrics -import pandas as pd - - -from pdebench.models.inverse.inverse import ProbRasterLatent, ElementStandardScaler, InitialConditionInterp -from pdebench.models.inverse.utils import plot_ic_solution_mcmc +from pdebench.models.unet.utils import UNetDatasetSingle +from pyro.infer import MCMC, NUTS +from torch import nn +from tqdm import tqdm -from torch.distributions.normal import Normal +logger = logging.getLogger(__name__) -from tqdm import tqdm +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') -def load_model(model,model_path, device): +def load_model(model, model_path, device): checkpoint = torch.load(model_path, map_location=device) - model.load_state_dict(checkpoint['model_state_dict']) + model.load_state_dict(checkpoint["model_state_dict"]) model.to(device) model.eval() return model -@hydra.main(config_path='../config', config_name='config') +@hydra.main(config_path="../config", config_name="config") def main(cfg: DictConfig): - print(cfg.args.filename) - print(cfg.args) - - # we use the test data - if cfg.args.model_name in ['FNO']: - inverse_data = FNODatasetSingle(cfg.args.filename, - saved_folder = cfg.args.base_path, - reduced_resolution=cfg.args.reduced_resolution, - reduced_resolution_t=cfg.args.reduced_resolution_t, - reduced_batch=cfg.args.reduced_batch, - initial_step=cfg.args.initial_step, - if_test=True, - num_samples_max = cfg.args.num_samples_max - ) - - _data, _, _ = next(iter(inverse_loader)) + logger.info(cfg.args.filename) + logger.info(cfg.args) + + # we use the test data + if cfg.args.model_name in ["FNO"]: + inverse_data = FNODatasetSingle( + cfg.args.filename, + saved_folder=cfg.args.base_path, + reduced_resolution=cfg.args.reduced_resolution, + reduced_resolution_t=cfg.args.reduced_resolution_t, + reduced_batch=cfg.args.reduced_batch, + initial_step=cfg.args.initial_step, + if_test=True, + num_samples_max=cfg.args.num_samples_max, + ) + + _data, _, _ = next(iter(inverse_data)) dimensions = len(_data.shape) spatial_dim = dimensions - 3 - if cfg.args.model_name in ['UNET','Unet']: - inverse_data = UNetDatasetSingle(cfg.args.filename, - saved_folder = cfg.args.base_path, - reduced_resolution=cfg.args.reduced_resolution, - reduced_resolution_t=cfg.args.reduced_resolution_t, - reduced_batch=cfg.args.reduced_batch, - initial_step=cfg.args.initial_step, - if_test=True, - num_samples_max = cfg.args.num_samples_max) - - inverse_loader = torch.utils.data.DataLoader(inverse_data, batch_size=1,shuffle=False) - _data, _ = next(iter(inverse_loader)) + if cfg.args.model_name in ["UNET", "Unet"]: + inverse_data = UNetDatasetSingle( + cfg.args.filename, + saved_folder=cfg.args.base_path, + reduced_resolution=cfg.args.reduced_resolution, + reduced_resolution_t=cfg.args.reduced_resolution_t, + reduced_batch=cfg.args.reduced_batch, + initial_step=cfg.args.initial_step, + if_test=True, + num_samples_max=cfg.args.num_samples_max, + ) + + inverse_loader = torch.utils.data.DataLoader( + inverse_data, batch_size=1, shuffle=False + ) + _data, _ = next(iter(inverse_loader)) dimensions = len(_data.shape) spatial_dim = dimensions - 3 - initial_step = cfg.args.initial_step t_train = cfg.args.t_train - - model_name = cfg.args.filename[:-5] + '_' + cfg.args.model_name + + model_name = cfg.args.filename[:-5] + "_" + cfg.args.model_name model_path = cfg.args.base_path + model_name + ".pt" - if cfg.args.model_name in ['FNO']: + if cfg.args.model_name in ["FNO"]: if dimensions == 4: - print(cfg.args.num_channels) - model = FNO1d(num_channels=cfg.args.num_channels, - width=cfg.args.width, - modes=cfg.args.modes, - initial_step=cfg.args.initial_step).to(device) + logger.info(cfg.args.num_channels) + model = FNO1d( + num_channels=cfg.args.num_channels, + width=cfg.args.width, + modes=cfg.args.modes, + initial_step=cfg.args.initial_step, + ).to(device) if dimensions == 5: - model = FNO2d(num_channels=cfg.args.num_channels, - width=cfg.args.width, - modes1=cfg.args.modes, - modes2=cfg.args.modes, - initial_step=cfg.args.initial_step).to(device) + model = FNO2d( + num_channels=cfg.args.num_channels, + width=cfg.args.width, + modes1=cfg.args.modes, + modes2=cfg.args.modes, + initial_step=cfg.args.initial_step, + ).to(device) if dimensions == 6: - model = FNO2d(num_channels=cfg.args.num_channels, - width=cfg.args.width, - modes1=cfg.args.modes, - modes2=cfg.args.modes, - modes3=cfg.args.modes, - initial_step=cfg.args.initial_step).to(device) - - if cfg.args.model_name in ['UNET','Unet']: + model = FNO3d( + num_channels=cfg.args.num_channels, + width=cfg.args.width, + modes1=cfg.args.modes, + modes2=cfg.args.modes, + modes3=cfg.args.modes, + initial_step=cfg.args.initial_step, + ).to(device) + + if cfg.args.model_name in ["UNET", "Unet"]: if dimensions == 4: model = UNet1d(cfg.args.in_channels, cfg.args.out_channels).to(device) elif dimensions == 5: model = UNet2d(cfg.args.in_channels, cfg.args.out_channels).to(device) elif dimensions == 6: - model = UNet3d(cfg.args.in_channels, cfg.args.out_channels).to(device) + model = UNet3d(cfg.args.in_channels, cfg.args.out_channels).to(device) - model = load_model(model,model_path, device) + model = load_model(model, model_path, device) model.eval() - if cfg.args.inverse_model_type in ['ProbRasterLatent']: - assert(spatial_dim==1), "give me time" - if spatial_dim==1: - ns,nx,nt,nc = _data.shape + if cfg.args.inverse_model_type in ["ProbRasterLatent"]: + assert spatial_dim == 1, "give me time" + if spatial_dim == 1: + ns, nx, nt, nc = _data.shape model_inverse = ProbRasterLatent( model.to(device), - dims=[nx,1], - latent_dims = [1,cfg.args.in_channels_hid,1], - prior_scale = 0.1, - obs_scale = 0.01, - prior_std = 0.01, - device=device - ) - - if cfg.args.inverse_model_type in ['InitialConditionInterp']: + dims=[nx, 1], + latent_dims=[1, cfg.args.in_channels_hid, 1], + prior_scale=0.1, + obs_scale=0.01, + prior_std=0.01, + device=device, + ) + + if cfg.args.inverse_model_type in ["InitialConditionInterp"]: loss_fn = nn.MSELoss(reduction="mean") - input_dims = list(_data.shape[1:1+spatial_dim]) - latent_dims = len(input_dims)*[cfg.args.in_channels_hid] - if cfg.args.num_channels> 1: - input_dims=input_dims+[cfg.args.num_channels] - latent_dims=latent_dims+[cfg.args.num_channels] - print(input_dims,latent_dims) - model_ic = InitialConditionInterp(input_dims,latent_dims).to(device) - model.to(device) + input_dims = list(_data.shape[1 : 1 + spatial_dim]) + latent_dims = len(input_dims) * [cfg.args.in_channels_hid] + if cfg.args.num_channels > 1: + input_dims = [*input_dims, cfg.args.num_channels] + latent_dims = [*latent_dims, cfg.args.num_channels] + model_ic = InitialConditionInterp(input_dims, latent_dims).to(device) + model.to(device) scaler = ElementStandardScaler() loss_fn = nn.MSELoss(reduction="mean") - inverse_u0_l2_full,inverse_y_l2_full = 0,0 + inverse_u0_l2_full, inverse_y_l2_full = 0, 0 all_metric = [] t1 = default_timer() - for ks,sample in enumerate(inverse_loader): - if cfg.args.model_name in ['FNO']: + for ks, sample in enumerate(inverse_loader): + if cfg.args.model_name in ["FNO"]: (xx, yy, grid) = sample xx = xx.to(device) yy = yy.to(device) grid = grid.to(device) - model_ = lambda x, grid: model(x,grid) - if cfg.args.model_name in ['UNET','Unet']: + def model_(x, grid): + return model(x, grid) + + if cfg.args.model_name in ["UNET", "Unet"]: (xx, yy) = sample grid = None xx = xx.to(device) yy = yy.to(device) - model_ = lambda x, grid: model(x.permute([0, 2, 1])).permute([0, 2, 1]) - num_samples = ks + 1 - loss = 0 + def model_(x, grid): + return model(x.permute([0, 2, 1])).permute([0, 2, 1]) + num_samples = ks + 1 - x = xx[..., 0 , :] - y = yy[..., t_train:t_train+1 , :] + x = xx[..., 0, :] + y = yy[..., t_train : t_train + 1, :] - if ks==0: - print(x.shape,y.shape) + if ks == 0: + msg = f"{x.shape}, {y.shape}" + logger.info(msg) - #scale the input and output + # scale the input and output x = scaler.fit_transform(x) y = scaler.transform(y) - if cfg.args.inverse_model_type in ['ProbRasterLatent']: - #Create model + if cfg.args.inverse_model_type in ["ProbRasterLatent"]: + # Create model model_inverse.to(device) - nuts_kernel = NUTS(model_inverse, full_mass=False, max_tree_depth=5, jit_compile=True) # high performacne config - - mcmc = MCMC(nuts_kernel, num_samples=cfg.args.mcmc_num_samples, warmup_steps=cfg.args.mcmc_warmup_steps, num_chains=cfg.args.mcmc_num_chains,disable_progbar=True) + nuts_kernel = NUTS( + model_inverse, full_mass=False, max_tree_depth=5, jit_compile=True + ) # high performacne config + + mcmc = MCMC( + nuts_kernel, + num_samples=cfg.args.mcmc_num_samples, + warmup_steps=cfg.args.mcmc_warmup_steps, + num_chains=cfg.args.mcmc_num_chains, + disable_progbar=True, + ) mcmc.run(grid, y) - mc_samples = {k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items()} + mc_samples = { + k: v.detach().cpu().numpy() for k, v in mcmc.get_samples().items() + } # get the initial solution - latent = torch.tensor(mc_samples['latent']) + latent = torch.tensor(mc_samples["latent"]) u0 = model_inverse.latent2source(latent[0]).to(device) pred_u0 = model(u0, grid) - if cfg.args.inverse_model_type in ['InitialConditionInterp']: - optimizer = torch.optim.Adam(model_ic.parameters(), lr=cfg.args.inverse_learning_rate, weight_decay=1e-4) + if cfg.args.inverse_model_type in ["InitialConditionInterp"]: + optimizer = torch.optim.Adam( + model_ic.parameters(), + lr=cfg.args.inverse_learning_rate, + weight_decay=1e-4, + ) # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma) if cfg.args.inverse_verbose_flag: _iter = tqdm(range(cfg.args.inverse_epochs)) else: _iter = range(cfg.args.inverse_epochs) - for epoch in _iter: - if cfg.args.num_channels>1: + for _ in _iter: + if cfg.args.num_channels > 1: u0 = model_ic().unsqueeze(0) else: u0 = model_ic().unsqueeze(0).unsqueeze(-1) - - pred_u0 = model_(u0,grid) - - loss_u0 = loss_fn(pred_u0,y) + + pred_u0 = model_(u0, grid) + + loss_u0 = loss_fn(pred_u0, y) optimizer.zero_grad() loss_u0.backward() optimizer.step() t2 = default_timer() if cfg.args.inverse_verbose_flag: - _iter.set_description(f"loss={loss_u0.item()}, t2-t1= {t2-t1}") + _iter.set_description(f"loss={loss_u0.item()}, t2-t1= {t2-t1}") - #compute losses + # compute losses loss_u0 = loss_fn(u0.reshape(1, -1), x.reshape(1, -1)).item() loss_y = loss_fn(pred_u0.reshape(1, -1), y.reshape(1, -1)).item() inverse_u0_l2_full += loss_u0 inverse_y_l2_full += loss_y - metric = inverse_metrics(u0,x,pred_u0,y) - metric['sample'] = ks + metric = inverse_metrics(u0, x, pred_u0, y) + metric["sample"] = ks + + all_metric += [metric] - all_metric+=[metric] - t2 = default_timer() - print('samples: {}, loss_u0: {:.5f},loss_y: {:.5f}, t2-t1: {:.5f}, mse_inverse_u0_L2: {:.5f}, mse_inverse_y_L2: {:.5f}'\ - .format(ks+1, loss_u0, loss_y, t2 - t1, inverse_u0_l2_full/num_samples, inverse_y_l2_full/num_samples)) + msg = ", ".join( + [ + f"samples: {ks + 1}", + f"loss_u0: {loss_u0:.5f}", + f"loss_y: {loss_y:.5f}", + f"t2-t1: {t2 - t1:.5f}", + f"mse_inverse_u0_L2: {inverse_u0_l2_full / num_samples:.5f}", + f"mse_inverse_y_L2: {inverse_y_l2_full / num_samples:.5f}", + ] + ) + logger.info(msg) df_metric = pd.DataFrame(all_metric) - inverse_metric_filename = cfg.args.base_path + cfg.args.filename[:-5] + '_' + cfg.args.model_name +'_'+cfg.args.inverse_model_type + ".csv" - print("saving in :", inverse_metric_filename) + inverse_metric_filename = ( + cfg.args.base_path + + cfg.args.filename[:-5] + + "_" + + cfg.args.model_name + + "_" + + cfg.args.inverse_model_type + + ".csv" + ) + msg = f"saving in : {inverse_metric_filename}" + logger.info(msg) df_metric.to_csv(inverse_metric_filename) - inverse_metric_filename = cfg.args.base_path + cfg.args.filename[:-5] + '_' + cfg.args.model_name +'_'+cfg.args.inverse_model_type+ ".pickle" - print("saving in :", inverse_metric_filename) + inverse_metric_filename = ( + cfg.args.base_path + + cfg.args.filename[:-5] + + "_" + + cfg.args.model_name + + "_" + + cfg.args.inverse_model_type + + ".pickle" + ) + msg = f"saving in : {inverse_metric_filename}" + logger.info(msg) df_metric.to_pickle(inverse_metric_filename) - inverse_metric_filename = cfg.args.base_path + cfg.args.filename[:-5] + '_' + cfg.args.model_name +'_'+cfg.args.inverse_model_type+ "_stats.csv" - print("saving in :", inverse_metric_filename) + inverse_metric_filename = ( + cfg.args.base_path + + cfg.args.filename[:-5] + + "_" + + cfg.args.model_name + + "_" + + cfg.args.inverse_model_type + + "_stats.csv" + ) + msg = f"saving in : {inverse_metric_filename}" + logger.info(msg) df_metric = df_metric.describe() df_metric.to_csv(inverse_metric_filename) -if __name__ == '__main__': - main() \ No newline at end of file + +if __name__ == "__main__": + main() diff --git a/pdebench/models/inverse/utils.py b/pdebench/models/inverse/utils.py index 7e4f0ce..9bcfda8 100644 --- a/pdebench/models/inverse/utils.py +++ b/pdebench/models/inverse/utils.py @@ -1,10 +1,9 @@ -# -*- coding: utf-8 -*- """ File: utils.py Authors: Francesco Alesiani (makoto.takamoto@neclab.eu) - Dan MacKinlay (Dan.MacKinlay@data61.csiro.au) + Dan MacKinlay (Dan.MacKinlay@data61.csiro.au) NEC Laboratories Europe GmbH, Copyright (c) , All rights reserved. @@ -146,72 +145,92 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ +from __future__ import annotations +import logging + +import hydra import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from omegaconf import DictConfig +from scipy.signal import welch -def plot_ic_solution_mcmc(latent,x,y,grid,model_inverse,model,device,fname_save="IC_inverse_problem_mcmc.pdf"): +logger = logging.getLogger(__name__) + + +def plot_ic_solution_mcmc( + latent, + x, + y, + grid, + model_inverse, + model, + device, + fname_save="IC_inverse_problem_mcmc.pdf", +): """ Plots the prediction of the initial condition estimated using MCMC from the latent with the model "model" y = model(x) - y[i] = model(latent[i]), i =0, ... + y[i] = model(latent[i]), i =0, ... June 2022, F.Alesiani - """ - fig, axes = plt.subplots(1,2,figsize=(15,7)) - ax = axes[0] + """ + fig, axes = plt.subplots(1, 2, figsize=(15, 7)) + ax = axes[0] u0 = model_inverse.latent2source(latent[0]).to(device) pred_u0 = model(u0, grid) - ax.plot(u0.detach().cpu().flatten(),'r',label="Predicted Initial Condition") + ax.plot(u0.detach().cpu().flatten(), "r", label="Predicted Initial Condition") for _latent in latent: u0 = model_inverse.latent2source(_latent).to(device) - ax.plot(u0.detach().cpu().flatten(),'r',alpha=0.1) - ax.plot(x.detach().cpu().flatten(),'b--',label="True Initial Condition") + ax.plot(u0.detach().cpu().flatten(), "r", alpha=0.1) + ax.plot(x.detach().cpu().flatten(), "b--", label="True Initial Condition") ax.legend() # plt.show() - ax = axes[1] - ax.plot(pred_u0.detach().cpu().flatten(),'r',label="Predicted forward value") - ax.plot(y.detach().cpu().flatten(),'b--',label="True forward value") + ax = axes[1] + ax.plot(pred_u0.detach().cpu().flatten(), "r", label="Predicted forward value") + ax.plot(y.detach().cpu().flatten(), "b--", label="True forward value") for _latent in latent: u0 = model_inverse.latent2source(_latent).to(device) pred_u0 = model(u0, grid) - ax.plot(pred_u0.detach().cpu().flatten(),'r',alpha=0.1) + ax.plot(pred_u0.detach().cpu().flatten(), "r", alpha=0.1) ax.legend() if fname_save: - plt.savefig(fname_save, bbox_inches='tight') + plt.savefig(fname_save, bbox_inches="tight") - -def plot_ic_solution_grad(model_ic,x,y,grid,model,device,fname_save="IC_inverse_problem_grad.pdf"): +def plot_ic_solution_grad( + model_ic, x, y, grid, model, device, fname_save="IC_inverse_problem_grad.pdf" +): """ Plots the prediction of the initial condition estimated using model_ic with the model "model" y = model(x) y' = model(model_ic()) June 2022, F.Alesiani - """ + """ - fig, axes = plt.subplots(1,2,figsize=(15,7)) - ax = axes[0] + fig, axes = plt.subplots(1, 2, figsize=(15, 7)) + ax = axes[0] u0 = model_ic().to(device).unsqueeze(0).unsqueeze(-1) pred_u0 = model(u0, grid) - ax.plot(u0.detach().cpu().flatten(),'r',label="Predicted Initial Condition") - ax.plot(x.detach().cpu().flatten(),'b--',label="True Initial Condition") + ax.plot(u0.detach().cpu().flatten(), "r", label="Predicted Initial Condition") + ax.plot(x.detach().cpu().flatten(), "b--", label="True Initial Condition") ax.legend() # plt.show() - ax = axes[1] - ax.plot(pred_u0.detach().cpu().flatten(),'r',label="Predicted forward value") - ax.plot(y.detach().cpu().flatten(),'b--',label="True forward value") + ax = axes[1] + ax.plot(pred_u0.detach().cpu().flatten(), "r", label="Predicted forward value") + ax.plot(y.detach().cpu().flatten(), "b--", label="True forward value") ax.legend() if fname_save: - plt.savefig(fname_save, bbox_inches='tight') - + plt.savefig(fname_save, bbox_inches="tight") -from scipy.signal import welch -import matplotlib.pyplot as plt -def plot_ic_solution_grad_psd(model_ic,x,y,grid,model,device,fname_save="IC_inverse_problem_grad_psd.pdf"): +def plot_ic_solution_grad_psd( + model_ic, x, y, grid, model, device, fname_save="IC_inverse_problem_grad_psd.pdf" +): """ Plots the prediction of the initial condition estimated using model_ic with the model "model" y = model(x) @@ -219,19 +238,19 @@ def plot_ic_solution_grad_psd(model_ic,x,y,grid,model,device,fname_save="IC_inve It also shows the power density June 2022, F.Alesiani - """ - fig, axes = plt.subplots(1,3,figsize=(22,7)) - ax = axes[0] + """ + fig, axes = plt.subplots(1, 3, figsize=(22, 7)) + ax = axes[0] u0 = model_ic().to(device).unsqueeze(0).unsqueeze(-1) pred_u0 = model(u0, grid) - ax.plot(u0.detach().cpu().flatten(),'r',label="Predicted Initial Condition") - ax.plot(x.detach().cpu().flatten(),'b--',label="True Initial Condition") + ax.plot(u0.detach().cpu().flatten(), "r", label="Predicted Initial Condition") + ax.plot(x.detach().cpu().flatten(), "b--", label="True Initial Condition") ax.legend() # plt.show() - ax = axes[1] - ax.plot(pred_u0.detach().cpu().flatten(),'r',label="Predicted forward value") - ax.plot(y.detach().cpu().flatten(),'b--',label="True forward value") + ax = axes[1] + ax.plot(pred_u0.detach().cpu().flatten(), "r", label="Predicted forward value") + ax.plot(y.detach().cpu().flatten(), "b--", label="True forward value") ax.legend() _u0 = u0.detach().cpu().flatten() @@ -239,76 +258,90 @@ def plot_ic_solution_grad_psd(model_ic,x,y,grid,model,device,fname_save="IC_inve fz = u0.shape[1] - fu,puu = welch(_u0,fz) - fx,pxx = welch(_x,fz) + fu, puu = welch(_u0, fz) + fx, pxx = welch(_x, fz) - ax = axes[2] - ax.semilogy(fu,puu,'r',label="predicted u0") - ax.semilogy(fx,pxx,'b--',label="x true") - ax.set_xlabel('spatial frequency') - ax.set_ylabel('PSD') + ax = axes[2] + ax.semilogy(fu, puu, "r", label="predicted u0") + ax.semilogy(fx, pxx, "b--", label="x true") + ax.set_xlabel("spatial frequency") + ax.set_ylabel("PSD") ax.legend() if fname_save: - plt.savefig(fname_save, bbox_inches='tight') - - - -import sys, os -import hydra -from omegaconf import DictConfig -from omegaconf import OmegaConf -from omegaconf import open_dict -import pandas as pd -import numpy as np + plt.savefig(fname_save, bbox_inches="tight") -def get_metric_name(filename,model_name, base_path,inverse_model_type): +def get_metric_name(filename, model_name, base_path, inverse_model_type): """ returns the name convention for the result file June 2022, F.Alesiani """ - inverse_metric_filename = base_path + filename[:-5] + '_' + model_name +'_'+ inverse_model_type + ".pickle" - return inverse_metric_filename - -def read_results(model_names,inverse_model_type, base_path, filenames,shortfilenames, verbose=False): + return ( + base_path + + filename[:-5] + + "_" + + model_name + + "_" + + inverse_model_type + + ".pickle" + ) + + +def read_results( + model_names, inverse_model_type, base_path, filenames, shortfilenames, verbose=False +): """ - reads and merges the result files. + reads and merges the result files. Shortnames are used for the name of the dataset as alternative to the file name. June 2022, F.Alesiani """ dfs = [] for model_name in model_names: - for filename,shortfilename in zip(filenames,shortfilenames): + for filename, shortfilename in zip(filenames, shortfilenames): # print(filename) - inverse_metric_filename = get_metric_name(filename,model_name, base_path,inverse_model_type) - if verbose: print ("reading resul file: ",inverse_metric_filename) - df = pd.read_pickle(inverse_metric_filename) - df['model'] = model_name - df['pde'] = shortfilename - dfs+=[df] - keys = ['pde','model'] - df = pd.concat(dfs,axis=0) - return df, keys - -@hydra.main(config_path='../config', config_name='results') + inverse_metric_filename = get_metric_name( + filename, model_name, base_path, inverse_model_type + ) + if verbose: + msg = f"reading result file: {inverse_metric_filename}" + logger.info(msg) + + dframe = pd.read_pickle(inverse_metric_filename) + dframe["model"] = model_name + dframe["pde"] = shortfilename + dfs += [dframe] + keys = ["pde", "model"] + dframe = pd.concat(dfs, axis=0) + return dframe, keys + + +@hydra.main(config_path="../config", config_name="results") def process_results(cfg: DictConfig): """ - reads and merges the result files and aggregate the results with the selected values. The results are aggregated by datafile. + reads and merges the result files and aggregate the results with the selected values. The results are aggregated by datafile. June 2022, F.Alesiani - """ - print(cfg.args) - - df, keys = read_results(cfg.args.model_names,cfg.args.inverse_model_type, cfg.args.base_path, cfg.args.filenames, cfg.args.shortfilenames) + """ + logger.info(cfg.args) + + df, keys = read_results( + cfg.args.model_names, + cfg.args.inverse_model_type, + cfg.args.base_path, + cfg.args.filenames, + cfg.args.shortfilenames, + ) df1p3 = df[keys + list(cfg.args.results_values)] - df2p3 = df1p3.groupby(by=keys).agg([np.mean,np.std]).reset_index() - print("saving results into: ", cfg.args.base_path + cfg.args.result_filename) - df2p3.to_csv(cfg.args.base_path + cfg.args.result_filename) + df2p3 = df1p3.groupby(by=keys).agg([np.mean, np.std]).reset_index() + msg = "saving results into: {cfg.args.base_path + cfg.args.result_filename}" + logger.info(msg) + df2p3.to_csv(cfg.args.base_path + cfg.args.result_filename) if __name__ == "__main__": process_results() - print("Done.") \ No newline at end of file + msg = "Done." + logger.info(msg) diff --git a/pdebench/models/metrics.py b/pdebench/models/metrics.py index 2b19051..e33881e 100644 --- a/pdebench/models/metrics.py +++ b/pdebench/models/metrics.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- """ @@ -146,42 +145,67 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import torch -import numpy as np +from __future__ import annotations + +import logging import math as mt + import matplotlib.pyplot as plt +import numpy as np +import torch from mpl_toolkits.axes_grid1 import make_axes_locatable +from torch import nn + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +logger = logging.getLogger(__name__) -def metric_func(pred, target, if_mean=True, Lx=1., Ly=1., Lz=1., iLow=4, iHigh=12): + +def metric_func( + pred, target, if_mean=True, Lx=1.0, Ly=1.0, Lz=1.0, iLow=4, iHigh=12, initial_step=1 +): """ code for calculate metrics discussed in the Brain-storming session - MSE, normalized MSE, max error, MSE at the boundaries, conserved variables, MSE in Fourier space, temporal sensitivity + RMSE, normalized RMSE, max error, RMSE at the boundaries, conserved variables, RMSE in Fourier space, temporal sensitivity """ + pred, target = pred.to(device), target.to(device) # (batch, nx^i..., timesteps, nc) + # slice out `initial context` timesteps + pred = pred[..., initial_step:, :] + target = target[..., initial_step:, :] idxs = target.size() - if len(idxs) == 4: + if len(idxs) == 4: # 1D pred = pred.permute(0, 3, 1, 2) target = target.permute(0, 3, 1, 2) - if len(idxs) == 5: + if len(idxs) == 5: # 2D pred = pred.permute(0, 4, 1, 2, 3) target = target.permute(0, 4, 1, 2, 3) - elif len(idxs) == 6: + elif len(idxs) == 6: # 3D pred = pred.permute(0, 5, 1, 2, 3, 4) target = target.permute(0, 5, 1, 2, 3, 4) idxs = target.size() nb, nc, nt = idxs[0], idxs[1], idxs[-1] - # MSE - err_mean = torch.sqrt(torch.mean((pred.view([nb, nc, -1, nt]) - target.view([nb, nc, -1, nt])) ** 2, dim=2)) - err_MSE = torch.mean(err_mean, axis=0) + # RMSE + err_mean = torch.sqrt( + torch.mean( + (pred.view([nb, nc, -1, nt]) - target.view([nb, nc, -1, nt])) ** 2, dim=2 + ) + ) + err_RMSE = torch.mean(err_mean, axis=0) nrm = torch.sqrt(torch.mean(target.view([nb, nc, -1, nt]) ** 2, dim=2)) - err_nMSE = torch.mean(err_mean / nrm, dim=0) - - err_CSV = torch.sqrt(torch.mean( - (torch.sum(pred.view([nb, nc, -1, nt]), dim=2) - torch.sum(target.view([nb, nc, -1, nt]), dim=2)) ** 2, - dim=0)) + err_nRMSE = torch.mean(err_mean / nrm, dim=0) + + err_CSV = torch.sqrt( + torch.mean( + ( + torch.sum(pred.view([nb, nc, -1, nt]), dim=2) + - torch.sum(target.view([nb, nc, -1, nt]), dim=2) + ) + ** 2, + dim=0, + ) + ) if len(idxs) == 4: nx = idxs[2] err_CSV /= nx @@ -192,20 +216,27 @@ def metric_func(pred, target, if_mean=True, Lx=1., Ly=1., Lz=1., iLow=4, iHigh=1 nx, ny, nz = idxs[2:5] err_CSV /= nx * ny * nz # worst case in all the data - err_Max = torch.max(torch.max( - torch.abs(pred.view([nb, nc, -1, nt]) - target.view([nb, nc, -1, nt])), dim=2)[0], dim=0)[0] + err_Max = torch.max( + torch.max( + torch.abs(pred.view([nb, nc, -1, nt]) - target.view([nb, nc, -1, nt])), + dim=2, + )[0], + dim=0, + )[0] if len(idxs) == 4: # 1D err_BD = (pred[:, :, 0, :] - target[:, :, 0, :]) ** 2 err_BD += (pred[:, :, -1, :] - target[:, :, -1, :]) ** 2 - err_BD = torch.mean(torch.sqrt(err_BD / 2.), dim=0) + err_BD = torch.mean(torch.sqrt(err_BD / 2.0), dim=0) elif len(idxs) == 5: # 2D nx, ny = idxs[2:4] err_BD_x = (pred[:, :, 0, :, :] - target[:, :, 0, :, :]) ** 2 err_BD_x += (pred[:, :, -1, :, :] - target[:, :, -1, :, :]) ** 2 err_BD_y = (pred[:, :, :, 0, :] - target[:, :, :, 0, :]) ** 2 err_BD_y += (pred[:, :, :, -1, :] - target[:, :, :, -1, :]) ** 2 - err_BD = (torch.sum(err_BD_x, dim=-2) + torch.sum(err_BD_y, dim=-2)) / (2 * nx + 2 * ny) + err_BD = (torch.sum(err_BD_x, dim=-2) + torch.sum(err_BD_y, dim=-2)) / ( + 2 * nx + 2 * ny + ) err_BD = torch.mean(torch.sqrt(err_BD), dim=0) elif len(idxs) == 6: # 3D nx, ny, nz = idxs[2:5] @@ -215,17 +246,21 @@ def metric_func(pred, target, if_mean=True, Lx=1., Ly=1., Lz=1., iLow=4, iHigh=1 err_BD_y += (pred[:, :, :, -1, :] - target[:, :, :, -1, :]) ** 2 err_BD_z = (pred[:, :, :, :, 0] - target[:, :, :, :, 0]) ** 2 err_BD_z += (pred[:, :, :, :, -1] - target[:, :, :, :, -1]) ** 2 - err_BD = torch.sum(err_BD_x.view([nb, -1, nt]), dim=-2) \ - + torch.sum(err_BD_y.view([nb, -1, nt]), dim=-2) \ - + torch.sum(err_BD_z.view([nb, -1, nt]), dim=-2) + err_BD = ( + torch.sum(err_BD_x.contiguous().view([nb, -1, nt]), dim=-2) + + torch.sum(err_BD_y.contiguous().view([nb, -1, nt]), dim=-2) + + torch.sum(err_BD_z.contiguous().view([nb, -1, nt]), dim=-2) + ) err_BD = err_BD / (2 * nx * ny + 2 * ny * nz + 2 * nz * nx) - err_BD = torch.mean(torch.sqrt(err_BD), dim=0) + err_BD = torch.sqrt(err_BD) if len(idxs) == 4: # 1D nx = idxs[2] pred_F = torch.fft.rfft(pred, dim=2) target_F = torch.fft.rfft(target, dim=2) - _err_F = torch.sqrt(torch.mean(torch.abs(pred_F - target_F) ** 2, axis=0)) / nx * Lx + _err_F = ( + torch.sqrt(torch.mean(torch.abs(pred_F - target_F) ** 2, axis=0)) / nx * Lx + ) if len(idxs) == 5: # 2D pred_F = torch.fft.fftn(pred, dim=[2, 3]) target_F = torch.fft.fftn(target, dim=[2, 3]) @@ -234,7 +269,7 @@ def metric_func(pred, target, if_mean=True, Lx=1., Ly=1., Lz=1., iLow=4, iHigh=1 err_F = torch.zeros([nb, nc, min(nx // 2, ny // 2), nt]).to(device) for i in range(nx // 2): for j in range(ny // 2): - it = mt.floor(mt.sqrt(i ** 2 + j ** 2)) + it = mt.floor(mt.sqrt(i**2 + j**2)) if it > min(nx // 2, ny // 2) - 1: continue err_F[:, :, it] += _err_F[:, :, i, j] @@ -248,442 +283,533 @@ def metric_func(pred, target, if_mean=True, Lx=1., Ly=1., Lz=1., iLow=4, iHigh=1 for i in range(nx // 2): for j in range(ny // 2): for k in range(nz // 2): - it = mt.floor(mt.sqrt(i ** 2 + j ** 2 + k ** 2)) + it = mt.floor(mt.sqrt(i**2 + j**2 + k**2)) if it > min(nx // 2, ny // 2, nz // 2) - 1: continue err_F[:, :, it] += _err_F[:, :, i, j, k] _err_F = torch.sqrt(torch.mean(err_F, axis=0)) / (nx * ny * nz) * Lx * Ly * Lz err_F = torch.zeros([nc, 3, nt]).to(device) - err_F[:,0] += torch.mean(_err_F[:,:iLow], dim=1) # low freq - err_F[:,1] += torch.mean(_err_F[:,iLow:iHigh], dim=1) # middle freq - err_F[:,2] += torch.mean(_err_F[:,iHigh:], dim=1) # high freq + err_F[:, 0] += torch.mean(_err_F[:, :iLow], dim=1) # low freq + err_F[:, 1] += torch.mean(_err_F[:, iLow:iHigh], dim=1) # middle freq + err_F[:, 2] += torch.mean(_err_F[:, iHigh:], dim=1) # high freq if if_mean: - return torch.mean(err_MSE, dim=[0, -1]), \ - torch.mean(err_nMSE, dim=[0, -1]), \ - torch.mean(err_CSV, dim=[0, -1]), \ - torch.mean(err_Max, dim=[0, -1]), \ - torch.mean(err_BD, dim=[0, -1]), \ - torch.mean(err_F, dim=[0, -1]) - else: - return err_MSE, err_nMSE, err_CSV, err_Max, err_BD, err_F - -def metrics(val_loader, model, Lx, Ly, Lz, plot, channel_plot, model_name, x_min, - x_max, y_min, y_max, t_min, t_max, mode='FNO', initial_step=None, ): - if mode=='Unet': + return ( + torch.mean(err_RMSE, dim=[0, -1]), + torch.mean(err_nRMSE, dim=[0, -1]), + torch.mean(err_CSV, dim=[0, -1]), + torch.mean(err_Max, dim=[0, -1]), + torch.mean(err_BD, dim=[0, -1]), + torch.mean(err_F, dim=[0, -1]), + ) + return err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F + + +def metrics( + val_loader, + model, + Lx, + Ly, + Lz, + plot, + channel_plot, + model_name, + x_min, + x_max, + y_min, + y_max, + t_min, + t_max, + mode="FNO", + initial_step=None, +): + if mode == "Unet": with torch.no_grad(): - itot = 0 - for xx, yy in val_loader: - xx = xx.to(device) - yy = yy.to(device) + for itot, (xx, yy) in enumerate(val_loader): + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 pred = yy[..., :initial_step, :] inp_shape = list(xx.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - for t in range(initial_step, yy.shape[-2]): + for _t in range(initial_step, yy.shape[-2]): inp = xx.reshape(inp_shape) temp_shape = [0, -1] - temp_shape.extend([i for i in range(1,len(inp.shape)-1)]) + temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) - - y = yy[..., t:t+1, :] - + temp_shape = [0] - temp_shape.extend([i for i in range(2,len(inp.shape))]) + temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) im = model(inp).permute(temp_shape).unsqueeze(-2) pred = torch.cat((pred, im), -2) - xx = torch.cat((xx[..., 1:, :], im), dim=-2) - - _err_MSE, _err_nMSE, _err_CSV, _err_Max, _err_BD, _err_F \ - = metric_func(pred, yy, if_mean=True, Lx=Lx, Ly=Ly, Lz=Lz) + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 + + ( + _err_RMSE, + _err_nRMSE, + _err_CSV, + _err_Max, + _err_BD, + _err_F, + ) = metric_func( + pred, + yy, + if_mean=True, + Lx=Lx, + Ly=Ly, + Lz=Lz, + initial_step=initial_step, + ) if itot == 0: - err_MSE, err_nMSE, err_CSV, err_Max, err_BD, err_F \ - = _err_MSE, _err_nMSE, _err_CSV, _err_Max, _err_BD, _err_F + err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F = ( + _err_RMSE, + _err_nRMSE, + _err_CSV, + _err_Max, + _err_BD, + _err_F, + ) pred_plot = pred[:1] target_plot = yy[:1] val_l2_time = torch.zeros(yy.shape[-2]).to(device) else: - err_MSE += _err_MSE - err_nMSE += _err_nMSE + err_RMSE += _err_RMSE + err_nRMSE += _err_nRMSE err_CSV += _err_CSV err_Max += _err_Max err_BD += _err_BD err_F += _err_F - - mean_dim = [i for i in range(len(yy.shape)-2)] + + mean_dim = list(range(len(yy.shape) - 2)) mean_dim.append(-1) mean_dim = tuple(mean_dim) - val_l2_time += torch.sqrt(torch.mean((pred-yy)**2, dim=mean_dim)) - - itot += 1 + val_l2_time += torch.sqrt( + torch.mean((pred - yy) ** 2, dim=mean_dim) + ) - elif mode=='FNO': + elif mode == "FNO": with torch.no_grad(): itot = 0 - for xx, yy, grid in val_loader: - xx = xx.to(device) - yy = yy.to(device) - grid = grid.to(device) + for itot, (xx, yy, grid) in enumerate(val_loader): + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 + grid = grid.to(device) # noqa: PLW2901 pred = yy[..., :initial_step, :] inp_shape = list(xx.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - for t in range(initial_step, yy.shape[-2]): + for _t in range(initial_step, yy.shape[-2]): inp = xx.reshape(inp_shape) - y = yy[..., t:t + 1, :] im = model(inp, grid) pred = torch.cat((pred, im), -2) - xx = torch.cat((xx[..., 1:, :], im), dim=-2) - - _err_MSE, _err_nMSE, _err_CSV, _err_Max, _err_BD, _err_F \ - = metric_func(pred, yy, if_mean=True, Lx=Lx, Ly=Ly, Lz=Lz) + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 + + ( + _err_RMSE, + _err_nRMSE, + _err_CSV, + _err_Max, + _err_BD, + _err_F, + ) = metric_func( + pred, + yy, + if_mean=True, + Lx=Lx, + Ly=Ly, + Lz=Lz, + initial_step=initial_step, + ) if itot == 0: - err_MSE, err_nMSE, err_CSV, err_Max, err_BD, err_F \ - = _err_MSE, _err_nMSE, _err_CSV, _err_Max, _err_BD, _err_F + err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F = ( + _err_RMSE, + _err_nRMSE, + _err_CSV, + _err_Max, + _err_BD, + _err_F, + ) pred_plot = pred[:1] target_plot = yy[:1] val_l2_time = torch.zeros(yy.shape[-2]).to(device) else: - err_MSE += _err_MSE - err_nMSE += _err_nMSE + err_RMSE += _err_RMSE + err_nRMSE += _err_nRMSE err_CSV += _err_CSV err_Max += _err_Max err_BD += _err_BD err_F += _err_F - - mean_dim = [i for i in range(len(yy.shape)-2)] + + mean_dim = list(range(len(yy.shape) - 2)) mean_dim.append(-1) mean_dim = tuple(mean_dim) - val_l2_time += torch.sqrt(torch.mean((pred-yy)**2, dim=mean_dim)) - - itot += 1 + val_l2_time += torch.sqrt( + torch.mean((pred - yy) ** 2, dim=mean_dim) + ) elif mode == "PINN": raise NotImplementedError + err_RMSE = np.array(err_RMSE.data.cpu() / itot) + err_nRMSE = np.array(err_nRMSE.data.cpu() / itot) + err_CSV = np.array(err_CSV.data.cpu() / itot) + err_Max = np.array(err_Max.data.cpu() / itot) + err_BD = np.array(err_BD.data.cpu() / itot) + err_F = np.array(err_F.data.cpu() / itot) + logger.info(f"RMSE: {err_RMSE:.5f}") + logger.info(f"normalized RMSE: {err_nRMSE:.5f}") + logger.info(f"RMSE of conserved variables: {err_CSV:.5f}") + logger.info(f"Maximum value of rms error: {err_Max:.5f}") + logger.info(f"RMSE at boundaries: {err_BD:.5f}") + logger.info(f"RMSE in Fourier space: {err_F}") + + val_l2_time = val_l2_time / itot - err_MSE = np.array(err_MSE.data.cpu()/itot) - err_nMSE = np.array(err_nMSE.data.cpu()/itot) - err_CSV = np.array(err_CSV.data.cpu()/itot) - err_Max = np.array(err_Max.data.cpu()/itot) - err_BD = np.array(err_BD.data.cpu()/itot) - err_F = np.array(err_F.data.cpu()/itot) - print('MSE: {0:.5f}'.format(err_MSE)) - print('normalized MSE: {0:.5f}'.format(err_nMSE)) - print('MSE of conserved variables: {0:.5f}'.format(err_CSV)) - print('Maximum value of rms error: {0:.5f}'.format(err_Max)) - print('MSE at boundaries: {0:.5f}'.format(err_BD)) - print('MSE in Fourier space: {0}'.format(err_F)) - - val_l2_time = val_l2_time/itot - if plot: dim = len(yy.shape) - 3 plt.ioff() if dim == 1: - - fig, ax = plt.subplots(figsize=(6.5,6)) - h = ax.imshow(pred_plot[...,channel_plot].squeeze().detach().cpu(), - extent=[t_min, t_max, x_min, x_max], origin='lower', aspect='auto') - h.set_clim(target_plot[...,channel_plot].min(), target_plot[...,channel_plot].max()) + fig, ax = plt.subplots(figsize=(6.5, 6)) + h = ax.imshow( + pred_plot[..., channel_plot].squeeze().detach().cpu(), + extent=[t_min, t_max, x_min, x_max], + origin="lower", + aspect="auto", + ) + h.set_clim( + target_plot[..., channel_plot].min(), + target_plot[..., channel_plot].max(), + ) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) cbar = fig.colorbar(h, cax=cax) - ax.set_title("Prediction", fontsize=20) - ax.set_ylabel("$x$", fontsize=20) - ax.set_xlabel("$t$", fontsize=20) + cbar.ax.tick_params(labelsize=30) + ax.set_title("Prediction", fontsize=30) + ax.tick_params(axis="x", labelsize=30) + ax.tick_params(axis="y", labelsize=30) + ax.set_ylabel("$x$", fontsize=30) + ax.set_xlabel("$t$", fontsize=30) plt.tight_layout() - filename = model_name + '_pred.pdf' + filename = model_name + "_pred.pdf" plt.savefig(filename) - - fig, ax = plt.subplots(figsize=(6.5,6)) - h = ax.imshow(target_plot[...,channel_plot].squeeze().detach().cpu(), - extent=[t_min, t_max, x_min, x_max], origin='lower', aspect='auto') - h.set_clim(target_plot[...,channel_plot].min(), target_plot[...,channel_plot].max()) + + fig, ax = plt.subplots(figsize=(6.5, 6)) + h = ax.imshow( + target_plot[..., channel_plot].squeeze().detach().cpu(), + extent=[t_min, t_max, x_min, x_max], + origin="lower", + aspect="auto", + ) + h.set_clim( + target_plot[..., channel_plot].min(), + target_plot[..., channel_plot].max(), + ) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) cbar = fig.colorbar(h, cax=cax) - ax.set_title("Data", fontsize=20) - ax.set_ylabel("$x$", fontsize=20) - ax.set_xlabel("$t$", fontsize=20) + cbar.ax.tick_params(labelsize=30) + ax.set_title("Data", fontsize=30) + ax.tick_params(axis="x", labelsize=30) + ax.tick_params(axis="y", labelsize=30) + ax.set_ylabel("$x$", fontsize=30) + ax.set_xlabel("$t$", fontsize=30) plt.tight_layout() - filename = model_name + '_data.pdf' + filename = model_name + "_data.pdf" plt.savefig(filename) - + elif dim == 2: - - fig, ax = plt.subplots(figsize=(6.5,6)) - h = ax.imshow(pred_plot[...,-1,channel_plot].squeeze().t().detach().cpu(), - extent=[x_min, x_max, y_min, y_max], origin='lower', aspect='auto') - h.set_clim(target_plot[...,-1,channel_plot].min(), target_plot[...,-1,channel_plot].max()) + fig, ax = plt.subplots(figsize=(6.5, 6)) + h = ax.imshow( + pred_plot[..., -1, channel_plot].squeeze().t().detach().cpu(), + extent=[x_min, x_max, y_min, y_max], + origin="lower", + aspect="auto", + ) + h.set_clim( + target_plot[..., -1, channel_plot].min(), + target_plot[..., -1, channel_plot].max(), + ) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) cbar = fig.colorbar(h, cax=cax) - ax.set_title("Prediction", fontsize=20) - ax.set_ylabel("$y$", fontsize=20) - ax.set_xlabel("$x$", fontsize=20) + cbar.ax.tick_params(labelsize=30) + ax.set_title("Prediction", fontsize=30) + ax.tick_params(axis="x", labelsize=30) + ax.tick_params(axis="y", labelsize=30) + ax.set_ylabel("$y$", fontsize=30) + ax.set_xlabel("$x$", fontsize=30) plt.tight_layout() - filename = model_name + '_pred.pdf' + filename = model_name + "_pred.pdf" plt.savefig(filename) - - fig, ax = plt.subplots(figsize=(6.5,6)) - h = ax.imshow(target_plot[...,-1,channel_plot].squeeze().t().detach().cpu(), - extent=[x_min, x_max, y_min, y_max], origin='lower', aspect='auto') - h.set_clim(target_plot[...,-1,channel_plot].min(), target_plot[...,-1,channel_plot].max()) + + fig, ax = plt.subplots(figsize=(6.5, 6)) + h = ax.imshow( + target_plot[..., -1, channel_plot].squeeze().t().detach().cpu(), + extent=[x_min, x_max, y_min, y_max], + origin="lower", + aspect="auto", + ) + h.set_clim( + target_plot[..., -1, channel_plot].min(), + target_plot[..., -1, channel_plot].max(), + ) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) cbar = fig.colorbar(h, cax=cax) - ax.set_title("Data", fontsize=20) - ax.set_ylabel("$y$", fontsize=20) - ax.set_xlabel("$x$", fontsize=20) + cbar.ax.tick_params(labelsize=30) + ax.set_title("Data", fontsize=30) + ax.tick_params(axis="x", labelsize=30) + ax.tick_params(axis="y", labelsize=30) + ax.set_ylabel("$y$", fontsize=30) + ax.set_xlabel("$x$", fontsize=30) plt.tight_layout() - filename = model_name + '_data.pdf' + filename = model_name + "_data.pdf" plt.savefig(filename) - - # plt.figure(figsize=(5,5)) + + # plt.figure(figsize=(8,8)) # plt.semilogy(torch.arange(initial_step,yy.shape[-2]), # val_l2_time[initial_step:].detach().cpu()) - # plt.xlabel('$t$', fontsize=20) - # plt.ylabel('$MSE$', fontsize=20) - # plt.title('MSE vs unrolled time steps', fontsize=20) + # plt.xlabel('$t$', fontsize=30) + # plt.ylabel('$MSE$', fontsize=30) + # plt.title('MSE vs unrolled time steps', fontsize=30) # plt.tight_layout() # filename = model_name + '_mse_time.pdf' # plt.savefig(filename) - - filename = model_name + 'mse_time.npz' - np.savez(filename, t=torch.arange(initial_step,yy.shape[-2]), - mse=val_l2_time[initial_step:].detach().cpu()) - return err_MSE, err_nMSE, err_CSV, err_Max, err_BD, err_F + filename = model_name + "mse_time.npz" + np.savez( + filename, + t=torch.arange(initial_step, yy.shape[-2]).cpu(), + mse=val_l2_time[initial_step:].detach().cpu(), + ) + + return err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F -# LpLoss Function -class LpLoss(object): +# LpLoss Function +class LpLoss: """ - Lp loss function + Lp loss function """ - def __init__(self, p=2, reduction='mean'): - super(LpLoss, self).__init__() - #Dimension and Lp-norm type are postive + + def __init__(self, p=2, reduction="mean"): + super().__init__() + # Dimension and Lp-norm type are positive assert p > 0 self.p = p self.reduction = reduction + def __call__(self, x, y, eps=1e-20): num_examples = x.size()[0] - _diff = x.view(num_examples,-1) - y.view(num_examples,-1) + _diff = x.view(num_examples, -1) - y.view(num_examples, -1) _diff = torch.norm(_diff, self.p, 1) - _norm = eps + torch.norm(y.view(num_examples,-1), self.p, 1) - if self.reduction in ['mean']: - return torch.mean(_diff/_norm) - if self.reduction in ['sum']: - return torch.sum(_diff/_norm) - return _diff/_norm - -# FftLoss Function -class FftLpLoss(object): + _norm = eps + torch.norm(y.view(num_examples, -1), self.p, 1) + if self.reduction in ["mean"]: + return torch.mean(_diff / _norm) + if self.reduction in ["sum"]: + return torch.sum(_diff / _norm) + return _diff / _norm + + +# FftLoss Function +class FftLpLoss: """ loss function in Fourier space June 2022, F.Alesiani """ - def __init__(self, p=2, reduction='mean'): - super(FftLpLoss, self).__init__() - #Dimension and Lp-norm type are postive + + def __init__(self, p=2, reduction="mean"): + super().__init__() + # Dimension and Lp-norm type are positive assert p > 0 self.p = p self.reduction = reduction - def __call__(self, x, y, flow=None,fhigh=None, eps=1e-20): + + def __call__(self, x, y, flow=None, fhigh=None, eps=1e-20): num_examples = x.size()[0] others_dims = x.shape[1:] - dims = list(range(1,len(x.shape))) - xf = torch.fft.fftn(x,dim=dims) - yf = torch.fft.fftn(y,dim=dims) - if flow is None: flow = 0 - if fhigh is None: fhigh = np.max(xf.shape[1:]) - - if len(others_dims) ==1: - xf = xf[:,flow:fhigh] - yf = yf[:,flow:fhigh] - if len(others_dims) ==2: - xf = xf[:,flow:fhigh,flow:fhigh] - yf = yf[:,flow:fhigh,flow:fhigh] - if len(others_dims) ==3: - xf = xf[:,flow:fhigh,flow:fhigh,flow:fhigh] - yf = yf[:,flow:fhigh,flow:fhigh,flow:fhigh] - if len(others_dims) ==4: - xf = xf[:,flow:fhigh,flow:fhigh,flow:fhigh,flow:fhigh] - yf = yf[:,flow:fhigh,flow:fhigh,flow:fhigh,flow:fhigh] + dims = list(range(1, len(x.shape))) + xf = torch.fft.fftn(x, dim=dims) + yf = torch.fft.fftn(y, dim=dims) + if flow is None: + flow = 0 + if fhigh is None: + fhigh = np.max(xf.shape[1:]) + + if len(others_dims) == 1: + xf = xf[:, flow:fhigh] + yf = yf[:, flow:fhigh] + if len(others_dims) == 2: + xf = xf[:, flow:fhigh, flow:fhigh] + yf = yf[:, flow:fhigh, flow:fhigh] + if len(others_dims) == 3: + xf = xf[:, flow:fhigh, flow:fhigh, flow:fhigh] + yf = yf[:, flow:fhigh, flow:fhigh, flow:fhigh] + if len(others_dims) == 4: + xf = xf[:, flow:fhigh, flow:fhigh, flow:fhigh, flow:fhigh] + yf = yf[:, flow:fhigh, flow:fhigh, flow:fhigh, flow:fhigh] _diff = xf - yf.reshape(xf.shape) - _diff = torch.norm(_diff.reshape(num_examples,-1), self.p, 1) - _norm = eps + torch.norm(yf.reshape(num_examples,-1), self.p, 1) - - if self.reduction in ['mean']: - return torch.mean(_diff/_norm) - if self.reduction in ['sum']: - return torch.sum(_diff/_norm) - return _diff/_norm - -import torch.nn.functional as F -# FftLoss Function -class FftMseLoss(object): + _diff = torch.norm(_diff.reshape(num_examples, -1), self.p, 1) + _norm = eps + torch.norm(yf.reshape(num_examples, -1), self.p, 1) + + if self.reduction in ["mean"]: + return torch.mean(_diff / _norm) + if self.reduction in ["sum"]: + return torch.sum(_diff / _norm) + return _diff / _norm + + +# FftLoss Function +class FftMseLoss: """ loss function in Fourier space June 2022, F.Alesiani """ - def __init__(self, reduction='mean'): - super(FftMseLoss, self).__init__() - #Dimension and Lp-norm type are postive + + def __init__(self, reduction="mean"): + super().__init__() + # Dimension and Lp-norm type are positive self.reduction = reduction - def __call__(self, x, y, flow=None,fhigh=None, eps=1e-20): + + def __call__(self, x, y, flow=None, fhigh=None): num_examples = x.size()[0] others_dims = x.shape[1:-2] for d in others_dims: - assert (d>1), "we expect the dimension to be the same and greater the 1" + assert d > 1, "we expect the dimension to be the same and greater the 1" # print(others_dims) - dims = list(range(1,len(x.shape)-1)) - xf = torch.fft.fftn(x,dim=dims) - yf = torch.fft.fftn(y,dim=dims) - if flow is None: flow = 0 - if fhigh is None: fhigh = np.max(xf.shape[1:]) - - if len(others_dims) ==1: - xf = xf[:,flow:fhigh] - yf = yf[:,flow:fhigh] - if len(others_dims) ==2: - xf = xf[:,flow:fhigh,flow:fhigh] - yf = yf[:,flow:fhigh,flow:fhigh] - if len(others_dims) ==3: - xf = xf[:,flow:fhigh,flow:fhigh,flow:fhigh] - yf = yf[:,flow:fhigh,flow:fhigh,flow:fhigh] - if len(others_dims) ==4: - xf = xf[:,flow:fhigh,flow:fhigh,flow:fhigh,flow:fhigh] - yf = yf[:,flow:fhigh,flow:fhigh,flow:fhigh,flow:fhigh] + dims = list(range(1, len(x.shape) - 1)) + xf = torch.fft.fftn(x, dim=dims) + yf = torch.fft.fftn(y, dim=dims) + if flow is None: + flow = 0 + if fhigh is None: + fhigh = np.max(xf.shape[1:]) + + if len(others_dims) == 1: + xf = xf[:, flow:fhigh] + yf = yf[:, flow:fhigh] + if len(others_dims) == 2: + xf = xf[:, flow:fhigh, flow:fhigh] + yf = yf[:, flow:fhigh, flow:fhigh] + if len(others_dims) == 3: + xf = xf[:, flow:fhigh, flow:fhigh, flow:fhigh] + yf = yf[:, flow:fhigh, flow:fhigh, flow:fhigh] + if len(others_dims) == 4: + xf = xf[:, flow:fhigh, flow:fhigh, flow:fhigh, flow:fhigh] + yf = yf[:, flow:fhigh, flow:fhigh, flow:fhigh, flow:fhigh] _diff = xf - yf - _diff = _diff.reshape(num_examples,-1).abs()**2 - if self.reduction in ['mean']: + _diff = _diff.reshape(num_examples, -1).abs() ** 2 + if self.reduction in ["mean"]: return torch.mean(_diff).abs() - if self.reduction in ['sum']: + if self.reduction in ["sum"]: return torch.sum(_diff).abs() return _diff.abs() -import torch.nn as nn -def inverse_metrics(u0,x,pred_u0,y): +def inverse_metrics(u0, x, pred_u0, y): """ computes all the metrics in the base and fourier space - u0: esimated initial condition, - pred_u0: prediction from the esimated initial condition, pred_u0 = model(u0) + u0: estimated initial condition, + pred_u0: prediction from the estimated initial condition, pred_u0 = model(u0) x: true initial condition y: true prediction, y = model(x) June 2022, F.Alesiani """ - + mseloss_fn = nn.MSELoss(reduction="mean") - l2loss_fn = LpLoss(p=2,reduction='mean') - l3loss_fn = LpLoss(p=3,reduction='mean') - + l2loss_fn = LpLoss(p=2, reduction="mean") + l3loss_fn = LpLoss(p=3, reduction="mean") + fftmseloss_fn = FftMseLoss(reduction="mean") - fftl2loss_fn = FftLpLoss(p=2,reduction="mean") - fftl3loss_fn = FftLpLoss(p=3,reduction="mean") + fftl2loss_fn = FftLpLoss(p=2, reduction="mean") + fftl3loss_fn = FftLpLoss(p=3, reduction="mean") - #initial condition + # initial condition mseloss_u0 = mseloss_fn(u0.view(1, -1), x.view(1, -1)).item() l2loss_u0 = l2loss_fn(u0.view(1, -1), x.view(1, -1)).item() l3loss_u0 = l3loss_fn(u0.view(1, -1), x.view(1, -1)).item() - - - fmid = u0.shape[1]//4 + + fmid = u0.shape[1] // 4 fftmseloss_u0 = fftmseloss_fn(u0, x).item() fftmseloss_low_u0 = fftmseloss_fn(u0, x, 0, fmid).item() - fftmseloss_mid_u0 = fftmseloss_fn(u0, x, fmid, 2*fmid).item() - fftmseloss_hi_u0 = fftmseloss_fn(u0, x, 2*fmid).item() - + fftmseloss_mid_u0 = fftmseloss_fn(u0, x, fmid, 2 * fmid).item() + fftmseloss_hi_u0 = fftmseloss_fn(u0, x, 2 * fmid).item() + fftl2loss_u0 = fftl2loss_fn(u0, x).item() fftl2loss_low_u0 = fftl2loss_fn(u0, x, 0, fmid).item() - fftl2loss_mid_u0 = fftl2loss_fn(u0, x, fmid, 2*fmid).item() - fftl2loss_hi_u0 = fftl2loss_fn(u0, x, 2*fmid).item() + fftl2loss_mid_u0 = fftl2loss_fn(u0, x, fmid, 2 * fmid).item() + fftl2loss_hi_u0 = fftl2loss_fn(u0, x, 2 * fmid).item() fftl3loss_u0 = fftl3loss_fn(u0, x).item() fftl3loss_low_u0 = fftl3loss_fn(u0, x, 0, fmid).item() - fftl3loss_mid_u0 = fftl3loss_fn(u0, x, fmid, 2*fmid).item() - fftl3loss_hi_u0 = fftl3loss_fn(u0, x, 2*fmid).item() + fftl3loss_mid_u0 = fftl3loss_fn(u0, x, fmid, 2 * fmid).item() + fftl3loss_hi_u0 = fftl3loss_fn(u0, x, 2 * fmid).item() - #prediction + # prediction mseloss_pred_u0 = mseloss_fn(pred_u0.reshape(1, -1), y.reshape(1, -1)).item() l2loss_pred_u0 = l2loss_fn(pred_u0.reshape(1, -1), y.reshape(1, -1)).item() l3loss_pred_u0 = l3loss_fn(pred_u0.reshape(1, -1), y.reshape(1, -1)).item() - fmid = pred_u0.shape[1]//4 + fmid = pred_u0.shape[1] // 4 pred_u0 = pred_u0.squeeze(-1) y = y.squeeze(-1) - + fftmseloss_pred_u0 = fftmseloss_fn(pred_u0, y).item() fftmseloss_low_pred_u0 = fftmseloss_fn(pred_u0, y, 0, fmid).item() - fftmseloss_mid_pred_u0 = fftmseloss_fn(pred_u0, y, fmid, 2*fmid).item() - fftmseloss_hi_pred_u0 = fftmseloss_fn(pred_u0, y, 2*fmid).item() - + fftmseloss_mid_pred_u0 = fftmseloss_fn(pred_u0, y, fmid, 2 * fmid).item() + fftmseloss_hi_pred_u0 = fftmseloss_fn(pred_u0, y, 2 * fmid).item() + fftl2loss_pred_u0 = fftl2loss_fn(pred_u0, y).item() fftl2loss_low_pred_u0 = fftl2loss_fn(pred_u0, y, 0, fmid).item() - fftl2loss_mid_pred_u0 = fftl2loss_fn(pred_u0, y, fmid, 2*fmid).item() - fftl2loss_hi_pred_u0= fftl2loss_fn(pred_u0, y, 2*fmid).item() + fftl2loss_mid_pred_u0 = fftl2loss_fn(pred_u0, y, fmid, 2 * fmid).item() + fftl2loss_hi_pred_u0 = fftl2loss_fn(pred_u0, y, 2 * fmid).item() fftl3loss_pred_u0 = fftl3loss_fn(pred_u0, y).item() fftl3loss_low_pred_u0 = fftl3loss_fn(pred_u0, y, 0, fmid).item() - fftl3loss_mid_pred_u0 = fftl3loss_fn(pred_u0, y, fmid, 2*fmid).item() - fftl3loss_hi_pred_u0 = fftl3loss_fn(pred_u0, y, 2*fmid).item() - - metric = { - 'mseloss_u0': mseloss_u0 - ,'l2loss_u0': l2loss_u0 - ,'l3loss_u0': l3loss_u0 - - ,'mseloss_pred_u0': mseloss_pred_u0 - ,'l2loss_pred_u0': l2loss_pred_u0 - ,'l3loss_pred_u0': l3loss_pred_u0 - - - ,'fftmseloss_u0': fftmseloss_u0 - ,'fftmseloss_low_u0': fftmseloss_low_u0 - ,'fftmseloss_mid_u0': fftmseloss_mid_u0 - ,'fftmseloss_hi_u0': fftmseloss_hi_u0 - - - - ,'fftmseloss_pred_u0': fftmseloss_pred_u0 - ,'fftmseloss_low_pred_u0': fftmseloss_low_pred_u0 - ,'fftmseloss_mid_pred_u0': fftmseloss_mid_pred_u0 - ,'fftmseloss_hi_pred_u0': fftmseloss_hi_pred_u0 - - ,'fftl2loss_u0': fftl2loss_u0 - ,'fftl2loss_low_u0': fftl2loss_low_u0 - ,'fftl2loss_mid_u0': fftl2loss_mid_u0 - ,'fftl2loss_hi_u0': fftl2loss_hi_u0 - - ,'fftl2loss_pred_u0': fftl2loss_pred_u0 - ,'fftl2loss_low_pred_u0': fftl2loss_low_pred_u0 - ,'fftl2loss_mid_pred_u0': fftl2loss_mid_pred_u0 - ,'fftl2loss_hi_pred_u0': fftl2loss_hi_pred_u0 - - ,'fftl3loss_u0': fftl3loss_u0 - ,'fftl3loss_low_u0': fftl3loss_low_u0 - ,'fftl3loss_mid_u0': fftl3loss_mid_u0 - ,'fftl3loss_hi_u0': fftl3loss_hi_u0 - - ,'fftl3loss_pred_u0': fftl3loss_pred_u0 - ,'fftl3loss_low_pred_u0': fftl3loss_low_pred_u0 - ,'fftl3loss_mid_pred_u0': fftl3loss_mid_pred_u0 - ,'fftl3loss_hi_pred_u0': fftl3loss_hi_pred_u0 - } - - return metric \ No newline at end of file + fftl3loss_mid_pred_u0 = fftl3loss_fn(pred_u0, y, fmid, 2 * fmid).item() + fftl3loss_hi_pred_u0 = fftl3loss_fn(pred_u0, y, 2 * fmid).item() + + return { + "mseloss_u0": mseloss_u0, + "l2loss_u0": l2loss_u0, + "l3loss_u0": l3loss_u0, + "mseloss_pred_u0": mseloss_pred_u0, + "l2loss_pred_u0": l2loss_pred_u0, + "l3loss_pred_u0": l3loss_pred_u0, + "fftmseloss_u0": fftmseloss_u0, + "fftmseloss_low_u0": fftmseloss_low_u0, + "fftmseloss_mid_u0": fftmseloss_mid_u0, + "fftmseloss_hi_u0": fftmseloss_hi_u0, + "fftmseloss_pred_u0": fftmseloss_pred_u0, + "fftmseloss_low_pred_u0": fftmseloss_low_pred_u0, + "fftmseloss_mid_pred_u0": fftmseloss_mid_pred_u0, + "fftmseloss_hi_pred_u0": fftmseloss_hi_pred_u0, + "fftl2loss_u0": fftl2loss_u0, + "fftl2loss_low_u0": fftl2loss_low_u0, + "fftl2loss_mid_u0": fftl2loss_mid_u0, + "fftl2loss_hi_u0": fftl2loss_hi_u0, + "fftl2loss_pred_u0": fftl2loss_pred_u0, + "fftl2loss_low_pred_u0": fftl2loss_low_pred_u0, + "fftl2loss_mid_pred_u0": fftl2loss_mid_pred_u0, + "fftl2loss_hi_pred_u0": fftl2loss_hi_pred_u0, + "fftl3loss_u0": fftl3loss_u0, + "fftl3loss_low_u0": fftl3loss_low_u0, + "fftl3loss_mid_u0": fftl3loss_mid_u0, + "fftl3loss_hi_u0": fftl3loss_hi_u0, + "fftl3loss_pred_u0": fftl3loss_pred_u0, + "fftl3loss_low_pred_u0": fftl3loss_low_pred_u0, + "fftl3loss_mid_pred_u0": fftl3loss_mid_pred_u0, + "fftl3loss_hi_pred_u0": fftl3loss_hi_pred_u0, + } diff --git a/pdebench/models/pinn/pde_definitions.py b/pdebench/models/pinn/pde_definitions.py index f4740d7..834fd37 100644 --- a/pdebench/models/pinn/pde_definitions.py +++ b/pdebench/models/pinn/pde_definitions.py @@ -1,6 +1,7 @@ +from __future__ import annotations + import deepxde as dde import numpy as np -import torch def reaction_1(u1, u2): @@ -14,7 +15,6 @@ def reaction_2(u1, u2): def pde_diffusion_reaction(x, y): - d1 = 1e-3 d2 = 5e-3 @@ -55,8 +55,8 @@ def pde_diffusion_sorption(x, y): ) return du1_t - D / retardation_factor * du1_xx - - + + def pde_swe1d(): raise NotImplementedError @@ -84,3 +84,132 @@ def pde_swe2d(x, y): eq3 = v_t + u * v_x + v * v_y + g * h_y return eq1 + eq2 + eq3 + + +def pde_adv1d(x, y, beta): + dy_x = dde.grad.jacobian(y, x, i=0, j=0) + dy_t = dde.grad.jacobian(y, x, i=0, j=1) + return dy_t + beta * dy_x + + +def pde_diffusion_reaction_1d(x, y, nu, rho): + dy_t = dde.grad.jacobian(y, x, i=0, j=1) + dy_xx = dde.grad.hessian(y, x, i=0, j=0) + return dy_t - nu * dy_xx - rho * y * (1.0 - y) + + +def pde_burgers1D(x, y, nu): + dy_x = dde.grad.jacobian(y, x, i=0, j=0) + dy_t = dde.grad.jacobian(y, x, i=0, j=1) + dy_xx = dde.grad.hessian(y, x, i=0, j=0) + return dy_t + y * dy_x - nu / np.pi * dy_xx + + +def pde_CFD1d(x, y, gamma): + h = y[..., 0].unsqueeze(1) # rho + u = y[..., 1].unsqueeze(1) # v + p = y[..., 2].unsqueeze(1) # p + E = p / (gamma - 1.0) + 0.5 * h * u**2 + E = E.unsqueeze(1) + Fx = u * (E + p) + Fx = Fx.unsqueeze(1) + + # non conservative form + hu_x = dde.grad.jacobian(h * u, x, i=0, j=0) + h_t = dde.grad.jacobian(y, x, i=0, j=1) + u_x = dde.grad.jacobian(y, x, i=1, j=0) + u_t = dde.grad.jacobian(y, x, i=1, j=1) + p_x = dde.grad.jacobian(y, x, i=2, j=0) + Fx_x = dde.grad.jacobian(Fx, x, i=0, j=0) + E_t = dde.grad.jacobian(E, x, i=0, j=1) + + eq1 = h_t + hu_x + eq2 = h * (u_t + u * u_x) - p_x + eq3 = E_t + Fx_x + + return eq1 + eq2 + eq3 + + +def pde_CFD2d(x, y, gamma): + h = y[..., 0].unsqueeze(1) # rho + ux = y[..., 1].unsqueeze(1) # vx + uy = y[..., 2].unsqueeze(1) # vy + p = y[..., 3].unsqueeze(1) # p + E = p / (gamma - 1.0) + 0.5 * h * (ux**2 + uy**2) + E = E.unsqueeze(1) + Fx = ux * (E + p) + Fx = Fx.unsqueeze(1) + Fy = uy * (E + p) + Fy = Fy.unsqueeze(1) + + # non conservative form + hu_x = dde.grad.jacobian(h * ux, x, i=0, j=0) + hu_y = dde.grad.jacobian(h * uy, x, i=0, j=1) + h_t = dde.grad.jacobian(y, x, i=0, j=2) + ux_x = dde.grad.jacobian(y, x, i=1, j=0) + ux_y = dde.grad.jacobian(y, x, i=1, j=1) + ux_t = dde.grad.jacobian(y, x, i=1, j=2) + uy_x = dde.grad.jacobian(y, x, i=2, j=0) + uy_y = dde.grad.jacobian(y, x, i=2, j=1) + uy_t = dde.grad.jacobian(y, x, i=2, j=2) + p_x = dde.grad.jacobian(y, x, i=3, j=0) + p_y = dde.grad.jacobian(y, x, i=3, j=1) + Fx_x = dde.grad.jacobian(Fx, x, i=0, j=0) + Fy_y = dde.grad.jacobian(Fy, x, i=0, j=1) + E_t = dde.grad.jacobian(E, x, i=0, j=2) + + eq1 = h_t + hu_x + hu_y + eq2 = h * (ux_t + ux * ux_x + uy * ux_y) - p_x + eq3 = h * (uy_t + ux * uy_x + uy * uy_y) - p_y + eq4 = E_t + Fx_x + Fy_y + + return eq1 + eq2 + eq3 + eq4 + + +def pde_CFD3d(x, y, gamma): + h = y[..., 0].unsqueeze(1) # rho + ux = y[..., 1].unsqueeze(1) # vx + uy = y[..., 2].unsqueeze(1) # vy + uz = y[..., 3].unsqueeze(1) # vz + p = y[..., 4].unsqueeze(1) # p + E = p / (gamma - 1.0) + 0.5 * h * (ux**2 + uy**2 + uz**2) + E = E.unsqueeze(1) + Fx = ux * (E + p) + Fx = Fx.unsqueeze(1) + Fy = uy * (E + p) + Fy = Fy.unsqueeze(1) + Fz = uz * (E + p) + Fz = Fz.unsqueeze(1) + + # non conservative form + hu_x = dde.grad.jacobian(h * ux, x, i=0, j=0) + hu_y = dde.grad.jacobian(h * uy, x, i=0, j=1) + hu_z = dde.grad.jacobian(h * uy, x, i=0, j=2) + h_t = dde.grad.jacobian(y, x, i=0, j=3) + ux_x = dde.grad.jacobian(y, x, i=1, j=0) + ux_y = dde.grad.jacobian(y, x, i=1, j=1) + ux_z = dde.grad.jacobian(y, x, i=1, j=2) + ux_t = dde.grad.jacobian(y, x, i=1, j=3) + uy_x = dde.grad.jacobian(y, x, i=2, j=0) + uy_y = dde.grad.jacobian(y, x, i=2, j=1) + uy_z = dde.grad.jacobian(y, x, i=2, j=2) + uy_t = dde.grad.jacobian(y, x, i=2, j=3) + uz_x = dde.grad.jacobian(y, x, i=3, j=0) + uz_y = dde.grad.jacobian(y, x, i=3, j=1) + uz_z = dde.grad.jacobian(y, x, i=3, j=2) + uz_t = dde.grad.jacobian(y, x, i=3, j=3) + p_x = dde.grad.jacobian(y, x, i=4, j=0) + p_y = dde.grad.jacobian(y, x, i=4, j=1) + p_z = dde.grad.jacobian(y, x, i=4, j=2) + Fx_x = dde.grad.jacobian(Fx, x, i=0, j=0) + Fy_y = dde.grad.jacobian(Fy, x, i=0, j=1) + Fz_z = dde.grad.jacobian(Fz, x, i=0, j=2) + E_t = dde.grad.jacobian(E, x, i=0, j=3) + + eq1 = h_t + hu_x + hu_y + hu_z + eq2 = h * (ux_t + ux * ux_x + uy * ux_y + uz * ux_z) - p_x + eq3 = h * (uy_t + ux * uy_x + uy * uy_y + uz * uy_z) - p_y + eq4 = h * (uz_t + ux * uz_x + uy * uz_y + uz * uz_z) - p_z + eq5 = E_t + Fx_x + Fy_y + Fz_z + + return eq1 + eq2 + eq3 + eq4 + eq5 diff --git a/pdebench/models/pinn/train.py b/pdebench/models/pinn/train.py index 4708011..098d701 100644 --- a/pdebench/models/pinn/train.py +++ b/pdebench/models/pinn/train.py @@ -1,29 +1,34 @@ """Backend supported: tensorflow.compat.v1, tensorflow, pytorch""" -import deepxde as dde -import numpy as np + +from __future__ import annotations + import pickle +from pathlib import Path + +import deepxde as dde import matplotlib.pyplot as plt -import sys +import numpy as np import torch - -from typing import Tuple - -sys.path.append(".") -from .utils import ( - PINNDatasetRadialDambreak, - PINNDatasetDiffReact, - PINNDataset2D, - PINNDatasetDiffSorption, - PINNDatasetBump, -) -from .pde_definitions import ( +from pdebench.models.metrics import metric_func +from pdebench.models.pinn.pde_definitions import ( + pde_adv1d, + pde_burgers1D, + pde_CFD1d, + pde_CFD2d, pde_diffusion_reaction, - pde_swe2d, + pde_diffusion_reaction_1d, pde_diffusion_sorption, - pde_swe1d, + pde_swe2d, +) +from pdebench.models.pinn.utils import ( + PINNDataset1Dpde, + PINNDataset2D, + PINNDataset2Dpde, + PINNDataset3Dpde, + PINNDatasetDiffReact, + PINNDatasetDiffSorption, + PINNDatasetRadialDambreak, ) - -from metrics import metrics, metric_func def setup_diffusion_sorption(filename, seed): @@ -127,8 +132,7 @@ def setup_diffusion_reaction(filename, seed): return model, dataset -def setup_swe_2d(filename, seed) -> Tuple[dde.Model, PINNDataset2D]: - +def setup_swe_2d(filename, seed) -> tuple[dde.Model, PINNDataset2D]: dataset = PINNDatasetRadialDambreak(filename, seed) # TODO: read from dataset config file @@ -176,7 +180,238 @@ def setup_swe_2d(filename, seed) -> Tuple[dde.Model, PINNDataset2D]: return model, dataset -def run_training(scenario, epochs, learning_rate, model_update, flnm, seed): +def _boundary_r(x, on_boundary, xL, xR): + return (on_boundary and np.isclose(x[0], xL)) or ( + on_boundary and np.isclose(x[0], xR) + ) + + +def setup_pde1D( + filename="1D_Advection_Sols_beta0.1.hdf5", + root_path="data", + val_batch_idx=-1, + input_ch=2, + output_ch=1, + hidden_ch=40, + xL=0.0, + xR=1.0, + if_periodic_bc=True, + aux_params=None, +): + # TODO: read from dataset config file + if aux_params is None: + aux_params = [0.1] + geom = dde.geometry.Interval(xL, xR) + boundary_r = lambda x, on_boundary: _boundary_r(x, on_boundary, xL, xR) + if filename[0] == "R": + timedomain = dde.geometry.TimeDomain(0, 1.0) + pde = lambda x, y: pde_diffusion_reaction_1d(x, y, aux_params[0], aux_params[1]) + elif filename.split("_")[1][0] == "A": + timedomain = dde.geometry.TimeDomain(0, 2.0) + pde = lambda x, y: pde_adv1d(x, y, aux_params[0]) + elif filename.split("_")[1][0] == "B": + timedomain = dde.geometry.TimeDomain(0, 2.0) + pde = lambda x, y: pde_burgers1D(x, y, aux_params[0]) + elif filename.split("_")[1][0] == "C": + timedomain = dde.geometry.TimeDomain(0, 1.0) + pde = lambda x, y: pde_CFD1d(x, y, aux_params[0]) + geomtime = dde.geometry.GeometryXTime(geom, timedomain) + + dataset = PINNDataset1Dpde( + filename, root_path=root_path, val_batch_idx=val_batch_idx + ) + # prepare initial condition + initial_input, initial_u = dataset.get_initial_condition() + if filename.split("_")[1][0] == "C": + ic_data_d = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[:, 0].unsqueeze(1), component=0 + ) + ic_data_v = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[:, 1].unsqueeze(1), component=1 + ) + ic_data_p = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[:, 2].unsqueeze(1), component=2 + ) + else: + ic_data_u = dde.icbc.PointSetBC(initial_input.cpu(), initial_u, component=0) + # prepare boundary condition + if if_periodic_bc: + if filename.split("_")[1][0] == "C": + bc_D = dde.icbc.PeriodicBC(geomtime, 0, boundary_r) + bc_V = dde.icbc.PeriodicBC(geomtime, 1, boundary_r) + bc_P = dde.icbc.PeriodicBC(geomtime, 2, boundary_r) + + data = dde.data.TimePDE( + geomtime, + pde, + [ic_data_d, ic_data_v, ic_data_p, bc_D, bc_V, bc_P], + num_domain=1000, + num_boundary=1000, + num_initial=5000, + ) + else: + bc = dde.icbc.PeriodicBC(geomtime, 0, boundary_r) + data = dde.data.TimePDE( + geomtime, + pde, + [ic_data_u, bc], + num_domain=1000, + num_boundary=1000, + num_initial=5000, + ) + else: + ic = dde.icbc.IC( + geomtime, + lambda x: -np.sin(np.pi * x[:, 0:1]), + lambda _, on_initial: on_initial, + ) + bd_input, bd_uL, bd_uR = dataset.get_boundary_condition() + bc_data_uL = dde.icbc.PointSetBC(bd_input.cpu(), bd_uL, component=0) + bc_data_uR = dde.icbc.PointSetBC(bd_input.cpu(), bd_uR, component=0) + + data = dde.data.TimePDE( + geomtime, + pde, + [ic, bc_data_uL, bc_data_uR], + num_domain=1000, + num_boundary=1000, + num_initial=5000, + ) + net = dde.nn.FNN( + [input_ch] + [hidden_ch] * 6 + [output_ch], "tanh", "Glorot normal" + ) + model = dde.Model(data, net) + + return model, dataset + + +def setup_CFD2D( + filename="2D_CFD_RAND_Eta1.e-8_Zeta1.e-8_periodic_Train.hdf5", + root_path="data", + val_batch_idx=-1, + input_ch=2, + output_ch=4, + hidden_ch=40, + xL=0.0, + xR=1.0, + yL=0.0, + yR=1.0, + if_periodic_bc=True, + aux_params=None, +): + # TODO: read from dataset config file + if aux_params is None: + aux_params = [1.6667] + geom = dde.geometry.Rectangle((-1, -1), (1, 1)) + timedomain = dde.geometry.TimeDomain(0.0, 1.0) + pde = lambda x, y: pde_CFD2d(x, y, aux_params[0]) + geomtime = dde.geometry.GeometryXTime(geom, timedomain) + + dataset = PINNDataset2Dpde( + filename, root_path=root_path, val_batch_idx=val_batch_idx + ) + # prepare initial condition + initial_input, initial_u = dataset.get_initial_condition() + ic_data_d = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 0].unsqueeze(1), component=0 + ) + ic_data_vx = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 1].unsqueeze(1), component=1 + ) + ic_data_vy = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 2].unsqueeze(1), component=2 + ) + ic_data_p = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 3].unsqueeze(1), component=3 + ) + # prepare boundary condition + # bc = dde.icbc.PeriodicBC(geomtime, lambda x: 0, lambda _, on_boundary: on_boundary) + data = dde.data.TimePDE( + geomtime, + pde, + [ic_data_d, ic_data_vx, ic_data_vy, ic_data_p], # , bc], + num_domain=1000, + num_boundary=1000, + num_initial=5000, + ) + net = dde.nn.FNN( + [input_ch] + [hidden_ch] * 6 + [output_ch], "tanh", "Glorot normal" + ) + model = dde.Model(data, net) + + return model, dataset + + +def setup_CFD3D( + filename="3D_CFD_RAND_Eta1.e-8_Zeta1.e-8_periodic_Train.hdf5", + root_path="data", + val_batch_idx=-1, + input_ch=2, + output_ch=4, + hidden_ch=40, + aux_params=None, +): + # TODO: read from dataset config file + if aux_params is None: + aux_params = [1.6667] + geom = dde.geometry.Cuboid((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)) + timedomain = dde.geometry.TimeDomain(0.0, 1.0) + pde = lambda x, y: pde_CFD2d(x, y, aux_params[0]) + geomtime = dde.geometry.GeometryXTime(geom, timedomain) + + dataset = PINNDataset3Dpde( + filename, root_path=root_path, val_batch_idx=val_batch_idx + ) + # prepare initial condition + initial_input, initial_u = dataset.get_initial_condition() + ic_data_d = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 0].unsqueeze(1), component=0 + ) + ic_data_vx = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 1].unsqueeze(1), component=1 + ) + ic_data_vy = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 2].unsqueeze(1), component=2 + ) + ic_data_vz = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 3].unsqueeze(1), component=3 + ) + ic_data_p = dde.icbc.PointSetBC( + initial_input.cpu(), initial_u[..., 4].unsqueeze(1), component=4 + ) + # prepare boundary condition + bc = dde.icbc.PeriodicBC(geomtime, lambda x: 0, lambda _, on_boundary: on_boundary) + data = dde.data.TimePDE( + geomtime, + pde, + [ic_data_d, ic_data_vx, ic_data_vy, ic_data_vz, ic_data_p, bc], + num_domain=1000, + num_boundary=1000, + num_initial=5000, + ) + net = dde.nn.FNN( + [input_ch] + [hidden_ch] * 6 + [output_ch], "tanh", "Glorot normal" + ) + model = dde.Model(data, net) + + return model, dataset + + +def _run_training( + scenario, + epochs, + learning_rate, + model_update, + flnm, + input_ch, + output_ch, + root_path, + val_batch_idx, + if_periodic_bc, + aux_params, + if_single_run, + seed, +): if scenario == "swe2d": model, dataset = setup_swe_2d(filename=flnm, seed=seed) n_components = 1 @@ -186,11 +421,43 @@ def run_training(scenario, epochs, learning_rate, model_update, flnm, seed): elif scenario == "diff-sorp": model, dataset = setup_diffusion_sorption(filename=flnm, seed=seed) n_components = 1 + elif scenario == "pde1D": + model, dataset = setup_pde1D( + filename=flnm, + root_path=root_path, + input_ch=input_ch, + output_ch=output_ch, + val_batch_idx=val_batch_idx, + if_periodic_bc=if_periodic_bc, + aux_params=aux_params, + ) + n_components = 3 if flnm.split("_")[1][0] == "C" else 1 + elif scenario == "CFD2D": + model, dataset = setup_CFD2D( + filename=flnm, + root_path=root_path, + input_ch=input_ch, + output_ch=output_ch, + val_batch_idx=val_batch_idx, + aux_params=aux_params, + ) + n_components = 4 + elif scenario == "CFD3D": + model, dataset = setup_CFD3D( + filename=flnm, + root_path=root_path, + input_ch=input_ch, + output_ch=output_ch, + val_batch_idx=val_batch_idx, + aux_params=aux_params, + ) + n_components = 5 else: - raise NotImplementedError(f"PINN training not implemented for {scenario}") + msg = f"PINN training not implemented for {scenario}" + raise NotImplementedError(msg) # filename - model_name = flnm + "_PINN" + model_name = flnm + "_PINN" if if_single_run else flnm[:-5] + "_PINN" checker = dde.callbacks.ModelCheckpoint( f"{model_name}.pt", save_better_only=True, period=5000 @@ -216,32 +483,99 @@ def run_training(scenario, epochs, learning_rate, model_update, flnm, seed): test_gt, n_last_time_steps=20, n_components=n_components ) - errs = metric_func(test_pred, test_gt) - errors = [np.array(err.cpu()) for err in errs] - print(errors) - pickle.dump(errors, open(model_name + ".pickle", "wb")) - - # plot sample - plot_input = dataset.generate_plot_input(time=1.0) - dim = dataset.config["plot"]["dim"] - xdim = dataset.config["sim"]["xdim"] - if dim == 2: - ydim = dataset.config["sim"]["ydim"] - y_pred = model.predict(plot_input)[:, 0] - if dim == 1: - plt.figure() - plt.plot(y_pred) - elif dim == 2: - im_data = y_pred.reshape(xdim, ydim) - plt.figure() - plt.imshow(im_data) - - plt.savefig(f"{model_name}.png") - - # TODO: implement function to get specific timestep from dataset - # y_true = dataset[:][1][-xdim * ydim :] - - # print("L2 relative error:", dde.metrics.l2_relative_error(y_true.cpu(), y_pred)) + if if_single_run: + errs = metric_func(test_pred, test_gt) + errors = [np.array(err.cpu()) for err in errs] + with Path(model_name + ".pickle").open("wb") as f: + pickle.dump(errors, f) + + # plot sample + plot_input = dataset.generate_plot_input(time=1.0) + if scenario == "pde1D": + xdim = dataset.xdim + dim = 1 + else: + dim = dataset.config["plot"]["dim"] + xdim = dataset.config["sim"]["xdim"] + if dim == 2: + ydim = dataset.config["sim"]["ydim"] + + y_pred = model.predict(plot_input)[:, 0] + if dim == 1: + plt.figure() + plt.plot(y_pred) + elif dim == 2: + im_data = y_pred.reshape(xdim, ydim) + plt.figure() + plt.imshow(im_data) + + plt.savefig(f"{model_name}.png") + return None + + # TODO: implement function to get specific timestep from dataset + # y_true = dataset[:][1][-xdim * ydim :] + return test_pred, test_gt, model_name + + +def run_training( + scenario, + epochs, + learning_rate, + model_update, + flnm, + input_ch=1, + output_ch=1, + root_path="../data/", + val_num=10, + if_periodic_bc=True, + aux_params=None, + seed="0000", +): + if aux_params is None: + aux_params = [None] + if val_num == 1: # single job + _run_training( + scenario, + epochs, + learning_rate, + model_update, + flnm, + input_ch, + output_ch, + root_path, + -val_num, + if_periodic_bc, + aux_params, + if_single_run=True, + seed=seed, + ) + else: + for val_batch_idx in range(-1, -val_num, -1): + test_pred, test_gt, model_name = _run_training( + scenario, + epochs, + learning_rate, + model_update, + flnm, + input_ch, + output_ch, + root_path, + val_batch_idx, + if_periodic_bc, + aux_params, + if_single_run=False, + seed=seed, + ) + if val_batch_idx == -1: + pred, target = test_pred.unsqueeze(0), test_gt.unsqueeze(0) + else: + pred = torch.cat([pred, test_pred.unsqueeze(0)], 0) + target = torch.cat([target, test_gt.unsqueeze(0)], 0) + + errs = metric_func(test_pred, test_gt) + errors = [np.array(err.cpu()) for err in errs] + with Path(model_name + ".pickle").open("wb") as f: + pickle.dump(errors, f) if __name__ == "__main__": @@ -258,7 +592,7 @@ def run_training(scenario, epochs, learning_rate, model_update, flnm, seed): epochs=100, learning_rate=1e-3, model_update=500, - flnm="2D_diff-react_NA_NA_0000.h5", + flnm="2D_diff-react_NA_NA.h5", seed="0000", ) # run_training( diff --git a/pdebench/models/pinn/utils.py b/pdebench/models/pinn/utils.py index 92dba36..91d39a7 100644 --- a/pdebench/models/pinn/utils.py +++ b/pdebench/models/pinn/utils.py @@ -1,18 +1,18 @@ -# -*- coding: utf-8 -*- """ Created on Wed Apr 20 09:43:15 2022 @author: timot """ +from __future__ import annotations + +from pathlib import Path + +import h5py import numpy as np import torch -from torch.utils.data import Dataset -from torch.utils.data import DataLoader -import os -import h5py -from omegaconf import DictConfig, OmegaConf import yaml +from torch.utils.data import Dataset class PINNDataset1D(Dataset): @@ -24,8 +24,8 @@ def __init__(self, filename, seed): self.seed = seed # load data file - root_path = os.path.abspath("../data") - data_path = os.path.join(root_path, filename) + root_path = Path("../data").resolve() + data_path = root_path / filename with h5py.File(data_path, "r") as h5_file: seed_group = h5_file[seed] @@ -51,7 +51,7 @@ def __init__(self, filename, seed): # permute from [t, x] -> [x, t] permute_idx = list(range(1, len(self.data_output.shape) - 1)) - permute_idx.extend(list([0, -1])) + permute_idx.extend([0, -1]) self.data_output = self.data_output.permute(permute_idx) def get_test_data(self, n_last_time_steps, n_components=1): @@ -97,8 +97,7 @@ def generate_plot_input(self, time=1.0): # xx, yy = np.meshgrid(x_space, y_space) tt = np.ones_like(x_space) * time - val_input = np.vstack((x_space, tt)).T - return val_input + return np.vstack((x_space, tt)).T def __len__(self): return len(self.data_output) @@ -119,8 +118,8 @@ def __init__(self, filename, seed): self.seed = seed # load data file - root_path = os.path.abspath("../data") - data_path = os.path.join(root_path, filename) + root_path = Path("../data").resolve() + data_path = root_path / filename with h5py.File(data_path, "r") as h5_file: seed_group = h5_file[seed] @@ -148,7 +147,7 @@ def __init__(self, filename, seed): # permute from [t, x, y] -> [x, y, t] permute_idx = list(range(1, len(self.data_output.shape) - 1)) - permute_idx.extend(list([0, -1])) + permute_idx.extend([0, -1]) self.data_output = self.data_output.permute(permute_idx) def generate_plot_input(self, time=1.0): @@ -164,8 +163,7 @@ def generate_plot_input(self, time=1.0): ) xx, yy = np.meshgrid(x_space, y_space) tt = np.ones_like(xx) * time - val_input = np.vstack((np.ravel(xx), np.ravel(yy), np.ravel(tt))).T - return val_input + return np.vstack((np.ravel(xx), np.ravel(yy), np.ravel(tt))).T def __len__(self): return len(self.data_output) @@ -244,12 +242,10 @@ def initial_h(coords): h_out = 1.0 dam_radius = self.config["sim"]["dam_radius"] - h_initial = np.expand_dims( + return np.expand_dims( h_in * (r <= dam_radius) + h_out * (r > dam_radius), 1 ) - return h_initial - return initial_h @@ -264,26 +260,15 @@ def __init__(self, filename, seed): def get_initial_condition(self): Nx = len(self.data_grid_x) Ny = len(self.data_grid_y) - Nt = len(self.data_grid_t) - np.random.seed(self.config["sim"]["seed"]) + rng = np.random.default_rng(self.config["sim"]["seed"]) - u0 = np.random.randn(Nx * Ny) - v0 = np.random.randn(Nx * Ny) + u0 = rng.standard_normal(Nx * Ny) + v0 = rng.standard_normal(Nx * Ny) u0 = u0.reshape(Nx * Ny) v0 = v0.reshape(Nx * Ny) - x_space = np.linspace( - self.config["sim"]["x_left"], - self.config["sim"]["x_right"], - self.config["sim"]["xdim"], - ) - y_space = np.linspace( - self.config["sim"]["y_bottom"], - self.config["sim"]["y_top"], - self.config["sim"]["ydim"], - ) xx, yy = np.meshgrid(self.data_grid_x.cpu(), self.data_grid_y.cpu()) tt = np.zeros_like(xx) ic_input = np.vstack((np.ravel(xx), np.ravel(yy), np.ravel(tt))).T @@ -313,8 +298,430 @@ def get_initial_condition(self): # Generate initial condition Nx = self.config["sim"]["xdim"] - np.random.seed(self.config["sim"]["seed"]) + rng = np.random.default_rng(self.config["sim"]["seed"]) - u0 = np.ones(Nx) * np.random.uniform(0, 0.2) + u0 = np.ones(Nx) * rng.uniform(0, 0.2) return (self.data_input[:Nx, :], np.expand_dims(u0, 1)) + + +class PINNDataset1Dpde(Dataset): + def __init__(self, filename, root_path="data", val_batch_idx=-1): + """ + :param filename: filename that contains the dataset + :type filename: STR + """ + + # load data file + data_path = Path(root_path) / filename + h5_file = h5py.File(data_path, "r") + + # build input data from individual dimensions + # dim x = [x] + self.data_grid_x = torch.tensor(h5_file["x-coordinate"], dtype=torch.float) + self.dx = self.data_grid_x[1] - self.data_grid_x[0] + self.xL = self.data_grid_x[0] - 0.5 * self.dx + self.xR = self.data_grid_x[-1] + 0.5 * self.dx + self.xdim = self.data_grid_x.size(0) + # # dim t = [t] + self.data_grid_t = torch.tensor(h5_file["t-coordinate"], dtype=torch.float) + + # main data + keys = list(h5_file.keys()) + keys.sort() + if "tensor" in keys: + self.data_output = torch.tensor( + np.array(h5_file["tensor"][val_batch_idx]), dtype=torch.float + ) + # permute from [t, x] -> [x, t] + self.data_output = self.data_output.T + + # for init/boundary conditions + self.init_data = self.data_output[..., 0, None] + self.bd_data_L = self.data_output[0, :, None] + self.bd_data_R = self.data_output[-1, :, None] + + else: + _data1 = np.array(h5_file["density"][val_batch_idx]) + _data2 = np.array(h5_file["Vx"][val_batch_idx]) + _data3 = np.array(h5_file["pressure"][val_batch_idx]) + _data = np.concatenate( + [_data1[..., None], _data2[..., None], _data3[..., None]], axis=-1 + ) + # permute from [t, x] -> [x, t] + _data = np.transpose(_data, (1, 0, 2)) + + self.data_output = torch.tensor(_data, dtype=torch.float) + del (_data, _data1, _data2, _data3) + + # for init/boundary conditions + self.init_data = self.data_output[:, 0] + self.bd_data_L = self.data_output[0] + self.bd_data_R = self.data_output[-1] + + self.tdim = self.data_output.size(1) + self.data_grid_t = self.data_grid_t[: self.tdim] + + XX, TT = torch.meshgrid( + [self.data_grid_x, self.data_grid_t], + indexing="ij", + ) + + self.data_input = torch.vstack([XX.ravel(), TT.ravel()]).T + + h5_file.close() + if "tensor" in keys: + self.data_output = self.data_output.reshape(-1, 1) + else: + self.data_output = self.data_output.reshape(-1, 3) + + def get_initial_condition(self): + # return (self.data_grid_x[:, None], self.init_data) + return (self.data_input[:: self.tdim, :], self.init_data) + + def get_boundary_condition(self): + # return (self.data_grid_t[:self.nt, None], self.bd_data_L, self.bd_data_R) + return (self.data_input[: self.xdim, :], self.bd_data_L, self.bd_data_R) + + def get_test_data(self, n_last_time_steps, n_components=1): + n_x = len(self.data_grid_x) + n_t = len(self.data_grid_t) + + # start_idx = n_x * n_y * (n_t - n_last_time_steps) + test_input_x = self.data_input[:, 0].reshape((n_x, n_t)) + test_input_t = self.data_input[:, 1].reshape((n_x, n_t)) + test_output = self.data_output.reshape((n_x, n_t, n_components)) + + # extract last n time steps + test_input_x = test_input_x[:, -n_last_time_steps:] + test_input_t = test_input_t[:, -n_last_time_steps:] + test_output = test_output[:, -n_last_time_steps:, :] + + test_input = torch.vstack([test_input_x.ravel(), test_input_t.ravel()]).T + + # stack depending on number of output components + test_output_stacked = test_output[..., 0].ravel() + if n_components > 1: + for i in range(1, n_components): + test_output_stacked = torch.vstack( + [test_output_stacked, test_output[..., i].ravel()] + ) + else: + test_output_stacked = test_output_stacked.unsqueeze(1) + + test_output = test_output_stacked.T + + return test_input, test_output + + def unravel_tensor(self, raveled_tensor, n_last_time_steps, n_components=1): + n_x = len(self.data_grid_x) + return raveled_tensor.reshape((1, n_x, n_last_time_steps, n_components)) + + def generate_plot_input(self, time=1.0): + x_space = np.linspace(self.xL, self.xR, self.xdim) + # xx, yy = np.meshgrid(x_space, y_space) + + tt = np.ones_like(x_space) * time + return np.vstack((x_space, tt)).T + + def __len__(self): + return len(self.data_output) + + def __getitem__(self, idx): + return self.data_input[idx, :], self.data_output[idx] + + +class PINNDataset2Dpde(Dataset): + def __init__(self, filename, root_path="data", val_batch_idx=-1, rdc_x=9, rdc_y=9): + """ + :param filename: filename that contains the dataset + :type filename: STR + """ + + # load data file + data_path = Path(root_path) / filename + h5_file = h5py.File(data_path, "r") + + # build input data from individual dimensions + # dim x = [x] + self.data_grid_x = torch.tensor(h5_file["x-coordinate"], dtype=torch.float) + self.data_grid_x = self.data_grid_x[::rdc_x] + self.dx = self.data_grid_x[1] - self.data_grid_x[0] + self.xL = self.data_grid_x[0] - 0.5 * self.dx + self.xR = self.data_grid_x[-1] + 0.5 * self.dx + self.xdim = self.data_grid_x.size(0) + # dim y = [y] + self.data_grid_y = torch.tensor(h5_file["y-coordinate"], dtype=torch.float) + self.data_grid_y = self.data_grid_y[::rdc_y] + self.dy = self.data_grid_y[1] - self.data_grid_y[0] + self.yL = self.data_grid_y[0] - 0.5 * self.dy + self.yR = self.data_grid_y[-1] + 0.5 * self.dy + self.ydim = self.data_grid_y.size(0) + # # dim t = [t] + self.data_grid_t = torch.tensor(h5_file["t-coordinate"], dtype=torch.float) + + # main data + _data1 = np.array(h5_file["density"][val_batch_idx]) + _data2 = np.array(h5_file["Vx"][val_batch_idx]) + _data3 = np.array(h5_file["Vy"][val_batch_idx]) + _data4 = np.array(h5_file["pressure"][val_batch_idx]) + _data = np.concatenate( + [ + _data1[..., None], + _data2[..., None], + _data3[..., None], + _data4[..., None], + ], + axis=-1, + ) + # permute from [t, x, y, v] -> [x, y, t, v] + _data = np.transpose(_data, (1, 2, 0, 3)) + _data = _data[::rdc_x, ::rdc_y] + + self.data_output = torch.tensor(_data, dtype=torch.float) + del (_data, _data1, _data2, _data3, _data4) + + # for init/boundary conditions + self.init_data = self.data_output[..., 0, :] + self.bd_data_xL = self.data_output[0] + self.bd_data_xR = self.data_output[-1] + self.bd_data_yL = self.data_output[:, 0] + self.bd_data_yR = self.data_output[:, -1] + + self.tdim = self.data_output.size(2) + self.data_grid_t = self.data_grid_t[: self.tdim] + + XX, YY, TT = torch.meshgrid( + [self.data_grid_x, self.data_grid_y, self.data_grid_t], + indexing="ij", + ) + + self.data_input = torch.vstack([XX.ravel(), YY.ravel(), TT.ravel()]).T + + h5_file.close() + self.data_output = self.data_output.reshape(-1, 4) + + def get_initial_condition(self): + # return (self.data_grid_x[:, None], self.init_data) + return (self.data_input[:: self.tdim, :], self.init_data) + + def get_boundary_condition(self): + # return (self.data_grid_t[:self.nt, None], self.bd_data_L, self.bd_data_R) + return ( + self.data_input[: self.xdim * self.ydim, :], + self.bd_data_xL, + self.bd_data_xR, + self.bd_data_yL, + self.bd_data_yR, + ) + + def get_test_data(self, n_last_time_steps, n_components=1): + n_x = len(self.data_grid_x) + n_y = len(self.data_grid_y) + n_t = len(self.data_grid_t) + + # start_idx = n_x * n_y * (n_t - n_last_time_steps) + test_input_x = self.data_input[:, 0].reshape((n_x, n_y, n_t)) + test_input_y = self.data_input[:, 1].reshape((n_x, n_y, n_t)) + test_input_t = self.data_input[:, 4].reshape((n_x, n_y, n_t)) + test_output = self.data_output.reshape((n_x, n_y, n_t, n_components)) + + # extract last n time steps + test_input_x = test_input_x[:, :, -n_last_time_steps:] + test_input_y = test_input_y[:, :, -n_last_time_steps:] + test_input_t = test_input_t[:, :, -n_last_time_steps:] + test_output = test_output[:, :, -n_last_time_steps:, :] + + test_input = torch.vstack( + [test_input_x.ravel(), test_input_y.ravel(), test_input_t.ravel()] + ).T + + # stack depending on number of output components + test_output_stacked = test_output[..., 0].ravel() + if n_components > 1: + for i in range(1, n_components): + test_output_stacked = torch.vstack( + [test_output_stacked, test_output[..., i].ravel()] + ) + else: + test_output_stacked = test_output_stacked.unsqueeze(1) + + test_output = test_output_stacked.T + + return test_input, test_output + + def unravel_tensor(self, raveled_tensor, n_last_time_steps, n_components=1): + n_x = len(self.data_grid_x) + n_y = len(self.data_grid_y) + return raveled_tensor.reshape((1, n_x, n_y, n_last_time_steps, n_components)) + + def generate_plot_input(self, time=1.0): # noqa: ARG002 + return None + + def __len__(self): + return len(self.data_output) + + def __getitem__(self, idx): + return self.data_input[idx, :], self.data_output[idx].unsqueeze(1) + + +class PINNDataset3Dpde(Dataset): + def __init__( + self, filename, root_path="data", val_batch_idx=-1, rdc_x=2, rdc_y=2, rdc_z=2 + ): + """ + :param filename: filename that contains the dataset + :type filename: STR + """ + + # load data file + data_path = Path(root_path) / filename + h5_file = h5py.File(data_path, "r") + + # build input data from individual dimensions + # dim x = [x] + self.data_grid_x = torch.tensor(h5_file["x-coordinate"], dtype=torch.float) + self.data_grid_x = self.data_grid_x[::rdc_x] + self.dx = self.data_grid_x[1] - self.data_grid_x[0] + self.xL = self.data_grid_x[0] - 0.5 * self.dx + self.xR = self.data_grid_x[-1] + 0.5 * self.dx + self.xdim = self.data_grid_x.size(0) + # dim y = [y] + self.data_grid_y = torch.tensor(h5_file["y-coordinate"], dtype=torch.float) + self.data_grid_y = self.data_grid_y[::rdc_y] + self.dy = self.data_grid_y[1] - self.data_grid_y[0] + self.yL = self.data_grid_y[0] - 0.5 * self.dy + self.yR = self.data_grid_y[-1] + 0.5 * self.dy + self.ydim = self.data_grid_y.size(0) + # dim z = [z] + self.data_grid_z = torch.tensor(h5_file["z-coordinate"], dtype=torch.float) + self.data_grid_z = self.data_grid_z[::rdc_z] + self.dz = self.data_grid_z[1] - self.data_grid_z[0] + self.zL = self.data_grid_z[0] - 0.5 * self.dz + self.zR = self.data_grid_z[-1] + 0.5 * self.dz + self.zdim = self.data_grid_z.size(0) + # # dim t = [t] + self.data_grid_t = torch.tensor(h5_file["t-coordinate"], dtype=torch.float) + + # main data + _data1 = np.array(h5_file["density"][val_batch_idx]) + _data2 = np.array(h5_file["Vx"][val_batch_idx]) + _data3 = np.array(h5_file["Vy"][val_batch_idx]) + _data4 = np.array(h5_file["Vz"][val_batch_idx]) + _data5 = np.array(h5_file["pressure"][val_batch_idx]) + _data = np.concatenate( + [ + _data1[..., None], + _data2[..., None], + _data3[..., None], + _data4[..., None], + _data5[..., None], + ], + axis=-1, + ) + # permute from [t, x, y, z, v] -> [x, y, z, t, v] + _data = np.transpose(_data, (1, 2, 3, 0, 4)) + _data = _data[::rdc_x, ::rdc_y, ::rdc_z] + + self.data_output = torch.tensor(_data, dtype=torch.float) + del (_data, _data1, _data2, _data3, _data4, _data5) + + # for init/boundary conditions + self.init_data = self.data_output[..., 0, :] + self.bd_data_xL = self.data_output[0] + self.bd_data_xR = self.data_output[-1] + self.bd_data_yL = self.data_output[:, 0] + self.bd_data_yR = self.data_output[:, -1] + self.bd_data_zL = self.data_output[:, :, 0] + self.bd_data_zR = self.data_output[:, :, -1] + + self.tdim = self.data_output.size(3) + self.data_grid_t = self.data_grid_t[: self.tdim] + + XX, YY, ZZ, TT = torch.meshgrid( + [self.data_grid_x, self.data_grid_y, self.data_grid_z, self.data_grid_t], + indexing="ij", + ) + + self.data_input = torch.vstack( + [XX.ravel(), YY.ravel(), ZZ.ravel(), TT.ravel()] + ).T + + h5_file.close() + self.data_output = self.data_output.reshape(-1, 5) + + def get_initial_condition(self): + # return (self.data_grid_x[:, None], self.init_data) + return (self.data_input[:: self.tdim, :], self.init_data) + + def get_boundary_condition(self): + # return (self.data_grid_t[:self.nt, None], self.bd_data_L, self.bd_data_R) + return ( + self.data_input[: self.xdim * self.ydim * self.zdim, :], + self.bd_data_xL, + self.bd_data_xR, + self.bd_data_yL, + self.bd_data_yR, + self.bd_data_zL, + self.bd_data_zR, + ) + + def get_test_data(self, n_last_time_steps, n_components=1): + n_x = len(self.data_grid_x) + n_y = len(self.data_grid_y) + n_z = len(self.data_grid_z) + n_t = len(self.data_grid_t) + + # start_idx = n_x * n_y * (n_t - n_last_time_steps) + test_input_x = self.data_input[:, 0].reshape((n_x, n_y, n_z, n_t)) + test_input_y = self.data_input[:, 1].reshape((n_x, n_y, n_z, n_t)) + test_input_z = self.data_input[:, 2].reshape((n_x, n_y, n_z, n_t)) + test_input_t = self.data_input[:, 3].reshape((n_x, n_y, n_z, n_t)) + test_output = self.data_output.reshape((n_x, n_y, n_z, n_t, n_components)) + + # extract last n time steps + test_input_x = test_input_x[:, :, :, -n_last_time_steps:] + test_input_y = test_input_y[:, :, :, -n_last_time_steps:] + test_input_z = test_input_z[:, :, :, -n_last_time_steps:] + test_input_t = test_input_t[:, :, :, -n_last_time_steps:] + test_output = test_output[:, :, :, -n_last_time_steps:, :] + + test_input = torch.vstack( + [ + test_input_x.ravel(), + test_input_y.ravel(), + test_input_z.ravel(), + test_input_t.ravel(), + ] + ).T + + # stack depending on number of output components + test_output_stacked = test_output[..., 0].ravel() + if n_components > 1: + for i in range(1, n_components): + test_output_stacked = torch.vstack( + [test_output_stacked, test_output[..., i].ravel()] + ) + else: + test_output_stacked = test_output_stacked.unsqueeze(1) + + test_output = test_output_stacked.T + + return test_input, test_output + + def unravel_tensor(self, raveled_tensor, n_last_time_steps, n_components=1): + n_x = len(self.data_grid_x) + n_y = len(self.data_grid_y) + n_z = len(self.data_grid_z) + return raveled_tensor.reshape( + (1, n_x, n_y, n_z, n_last_time_steps, n_components) + ) + + def generate_plot_input(self, time=1.0): # noqa: ARG002 + return None + + def __len__(self): + return len(self.data_output) + + def __getitem__(self, idx): + return self.data_input[idx, :], self.data_output[idx].unsqueeze(1) diff --git a/pdebench/models/run_forward_1D.sh b/pdebench/models/run_forward_1D.sh index 19d6916..c910093 100644 --- a/pdebench/models/run_forward_1D.sh +++ b/pdebench/models/run_forward_1D.sh @@ -1,56 +1,73 @@ +#!/bin/bash ## 'FNO' # Advection -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='FNO' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='FNO' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='FNO' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='FNO' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='FNO' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.1.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.1.hdf5' ++args.model_name='FNO'++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='FNO' ++args.if_training=False # Reaction Diffusion -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='FNO' ++args.reduced_resolution_t=1 ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='FNO' ++args.if_training=False # Burgers Eq. -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='FNO' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='FNO' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='FNO' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='FNO' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='FNO' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='FNO' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='FNO' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='FNO' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='FNO' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='FNO' ++args.if_training=False ## Unet # Advection -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='Unet' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='Unet' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='Unet' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='Unet' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='Unet' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.1.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.1.hdf5' ++args.model_name='Unet'++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Adv.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='Unet' ++args.if_training=False # Reaction Diffusion -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='Unet' ++args.reduced_resolution_t=1 ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_ReacDiff.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.model_name='Unet' ++args.if_training=False # Burgers Eq. -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='Unet' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='Unet' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='Unet' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='Unet' -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='Unet' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='Unet' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='Unet' ++args.if_training=False -CUDA_VISIBLE_DEVICES='2' python3 train_models_forward.py +args=config.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='Unet' +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.model_name='Unet' ++args.if_training=False +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_Bgs.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='Unet' ++args.if_training=False +# 'PINN' +# Advection +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Advection_Sols_beta0.1.hdf5' ++args.aux_params=[0.1] +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Advection_Sols_beta0.4.hdf5' ++args.aux_params=[0.4] +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.aux_params=[1.] +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.aux_params=[4.] +# Reaction Diffusion +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.aux_params="[0.5,1.]" ++args.val_time=0.5 +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.aux_params="[0.5,10.]" ++args.val_time=0.5 +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.aux_params="[2.,1.]" ++args.val_time=0.5 +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.aux_params="[2.,10.]" ++args.val_time=0.5 +# Burgers Eq. +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.aux_params="[0.001]" +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.aux_params="[0.01]" +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.aux_params="[0.1]" +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.aux_params="[1.]" diff --git a/pdebench/models/run_inverse.sh b/pdebench/models/run_inverse.sh index 2242dde..fe828e7 100644 --- a/pdebench/models/run_inverse.sh +++ b/pdebench/models/run_inverse.sh @@ -1,27 +1,25 @@ -# /bin/bash +#! /bin/bash # F.Alesiani, 2022, June 6th # Train forward model HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/Advection/Train/1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='FNO' HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/Burgers/Train/1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='FNO' HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/ReactionDiffusion/Train/ReacDiff_Nu1.0_Rho2.0.hdf5' ++args.model_name='FNO' -HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='FNO' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 +HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='FNO' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/Advection/Train/1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='Unet' HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/Burgers/Train/1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='Unet' HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/ReactionDiffusion/Train/ReacDiff_Nu1.0_Rho2.0.hdf5' ++args.model_name='Unet' -HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='Unet' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 +HYDRA_FULL_ERROR=1 python3 train_models_inverse.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='Unet' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 # Inverse HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/Advection/Train/1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='FNO' HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/Burgers/Train/1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='FNO' HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/ReactionDiffusion/Train/ReacDiff_Nu1.0_Rho2.0.hdf5' ++args.model_name='FNO' -HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='FNO' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 +HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='FNO' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/Advection/Train/1D_Advection_Sols_beta4.0.hdf5' ++args.model_name='Unet' HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/Burgers/Train/1D_Burgers_Sols_Nu1.0.hdf5' ++args.model_name='Unet' HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/ReactionDiffusion/Train/ReacDiff_Nu1.0_Rho2.0.hdf5' ++args.model_name='Unet' -HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='Unet' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 - - +HYDRA_FULL_ERROR=1 python3 inverse/train.py ++args.filename='/1D/CFD/Train/1D_CFD_Shock_trans_Train.hdf5' ++args.model_name='Unet' ++args.in_channels=3 ++args.out_channels=3 ++args.num_channels=3 ++args.final_time=5 diff --git a/pdebench/models/train_models_forward.py b/pdebench/models/train_models_forward.py index 98b1276..001c79a 100644 --- a/pdebench/models/train_models_forward.py +++ b/pdebench/models/train_models_forward.py @@ -145,26 +145,23 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import sys, os -import hydra -from omegaconf import DictConfig -import operator -from functools import reduce -from functools import partial +from __future__ import annotations -from timeit import default_timer +import logging + +import hydra +from omegaconf import DictConfig -sys.path.append(".") -from fno.train import run_training as run_training_FNO -from pinn.train import run_training as run_training_PINN -from unet.train import run_training as run_training_Unet +logger = logging.getLogger(__name__) -@hydra.main(config_path="config", config_name="config") +@hydra.main(version_base="1.2", config_path="config", config_name="config_rdb") def main(cfg: DictConfig): if cfg.args.model_name == "FNO": - print("FNO") + from pdebench.models.fno.train import run_training as run_training_FNO + + logger.info("FNO") run_training_FNO( if_training=cfg.args.if_training, continue_training=cfg.args.continue_training, @@ -173,6 +170,7 @@ def main(cfg: DictConfig): width=cfg.args.width, initial_step=cfg.args.initial_step, t_train=cfg.args.t_train, + training_type=cfg.args.training_type, num_channels=cfg.args.num_channels, batch_size=cfg.args.batch_size, epochs=cfg.args.epochs, @@ -182,6 +180,7 @@ def main(cfg: DictConfig): model_update=cfg.args.model_update, flnm=cfg.args.filename, single_file=cfg.args.single_file, + base_path=cfg.args.data_path, reduced_resolution=cfg.args.reduced_resolution, reduced_resolution_t=cfg.args.reduced_resolution_t, reduced_batch=cfg.args.reduced_batch, @@ -195,7 +194,9 @@ def main(cfg: DictConfig): t_max=cfg.args.t_max, ) elif cfg.args.model_name == "Unet": - print("Unet") + from pdebench.models.unet.train import run_training as run_training_Unet + + logger.info("Unet") run_training_Unet( if_training=cfg.args.if_training, continue_training=cfg.args.continue_training, @@ -215,6 +216,7 @@ def main(cfg: DictConfig): model_update=cfg.args.model_update, flnm=cfg.args.filename, single_file=cfg.args.single_file, + base_path=cfg.args.data_path, reduced_resolution=cfg.args.reduced_resolution, reduced_resolution_t=cfg.args.reduced_resolution_t, reduced_batch=cfg.args.reduced_batch, @@ -228,7 +230,10 @@ def main(cfg: DictConfig): t_max=cfg.args.t_max, ) elif cfg.args.model_name == "PINN": - print("PINN") + # not importing globally as DeepXDE changes some global PyTorch settings + from pdebench.models.pinn.train import run_training as run_training_PINN + + logger.info("PINN") run_training_PINN( scenario=cfg.args.scenario, epochs=cfg.args.epochs, @@ -236,9 +241,14 @@ def main(cfg: DictConfig): model_update=cfg.args.model_update, flnm=cfg.args.filename, seed=cfg.args.seed, + input_ch=cfg.args.input_ch, + output_ch=cfg.args.output_ch, + root_path=cfg.args.root_path, + val_num=cfg.args.val_num, + if_periodic_bc=cfg.args.if_periodic_bc, + aux_params=cfg.args.aux_params, ) if __name__ == "__main__": main() - print("Done.") diff --git a/pdebench/models/train_models_inverse.py b/pdebench/models/train_models_inverse.py index 8ecb621..dc7483d 100644 --- a/pdebench/models/train_models_inverse.py +++ b/pdebench/models/train_models_inverse.py @@ -144,93 +144,93 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import sys, os -import hydra -from omegaconf import DictConfig -import operator -from functools import reduce -from functools import partial +from __future__ import annotations -from timeit import default_timer +import logging + +import hydra +from omegaconf import DictConfig +from pdebench.models.fno.train import run_training as run_training_FNO +from pdebench.models.pinn.train import run_training as run_training_PINN +from pdebench.models.unet.train import run_training as run_training_Unet -sys.path.append('.') -from fno.train import run_training as run_training_FNO -from pinn.train import run_training as run_training_PINN -from unet.train import run_training as run_training_Unet +logger = logging.getLogger(__name__) -@hydra.main(config_path='config', config_name='config') +@hydra.main(config_path="config", config_name="config") def main(cfg: DictConfig): - print(cfg.args) - if cfg.args.model_name=='FNO': - print('FNO') - run_training_FNO(if_training=cfg.args.if_training, - continue_training=cfg.args.continue_training, - num_workers=cfg.args.num_workers, - modes=cfg.args.modes, - width=cfg.args.width, - initial_step=cfg.args.initial_step, - t_train=cfg.args.t_train, - num_channels=cfg.args.num_channels, - batch_size=cfg.args.batch_size, - epochs=cfg.args.epochs, - learning_rate=cfg.args.learning_rate, - scheduler_step=cfg.args.scheduler_step, - scheduler_gamma=cfg.args.scheduler_gamma, - model_update=cfg.args.model_update, - flnm=cfg.args.filename, - single_file=cfg.args.single_file, - reduced_resolution=cfg.args.reduced_resolution, - reduced_resolution_t=cfg.args.reduced_resolution_t, - reduced_batch=cfg.args.reduced_batch, - plot=cfg.args.plot, - channel_plot=cfg.args.channel_plot, - x_min=cfg.args.x_min, - x_max=cfg.args.x_max, - y_min=cfg.args.y_min, - y_max=cfg.args.y_max, - t_min=cfg.args.t_min, - t_max=cfg.args.t_max, - base_path = cfg.args.base_path, - training_type = cfg.args.training_type - ) - elif cfg.args.model_name=='Unet': - print('Unet') - run_training_Unet(if_training=cfg.args.if_training, - continue_training=cfg.args.continue_training, - num_workers=cfg.args.num_workers, - initial_step=cfg.args.initial_step, - t_train=cfg.args.t_train, - in_channels=cfg.args.in_channels, - out_channels=cfg.args.out_channels, - batch_size=cfg.args.batch_size, - unroll_step=cfg.args.unroll_step, - ar_mode=cfg.args.ar_mode, - pushforward=cfg.args.pushforward, - epochs=cfg.args.epochs, - learning_rate=cfg.args.learning_rate, - scheduler_step=cfg.args.scheduler_step, - scheduler_gamma=cfg.args.scheduler_gamma, - model_update=cfg.args.model_update, - flnm=cfg.args.filename, - single_file=cfg.args.single_file, - reduced_resolution=cfg.args.reduced_resolution, - reduced_resolution_t=cfg.args.reduced_resolution_t, - reduced_batch=cfg.args.reduced_batch, - plot=cfg.args.plot, - channel_plot=cfg.args.channel_plot, - x_min=cfg.args.x_min, - x_max=cfg.args.x_max, - y_min=cfg.args.y_min, - y_max=cfg.args.y_max, - t_min=cfg.args.t_min, - t_max=cfg.args.t_max, - base_path = cfg.args.base_path, - training_type = cfg.args.training_type - ) + logger.info(cfg.args) + if cfg.args.model_name == "FNO": + logger.info("FNO") + run_training_FNO( + if_training=cfg.args.if_training, + continue_training=cfg.args.continue_training, + num_workers=cfg.args.num_workers, + modes=cfg.args.modes, + width=cfg.args.width, + initial_step=cfg.args.initial_step, + t_train=cfg.args.t_train, + num_channels=cfg.args.num_channels, + batch_size=cfg.args.batch_size, + epochs=cfg.args.epochs, + learning_rate=cfg.args.learning_rate, + scheduler_step=cfg.args.scheduler_step, + scheduler_gamma=cfg.args.scheduler_gamma, + model_update=cfg.args.model_update, + flnm=cfg.args.filename, + single_file=cfg.args.single_file, + reduced_resolution=cfg.args.reduced_resolution, + reduced_resolution_t=cfg.args.reduced_resolution_t, + reduced_batch=cfg.args.reduced_batch, + plot=cfg.args.plot, + channel_plot=cfg.args.channel_plot, + x_min=cfg.args.x_min, + x_max=cfg.args.x_max, + y_min=cfg.args.y_min, + y_max=cfg.args.y_max, + t_min=cfg.args.t_min, + t_max=cfg.args.t_max, + base_path=cfg.args.base_path, + training_type=cfg.args.training_type, + ) + elif cfg.args.model_name == "Unet": + logger.info("Unet") + run_training_Unet( + if_training=cfg.args.if_training, + continue_training=cfg.args.continue_training, + num_workers=cfg.args.num_workers, + initial_step=cfg.args.initial_step, + t_train=cfg.args.t_train, + in_channels=cfg.args.in_channels, + out_channels=cfg.args.out_channels, + batch_size=cfg.args.batch_size, + unroll_step=cfg.args.unroll_step, + ar_mode=cfg.args.ar_mode, + pushforward=cfg.args.pushforward, + epochs=cfg.args.epochs, + learning_rate=cfg.args.learning_rate, + scheduler_step=cfg.args.scheduler_step, + scheduler_gamma=cfg.args.scheduler_gamma, + model_update=cfg.args.model_update, + flnm=cfg.args.filename, + single_file=cfg.args.single_file, + reduced_resolution=cfg.args.reduced_resolution, + reduced_resolution_t=cfg.args.reduced_resolution_t, + reduced_batch=cfg.args.reduced_batch, + plot=cfg.args.plot, + channel_plot=cfg.args.channel_plot, + x_min=cfg.args.x_min, + x_max=cfg.args.x_max, + y_min=cfg.args.y_min, + y_max=cfg.args.y_max, + t_min=cfg.args.t_min, + t_max=cfg.args.t_max, + base_path=cfg.args.base_path, + training_type=cfg.args.training_type, + ) elif cfg.args.model_name == "PINN": - print("PINN") + logger.info("PINN") run_training_PINN( scenario=cfg.args.scenario, epochs=cfg.args.epochs, @@ -240,6 +240,6 @@ def main(cfg: DictConfig): seed=cfg.args.seed, ) + if __name__ == "__main__": main() - print("Done.") \ No newline at end of file diff --git a/pdebench/models/unet/train.py b/pdebench/models/unet/train.py index cf9b57f..858f4d0 100644 --- a/pdebench/models/unet/train.py +++ b/pdebench/models/unet/train.py @@ -1,357 +1,401 @@ -# -*- coding: utf-8 -*- +from __future__ import annotations -import sys -import torch -import numpy as np +import logging import pickle -import torch.nn as nn -import torch.nn.functional as F - -import operator -from functools import reduce -from functools import partial - +from pathlib import Path from timeit import default_timer +import numpy as np +import torch +from pdebench.models.metrics import metrics +from pdebench.models.unet.unet import UNet1d, UNet2d, UNet3d +from pdebench.models.unet.utils import UNetDatasetMult, UNetDatasetSingle +from torch import nn + # torch.manual_seed(0) # np.random.seed(0) +logger = logging.getLogger(__name__) + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def run_training( + if_training, + continue_training, + num_workers, + initial_step, + t_train, + in_channels, + out_channels, + batch_size, + unroll_step, + ar_mode, + pushforward, + epochs, + learning_rate, + scheduler_step, + scheduler_gamma, + model_update, + flnm, + single_file, + reduced_resolution, + reduced_resolution_t, + reduced_batch, + plot, + channel_plot, + x_min, + x_max, + y_min, + y_max, + t_min, + t_max, + base_path="../data/", + training_type="autoregressive", +): + msg = f"Epochs = {epochs}, learning rate = {learning_rate}, scheduler step = {scheduler_step}, scheduler gamma = {scheduler_gamma}" + logger.info(msg) -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - -sys.path.append('.') -from .unet import UNet1d, UNet2d, UNet3d -from .utils import UNetDatasetSingle, UNetDatasetMult -from metrics import metrics - -def run_training(if_training, - continue_training, - num_workers, - initial_step, - t_train, - in_channels, - out_channels, - batch_size, - unroll_step, - ar_mode, - pushforward, - epochs, - learning_rate, - scheduler_step, - scheduler_gamma, - model_update, - flnm, - single_file, - reduced_resolution, - reduced_resolution_t, - reduced_batch, - plot, - channel_plot, - x_min, - x_max, - y_min, - y_max, - t_min, - t_max, - base_path='../data/', - training_type='autoregressive' - ): - - print(f'Epochs = {epochs}, learning rate = {learning_rate}, scheduler step = {scheduler_step}, scheduler gamma = {scheduler_gamma}') - ################################################################ # load data ################################################################ - + if single_file: # filename - model_name = flnm[:-5] + '_Unet' - + model_name = flnm[:-5] + "_Unet" + # Initialize the dataset and dataloader - train_data = UNetDatasetSingle(flnm, - saved_folder=base_path, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - initial_step=initial_step) - val_data = UNetDatasetSingle(flnm, - saved_folder=base_path, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - initial_step=initial_step, - if_test=True) - + train_data = UNetDatasetSingle( + flnm, + saved_folder=base_path, + reduced_resolution=reduced_resolution, + reduced_resolution_t=reduced_resolution_t, + reduced_batch=reduced_batch, + initial_step=initial_step, + ) + val_data = UNetDatasetSingle( + flnm, + saved_folder=base_path, + reduced_resolution=reduced_resolution, + reduced_resolution_t=reduced_resolution_t, + reduced_batch=reduced_batch, + initial_step=initial_step, + if_test=True, + ) + else: # filename - model_name = flnm + '_Unet' - - train_data = UNetDatasetMult(flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - saved_folder=base_path) - val_data = UNetDatasetMult(flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, - if_test=True, - saved_folder=base_path) - - train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, - num_workers=num_workers, shuffle=True) - val_loader = torch.utils.data.DataLoader(val_data, batch_size=batch_size, - num_workers=num_workers, shuffle=False) - + model_name = flnm + "_Unet" + + train_data = UNetDatasetMult( + flnm, + reduced_resolution=reduced_resolution, + reduced_resolution_t=reduced_resolution_t, + reduced_batch=reduced_batch, + saved_folder=base_path, + ) + val_data = UNetDatasetMult( + flnm, + reduced_resolution=reduced_resolution, + reduced_resolution_t=reduced_resolution_t, + reduced_batch=reduced_batch, + if_test=True, + saved_folder=base_path, + ) + + train_loader = torch.utils.data.DataLoader( + train_data, batch_size=batch_size, num_workers=num_workers, shuffle=True + ) + val_loader = torch.utils.data.DataLoader( + val_data, batch_size=batch_size, num_workers=num_workers, shuffle=False + ) + ################################################################ # training and evaluation ################################################################ - - #model = UNet2d(in_channels, out_channels).to(device) + + # model = UNet2d(in_channels, out_channels).to(device) _, _data = next(iter(val_loader)) dimensions = len(_data.shape) - print('Spatial Dimension', dimensions - 3) - if training_type in ['autoregressive']: + msg = f"Spatial Dimension: {dimensions - 3}" + logger.info(msg) + if training_type in ["autoregressive"]: if dimensions == 4: - model = UNet1d(in_channels*initial_step, out_channels).to(device) + model = UNet1d(in_channels * initial_step, out_channels).to(device) elif dimensions == 5: - model = UNet2d(in_channels*initial_step, out_channels).to(device) + model = UNet2d(in_channels * initial_step, out_channels).to(device) elif dimensions == 6: - model = UNet3d(in_channels*initial_step, out_channels).to(device) - if training_type in ['single']: + model = UNet3d(in_channels * initial_step, out_channels).to(device) + if training_type in ["single"]: if dimensions == 4: model = UNet1d(in_channels, out_channels).to(device) elif dimensions == 5: model = UNet2d(in_channels, out_channels).to(device) elif dimensions == 6: model = UNet3d(in_channels, out_channels).to(device) - + # Set maximum time step of the data to train - if t_train > _data.shape[-2]: - t_train = _data.shape[-2] + t_train = min(t_train, _data.shape[-2]) # Set maximum of unrolled time step for the pushforward trick if t_train - unroll_step < 1: unroll_step = t_train - 1 - if training_type in ['autoregressive']: + if training_type in ["autoregressive"]: if ar_mode: if pushforward: - model_name = model_name + '-PF-' + str(unroll_step) + model_name = model_name + "-PF-" + str(unroll_step) if not pushforward: unroll_step = _data.shape[-2] - model_name = model_name + '-AR' + model_name = model_name + "-AR" else: - model_name = model_name + '-1-step' - + model_name = model_name + "-1-step" + model_path = model_name + ".pt" - + total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) - print(f'Total parameters = {total_params}') - - optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4) - scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma) - + msg = f"Total parameters = {total_params}" + logger.info(msg) + + optimizer = torch.optim.Adam( + model.parameters(), lr=learning_rate, weight_decay=1e-4 + ) + scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, step_size=scheduler_step, gamma=scheduler_gamma + ) + loss_fn = nn.MSELoss(reduction="mean") - loss_val_min = np.infty - + loss_val_min = np.inf + start_epoch = 0 if not if_training: checkpoint = torch.load(model_path, map_location=device) - model.load_state_dict(checkpoint['model_state_dict']) + model.load_state_dict(checkpoint["model_state_dict"]) model.to(device) model.eval() - Lx, Ly, Lz = 1., 1., 1. - errs = metrics(val_loader, model, Lx, Ly, Lz, plot, channel_plot, - model_name, x_min, x_max, y_min, y_max, - t_min, t_max, mode='Unet', initial_step=initial_step) - pickle.dump(errs, open(model_name+'.pickle', "wb")) - + Lx, Ly, Lz = 1.0, 1.0, 1.0 + errs = metrics( + val_loader, + model, + Lx, + Ly, + Lz, + plot, + channel_plot, + model_name, + x_min, + x_max, + y_min, + y_max, + t_min, + t_max, + mode="Unet", + initial_step=initial_step, + ) + pickle.dump(errs, Path.open(model_name + ".pickle", "wb")) + return # If desired, restore the network by loading the weights saved in the .pt # file if continue_training: - print('Restoring model (that is the network\'s weights) from file...') + msg = "Restoring model (that is the network's weights) from file..." + logger.info(msg) checkpoint = torch.load(model_path, map_location=device) - model.load_state_dict(checkpoint['model_state_dict']) + model.load_state_dict(checkpoint["model_state_dict"]) model.to(device) model.train() - + # Load optimizer state dict - optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.to(device) - - start_epoch = checkpoint['epoch'] - loss_val_min = checkpoint['loss'] - print('start training...') - + start_epoch = checkpoint["epoch"] + loss_val_min = checkpoint["loss"] + + msg = "start training..." + logger.info(msg) + if ar_mode: - for ep in range(start_epoch, epochs): model.train() t1 = default_timer() train_l2_step = 0 train_l2_full = 0 - + for xx, yy in train_loader: loss = 0 - + # xx: input tensor (first few time steps) [b, x1, ..., xd, t_init, v] # yy: target tensor [b, x1, ..., xd, t, v] # grid: meshgrid [b, x1, ..., xd, dims] - xx = xx.to(device) - yy = yy.to(device) - - if training_type in ['autoregressive']: + xx_tensor = xx.to(device) + yy_tensor = yy.to(device) + if training_type in ["autoregressive"]: # Initialize the prediction tensor pred = yy[..., :initial_step, :] - + # Extract shape of the input tensor for reshaping (i.e. stacking the # time and channels dimension together) - inp_shape = list(xx.shape) + inp_shape = list(xx_tensor.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - + # Autoregressive loop for t in range(initial_step, t_train): - - if t < t_train-unroll_step: + if t < t_train - unroll_step: with torch.no_grad(): # Reshape input tensor into [b, x1, ..., xd, t_init*v] - inp = xx.reshape(inp_shape) + inp = xx_tensor.reshape(inp_shape) temp_shape = [0, -1] - temp_shape.extend([i for i in range(1,len(inp.shape)-1)]) + temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) - + # Extract target at current time step - y = yy[..., t:t+1, :] - + y = yy_tensor[..., t : t + 1, :] + # Model run temp_shape = [0] - temp_shape.extend([i for i in range(2,len(inp.shape))]) + temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) im = model(inp).permute(temp_shape).unsqueeze(-2) - + # Concatenate the prediction at current time step into the # prediction tensor pred = torch.cat((pred, im), -2) - + # Concatenate the prediction at the current time step to be used # as input for the next time step - xx = torch.cat((xx[..., 1:, :], im), dim=-2) - + xx_tensor = torch.cat( + (xx_tensor[..., 1:, :], im), dim=-2 + ) + else: # Reshape input tensor into [b, x1, ..., xd, t_init*v] - inp = xx.reshape(inp_shape) + inp = xx_tensor.reshape(inp_shape) temp_shape = [0, -1] - temp_shape.extend([i for i in range(1,len(inp.shape)-1)]) + temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) - + # Extract target at current time step - y = yy[..., t:t+1, :] - + y = yy_tensor[..., t : t + 1, :] + # Model run temp_shape = [0] - temp_shape.extend([i for i in range(2,len(inp.shape))]) + temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) im = model(inp).permute(temp_shape).unsqueeze(-2) - + # Loss calculation - loss += loss_fn(im.reshape(batch_size, -1), y.reshape(batch_size, -1)) - + loss += loss_fn( + im.reshape(batch_size, -1), y.reshape(batch_size, -1) + ) + # Concatenate the prediction at current time step into the # prediction tensor pred = torch.cat((pred, im), -2) - + # Concatenate the prediction at the current time step to be used # as input for the next time step - xx = torch.cat((xx[..., 1:, :], im), dim=-2) - + xx_tensor = torch.cat((xx_tensor[..., 1:, :], im), dim=-2) + train_l2_step += loss.item() - l2_full = loss_fn(pred.reshape(batch_size, -1), yy.reshape(batch_size, -1)) + _batch = yy_tensor.size(0) + _yy = yy_tensor[..., :t_train, :] + l2_full = loss_fn(pred.reshape(_batch, -1), _yy.reshape(_batch, -1)) train_l2_full += l2_full.item() - + optimizer.zero_grad() loss.backward() optimizer.step() - if training_type in ['single']: - x = xx[..., 0 , :] - y = yy[..., t_train:t_train+1 , :] + if training_type in ["single"]: + x = xx[..., 0, :] + y = yy[..., t_train - 1 : t_train, :] pred = model(x.permute([0, 2, 1])).permute([0, 2, 1]) _batch = yy.size(0) loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1)) - + train_l2_step += loss.item() train_l2_full += loss.item() - + optimizer.zero_grad() loss.backward() optimizer.step() - if ep % model_update == 0: val_l2_step = 0 val_l2_full = 0 with torch.no_grad(): for xx, yy in val_loader: loss = 0 - xx = xx.to(device) - yy = yy.to(device) - - if training_type in ['autoregressive']: - pred = yy[..., :initial_step, :] + xx_tensor = xx.to(device) + yy_tensor = yy.to(device) + + if training_type in ["autoregressive"]: + pred = yy_tensor[..., :initial_step, :] inp_shape = list(xx.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - + for t in range(initial_step, t_train): - inp = xx.reshape(inp_shape) + inp = xx_tensor.reshape(inp_shape) temp_shape = [0, -1] - temp_shape.extend([i for i in range(1,len(inp.shape)-1)]) + temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) - y = yy[..., t:t+1, :] + y = yy_tensor[..., t : t + 1, :] temp_shape = [0] - temp_shape.extend([i for i in range(2,len(inp.shape))]) + temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) im = model(inp).permute(temp_shape).unsqueeze(-2) - loss += loss_fn(im.reshape(batch_size, -1), y.reshape(batch_size, -1)) - + loss += loss_fn( + im.reshape(batch_size, -1), + y.reshape(batch_size, -1), + ) + pred = torch.cat((pred, im), -2) - - xx = torch.cat((xx[..., 1:, :], im), dim=-2) - + + xx_tensor = torch.cat( + (xx_tensor[..., 1:, :], im), dim=-2 + ) + val_l2_step += loss.item() - val_l2_full += loss_fn(pred.reshape(batch_size, -1), yy.reshape(batch_size, -1)).item() - - if training_type in ['single']: - x = xx[..., 0 , :] - y = yy[..., t_train:t_train+1 , :] + _batch = yy.size(0) + _pred = pred[..., initial_step:t_train, :] + _yy = yy_tensor[..., initial_step:t_train, :] + val_l2_full += loss_fn( + _pred.reshape(_batch, -1), _yy.reshape(_batch, -1) + ).item() + + if training_type in ["single"]: + x = xx[..., 0, :] + y = yy[..., t_train - 1 : t_train, :] pred = model(x.permute([0, 2, 1])).permute([0, 2, 1]) _batch = yy.size(0) loss += loss_fn(pred.reshape(_batch, -1), y.reshape(_batch, -1)) - + val_l2_step += loss.item() val_l2_full += loss.item() - if val_l2_full < loss_val_min: + if val_l2_full < loss_val_min: loss_val_min = val_l2_full - torch.save({ - 'epoch': ep, - 'model_state_dict': model.state_dict(), - 'optimizer_state_dict': optimizer.state_dict(), - 'loss': loss_val_min - }, model_path) - + torch.save( + { + "epoch": ep, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "loss": loss_val_min, + }, + model_path, + ) + t2 = default_timer() scheduler.step() - print('epoch: {0}, loss: {1:.5f}, t2-t1: {2:.5f}, trainL2: {3:.5f}, testL2: {4:.5f}'\ - .format(ep, loss.item(), t2 - t1, train_l2_step, val_l2_step)) + msg = f"epoch: {ep}, loss: {loss.item():.5f}, t2-t1: {t2 - t1:.5f}, trainL2: {train_l2_step:.5f}, testL2: {val_l2_step:.5f}" + logger.info(msg) else: for ep in range(start_epoch, epochs): @@ -359,110 +403,125 @@ def run_training(if_training, t1 = default_timer() train_l2_step = 0 train_l2_full = 0 - + for xx, yy in train_loader: loss = 0 - + # xx: input tensor (first few time steps) [b, x1, ..., xd, t_init, v] # yy: target tensor [b, x1, ..., xd, t, v] - xx = xx.to(device) - yy = yy.to(device) - + xx_tensor = xx.to(device) + yy_tensor = yy.to(device) + # Initialize the prediction tensor - pred = yy[..., :initial_step, :] - + pred = yy_tensor[..., :initial_step, :] + # Extract shape of the input tensor for reshaping (i.e. stacking the # time and channels dimension together) - inp_shape = list(xx.shape) + inp_shape = list(xx_tensor.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - + # Autoregressive loop for t in range(initial_step, t_train): - # Reshape input tensor into [b, x1, ..., xd, t_init*v] - inp = yy[..., t-initial_step:t, :].reshape(inp_shape) + inp = yy_tensor[..., t - initial_step : t, :].reshape(inp_shape) temp_shape = [0, -1] - temp_shape.extend([i for i in range(1,len(inp.shape)-1)]) + temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) inp = torch.normal(inp, 0.001) - + # Extract target at current time step - y = yy[..., t:t+1, :] - + y = yy_tensor[..., t : t + 1, :] + # Model run temp_shape = [0] - temp_shape.extend([i for i in range(2,len(inp.shape))]) + temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) im = model(inp).permute(temp_shape).unsqueeze(-2) - + # Loss calculation - loss += loss_fn(im.reshape(batch_size, -1), y.reshape(batch_size, -1)) - + loss += loss_fn( + im.reshape(batch_size, -1), y.reshape(batch_size, -1) + ) + # Concatenate the prediction at current time step into the # prediction tensor pred = torch.cat((pred, im), -2) - + # Concatenate the prediction at the current time step to be used # as input for the next time step # xx = torch.cat((xx[..., 1:, :], im), dim=-2) - + train_l2_step += loss.item() - l2_full = loss_fn(pred.reshape(batch_size, -1), yy.reshape(batch_size, -1)) + _batch = yy.size(0) + _yy = yy_tensor[..., :t_train, :] # if t_train is not -1 + l2_full = loss_fn(pred.reshape(_batch, -1), _yy.reshape(_batch, -1)) train_l2_full += l2_full.item() - + optimizer.zero_grad() loss.backward() optimizer.step() - + if ep % model_update == 0 or ep == epochs: val_l2_step = 0 val_l2_full = 0 with torch.no_grad(): for xx, yy in val_loader: loss = 0 - xx = xx.to(device) - yy = yy.to(device) - - pred = yy[..., :initial_step, :] - inp_shape = list(xx.shape) + xx_tensor = xx.to(device) + yy_tensor = yy.to(device) + + pred = yy_tensor[..., :initial_step, :] + inp_shape = list(xx_tensor.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - + for t in range(initial_step, t_train): - inp = yy[..., t-initial_step:t, :].reshape(inp_shape) + inp = yy_tensor[..., t - initial_step : t, :].reshape( + inp_shape + ) temp_shape = [0, -1] - temp_shape.extend([i for i in range(1,len(inp.shape)-1)]) + temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) - y = yy[..., t:t+1, :] + y = yy_tensor[..., t : t + 1, :] temp_shape = [0] - temp_shape.extend([i for i in range(2,len(inp.shape))]) + temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) im = model(inp).permute(temp_shape).unsqueeze(-2) - loss += loss_fn(im.reshape(batch_size, -1), y.reshape(batch_size, -1)) - + loss += loss_fn( + im.reshape(batch_size, -1), y.reshape(batch_size, -1) + ) + pred = torch.cat((pred, im), -2) - + val_l2_step += loss.item() - val_l2_full += loss_fn(pred.reshape(batch_size, -1), yy.reshape(batch_size, -1)).item() - - if val_l2_full < loss_val_min: + _batch = yy.size(0) + _pred = pred[..., initial_step:t_train, :] + # if t_train is not -1 + _yy = yy_tensor[..., initial_step:t_train, :] + val_l2_full += loss_fn( + _pred.reshape(_batch, -1), _yy.reshape(_batch, -1) + ).item() + + if val_l2_full < loss_val_min: loss_val_min = val_l2_full - torch.save({ - 'epoch': ep, - 'model_state_dict': model.state_dict(), - 'optimizer_state_dict': optimizer.state_dict(), - 'loss': loss_val_min - }, model_path) - - + torch.save( + { + "epoch": ep, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "loss": loss_val_min, + }, + model_path, + ) + t2 = default_timer() scheduler.step() - print('epoch: {0}, loss: {1:.5f}, t2-t1: {2:.5f}, trainL2: {3:.5f}, testL2: {4:.5f}'\ - .format(ep, loss.item(), t2 - t1, train_l2_step, val_l2_step)) + msg = f"epoch: {ep}, loss: {loss.item():.5f}, t2-t1: {t2 - t1:.5f}, trainL2: {train_l2_step:.5f}, testL2: {val_l2_step:.5f}" + logger.info(msg) if __name__ == "__main__": - run_training() - print("Done.") \ No newline at end of file + msg = "Done." + logger.info(msg) diff --git a/pdebench/models/unet/unet.py b/pdebench/models/unet/unet.py index 7bd828b..38e6a7a 100644 --- a/pdebench/models/unet/unet.py +++ b/pdebench/models/unet/unet.py @@ -1,5 +1,3 @@ -#!/usr/bin/env python3 - """ U-Net. Implementation taken and modified from https://github.com/mateuszbuda/brain-segmentation-pytorch @@ -18,16 +16,18 @@ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. """ + +from __future__ import annotations + from collections import OrderedDict import torch -import torch.nn as nn +from torch import nn class UNet1d(nn.Module): - def __init__(self, in_channels=3, out_channels=1, init_features=32): - super(UNet1d, self).__init__() + super().__init__() features = init_features self.encoder1 = UNet1d._block(in_channels, features, name="enc1") @@ -119,9 +119,8 @@ def _block(in_channels, features, name): class UNet2d(nn.Module): - def __init__(self, in_channels=3, out_channels=1, init_features=32): - super(UNet2d, self).__init__() + super().__init__() features = init_features self.encoder1 = UNet2d._block(in_channels, features, name="enc1") @@ -210,12 +209,11 @@ def _block(in_channels, features, name): ] ) ) - -class UNet3d(nn.Module): +class UNet3d(nn.Module): def __init__(self, in_channels=3, out_channels=1, init_features=32): - super(UNet3d, self).__init__() + super().__init__() features = init_features self.encoder1 = UNet3d._block(in_channels, features, name="enc1") @@ -303,4 +301,4 @@ def _block(in_channels, features, name): (name + "tanh2", nn.Tanh()), ] ) - ) \ No newline at end of file + ) diff --git a/pdebench/models/unet/utils.py b/pdebench/models/unet/utils.py index 669035d..a272f19 100644 --- a/pdebench/models/unet/utils.py +++ b/pdebench/models/unet/utils.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- """ @@ -148,126 +147,252 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -import torch -from torch.utils.data import Dataset -from torch.utils.data import DataLoader -import os -import glob +from __future__ import annotations + +import math as mt +from pathlib import Path + import h5py import numpy as np -import math as mt +import torch +from torch.utils.data import Dataset + class UNetDatasetSingle(Dataset): - def __init__(self, filename, - initial_step=10, - saved_folder='../data/', - reduced_resolution=1, - reduced_resolution_t=1, - reduced_batch=1, - if_test=False, - test_ratio=0.1, - num_samples_max = -1): + def __init__( + self, + filename, + initial_step=10, + saved_folder="../data/", + reduced_resolution=1, + reduced_resolution_t=1, + reduced_batch=1, + if_test=False, + test_ratio=0.1, + num_samples_max=-1, + ): """ - + :param filename: filename that contains the dataset :type filename: STR :param filenum: array containing indices of filename included in the dataset :type filenum: ARRAY """ - + # Define path to files - root_path = os.path.abspath(saved_folder + filename) - assert filename[-2:] != 'h5', 'HDF5 data is assumed!!' - - with h5py.File(root_path, 'r') as f: + root_path = Path(saved_folder + filename).resolve() + assert filename[-2:] != "h5", "HDF5 data is assumed!!" + + with h5py.File(root_path, "r") as f: keys = list(f.keys()) keys.sort() - if 'tensor' not in keys: - _data = np.array(f['density'], dtype=np.float32) # batch, time, x,... + if "tensor" not in keys: + _data = np.array(f["density"], dtype=np.float32) # batch, time, x,... idx_cfd = _data.shape - if len(idx_cfd)==3: # 1D - self.data = np.zeros([idx_cfd[0]//reduced_batch, - idx_cfd[2]//reduced_resolution, - mt.ceil(idx_cfd[1]/reduced_resolution_t), - 3], - dtype=np.float32) - #density - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] + if len(idx_cfd) == 3: # 1D + self.data = np.zeros( + [ + idx_cfd[0] // reduced_batch, + idx_cfd[2] // reduced_resolution, + mt.ceil(idx_cfd[1] / reduced_resolution_t), + 3, + ], + dtype=np.float32, + ) + # density + _data = _data[ + ::reduced_batch, ::reduced_resolution_t, ::reduced_resolution + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data[:, :, :], (0, 2, 1)) - self.data[...,0] = _data # batch, x, t, ch + self.data[..., 0] = _data # batch, x, t, ch # pressure - _data = np.array(f['pressure'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] + _data = np.array( + f["pressure"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, ::reduced_resolution_t, ::reduced_resolution + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data[:, :, :], (0, 2, 1)) - self.data[...,1] = _data # batch, x, t, ch + self.data[..., 1] = _data # batch, x, t, ch # Vx - _data = np.array(f['Vx'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] + _data = np.array(f["Vx"], dtype=np.float32) # batch, time, x,... + _data = _data[ + ::reduced_batch, ::reduced_resolution_t, ::reduced_resolution + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data[:, :, :], (0, 2, 1)) - self.data[...,2] = _data # batch, x, t, ch - - if len(idx_cfd)==4: # 2D - self.data = np.zeros([idx_cfd[0]//reduced_batch, - idx_cfd[2]//reduced_resolution, - idx_cfd[3]//reduced_resolution, - mt.ceil(idx_cfd[1]/reduced_resolution_t), - 4], - dtype=np.float32) + self.data[..., 2] = _data # batch, x, t, ch + + if len(idx_cfd) == 4: # 2D + self.data = np.zeros( + [ + idx_cfd[0] // reduced_batch, + idx_cfd[2] // reduced_resolution, + idx_cfd[3] // reduced_resolution, + mt.ceil(idx_cfd[1] / reduced_resolution_t), + 4, + ], + dtype=np.float32, + ) # density - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,0] = _data # batch, x, t, ch + self.data[..., 0] = _data # batch, x, t, ch # pressure - _data = np.array(f['pressure'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] + _data = np.array( + f["pressure"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,1] = _data # batch, x, t, ch + self.data[..., 1] = _data # batch, x, t, ch # Vx - _data = np.array(f['Vx'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] + _data = np.array(f["Vx"], dtype=np.float32) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,2] = _data # batch, x, t, ch + self.data[..., 2] = _data # batch, x, t, ch # Vy - _data = np.array(f['Vy'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution,::reduced_resolution] + _data = np.array(f["Vy"], dtype=np.float32) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data, (0, 2, 3, 1)) - self.data[...,3] = _data # batch, x, t, ch + self.data[..., 3] = _data # batch, x, t, ch + + if len(idx_cfd) == 5: # 3D + self.data = np.zeros( + [ + idx_cfd[0] // reduced_batch, + idx_cfd[2] // reduced_resolution, + idx_cfd[3] // reduced_resolution, + idx_cfd[4] // reduced_resolution, + mt.ceil(idx_cfd[1] / reduced_resolution_t), + 5, + ], + dtype=np.float32, + ) + # density + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 0] = _data # batch, x, t, ch + # pressure + _data = np.array( + f["pressure"], dtype=np.float32 + ) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 1] = _data # batch, x, t, ch + # Vx + _data = np.array(f["Vx"], dtype=np.float32) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 2] = _data # batch, x, t, ch + # Vy + _data = np.array(f["Vy"], dtype=np.float32) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 3] = _data # batch, x, t, ch + # Vz + _data = np.array(f["Vz"], dtype=np.float32) # batch, time, x,... + _data = _data[ + ::reduced_batch, + ::reduced_resolution_t, + ::reduced_resolution, + ::reduced_resolution, + ::reduced_resolution, + ] + ## convert to [x1, ..., xd, t, v] + _data = np.transpose(_data, (0, 2, 3, 4, 1)) + self.data[..., 4] = _data # batch, x, t, ch else: # scalar equations ## data dim = [t, x1, ..., xd, v] - _data = np.array(f['tensor'], dtype=np.float32) # batch, time, x,... + _data = np.array(f["tensor"], dtype=np.float32) # batch, time, x,... if len(_data.shape) == 3: # 1D - _data = _data[::reduced_batch,::reduced_resolution_t,::reduced_resolution] + _data = _data[ + ::reduced_batch, ::reduced_resolution_t, ::reduced_resolution + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data[:, :, :], (0, 2, 1)) self.data = _data[:, :, :, None] # batch, x, t, ch if len(_data.shape) == 4: # 2D Darcy flow # u: label - _data = _data[::reduced_batch,:,::reduced_resolution,::reduced_resolution] + _data = _data[ + ::reduced_batch, :, ::reduced_resolution, ::reduced_resolution + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1)) - #if _data.shape[-1]==1: # if nt==1 + # if _data.shape[-1]==1: # if nt==1 # _data = np.tile(_data, (1, 1, 1, 2)) self.data = _data # nu: input - _data = np.array(f['nu'], dtype=np.float32) # batch, time, x,... - _data = _data[::reduced_batch, None,::reduced_resolution,::reduced_resolution] + _data = np.array(f["nu"], dtype=np.float32) # batch, time, x,... + _data = _data[ + ::reduced_batch, + None, + ::reduced_resolution, + ::reduced_resolution, + ] ## convert to [x1, ..., xd, t, v] _data = np.transpose(_data[:, :, :, :], (0, 2, 3, 1)) self.data = np.concatenate([_data, self.data], axis=-1) self.data = self.data[:, :, :, :, None] # batch, x, y, t, ch - if num_samples_max>0: - num_samples_max = min(num_samples_max,self.data.shape[0]) + if num_samples_max > 0: + num_samples_max = min(num_samples_max, self.data.shape[0]) else: num_samples_max = self.data.shape[0] @@ -279,29 +404,27 @@ def __init__(self, filename, # Time steps used as initial conditions self.initial_step = initial_step - - self.data = torch.tensor(self.data) + self.data = torch.tensor(self.data) def __len__(self): return len(self.data) - + def __getitem__(self, idx): - - return self.data[idx,...,:self.initial_step,:], self.data[idx] - + return self.data[idx, ..., : self.initial_step, :], self.data[idx] + class UNetDatasetMult(Dataset): - def __init__(self, filename, - initial_step=10, - saved_folder='../data/', - reduced_resolution=1, - reduced_resolution_t=1, - reduced_batch=1, - if_test=False, test_ratio=0.1 - ): + def __init__( + self, + filename, + initial_step=10, + saved_folder="../data/", + if_test=False, + test_ratio=0.1, + ): """ - + :param filename: filename that contains the dataset :type filename: STR :param filenum: array containing indices of filename included in the dataset @@ -310,39 +433,38 @@ def __init__(self, filename, :type initial_step: INT, optional """ - + # Define path to files - self.file_path = os.path.abspath(saved_folder + filename + ".h5") - + self.file_path = Path(saved_folder + filename + ".h5").resolve() + # Extract list of seeds - with h5py.File(self.file_path, 'r') as h5_file: + with h5py.File(self.file_path, "r") as h5_file: data_list = sorted(h5_file.keys()) - test_idx = int(len(data_list) * (1-test_ratio)) + test_idx = int(len(data_list) * (1 - test_ratio)) if if_test: self.data_list = np.array(data_list[test_idx:]) else: self.data_list = np.array(data_list[:test_idx]) - + # Time steps used as initial conditions self.initial_step = initial_step def __len__(self): return len(self.data_list) - + def __getitem__(self, idx): - # Open file and read data - with h5py.File(self.file_path, 'r') as h5_file: + with h5py.File(self.file_path, "r") as h5_file: seed_group = h5_file[self.data_list[idx]] - + # data dim = [t, x1, ..., xd, v] - data = np.array(seed_group["data"], dtype='f') + data = np.array(seed_group["data"], dtype="f") data = torch.tensor(data, dtype=torch.float) - + # convert to [x1, ..., xd, t, v] - permute_idx = list(range(1,len(data.shape)-1)) - permute_idx.extend(list([0, -1])) + permute_idx = list(range(1, len(data.shape) - 1)) + permute_idx.extend([0, -1]) data = data.permute(permute_idx) - - return data[...,:self.initial_step,:], data \ No newline at end of file + + return data[..., : self.initial_step, :], data diff --git a/pdebench_examples.PNG b/pdebench_examples.PNG new file mode 100644 index 0000000..85f659e Binary files /dev/null and b/pdebench_examples.PNG differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3b20d32 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,208 @@ +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" + +[project] +requires-python = ">=3.9,<3.11" +dynamic = ["readme"] +name = "pdebench" +version = "0.1.0" +description = "PDEBench: An Extensive Benchmark for Scientific Machine Learning" +authors = [ + {name = "Makoto Takamoto", email = "Makoto.Takamoto@neclab.eu"}, + {name = "Timothy Praditia", email = "timothy.praditia@iws.uni-stuttgart.de"}, + {name = "Raphael Leiteritz", email = "raphael.leiteritz@ipvs.uni-stuttgart.de"}, + {name = "Dan MacKinlay"}, + {name = "Francesco Alesiani"}, + {name = "Dirk Pflรผger"}, + {name = "Mathias Niepert"}, +] +license = {file = "LICENSE.txt"} +dependencies = [ + "scipy", + "numpy<2", + "matplotlib", + "h5py", + "pandas", + "python-dotenv", + "hydra-core", + "torch~=1.13.0", + "torchvision~=0.14.1", + "deepxde~=1.1.3", + "pyro-ppl", + "tqdm", +] + +[project.optional-dependencies] +datagen310 = [ + "six", + "clawpack@git+https://github.com/clawpack/clawpack.git@d619d6835ce128a0421aa52d70d2a6c9d9d1ce93", + "dash", + "phiflow", + "imageio", + "einops", + "torch @ https://download.pytorch.org/whl/cu117/torch-1.13.1%2Bcu117-cp310-cp310-linux_x86_64.whl", + "torchvision @ https://download.pytorch.org/whl/cu117/torchvision-0.14.1%2Bcu117-cp310-cp310-linux_x86_64.whl", + "jax==0.4.11", + "jaxlib @ https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.11+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl", +] +datagen39 = [ + "six", + "clawpack@git+https://github.com/clawpack/clawpack.git@d619d6835ce128a0421aa52d70d2a6c9d9d1ce93", + "dash", + "phiflow", + "imageio", + "einops", + "torch @ https://download.pytorch.org/whl/cu117/torch-1.13.1%2Bcu117-cp39-cp39-linux_x86_64.whl", + "torchvision @ https://download.pytorch.org/whl/cu117/torchvision-0.14.1%2Bcu117-cp39-cp39-linux_x86_64.whl", + "jax==0.4.11", + "jaxlib @ https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.11+cuda11.cudnn86-cp39-cp39-manylinux2014_x86_64.whl" +] + +test = [ + "pytest >=6", + "pytest-cov >=3", + "jax", # cpu only + "nox" +] + +docs = [ + "sphinx>=7.0", + "myst_parser>=0.13", + "sphinx_copybutton", + "sphinx_autodoc_typehints", + "furo>=2023.08.17" +] + +dev = ["anybadge", + "ruff", + "pytest", + "pytest-coverage", + "pytest-mypy", + "hatchling", + "nox", + "pre-commit"] + +[project.urls] +Homepage = "https://github.com/pdebench/PDEBenchm" +Documentation = "https://github.com/pdebench/PDEBench" +Repository = "https://github.com/pdebench/PDEBench" + +[tool.hatch] +version.source = "vcs" +build.hooks.vcs.version-file = "pdebench/_version.py" +build.include = [ + "pdebench" +] +metadata.allow-direct-references = true + +[tool.hatch.envs.default] +features = ["test"] +scripts.test = "pytest {args}" + + +[tool.pytest.ini_options] +minversion = "6.0" +addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"] +xfail_strict = true +filterwarnings = [ + "error", +] +log_cli_level = "INFO" +testpaths = [ + "tests", +] + + +[tool.coverage] +run.source = ["pdebench"] +report.exclude_also = [ + '\.\.\.', + 'if typing.TYPE_CHECKING:', +] + +[tool.mypy] +files = ["pdebench", "tests"] +python_version = "3.8" +warn_unused_configs = true +strict = true +enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"] +warn_unreachable = true +disallow_untyped_defs = false +disallow_incomplete_defs = false +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "pdebench.*" +disallow_untyped_defs = true +disallow_incomplete_defs = true + +[tool.ruff] +src = ["pdebench"] + +[tool.ruff.format] +exclude = ["*.pyi"] + +[tool.ruff.lint] +extend-select = [ + "B", # flake8-bugbear + "I", # isort + "ARG", # flake8-unused-arguments + "C4", # flake8-comprehensions + "EM", # flake8-errmsg + "ICN", # flake8-import-conventions + "G", # flake8-logging-format + "PGH", # pygrep-hooks + "PIE", # flake8-pie + "PL", # pylint + "PT", # flake8-pytest-style + "PTH", # flake8-use-pathlib + "RET", # flake8-return + "RUF", # Ruff-specific + "SIM", # flake8-simplify + "T20", # flake8-print + "UP", # pyupgrade + "YTT", # flake8-2020 + "EXE", # flake8-executable + "NPY", # NumPy specific rules + "PD", # pandas-vet +] +ignore = [ + "PLR09", # Too many <...> + "PLR2004", # Magic value used in comparison + "ISC001", # Conflicts with formatter + "UP007", + "ARG001", # too many false positives + "ARG005", + "E731", # do not assign a lambda expression, use a def + "G004", + "PD008", # dataframe confused with jax array + "PD011", # dataframe confused with jax array +] +isort.required-imports = ["from __future__ import annotations"] +# Uncomment if using a _compat.typing backport +# typing-modules = ["pdebench._compat.typing"] + +[tool.ruff.lint.per-file-ignores] +"tests/**" = ["T20"] +"noxfile.py" = ["T20"] + + +[tool.pylint] +py-version = "3.8" +ignore-paths = [".*/_version.py"] +reports.output-format = "colorized" +similarities.ignore-imports = "yes" +messages_control.disable = [ + "design", + "fixme", + "line-too-long", + "missing-module-docstring", + "wrong-import-position", +] + +[tool.setuptools.dynamic] +readme = {file = ["README.md"], content-type = "text/markdown"} + +[project.scripts] +velocity2vorticity = "pdebench.data_gen.velocity2vorticity:convert_velocity" diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 0ff5383..0000000 --- a/requirements.txt +++ /dev/null @@ -1,31 +0,0 @@ -numpy~=1.22.3 # Basic mathematical library for array operations -scipy~=1.8.0 # Scientific library -matplotlib~=3.5.2 # Plotting library -h5py~=3.6.0 # Data storage format -pandas~=1.4.2 # Dataframe library -python-dotenv~=0.20.0 # Configuration library -hydra-core~=1.2.0 # Configuration library -hydra-joblib-launcher~=1.2.0 # parallel jobs locally -hydra-submitit-launcher~=1.1.6 # parallel jobs on HPC cluster -pytorch-lightning~=1.6.3 # only for LightningDataModule -omegaconf~=2.2.1 # Configuration library -clawpack~=5.6.1 # Finite Volume library -deepxde~=1.1.3 # Physics-informed ML library -torch~=1.11.0 # PyTorch library for ML applications -jax~=0.3.13 # Numerical library -jaxlib~=0.3.10 # Numerical library part 2 (CPU) -dash~=2.2.0 # dashboard visualisation for phiflow -phiflow~=2.0.3 # PDE simulator -pyDaRUS~=1.0.5 # DaRUS API package to upload and download dataset -pyDataverse@ git+https://github.com/JR-1991/pyDataverse.git@0fcfcd3fbc6bf1aec869899f715a51dca25e91be -imageio~=2.19.2 -tqdm~=4.64.0 -einops~=0.4.1 -tensorboard~=2.9.0 -tensorboard-plugin-profile~=2.8.0 - -# Optional, if you want to use jupyter -jupyter~=1.0.0 -ipykernel~=6.9.0 -jupyter-dash~=0.4.1 -nbstripout~=0.5.0 # trim outputs from jupyter notebooks \ No newline at end of file diff --git a/setup.py b/setup.py deleted file mode 100644 index cb2e6db..0000000 --- a/setup.py +++ /dev/null @@ -1,33 +0,0 @@ -from setuptools import setup - -with open("README.md", "r", encoding="utf-8") as fh: - long_description = fh.read() - -setup( - name='pdebench', - version='0.0.1', - install_requires=['numpy', - 'scipy', - 'matplotlib', - 'h5py', - 'pandas', - 'python-dotenv', - 'hydra-core', - 'omegaconf', - 'clawpack', - 'deepxde', - 'torch', - 'pyDaRUS' - ], - packages=['pdebench'], - author='Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflรผger, Mathias Niepert', - author_email='Makoto.Takamoto@neclab.eu, timothy.praditia@iws.uni-stuttgart.de', - url='https://github.com/pdebench/PDEBench', - zip_safe=False, - description='PDEBench: An Extensive Benchmark for Scientific Machine Learning', - long_description=long_description, - long_description_content_type="text/markdown", - classifiers=[ - "License :: see licence file", - ], -) \ No newline at end of file diff --git a/tests/test_vorticity.py b/tests/test_vorticity.py new file mode 100644 index 0000000..d941983 --- /dev/null +++ b/tests/test_vorticity.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import jax.numpy as jnp +import numpy as np +import pytest +from pdebench.data_gen.src.vorticity import ( + compute_spectral_vorticity_jnp, + compute_spectral_vorticity_np, +) + + +@pytest.fixture +def generate_random_spectral_velvor() -> tuple[np.ndarray, np.ndarray]: + """Generate random 5D velocity- and corresponding vorticity field + + :return: Velocity- and vorticity field + :rtype: tuple[np.ndarray, np.ndarray] + """ + generator = np.random.default_rng(seed=None) + vel = generator.uniform(size=(10, 16, 32, 32, 3)) + vx = vel[..., 0] + vy = vel[..., 1] + vz = vel[..., 2] + + fxy = np.fft.fft(vx, axis=2) + fxz = np.fft.fft(vx, axis=3) + fyx = np.fft.fft(vy, axis=1) + fyz = np.fft.fft(vy, axis=3) + fzx = np.fft.fft(vz, axis=1) + fzy = np.fft.fft(vz, axis=2) + + kappa_xy = 2.0 * np.pi * np.fft.fftfreq(vel.shape[2], 1.0 / vel.shape[2]) + kappa_xz = 2.0 * np.pi * np.fft.fftfreq(vel.shape[3], 1.0 / vel.shape[3]) + kappa_yx = 2.0 * np.pi * np.fft.fftfreq(vel.shape[1], 1.0 / vel.shape[1]) + kappa_yz = 2.0 * np.pi * np.fft.fftfreq(vel.shape[3], 1.0 / vel.shape[3]) + kappa_zx = 2.0 * np.pi * np.fft.fftfreq(vel.shape[1], 1.0 / vel.shape[1]) + kappa_zy = 2.0 * np.pi * np.fft.fftfreq(vel.shape[2], 1.0 / vel.shape[2]) + + vxy = np.fft.ifft(1j * kappa_xy[None, None, :, None] * fxy, axis=2).real + vyx = np.fft.ifft(1j * kappa_yx[None, :, None, None] * fyx, axis=1).real + vxz = np.fft.ifft(1j * kappa_xz[None, None, None, :] * fxz, axis=3).real + vzx = np.fft.ifft(1j * kappa_zx[None, :, None, None] * fzx, axis=1).real + vyz = np.fft.ifft(1j * kappa_yz[None, None, None, :] * fyz, axis=3).real + vzy = np.fft.ifft(1j * kappa_zy[None, None, :, None] * fzy, axis=2).real + + omegax = vzy - vyz + omegay = vxz - vzx + omegaz = vyx - vxy + + omega = np.concatenate( + [omegax[..., None], omegay[..., None], omegaz[..., None]], axis=-1 + ) + + return vel, omega + + +def test_vorticity_np(generate_random_spectral_velvor) -> None: + """Test approximated vorticity by spectral derivation""" + vel, vort = generate_random_spectral_velvor + dx = 1.0 / vel.shape[1] + dy = 1.0 / vel.shape[2] + dz = 1.0 / vel.shape[3] + + vort_np = compute_spectral_vorticity_np(vel, dx, dy, dz) + np.testing.assert_almost_equal(vort_np, vort) + + +def test_vorticity_jnp(generate_random_spectral_velvor) -> None: + """Test approximated vorticity by spectral derivation""" + vel, vort = generate_random_spectral_velvor + dx = 1.0 / vel.shape[1] + dy = 1.0 / vel.shape[2] + dz = 1.0 / vel.shape[3] + + vort_jnp = compute_spectral_vorticity_jnp(jnp.array(vel), dx, dy, dz) + np.testing.assert_almost_equal(np.array(vort_jnp), vort, decimal=4)