Mastering NVIDIA PhysicsNeMo for Darcy Flow and Neural Operators

A Step-by-Step Coding Tutorial on NVIDIA PhysicsNeMo: Darcy Flow, FNOs, PINNs, Surrogate Models, and Inference Benchmarking

NVIDIA PhysicsNeMo provides a specialized framework for scientific machine learning, enabling the resolution of complex partial differential equations through neural operators. This implementation solves the 2D Darcy Flow problem by mapping permeability fields to pressure fields using Fourier Neural Operators. The tutorial establishes a complete workflow from synthetic data generation using Gaussian Random Fields to real-time inference benchmarking.

Why This Matters

Traditional numerical methods for solving partial differential equations, such as finite differences or finite elements, are computationally intensive and often require iterative solvers that do not scale well with resolution. Physics-informed machine learning models like FNOs and PINNs offer a paradigm shift by learning the underlying operator or satisfying the PDE residual, providing near-instantaneous inference once trained. This transition from numerical approximation to learned operators is critical for real-time subsurface flow modeling and heat conduction simulations where traditional solvers fail to meet latency requirements. By integrating physical laws directly into the loss function or architecture, these models maintain high fidelity even in data-scarce regimes.

Key Insights

• The 2D Darcy Flow equation, -∇·(k∇u) = f, serves as a primary benchmark for evaluating the accuracy of neural operators in diffusion-dominated physics.
• Fourier Neural Operators (FNO) utilize the Fast Fourier Transform (FFT) to perform convolutions in the frequency domain, effectively capturing global dependencies with learnable spectral weights.
• Physics-Informed Neural Networks (PINNs) leverage automatic differentiation to compute PDE residuals, allowing the model to optimize for physical consistency alongside data accuracy.
• Spectral Convolution layers in FNOs are designed to keep only low-frequency modes, which acts as a regularization technique for smooth physical fields.
• Gaussian Random Fields (GRF) are employed to generate complex, heterogeneous permeability fields (k) that serve as realistic inputs for training surrogate models.
• Inference benchmarking demonstrates that neural operators can process 100 samples in milliseconds, providing a massive throughput advantage over classical iterative Jacobi methods.

Working Examples

Implementation of a 2D Spectral Convolution layer, the core component of a Fourier Neural Operator.

class SpectralConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, modes1, modes2):
        super().__init__()
        self.scale = 1 / (in_channels * out_channels)
        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, modes1, modes2, dtype=torch.cfloat))
        self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, modes1, modes2, dtype=torch.cfloat))

    def compl_mul2d(self, input, weights):
        return torch.einsum("bixy,ioxy->boxy", input, weights)

    def forward(self, x):
        batch_size = x.shape[0]
        x_ft = torch.fft.rfft2(x)
        out_ft = torch.zeros(batch_size, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, dtype=torch.cfloat, device=x.device)
        out_ft[:, :, :self.modes1, :self.modes2] = self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
        out_ft[:, :, -self.modes1:, :self.modes2] = self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x

Practical Applications

• Subsurface flow modeling for reservoir engineering, where FNOs predict pressure distribution based on permeability maps. Pitfall: High-frequency noise in input data can lead to artifacts if spectral modes are not properly truncated.
• Real-time heat conduction analysis using PINNs to ensure thermal safety in electronic components. Pitfall: Imbalanced loss weights between data and physics residuals can cause the model to ignore boundary conditions.
• Surrogate modeling for CFD simulations where a U-Net style Convolutional Surrogate provides a baseline for operator learning. Pitfall: Standard CNNs may struggle with resolution-invariance compared to neural operators.

References:

• https://www.marktechpost.com/2026/04/13/a-step-by-step-coding-tutorial-on-nvidia-physicsnemo-darcy-flow-fnos-pinns-surrogate-models-and-inference-benchmarking/