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/
├── seamount/
│   └── ...
└── box/
    └── ...

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

Verification

To verify your installation:

import torch
from neuralpom.cores.explicit.core import ExplicitCore

# Create a minimal core
core = ExplicitCore(im=64, jm=64, kb=21)
core.initialize()
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])