Chapter 6: Physics-Informed AI and Simulation

•   u(x,t) = temperature
•   x = position
•   t = time
•   α = thermal diffusivity (material constant)

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

class PINN(nn.Module):
    """Physics-Informed Neural Network for PDEs"""
    # Input dimension = 2 (because input is (x, t))
    # 4 hidden layers, each 50 neurons
    # Output dimension = 1 (temperature u)
    def __init__(self, layers=[2, 50, 50, 50, 50, 1]):
        super().__init__()
        
        self.activation = nn.Tanh()
        self.layers = nn.ModuleList()
        
        # Build network
        for i in range(len(layers) - 1):
            self.layers.append(nn.Linear(layers[i], layers[i+1]))
    
    def forward(self, x, t):
        """
        x: spatial coordinate
        t: time coordinate
        """
        inputs = torch.cat([x, t], dim=1)
        #This forms a 2-column tensor [[x, t], [x, t], ...].

        # Forward pass through hidden layers
        for layer in self.layers[:-1]:
            inputs = self.activation(layer(inputs))
        
        # Output layer (no activation)
        return self.layers[-1](inputs)

def compute_pde_residual(model, x, t, alpha=0.01):
    """
    Compute PDE residual for heat equation:
    ∂u/∂t = α ∂²u/∂x²
    """
    #This is the heart of PINNs: use autograd to compute derivatives, then penalize PDE violations.

    u = model(x, t)
    
    # Compute derivatives using autograd
    u_t = torch.autograd.grad(
        u, t,
        torch.ones_like(u),
        create_graph=True,
        retain_graph=True
    )[0]
    #compute derivatives using PyTorch autograd
    #Because x_colloc and t_colloc were created with requires_grad=True, PyTorch can differentiate through the network.

    
    u_x = torch.autograd.grad(
        u, x,
        torch.ones_like(u),
        create_graph=True,
        retain_graph=True
    )[0]
    
    u_xx = torch.autograd.grad(
        u_x, x,
        torch.ones_like(u_x),
        create_graph=True,
        retain_graph=True
    )[0]
    
    # PDE residual
    residual = u_t - alpha * u_xx
    
    return (residual ** 2).mean()

def analytical_solution(x, t, alpha=0.01):
    """Analytical solution for validation"""
    return np.sin(np.pi * x) * np.exp(-np.pi**2 * alpha * t)

# Initialize model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
pinn = PINN().to(device)
optimizer = torch.optim.Adam(pinn.parameters(), lr=1e-3)
#torch.optim.Adam   •   Adam = Adaptive Moment Estimation
print(f"PINN parameters: {sum(p.numel() for p in pinn.parameters()):,}")
# Output: PINN parameters: 7,851

# Generate training points
n_colloc, n_bc = 5000, 200

# Collocation points (enforce PDE)
x_colloc = torch.rand(n_colloc, 1, requires_grad=True).to(device)
t_colloc = torch.rand(n_colloc, 1, requires_grad=True).to(device)

# Initial condition points
x_ic = torch.rand(n_bc, 1).to(device)
t_ic = torch.zeros(n_bc, 1).to(device)
u_ic = torch.sin(np.pi * x_ic)

# Boundary condition points
x_bc = torch.cat([torch.zeros(n_bc // 2, 1), torch.ones(n_bc // 2, 1)]).to(device)
t_bc = torch.rand(n_bc, 1).to(device)
u_bc = torch.zeros(n_bc, 1).to(device)

print(f"Training points: {n_colloc + 2*n_bc}")
# Output: Training points: 5400

# Training loop
import time

print("Training PINN...\n")
epochs = 3000
pde_losses, ic_losses, bc_losses = [], [], []
start_time = time.time()

for epoch in range(epochs):
    optimizer.zero_grad()
    
    # Compute losses
    loss_pde = compute_pde_residual(pinn, x_colloc, t_colloc, alpha=0.01)
    loss_ic = torch.nn.functional.mse_loss(pinn(x_ic, t_ic), u_ic)
    loss_bc = torch.nn.functional.mse_loss(pinn(x_bc, t_bc), u_bc)
    
    # Combined loss (weight initial/boundary conditions more)
    loss = loss_pde + 10 * loss_ic + 10 * loss_bc
    
    # Backward pass
    loss.backward()
    optimizer.step()
    
    # Track losses
    pde_losses.append(loss_pde.item())
    ic_losses.append(loss_ic.item())
    bc_losses.append(loss_bc.item())
    
    if (epoch + 1) % 500 == 0:
        print(f"Epoch {epoch+1}: PDE={loss_pde.item():.6f}, "
              f"IC={loss_ic.item():.6f}, BC={loss_bc.item():.6f}")

#Each epoch computes:
#   loss_pde: PDE residual inside the domain
#   loss_ic: mismatch at initial condition
#   loss_bc: mismatch at boundary conditions

training_time = time.time() - start_time
print(f"\n✓ Training complete! Time: {training_time/60:.1f} min")

# Evaluation
'''It creates a meshgrid over [0,1]×[0,1], predicts U_pred, and compares to:
    U_analytical = sin(pi x) * exp(-pi^2 alpha t)
    Then reports:
    •   RMSE
    •   max absolute error
    •   final PDE/IC/BC losses
'''
n_test = 100
x_test = np.linspace(0, 1, n_test)
t_test = np.linspace(0, 1, n_test)
X_test, T_test = np.meshgrid(x_test, t_test)

# PINN predictions
with torch.no_grad():
    X_flat = torch.FloatTensor(X_test.flatten()[:, None]).to(device)
    T_flat = torch.FloatTensor(T_test.flatten()[:, None]).to(device)
    U_pred = pinn(X_flat, T_flat).cpu().numpy().reshape(n_test, n_test)

# Analytical solution
U_analytical = analytical_solution(X_test, T_test, alpha=0.01)
error = np.abs(U_pred - U_analytical)
rmse = np.sqrt(np.mean(error ** 2))

print("\n" + "="*70)
print("PINN RESULTS")
print("="*70)
print(f"Training time: {training_time/60:.1f} minutes")
print(f"Final PDE loss: {pde_losses[-1]:.6f}")
print(f"Final IC loss: {ic_losses[-1]:.6f}")
print(f"Final BC loss: {bc_losses[-1]:.6f}")
print(f"\nRMSE vs analytical: {rmse:.6f}")
print(f"Max error: {error.max():.6f}")
print("="*70)
if rmse < 0.01:
    print("\n✓ EXCELLENT: PINN solves PDE accurately!")
print("="*70)

Training PINN...

Epoch 500: PDE=0.001629, IC=0.000023, BC=0.000059
...
Epoch 3000: PDE=0.000088, IC=0.000003, BC=0.000002

✓ Training complete! Time: 0.4 min

======================================================================
PINN RESULTS
======================================================================
Training time: 0.4 minutes
Final PDE loss: 0.000088
Final IC loss: 0.000003
Final BC loss: 0.000002

RMSE vs analytical: 0.002307
Max error: 0.004397
======================================================================

✓ EXCELLENT: PINN solves PDE accurately!
======================================================================

loss = loss_pde + 10 * loss_ic + 10 * loss_bc

# Without weighting: BC violations may persist
loss = loss_pde + loss_ic + loss_bc  # BC often violated

# With weighting: BC enforced more strictly  
loss = loss_pde + 10 * loss_ic + 10 * loss_bc  # BC satisfied first

# Example: Adaptive weighting implementation
def compute_adaptive_weights(model, losses, epsilon=1e-8):
    """
    Adjust weights based on gradient magnitudes.
    Prevents any single loss from dominating training.
    """
    grads = []
    for loss in losses:
        model.zero_grad()
        loss.backward(retain_graph=True)
        grad_norm = sum(p.grad.norm()**2 for p in model.parameters() 
                       if p.grad is not None)**0.5
        grads.append(grad_norm.item())
    
    max_grad = max(grads)
    weights = [max_grad / (g + epsilon) for g in grads]
    return weights

fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# PINN solution
c1 = axes[0].contourf(X_test, T_test, U_pred, levels=20, cmap='coolwarm')
axes[0].set_title('PINN Solution', fontweight='bold')
axes[0].set_xlabel('x')
axes[0].set_ylabel('t')
plt.colorbar(c1, ax=axes[0])

# Analytical solution
c2 = axes[1].contourf(X_test, T_test, U_analytical, levels=20, cmap='coolwarm')
axes[1].set_title('Analytical Solution', fontweight='bold')
axes[1].set_xlabel('x')
axes[1].set_ylabel('t')
plt.colorbar(c2, ax=axes[1])

# Error
c3 = axes[2].contourf(X_test, T_test, error, levels=20, cmap='Reds')
axes[2].set_title('Absolute Error', fontweight='bold')
axes[2].set_xlabel('x')
axes[2].set_ylabel('t')
plt.colorbar(c3, ax=axes[2])

plt.tight_layout()
plt.savefig('pinn_heat_equation.png', dpi=300, bbox_inches='tight')

### The complete running codes with visulization are available on the Google colab code: https://github.com/jpliu168/Generative_AI_For_Science
###
# ============================================================
# 3. Simple PINN architecture (tanh activation)
# ============================================================

class SimpleNSPINN(nn.Module):
    """
    PINN for 2D incompressible Navier–Stokes.

    Inputs : (x, y, t)
    Outputs: (u, v, p)
    """

    def __init__(self, in_dim=3, out_dim=3, hidden_dim=64, num_layers=4):
        super().__init__()

        layers = [nn.Linear(in_dim, hidden_dim), nn.Tanh()]
        for _ in range(num_layers - 1):
            layers += [nn.Linear(hidden_dim, hidden_dim), nn.Tanh()]
        layers += [nn.Linear(hidden_dim, out_dim)]
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor, y: torch.Tensor, t: torch.Tensor):
        inputs = torch.cat([x, y, t], dim=1)
        out = self.net(inputs)
        u = out[:, 0:1]
        v = out[:, 1:2]
        p = out[:, 2:3]
        return u, v, p


model = SimpleNSPINN().to(device)
print(f"Number of parameters: {sum(p.numel() for p in model.parameters()):,}")

# ============================================================
# 4. Navier–Stokes residuals and physics loss
# ============================================================

def navier_stokes_residuals(model, x, y, t, nu=NU):
    """
    Compute residuals of 2D incompressible Navier–Stokes:

    Momentum (x):
      ∂u/∂t + u ∂u/∂x + v ∂u/∂y = -∂p/∂x + ν(∂²u/∂x² + ∂²u/∂y²)
    Momentum (y):
      ∂v/∂t + u ∂v/∂x + v ∂v/∂y = -∂p/∂y + ν(∂²v/∂x² + ∂²v/∂y²)
    Continuity:
      ∂u/∂x + ∂v/∂y = 0
    """
    x.requires_grad_(True)
    y.requires_grad_(True)
    t.requires_grad_(True)

    u, v, p = model(x, y, t)

    # First derivatives
    u_t = grad(u, t)
    u_x = grad(u, x)
    u_y = grad(u, y)

    v_t = grad(v, t)
    v_x = grad(v, x)
    v_y = grad(v, y)

    p_x = grad(p, x)
    p_y = grad(p, y)

    # Second derivatives
    u_xx = grad(u_x, x)
    u_yy = grad(u_y, y)

    v_xx = grad(v_x, x)
    v_yy = grad(v_y, y)

    # Residuals
    f_u = u_t + u * u_x + v * u_y + p_x - nu * (u_xx + u_yy)
    f_v = v_t + u * v_x + v * v_y + p_y - nu * (v_xx + v_yy)
    f_cont = u_x + v_y

    return f_u, f_v, f_cont

def physics_loss(model, n_collocation=5000):
    """
    Sample collocation points in (x,y,t) ∈ [0,1]^3 and
    compute mean squared residual of NS + continuity.
    """
    x = torch.rand(n_collocation, 1, device=device)
    y = torch.rand(n_collocation, 1, device=device)
    t = torch.rand(n_collocation, 1, device=device)

    f_u, f_v, f_cont = navier_stokes_residuals(model, x, y, t, nu=NU)
    loss_pde = (f_u**2).mean() + (f_v**2).mean() + (f_cont**2).mean()
    return loss_pde

# ============================================================
# 5. Data loss: match analytical solution
# ============================================================

def data_loss(model, n_data=2000):
    """
    Sample random points and penalize deviation from
    the manufactured analytic solution.
    """
    x = torch.rand(n_data, 1)
    y = torch.rand(n_data, 1)
    t = torch.rand(n_data, 1)

    u_true_np, v_true_np, p_true_np = analytical_vortex_solution_np(
        x.numpy(), y.numpy(), t.numpy(), nu=NU
    )

    u_true = torch.tensor(u_true_np, dtype=torch.float32)
    v_true = torch.tensor(v_true_np, dtype=torch.float32)
    p_true = torch.tensor(p_true_np, dtype=torch.float32)

    # Move to device
    x = x.to(device)
    y = y.to(device)
    t = t.to(device)
    u_true = u_true.to(device)
    v_true = v_true.to(device)
    p_true = p_true.to(device)

    u_pred, v_pred, p_pred = model(x, y, t)

    loss_u = (u_pred - u_true)**2
    loss_v = (v_pred - v_true)**2
    loss_p = (p_pred - p_true)**2

    return loss_u.mean() + loss_v.mean() + loss_p.mean()

# ============================================================
# 6. Training loop: data-first, then physics+data
# ============================================================

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.5)

n_epochs = 3000

for epoch in range(1, n_epochs + 1):
    # Phase 1 (0–999): data only, learn the vortex shape
    if epoch < 1000:
        lambda_pde = 0.0
        lambda_data = 1.0
    # Phase 2 (1000+): gently enforce PDE
    else:
        lambda_pde = 0.1
        lambda_data = 1.0

    optimizer.zero_grad()

    loss_dat = data_loss(model, n_data=2000)
    loss_pde = physics_loss(model, n_collocation=4000)

    loss = lambda_pde * loss_pde + lambda_data * loss_dat
    loss.backward()
    optimizer.step()
    scheduler.step()

    if epoch % 100 == 0:
        print(
            f"Epoch {epoch:4d} | "
            f"Total: {loss.item():.4e} | "
            f"PDE: {loss_pde.item():.4e} | "
            f"Data: {loss_dat.item():.4e}"
        )

...

def train_ablation_study(model_class, n_epochs=3000):
    """
    Compare physics-only, data-only, and combined training.
    Returns metrics for each regime.
    """
    results = {}
    
    # === PHYSICS-ONLY: No data supervision ===
    print("Training PHYSICS-ONLY model...")
    model_phys = model_class().to(device)
    optimizer_phys = torch.optim.Adam(model_phys.parameters(), lr=1e-3)
    
    phys_history = {'pde': [], 'test_rmse': []}
    
    for epoch in range(n_epochs):
        optimizer_phys.zero_grad()
        loss_pde = physics_loss(model_phys, n_collocation=5000)
        loss_pde.backward()
        optimizer_phys.step()
        
        if epoch % 100 == 0:
            test_rmse = evaluate_on_test_set(model_phys)
            phys_history['pde'].append(loss_pde.item())
            phys_history['test_rmse'].append(test_rmse)
    
    results['physics_only'] = {
        'final_pde': phys_history['pde'][-1],
        'final_rmse': phys_history['test_rmse'][-1],
        'history': phys_history
    }
    
    # === DATA-ONLY: No physics constraints ===
    print("Training DATA-ONLY model...")
    model_data = model_class().to(device)
    optimizer_data = torch.optim.Adam(model_data.parameters(), lr=1e-3)
    
    data_history = {'data': [], 'test_rmse': [], 'pde_residual': []}
    
    for epoch in range(n_epochs):
        optimizer_data.zero_grad()
        loss_dat = data_loss(model_data, n_data=2000)
        loss_dat.backward()
        optimizer_data.step()
        
        if epoch % 100 == 0:
            test_rmse = evaluate_on_test_set(model_data)
            # Check PDE residual even though not trained on it
            with torch.no_grad():
                pde_res = physics_loss(model_data, n_collocation=1000).item()
            data_history['data'].append(loss_dat.item())
            data_history['test_rmse'].append(test_rmse)
            data_history['pde_residual'].append(pde_res)
    
    results['data_only'] = {
        'final_data': data_history['data'][-1],
        'final_rmse': data_history['test_rmse'][-1],
        'final_pde_residual': data_history['pde_residual'][-1],
        'history': data_history
    }
    
    # === COMBINED: Physics + Data (Phased Training) ===
    print("Training COMBINED model...")
    model_combined = model_class().to(device)
    optimizer_combined = torch.optim.Adam(model_combined.parameters(), lr=1e-3)
    
    combined_history = {'pde': [], 'data': [], 'test_rmse': []}
    
    for epoch in range(n_epochs):
        optimizer_combined.zero_grad()
        
        loss_pde = physics_loss(model_combined, n_collocation=4000)
        loss_dat = data_loss(model_combined, n_data=2000)
        
        # Phased training: data first, then add physics
        if epoch < 1000:
            loss = loss_dat  # Data only initially
        else:
            loss = 0.1 * loss_pde + loss_dat  # Add physics gradually
        
        loss.backward()
        optimizer_combined.step()
        
        if epoch % 100 == 0:
            test_rmse = evaluate_on_test_set(model_combined)
            combined_history['pde'].append(loss_pde.item())
            combined_history['data'].append(loss_dat.item())
            combined_history['test_rmse'].append(test_rmse)
    
    results['combined'] = {
        'final_pde': combined_history['pde'][-1],
        'final_data': combined_history['data'][-1],
        'final_rmse': combined_history['test_rmse'][-1],
        'history': combined_history
    }
    
    return results

# Run ablation study
ablation_results = train_ablation_study(SimpleNSPINN, n_epochs=3000)

# Quantifying data contribution
def compute_data_contribution(phys_only_rmse, combined_rmse):
    """
    Measure how much data improves upon physics-only baseline.
    """
    improvement = (phys_only_rmse - combined_rmse) / phys_only_rmse * 100
    return improvement

# Example calculation:
# Physics-only RMSE: 4.2e-2
# Combined RMSE: 2.3e-3
# Improvement: (4.2e-2 - 2.3e-3) / 4.2e-2 = 94.5%

print(f"Data contribution: {compute_data_contribution(0.042, 0.0023):.1f}% error reduction")
# Output: Data contribution: 94.5% error reduction

def detect_pinn_overfitting(history, patience=100):
    """
    Detect overfitting in PINN training by monitoring
    the divergence between data and physics losses.
    
    Returns list of warning messages if overfitting detected.
    """
    warnings = []
    
    # Check 1: Data loss decreasing while PDE loss increasing
    recent_data = history['data'][-patience:]
    recent_pde = history['pde'][-patience:]
    
    data_trend = np.polyfit(range(len(recent_data)), recent_data, 1)[0]
    pde_trend = np.polyfit(range(len(recent_pde)), recent_pde, 1)[0]
    
    if data_trend < 0 and pde_trend > 0:
        warnings.append("⚠️ Data-physics divergence: model fitting data at expense of physics")
    
    # Check 2: PDE residual much larger than data loss
    if recent_pde[-1] > 100 * recent_data[-1]:
        warnings.append("⚠️ Large PDE residual: model ignoring physics constraints")
    
    # Check 3: Validation loss increasing (if available)
    if 'val_loss' in history:
        recent_val = history['val_loss'][-patience:]
        val_trend = np.polyfit(range(len(recent_val)), recent_val, 1)[0]
        if val_trend > 0:
            warnings.append("⚠️ Validation loss increasing: generalization degrading")
    
    return warnings

Collocation points:    ●     ●     ●     ●     ●
True solution:       ─────────────────────────────
Network solution:    ~~~○~~~○~~~○~~~○~~~○~~~
                        ↑ Physics violated between points

# Problematic: data concentrated in one region
x_data = torch.rand(1000, 1) * 0.3  # Only covers [0, 0.3]
# Network may overfit to this region, fail in [0.3, 1.0]

# Better: data spread across entire domain
x_data = torch.rand(1000, 1)  # Uniform coverage of [0, 1]

# Risky: heavy data weighting may cause overfitting
loss = 0.01 * loss_pde + 100 * loss_data  # Physics underweighted

# More balanced: comparable loss magnitudes
loss = loss_pde + loss_data  # Neither dominates

def create_pinn_data_splits(x_data, y_data, train_ratio=0.7, val_ratio=0.15):
    """
    Create proper data splits for PINN training.
    
    Note: Collocation points (for physics) don't need splitting—they're 
    for enforcing PDEs, not fitting data. Only observation data is split.
    """
    n = len(x_data)
    indices = np.random.permutation(n)
    
    n_train = int(n * train_ratio)
    n_val = int(n * val_ratio)
    
    train_idx = indices[:n_train]
    val_idx = indices[n_train:n_train + n_val]
    test_idx = indices[n_train + n_val:]
    
    return {
        'train': (x_data[train_idx], y_data[train_idx]),
        'val': (x_data[val_idx], y_data[val_idx]),
        'test': (x_data[test_idx], y_data[test_idx]),
    }

# Usage
splits = create_pinn_data_splits(x_observations, y_observations)
print(f"Train: {len(splits['train'][0])}, Val: {len(splits['val'][0])}, Test: {len(splits['test'][0])}")

def train_pinn_with_early_stopping(model, train_data, val_data, 
                                    max_epochs=5000, patience=200):
    """
    Stop training when validation loss stops improving.
    Prevents overfitting to training data.
    """
    best_val_loss = float('inf')
    patience_counter = 0
    best_model_state = None
    
    for epoch in range(max_epochs):
        # Training step
        model.train()
        train_loss = compute_train_loss(model, train_data)
        train_loss.backward()
        optimizer.step()
        
        # Validation evaluation (no gradient)
        model.eval()
        with torch.no_grad():
            val_loss = compute_val_loss(model, val_data)
        
        # Check for improvement
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_model_state = model.state_dict().copy()
            patience_counter = 0
        else:
            patience_counter += 1
        
        # Early stopping check
        if patience_counter >= patience:
            print(f"Early stopping at epoch {epoch}")
            model.load_state_dict(best_model_state)
            break
    
    return model, best_val_loss

class RegularizedPINN(nn.Module):
    """PINN with dropout and weight decay for regularization."""
    
    def __init__(self, layers, dropout_rate=0.1):
        super().__init__()
        self.layers = nn.ModuleList()
        self.dropouts = nn.ModuleList()
        
        for i in range(len(layers) - 1):
            self.layers.append(nn.Linear(layers[i], layers[i+1]))
            if i < len(layers) - 2:  # No dropout on output layer
                self.dropouts.append(nn.Dropout(dropout_rate))
    
    def forward(self, x):
        for i, layer in enumerate(self.layers[:-1]):
            x = torch.tanh(layer(x))
            if self.training:  # Only apply dropout during training
                x = self.dropouts[i](x)
        return self.layers[-1](x)

# Use with weight decay in optimizer
optimizer = torch.optim.Adam(
    model.parameters(), 
    lr=1e-3, 
    weight_decay=1e-4  # L2 regularization
)

def physics_regularization(model, x_colloc):
    """
    Additional physics-based regularization terms.
    These penalize behaviors that are physically unrealistic.
    """
    u, v, p = model(x_colloc[:, 0:1], x_colloc[:, 1:2], x_colloc[:, 2:3])
    
    # Mass conservation (continuity equation)
    u_x = grad(u, x_colloc[:, 0:1])
    v_y = grad(v, x_colloc[:, 1:2])
    continuity_violation = (u_x + v_y) ** 2
    
    # Smoothness penalty (physical solutions are typically smooth)
    # Penalize large second derivatives
    u_xx = grad(u_x, x_colloc[:, 0:1])
    u_yy = grad(grad(u, x_colloc[:, 1:2]), x_colloc[:, 1:2])
    smoothness_penalty = (u_xx ** 2 + u_yy ** 2).mean()
    
    return 0.1 * continuity_violation.mean() + 0.01 * smoothness_penalty

def comprehensive_pinn_validation(model, test_data, physics_bounds, x_colloc):
    """
    Comprehensive validation checking both accuracy and physical consistency.
    
    Returns a report with multiple diagnostic metrics.
    """
    model.eval()
    results = {}
    
    with torch.no_grad():
        # 1. Data accuracy
        pred = model(test_data['x'], test_data['y'], test_data['t'])
        results['rmse'] = torch.sqrt(((pred - test_data['true'])**2).mean()).item()
        
        # 2. Physics consistency (PDE residual on test points)
        residual = compute_pde_residual(model, x_colloc)
        results['pde_residual'] = residual.item()
        
        # 3. Physical bounds check
        results['bounds_violated'] = (
            (pred < physics_bounds['min']).any().item() or 
            (pred > physics_bounds['max']).any().item()
        )
        
        # 4. Conservation law check
        continuity = check_continuity(model, x_colloc)
        results['mass_conservation_error'] = continuity.item()
    
    # Generate report
    print("="*60)
    print("PINN VALIDATION REPORT")
    print("="*60)
    print(f"Test RMSE: {results['rmse']:.6f}")
    print(f"PDE Residual: {results['pde_residual']:.6f}")
    print(f"Physical Bounds Violated: {results['bounds_violated']}")
    print(f"Mass Conservation Error: {results['mass_conservation_error']:.6f}")
    print("="*60)
    
    # Overall assessment
    if results['rmse'] < 0.01 and results['pde_residual'] < 0.001:
        print("✓ Model appears well-trained and physically consistent")
    elif results['rmse'] < 0.01 and results['pde_residual'] > 0.01:
        print("⚠️ WARNING: Low data error but high physics residual - possible overfitting!")
    else:
        print("✗ Model needs further training or architecture adjustment")
    
    return results

url: https://cds.climate.copernicus.eu/api/v2
key: YOUR_UID:YOUR_API_KEY

pip install cdsapi xarray netCDF4

import cdsapi

def download_era5_sample():
    """
    Download ERA5 data for atmospheric PINN training.
    
    This downloads temperature, humidity, and wind components
    on pressure levels for a European domain.
    """
    c = cdsapi.Client()
    
    c.retrieve(
        'reanalysis-era5-pressure-levels',
        {
            'product_type': 'reanalysis',
            'format': 'netcdf',
            'variable': [
                'temperature',
                'specific_humidity',
                'u_component_of_wind',
                'v_component_of_wind',
            ],
            'pressure_level': [
                '500', '700', '850', '925', '1000',
            ],
            'year': '2020',
            'month': '07',
            'day': ['01', '02', '03'],
            'time': ['00:00', '06:00', '12:00', '18:00'],
            'area': [50, -10, 30, 20],  # North, West, South, East
            'grid': [1.0, 1.0],  # 1 degree resolution
        },
        'era5_data.nc'
    )
    print("Download complete: era5_data.nc")

# Run the download (takes a few minutes)
# download_era5_sample()

import xarray as xr
import numpy as np

# Open and inspect the dataset
ds = xr.open_dataset('era5_data.nc')

print("=== ERA5 Dataset Structure ===")
print(f"\nVariables: {list(ds.data_vars)}")
print(f"Dimensions: {dict(ds.dims)}")
print(f"Coordinates: {list(ds.coords)}")

# Examine each dimension
print(f"\n=== Coordinate Details ===")
print(f"Latitude: {ds.latitude.values.min():.1f}°N to {ds.latitude.values.max():.1f}°N")
print(f"Longitude: {ds.longitude.values.min():.1f}°E to {ds.longitude.values.max():.1f}°E")
print(f"Pressure levels: {ds.pressure_level.values} hPa")
print(f"Time steps: {len(ds.valid_time)} ({ds.valid_time.values[0]} to {ds.valid_time.values[-1]})")

# Calculate total data points
n_points = np.prod([ds.dims[d] for d in ['valid_time', 'pressure_level', 'latitude', 'longitude']])
print(f"\nTotal data points: {n_points:,}")

ds.close()

=== ERA5 Dataset Structure ===

Variables: ['t', 'q', 'u', 'v']
Dimensions: {'valid_time': 12, 'pressure_level': 5, 'latitude': 21, 'longitude': 31}
Coordinates: ['number', 'valid_time', 'pressure_level', 'latitude', 'longitude']

=== Coordinate Details ===
Latitude: 30.0°N to 50.0°N
Longitude: -10.0°E to 20.0°E
Pressure levels: [ 500  700  850  925 1000] hPa
Time steps: 12 (2020-07-01T00:00:00 to 2020-07-03T18:00:00)

Total data points: 39,060

import numpy as np
import torch
import xarray as xr

def load_era5_data(filepath='era5_data.nc', n_samples=10000):
    """
    Load ERA5 data and prepare it for PINN training.
    
    Parameters:
    -----------
    filepath : str
        Path to the ERA5 NetCDF file
    n_samples : int
        Number of samples to use (random subset if data is larger)
    
    Returns:
    --------
    dict : Contains normalized inputs, raw outputs, and normalization parameters
    """
    print(f"Loading ERA5 data from {filepath}...")
    ds = xr.open_dataset(filepath)
    
    # Extract coordinates
    lat = ds['latitude'].values
    lon = ds['longitude'].values
    pressure = ds['pressure_level'].values
    time = ds['valid_time'].values
    
    # Extract variables
    T_data = ds['t'].values      # Temperature (K)
    q_data = ds['q'].values      # Specific humidity (kg/kg)
    u_data = ds['u'].values      # U-wind (m/s)
    
    print(f"\nData shapes: T={T_data.shape}, q={q_data.shape}, u={u_data.shape}")
    print(f"Pressure levels: {pressure} hPa")
    
    # Create coordinate meshgrid
    # ERA5 shape: (time, level, lat, lon)
    n_time, n_levels, n_lat, n_lon = T_data.shape
    time_idx = np.arange(n_time)
    
    TIME_IDX, LEVEL, LAT, LON = np.meshgrid(
        time_idx, pressure, lat, lon, indexing='ij'
    )
    
    # Flatten all arrays
    lat_flat = LAT.flatten()
    lon_flat = LON.flatten()
    p_flat = LEVEL.flatten()
    t_flat = TIME_IDX.flatten().astype(float)
    T_flat = T_data.flatten()
    q_flat = q_data.flatten()
    u_flat = u_data.flatten()
    
    # Remove any NaN values
    valid_mask = ~(np.isnan(T_flat) | np.isnan(q_flat) | np.isnan(u_flat))
    lat_flat = lat_flat[valid_mask]
    lon_flat = lon_flat[valid_mask]
    p_flat = p_flat[valid_mask]
    t_flat = t_flat[valid_mask]
    T_flat = T_flat[valid_mask]
    q_flat = q_flat[valid_mask]
    u_flat = u_flat[valid_mask]
    
    total_points = len(lat_flat)
    print(f"Total valid data points: {total_points:,}")
    
    # Random sampling if needed
    if total_points > n_samples:
        print(f"Randomly sampling {n_samples:,} points...")
        idx = np.random.choice(total_points, n_samples, replace=False)
        lat_flat = lat_flat[idx]
        lon_flat = lon_flat[idx]
        p_flat = p_flat[idx]
        t_flat = t_flat[idx]
        T_flat = T_flat[idx]
        q_flat = q_flat[idx]
        u_flat = u_flat[idx]
    
    # Store raw values for physics calculations
    lat_raw = lat_flat.copy()
    lon_raw = lon_flat.copy()
    p_raw = p_flat.copy()
    t_raw = t_flat.copy()
    
    # Normalize inputs to approximately [-1, 1] range
    # This improves neural network training stability
    lat_mean, lat_std = lat_flat.mean(), lat_flat.std()
    lon_mean, lon_std = lon_flat.mean(), lon_flat.std()
    p_mean, p_std = p_flat.mean(), p_flat.std()
    t_mean, t_std = t_flat.mean(), max(t_flat.std(), 1.0)
    
    lat_norm = (lat_flat - lat_mean) / lat_std
    lon_norm = (lon_flat - lon_mean) / lon_std
    p_norm = (p_flat - p_mean) / p_std
    t_norm = (t_flat - t_mean) / t_std
    
    # Store normalization parameters for later use
    norm_params = {
        'lat': (lat_mean, lat_std),
        'lon': (lon_mean, lon_std),
        'p': (p_mean, p_std),
        't': (t_mean, t_std),
        'T': (T_flat.mean(), T_flat.std()),
        'q': (q_flat.mean(), q_flat.std()),
        'u': (u_flat.mean(), u_flat.std()),
    }
    
    print(f"\n=== Data Statistics ===")
    print(f"Temperature: {T_flat.mean():.1f} ± {T_flat.std():.1f} K")
    print(f"Humidity: {q_flat.mean()*1000:.2f} ± {q_flat.std()*1000:.2f} g/kg")
    print(f"Wind: {u_flat.mean():.1f} ± {u_flat.std():.1f} m/s")
    
    ds.close()
    
    return {
        'lat': lat_norm,
        'lon': lon_norm,
        'p': p_norm,
        't': t_norm,
        'T': T_flat,
        'q': q_flat,
        'u': u_flat,
        'norm_params': norm_params,
        'lat_raw': lat_raw,
        'lon_raw': lon_raw,
        'p_raw': p_raw,
        't_raw': t_raw,
    }

import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class AtmosphericPINN(nn.Module):
    """
    Physics-Informed Neural Network for Atmospheric Variables.
    
    Inputs: (latitude, longitude, pressure, time) - all normalized
    Outputs: (Temperature, Specific Humidity, Zonal Wind)
    
    The network learns a continuous mapping from space-time coordinates
    to atmospheric state variables, constrained by physical laws.
    """
    
    def __init__(self, hidden_dim=128, num_layers=4):
        super().__init__()
        
        # Build the network with tanh activations
        # Tanh works well for PINNs as it's smooth and bounded
        layers = [nn.Linear(4, hidden_dim), nn.Tanh()]
        for _ in range(num_layers - 1):
            layers += [nn.Linear(hidden_dim, hidden_dim), nn.Tanh()]
        layers += [nn.Linear(hidden_dim, 3)]
        
        self.net = nn.Sequential(*layers)
        
        # Learnable output scaling parameters
        # These help the network output values in the correct physical range
        self.T_scale = nn.Parameter(torch.tensor(20.0))   # Temperature range ~20K
        self.T_offset = nn.Parameter(torch.tensor(280.0)) # Base temperature
        self.q_scale = nn.Parameter(torch.tensor(0.01))   # Humidity range
        self.u_scale = nn.Parameter(torch.tensor(10.0))   # Wind range
        
    def forward(self, lat, lon, p, t):
        """
        Forward pass: coordinates → atmospheric variables
        
        Parameters:
        -----------
        lat, lon, p, t : torch.Tensor
            Normalized coordinates, each shape (N, 1)
            
        Returns:
        --------
        T : Temperature in Kelvin
        q : Specific humidity in kg/kg
        u : Zonal wind in m/s
        """
        inputs = torch.cat([lat, lon, p, t], dim=1)
        outputs = self.net(inputs)
        
        # Apply output scaling
        T = outputs[:, 0:1] * self.T_scale + self.T_offset
        q = torch.sigmoid(outputs[:, 1:2]) * self.q_scale  # Bounded [0, q_scale]
        u = outputs[:, 2:3] * self.u_scale
        
        return T, q, u

def grad(outputs, inputs):
    """
    Compute gradient using automatic differentiation.
    This is essential for calculating physics residuals.
    """
    return torch.autograd.grad(
        outputs, inputs,
        grad_outputs=torch.ones_like(outputs),
        create_graph=True,
        retain_graph=True,
    )[0]


def atmospheric_physics_loss(model, lat, lon, p, t, p_raw, norm_params):
    """
    Compute physics-based loss terms for atmospheric constraints.
    
    Parameters:
    -----------
    model : AtmosphericPINN
        The neural network model
    lat, lon, p, t : torch.Tensor
        Normalized input coordinates (require gradients)
    p_raw : torch.Tensor
        Raw pressure values in hPa (for physics calculations)
    norm_params : dict
        Normalization parameters for coordinate transformation
        
    Returns:
    --------
    total_loss : torch.Tensor
        Combined physics loss
    loss_dict : dict
        Individual loss components for monitoring
    """
    # Enable gradient computation for inputs
    lat = lat.requires_grad_(True)
    lon = lon.requires_grad_(True)
    p = p.requires_grad_(True)
    t = t.requires_grad_(True)
    
    # Forward pass
    T, q, u = model(lat, lon, p, t)
    
    # Physical constants
    g = 9.81    # Gravitational acceleration (m/s²)
    R = 287.0   # Gas constant for dry air (J/(kg·K))
    
    # Get normalization scales
    p_mean, p_std = norm_params['p']
    t_mean, t_std = norm_params['t']
    lon_mean, lon_std = norm_params['lon']
    
    # === 1. HYDROSTATIC BALANCE ===
    # ∂T/∂p should follow: -(g/R) * (T/p)
    T_p = grad(T, p)
    T_p_physical = T_p / p_std  # Convert to physical units
    
    p_pa = p_raw * 100  # hPa to Pa
    hydrostatic_expected = -(g / R) * (T / p_pa.unsqueeze(1))
    hydrostatic_residual = T_p_physical - hydrostatic_expected
    
    # === 2. ENERGY CONSERVATION ===
    # ∂T/∂t + u·∂T/∂lon ≈ 0
    T_t = grad(T, t)
    T_lon = grad(T, lon)
    
    # Scale factors for physical units
    # Assume time steps are ~6 hours
    t_scale = t_std * 6 * 3600  # Convert to seconds
    lon_scale = lon_std * 111000 * np.cos(np.radians(40))  # Degrees to meters
    
    energy_residual = T_t / t_scale + u * T_lon / lon_scale
    
    # === 3. MOISTURE CONSERVATION ===
    # ∂q/∂t + u·∂q/∂lon ≈ 0
    q_t = grad(q, t)
    q_lon = grad(q, lon)
    
    moisture_residual = q_t / t_scale + u * q_lon / lon_scale
    
    # Compute MSE losses with appropriate scaling
    loss_hydro = (hydrostatic_residual ** 2).mean()
    loss_energy = (energy_residual ** 2).mean() * 1e6
    loss_moisture = (moisture_residual ** 2).mean() * 1e8
    
    total_loss = loss_hydro + loss_energy + loss_moisture
    
    return total_loss, {
        'hydrostatic': loss_hydro.item(),
        'energy': loss_energy.item(),
        'moisture': loss_moisture.item(),
    }

def train_atmospheric_pinn(data, n_epochs=3000):
    """
    Train the Atmospheric PINN with phased physics introduction.
    
    Parameters:
    -----------
    data : dict
        Output from load_era5_data()
    n_epochs : int
        Total training epochs
        
    Returns:
    --------
    model : AtmosphericPINN
        Trained model
    history : dict
        Training history for analysis
    """
    # Convert data to tensors
    lat = torch.tensor(data['lat'], dtype=torch.float32).unsqueeze(1).to(device)
    lon = torch.tensor(data['lon'], dtype=torch.float32).unsqueeze(1).to(device)
    p = torch.tensor(data['p'], dtype=torch.float32).unsqueeze(1).to(device)
    t = torch.tensor(data['t'], dtype=torch.float32).unsqueeze(1).to(device)
    
    T_true = torch.tensor(data['T'], dtype=torch.float32).unsqueeze(1).to(device)
    q_true = torch.tensor(data['q'], dtype=torch.float32).unsqueeze(1).to(device)
    u_true = torch.tensor(data['u'], dtype=torch.float32).unsqueeze(1).to(device)
    
    p_raw = torch.tensor(data['p_raw'], dtype=torch.float32).to(device)
    norm_params = data['norm_params']
    
    # Initialize model and optimizer
    model = AtmosphericPINN(hidden_dim=128, num_layers=4).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.5)
    
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {n_params:,}")
    print(f"Training samples: {len(lat):,}")
    print("=" * 70)
    
    history = {'total': [], 'data': [], 'physics': [], 'lambda': []}
    
    for epoch in range(1, n_epochs + 1):
        # Sample mini-batch
        idx = torch.randperm(len(lat))[:2000]
        lat_batch = lat[idx]
        lon_batch = lon[idx]
        p_batch = p[idx]
        t_batch = t[idx]
        T_batch = T_true[idx]
        q_batch = q_true[idx]
        u_batch = u_true[idx]
        p_raw_batch = p_raw[idx]
        
        # Phased physics introduction
        if epoch <= 500:
            # Phase 1: Data only
            lambda_physics = 0.0
        elif epoch <= 1500:
            # Phase 2: Gradual introduction
            lambda_physics = 0.01 * (epoch - 500) / 1000
        else:
            # Phase 3: Full physics
            lambda_physics = 0.01
        
        optimizer.zero_grad()
        
        # === Data Loss ===
        T_pred, q_pred, u_pred = model(lat_batch, lon_batch, p_batch, t_batch)
        
        loss_T = ((T_pred - T_batch) ** 2).mean()
        loss_q = ((q_pred - q_batch) ** 2).mean() * 1e4  # Scale up small values
        loss_u = ((u_pred - u_batch) ** 2).mean()
        
        data_loss = loss_T + loss_q + loss_u
        
        # === Physics Loss ===
        if lambda_physics > 0:
            colloc_idx = torch.randperm(len(lat))[:1000]
            physics_loss, _ = atmospheric_physics_loss(
                model,
                lat[colloc_idx], lon[colloc_idx],
                p[colloc_idx], t[colloc_idx],
                p_raw[colloc_idx], norm_params
            )
        else:
            physics_loss = torch.tensor(0.0)
        
        # === Total Loss ===
        total_loss = data_loss + lambda_physics * physics_loss
        
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        
        # Record history
        history['total'].append(total_loss.item())
        history['data'].append(data_loss.item())
        history['physics'].append(physics_loss.item() if isinstance(physics_loss, torch.Tensor) else 0)
        history['lambda'].append(lambda_physics)
        
        # Logging
        if epoch % 200 == 0:
            print(f"Epoch {epoch:4d} | "
                  f"Total: {total_loss.item():.4e} | "
                  f"Data: {data_loss.item():.4e} | "
                  f"Physics: {physics_loss.item() if isinstance(physics_loss, torch.Tensor) else 0:.4e} | "
                  f"λ: {lambda_physics:.4f}")
    
    print("=" * 70)
    print("Training complete!")
    
    return model, history

import matplotlib.pyplot as plt

def evaluate_and_plot(model, data, pressure_level=850):
    """
    Evaluate the trained PINN and create visualizations.
    
    Parameters:
    -----------
    model : AtmosphericPINN
        Trained model
    data : dict
        Data dictionary with normalization parameters
    pressure_level : float
        Pressure level (hPa) for horizontal plots
    """
    model.eval()
    norm_params = data['norm_params']
    
    # Get normalization parameters
    lat_mean, lat_std = norm_params['lat']
    lon_mean, lon_std = norm_params['lon']
    p_mean, p_std = norm_params['p']
    
    # === HORIZONTAL MAPS ===
    # Create grid at specified pressure level
    n_grid = 50
    lat_raw = np.linspace(30, 50, n_grid)
    lon_raw = np.linspace(-10, 20, n_grid)
    LAT_raw, LON_raw = np.meshgrid(lat_raw, lon_raw)
    
    # Normalize coordinates
    lat_norm = (LAT_raw - lat_mean) / lat_std
    lon_norm = (LON_raw - lon_mean) / lon_std
    p_norm = (pressure_level - p_mean) / p_std
    
    # Create tensors
    lat_eval = torch.tensor(lat_norm.flatten(), dtype=torch.float32).unsqueeze(1).to(device)
    lon_eval = torch.tensor(lon_norm.flatten(), dtype=torch.float32).unsqueeze(1).to(device)
    p_eval = torch.full_like(lat_eval, p_norm)
    t_eval = torch.zeros_like(lat_eval)
    
    with torch.no_grad():
        T_pred, q_pred, u_pred = model(lat_eval, lon_eval, p_eval, t_eval)
    
    T_grid = T_pred.cpu().numpy().reshape(n_grid, n_grid)
    q_grid = q_pred.cpu().numpy().reshape(n_grid, n_grid) * 1000  # Convert to g/kg
    u_grid = u_pred.cpu().numpy().reshape(n_grid, n_grid)
    
    # Plot horizontal maps
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    
    im1 = axes[0].contourf(LON_raw, LAT_raw, T_grid, levels=20, cmap='RdYlBu_r')
    axes[0].set_xlabel('Longitude (°E)')
    axes[0].set_ylabel('Latitude (°N)')
    axes[0].set_title('Temperature (K)')
    plt.colorbar(im1, ax=axes[0])
    
    im2 = axes[1].contourf(LON_raw, LAT_raw, q_grid, levels=20, cmap='YlGnBu')
    axes[1].set_xlabel('Longitude (°E)')
    axes[1].set_ylabel('Latitude (°N)')
    axes[1].set_title('Specific Humidity (g/kg)')
    plt.colorbar(im2, ax=axes[1])
    
    im3 = axes[2].contourf(LON_raw, LAT_raw, u_grid, levels=20, cmap='coolwarm')
    axes[2].set_xlabel('Longitude (°E)')
    axes[2].set_ylabel('Latitude (°N)')
    axes[2].set_title('Zonal Wind (m/s)')
    plt.colorbar(im3, ax=axes[2])
    
    plt.suptitle(f'Atmospheric PINN Predictions at {pressure_level} hPa', fontsize=14)
    plt.tight_layout()
    plt.savefig('atmospheric_pinn_maps.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # === VERTICAL PROFILES ===
    p_levels = np.linspace(300, 1000, 50)
    p_norm_profile = (p_levels - p_mean) / p_std
    
    lat_point = torch.zeros(50, 1, device=device)
    lon_point = torch.zeros(50, 1, device=device)
    p_profile = torch.tensor(p_norm_profile, dtype=torch.float32).unsqueeze(1).to(device)
    t_point = torch.zeros(50, 1, device=device)
    
    with torch.no_grad():
        T_profile, q_profile, u_profile = model(lat_point, lon_point, p_profile, t_point)
    
    fig2, axes2 = plt.subplots(1, 3, figsize=(12, 4))
    
    axes2[0].plot(T_profile.cpu().numpy(), p_levels, 'b-', linewidth=2)
    axes2[0].invert_yaxis()
    axes2[0].set_xlabel('Temperature (K)')
    axes2[0].set_ylabel('Pressure (hPa)')
    axes2[0].set_title('Temperature Profile')
    axes2[0].grid(True, alpha=0.3)
    
    axes2[1].plot(q_profile.cpu().numpy() * 1000, p_levels, 'g-', linewidth=2)
    axes2[1].invert_yaxis()
    axes2[1].set_xlabel('Specific Humidity (g/kg)')
    axes2[1].set_ylabel('Pressure (hPa)')
    axes2[1].set_title('Humidity Profile')
    axes2[1].grid(True, alpha=0.3)
    
    axes2[2].plot(u_profile.cpu().numpy(), p_levels, 'r-', linewidth=2)
    axes2[2].invert_yaxis()
    axes2[2].set_xlabel('Zonal Wind (m/s)')
    axes2[2].set_ylabel('Pressure (hPa)')
    axes2[2].set_title('Wind Profile')
    axes2[2].grid(True, alpha=0.3)
    
    plt.suptitle('Vertical Profiles at Domain Center', fontsize=14)
    plt.tight_layout()
    plt.savefig('atmospheric_pinn_profiles.png', dpi=300, bbox_inches='tight')
    plt.show()