spinn's Introduction

JAX implementation of Separable PINN

Project Page | Paper

Junwoo Cho*, Seungtae Nam*, Hyunmo Yang, Youngjoon Hong, Seok-Bae Yun, Eunbyun Park†
*Equal contribution, †Corresponding author.

Conference on Neural Information Processing Systems (NeurIPS 2023)
Spotlight presentation

Navier_Stokes.mp4

Architecture overview

SPINN consists of multiple MLPs, each of which takes an individual 1-dimensional coordinate as an input.
The output is constructed by a simple product and summation.

Environment Setup

0. If you're using Google Colab, just run the code

1. Install Docker and NVIDIA Container Toolkit

please follow the official document for installation.
if you already installed both of them, please skip this part.

2. Build the docker image

run the command below at "/your/path/to/SPINN".
don't forget to include the dot at the end.

docker build -t spinn_environment .

3. Run the docker image

run the command below at "/your/path/to/SPINN".

docker run -it -v $(pwd):/workspace -p 8888:8888 --gpus all --ipc host --name spinn spinn_environment:latest

4. (for demo code) Launch Jupyter and run the code

run the command below inside the container.

jupyter notebook --allow-root --ip 0.0.0.0 --port 8888

Training

Run the command below:

XLA_PYTHON_CLIENT_PREALLOCATE=false CUDA_VISIBLE_DEVICES=0 python <EQUATIONnd>.py

You can also use our configurations by running the script files:

bash ./scripts/<EQUATIONnd_MODEL>.sh

Navier-Stokes Reference Data

Find the original NS data from here: https://github.com/PredictiveIntelligenceLab/CausalPINNs/tree/main/data

Citation

@article{cho2023separable,
  title={Separable Physics-Informed Neural Networks},
  author={Cho, Junwoo and Nam, Seungtae and Yang, Hyunmo and Yun, Seok-Bae and Hong, Youngjoon and Park, Eunbyung},
  journal={Advances in Neural Information Processing Systems},
  year={2023}
}

Recommend Projects