Physics-Informed Neural Networks for Navier-Stokes Equations

A comprehensive implementation of Physics-Informed Neural Networks (PINNs) and Separable PINNs (SPINNs) for solving incompressible Navier-Stokes equations and other PDEs.

Overview

This repository contains two complementary implementations:

Method	Framework	Paper
PINN (NSFnets)	PyTorch	Jin et al., JCP 2021
SPINN	JAX/Flax	Cho et al., NeurIPS 2023

Key Features

Two Navier-Stokes Formulations

• Velocity-Pressure (VP) formulation
• Vorticity-Velocity (VV) formulation

Benchmark Problems

• Kovasznay flow (steady, 2D)
• Cylinder wake (unsteady, 2D)
• Taylor-Green vortex
• Ethier-Steinman flow (3D)

Comprehensive Experiments

• Architecture studies (network depth/width)
• Convergence analysis
• Dynamic weight scheduling
• Adaptive refinement (RAR)

Project Structure

physics-informed-neural-networks/
├── pinn/                    # PyTorch PINN implementation
│   ├── models/              # PINN, VPNSFnet, VVNSFnet
│   ├── utils/               # Training, sampling, losses
│   └── benchmarks/          # Analytical solutions
│
├── spinn/                   # JAX SPINN implementation
│   ├── networks/            # Separable PINN architectures
│   ├── utils/               # Training utilities
│   ├── poisson2d.py         # 2D Poisson equation
│   ├── helmholtz3d.py       # 3D Helmholtz equation
│   ├── navier_stokes3d.py   # 3D Navier-Stokes
│   └── ...
│
├── experiments/             # Runnable experiments
│   ├── kovasznay/           # Kovasznay flow experiments
│   ├── cylinder_wake/       # Cylinder flow experiments
│   └── ...
│
├── results/                 # Results and visualizations
├── docs/                    # Documentation
└── scripts/                 # Helper scripts

Installation

# Clone the repository
git clone https://github.com/yourusername/physics-informed-neural-networks.git
cd physics-informed-neural-networks

# Create virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # Linux/Mac
# or: venv\Scripts\activate  # Windows

# Install dependencies
pip install -r requirements.txt

GPU Support

PyTorch (PINN):

pip install torch --index-url https://download.pytorch.org/whl/cu118

JAX (SPINN):

pip install jax[cuda12_pip] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Quick Start

PINN - Kovasznay Flow

from pinn.models import VPNSFnet

# Create model (6 hidden layers, 50 neurons each)
model = VPNSFnet(hidden_layers=6, neurons_per_layer=50, nu=0.025)

# Set domain bounds
model.set_domain_bounds(lb=[0.0, -0.5], ub=[1.0, 0.5])

# Train (see experiments/kovasznay/ for full example)

SPINN - Poisson Equation

cd spinn
XLA_PYTHON_CLIENT_PREALLOCATE=false CUDA_VISIBLE_DEVICES=0 python poisson2d.py

Experiment Tracking

Experiments are tracked with Comet ML:

Experiment	Comet ML Dashboard
Kovasznay Flow	pinn-kovasznay-paper-02
SPINN Reproduction	spinn-reproduction
Cylinder Wake	pinn-cylinder-wake

To use Comet ML, copy .env.example to .env and add your API key.

Results Summary

Kovasznay Flow — PINN vs SPINN Comparison

PINN Results (PyTorch):

Architecture	Error u	Error v	Error p
2×50	0.0566 ± 0.0125	0.4318 ± 0.0924	0.2116 ± 0.0448
3×50	0.0297 ± 0.0046	0.2364 ± 0.0380	0.1057 ± 0.0096
4×50	0.0221 ± 0.0064	0.1662 ± 0.0425	0.0791 ± 0.0263
5×50	0.0180 ± 0.0051	0.1497 ± 0.0372	0.0600 ± 0.0130
6×50	0.0155 ± 0.0028	0.1136 ± 0.0230	0.0557 ± 0.0056

SPINN Results (JAX):

Architecture	Error u	Error v	Error p
10×100_r100_51×51	0.0195 ± 0.0040	0.1250 ± 0.0310	0.0650 ± 0.0150
10×250_r250_51×51	0.0210 ± 0.0055	0.1380 ± 0.0350	0.0720 ± 0.0180
7×100_r100_51×51	0.0175 ± 0.0035	0.1180 ± 0.0250	0.0580 ± 0.0120
4×50_r50_51×51	0.0280 ± 0.0060	0.1750 ± 0.0400	0.0850 ± 0.0200

Both PINN and SPINN achieve comparable accuracy on Kovasznay flow. PINN (6×50) slightly outperforms in error metrics, while SPINN offers advantages in higher-dimensional problems through separable architectures.

SPINN training curves for Kovasznay flow: loss evolution and error metrics across iterations

---

Cylinder Wake — NSFnet Vorticity Predictions

The cylinder wake problem demonstrates the capability of NSFnets to capture complex unsteady flow dynamics with vortex shedding.

VV-NSFnet: Vorticity field predictions vs reference solution at different time steps

VP-NSFnet: Velocity-Pressure formulation vorticity visualization

See results/kovasznay_pinn_spinn_comparison.md for detailed analysis.

Running Experiments

Single Experiment

cd experiments
python kovasznay/kovasznay_paper.py --arch 6x50 --device cuda

All Experiments (Parallel)

cd scripts
./RUN_NOW.sh

Monitor Progress

# View live training
tmux attach -t nsfnets

# Check logs
tail -f logs/parallel_all_goals.log

Documentation

• PROJECT_SUMMARY.md - Complete technical overview
• goals.md - Experiment descriptions (12 goals)
• TRAINING_GUIDE.md - Training best practices

Citations

If you use this code, please cite the original papers:

NSFnets (PINN):

@article{jin2021nsfnets,
  title={NSFnets (Navier-Stokes flow nets): Physics-informed neural networks 
         for the incompressible Navier-Stokes equations},
  author={Jin, Xiaowei and Cai, Shengze and Li, Hui and Karniadakis, George Em},
  journal={Journal of Computational Physics},
  volume={426},
  pages={109951},
  year={2021},
  publisher={Elsevier}
}

SPINN:

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

License

This project is released under the MIT License.

Acknowledgments

• SPINN implementation from SPINN

---

Author: Ivan Novosad Date: November 2024 - January 2025