Fourier Neural Operator (FNO)

One of the most popular neural operator architectures is the Fourier Neural Operator, or FNO. This notebook demonstrates the use of FourierLayer and FourierNeuralOperator in continuiti.

The FNO architecture was proposed in Z. Li et al., 2020 and gained a lot of attention because it is computationally very efficient. However, due to the Fast Fourier Transform (FFT) being used, it is in general restricted to simple geometries like rectangles.

import torch
from continuiti.data import OperatorDataset
from continuiti.discrete import RegularGridSampler
from continuiti.operators import FourierNeuralOperator
from continuiti.operators.fourierlayer import FourierLayer
from continuiti.trainer import Trainer

1D

Let's start with a simple one-dimensional function training a single FourierLayer on it.

# Input function
u_func = lambda x: torch.exp(-x**2)*10

# Target function
v_func = lambda y: torch.exp(-y**2)*2*10 * (-y)

def get_equidistant(N):
    return torch.arange(-N/2, N/2) / N * torch.pi * 2

num_sensors = 31
num_evaluations = 31

x = get_equidistant(num_sensors)
y = get_equidistant(num_evaluations)

u = u_func(x)
v = v_func(y)

Dataset

n_observations = 1
u_dim = x_dim = y_dim = v_dim = 1

x = x.reshape(n_observations, x_dim, num_sensors)
u = u.reshape(n_observations, u_dim, num_sensors)
y = y.reshape(n_observations, y_dim, num_evaluations)
v = v.reshape(n_observations, v_dim, num_evaluations)

dataset = OperatorDataset(x, u, y, v)

Fourier Layer

We instantiate the Fourier layer with the shapes determined by the dataset.

fourier = FourierLayer(dataset.shapes)

Training

trainer = Trainer(fourier)
trainer.fit(dataset, epochs=10_000)

We can evaluate the trained FourierLayer with any other resolution, e.g., to plot the output function an a finer resolution.

num_plot = 1000
y_plot = get_equidistant(num_plot).reshape(1, 1, -1)
v_pred = fourier(x, u, y_plot)

Modes

In the same way as the FourierLayer allows you to change the number of evaluation points of \(y\) for each forward pass, the dimensionality of the input \(u\) can also be changed.

In the case of \(u\) being larger than what was specified during initialization, the FourierLayer removes the high frequencies (or modes). In case \(u\) is smaller, zero-values large frequencies are added.

x_fine = get_equidistant(128).reshape(1, 1, -1)
u_fine = u_func(x_fine)
v_pred = fourier(x_fine, u_fine, y_plot)

Per default, the Fourier layer considers as many modes as are suggested by the shapes.u.num parameter. However, it can be set to any other value and determines the number of weights in the layer.

vs = {}
for modes in [3, 4, 10, 100]:
    fourier = FourierLayer(dataset.shapes, num_modes=(modes,))
    Trainer(fourier).fit(dataset, epochs=10_000)

    v_pred = fourier(x, u, y_plot)
    vs[modes] = v_pred

2D

The FourierLayer also works on multi-dimensional data and we demonstrate these capabilities along with the multi-layer FourierNeuralOperator (FNO) in the following.

sampler = RegularGridSampler([-3., -3.], [3., 3.])
N = 100
x = y = sampler(N**2)

u = multivariate_normal_2d(x)
v = double_multivariate_normal_2d(y)

Note

The definitions of multivariate_normal_2d and double_multivariate_normal_2d are hidden for readability, but both functions are visualized below.

Dataset

x = x.reshape(1, 2, N, N)
u = u.reshape(1, 1, N, N)
y = y.reshape(1, 2, N, N)
v = v.reshape(1, 1, N, N)
dataset = OperatorDataset(x, u, y, v)

Fourier Neural Operator

Let's use an FNO! It contains depth many FourierLayers with a latent dimension of width.

model = FourierNeuralOperator(shapes=dataset.shapes, depth=3, width=3)

Training

Trainer(model).fit(dataset, epochs=10_000)
v_pred = model(x, u, y)

---

Last update: 2024-08-20
Created: 2024-08-20