MeshfreeFlowNet: A Physics-Constrained Deep Continuous Space-Time Super-Resolution Framework

Chiyu Max Jiang, Soheil Esmaeilzadeh, Kamyar Azizzadenesheli, Karthik Kashinath, Mustafa Mustafa, Hamdi A. Tchelepi, Philip Marcus, Prabhat, Anima Anandkumar

Introduction

In recent years, along with the significant growth of computational resources, there has been an increasing attempt towards accurately resolving the wide range of length and time scales present within different physical systems. These length and time scales often differ by several orders of magnitude. Such time and length scale disparities are commonly observed in different physical phenomena such as turbulent flows, convective diffusive systems, subsurface systems, multiphase flows, chemical processes, and climate systems . Yet many key physical quantities of interest are highly dependent on correctly resolving such fine scales, such as the energy spectrum in turbulent flows.

From a numerical perspective, resolving the wide range of spatio-temporal scales within such physical systems is challenging since extremely small spatial and temporal numerical stencils would be required. In order to alleviate the computational burden of fully resolving such a wide range of spatial and temporal scales, multiscale computational approaches have been developed. For instance, in the subsurface flow problem, the main idea of the multiscale approach is to build a set of operators that map between the unknowns associated with the computational cells in a fine-grid and the unknowns on a coarser grid. The operators are computed numerically by solving localized flow problems. The multiscale basis functions have subgrid-scale resolutions, ensuring that the fine-scale heterogeneity is correctly accounted for in a systematic manner . Similarly, in turbulent flows , climate forecast , multiphase systems , reactive transport , complex fluids and biological systems , multiscale approaches attempt to relate the coarse-scale solutions to the fine-scale local solutions taking into account the fine-scale physics and dynamics. In these multiscale approaches, although the small scale physics and dynamics are being captured and accounted for at larger scales, the fine-grid solutions are extremely costly and hence impractical to solve for. Reconstructing the fine-scale subgrid solutions by having coarse-scale solutions in space and/or time still remains a challenge.

In this manuscript, we refer to such a process of reconstructing the fine-scale subgrid solutions from the coarse-scale solutions as super-resolution, where the high-resolution solutions are reconstructed using the low-resolution physical solutions in space and/or time. One key insight is that there are inherent statistical correlations between such pairs of low-resolution and high-resolution solutions that can be leveraged to reconstruct high-resolution solutions from the low-resolution inputs. Furthermore, this super-resolution process can be effectively modeled using deep learning models that learn such statistical correlations in a self-supervised manner from low-resolution and high-resolution pairs that exist in a training dataset. A successful super-resolution model should be able to efficiently represent the high-resolution outputs, effectively scale to large spatio-temporal domains as in a realistic problem setting, and allow for incorporating physical constraints in the form of PDEs to regularize the outputs to physically valid solutions. Furthermore, in order for a learning-based methodology to be applied to real problems such as Computational Fluid Dynamics (CFD) simulations and climate modeling, the methodology needs to address the High Performance Computing (HPC) challenges of scalability to large scale computing systems, for both the training and inference stages.

To this end, we propose MeshfreeFlowNet, a novel physics-constrained, deep learning based super-resolution framework to generate continuous (grid-free) spatio-temporal solutions from the low-resolution inputs. MeshfreeFlowNet first maps the low-resolution inputs to a localized latent context grid using a convolutional encoder, which can then be continuously queried at arbitrary resolutions. In summary, our main contributions are as follows:

We propose MeshfreeFlowNet, a novel and efficient physics-constrained deep learning model for super-resolution tasks.

We implement a set of physics-based metrics that allow for an objective assessment of the reconstruction quality of the super-resolved high-resolution turbulent flows.

We empirically assess the effectiveness of the MeshfreeFlowNet framework on the task of super-resolving turbulent flows in the Rayleigh–Bénard convection problem, showing a consistent and significant improvement in recovering key physical quantities of interest in super-resolved solutions over competing baselines.

We demonstrate the scalability of our framework to larger and more challenging problems by providing a large scale implementation of MeshfreeFlowNet that scales to up to 128 GPUs while retaining a ∼97%\sim 97\% scaling efficiency.

Related Works

Recently, deep learning models have been applied for studying fluid flow problems in different applications. In particular, a certain class of research works has studied flow problems on a grid representation where dynamic flow properties (e.g., velocity, pressure, etc.) could be considered as structured data similar to images. For instance, Guo et al. 2016 used convolutional architectures for approximating non-uniform steady laminar flow around 2D and 3D objects. Similarly, Bhatnagar et al. 2019 used convolutional neural networks for prediction of the flow fields around airfoil objects and for studying aerodynamic forces in a turbulent flow regime. Zhu and Zabaras 2018 used Bayesian deep convolutional neural networks to build surrogate flow models for the purpose of uncertainty quantification for flow applications in heterogeneous porous media. Alternatively, a different area of research has attempted to solve PDEs governing different physical phenomena by using simple fully connected deep neural networks. In order to enforce physical constraints, the PDEs under consideration are added as a term to the loss function at the training stage. For instance, E and Yu 2018 used fully connected deep neural networks with skip connections for solving different classes of PDEs such as Poisson equations and eigenvalue problems. Bar and Sochen 2019 used fully connected deep neural networks for solving forward and inverse problems in the context of electrical impedance tomography governed by elliptical PDEs. Long et al. 2017 proposed a framework to uncover unknown physics in the form of PDEs by learning constrained convolutional kernels. Raissi et al. 2019 applied fully connected neural networks for solving PDEs in a forward or inverse setup. Smith et al. 2020 deployed residual neural networks to solve the Eikonal PDE equation and perform ray tracing for an application in earthquake hypo-center inversion and seismic tomography. Using their proposed framework they solved PDEs in different contexts such as fluids, quantum mechanics, seismology, reaction-diffusion systems, and the propagation of nonlinear shallow-water waves. More recently, alternative representations for spatial functions have been explored. Li et al. 2020 proposed using message passing on graph networks in order to map between input data of PDEs and their solutions. They showed that the learned networks can generalize between different PDE approximation methods (e.g., Finite Difference and Finite Element methods) and between approximations that correspond to different levels of discretization. Similar to our MeshfreeFlowNet framework in using a latent context grid that can be continuously decoded into an output field, Jiang et al. 2020a used such representations for the computer vision task of reconstructing 3D scenes, where latent context grids are decoded into continuous implicit functions that represent the surfaces of different geometries. In the area of fluid dynamics, for turbulent flow predictions, Wang et al. 2019 developed a physics-informed deep learning framework. They introduced the use of trainable spectral filters coupled with Reynolds-averaged Navier-Stokes (RANS) and Large Eddy Simulation (LES) models followed by a convolutional architecture in order to predict turbulent flows. Jiang et al. 2020b presented a differentiable physics layer within the neural networks using spectral methods in order to enforce hard linear constraints. By performing a linear projection in the spectral space, they were able to enforce a divergence-free condition for the velocity. Accordingly, they could conserve mass within the physical system both locally and globally, allowing the super-resolution of turbulent flows to be statistically accurate. In computer vision applications, the super-resolution (SR) process has been proposed in the context of reconstructing high-resolution (HR) videos/images from the low-resolution (LR) ones. In the literature, different classical super-resolution methods have been proposed such as prediction-based methods , edge-based methods , statistical methods , patch-based methods and sparse representation methods . Most recently, deep learning based super-resolution models have also been actively explored and differ mainly in their types of specific intended application, architectures, loss functions, and learning principles .

Preliminaries

In this section, we provide a detailed explanation of the problem set up, notations, data sets, along with the learning paradigm and evaluation metrics.

Consider a space and time dependent partial differential equation (PDE) with a unique solution, written in the general form as

For a partial differential equation (Eqn. 1), a problem instance is realized with a given source and boundary condition functions ss and bb. For a pair (s,b)\left(s,b\right), we consider H{\mathcal{H}} to be an operator that generates the high-resolution solution yH\bm{y}_{\mathcal{H}}, i.e., yH=H(s,b)\bm{y}_{\mathcal{H}}={\mathcal{H}}(s,b), and L{\mathcal{L}} to be an operator that produces the low-resolution solution yL\bm{y}_{\mathcal{L}}, i.e., yL=L(s,b,yH)\bm{y}_{\mathcal{L}}={\mathcal{L}}(s,b,\bm{y}_{\mathcal{H}}). Consider a compact normed space of operators ΠF\Pi_{\mathcal{F}}, a set of operators F∈ΠF\mathcal{F}\in\Pi_{\mathcal{F}}, mapping low-resolution solutions to high-resolution ones. For a given compact normed space of functions Πu,b\Pi_{u,b}, let ε(ΠF,Πs,b)\varepsilon\left(\Pi_{\mathcal{F}},\Pi_{s,b}\right) denote the approximation gap with respect to ΠF\Pi_{\mathcal{F}}, i.e.,

with continuous approximation error ∥H(s,b)−F(L(s,b,y))∥\|{\mathcal{H}}(s,b)-\mathcal{F}\left({\mathcal{L}}(s,b,\bm{y})\right)\| in F,s,bF,s,b. Here the norm is with respect to a desired Lp(μ)L^{p}(\mu) space where p≥1p\geq 1 and μ\mu is a preferred measure Note that in the max-min game of the Eqn. 2, the action of the environment player, (s,b)(s,b), is revealed to the approximator player to choose F\mathcal{F}. Therefore, since the game is in the favor of the approximating player, the value of the game, ε(ΠF,Πs,b)\varepsilon\left(\Pi_{\mathcal{F}},\Pi_{s,b}\right) can be much smaller than its min-max counterpart. We require the min-max value to be desirably small when we aim to have a single model F\mathcal{F} to perform well on a set of designated problems.. In this work, given a problem instance specified with (s,b)(s,b), we are interested in learning F\mathcal{F} using parametric models. In order to map a low-resolution solution to its corresponding super-resolution with low approximation error, we require ε(ΠF,Πs,b)\varepsilon\left(\Pi_{\mathcal{F}},\Pi_{s,b}\right) to be small and comparable with a desired error level. Therefore, given a problem set Πu,b\Pi_{u,b}, a proper and rich class of operators, ΠF\Pi_{\mathcal{F}}, allows for a desirable approximation error.

Dataset Overview

In this work, we generate the dataset as the solution to a classical fluid dynamics system with a chaotic nature. We consider the well-known Rayleigh–Bénard instability problem in 2D where a static bulk fluid (kinematic viscosity ν\nu, thermal diffusivity α\alpha) is initially occupying the space between two horizontal plates (see Fig. 1). The lower and upper plates are considered to be respectively hot and cold with temperatures of THT_{H} and TCT_{C}. Gradually the temperature of the bulk fluid adjacent to the hot (cold) plate increases (decreases) and due to the buoyancy effects and density gradient the bulk fluid ascends (descends), leading to the formation of vortices and growth of flow instability in a chaotic and turbulent regime. The governing partial differential equations for the Rayleigh–Bénard instability problem are

where P∗=(Ra Pr)−1/2P^{*}=(Ra\,Pr)^{-1/2} and R∗=(Ra/Pr)−1/2R^{*}=(Ra/Pr)^{-1/2} with RaRa and PrPr being the Rayleigh and Prandtl numbers respectively defined as Ra=gαΔTL3ν−1κ−1Ra=g\alpha\Delta TL^{3}\nu^{-1}\kappa^{-1} and Pr=νκ−1Pr=\nu\kappa^{-1}, with gg, α\alpha, ν\nu, κ\kappa, ΔT\Delta T, and LL respectively being the gravity acceleration, thermal expansion coefficient, kinematic viscosity, thermal diffusivity, temperature difference between hot and cold plates, and the separation length between the plates. From a physical perspective RaRa quantifies the balance between the gravitational forces and viscous damping, and PrPr quantifies the ratio between momentum diffusivity and thermal diffusivity (balance between heat convection and conduction).

We use the Dedalus framework in order to numerically solve the system of Equations (3a)-(3c) using the spectral method approach. We solve Equations (3a)-(3c) for a duration of tft_{f} in time with a time step size of Δt\Delta t. In a coordinate system of xx-axis, and zz-axis, we consider a plate length LxL_{x} and a separation distance LzL_{z}, and discretize the domain with nxn_{x} and nzn_{z} points respectively in the xx and zz directions.

For the simulation cases, we consider Rayleigh and Prandtl numbers respectively in the range of Ra∈[104, 108]Ra\in[10^{4},\,10^{8}] and Pr∈[0.1, 10]Pr\in[0.1,\,10]. Upon solving the system of Eqns. (3a)-(3c), we create a high-resolution Rayleigh–Bénard simulation dataset DH\mathcal{D}_{H}, unless otherwise mentioned, with a spatial resolution of nxn_{x} = 4×nz4\times n_{z} = 512512, and a temporal resolution of ntn_{t} = 400400 (upon adaptive time stepping). We consider a normalized domain size of unit length in zz-direction with a domain aspect ratio of 4, i.e., Lx=4×Lz=4L_{x}=4\times L_{z}=4 [mm], and solve the Rayleigh–Bénard problem for a duration of 50 [ss]. Then we create a low-resolution dataset DL\mathcal{D}_{L}, by downsampling the high-resolution data in both space and time. We use downsampling factors of dt=4d_{t}=4 and ds=8d_{s}=8 for creating the low-resolution data in the temporal and spatial dimensions respectively.

Evaluation Metrics

In this work, we use multiple physical metrics, each accounting for different aspects of the flow field, in order to report the evaluation of the MeshfreeFlowNet model for super-resolving low-resolution data. As the specific metrics of evaluation, we report the Normalized Mean Absolute Error (NMAE) and R2 score for the physical metrics between the ground truth and the predicted high-resolution data. Such physical metrics are listed in the following.

Total Kinetic Energy (EtotE_{tot}) : the kinetic energy per unit mass associated with the flow is defined as the total kinetic energy and can be expressed as Etot=12⟨uiui⟩E_{tot}=\frac{1}{2}\langle u_{i}u_{i}\rangle .

Root Mean Square Velocity (urmsu_{rms}) : the square root of the scaled total kinetic energy is defined as the root mean square (RMS) velocity as urms=(2/3)Etotu_{rms}=\sqrt{({2}/{3})E_{tot}} .

Dissipation (ε\varepsilon) : is the rate at which turbulence kinetic energy is converted into thermal internal energy by molecular viscosity and defined as ε=2ν⟨SijSij⟩\varepsilon=2\nu\langle S_{ij}S_{ij}\rangle with SijS_{ij} and ν\nu respectively being the rate of strain tensor and the kinematic viscosity.

Taylor Microscale (λ\lambda) : is the intermediate length scale at which viscous forces significantly affect the dynamics of turbulent eddies and defined as λ=15 ν urms2 ε−1\lambda=\sqrt{15\,\nu\,u_{rms}^{2}\,\varepsilon^{-1}}. Length scales larger than the Taylor microscale (i.e., inertial range) are not strongly affected by viscosity. Below the Taylor microscale (i.e., dissipation range) the turbulent motions are subject to strong viscous forces and kinetic energy is dissipated into heat.

Taylor-scale Reynolds (ReλRe_{\lambda}) : is defined as the ratio of RMS inertial forces to viscous forces and is expressed as Reλ=urms λ ν−1Re_{\lambda}=u_{rms}\,\lambda\,\nu^{-1} .

Kolmogorov Time (τη\tau_{\eta}) and Length (η\eta) Scales : Kolmogorov microscales are the smallest scales in turbulent flows where viscosity dominates and the turbulent kinetic energy dissipates into heat. Kolmogorov time and length scale can be respectively expressed as τη=ν/ε\tau_{\eta}=\sqrt{\nu/\varepsilon} and η=ν3/4 ε−1/4\eta=\nu^{3/4}\,\varepsilon^{-1/4} .

Turbulent Integral Scale (LL) : is a measure of the average spatial extent or coherence of the fluctuations and can be expressed as L=π2urms2∫E(k)kdkL=\frac{\pi}{2u^{2}_{rms}}\int{\frac{E(k)}{k}}dk .

Large Eddy Turnover Time (TLT_{L}) : is defined as the typical time scale for an eddy of length scale LL to undergo significant distortion and is also the typical time scale for the transfer of energy from scale LL to smaller scales, since this distortion is the mechanism for energy transfer and expressed as TL=L/urmsT_{L}=L/u_{rms} .

Systems, Platforms and Configuration

In order to further illustrate the feasibility of utilizing our proposed MeshfreeFlowNet framework for large physical systems requiring processing orders of magnitude more computation, here we study the scalability of MeshfreeFlowNet. We scale our model on a GPU stack on the Cori supercomputer at the National Energy Research Scientific Computing Center (NERSC). We use data distributed parallelism with synchronous gradient descent, and test performance up to 16 nodes with a total of 128128 GPUs.

In a data-parallel distribution strategy, the model is replicated over all GPUs. At every step, each GPU receives its own random batch of the dataset to calculate the local gradients. Gradients are averaged across all devices with an all_reduce operation. To achieve better scaling efficiency, the communication of one layer’s gradients is overlapped with the backprop computation of the previous layer. PyTorch torch.distributed package provides communication primitives for multiprocess parallelism with different collective communication backends. We use NVIDIA’s Collective Communications Library (NCCL) backend which gave us the best performance when running on GPUs, both within and across multiple nodes. The PyTorch torch.nn.parallel.DistributedDataParallel wrapper, which we use for the results in this paper, builds on top of the torhch.distributed package to provide efficient data-parallel distributed training. With this setup, we achieve more than 96%96\% throughput efficiency on 16 Cori GPU nodes, 128 GPUs in total. Cori has 8 V100 (Volta) GPUs per node, the GPUs are interconnected with NVLinks in hybrid cube-mesh topology . The nodes are equipped with Mellanox MT27800 (ConnectX-5) EDR InfiniBand network cards.

MeshfreeFlowNet

In this work, we propose the MeshfreeFlowNet framework as a novel computational algorithm for constructing the super-resolution solutions to partial differential equations using their low-resolution counterpart solutions.

The Context Generation Network is a convolutional encoder that produces a Latent Context Grid from the low-resolution physical input DL\mathcal{D}_{L}. Denote this network as

In this work, we implement the Context Generation Network as a 3D variant of the U-Net architecture which was originally proposed by Ronneberger et al. 2015. U-Net has successfully been applied for different computer vision tasks that involve dense localized predictions such as image style transfer , image segmentation , image enhancement , image coloring , and image generation

Different from the original U-Net architecture, we replace the 2D convolutions with 3D counterparts and utilize residue blocks instead of individual convolution layers for better training properties of the network. The U-Net comprises of a contractive part followed by an expansive part. The contractive part is composed of multiple stages of convolutional residue blocks, each followed by a max-pooling layer (of stride 2). Each residue block consists of 3 convolution layers (1x1, 3x3, 1x1) interleaved with batch normalization layers and ReLU activation. The expansive part mirrors the contractive part, replacing max-pooling with nearest neighbor upsampling. In between the layers with similar grid sizes within the contractive and expansive parts, a skip connection concatenates the features from the contractive layer with the features in the expansive layer as the input to the subsequent layers in the contractive part in order to preserve the localized contextual information.

Continuous Decoding Network

One unique property of our super-resolution methodology is that the output is continuous instead of discrete. This removes the limitations in output resolution, and additionally, it allows for an effective computation of the gradients of predicted output physical quantities, enabling an easy way of enforcing PDE-based physical constraints.

The continuous decoding network can be implemented using a simple Multilayer Perceptron, see Fig. 6. For each query, denote the spatio-temporal query location to be xi\bm{x}_{i} and the latent context grid to be G:={(xj,cj);j≤∣∣G∣∣}\mathcal{G}:=\{(\bm{x}_{j},\bm{c}_{j});j\leq||\mathcal{G}||\}, where (xj,cj)(\bm{x}_{j},\bm{c}_{j}) are the spatio-temporal coordinates and the latent context vector for the jj-th vertex of the grid. Denote Ni\mathcal{N}_{i} as the set of neighboring vertices that bound xix_{i}, where for a (dd+1) dimensional spatio-temporal grid ∣∣Ni∣∣=2d+1||\mathcal{N}_{i}||=2^{d+1}. Denote the continuous decoding network, implemented as a Multilayer Perception as

where θ2\theta_{2} is the set of trainable parameters of the Multilayer Perception network. The query value at x\bm{x} with respect to the shared network Φθ2(x)\Phi_{\theta_{2}}(\bm{x}) and the latent context grid G\mathcal{G} can be calculated as

where ∑j∈Niwj=1\sum_{j\in\mathcal{N}_{i}}w_{j}=1, wjw_{j} is the trilinear interpolation weight with respect to the bounding vertex jj, Δx:={Δx,Δz,Δt}\Delta\bm{x}:=\{\Delta x,\Delta z,\Delta t\} is the stencil size corresponding to the discretization grid vertices.

Since the Continuous Decoding Network is implemented as an MLP, arbitrary spatio-temporal derivatives of the output quantities: Γy\Gamma\bm{y} can be effectively computed via backpropagation through the MLP. Denote the approximation of the derivative operator to be ΓΦ\Gamma_{\Phi}. We combine the partial derivatives to compute the equation loss as the norm of the residue of the governing equations. In the continuous decoding network we consider two infinitely differentiable activation functions namely Softplus and Swish, where we have found the results obtained by Swish to outperform the ones obtained by Softplus.

Loss Function

We use a weighted combination of two losses to train our MeshfreeFlowNet network: the norm of the difference between the predicted physical outputs and the ground truth physical outputs, which we refer to as Prediction Loss, and the norm of the residues of the governing PDEs, which we refer to as Equation Loss.

Denote the set of sample locations within a mini-batch B\mathcal{B} of training samples to be {(xji,yji,DLi);i∈B,j∈Bi}\{(\bm{x}_{j}^{i},\bm{y}_{j}^{i},\mathcal{D}_{L}^{i});i\in\mathcal{B},j\in\mathcal{B}^{i}\} where Bi\mathcal{B}^{i} is the mini-batch of point samples for the ii-th low-resolution input, y\bm{y} is the vector that represents the ground truth physical output quantities. In the case of the Rayleigh-Bénard Convection example in this study, we have y:={P,T,u,v}\bm{y}:=\{P,T,u,v\} where P,T,u,vP,T,u,v are the pressure, temperature, x-velocity and y-velocity terms respectively. The super-resolution for the learned model queried at x\bm{x} conditioning on low-resolution input DL\mathcal{D}_{L} is

where y^\hat{\bm{y}} is the predicted output vector. The prediction loss Lp\mathcal{L}_{p} for a mini-batch can be formulated as

where ∣∣⋅∣∣l||\cdot||_{l} is the Frobenius ll-norm of the difference. We use the L1 Norm for computing the prediction loss. The Equation loss Le\mathcal{L}_{e} for a mini-batch can be formulated as

which is the norm of the PDE equation residue, using the PDE definition in Eqn. 1. Finally, a single loss term for training the network can be represented as a weighted sum of the two loss terms as

where γ\gamma is a hyperparameter for weighting the equation loss.

Experiments

In all the experiments, we use an Adam optimizer with learning rate of 10−210^{-2}, and l1l_{1} regularization for the loss function, 3000 random samples per epoch, and train for 100 epochs.

In this part, we investigate the influence of the importance given to the Equation loss and the prediction loss (see, Section 4.3) on the performance of MeshfreeFlowNet. As presented in Eqn. 10, the total loss (L\mathcal{L}) comprises of the Prediction loss (Lp\mathcal{L}_{p}) and the Equation loss (Le\mathcal{L}_{e}) where the Equation loss is weighted with a scaling coefficient γ\gamma. Accordingly, we study the influence of the hyperparameter γ\gamma in the loss function given in Eqn. 10 on the accuracy of MeshfreeFlowNet. For this purpose, we consider γ ∈ {0,0.0125,0.025,0.05,0.1,0.2,0.4,0.8,1.0}\gamma\,\in\,\{0,0.0125,0.025,0.05,0.1,0.2,0.4,0.8,1.0\}. Considering both Softplus and Swish activation functions in the continuous decoding network, we have found that the results obtained by the latter one outperform the ones obtained by the former one, for which, the values of the evaluation metrics for MeshfreeFlowNet trained with each of the γ\gamma values are presented in Table 1. γ=0\gamma=0 indicates a loss function which only depends on the Prediction loss, and the physical aspects (PDE imposed constraints) of the predicted high-resolution data are not accounted for. As presented in Table 1, the best performance on the validation set is achieved for a loss function with the Equation loss weighting coefficient of γ=0.05\gamma=0.05. Allover this work we refer to this optimum weighting coefficient as γ∗\gamma^{*} and perform the training tasks using the loss function in Eqn. 10 where the Equation loss is weighted with γ∗=0.05\gamma^{*}=0.05. As presented in Table 1, a model that is trained with γ=0\gamma=0, which only focuses on the data (i.e., uses Prediction loss only) and does not account for the physics and the PDE constraints (i.e., ignores Equation loss) underperforms compared to the trained model with γ=γ∗\gamma=\gamma^{*}. On the other hand models trained with a significant focus on the physical constraints only (i.e., large γ\gamma) underperform in super-resolving the low-resolution data. In general, a balance between the focus of the MeshfreeFlowNet on the model and the physical constraints leads to an optimal super-resolution performance. In achieving that, in Eqn. 10, the Prediction loss (Lp\mathcal{L}_{p}) captures the global structure of the data and the Equation loss (Le\mathcal{L}_{e}) further guides the model in accurately reconstructing the local structure of the data.

MeshfreeFlowNet vs. Baseline Models

In this section, we present a comparison between the performance of our proposed MeshfreeFlowNet framework for super-resolving low-resolution data against two baselines, namely: a classic trilinear interpolation algorithm (Baseline (I)), and a deep learning based 3D U-Net model (Baseline (II)). Specifically for the 3D U-Net model for Baseline (II), we use the same U-Net backbone as in our MeshfreeFlowNet framework, with the difference being that while MeshfreeFlowNet uses the 3D U-Net to generate a latent context grid, the baseline U-Net continues with 3D up-sampling and transpose-convolution operations up to the target high-resolution space (see, Fig. 7). Table 2 presents the comparison of the MeshfreeFlowNet with the baseline models. As shown in Table 2, the Baseline (I) is a purely interpolation-based approach and fails to reconstruct the high-resolution data and resolve the fine-scale details, leading to large errors in flow-based evaluation metrics that characterize the flow dynamics. The extremely large normalized mean absolute error (NMAE) of the calculated rmsrms velocity for the fine-scale solution found by Baseline (I) indicates that a merely interpolative scheme cannot accurately reconstruct the fine-scale local dynamics (e.g., the flow velocities). Moreover, the NMAE of the total kinetic energy (Etot)E_{tot}), as a global characterizing parameter of the turbulent dynamics, is also very large for the high-resolution solution by Baseline (I). On the other hand, the deep learning based Baseline (II) which utilizes the 3D U-Net model directly maps the low-resolution data to the high-resolution space, achieves better performance compared to Baseline (I). However, as presented in Table 2, our MeshfreeFlowNet model performs significantly better than Baselines (I) and (II). The specification of γ\gamma in Table 2 refers to the weighting coefficient of the Equation loss component in the loss function in Eqn. 10. In Table 2, γ=0\gamma=0 indicates that only the prediction loss has been considered for training the MeshfreeFlowNet model, whereas γ=γ∗\gamma=\gamma^{*} refers to the optimum value for the weighting coefficient of the Equation loss from the ablation study (see, Table 1, Sec. 5.1).

Generalizability of MeshfreeFlowNet

We further evaluate the robustness of MeshfreeFlowNet for resolution enhancement of low-resolution datasets that have physical initial and boundary conditions different from the datasets the MeshfreeFlowNet model has been trained on. We refer to such initial and boundary conditions as unseen initial/boundary conditions. In order to investigate the generalizability of MeshfreeFlowNet on unseen initial and boundary conditions, we study the effect of each condition separately in the following setups.

We investigate the robustness of a trained MeshfreeFlowNet for enhancing the resolution of a low-resolution unseen data with physical initial conditions different than the training datasets that the MeshfreeFlowNet has been trained on. Table 3 shows the performance of the MeshfreeFlowNet on a dataset with unseen initial conditions. The first row of Table 3 shows the values of the evaluation metrics for unseen test data when MeshfreeFlowNet is trained only on one dataset whereas the second row shows the values of the evaluation metrics when MeshfreeFlowNet is trained on 10 datasets each with a different initial condition. As can be observed from the results, the performance of MeshfreeFlowNet on unseen cases can be improved by training on a more diverse set of initial conditions.

Unseen Physical Boundary Conditions

We further investigate the robustness of a trained MeshfreeFlowNet for super-resolving a low-resolution unseen dataset with physical boundary conditions different from the training datasets that the MeshfreeFlowNet has been trained on. As the use of more datasets in Sec. 5.3.1 was shown to be effective in improving the performance of MeshfreeFlowNet for super-resolution, here we use a dataset that comprises of 10 different sets of boundary conditions (i.e. different Rayleigh numbers of RaRa ∈×105\in\times 10^{5}, corresponding to Reynolds numbers of up to 10,000). Table 4 shows the performance of such a trained MeshfreeFlowNet for 5 different test datasets each with a different Rayleigh number boundary condition. In Table 4, MeshfreeFlowNet’s performance is evaluated for a Rayleigh number within the range of boundary conditions of the training sets (i.e., Ra=5×106Ra=5\times 10^{6}), for Rayleigh numbers slightly below and above the range of boundary conditions of the training sets (i.e., Ra=1×105Ra=1\times 10^{5} and Ra=1×107Ra=1\times 10^{7} respectively), and for Rayleigh numbers far below and above the range of boundary conditions of the training sets (i.e., Ra=1×104Ra=1\times 10^{4} and Ra=1×108Ra=1\times 10^{8} respectively). As presented in Table 4, MeshfreeFlowNet achieves a good performance not only on the boundary conditions within the range of Rayleigh number boundary conditions it has been trained on, but also on the unseen boundary conditions far out of the range of Rayleigh number boundary conditions it has been trained on. This illustrates the fact that a trained MeshfreeFlowNet model can generalize well to unseen boundary conditions.

Scalability of MeshfreeFlowNet

Last but not least, we study the scalability of the MeshfreeFlowNet model to study its applicability to larger problems that require orders-of-magnitude more compute. The scaling results are presented in Fig. 9. Figs. 9(a) demonstrates that an almost-ideal scaling performance can be achieved for the MeshfreeFlowNet model on up to 128 GPU workers, achieving approximately 96.80%96.80\% scaling efficiency. In Fig. 9(b), we show the convergence for the model training loss with respect to the number of epochs. In Fig. 9(c), we show the loss convergence with respect to total wall time, where an increasing number of GPU workers leads to a drastic decrease in total training time. As the models achieve similar levels of losses after 100 epochs, yet the training throughput scales almost linearly with the number of GPUs, we see a close to an ideal level of scaling for convergence speed on up to 16 GPUs. A small anomaly regarding the loss curve for 128 GPUs where the loss does not decrease to an ideal level after 100 epochs shows that a very large batch size could lead to diminishing scalability with respect to model convergence. This has been observed in numerous machine learning scaling studies and requires further investigations within the community.

Conclusion and Future Work

In this work, for the first time, we presented the MeshfreeFlowNet, a physics-constrained super-resolution framework, that can produce continuous super-resolution outputs, would allow imposing arbitrary combinations of PDE constraints and could be evaluated on arbitrary-sized spatio-temporal domains due to its fully-convolutional nature. We further demonstrated that MeshfreeFlowNet can recover a wide range of important physical flow quantities (e.g., including Turbulent Kinetic Energy, Kolmogorov Time and Length Scales, etc.) by accurately super-resolving turbulent flows significantly better than traditional (trilinear interpolation) and deep learning based (3D U-Net) baselines. We further illustrated the scalability of MeshfreeFlowNet to a large cluster of GPU nodes with a high speed interconnect, demonstrating its applicability to problems that require orders of magnitude more computational resources.

Future work includes exploring the applicability of MeshfreeFlowNet to other physical applications beyond 2D Rayleigh Bernard convection. One interesting direction to pursue is to explore the use of 4D spatio-temporal convolution operators such as those proposed by Choy et al. 2019 to further extend this framework to 4D space-time simulations. That will open up a wide range of applications in Turbulence modeling, where statistical priors between the low-resolution simulation and subgrid-scale physics can be learned from 3+13+1D Direct Numerical Simulations (DNS). The scalability of the model on HPC clusters will be critical in learning from such large scale datasets, where single node training will be prohibitively slow. The fully convolutional nature of the MeshfreeFlowNet framework, along with the demonstrated scalability makes it well poised for such challenges. Moreover, due to the generalizability of our PDE constrained framework, it would be interesting to apply this framework on applications beyond turbulent flows.

Acknowledgements

This research used resources of the National Energy Research Scientific Computing Center (NERSC), a DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-05CH11231. The research was performed at the Lawrence Berkeley National Laboratory for the U.S. Department of Energy under Contract No. DE340AC02- 05CH11231. K. Kashinath is supported by the Intel Big Data Center at NERSC. K. Azizzadenesheli gratefully acknowledges the financial support of Raytheon and Amazon Web Services. A. Anandkumar is supported in part by Bren endowed chair, DARPA PAIHR00111890035 and LwLL grants, Raytheon, Microsoft, Google, and Adobe faculty fellowships. We also acknowledge the Industrial Consortium on Reservoir Simulation Research at Stanford University (SUPRI-B).

References