Turbulence Super Resolution

Sym v1.0.0

This example uses Modulus Sym to train a super-resolution surrogate model for predicting high-fidelity homogeneous isotropic turbulence fields from filtered low-resolution observations provided by the Johns Hopkins Turbulence Database. This model will combine standard data-driven learning as well as how to define custom data-driven loss functions that are uniquely catered to a specific problem. In this example you will learn the following:

  1. How to use data-driven convolutional neural network models in Modulus Sym

  2. How to define custom data-driven loss and constraint

  3. How to define custom data-driven validator

  4. Modulus Sym features for structured/grid data

  5. Adding custom parameters to the problem configuration file

Note

This tutorial assumes that you have completed the Introductory Example tutorial on Lid Driven Cavity flow and have familiarized yourself with the basics of Modulus Sym. This also assumes that you have a basic understanding of the convolution models in Modulus Sym Pix2Pix Net and Super Resolution Net.

Warning

The Python package pyJHTDB is required for this example to download and process the training and validation datasets. Install using pip install pyJHTDB.

The objective of this problem is to learn the mapping between a low-resolution filtered 3D flow field to a high-fidelity solution. The flow field will be samples of a forced isotropic turbulence direct numerical simulation originally simulated with a resolution of \(1024^{3}\). This simulation solves the forced Navier-Stokes equations:

(205)\[\frac{\partial \textbf{u}}{\partial t} + \textbf{u} \cdot \nabla \textbf{u} = -\nabla p /\rho + \nu \nabla^{2}\textbf{u} + \textbf{f}.\]

The forcing term \(\textbf{f}\) is used to inject energy into the simulation to maintain a constant total energy. This dataset contains 5028 time steps spanning from 0 to 10.05 seconds which are sampled every 10 time steps from the original pseudo-spectral simulation.

super_res_dataset_sample.png

Fig. 137 Snap shot of \(128^{3}\) isotropic turbulence velocity fields

The objective is to build a surrogate model to learn the mapping between a low-resolution velocity field \(\textbf{U}_{l} = \left\{u_{l}, v_{l}, w_{l}\right\}\) to the true high-resolution field \(\textbf{U}_{h} = \left\{u_{h}, v_{h}, w_{h}\right\}\) for any low-resolution sample in this isotropic turbulence dataset \(\textbf{U}_{l} \sim p(\textbf{U}_{l})\). Due to the size of the full simulation domain, this tutorial focuses on predicting smaller volumes such that the surrogate learns with a low resolution dimensionality of \(32^{3}\) to a high-resolution dimensionality of \(128^{3}\). Use the Super Resolution Net in this tutorial, but Pix2Pix Net is also integrated into this problem and can be used instead if desired.

super_res_surrogate.png

Fig. 138 Super resolution network for predicting high-resolution turbulent flow from low-resolution input

This example demonstrates how to write your own data-driven constraint. Modulus Sym ships with a standard supervised learning constraint for structured data called SupervisedGridConstraint used in the Darcy Flow with Fourier Neural Operator example. However, if you want to have a problem specific loss you can extend the base GridConstraint. Here, you will set up a constraint that can pose a loss between various measures for the fluid flow including velocity, continuity, vorticity, enstrophy and strain rate.

Note

The python script for this problem can be found at examples/super_resolution/super_resolution.py.

Copy
Copied!
            

class SuperResolutionConstraint(Constraint): def __init__( self, nodes: List[Node], invar: Dict[str, np.array], outvar: Dict[str, np.array], batch_size: int, loss_weighting: Dict[str, int], dx: float = 1.0, lambda_weighting: Dict[str, Union[np.array, sp.Basic]] = None, num_workers: int = 0, ): dataset = DictGridDataset( invar=invar, outvar=outvar, lambda_weighting=lambda_weighting ) super().__init__( nodes=nodes, dataset=dataset, loss=PointwiseLossNorm(), batch_size=batch_size, shuffle=True, drop_last=True, num_workers=num_workers, ) self.dx = dx self.ops = FlowOps().to(self.device) self.loss_weighting = {} self.fields = set("U") for key, value in loss_weighting.items(): if float(value) > 0: self.fields = set(key).union(self.fields) self.loss_weighting[key] = value

An important part to note is that you can control which losses you want to contribute using a loss_weighting dictionary which will be provided from the config file. Each one of these measures are calculated in examples/super_resolution/ops.py, which can be referenced for more information. However, as a general concept, finite difference methods are used to calculate gradients of the flow field and the subsequent measures.

Copy
Copied!
            

def calc_flow_stats(self, data_var): output = {"U": data_var["U"]} vel_output = {} cont_output = {} vort_output = {} enst_output = {} strain_output = {} # compute derivatives if len(self.fields) > 1: grad_output = self.ops.get_velocity_grad( data_var["U"], dx=self.dx, dy=self.dx, dz=self.dx ) # compute continuity if "continuity" in self.fields: cont_output = self.ops.get_continuity_residual(grad_output) # compute vorticity if "omega" in self.fields or "enstrophy" in self.fields: vort_output = self.ops.get_vorticity(grad_output) # compute enstrophy if "enstrophy" in self.fields: enst_output = self.ops.get_enstrophy(vort_output) # compute strain rate if "strain" in self.fields: strain_output = self.ops.get_strain_rate_mag(grad_output) if "dU" in self.fields: # Add to output dictionary grad_output = torch.cat( [ grad_output[key] for key in [ "u__x", "u__y", "u__z", "v__x", "v__y", "v__z", "w__x", "w__y", "w__z", ] ], dim=1, ) vel_output = {"dU": grad_output} if "omega" in self.fields: vort_output = torch.cat( [vort_output[key] for key in ["omega_x", "omega_y", "omega_z"]], dim=1 ) vort_output = {"omega": vort_output} output.update(vel_output) output.update(cont_output) output.update(vort_output) output.update(enst_output) output.update(strain_output) return output

Override the loss calculation with the custom method which calculates the relative MSE between predicted and target velocity fields using flow measures defined in the weight dictionary, self.loss_weighting.

Copy
Copied!
            

def loss(self, step: int) -> Dict[str, torch.Tensor]: # Calc flow related stats pred_outvar = self.calc_flow_stats(self._pred_outvar) target_vars = self.calc_flow_stats(self._target_vars) # compute losses losses = {} for key in target_vars.keys(): mean = (target_vars[key] ** 2).mean() losses[key] = ( self.loss_weighting[key] * (((pred_outvar[key] - target_vars[key]) ** 2) / mean).mean() ) return losses

The resulting complete loss for this problem is the following:

(206)\[\begin{split}\mathcal{L} = RMSE(\hat{U}_{h}, U_{h}) + \lambda_{dU}RMSE(\hat{dU}_{h}, dU_{h}) + \lambda_{cont}RMSE(\nabla\cdot\hat{U}_{h}, \nabla\cdot U_{h}) + \lambda_{\omega}RMSE(\hat{\omega}_{h}, \omega_{h}) \\ + \lambda_{strain}RMSE(|\hat{D}|_{h}, |D|_{h}) + \lambda_{enst}RMSE(\hat{\epsilon}_{h}, \epsilon_{h}),\end{split}\]

in which \(\hat{U}_{h}\) is the prediction from the neural network and \(U_{h}\) is the target. \(dU\) is the velocity tensor, \(\omega\) is the vorticity, \(|D|\) is the magnitude of the strain rate and \(\epsilon\) is the flow’s enstrophy. All of these can be turned on and off in the configuration file under the custom.loss_weights config group.

Similarly, because the input and output are of different dimensionality, the built in GridValidator in Modulus Sym will not work since it expects all tensors to be the same size. You can easily extend this to write out the high-resolution outputs and low-resolution outputs into separate VTK uniform grid files.

Copy
Copied!
            

class SuperResolutionValidator(GridValidator): def __init__(self, *args, log_iter: bool = False, **kwargs): super().__init__(*args, **kwargs) self.log_iter = log_iter self.device = DistributedManager().device def save_results(self, name, results_dir, writer, save_filetypes, step): invar_cpu = {key: [] for key in self.dataset.invar_keys} true_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} pred_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} # Loop through mini-batches for i, (invar0, true_outvar0, lambda_weighting) in enumerate(self.dataloader): # Move data to device (may need gradients in future, if so requires_grad=True) invar = Constraint._set_device( invar0, device=self.device, requires_grad=self.requires_grad ) true_outvar = Constraint._set_device( true_outvar0, device=self.device, requires_grad=self.requires_grad ) pred_outvar = self.forward(invar) # Collect minibatch info into cpu dictionaries invar_cpu = { key: value + [invar[key].cpu().detach()] for key, value in invar_cpu.items() } true_outvar_cpu = { key: value + [true_outvar[key].cpu().detach()] for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: value + [pred_outvar[key].cpu().detach()] for key, value in pred_outvar_cpu.items() } # Concat mini-batch tensors invar_cpu = {key: torch.cat(value) for key, value in invar_cpu.items()} true_outvar_cpu = { key: torch.cat(value) for key, value in true_outvar_cpu.items() } pred_outvar_cpu = { key: torch.cat(value) for key, value in pred_outvar_cpu.items() } # compute losses on cpu losses = GridValidator._l2_relative_error(true_outvar_cpu, pred_outvar_cpu) # convert to numpy arrays invar = {k: v.numpy() for k, v in invar_cpu.items()} true_outvar = {k: v.numpy() for k, v in true_outvar_cpu.items()} pred_outvar = {k: v.numpy() for k, v in pred_outvar_cpu.items()} # save batch to vtk file named_target_outvar = {"true_" + k: v for k, v in true_outvar.items()} named_pred_outvar = {"pred_" + k: v for k, v in pred_outvar.items()} for b in range(min(4, next(iter(invar.values())).shape[0])): if self.log_iter: grid_to_vtk( {**named_target_outvar, **named_pred_outvar}, results_dir + name + f"_{b}_hr" + f"{step:06}", batch_index=b, ) else: grid_to_vtk( {**named_target_outvar, **named_pred_outvar}, results_dir + name + f"_{b}_hr", batch_index=b, ) grid_to_vtk(invar, results_dir + name + f"_{b}_lr", batch_index=b) # add tensorboard plots if self.plotter is not None: self.plotter._add_figures( name, results_dir, writer, step, invar, true_outvar, pred_outvar, ) # add tensorboard scalars for k, loss in losses.items(): writer.add_scalar("val/" + name + "/" + k, loss, step, new_style=True)

Here the grid_to_vtk function in Modulus Sym is used, which writes tensor data to VTK Image Datasets (Uniform grids), which can then be viewed in Paraview. When your data is structured grid_to_vtk is preferred to var_to_polyvtk due to the lower memory footprint of a VTK Image Dataset vs VTK Poly Dataset.

Before proceeding, it is important to recognize that this problem needs to download the dataset upon its first run. To download the dataset from Johns Hopkins Turbulence Database you will need to request an access token. Information regarding this process can be found on the database website. Once aquired, please overwrite the default token in the config in the specified location. Utilities used to download the data can be located in examples/super_resolution/jhtdb_utils.py, but will not be discussed in this tutorial.

Warning

By registering for an access token and using the Johns Hopkins Turbulence Database, you are agreeing to the terms and conditions of the dataset itself. This example will not work without an access token.

Warning

The default training dataset size is 512 sampled and the validation dataset size is 16 samples. The download can take several hours depending on your internet connection. The total memory footprint of the data is around 13.5Gb. A smaller dataset can be set in the config file.

Configuration

The config file for this example is as follows. Note that both the super-resolution and pix2pix encoder-decoder architecture configuration are included to test.

Copy
Copied!
            

# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. defaults: - modulus_default - /arch/super_res_cfg@arch.super_res - /arch/pix2pix_cfg@arch.pix2pix - scheduler: tf_exponential_lr - optimizer: adam - loss: sum - _self_ jit: True cuda_graphs: False # Graphs does not work with super res network arch: super_res: scaling_factor: 4 pix2pix: batch_norm: True n_downsampling: 1 n_blocks: 9 dimension: 3 scaling_factor: 4 scheduler: decay_rate: 0.95 decay_steps: 2000 optimizer: lr: 0.0001 training: rec_validation_freq: 250 rec_constraint_freq: 250 save_network_freq: 250 print_stats_freq: 25 max_steps: 20000 batch_size: train: 4 valid: 4 custom: jhtdb: n_train: 512 n_valid: 16 domain_size: 128 access_token: "edu.jhu.pha.turbulence.testing-201311" #Replace with your own token here loss_weights: U: 1.0 dU: 0 continuity: 0 omega: 0.1 enstrophy: 0 strain: 0

The custom config group can be used to store case specific parameters that will not be used inside of Modulus Sym. Here you can use this group to define parameters related to the dataset size, the domain size of the fluid volumes and the database access token which you should replace with your own!

Note

To just test the model with a toy dataset without a database access token you are recommended to use the below settings:

Copy
Copied!
            

jhtdb: n_train: 4 n_valid: 1 domain_size: 16 access_token: "edu.jhu.pha.turbulence.testing-201311"

Loading Data

To load the dataset into memory, you will use the following utilities:

Copy
Copied!
            

# load jhtdb datasets invar, outvar = make_jhtdb_dataset( nr_samples=cfg.custom.jhtdb.n_train, domain_size=cfg.custom.jhtdb.domain_size, lr_factor=cfg.arch.super_res.scaling_factor, token=cfg.custom.jhtdb.access_token, data_dir=to_absolute_path("datasets/jhtdb_training"), time_range=[1, 768], dataset_seed=123, ) invar_valid, outvar_valid = make_jhtdb_dataset( nr_samples=cfg.custom.jhtdb.n_valid, domain_size=cfg.custom.jhtdb.domain_size, lr_factor=cfg.arch.super_res.scaling_factor, token=cfg.custom.jhtdb.access_token, data_dir=to_absolute_path("datasets/jhtdb_valid"), time_range=[768, 1024], dataset_seed=124, )

This will download and cache the dataset locally, so you will not need to download it with every run.

Initializing the Model

Here you initialize the model following the standard Modulus Sym process. Note that the input and output keys have a size=3, which tells Modulus Sym that these variables have 3 dimensions (velocity components).

Copy
Copied!
            

model = instantiate_arch( input_keys=[Key("U_lr", size=3)], output_keys=[Key("U", size=3)], cfg=cfg.arch.super_res, ) nodes = [model.make_node(name="super_res")]

Adding Data Constraints

Copy
Copied!
            

# make super resolution domain jhtdb_domain = Domain() # make data driven constraint jhtdb_constraint = SuperResolutionConstraint( nodes=nodes, invar=invar, outvar=outvar, batch_size=cfg.batch_size.train, loss_weighting=cfg.custom.loss_weights, lambda_weighting=None, dx=2 * np.pi / 1024.0, ) jhtdb_domain.add_constraint(jhtdb_constraint, "constraint")

Adding Data Validator

Copy
Copied!
            

# make validator dataset = DictGridDataset(invar_valid, outvar_valid) jhtdb_validator = SuperResolutionValidator( dataset=dataset, nodes=nodes, batch_size=cfg.batch_size.valid, log_iter=False, ) jhtdb_domain.add_validator(jhtdb_validator, "validator")

NVIDIA recommends that your first run be on a single GPU to download the dataset. Only root process will download the data from the online database while the others will be idle.

Copy
Copied!
            

python super_resolution.py

However, parallel training is suggested for this problem once the dataset is downloaded. This example was trained on 4 V100 GPUs which can be run via Open MPI using the following command:

Copy
Copied!
            

mpirun -np 4 python super_resolution.py

Results and Post-processing

Since this example illustrated how to set up a custom data-driven loss that is controllable through the config, you can compare the impact of several different loss components on the model’s performance. The TensorBoard plot is shown below with the validation dataset loss being the bottom most graph. Given the number of potential loss components, only a handful are compared here:

  • U=1.0: \(\mathcal{L} = RMSE(\hat{U}_{h}, U_{h})\)

  • U=1.0, omega=0.1: \(\mathcal{L} = RMSE(\hat{U}_{h}, U_{h}) + 0.1RMSE(\hat{\omega}_{h}, \omega_{h})\)

  • U=1.0, omega=0.1, dU=0.1: \(\mathcal{L} = RMSE(\hat{U}_{h}, U_{h}) + 0.1RMSE(\hat{\omega}_{h}, \omega_{h}) + 0.1RMSE(\hat{dU}_{h}, dU_{h})\)

  • U=1.0, omega=0.1, dU=0.1, contin=0.1: \(\mathcal{L} = RMSE(\hat{U}_{h}, U_{h}) + 0.1RMSE(\hat{\omega}_{h}, \omega_{h}) + 0.1RMSE(\hat{dU}_{h}, dU_{h}) + 0.1RMSE(\nabla\cdot\hat{U}_{h}, \nabla\cdot U_{h})\)

The validation error is the L2 relative error between the predicted and true high-resolution velocity fields. You can see that the inclusion of vorticity in the loss equation increases the model’s accuracy, however the inclusion of other terms does not. Loss combinations of additional fluid measures has proven successful in past works 1 2. However, additional losses can potentially make the optimization more difficult for the model and adversely impact accuracy.

super_res_tensorboard.png

Fig. 139 Tensorboard plot comparing different loss functions for turbulence super-resolution

The output VTK files can be found in the 'outputs/super_resolution/validators' folder which you can then view in Paraview. The volumetric plots of the velocity magnitude fields are shown below where you can see the model dramatically improves the low-resolution velocity field.

super_res_validation_0.png

Fig. 140 Velocity magnitude for a validation case using the super resolution model for predicting turbulence

super_res_validation_1.png

Fig. 141 Velocity magnitude for a validation case using the super resolution model for predicting turbulence

References

[1]

Geneva, Nicholas and Zabaras, Nicholas. “Multi-fidelity generative deep learning turbulent flows” Foundations of Data Science (2020).

[2]

Subramaniam, Akshay et al. “Turbulence enrichment using physics-informed generative adversarial networks” arXiv preprint arXiv:2003.01907 (2020).

© Copyright 2023, NVIDIA Modulus Team. Last updated on Aug 8, 2023.