A Step-by-Step Coding Tutorial on NVIDIA PhysicsNeMo: Darcy Flow, FNOs, PINNs, Surrogate Models, and Inference Benchmarking
In this tutorial, we implement NVIDIA PhysicsNeMo on Colab and construct a sensible workflow for physics-informed machine studying. We begin by organising the setting, producing information for the 2D Darcy Flow drawback, and visualizing the bodily fields to obviously perceive the educational job. From there, we implement and practice highly effective fashions such because the Fourier Neural Operator and a convolutional surrogate baseline, whereas additionally exploring the concepts behind Physics-Informed Neural Networks. Also, we examine architectures, consider predictions, benchmark inference, and save skilled fashions, offering a complete hands-on view of how PhysicsNeMo can be utilized for scientific machine studying issues.
print("="*80)
print("SECTION 1: INSTALLATION AND SETUP")
print("="*80)
import subprocess
import sys
def install_packages():
"""Install required packages for the tutorial."""
packages = [
"nvidia-physicsnemo",
"matplotlib",
"numpy",
"h5py",
"scipy",
"tqdm",
]
for package deal in packages:
print(f"Installing {package deal}...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])
print("n✓ All packages put in efficiently!")
install_packages()
print("n" + "="*80)
print("SECTION 2: IMPORTS AND CONFIGURATION")
print("="*80)
import torch
import torch.nn as nn
import torch.nn.useful as F
from torch.utils.information import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple, Optional
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')
print(f"Using system: {system}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
attempt:
import physicsnemo
print(f"PhysicsNeMo model: {physicsnemo.__version__}")
PHYSICSNEMO_AVAILABLE = True
besides ImportError:
print("PhysicsNeMo not put in. Using customized implementations.")
PHYSICSNEMO_AVAILABLE = False
def set_seed(seed: int = 42):
"""Set random seeds for reproducibility."""
torch.manual_seed(seed)
np.random.seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
set_seed(42)
print("n" + "="*80)
print("SECTION 3: DATA GENERATION - 2D DARCY FLOW")
print("="*80)
"""
The 2D Darcy Flow equation is a basic benchmark for neural operators:
-∇·(ok(x,y)∇u(x,y)) = f(x,y) in Ω = [0,1]²
the place:
- ok(x,y) is the permeability discipline (enter)
- u(x,y) is the stress discipline (output)
- f(x,y) is the supply time period
This is a basic drawback in subsurface circulation modeling, warmth conduction,
and different diffusion-dominated physics.
"""
class DarcyFlowDataGenerator:
"""
Generate artificial 2D Darcy Flow information for coaching neural operators.
Uses Gaussian Random Fields for permeability and finite variations
to unravel the Darcy equation.
"""
def __init__(
self,
decision: int = 64,
length_scale: float = 0.1,
variance: float = 1.0
):
self.decision = decision
self.length_scale = length_scale
self.variance = variance
self.dx = 1.0 / (decision - 1)
x = np.linspace(0, 1, decision)
y = np.linspace(0, 1, decision)
self.X, self.Y = np.meshgrid(x, y)
self._setup_grf()
def _setup_grf(self):
"""Setup Gaussian Random Field covariance."""
n = self.decision
x_flat = self.X.flatten()
y_flat = self.Y.flatten()
dist_sq = (
(x_flat[:, None] - x_flat[None, :]) ** 2 +
(y_flat[:, None] - y_flat[None, :]) ** 2
)
self.cov_matrix = self.variance * np.exp(
-dist_sq / (2 * self.length_scale ** 2)
)
self.cov_matrix += 1e-6 * np.eye(n * n)
self.L = np.linalg.cholesky(self.cov_matrix)
def generate_permeability(self, n_samples: int = 1) -> np.ndarray:
"""Generate random permeability fields utilizing GRF."""
n = self.decision
z = np.random.randn(n * n, n_samples)
samples = self.L @ z
ok = np.exp(samples.T.reshape(n_samples, n, n))
return ok
def solve_darcy(
self,
permeability: np.ndarray,
supply: Optional[np.ndarray] = None,
n_iterations: int = 5000,
tol: float = 1e-6
) -> np.ndarray:
"""
Solve Darcy equation utilizing iterative Jacobi methodology.
-∇·(ok∇u) = f with u=0 on boundary (Dirichlet BC)
"""
n = self.decision
dx = self.dx
if supply is None:
supply = np.ones((n, n))
u = np.zeros((n, n))
u_new = np.zeros_like(u)
ok = permeability
for iteration in vary(n_iterations):
for i in vary(1, n - 1):
for j in vary(1, n - 1):
k_sum = k_e + k_w + k_n + k_s
u_new[i, j] = (
k_e * u[i, j + 1] + k_w * u[i, j - 1] +
k_n * u[i - 1, j] + k_s * u[i + 1, j] +
dx ** 2 * supply[i, j]
) / (k_sum + 1e-10)
if np.max(np.abs(u_new - u)) < tol:
break
u = u_new.copy()
return u
def generate_dataset(
self,
n_samples: int = 100,
show_progress: bool = True
) -> Tuple[np.ndarray, np.ndarray]:
"""Generate a dataset of (permeability, stress) pairs."""
permeabilities = self.generate_permeability(n_samples)
pressures = np.zeros_like(permeabilities)
iterator = tqdm(vary(n_samples), desc="Generating information") if show_progress else vary(n_samples)
for i in iterator:
pressures[i] = self.solve_darcy(permeabilities[i])
return permeabilities, pressures
class DarcyDataset(Dataset):
"""PyTorch Dataset for Darcy Flow information."""
def __init__(
self,
permeability: np.ndarray,
stress: np.ndarray,
normalize: bool = True
):
self.permeability = torch.FloatTensor(permeability)
self.stress = torch.FloatTensor(stress)
if normalize:
self.perm_mean = self.permeability.imply()
self.perm_std = self.permeability.std()
self.press_mean = self.stress.imply()
self.press_std = self.stress.std()
self.permeability = (self.permeability - self.perm_mean) / self.perm_std
self.stress = (self.stress - self.press_mean) / self.press_std
self.permeability = self.permeability.unsqueeze(1)
self.stress = self.stress.unsqueeze(1)
def __len__(self):
return len(self.permeability)
def __getitem__(self, idx):
return self.permeability[idx], self.stress[idx]
print("nGenerating Darcy Flow dataset...")
generator = DarcyFlowDataGenerator(decision=32, length_scale=0.15)
n_train = 200
n_test = 50
print(f"Generating {n_train} coaching samples...")
perm_train, press_train = generator.generate_dataset(n_train)
print(f"Generating {n_test} take a look at samples...")
perm_test, press_test = generator.generate_dataset(n_test)
train_dataset = DarcyDataset(perm_train, press_train)
test_dataset = DarcyDataset(perm_test, press_test, normalize=True)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)
print(f"n✓ Dataset created!")
print(f" Training samples: {len(train_dataset)}")
print(f" Test samples: {len(test_dataset)}")
print(f" Resolution: {generator.decision}x{generator.decision}")We arrange the total Colab setting and put in all required to run the PhysicsNeMo tutorial easily. We import the core libraries, configure the system, test GPU availability, and make the workflow reproducible by setting a set random seeds. We additionally generate the Darcy Flow dataset, put together the PyTorch datasets and dataloaders, and set up the muse for every thing we practice later within the tutorial.
print("n" + "="*80)
print("SECTION 4: DATA VISUALIZATION")
print("="*80)
def visualize_darcy_samples(
permeability: np.ndarray,
stress: np.ndarray,
n_samples: int = 3
):
"""Visualize Darcy circulation samples."""
fig, axes = plt.subplots(n_samples, 2, figsize=(10, 4 * n_samples))
for i in vary(n_samples):
im1 = axes[i, 0].imshow(permeability[i], cmap='viridis', origin='decrease')
axes[i, 0].set_title(f'Permeability Field (Sample {i+1})')
axes[i, 0].set_xlabel('x')
axes[i, 0].set_ylabel('y')
plt.colorbar(im1, ax=axes[i, 0], label='ok(x,y)')
im2 = axes[i, 1].imshow(stress[i], cmap='sizzling', origin='decrease')
axes[i, 1].set_title(f'Pressure Field (Sample {i+1})')
axes[i, 1].set_xlabel('x')
axes[i, 1].set_ylabel('y')
plt.colorbar(im2, ax=axes[i, 1], label='u(x,y)')
plt.tight_layout()
plt.savefig('darcy_samples.png', dpi=150, bbox_inches='tight')
plt.present()
print("✓ Saved visualization to 'darcy_samples.png'")
visualize_darcy_samples(perm_train[:3], press_train[:3])
print("n" + "="*80)
print("SECTION 5: FOURIER NEURAL OPERATOR (FNO)")
print("="*80)
"""
The Fourier Neural Operator (FNO) learns mappings between perform areas
by parameterizing the integral kernel in Fourier house.
Key perception: Convolution in bodily house = multiplication in Fourier house
The FNO layer consists of:
1. FFT to rework to frequency area
2. Multiplication with learnable weights (protecting solely low-frequency modes)
3. Inverse FFT to rework again
4. Residual reference to an area linear transformation
"""
class SpectralConv2nd(nn.Module):
"""
2D Spectral Convolution Layer for FNO.
Performs convolution in Fourier house by:
1. Computing FFT of enter
2. Multiplying with advanced learnable weights
3. Computing inverse FFT
"""
def __init__(
self,
in_channels: int,
out_channels: int,
modes1: int,
modes2: int
):
tremendous().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.modes1 = modes1
self.modes2 = modes2
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, enter: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
"""Complex multiplication for batch of 2D tensors."""
return torch.einsum("bixy,ioxy->boxy", enter, weights)
def ahead(self, x: torch.Tensor) -> torch.Tensor:
batch_size = x.form[0]
x_ft = torch.fft.rfft2(x)
out_ft = torch.zeros(
batch_size, self.out_channels, x.measurement(-2), x.measurement(-1) // 2 + 1,
dtype=torch.cfloat, system=x.system
)
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.measurement(-2), x.measurement(-1)))
return x
class FNOBlock(nn.Module):
"""
FNO Block combining spectral convolution with native linear remodel.
output = σ(SpectralConv(x) + NativeLinear(x))
"""
def __init__(
self,
channels: int,
modes1: int,
modes2: int,
activation: str = 'gelu'
):
tremendous().__init__()
self.spectral_conv = SpectralConv2nd(channels, channels, modes1, modes2)
self.local_linear = nn.Conv2nd(channels, channels, 1)
self.activation = nn.GELU() if activation == 'gelu' else nn.ReLU()
def ahead(self, x: torch.Tensor) -> torch.Tensor:
return self.activation(self.spectral_conv(x) + self.local_linear(x))
class FourierNeuralOperator2D(nn.Module):
"""
Complete 2D Fourier Neural Operator for studying operators.
Architecture:
1. Lift enter to greater dimensional channel house
2. Apply a number of FNO blocks (spectral convolutions + residuals)
3. Project again to output house
This learns the mapping: ok(x,y) -> u(x,y) for Darcy circulation
"""
def __init__(
self,
in_channels: int = 1,
out_channels: int = 1,
modes1: int = 12,
modes2: int = 12,
width: int = 32,
n_layers: int = 4,
padding: int = 9
):
tremendous().__init__()
self.modes1 = modes1
self.modes2 = modes2
self.width = width
self.padding = padding
self.fc0 = nn.Linear(in_channels + 2, width)
self.fno_blocks = nn.ModuleList([
FNOBlock(width, modes1, modes2) for _ in range(n_layers)
])
self.fc1 = nn.Linear(width, 128)
self.fc2 = nn.Linear(128, out_channels)
def get_grid(self, form: Tuple, system: torch.system) -> torch.Tensor:
"""Create normalized grid coordinates."""
batch_size, size_x, size_y = form[0], form[2], form[3]
gridx = torch.linspace(0, 1, size_x, system=system)
gridy = torch.linspace(0, 1, size_y, system=system)
gridx, gridy = torch.meshgrid(gridx, gridy, indexing='ij')
grid = torch.stack([gridx, gridy], dim=-1)
grid = grid.unsqueeze(0).repeat(batch_size, 1, 1, 1)
return grid
def ahead(self, x: torch.Tensor) -> torch.Tensor:
batch_size = x.form[0]
grid = self.get_grid(x.form, x.system)
x = x.permute(0, 2, 3, 1)
x = torch.cat([x, grid], dim=-1)
x = self.fc0(x)
x = x.permute(0, 3, 1, 2)
if self.padding > 0:
x = F.pad(x, [0, self.padding, 0, self.padding])
for block in self.fno_blocks:
x = block(x)
if self.padding > 0:
x = x[..., :-self.padding, :-self.padding]
x = x.permute(0, 2, 3, 1)
x = F.gelu(self.fc1(x))
x = self.fc2(x)
x = x.permute(0, 3, 1, 2)
return x
print("nCreating Fourier Neural Operator mannequin...")
fno_model = FourierNeuralOperator2D(
in_channels=1,
out_channels=1,
modes1=8,
modes2=8,
width=32,
n_layers=4,
padding=5
).to(system)
n_params = sum(p.numel() for p in fno_model.parameters() if p.requires_grad)
print(f"✓ FNO Model created with {n_params:,} trainable parameters")We first visualize the generated Darcy Flow samples to obviously see the connection between the permeability discipline and the ensuing stress discipline. We then implement the core constructing blocks of the Fourier Neural Operator, together with the spectral convolution layer, the FNO block, and the whole 2D FNO structure.
print("n" + "="*80)
print("SECTION 6: PHYSICS-INFORMED NEURAL NETWORK (PINN)")
print("="*80)
"""
Physics-Informed Neural Networks (PINNs) incorporate bodily legal guidelines straight
into the loss perform. For the Darcy equation:
-∇·(ok∇u) = f
The PINN loss consists of:
1. Data loss: MSE between predictions and noticed information
2. Physics loss: Residual of the PDE at collocation factors
3. Boundary loss: Satisfaction of boundary situations
"""
class PINN_MLP(nn.Module):
"""
Multi-Layer Perceptron for PINN.
Takes (x, y, ok(x,y)) as enter and outputs u(x,y).
Uses sine activation (Fourier options) for higher convergence.
"""
def __init__(
self,
input_dim: int = 3,
output_dim: int = 1,
hidden_dims: List[int] = [64, 64, 64, 64],
use_fourier_features: bool = True,
n_frequencies: int = 32
):
tremendous().__init__()
self.use_fourier_features = use_fourier_features
self.n_frequencies = n_frequencies
if use_fourier_features:
self.B = nn.Parameter(
torch.randn(2, n_frequencies) * 2 * np.pi,
requires_grad=False
)
actual_input_dim = 2 * n_frequencies + 1
else:
actual_input_dim = input_dim
layers = []
dims = [actual_input_dim] + hidden_dims
for i in vary(len(dims) - 1):
layers.append(nn.Linear(dims[i], dims[i + 1]))
layers.append(nn.Tanh())
layers.append(nn.Linear(dims[-1], output_dim))
self.community = nn.Sequential(*layers)
self._init_weights()
def _init_weights(self):
"""Xavier initialization for higher convergence."""
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
if m.bias isn't None:
nn.init.zeros_(m.bias)
def fourier_embedding(self, xy: torch.Tensor) -> torch.Tensor:
"""Apply Fourier function embedding to positional inputs."""
proj = xy @ self.B
return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
def ahead(
self,
x: torch.Tensor,
y: torch.Tensor,
ok: torch.Tensor
) -> torch.Tensor:
"""
Forward cross.
Args:
x: x-coordinates (batch,)
y: y-coordinates (batch,)
ok: permeability values at (x,y) (batch,)
Returns:
u: predicted stress values (batch,)
"""
xy = torch.stack([x, y], dim=-1)
if self.use_fourier_features:
pos_features = self.fourier_embedding(xy)
inputs = torch.cat([pos_features, k.unsqueeze(-1)], dim=-1)
else:
inputs = torch.cat([xy, k.unsqueeze(-1)], dim=-1)
return self.community(inputs).squeeze(-1)
class DarcyPINNLoss:
"""
Physics-informed loss for Darcy equation.
Computes:
1. Data loss at statement factors
2. PDE residual loss: -∇·(ok∇u) - f = 0
3. Boundary loss: u = 0 on boundaries
"""
def __init__(
self,
lambda_data: float = 1.0,
lambda_pde: float = 1.0,
lambda_bc: float = 10.0
):
self.lambda_data = lambda_data
self.lambda_pde = lambda_pde
self.lambda_bc = lambda_bc
def compute_gradients(
self,
mannequin: nn.Module,
x: torch.Tensor,
y: torch.Tensor,
ok: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Compute first and second derivatives utilizing autograd."""
x = x.requires_grad_(True)
y = y.requires_grad_(True)
u = mannequin(x, y, ok)
grad_outputs = torch.ones_like(u)
u_x = torch.autograd.grad(
u, x, grad_outputs=grad_outputs, create_graph=True
)[0]
u_y = torch.autograd.grad(
u, y, grad_outputs=grad_outputs, create_graph=True
)[0]
u_xx = torch.autograd.grad(
u_x, x, grad_outputs=grad_outputs, create_graph=True
)[0]
u_yy = torch.autograd.grad(
u_y, y, grad_outputs=grad_outputs, create_graph=True
)[0]
return u, (u_x, u_y), (u_xx, u_yy)
def pde_residual(
self,
mannequin: nn.Module,
x: torch.Tensor,
y: torch.Tensor,
ok: torch.Tensor,
f: torch.Tensor = None
) -> torch.Tensor:
"""
Compute PDE residual: -∇·(ok∇u) - f
Simplified kind for fixed ok regionally:
-k(u_xx + u_yy) - f ≈ 0
"""
u, (u_x, u_y), (u_xx, u_yy) = self.compute_gradients(mannequin, x, y, ok)
if f is None:
f = torch.ones_like(u)
residual = -k * (u_xx + u_yy) - f
return residual
def __call__(
self,
mannequin: nn.Module,
x_data: torch.Tensor,
y_data: torch.Tensor,
k_data: torch.Tensor,
u_data: torch.Tensor,
x_pde: torch.Tensor,
y_pde: torch.Tensor,
k_pde: torch.Tensor,
x_bc: torch.Tensor,
y_bc: torch.Tensor,
k_bc: torch.Tensor
) -> Dict[str, torch.Tensor]:
"""Compute complete PINN loss."""
u_pred = mannequin(x_data, y_data, k_data)
loss_data = F.mse_loss(u_pred, u_data)
residual = self.pde_residual(mannequin, x_pde, y_pde, k_pde)
loss_pde = torch.imply(residual ** 2)
u_bc_pred = mannequin(x_bc, y_bc, k_bc)
loss_bc = torch.imply(u_bc_pred ** 2)
total_loss = (
self.lambda_data * loss_data +
self.lambda_pde * loss_pde +
self.lambda_bc * loss_bc
)
return {
'complete': total_loss,
'information': loss_data,
'pde': loss_pde,
'bc': loss_bc
}
print("nCreating Physics-Informed Neural Network...")
pinn_model = PINN_MLP(
input_dim=3,
output_dim=1,
hidden_dims=[128, 128, 128, 128],
use_fourier_features=True,
n_frequencies=64
).to(system)
n_params_pinn = sum(p.numel() for p in pinn_model.parameters() if p.requires_grad)
print(f"✓ PINN Model created with {n_params_pinn:,} trainable parameters")
print("n" + "="*80)
print("SECTION 7: TRAINING UTILITIES")
print("="*80)
class Trainer:
"""Generic coach for neural community fashions."""
def __init__(
self,
mannequin: nn.Module,
optimizer: torch.optim.Optimizer,
scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
system: torch.system = system
):
self.mannequin = mannequin
self.optimizer = optimizer
self.scheduler = scheduler
self.system = system
self.historical past = {'train_loss': [], 'val_loss': []}
def train_epoch(
self,
train_loader: DataLoader,
loss_fn: nn.Module = nn.MSELoss()
) -> float:
"""Train for one epoch."""
self.mannequin.practice()
total_loss = 0.0
for batch_x, batch_y in train_loader:
batch_x = batch_x.to(self.system)
batch_y = batch_y.to(self.system)
self.optimizer.zero_grad()
pred = self.mannequin(batch_x)
loss = loss_fn(pred, batch_y)
loss.backward()
self.optimizer.step()
total_loss += loss.merchandise()
return total_loss / len(train_loader)
@torch.no_grad()
def validate(
self,
val_loader: DataLoader,
loss_fn: nn.Module = nn.MSELoss()
) -> float:
"""Validate mannequin."""
self.mannequin.eval()
total_loss = 0.0
for batch_x, batch_y in val_loader:
batch_x = batch_x.to(self.system)
batch_y = batch_y.to(self.system)
pred = self.mannequin(batch_x)
loss = loss_fn(pred, batch_y)
total_loss += loss.merchandise()
return total_loss / len(val_loader)
def practice(
self,
train_loader: DataLoader,
val_loader: DataLoader,
n_epochs: int = 100,
loss_fn: nn.Module = nn.MSELoss(),
verbose: bool = True
):
"""Full coaching loop."""
best_val_loss = float('inf')
pbar = tqdm(vary(n_epochs), desc="Training") if verbose else vary(n_epochs)
for epoch in pbar:
train_loss = self.train_epoch(train_loader, loss_fn)
val_loss = self.validate(val_loader, loss_fn)
self.historical past['train_loss'].append(train_loss)
self.historical past['val_loss'].append(val_loss)
if self.scheduler:
self.scheduler.step()
if val_loss < best_val_loss:
best_val_loss = val_loss
self.best_state = {ok: v.cpu().clone() for ok, v in self.mannequin.state_dict().objects()}
if verbose:
pbar.set_postfix({
'practice': f'{train_loss:.6f}',
'val': f'{val_loss:.6f}',
'greatest': f'{best_val_loss:.6f}'
})
return self.historical past
def plot_training_history(historical past: Dict[str, List[float]], title: str = "Training History"):
"""Plot coaching curves."""
fig, ax = plt.subplots(figsize=(10, 4))
ax.semilogy(historical past['train_loss'], label='Train Loss', linewidth=2)
ax.semilogy(historical past['val_loss'], label='Validation Loss', linewidth=2)
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss (log scale)')
ax.set_title(title)
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('training_history.png', dpi=150, bbox_inches='tight')
plt.present()
print("✓ Saved coaching historical past to 'training_history.png'")We construct the Physics-Informed Neural Network elements and outline find out how to incorporate bodily legal guidelines straight into the educational course of. We create the PINN-based MLP, implement the Darcy PDE loss with information, residual, and boundary phrases, and initialize the PINN mannequin for experimentation. We additionally outline generic coaching utilities, together with the coach class and a loss-plotting perform, so we now have a reusable coaching framework for the fashions within the tutorial.
print("n" + "="*80)
print("SECTION 8: TRAINING FNO MODEL")
print("="*80)
optimizer_fno = torch.optim.AdamW(fno_model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler_fno = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_fno, T_max=100)
fno_trainer = Trainer(fno_model, optimizer_fno, scheduler_fno, system)
print("nTraining FNO mannequin...")
fno_history = fno_trainer.practice(
train_loader,
test_loader,
n_epochs=100,
loss_fn=nn.MSELoss(),
verbose=True
)
plot_training_history(fno_history, "FNO Training History")
print("n" + "="*80)
print("SECTION 9: EVALUATION AND VISUALIZATION")
print("="*80)
@torch.no_grad()
def evaluate_model(
mannequin: nn.Module,
test_loader: DataLoader,
system: torch.system = system
) -> Dict[str, float]:
"""Evaluate mannequin on take a look at set."""
mannequin.eval()
all_preds = []
all_targets = []
for batch_x, batch_y in test_loader:
batch_x = batch_x.to(system)
batch_y = batch_y.to(system)
pred = mannequin(batch_x)
all_preds.append(pred.cpu())
all_targets.append(batch_y.cpu())
preds = torch.cat(all_preds, dim=0)
targets = torch.cat(all_targets, dim=0)
mse = F.mse_loss(preds, targets).merchandise()
rmse = np.sqrt(mse)
mae = F.l1_loss(preds, targets).merchandise()
rel_l2 = torch.norm(preds - targets) / torch.norm(targets)
rel_l2 = rel_l2.merchandise()
return {
'MSE': mse,
'RMSE': rmse,
'MAE': mae,
'Relative L2': rel_l2,
'predictions': preds,
'targets': targets
}
fno_model.load_state_dict({ok: v.to(system) for ok, v in fno_trainer.best_state.objects()})
fno_metrics = evaluate_model(fno_model, test_loader)
print("nFNO Model Evaluation Metrics:")
print(f" MSE: {fno_metrics['MSE']:.6f}")
print(f" RMSE: {fno_metrics['RMSE']:.6f}")
print(f" MAE: {fno_metrics['MAE']:.6f}")
print(f" Relative L2: {fno_metrics['Relative L2']:.4f} ({fno_metrics['Relative L2']*100:.2f}%)")
def visualize_predictions(
mannequin: nn.Module,
test_dataset: Dataset,
n_samples: int = 3,
system: torch.system = system
):
"""Visualize mannequin predictions vs floor fact."""
mannequin.eval()
fig, axes = plt.subplots(n_samples, 4, figsize=(16, 4 * n_samples))
for i in vary(n_samples):
x, y_true = test_dataset[i]
x = x.unsqueeze(0).to(system)
with torch.no_grad():
y_pred = mannequin(x)
x = x.squeeze().cpu().numpy()
y_true = y_true.squeeze().cpu().numpy()
y_pred = y_pred.squeeze().cpu().numpy()
error = np.abs(y_true - y_pred)
im0 = axes[i, 0].imshow(x, cmap='viridis', origin='decrease')
axes[i, 0].set_title(f'Input Permeability (Sample {i+1})')
plt.colorbar(im0, ax=axes[i, 0])
im1 = axes[i, 1].imshow(y_true, cmap='sizzling', origin='decrease')
axes[i, 1].set_title('True Pressure')
plt.colorbar(im1, ax=axes[i, 1])
im2 = axes[i, 2].imshow(y_pred, cmap='sizzling', origin='decrease')
axes[i, 2].set_title('Predicted Pressure')
plt.colorbar(im2, ax=axes[i, 2])
im3 = axes[i, 3].imshow(error, cmap='Reds', origin='decrease')
axes[i, 3].set_title(f'Absolute Error (Max: {error.max():.4f})')
plt.colorbar(im3, ax=axes[i, 3])
for ax_row in axes:
for ax in ax_row:
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
plt.savefig('fno_predictions.png', dpi=150, bbox_inches='tight')
plt.present()
print("✓ Saved predictions to 'fno_predictions.png'")
visualize_predictions(fno_model, test_dataset, n_samples=3)
print("n" + "="*80)
print("SECTION 10: PHYSICSNEMO BUILT-IN MODELS")
print("="*80)
if PHYSICSNEMO_AVAILABLE:
attempt:
from physicsnemo.fashions.fno import FNO
from physicsnemo.fashions.mlp.fully_connected import FullyConnected
print("n✓ PhysicsNeMo fashions accessible!")
pn_fno = FNO(
in_channels=1,
out_channels=1,
decoder_layers=1,
decoder_layer_size=32,
dimension=2,
latent_channels=32,
num_fno_layers=4,
num_fno_modes=12,
padding=5
).to(system)
print(f" PhysicsNeMo FNO created with {sum(p.numel() for p in pn_fno.parameters()):,} parameters")
pn_mlp = FullyConnected(
in_features=32,
out_features=64,
num_layers=4,
layer_size=128
).to(system)
print(f" PhysicsNeMo MLP created with {sum(p.numel() for p in pn_mlp.parameters()):,} parameters")
besides ImportError as e:
print(f"Some PhysicsNeMo fashions not accessible: {e}")
else:
print("nPhysicsNeMo not put in. Install with: pip set up nvidia-physicsnemo")
print("Using customized implementations demonstrated on this tutorial.")We practice the Fourier Neural Operator utilizing the ready dataloaders, optimizer, and scheduler, and we observe its studying progress over epochs. We then consider the skilled FNO mannequin on the take a look at set, compute metrics, and visualize predictions in opposition to the bottom fact to grasp how nicely it solves the Darcy drawback. We additionally test whether or not PhysicsNeMo’s built-in fashions can be found and, in that case, instantiate them to attach our customized implementation to the official framework elements.
print("n" + "="*80)
print("SECTION 11: CONVOLUTIONAL BASELINE MODEL")
print("="*80)
class ConvolutionalSurrogate(nn.Module):
"""
Simple convolutional neural community baseline.
A commonplace U-Net type encoder-decoder for comparability with FNO.
"""
def __init__(
self,
in_channels: int = 1,
out_channels: int = 1,
base_channels: int = 32
):
tremendous().__init__()
self.enc1 = self._conv_block(in_channels, base_channels)
self.enc2 = self._conv_block(base_channels, base_channels * 2)
self.enc3 = self._conv_block(base_channels * 2, base_channels * 4)
self.bottleneck = self._conv_block(base_channels * 4, base_channels * 8)
self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, 2, stride=2)
self.dec3 = self._conv_block(base_channels * 8, base_channels * 4)
self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 2, stride=2)
self.dec2 = self._conv_block(base_channels * 4, base_channels * 2)
self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, 2, stride=2)
self.dec1 = self._conv_block(base_channels * 2, base_channels)
self.output = nn.Conv2nd(base_channels, out_channels, 1)
self.pool = nn.MaxPool2d(2)
def _conv_block(self, in_ch: int, out_ch: int) -> nn.Sequential:
return nn.Sequential(
nn.Conv2nd(in_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2nd(out_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True)
)
def ahead(self, x: torch.Tensor) -> torch.Tensor:
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
e3 = self.enc3(self.pool(e2))
b = self.bottleneck(self.pool(e3))
d3 = self.dec3(torch.cat([self.up3(b), e3], dim=1))
d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
return self.output(d1)
print("nCreating Convolutional baseline mannequin...")
conv_model = ConvolutionalSurrogate(in_channels=1, out_channels=1, base_channels=16).to(system)
n_params_conv = sum(p.numel() for p in conv_model.parameters() if p.requires_grad)
print(f"✓ Conv Model created with {n_params_conv:,} trainable parameters")
optimizer_conv = torch.optim.AdamW(conv_model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler_conv = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_conv, T_max=100)
conv_trainer = Trainer(conv_model, optimizer_conv, scheduler_conv, system)
print("nTraining Convolutional mannequin...")
conv_history = conv_trainer.practice(
train_loader,
test_loader,
n_epochs=100,
loss_fn=nn.MSELoss(),
verbose=True
)
conv_model.load_state_dict({ok: v.to(system) for ok, v in conv_trainer.best_state.objects()})
conv_metrics = evaluate_model(conv_model, test_loader)
print("nConv Model Evaluation Metrics:")
print(f" MSE: {conv_metrics['MSE']:.6f}")
print(f" RMSE: {conv_metrics['RMSE']:.6f}")
print(f" MAE: {conv_metrics['MAE']:.6f}")
print(f" Relative L2: {conv_metrics['Relative L2']:.4f} ({conv_metrics['Relative L2']*100:.2f}%)")
print("n" + "="*80)
print("SECTION 12: MODEL COMPARISON")
print("="*80)
def compare_models(
fashions: Dict[str, Tuple[nn.Module, Dict]],
test_loader: DataLoader,
system: torch.system = system
):
"""Compare a number of fashions."""
outcomes = {}
print("n" + "-"*60)
print(f"{'Model':<20} {'MSE':<12} {'RMSE':<12} {'Rel. L2':<12} {'Params':<12}")
print("-"*60)
for identify, (mannequin, trainer_state) in fashions.objects():
mannequin.load_state_dict({ok: v.to(system) for ok, v in trainer_state.objects()})
metrics = evaluate_model(mannequin, test_loader)
n_params = sum(p.numel() for p in mannequin.parameters())
outcomes[name] = metrics
print(f"{identify:<20} {metrics['MSE']:<12.6f} {metrics['RMSE']:<12.6f} "
f"{metrics['Relative L2']:<12.4f} {n_params:<12,}")
print("-"*60)
return outcomes
models_to_compare = {
'FNO': (fno_model, fno_trainer.best_state),
'Conv U-Net': (conv_model, conv_trainer.best_state)
}
comparison_results = compare_models(models_to_compare, test_loader)
def plot_comparison(outcomes: Dict):
"""Plot mannequin comparability."""
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
fashions = record(outcomes.keys())
metrics = ['MSE', 'RMSE', 'Relative L2']
colours = plt.cm.Set2(np.linspace(0, 1, len(fashions)))
for i, metric in enumerate(metrics):
values = [results[m][metric] for m in fashions]
axes[i].bar(fashions, values, shade=colours)
axes[i].set_title(metric)
axes[i].set_ylabel('Value')
axes[i].grid(axis='y', alpha=0.3)
for j, v in enumerate(values):
axes[i].textual content(j, v + 0.01 * max(values), f'{v:.4f}',
ha='heart', va='backside', fontsize=10)
plt.tight_layout()
plt.savefig('model_comparison.png', dpi=150, bbox_inches='tight')
plt.present()
print("✓ Saved comparability to 'model_comparison.png'")
plot_comparison(comparison_results)We implement a convolutional surrogate baseline mannequin to match an ordinary deep studying structure with the Fourier Neural Operator. We practice this convolutional mannequin, consider its efficiency, and print the primary error metrics to see how nicely it approximates the stress discipline. We then examine the FNO and convolutional fashions facet by facet and visualize their metrics, which helps us perceive the strengths of operator studying versus typical convolutional studying.
print("n" + "="*80)
print("SECTION 13: INFERENCE AND DEPLOYMENT")
print("="*80)
@torch.no_grad()
def inference_single(
mannequin: nn.Module,
permeability: np.ndarray,
system: torch.system = system
) -> np.ndarray:
"""
Run inference on a single permeability discipline.
Args:
mannequin: Trained mannequin
permeability: 2D numpy array of permeability values
system: Device to run inference on
Returns:
Predicted stress discipline as numpy array
"""
mannequin.eval()
perm_mean = train_dataset.perm_mean
perm_std = train_dataset.perm_std
x = (torch.FloatTensor(permeability) - perm_mean) / perm_std
x = x.unsqueeze(0).unsqueeze(0).to(system)
pred = mannequin(x)
press_mean = train_dataset.press_mean
press_std = train_dataset.press_std
pred = pred.squeeze().cpu().numpy()
pred = pred * press_std.numpy() + press_mean.numpy()
return pred
def benchmark_inference(
mannequin: nn.Module,
n_samples: int = 100,
decision: int = 32,
system: torch.system = system
) -> Dict[str, float]:
"""Benchmark inference velocity."""
import time
mannequin.eval()
x = torch.randn(n_samples, 1, decision, decision).to(system)
with torch.no_grad():
for _ in vary(10):
_ = mannequin(x[:1])
if torch.cuda.is_available():
torch.cuda.synchronize()
begin = time.time()
with torch.no_grad():
for i in vary(n_samples):
_ = mannequin(x[i:i+1])
if torch.cuda.is_available():
torch.cuda.synchronize()
total_time = time.time() - begin
return {
'total_time': total_time,
'time_per_sample': total_time / n_samples,
'samples_per_second': n_samples / total_time
}
print("nBenchmarking inference velocity...")
print("-" * 50)
for identify, (mannequin, _) in models_to_compare.objects():
benchmark = benchmark_inference(mannequin, n_samples=100)
print(f"{identify}:")
print(f" Time per pattern: {benchmark['time_per_sample']*1000:.3f} ms")
print(f" Throughput: {benchmark['samples_per_second']:.1f} samples/sec")
print("nnDemonstrating single pattern inference...")
test_perm = perm_test[0]
predicted_pressure = inference_single(fno_model, test_perm)
print(f"Input permeability form: {test_perm.form}")
print(f"Output stress form: {predicted_pressure.form}")
print(f"Pressure vary: [{predicted_pressure.min():.4f}, {predicted_pressure.max():.4f}]")
print("n" + "="*80)
print("SECTION 14: SAVING AND LOADING MODELS")
print("="*80)
def save_model(
mannequin: nn.Module,
path: str,
optimizer: torch.optim.Optimizer = None,
metadata: Dict = None
):
"""Save mannequin checkpoint."""
checkpoint = {
'model_state_dict': mannequin.state_dict(),
'model_class': mannequin.__class__.__name__,
}
if optimizer:
checkpoint['optimizer_state_dict'] = optimizer.state_dict()
if metadata:
checkpoint['metadata'] = metadata
torch.save(checkpoint, path)
print(f"✓ Model saved to {path}")
def load_model(
mannequin: nn.Module,
path: str,
optimizer: torch.optim.Optimizer = None,
system: torch.system = system
) -> Dict:
"""Load mannequin checkpoint."""
checkpoint = torch.load(path, map_location=system)
mannequin.load_state_dict(checkpoint['model_state_dict'])
if optimizer and 'optimizer_state_dict' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
print(f"✓ Model loaded from {path}")
return checkpoint.get('metadata', {})
save_model(
fno_model,
'fno_darcy_best.pt',
optimizer_fno,
metadata={
'epochs_trained': 100,
'test_rmse': fno_metrics['RMSE'],
'decision': 32
}
)
print("n" + "="*80)
print("SECTION 15: SUMMARY AND NEXT STEPS")
print("="*80)
print("""
╔══════════════════════════════════════════════════════════════════════════════╗
║ TUTORIAL SUMMARY ║
╠══════════════════════════════════════════════════════════════════════════════╣
║ ║
║ This tutorial coated: ║
║ ║
║ 1. ✓ Installing and organising PhysicsNeMo on Google Colab ║
║ 2. ✓ Understanding the 2D Darcy Flow drawback (PDE for subsurface circulation) ║
║ 3. ✓ Generating artificial coaching information ║
║ 4. ✓ Implementing Fourier Neural Operators (FNO) from scratch ║
║ 5. ✓ Understanding Physics-Informed Neural Networks (PINNs) ║
║ 6. ✓ Training neural operators for operator studying ║
║ 7. ✓ Comparing totally different mannequin architectures ║
║ 8. ✓ Benchmarking inference velocity ║
║ 9. ✓ Saving and loading skilled fashions ║
║ ║
╠══════════════════════════════════════════════════════════════════════════════╣
║ NEXT STEPS ║
╠══════════════════════════════════════════════════════════════════════════════╣
║ ║
║ → Try PhysicsNeMo's built-in fashions for higher efficiency ║
║ → Explore different neural operators: DeepONet, GraphNet ║
║ → Apply to real-world datasets from CFD simulations ║
║ → Implement physics-informed coaching for higher generalization ║
║ → Use distributed coaching for bigger fashions ║
║ ║
║ Resources: ║
║ • PhysicsNeMo Docs: https://docs.nvidia.com/physicsnemo ║
║ • GitHub: https://github.com/NVIDIA/physicsnemo ║
║ • Examples: https://github.com/NVIDIA/physicsnemo/tree/fundamental/examples ║
║ ║
╚══════════════════════════════════════════════════════════════════════════════╝
""")
print("n✓ Tutorial accomplished efficiently!")
print(f" Total GPU reminiscence used: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB" if torch.cuda.is_available() else " (No GPU used)")We focus on inference, deployment, and mannequin persistence so the tutorial strikes past coaching into sensible utilization. We benchmark inference velocity, run single-sample prediction, and examine the form and vary of the generated output to confirm that the skilled mannequin behaves as anticipated. Also, we save the most effective mannequin checkpoint and shut the tutorial with a abstract of what we achieved and the place we are able to go subsequent with PhysicsNeMo and physics-informed machine studying.
In conclusion, we constructed a full PhysicsNeMo-inspired pipeline to unravel a consultant scientific machine studying drawback in an interactive Colab setting. We generated artificial PDE information, skilled neural operator fashions, examined prediction high quality, in contrast totally different approaches, and ready the fashions for inference and reuse. Also, we strengthened our understanding of operator studying, surrogate modeling, and physics-informed deep studying, and we noticed how these concepts connect with real-world simulation duties. This tutorial offers a powerful basis for additional exploration with PhysicsNeMo, whether or not we wish to experiment with bigger datasets, extra superior neural operators, or extra practical physics-driven purposes.
Check out the Full Implementation Codes here. Also, be happy to observe us on Twitter and don’t overlook to hitch our 130k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to companion with us for selling your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar and many others.? Connect with us
The submit A Step-by-Step Coding Tutorial on NVIDIA PhysicsNeMo: Darcy Flow, FNOs, PINNs, Surrogate Models, and Inference Benchmarking appeared first on MarkTechPost.
