EPINF-NeuFlow

The pytorch implementation of the paper "EPINF: Efficient Physics Informed Neural Reconstruction of Diverse Fluid Dynamics from Sparse Observations"

This repo is also a pypi package NeuFlow that can be installed via pip:

python -m pip install neuflow

Environment Setup

System Requirements

• System
  • Windows 11 (Fully Tested)
  • Ubuntu 24.04.2 (Fully Tested)
• Python: 3.13
• PyTorch: 2.7.1 + CUDA 12.8

Pip Packages

python -m pip install --upgrade pip setuptools wheel
python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
python -m pip install lightning lightning[extra] dearpygui tyro matplotlib av huggingface_hub wandb opencv-python phiflow

set TCNN_CUDA_ARCHITECTURES=86
python -m pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch

python -m pip install ./cuda_extensions/freqencoder ./cuda_extensions/gridencoder ./cuda_extensions/raymarching ./cuda_extensions/shencoder

For Windows user, open X64 Native Tools Command Prompt for VS 2022 then activate your python venv if any, and set TCNN_CUDA_ARCHITECTURES to your GPU architecture.

Optional Packages

triton-windows - Activate torch.compile for Windows

python -m pip install -U "triton-windows<3.3"

Wandb Logger

To use the Wandb logger, you need to set up your Wandb account and login. You can do this by running the following command:

wandb login

Datasets

We provide two datasets for training and testing the model:

• the NeRF-Synthetic is available at Hugging Face - nerf_synthetic.
• The PI-Neuflow dataset is available at Hugging Face - PI-NeuFlow.

ALL datasets are downloaded AUTOMATICALLY when running the training script.

Training

NGP Model

python -O app.py ngp --dataset.dataset=chair --epochs=500 --gpu_id=0
python -O app.py ngp --dataset.dataset=drums --epochs=500 --gpu_id=0
python -O app.py ngp --dataset.dataset=ficus --epochs=500 --gpu_id=0
python -O app.py ngp --dataset.dataset=hotdog --epochs=500 --gpu_id=0
python -O app.py ngp --dataset.dataset=lego --epochs=500 --gpu_id=0
python -O app.py ngp --dataset.dataset=materials --epochs=500 --gpu_id=0
python -O app.py ngp --dataset.dataset=mic --epochs=500 --gpu_id=0
python -O app.py ngp --dataset.dataset=ship --epochs=500 --gpu_id=0

EPINF Model

python -O app.py epinf --dataset.dataset=sphere --epochs=500 --gpu_id=0
python -O app.py epinf --dataset.dataset=game --epochs=500 --gpu_id=0
python -O app.py epinf --dataset.dataset=torch2 --epochs=500 --gpu_id=0
python -O app.py epinf --dataset.dataset=fireplace --epochs=500 --gpu_id=0