Jan Hagnberger, Marimuthu Kalimuthu, Daniel Musekamp, Mathias Niepert
This repository contains the official PyTorch implementation of the VCNeF model from the paper "Vectorized Conditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations".
The VCNeF model requires and is tested with the following packages.
Please also see the requirements.txt file which contains all packages to run the provided examples.
The following example shows how to use the VCNeF model.
import torch
from vcnef.vcnef_1d import VCNeFModel as VCNeF1DModel
from vcnef.vcnef_2d import VCNeFModel as VCNeF2DModel
from vcnef.vcnef_3d import VCNeFModel as VCNeF3DModel
model = VCNeF2DModel(num_channels=4,
condition_on_pde_param=True,
pde_param_dim=2,
d_model=256,
n_heads=8,
n_transformer_blocks=1,
n_modulation_blocks=6)
# Random data with shape b, s_x, s_y, c
x = torch.rand(4, 64, 64, 4)
grid = torch.rand(4, 64, 64, 2)
pde_param = torch.rand(4, 2)
t = torch.arange(1, 21).repeat(4, 1) / 20
y_hat = model(x, grid, pde_param, t)Below is a listing of the directory structure of VCNeF.
examples.py: Contains lightweight examples of how to use VCNeF.
examples_pde_bench.py: Contains examples of how to use VCNeF with PDEBench data and the PDEBench training loop.
📂 vcnef: Contains the code for the VCNeF model.
📂 utils: Contains utils for the PDEBench example.
To use the PDEBench example examples_pde_bench.py, you have to download the PDEBench datasets. An overview of the avaiable data and how to download it can be found in the PDEBench repository. To use the downloaded datasets in the example, you have to adapt the path in base_path and the file name(s) in file_names.
The following illustation shows the architecture of the VCNeF model for solving 2D time-dependent PDEs (e.g., Navier-Stokes equations).
The code of VCNeF is based on the code of Linear Transformers and PDEBench. We would like to thank the authors of Linear Transformers and PDEBench for their work, which made our method possible.
MIT licensed, except where otherwise stated. Please see LICENSE file.
If you find our project useful, please consider citing it.
@inproceedings{vcnef2024,
author = {Hagnberger, Jan and Kalimuthu, Marimuthu and Musekamp, Daniel and Niepert, Mathias},
title = {{Vectorized Conditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations}},
year = {2024},
booktitle = {Proceedings of the 41st International Conference on Machine Learning (ICML 2024)}
}