underPINN

A modular, GPU-accelerated Physics-Informed Neural Network framework built on JAX + Flax + Optax

underPINN is a research-grade PINN engine that combines classical collocation-based PINNs with Finite Basis decomposition (FBPINN), attention-augmented networks, residual-based adaptive weighting/resampling, transfer learning (including windowed time-marching for long-horizon unsteady flows), shock capturing with learnable artificial viscosity, non-Newtonian (Carreau) blood rheology, neural operators (FNO, DeepONet, CViT), inverse problems, and a full restart/resume system — all JIT-compiled and differentiable via XLA on CPU, GPU, and TPU.

🚀 Get Started

Install underPINN and train your first PINN in five minutes.

Quick Start

🧪 Physics Examples

22 worked examples across 8 physics domains, from ODEs to 3-D turbulence.

Physics Examples

⚙️ Training System

TrainingConfig, callbacks, lax.scan fusion, and RAR-D adaptive resampling.

Training System

💾 Restart & Checkpoints

Fault-tolerant training that resumes exactly where it left off.

Restart / Resume System

Key numbers

22	Physics examples
5	PDE solver classes
19+	CLI-registered runners
500×	Less GPU dispatch overhead with lax.scan fusion
3-D	Unsteady Navier–Stokes support (pulsatile pipe flow)
Auto	Restart / resume on any interruption

Why underPINN?

Zero-boilerplate GPU memory management

JAX’s XLA allocator reserves ~90% of free VRAM the instant import jax runs. underPINN sets XLA_PYTHON_CLIENT_PREALLOCATE=false automatically in every entry point, so a 3-layer MLP uses ~200 MB instead of reserving 73 GB on an 80 GB A100. See GPU Memory Management for details.

YAML-driven, zero code changes

Every hyperparameter — network architecture, physics parameters, training schedule, loss weights — lives in a YAML file. Adding a brand-new physics case requires one script, one YAML file, and one line in the runner registry. See Command-Line Interface and Physics Examples.

Documentation contents

Getting Started

• Installation
  • Requirements
  • CPU / Development
  • GPU (CUDA 12)
  • TPU (Colab / Cloud TPU VM)
  • From Source (recommended)
  • Verify your installation
• GPU Memory Management
  • Why does nvidia-smi show 73 GB right after import?
  • underPINN disables this automatically
  • Manual control via environment variables
  • Programmatic override
  • Typical VRAM usage (preallocation disabled)
• Quick Start
  • 1. CLI — zero Python
  • 2. Python API
  • 3. YAML Config
  • 4. Checkpoint & Inference
  • 5. Transfer Learning
  • 6. GPU Acceleration — lax.scan + RAR-D
  • Running examples directly
• Command-Line Interface
  • Single run
  • Hyperparameter sweep
  • Inspect & list
  • Benchmark suite
  • Adding a new physics case

Core Systems

• Training System
  • TrainingConfig — full field reference
  • Network architectures
  • lax.scan XLA fusion
  • RAR-D adaptive collocation resampling
  • QR-DEIM-R adaptive collocation resampling
  • Shock capturing — artificial viscosity
  • Time-marching transfer learning
  • RBA — residual-based adaptive weighting
  • Gauss-Newton / natural-gradient training
  • Callbacks
  • Performance tips
• Restart / Resume System
  • How it works
  • Snapshot directory contents
  • Configuration
  • Verifying config changes before resuming
• Model Checkpointing & Inference
  • Save during training — ModelCheckpoint callback
  • Reload and predict on new inputs
  • Lower-level API
  • Post-processing utilities
• Transfer Learning
  • Parameter transfer (different ν / Re / diffusivity)
  • Temporal transfer (extended time horizon)
  • Time-marching transfer (windowed, long-horizon unsteady flows)
• Inverse Problems
  • How it works
• Neural Operators (PINO / DeepONet / CViT)
  • Architectures
  • Training loops
  • Operator-specific PDE residuals
  • Worked examples

Reference

• Physics Examples
  • Core PINN examples
  • Fluid dynamics
  • Compressible flow (shock capturing)
  • Transfer-learning examples
  • Neural operator examples
• PDE Library Reference
• Geometry Reference
  • Notable design details
  • Every geometry implements a common sampler interface
• API Reference
  • Core abstractions
  • PDE + Geometry convention
  • Package layout
  • Full repository tree
• Benchmark Suite
  • CLI usage
  • Programmatic usage
  • Outputs written to outputs/bench/
• Performance
  • Headline numbers
  • lax.scan — n_scan_steps reference
  • Tuning checklist
  • TPU
  • Cosine learning-rate decay
  • Training time reporting

Citing underPINN

If you use underPINN in research or publications, please cite:

@software{underPINN,
  author  = {Kumar Prashant, Senthilkumar Lohith, Ranjan Rajesh},
  title   = {underPINN: A Modular JAX Framework for Physics-Informed Neural Networks},
  year    = {2026},
  version = {v2608},
  url     = {https://github.com/Aeroscience-Computations-Analysis-Lab/underPINN.git}
}

underPINN is released under the GPL-3.0 License. See LICENSE.txt for the full text.