Transolver for Darcy Flow

This example demonstrates how to set up a data-driven model for a 2D Darcy flow using the Transolver inside of PhysicsNeMo.

Training progress can be tracked through MLFlow. This example runs on a single GPU.

Problem overview

This example is based on a 2D Darcy flow problem, which is often used to model flow diffusion through a porous medium. The equation describes the steady-state pressure and velocity field (related by the constitutive relation, based on Darcy's law) for some fluid of interest. As an applied example, one might consider a geological example, where the fluid is oil which is both a) generated by a source term $f(\mathbf{x})$ and b) diffusing through porous bedrock with permeability $k(\mathbf{x})$. For isotropic media, the governing equation is given as the following second-order elliptic PDE:

$$ \nabla \cdot \left(,k(\mathbf{x})\ \nabla p \left(\mathbf{x}\right),\right) = f(\mathbf{x}) $$

where $\mathbf{x} \in \Omega \subset \mathbb{R}^2$ is the spatial coordinate, $p(\mathbf{x})$ is the pressure, $k(\mathbf{x})$ is the local permeability and $f(\mathbf{x})$ is the source term.

Dataset

For the baseline case defined in train_transolver_darcy.py, the dataset is generated on-the-fly at train time via the Darcy2D Datapipe described in ./physicsnemo/datapipes/benchmarks/darcy.py.

For the modified case with a fixed (i.e., precomputed) dataset defined in train_transolver_darcy_fix.py, the dataset is pre-generated, and loaded via the Darcy2D_fix Datapipe. In the fixed case, extra data is needed for training and the data path should be added when Darcy_2D_fix dataset is constructed. You can download the data here.

The fix dataset training (which uses a fixed dataset) requires you to convert data from matlab to numpy format, for faster startup of the training. Just use the convert_mat_to_npz.py script to port your data.

Model overview and architecture

Getting Started

Prerequisites

Install the required dependencies by running below:

pip install -r requirements.txt

Training

To train the model following PhysicsNeMo's settings, simply run:

python train_transolver_darcy.py

Each batch is a new data generated by equation, which is different from commonly-used settings.

To reproduce the results in the paper, run:

python train_transolver_darcy_fix.py

In this case, the train set and test set are fixed after the construction of Dataset, corresponding to Transolver's setting.

Customization

To train Transolver on the same 2D Darcy flow problem with different physical parameters or ML hyperparameters, modify the config.yaml file.

To train Transolver on your own physics problem, modify the dataloader in train_transolver_darcy.py to use your own pre-computed data or on-the-fly solver.

Additional Information

More components are added for convenience. Validators calculate the loss between ground-truth and prediction, and visualize them in ./mlruns. Below is a simple example of visualization.

References

Transolver: A Fast Transformer Solver for PDEs on General Geometries