Installation

This page covers environment setup for NeuralPOM, from basic dependencies to multi-GPU training configuration.

System Requirements

Component	Minimum	Recommended
OS	Linux / macOS	Linux (Ubuntu 20.04+)
Python	3.10	3.10+
CUDA	11.8	12.1+
GPU Memory	24 GB	80 GB (A100)
RAM	64 GB	256 GB+

Basic Installation

1. Create a Conda Environment

conda create -n neuralpom python=3.10 -y
conda activate neuralpom

2. Install PyTorch

For CUDA 12.1:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

For CUDA 11.8:

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

3. Install NeuralPOM Dependencies

git clone https://github.com/ChiyodaMomo01/NeuralPOM.git
cd NeuralPOM
pip install -r requirements.txt

Core Dependencies

torch>=2.0.0
numpy>=1.24.0
scipy>=1.10.0
netCDF4>=1.6.0
tqdm>=4.65.0

For the implicit core with CG/PCG solver, no additional packages beyond PyTorch are required — the linear solvers are implemented natively in PyTorch.

Multi-GPU Setup (DDP)

NeuralPOM uses PyTorch Distributed Data Parallel (DDP) with the NCCL backend for multi-GPU training.

Environment Variables

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export NCCL_DEBUG=INFO
export OMP_NUM_THREADS=1

Launch Training

torchrun --nproc_per_node=8 \
  --master_port=29500 \
  scripts/training/train_npom.py \
  --config configs/train_decoder_only.json

Data Directory Structure

NeuralPOM expects the following data layout:

data/
├── double_gyre/
│   ├── train/
│   │   ├── member_00/
│   │   └── member_01/
│   └── test/
│       └── member_10/

Data paths and member IDs are configurable in the training JSON.

Verification

To verify your installation:

import torch
from neuralpom.hybrid.core_pool import build_initialized_core
from neuralpom.core_factory import resolve_core_class

# Create the production Double-gyre core.
core = build_initialized_core(
    resolve_core_class("implicit"),
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    im=90,
    jm=82,
)
print(f"Grid: {core.im}x{core.jm}x{core.kb}")
print(f"Fields: u={core.u.shape}, v={core.v.shape}, el={core.el.shape}")

Expected output:

Grid: 64x64x21
Fields: u=torch.Size([64, 65, 21]), v=torch.Size([65, 64, 21]), el=torch.Size([64, 64])