A simple finite-difference solver

An intro to our loss module’s finite difference utility demonstrating its use to create a simple numerical solver for the diffusion-advection equation.

Import the library

We first import our neuralop library and required dependencies.

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

from neuralop.losses.differentiation import FiniteDiff

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

Defining our problem

We aim to solve the 2D diffusion advection equation:

\(u_t + cx \cdot u_x + cy \cdot u_y = \nu (u_xx + u_yy) + f(x,y,t)\),

Where \(f(x,y,t)\) is a source term and \(cx\) and \(cy\) are advection speeds in x and y. We set simulation parameters below:

## Simulation parameters
Lx, Ly = 2.0, 2.0   # Domain lengths
nx, ny = 64, 64   # Grid resolution
T = 1.6    # Total simulation time
dt = 0.001  # Time step
nu = 0.02   # diffusion coefficient
cx, cy = 1.0, 0.6  # advection speeds

## Create grid
X = torch.linspace(0, Lx, nx, device=device).repeat(ny, 1).T
Y = torch.linspace(0, Ly, ny, device=device).repeat(nx, 1)
dx = Lx / (nx - 1)
dy = Ly / (ny - 1)
nt = int(T / dt)

## Initialize finite difference operator
fd = FiniteDiff(dim=2, h=(dx, dy))


## Initial condition and source term
u = (-torch.sin(2 * np.pi * Y) * torch.cos(2 * np.pi * X)
        + 0.3 * torch.exp(-((X - 0.75)**2 + (Y - 0.5)**2) / 0.02)
        - 0.3 * torch.exp(-((X - 1.25)**2 + (Y - 1.5)**2) / 0.02)).to(device)

def source_term(X, Y, t):
    return 0.2 * torch.sin(3 * np.pi * X) * torch.cos(3 * np.pi * Y) * torch.cos(4 * np.pi * t)

Simulate evolution using numerical solver

u_evolution = [u.clone()]

t = torch.tensor(0.0)
for _ in range(nt):

    # Compute derivatives
    u_x = fd.dx(u)
    u_y = fd.dy(u)
    u_xx = fd.dx(u_x)
    u_yy = fd.dy(u_y)

    # Evolve one step in time using Euler's method
    u = u + dt * (-cx * u_x - cy * u_y + nu * (u_xx + u_yy) + source_term(X, Y, t))
    t += dt
    u_evolution.append(u.clone())

u_evolution = torch.stack(u_evolution).cpu().numpy()

Animate our solution

num_frames = 100
frame_indices = torch.linspace(0, len(u_evolution) - 1, num_frames, dtype=torch.int).cpu().numpy()
u_frames = u_evolution[frame_indices]

fig, ax = plt.subplots(figsize=(6, 6))
cmap_u = ax.imshow(u_frames[0], extent=[0, Lx, 0, Ly], origin="lower", cmap="plasma")
ax.set_title("Advection-Diffusion: u")
plt.colorbar(cmap_u, ax=ax, shrink=0.75)

def update(frame):
    cmap_u.set_data(u_frames[frame])
    ax.set_title(f"Time: {frame_indices[frame] * dt:.3f}")
    ax.set_xticks([])
    ax.set_yticks([])
    return cmap_u,

ani = animation.FuncAnimation(fig, update, frames=len(u_frames), interval=50, blit=False)

Total running time of the script: (0 minutes 9.415 seconds)

Gallery generated by Sphinx-Gallery