{"id":113051,"date":"2026-03-12T10:30:00","date_gmt":"2026-03-12T17:30:00","guid":{"rendered":"https:\/\/2.zoppoz.workers.dev:443\/https\/developer.nvidia.com\/blog\/?p=113051"},"modified":"2026-04-02T11:35:39","modified_gmt":"2026-04-02T18:35:39","slug":"build-accelerated-differentiable-computational-physics-code-for-ai-with-nvidia-warp","status":"publish","type":"post","link":"https:\/\/2.zoppoz.workers.dev:443\/https\/developer.nvidia.com\/blog\/build-accelerated-differentiable-computational-physics-code-for-ai-with-nvidia-warp\/","title":{"rendered":"Build Accelerated, Differentiable Computational Physics Code for AI with NVIDIA Warp"},"content":{"rendered":"\n

Computer-aided engineering (CAE)<\/a> is shifting from human-driven workflows toward AI-driven ones, including physics foundation models that generalize across geometries and operating conditions. Unlike LLMs, these models depend on large volumes of high-fidelity, physics-compliant data.\u00a0<\/p>\n\n\n\n

Recent scaling-law work on computational fluid dynamics (CFD) surrogates<\/a> indicates that simulation-generated training data is often the limiting cost in practice. This pushes requirements onto the simulator, which must be GPU-native, fast, and able to plug directly into ML workflows.<\/p>\n\n\n\n

NVIDIA Warp<\/a> is a framework for accelerated simulation, data generation, and spatial computing that bridges CUDA and Python. Warp enables developers to write high-performance kernels as regular Python functions that are JIT-compiled into efficient code for execution on the GPU. Unlike the tensor-based frameworks, in which developers express computation as operations on entire N-dimensional arrays, developers author flexible kernels in the Warp framework that execute simultaneously across all elements of a computational grid. <\/p>\n\n\n\n

Simulation kernels are often expressed on computational grids and rely on data-dependent control flow like conditionals, early-outs, and selective updates that vary per element. In tensor frameworks, these patterns require composing Boolean masks that quickly become unwieldy and can waste computation on irrelevant elements. In a Warp kernel, each thread can branch, skip, or exit independently, expressing this logic naturally without masking workarounds.<\/p>\n\n\n\n

Furthermore, as this post will show, solvers written in Warp can be easily made differentiable through the Warp native support for automatic differentiation. They are straightforward to integrate with optimization or training workflows while remaining interoperable with frameworks like PyTorch, JAX, and NumPy for use cases spanning simulation, robotics, perception, and geometry processing.<\/p>\n\n\n\n

This post walks you through how to build a 2D Navier\u2013Stokes solver entirely in Warp. It explains how the Warp programming model maps onto a PDE solver. Then, it differentiates through the simulation to solve an optimal perturbation problem end-to-end. It closes with industrial case studies showcasing what Warp can enable in production workflows. For more information, see the 2D Navier\u2013Stokes solver example<\/a> and 2D Navier-Stokes optimal perturbation example<\/a> on the NVIDIA\/warp GitHub repo.<\/p>\n\n\n\n

How to write a 2D Navier\u2013<\/em>Stokes solver using Warp<\/strong><\/i><\/a><\/h2>\n\n\n\n

To keep the focus on Warp rather than on numerical methods, a textbook example of 2D decaying turbulence is used here, described by the vorticity-streamfunction formulation of the incompressible Navier-Stokes equations. The vorticity \\(\\omega\\) evolves according to the transport equation:<\/p>\n\n\n\n

\\(\\frac{\\partial \\omega}{\\partial t} + \\frac{\\partial \\psi}{\\partial y}\\frac{\\partial \\omega}{\\partial x} – \\frac{\\partial \\psi}{\\partial x}\\frac{\\partial \\omega}{\\partial y} = \\frac{1}{\\text{Re}}\\nabla^2 \\omega \\tag{1}\\)<\/p>\n\n\n\n

and the streamfunction \\(\\psi\\) is recovered from vorticity through the Poisson equation:<\/p>\n\n\n\n

\\(\\nabla^2 \\psi = -\\omega \\tag{2}\\)<\/p>\n\n\n\n

With periodic boundary conditions, the equation above reduces to an algebraic equation in Fourier space bypassing the need for iterative solvers:<\/p>\n\n\n\n

\\(\\hat{\\psi}_{m,n} = \\frac{\\hat{\\omega}_{m,n}}{k_x^2 + k_y^2} \\tag{3}\\)<\/p>\n\n\n\n

where \\((k_x, k_y)\\) is the wavenumber pair in the Fourier space. The solver makes use of the Fast Fourier Transform (FFT) algorithm to efficiently transform \\(\\omega\\) and \\(\\psi\\) to Fourier space and vice versa.<\/p>\n\n\n\n

Each timestep has two subcomponents (Figure 1). First, the vorticity transport equation is discretized on an \\(N \\times N\\) grid over an \\(L \\times L\\) square domain. The solution is marched forward in time by \\(\\Delta t\\) using a third-order strong stability-preserving Runge-Kutta (RK3) scheme<\/a> to obtain \\(\\omega(t+\\Delta t)\\). Second, the Poisson equation is solved in the Fourier space to obtain the updated \\(\\psi(t+\\Delta t)\\). <\/p>\n\n\n

\n
\"Flowchart\n\t\t\t\n\t\t\t\t\n\t\t\t<\/svg>\n\t\t<\/button>
Figure 1. Schematic of a single timestep loop for the solver<\/em><\/em><\/figcaption><\/figure>\n<\/div>\n\n\n

Thus, the forward solver has two building blocks that will be described in the subsequent sections:<\/p>\n\n\n\n