Installation

underPINN is pure Python. Install it in editable mode from the repository root so every example and the CLI resolve imports automatically.

Requirements

Package	Purpose
jax[cpu] >= 0.4.26	JIT compilation, autodiff, PRNG
flax >= 0.8.0	Neural network layers and parameter trees
optax >= 0.2.0	Adam, cosine decay, gradient clipping
numpy, scipy, matplotlib	Numerics, exact solutions, plotting
shapely >= 2.0	Arbitrary polygon geometry support
pyyaml >= 6.0	YAML config loading and merging
pandas, trimesh, einops	Reporting utilities and operator-learning helpers

CPU / Development

pip install jax flax optax matplotlib scipy shapely trimesh pandas einops pyyaml

GPU (CUDA 12)

pip install -U "jax[cuda12]"
pip install -r requirements-gpu.txt

TPU (Colab / Cloud TPU VM)

pip install -r requirements-tpu.txt
pip install -e . --no-deps

Note

--no-deps is required so pip cannot silently replace the TPU-provisioned JAX build with the CPU pin from setup.py. underPINN auto-detects the TPU backend and forces full-float32 matmuls — the default bfloat16 MXU precision corrupts second-order PDE Hessians used throughout the framework.

From Source (recommended)

git clone https://github.com/Aeroscience-Computations-Analysis-Lab/underPINN.git
cd underPINN
pip install -e .

Verify your installation

python -c "import jax; print(jax.devices())"

Expected output on a GPU-enabled machine:

[CudaDevice(id=0)]

See also

Once installed, head to Quick Start to train your first PINN, or read GPU Memory Management to understand how underPINN manages VRAM on shared GPU nodes.