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}
}