NUNO: A General Framework for Learning Parametric PDEs with Non-Uniform Data

TL;DR

NUNO introduces a framework for efficient learning of parametric PDEs from non-uniform data using domain decomposition, significantly improving accuracy and training speed across complex 2D and 3D problems.

Contribution

It presents a novel K-D tree-based domain decomposition method enabling neural operators to handle non-uniform data effectively, extending PDE learning to complex 3D geometries.

Findings

Reduced error rates by up to 60%

Training speeds improved by 2x to 30x

Successfully applied to complex 3D PDE problems

Abstract

The neural operator has emerged as a powerful tool in learning mappings between function spaces in PDEs. However, when faced with real-world physical data, which are often highly non-uniformly distributed, it is challenging to use mesh-based techniques such as the FFT. To address this, we introduce the Non-Uniform Neural Operator (NUNO), a comprehensive framework designed for efficient operator learning with non-uniform data. Leveraging a K-D tree-based domain decomposition, we transform non-uniform data into uniform grids while effectively controlling interpolation error, thereby paralleling the speed and accuracy of learning from non-uniform data. We conduct extensive experiments on 2D elasticity, (2+1)D channel flow, and a 3D multi-physics heatsink, which, to our knowledge, marks a novel exploration into 3D PDE problems with complex geometries. Our framework has reduced error rates by up to 60% and enhanced training speeds by 2x to 30x. The code is now available at https://github.com/thu-ml/NUNO.