Abstract: Existing neural operator architectures face challenges when solving multiphysics problems with coupled partial differential equations (PDEs), due to complex geometries, interactions between physical variables, and the lack of large amounts of high-resolution training data. To address these issues, we propose Codomain Attention Neural Operator (CoDA-NO), which tokenizes functions along the codomain or channel space, enabling self-supervised learning or pretraining of multiple PDE systems. Specifically, we extend positional encoding, self-attention, and normalization layers to the function space. CoDA-NO can learn representations of different PDE systems with a single model. We evaluate CoDA-NO's potential as a backbone for learning multiphysics PDEs over multiple systems by considering few-shot learning settings. On complex downstream tasks with limited data, such as fluid flow simulations and fluid-structure interactions, we found CoDA-NO to outperform existing methods on the few-shot learning task by over $36$%. Paper Link

Left: Architecture of the Codomain Attention Neural Operator. Right: Mechanism of codomain attention.

At each CoDA-NO layer, the input function is tokenized codomain-wise to generate token functions. Each token function is passed through the K, Q, and V operators to get key, query, and value functions. The output function is calculated via an extension of the self-attention mechanism to the function space.

The fluid-solid interaction dataset is available at Dataset link.

The configurations for all the experiments are at config/ssl_ns_elastic.yaml.

To run the experiments, download the datasets, update the "input_mesh_location" and "data_location" in the config file, update the wandb cradentials and execute the following command

python main.py --config "config name" --ntrain N

--config: which configuration to use from the config file 'config/ssl_ns_elastic.yaml`.

--ntrain: Number of training data points.

For training CoDA-NO architecture on NS and NS+EW datasets (both pre-training and fine-tuning) please execute the following scrips:

codano_ns.sh
codano_nses.sh

For training the baseline, execute the following scripts

fno_baseline.sh 

gnn_baseline.sh

deeponet_baseline.sh

unet_baseline.sh

vit_baseline.sh

If you find this paper and code useful in your research, please consider citing:

@article{rahman2024pretraining, title={Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs}, author={Rahman, Md Ashiqur and George, Robert Joseph and Elleithy, Mogab and Leibovici, Daniel and Li, Zongyi and Bonev, Boris and White, Colin and Berner, Julius and Yeh, Raymond A and Kossaifi, Jean and others}, journal={arXiv preprint arXiv:2403.12553}, year={2024} }