使用傅里叶神经算子求解 PDE

R2026b

此示例说明如何训练傅里叶神经算子 (FNO) 神经网络来输出偏微分方程 (PDE) 的解。

神经算子 [1] 是一种能够在函数空间之间进行映射的神经网络。例如，当给定系统的初始条件时，它可以学习输出 PDE 的解。傅里叶神经算子 (FNO) [2] 是一种能够利用傅里叶变换更高效地学习这些映射的神经算子。具体而言，对于空间域中的输入函数 $f$ ，神经网络中的各层学习变换 $F$ ，使得 $F (f) = f_{}^{ˆ}$ ，其中 $f_{}^{ˆ}$ 是 $f$ 的傅里叶变换。然后 FNO 可以在频域中进行运算。使用神经网络来预测 PDE 的解，可能比数值求解更快。

下图说明了数据在 FNO 中的流动过程。

Diagram of data flowing through FNO. The input is a function labeled "Initial Condition". The output is a function labeled "PDE Solution".

由于 FNO 学习的是函数空间之间的映射，您可以在比训练数据分辨率更高的条件下计算输出函数。该功能被称为零样本超分辨率 (ZSSR)。也就是说，ZSSR 能够在无需额外高分辨率样本的情况下，将输入和输出数据的分辨率提升到原始训练数据之上。下图说明了用于预测的低分辨率训练数据与高分辨率输入输出之间的差异。

Plots showing low resolution and high resolution inputs and outputs. The low-resolution functions are composed of a small number of straight lines between points. The high-resolution functions are smooth curves.

此示例训练一个 FNO 来输出一维伯格斯方程的解。

加载训练数据

加载一维伯格斯方程的数据。

伯格斯方程是一中 PDE，它在应用数学的多个领域（如流体力学、非线性声学、气体动力学和交通流）中都有出现。该 PDE 为：

$\frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} = \frac{1}{R e} \frac{\partial^{2} u}{\partial x^{2}}, x \in (0, 1), t \in (0, 1),$

其中 $R e$ 是雷诺数。初始条件为：

$u (x, 0) = u_{0} (x), x \in (0, 1)$

该方程在整个空间域上施加周期性边界条件。目标是学习从初始条件 $u_{0}$ 到时间 $t = 1$ 时的解的映射： $u_{0} ⟼ u (x, 1)$ 。

从 burgers_data_R10.mat 加载一维伯格斯方程的数据，其中包含初始速度 $u_{0}$ 和伯格斯方程的解 $u (x, 1)$ 。

dataFile = "burgers_data_R10.mat";
if ~exist(dataFile,"file")
    location = "https://ssd.mathworks.com/supportfiles/nnet/data/burgers1d/burgers_data_R10.mat";
    websave(dataFile,location); 
end

data = load(dataFile,"a","u");
u0 = data.a;
u1 = data.u;

查看其中一些训练数据。

numPlots = 3;
gridSize = size(u0,2);
figure
tiledlayout(numPlots,2)
for i = 1:numPlots
    nexttile
    plot(linspace(0,1,gridSize),u0(i,:));
    title("Observation " + i + newline + "Initial Condition")
    xlabel("x")
    ylabel("u")

    nexttile
    plot(linspace(0,1,gridSize),u1(i,:));
    title("PDE Solution")
    xlabel("x")
    ylabel("v")
end

Three pairs of plots showing initial conditions (u vs. x) and corresponding PDE solutions (v vs. x) for observations 1, 2, and 3

准备训练数据

使用 trainingPartitions 函数来拆分训练数据和测试数据，此函数作为支持文件包含在此示例中。要访问此函数，请以实时脚本形式打开此示例。将 80% 的数据用于训练，10% 用于验证，其余 10% 用于测试。

numObservations = size(u0,1);
[idxTrain,idxValidation,idxTest] = trainingPartitions(numObservations,[0.8 0.1 0.1]);
U0Train = u0(idxTrain,:);
U1Train = u1(idxTrain,:);

U0Validation = u0(idxValidation,:);
U1Validation = u1(idxValidation,:);
U0Test = u0(idxTest,:);

为减少 FNO 中的处理量，对训练数据和验证数据进行下采样，使其网格大小为 1024。

gridSizeDownsampled = 1024;
downsampleFactor = floor(gridSize./gridSizeDownsampled);

U0Train = U0Train(:,1:downsampleFactor:end);
U1Train = U1Train(:,1:downsampleFactor:end);
U0Validation = U0Validation(:,1:downsampleFactor:end);
U1Validation = U1Validation(:,1:downsampleFactor:end);

FNO 要求输入数据有一个对应于空间坐标（网格）的通道。创建该网格并将其与输入数据串联。

xMin = 0;
xMax = 1;
xGrid = linspace(xMin,xMax,gridSizeDownsampled);

xGridTrain = repmat(xGrid, [numel(idxTrain) 1]);
xGridValidation = repmat(xGrid, [numel(idxValidation) 1]);

U0Train = cat(3,U0Train,xGridTrain);
U0Validation = cat(3,U0Validation,xGridValidation);

查看训练集和验证集的大小。第一个维度、第二个维度和第三个维度分别是批量维度、空间维度和通道维度。

size(U0Train)

ans = 1×3

    1638    1024    2

size(U0Validation)

ans = 1×3

    204    1024    2

定义神经网络架构

使用 fnoNetwork 函数（自 R2026b 起）创建 FNO 网络。该网络由一个提升层、串联连接的多个傅里叶层和一个投影层组成。

Diagram of the neural network of layers connected in series. From left to right, there is the input, convolution, four fourier-ReLU blocks, convolution, ReLU, convolution, and output.

指定以下选项：

使用与训练预测变量（初始条件）和目标变量（PDE 解）一致的输入大小和输出大小。
对于一维 FNO 网络，请指定数据具有一个空间维度。
指定 4 个傅里叶层、16 个傅里叶模式、64 的隐藏大小以及 ReLU 激活。

inputSize = size(U0Train,3);
outputSize = size(U1Train,3);

net = fnoNetwork(inputSize,outputSize, ...
    Dimension=1, ...
    NumFourierLayers=4, ...
    NumModes=16, ...
    HiddenSize=64, ...
    Activation="relu");

在 R2026b 之前：使用 dlnetwork 对象手动定义神经网络架构。

训练选项

指定训练选项。在选项中进行选择需要经验分析。要通过运行试验探索不同训练选项配置，您可以使用Experiment Manager。

使用 SGDM 求解器进行训练。
使用分段学习率调度进行训练，初始学习率为 0.001，下降因子为 0.5。
使用 20 的小批量大小进行 40 轮训练。
每轮训练都会打乱数据。
由于输入数据的维度对应于批量维度、空间维度和通道维度，因此将输入数据格式指定为 "BSC"。
在图中监控训练进度。
指定验证数据，并每 100 次迭代验证一次网络。
禁用详尽输出。

schedule = piecewiseLearnRate(DropFactor=0.5);

options = trainingOptions("sgdm", ...
    InitialLearnRate=1e-3, ...
    LearnRateSchedule=schedule, ...
    MaxEpochs=40, ...
    MiniBatchSize=20, ...
    Shuffle="every-epoch", ...
    InputDataFormats="BSC", ...
    Plots="training-progress", ...
    ValidationData={U0Validation,U1Validation}, ...
    ValidationFrequency=100, ...
    Verbose=false);

定义相对 $L_{2}$ 损失函数

创建函数 relativeL2Loss，该函数以预测值和目标作为输入，并返回相对 $L_{2}$ 损失。相对 $L_{2}$ 损失函数是 $L_{2}$ 损失函数的一种变体，它通过使用预测值与目标之间差值的欧几里德范数，再除以目标的欧几里德范数来考虑尺度因素：

$R e l a t i v e L 2 (Y, T) = \sum_{n = 1}^{N} \frac{‖ T_{n} - Y_{n} ‖_{2}}{‖ T_{n} ‖_{2}}$

其中：

$Y$ 表示预测值。
$T$ 表示目标。
$N$ 表示观测值数量。

要从 dlarray 输入中删除格式，请使用 stripdims 函数。将维度顺序指定为 "SCB"，使空间维度排在第一位，以便与 vecnorm 一起使用。

function loss = relativeL2Loss(Y,T)

Y = stripdims(Y,"SCB");
T = stripdims(T,"SCB");

p = vecnorm(T-Y,2,1);
q = vecnorm(T,2,1);

loss = sum(p./q);

end

训练网络

使用 trainnet 函数训练神经网络。使用自定义损失函数 relativeL2Loss。默认情况下，trainnet 函数使用 GPU（如果有）。使用 GPU 需要 Parallel Computing Toolbox™ 许可证和受支持的 GPU 设备。有关受支持设备的信息，请参阅GPU 计算要求 (Parallel Computing Toolbox)。否则，该函数使用 CPU。要手动选择执行环境，请使用 ExecutionEnvironment 训练选项。

net = trainnet(U0Train,U1Train,net,@relativeL2Loss,options);

Training progress plot showing training and validation loss decreasing over 3240 iterations across 40 epochs

测试网络

通过将模型对保留测试集的预测值与真实标签进行比较，并计算均方根误差 (RMSE)，来测试模型的预测性能。

对测试数据进行下采样，并使用与训练数据相同的步骤串联网格。

xGridTest = repmat(xGrid, [numel(idxTest) 1]);

U0Test = U0Test(:,1:downsampleFactor:end);
U0Test = cat(3,U0Test,xGridTest);

U1Test = u1(idxTest,:);
U1Test = U1Test(:,1:downsampleFactor:end);

使用 testnet 函数测试神经网络。默认情况下，testnet 函数使用 GPU（如果有）。否则，该函数使用 CPU。要手动选择执行环境，请使用 testnet 函数的 ExecutionEnvironment 参量。

rmseTest = testnet(net,U0Test,U1Test,"rmse",InputDataFormats="BSC")

rmseTest = 
0.0425

为评估 ZSSR 的准确度，使用未下采样的测试数据来测试网络。

xGridZSSR = linspace(xMin,xMax,gridSize);
xGridTestZSSR = repmat(xGridZSSR, [numel(idxTest) 1]);

U0TestZSSR = u0(idxTest,:);
U0TestZSSR = cat(3,U0TestZSSR,xGridTestZSSR);
U1TestZSSR = u1(idxTest,:);

rmseTestZSSR = testnet(net,U0TestZSSR,U1TestZSSR,"rmse",InputDataFormats="BSC")

rmseTestZSSR = 
0.0425

使用新数据进行预测

加载新数据。出于示例目的，使用前三个测试观测值。

dataNew = u0(idxTest(1:3),:);
numObservationsNew = size(dataNew,1);

计算网格大小并将其与数据串联。

gridSizeNew = size(dataNew,2);
xGridNew = linspace(xMin,xMax,gridSizeNew);
xGridNew = repmat(xGridNew, [size(dataNew,1) 1]);
U0New = cat(3,dataNew,xGridNew);

使用 minibatchpredict 函数进行预测。默认情况下，minibatchpredict 函数使用 GPU（如果有）。否则，该函数使用 CPU。要手动选择执行环境，请使用 ExecutionEnvironment 选项。

Y = minibatchpredict(net,U0New,InputDataFormats="BSC");

在图中可视化预测。

figure
tiledlayout(numObservationsNew,2)
for i = 1:numObservationsNew
    nexttile
    plot(xGridNew(i,:),dataNew(i,:))
    title("Initial Condition")

    nexttile
    plot(xGridNew(i,:),Y(i,:))
    title("PDE Solution")
end

Three pairs of plots showing initial conditions and predicted PDE solutions for three test observations, with high-resolution output from the trained FNO

参考资料

Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.“Neural Operator:Graph Kernel Network for Partial Differential Equations.” arXiv, March 6, 2020. http://arxiv.org/abs/2003.03485.
Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.“Fourier Neural Operator for Parametric Partial Differential Equations,” 2020. https://openreview.net/forum?id=c8P9NQVtmnO.

另请参阅