Neural Operators (FNO / DeepONet)

pip install neuraloperator torch

import torch
import numpy as np
from neuraloperator.models import FNO1d
from neuraloperator.datasets import load_darcy_flow_small
import matplotlib.pyplot as plt

# Generate training data: Burgers equation solutions
# u_t + u*u_x = nu*u_xx, x in [0, 2pi], periodic BC
from scipy.integrate import solve_ivp
from scipy.fft import fft, ifft, fftfreq

def solve_burgers(u0, nu=0.01, T=1.0, N=256, dt=0.001):
    """Solve Burgers equation using pseudospectral method."""
    dx = 2*np.pi / N
    x = np.linspace(0, 2*np.pi, N, endpoint=False)
    k = fftfreq(N, d=dx) * 2*np.pi

    def rhs(t, u_hat):
        u = np.real(ifft(u_hat))
        u_x = np.real(ifft(1j * k * u_hat))
        return -fft(u * u_x) - nu * k**2 * u_hat

    u0_hat = fft(u0)
    sol = solve_ivp(rhs, (0, T), u0_hat, method='RK45',
                    rtol=1e-8, atol=1e-10)
    return np.real(ifft(sol.y[:, -1]))

# Generate dataset
N = 256
n_train = 500
n_test = 100
x = np.linspace(0, 2*np.pi, N, endpoint=False)

# Random initial conditions (sum of random Fourier modes)
np.random.seed(42)
inputs = []
outputs = []
for i in range(n_train + n_test):
    # Random IC: sum of low-frequency modes
    u0 = np.zeros(N)
    for k in range(1, 8):
        u0 += np.random.randn() * np.sin(k*x) + np.random.randn() * np.cos(k*x)
    u0 *= 0.5
    u_final = solve_burgers(u0)
    inputs.append(u0)
    outputs.append(u_final)

inputs = np.array(inputs)
outputs = np.array(outputs)

# Convert to PyTorch tensors
x_train = torch.tensor(inputs[:n_train], dtype=torch.float32).unsqueeze(-1)
y_train = torch.tensor(outputs[:n_train], dtype=torch.float32).unsqueeze(-1)
x_test = torch.tensor(inputs[n_train:], dtype=torch.float32).unsqueeze(-1)
y_test = torch.tensor(outputs[n_train:], dtype=torch.float32).unsqueeze(-1)

print(f"Training data: {x_train.shape} → {y_train.shape}")
print(f"Test data: {x_test.shape} → {y_test.shape}")

# Define FNO model
model = FNO1d(
    n_modes_height=16,        # Fourier modes to keep
    hidden_channels=64,       # channel width
    in_channels=1,            # input function dimension
    out_channels=1,           # output function dimension
    n_layers=4,               # number of FNO layers
)

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

# Training loop
n_epochs = 500
batch_size = 32

for epoch in range(n_epochs):
    model.train()
    perm = torch.randperm(n_train)
    total_loss = 0
    n_batches = 0

    for i in range(0, n_train, batch_size):
        idx = perm[i:i+batch_size]
        x_batch = x_train[idx]
        y_batch = y_train[idx]

        pred = model(x_batch)
        loss = torch.nn.functional.mse_loss(pred, y_batch)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        n_batches += 1

    scheduler.step()

    if (epoch + 1) % 50 == 0:
        # Test error
        model.eval()
        with torch.no_grad():
            pred_test = model(x_test)
            test_loss = torch.nn.functional.mse_loss(pred_test, y_test)
            # Relative L2 error
            rel_err = torch.mean(
                torch.norm(pred_test - y_test, dim=1) / torch.norm(y_test, dim=1)
            )
        print(f"Epoch {epoch+1}: train_loss={total_loss/n_batches:.4e}, "
              f"test_loss={test_loss:.4e}, rel_L2={rel_err:.4f}")

# DeepONet: separate networks for input function (branch) and query location (trunk)

class DeepONet(torch.nn.Module):
    def __init__(self, branch_input_dim, trunk_input_dim, hidden_dim=128, p=64):
        super().__init__()
        # Branch net: processes input function (sampled at fixed sensors)
        self.branch = torch.nn.Sequential(
            torch.nn.Linear(branch_input_dim, hidden_dim),
            torch.nn.Tanh(),
            torch.nn.Linear(hidden_dim, hidden_dim),
            torch.nn.Tanh(),
            torch.nn.Linear(hidden_dim, p),
        )
        # Trunk net: processes query location
        self.trunk = torch.nn.Sequential(
            torch.nn.Linear(trunk_input_dim, hidden_dim),
            torch.nn.Tanh(),
            torch.nn.Linear(hidden_dim, hidden_dim),
            torch.nn.Tanh(),
            torch.nn.Linear(hidden_dim, p),
        )
        self.bias = torch.nn.Parameter(torch.zeros(1))

    def forward(self, u_input, x_query):
        """
        u_input: (batch, n_sensors) — input function values at sensor locations
        x_query: (batch, n_query, dim) — query locations
        Returns: (batch, n_query) — predicted output function values
        """
        b = self.branch(u_input)  # (batch, p)
        t = self.trunk(x_query)   # (batch, n_query, p)
        # Dot product + bias
        out = torch.einsum('bp,bqp->bq', b, t) + self.bias
        return out

# Usage:
# n_sensors = 100 (fixed sensor locations for input function)
# model = DeepONet(branch_input_dim=100, trunk_input_dim=1)
# pred = model(u_sensors, x_query)  # u_sensors: (B, 100), x_query: (B, N, 1)

model.eval()
with torch.no_grad():
    pred = model(x_test)

# Plot 3 random test examples
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
for i, ax in enumerate(axes):
    idx = np.random.randint(n_test)
    ax.plot(x, x_test[idx, :, 0].numpy(), 'b-', label='Input IC')
    ax.plot(x, y_test[idx, :, 0].numpy(), 'k-', linewidth=2, label='True')
    ax.plot(x, pred[idx, :, 0].numpy(), 'r--', linewidth=2, label='FNO')
    ax.set_xlabel('x')
    ax.set_ylabel('u')
    ax.legend(fontsize=9)
    ax.grid(True, alpha=0.3)
    rel = torch.norm(pred[idx] - y_test[idx]) / torch.norm(y_test[idx])
    ax.set_title(f'Test {idx}: rel. error = {rel:.3f}')

plt.suptitle('FNO: Burgers Equation', fontsize=14)
plt.tight_layout()
plt.savefig('fno_predictions.png', dpi=150, bbox_inches='tight')