# Darcy Flow with Fourier Neural Operator¶

## Introduction¶

This tutorial sets up a data-driven model for a 2D Darcy flow using the Fourier Neural Operator (FNO) architecture inside of Modulus. It covers these details:

Loading grid data and setting up data-driven constraints

How to create a grid validator node

How to use Fourier Neural Operator architecture in Modulus

Note

This tutorial assumes that you are familiar with the basic functionality of Modulus and understand the FNO architecture. Please see the Lid Driven Cavity Background and Fourier Neural Operator sections for additional information.

Warning

The Python package gdown is required for this example if you do not already have the example data downloaded and converted.
Install using `pip install gdown`

.

## Problem Description¶

The Darcy PDE is a second order, elliptic PDE with the following form:

in which \(u(\textbf{x})\) is the flow pressure, \(k(\textbf{x})\) is the permeability field and \(f(\cdot)\) is the forcing function. The Darcy flow can parameterize a variety of systems including flow through porous media, elastic materials and heat conduction. Here you will define the domain as a 2D unit square \(D=\left\{x,y \in (0,1)\right\}\) with the boundary condition \(u(\textbf{x})=0, \textbf{x}\in\partial D\). Recall that FNO requires a structured Euclidean input s.t. \(D = \textbf{x}_{i}\) where \(i \in \mathbb{N}_{N\times N}\). Thus both the permeability and flow fields are discretized into a 2D matrix \(\textbf{K}, \textbf{U} \in \mathbb{R}^{N \times N}\).

This problem develops a surrogate model that learns the mapping between a permeability field and the pressure field, \(\textbf{K} \rightarrow \textbf{U}\), for a distribution of permeability fields \(\textbf{K} \sim p(\textbf{K})\). This is a key distinction of this problem from other examples, you are not learning just a single solution but rather a distribution.

## Case Setup¶

This example is a data-driven problem. This means that before starting any coding you need to make sure you have both the training and validation data.
The training and validation data-sets for this example can be found on the Fourier Neural Operator Github page.
Here is an automated script for downloading and converting this dataset.
This requires the package gdown which can easily installed through `pip install gdown`

.

Note

The python script for this problem can be found at `examples/darcy/darcy_FNO_lazy.py`

.

### Configuration¶

The configuration for this problem is fairly standard within Modulus. A specific FNO architecture is defined for this
example inside of the config file.
These settings were derived through an automated hyper-parameter sweep using Hydra multirun.
The most important parameter for FNO models is `dimension`

which tells Modulus to load a 1D, 2D or 3D FNO architecture.
`nr_fno_layers`

are the number of Fourier convolution layers in the model and `fno_layer_size`

are the size of the latent embedded features inside the model.

```
defaults :
- modulus_default
- arch:
- fno
- scheduler: tf_exponential_lr
- optimizer: adam
- loss: sum
- _self_
jit: false
arch:
fno:
dimension: 2
nr_fno_layers: 4
fno_layer_size: 32
fno_modes: 12
padding: 9
output_fc_layer_sizes:
- 128
scheduler:
decay_rate: 0.95
decay_steps: 1000
training:
rec_results_freq : 1000
max_steps : 10000
batch_size:
grid: 32
validation: 32
```

### Loading Data¶

For this data-driven problem the first step is to get the training data into a Modulus.
Prior to loading data, set any normalization value that you want to apply to the data.
For this dataset, calculate the scale and shift parameters for both the input permeability field and output pressure.
Then, set this normalization inside Modulus by providing a shift/scale to each key, `Key(name, scale=(shift, scale))`

.

```
print("Multi-GPU currently not supported for this example. Exiting.")
return
```

There are two approaches for loading data: First, use eager loading where you immediately read the entire dataset onto memory at one time. Alternatively, you can use lazy loading where the data is loaded on a per example basis as the model needs it for training. The former eliminates potential overhead from reading data from disk during training, however this cannot scale to large datasets. Lazy loading is used in this example for the training dataset to demonstrate this utility for larger problems.

```
if DistributedManager().distributed:
print("Multi-GPU currently not supported for this example. Exiting.")
return
# load training/ test data
input_keys = [Key("coeff", scale=(7.48360e00, 4.49996e00))]
output_keys = [Key("sol", scale=(5.74634e-03, 3.88433e-03))]
download_FNO_dataset("Darcy_241", outdir="datasets/")
train_path = to_absolute_path(
```

This data is in HDF5 format which is ideal for lazy loading using the HDF5DataFile object. Since the validation dataset is smaller, use eager loading and read the entire dataset into memory.

```
)
invar_train = {"coeff": HDF5DataFile(train_path, "coeff")} # lazy file loader
outvar_train = {"sol": HDF5DataFile(train_path, "sol")} # lazy file loader
invar_test, outvar_test = load_FNO_dataset(
"datasets/Darcy_241/piececonst_r241_N1024_smooth2.hdf5",
```

Note

The key difference when setting up eager versus lazy loading is the object passed in the variable dictionaries invar_train and outvar_train.
In eager loading these dictionaries should be of the type `Dict[str: np.array]`

, where each variable is a numpy array of data.
Lazy loading uses dictionaries of the type `Dict[str: DataFile]`

, consisting of `DataFile`

objects which are classes that are used to map between example index and the datafile.

### Initializing the Model¶

Initializing the model and domain follows the same steps as other examples.

```
for d in (invar_test, outvar_test):
for k in d:
print(f"{k}: {d[k].shape}")
# make list of nodes to unroll graph on
model = instantiate_arch(
input_keys=input_keys,
output_keys=output_keys,
cfg=cfg.arch.fno,
)
```

### Adding Data Constraints¶

For the physics-informed problems in Modulus, you typically need to define a geometry and
constraints based on boundary conditions and governing equations.
Here the only constraint is a `SupervisedGridConstraint`

which performs standard supervised
training on grid data.
This constraint supports the use of multiple workers, which are particularly important when using lazy loading.

```
# make domain
domain = Domain()
# add constraints to domain
supervised = SupervisedGridConstraint(
nodes=nodes,
invar=invar_train,
outvar=outvar_train,
batch_size=cfg.batch_size.grid,
cell_volumes=None,
lambda_weighting=None,
```

Note

Grid data refers to data that can be defined in a tensor like an image. Inside Modulus this grid of data typically represents a spatial domain
and should follow the standard dimensionality of `[batch, channel, xdim, ydim, zdim]`

where channel is the dimensionality of your state variables.
Both Fourier and convolutional models use grid-based data to efficiently learn and predict entire domains in one forward pass, which contrasts to the
pointwise predictions of standard PINN approaches.

### Adding Data Validator¶

The validation data is then added to the domain using `GridValidator`

which should be used
when dealing with structured data.
Recall that unlike the training constraint, we will use eager loading for the validator.
Thus, a dictionary of numpy arrays are passed to the constraint.

```
num_workers=4, # number of parallel data loaders
)
domain.add_constraint(supervised, "supervised")
# add validator
val = GridValidator(
invar_test,
outvar_test,
nodes,
```

## Training the Model¶

Start the training by executing the python script.

```
python darcy_FNO_lazy.py
```

Warning

Multi-GPU training is currently not supported for this problem.

### Results and Post-processing¶

The checkpoint directory is saved based on the results recording frequency
specified in the `rec_results_freq`

parameter of its derivatives. See Results Frequency for more information.
The network directory folder (in this case `'outputs/darcy_fno/validators'`

) contains several plots of different
validation predictions. Several are shown below, and you can see that the model is able to accurately predict the
pressure field for permeability fields it had not seen previously.