Deepxde学习笔记——DeepONet求解偏微分方程相关理论与代码实战

🚩本文内容

• 偏微分方程算子，万能近似定理，损失函数设计等理论内容
• 使用Deepxde对DeepONet代码编写
• 一般的PI-DeepONet流程

最近用Deepxde写了一个带参数的偏微分方程求解，感觉自己略有所得，想和着理论一起整理记录一点小小的心得。

理论内容

1. 算子

在偏微分方程里，我们给一个非线性的微分算子$\mathcal{N}$，其满足
$$\mathcal{N}(u,s)=0\text{\hspace{1em}(1)}$$
其中$u$是输入函数（注意是映射本身），$s$是我们想要的，但目前不知道的函数。那么DeepONet（The Deep Operator Network）学到的就是函数与函数之间的映射，如下：
$$G(u)=s\text{\hspace{1em}(2)}$$
在偏微分方程问题里，我们有
$$u=u(x_1,x_2,\dots ,x_n,t)\text{\hspace{1em}(3)}$$
那么(2)即
$$G(u)(x_1,x_2,\dots ,x_n,t)=s\text{\hspace{1em}(4)}$$
可能有失偏颇的概括说，Lu etal.(2019)等人提出的Deeponet在偏微分方程上所要实现的就是这点。

2. 对算子的万能近似定理（The Universal Approximation Theorem for Operator）

2.1 PINN理论回顾

介绍这个定理之前，首先回顾，在PINN（Physics Informed Neural Network）里面，我们利用深度神经网络所学习到的如下：
$$NN(X)=W_n{\sigma }_{n-1}(W_{n-1}{\sigma }_{n-2}(\dots W_2{\sigma }_1(W_{1}X+b_1)+b_2)+\dots )+b_{n-1})+b_n$$
PINN对初值条件（即IC，Initial Condition，之后用IC叙述），边界条件（即BC，Boundary Condition，之后用BC叙述）以及你在训练时所提供的数据点作为一个个约束（输入），去让PINN学到$W_i$和$b_i$，从而得到
$$NN(x,t)\approx u(x,t)$$

2.2 DeepONet的理论知识

我们可以看到，上式本质上是使用深度神经网络在每一层的叠加后，利用深度神经网络对函数进行近似学习。那么DeepONet其实也是一样的出发点：$\forall ϵ>0$, 有正数$n,p,m$,对以下的数或矩阵：$c_i^k,W_{ij}^k,b_i^k,b_k\in R,{\mathbf{W}}_{\mathbf{k}}\in R^n$,有如下万能近似定理：
$$\left|G(u)(y)-\sum _{k=1}^p\underset{Branch}{\underbrace{\sum _{i=1}^n{c}_i^k\sigma \left(\sum _{j=1}^m{W}_{ij}^{k}u(x_j)+b_i^k\right)}}\underset{Trunk}{\underbrace{\sigma ({\mathbf{W}}_{\mathbf{k}}\cdot \text{y}+b_k)}}\right|<ϵ\text{\hspace{1em}(THM1)}$$
这里需要注意的是，$\text{y}$是一个向量，和${\mathbf{W}}_{\mathbf{k}}$大小一致。上文PINN中$\sigma$是激活函数，这里的$\sigma$是一个连续的非多项式函数，$X$是一个Banach空间，其中的一个紧致集上的元素为$x_j$。Lu etal.(2019)等人的原文是根据紧致集上的拓扑映射给出的$G$，此处不表。
大致来说，（THM1）式子中，Branch代表的$u$就是类似PINN所学习到的$NN$，Trunk代表的$\text{y}$也是类似的，外面再套一层$G$，相当于用深度神经网络去做拟合连续做两次，对相应元素作点积求和得到的$G$。
所以这个式子可以化简：
$$G(u)(y)\approx \sum _{k=1}^q\underset{Branch}{\underbrace{b_k\left(u(x_1),u(x_2),...,u(x_m)\right)}}\cdot \underset{Trunk}{\underbrace{t_k(\text{y})}}\text{\hspace{1em}(5)}$$
故而有更广泛的对算子的万能近似定理（Generalized Universal Approximation Theoremfor Operator）这是基于Banach空间做内积提出的，这里只写出其形式，对于定义不多描述：
$$|G(u)(y)-\underset{\text{branch }}{\underbrace{⟨\mathbf{g}\left(u\left(x_1\right),u\left(x_2\right),\cdots ,u\left(x_m\right)\right)}},\underset{\text{trunk}}{\underbrace{\mathbf{f}(y)}}⟩\mid <ϵ\text{\hspace{1em}(THM2)}$$
其中：$\mathbf{g}:R^m\to R^p,\mathbf{f}:R^n\to R^p$

3. 损失函数设计

根据上面的理论知识，误差就在于这个$ϵ$了，在Deeponet中，$G(u)(y)$代表理想情况，实际上我们所训练到的是$G_{\theta }(u)(y)$
$$G_{\theta }(u)(y)=\sum _{k=1}^q\underset{Branch}{\underbrace{b_k\left(u(x_1),u(x_2),...,u(x_m)\right)}}\cdot \underset{Trunk}{\underbrace{t_k(\text{y})}}\text{\hspace{1em}(6)}$$
将其代入下面的算子损失函数：
$${\mathcal{L}}_{Operator}(\theta )=\frac{1}{NP}\sum _{i=1}^N\sum _{j=1}^P{\left|G_{\theta }(u^{(i)})y_j^{(i)}-G(u^{(i)})y_j^{(i)}\right|}^2$$
$${\mathcal{L}}_{Operator}(\theta )=\frac{1}{NP}\sum _{i=1}^N\sum _{j=1}^P{\left|\sum _{k=1}^q{b}_k\left(u(x_1),u(x_2),...,u(x_m)\right)\cdot t_k(y_j^{(i)})-G(u^{(i)})y_j^{(i)}\right|}^2$$
这里：

• $m$代表训练点的个数。
• $N$代表输入函数个数（函数空间中的）。
• $P$代表我们评估输出函数的那些点。

与PINN类似，我们所提供的IC,BC条件以及我们提供的数据点相当于一个个约束（输入）去逼近得到$G$。

那么从函数的构造我们得到了一个损失函数，因为偏微分方程是从物理问题提取出来的，所以在解决实际问题的时候，我们一般都会考虑到相应的物理定律。在PINN里，物理定律也是作为了损失函数的参考。
一般地，考虑一个广泛的非线性偏微分方程如下：
$$\{\begin{array}{l}\begin{array}{ll}\frac{\partial \mathbf{u}}{\partial t}+\mathcal{L}(\mathbf{u})=0,& (x,t)\in \Omega \times (0,T]\\ \mathbf{u}(x,t)=\mathbf{g}(x,t),& (x,t)\in \partial \Omega \times (0,T]\\ \mathbf{u}(x,0)=\mathbf{h}(x),& x\in \Omega \end{array}\end{array}\text{\hspace{1em}(GPINN)}$$
这里$\mathcal{L}$是非线性微分算子，$\Omega \in R^d$，$\mathbf{u}(x,t)$是方程(5)的精确解。
构建一个深度学习网络，对于$\tilde{\mathbf{u}}(x,t,\theta )$，其中$\theta \in R^k$为网络权重，$(x,t)$为网络输入，$\tilde{\mathbf{u}}=[\widetilde{u_1}(x,t),...,\widetilde{u_n}(x,t)]$为输出。
根据Karniadakis et al., 2019，标准损失如下：
$$\begin{array}{rl}G(\theta )& =\frac{1}{N_f}{\Vert \frac{\partial \tilde{\mathbf{u}}}{\partial t}(x,t,\theta )+\mathcal{L}(\tilde{\mathbf{u}}(x,t,\theta ))\Vert }_{\Omega \times (0,T],{\nu }_1}^2\\ & +\frac{1}{N_{IC}}\Vert \tilde{\mathbf{u}}(x,0,\theta )-\mathbf{h}(x){\Vert }_{\Omega ,{\nu }_2}^2\\ & +\frac{1}{N_{BC}}\Vert \tilde{\mathbf{u}}(x,t,\theta )-\mathbf{g}(x,t){\Vert }_{\partial \Omega \times (0,T],{\nu }_3}^2\end{array}$$
其中，$N_f、N_{IC}$和$N_{B}C$分别对应于根据内部条件、初始条件和边界条件各自的概率密度${\nu }_1$、${\nu }_2$和${\nu }_3$采样的点数。内部条件其实就是在给定定义域范围内部的数据点条件。
我们注意到这个损失函数的后两部分是从残差出发的。和PINN一样，Deeponet的operator下的损失函数也是从残差去设计的，至于其物理定律的损失函数，则是从(1)出发，实际情况下，operator的$G(u^{(i)})y_j^{(i)}$是从边界和初值条件取点，而物理定律的损失函数会用残差，这一点可以见编码流程设计。
$${\mathcal{L}}_{Physics}(\theta )=\frac{1}{NQm}\sum _{i=1}^N\sum _{j=1}^Q\sum _{k=1}^m{\left|\mathcal{N}(u^{(i)}(x_k),G_{\theta }(u^{(i)})(y_j^{(i)}))\right|}^2$$

$\mathcal{N}$是一个非线性微分算子， { y j } i = 1 Q \{y_j\}_{i=1}^{Q} {yj​}i=1Q​就是那些用于满足物理定律约束的数据点。
所以Deeponet在结合物理定律后的总误差如下：
$$\mathcal{L}(\theta )={\mathcal{L}}_{Operator}(\theta )+{\mathcal{L}}_{Physics}(\theta )$$

此时DeepONet在结合物理定律后，便是PI-DeepONet。

程序编码

1.1 编码前的准备与了解

那么到目前为止，我们弄明白了DeepONet的拟合本质，了解了其网络，输出和损失函数的设计。那么DeepONet的输入就和PINN是一样的只是$(x,t)$吗？答案是否定的，读者也许能发现上述损失函数的设计中出现了输入函数。DeepONet有两个网络，Branch网络和Trunk网络：
Branch网络：使用深度神经网络，编码离散的输入函数空间。
Trunk网络：使用深度神经网络，编码输出函数的定义域。
DeepONet真正的强大之处在于：其可以学习各种显式算子，如积分和分数拉普拉斯算子，以及表示确定性和随机微分方程的隐式算子。

这里贴一下deepxde库的最新指南pde_operator：deepxde库最新指南pde_operator里，下面的函数中：

class deepxde.data.pde_operator.PDEOperator(pde, function_space, evaluation_points, num_function, function_variables=None, num_test=None)

• pde：是dde.data.PDE或者dde.data.TimePDE的一种，待会阐述。
• function_space：便是Branch网络所定义的输入函数空间，函数空间有PowerSeries，GRF，chebyshev等几种。
• evaluation_points：是numpy array的形式，数组大小为(n_points, dimension)。它作为一组测试评估点，在输入函数空间中随机取样获得的输入函数中进行评估（相当于Branch网络的输入锚定点，让输入函数在它们的基础上计算值）。
• num_function：function_space中抽样的函数个数。
• function_variable：function_space中抽样的函数与哪几个自变量关联。这里注意，以下对于关联的解读仅为个人猜测：0是x，1是t，2开头是你所定义的参数或者辅助变量，具体0，1，2等谁指定谁需要看你在dde.geometry的定义。
• num_test：测试样本的数量，默认为None，表示测试数据与训练数据相同（即输入函数会被拿来测试）。

该函数定义了Branch网络的行为，那么Trunk网络是从如下函数定义的。
在deepxde.data.pde module里，下面的函数中：

class deepxde.data.pde.PDE(geometry, pde, bcs, num_domain=0, num_boundary=0, train_distribution='Hammersley', anchors=None, exclusions=None, solution=None, num_test=None, auxiliary_var_function=None)

我们一个个解读：

• geometry：dde.geometry的一种，待会阐述。
• pde：你所定义的pde函数，如何定义，待会阐述。
• bcs：BC条件，如何写BC条件待会阐述。
• num_domain：定义域划分部分个数（在定义域里抽样点个数）。
• num_boundary：定义域边界条件抽样点个数。
• train_distribution：抽样点的分布，有“uniform","Hammersley"等。
• anchors：我们所提供的训练点的Numpy array，以及取样点的 num_domain 和 num_boundary。
• exclusions：我们提供的Numpy array用于排除这些点的训练。
• solution：精确解，显式解参考。
• num_test：测试点个数，用于测试损失，如果"None"，则使用训练点进行测试。
• auxiliary_var_function：辅助变量的函数输入（需要提前定义好辅助变量的函数）。

那么TimePDE也类似的，只不过在bcs变成了ic_bcs用于IC,BC条件的输入，多了一个num_initial，在定义域初值条件的抽样点个数。

class deepxde.data.pde.TimePDE(geometryxtime, pde, ic_bcs, num_domain=0, num_boundary=0, num_initial=0, train_distribution='Hammersley', anchors=None, exclusions=None, solution=None, num_test=None, auxiliary_var_function=None)

接下来阐述deepxde.geometry：
1d是Interval，比如我想定义x的定义域为$[-1,6]$：

x = dde.geometry.Interval(-1,6)

2d是Rectangle，比如我想定义$x,y$的定义域分别为$[-1,6],[0,5]$，把它们用geom统一表示。

geom = dde.geometry.Rectangle([-1, 0], [6, 5])

3d是Cuboid，nd是Hypercube，时间是TimeDomain。

geom = dde.geometry.Cuboid([-1,0,0],[6,5,5])
t = dde.geometry.TimeDomain(0, 10)

针对TimePDE，我们需要把非t变量和t变量统一，例子如下：

geomtime = dde.geometry.GeometryXTime(geom,t)

如何定义pde？
以github官方：diff_rec_aligned_pideeponet.py为例子：

def pde(x, y, v):
    D = 0.01
    k = 0.01
    dy_t = dde.grad.jacobian(y, x, j=1)
    dy_xx = dde.grad.hessian(y, x, j=0)
    return dy_t - D * dy_xx + k * y**2 - v

这里，j=0是对x，j=1是对y，v是参数。要注意的是，pde定义里，x包含了所有的自变量，诸如x，t等，y是唯一因变量。
以advection_aligned_pideeponet.py为例：

def pde(x, y, v):
    dy_x = dde.grad.jacobian(y, x, j=0)
    dy_t = dde.grad.jacobian(y, x, j=1)
    return dy_t + dy_x

如果v作为参数不在pde式子里，就要在ic_bcs条件里。

那么如何写IC,BC条件？
deepxde的IC,BC条件是从匿名函数定义的，匿名函数即lambda x，其针对单个函数的传值，能快捷的计算出一个函数对于这个输入值的值，比如lambda x,y:x+y就定义了一个简单的加法。
这里，lambda后面可以空转，也就是lambda _，也可以写lambda x。icbc中无论BC还是IC都会给你自动找到x，也可以传入函数，故而以下几个写法均可：

bc = dde.icbc.DirichletBC(geom, lambda x: 0, lambda _, on_boundary: on_boundary)
bc = dde.icbc.DirichletBC(geom, lambda x: 0, lambda x, on_boundary: on_boundary)
bc = dde.icbc.DirichletBC(geom, boundary_condition, lambda _, on_boundary: on_boundary)
ic = dde.icbc.IC(geom, initial_condition, lambda _, on_initial: on_initial)
ic = dde.icbc.IC(geom, initial_condition, lambda x, on_initial: on_initial)
ic = dde.icbc.IC(geomtime, lambda _: 0, lambda _, on_initial: on_initial)

要注意这样定义的IC,BC全是软约束，神经网络训练时，不会强迫pde相关数据点去正好得到这些数，而是尽可能地靠近IC,BC条件所给的值。如果要使用硬约束，请检查github issue。
在这里贴一个个人写的分段函数初值条件：

def initial_condition(x, c):
    x = x[:, 0:1]
    return np.piecewise(
        x,
        [x < 0, (0 <= x) & (x < 1), x >= 1],
        [0, 1, 0]
    )

第一行的x = x[:, 0:1]表明在x（包括x，t等）的第一个变量（x）中，有如上初值条件。

最后的问题就是网络定义问题了，在deepxde.nn.pytorch.deeponet module里有如下内容，以pytorch为例，前者是普通数据，后者是用笛卡尔坐标定义的数据（卡氏积）。

class deepxde.nn.pytorch.deeponet.DeepONet(layer_sizes_branch, layer_sizes_trunk, activation, kernel_initializer, num_outputs=1, multi_output_strategy=None
class deepxde.nn.pytorch.deeponet.DeepONetCartesianProd(layer_sizes_branch, layer_sizes_trunk, activation, kernel_initializer, num_outputs=1, multi_output_strategy=None)

我们还是一一解读：

• layer_sizes_branch：Branch网络，用list构建，或者（dimension, f），其中dimension是输入维度，f 是网络函数。在multi_output_strategy里，除"split_branch"和 "split_trunk"外，所有策略的Branch和Trunk网络最后一层的宽度都应相同。
• layer_sizes_trunk：Trunk网络，用list构建。
• activation：激活函数。
• kernel_initializer：参数初始化方法，比如Xavier Glorot初始化方法：“Glorot normal”。
• num_outputs：神经网络输出个数。
• multi_output_strategy：多输出策略，默认为None，具体可参考deepxde:multi_output_strategy

1.2 个人建议与笔记

1. 使用DeepONet注意，如果是anchor自己提供值，在没有足够bc，ic值的情况下或者bc，ic有一定冲突的情况下，建议在PDE和TimePDE里给num_boundary和num_initial一定初始值。
2. dde.data.PDEOperator或者dde.data.PDEOperatorCartesianProd中，evaluation_points个数与Branch网络第一层输入神经元个数相同。建议通过linspace均匀取样测试。
3. pde定义中，IC，BC条件中有component的说法。当在pde定义的自变量有多个同类型坐标（除了x，t这种一个空间坐标一个时间坐标的），需要使用component针对一个变量，对其求导或者对其定义bc，ic条件。
   以stokes_aligned_pideeponet.py为例子：

# PDE equation
def pde(xy, uvp, aux):
    mu = 0.01
    # first order
    du_x = dde.grad.jacobian(uvp, xy, i=0, j=0)
    dv_y = dde.grad.jacobian(uvp, xy, i=1, j=1)
    dp_x = dde.grad.jacobian(uvp, xy, i=2, j=0)
    dp_y = dde.grad.jacobian(uvp, xy, i=2, j=1)
    # second order
    du_xx = dde.grad.hessian(uvp, xy, component=0, i=0, j=0)
    du_yy = dde.grad.hessian(uvp, xy, component=0, i=1, j=1)
    dv_xx = dde.grad.hessian(uvp, xy, component=1, i=0, j=0)
    dv_yy = dde.grad.hessian(uvp, xy, component=1, i=1, j=1)
    motion_x = mu * (du_xx + du_yy) - dp_x
    motion_y = mu * (dv_xx + dv_yy) - dp_y
    mass = du_x + dv_y
    return motion_x, motion_y, mass
   
# Boundary condition
# other boundary conditions will be enforced by output transform
def bc_slip_top_func(x, aux_var):
    # using (perturbation / 10 + 1) * x * (1 - x)
    return (aux_var / 10 + 1.) * dde.backend.as_tensor(x[:, 0:1] * (1 - x[:, 0:1]))
bc_slip_top = dde.icbc.DirichletBC(
    geom=geom,
    func=bc_slip_top_func,
    on_boundary=lambda x, on_boundary: np.isclose(x[1], 1.),
    component=0)

1.3 编码流程参考

DeepONet中，对于训练的神经网络，默认采用的前馈连接网络FNN。当然，可以自己定义网络，比如多层感知机，这里不表。
PI-DeepONet编码大致流程如下，假设给定$y=(x,t)$，u仅仅只是x的一个函数。

1. 选择N个不同的函数u： { u ( i ) ( x ) } i = 1 N = [ u ( 1 ) ( x ) , u ( 2 ) ( x ) , . . . , u ( N ) ( x ) ] \{u^{(i)}(\textbf{x})\}_{i=1}^{N}={[u^{(1)}(\textbf{x}),u^{(2)}(\textbf{x}),...,u^{(N)}(\textbf{x})]} {u(i)(x)}i=1N​=[u(1)(x),u(2)(x),...,u(N)(x)].

2. 在$x_1,\dots ,x_m$点评估$N$个u函数：$\left[\begin{array}{cccc}{u}^{(1)}(x_1)& u^{(1)}(x_2)& ...& u^{(1)}(x_m)\\ u^{(2)}(x_1)& u^{(2)}(x_2)& ...& u^{(2)}(x_m)\\ ⋮& ⋮& \ddots & ⋮\\ u^{(N)}(x_1)& u^{(N)}(x_2)& ...& u^{(N)}(x_m)\end{array}\right]$.

3. 将结果输入branch网络：$b_k\left(\begin{array}{cccc}{u}^{(1)}(x_1)& u^{(1)}(x_2)& ...& u^{(1)}(x_m)\\ u^{(2)}(x_1)& u^{(2)}(x_2)& ...& u^{(2)}(x_m)\\ ⋮& ⋮& \ddots & ⋮\\ u^{(N)}(x_1)& u^{(N)}(x_2)& ...& u^{(N)}(x_m)\end{array}\right)$.

4. 从边界和初始条件中选择$P$个点准备评估：${\text{y}}_u=y_{u1},y_{u2},...,y_{uP}$.

5. 将评估点输入trunk网络：$t_k(y_{u1},y_{u2},...,y_{uP})$.

6. 通过计算branch网络的输出与trunk网络的输出之间的点积来近似算子，即(THM2)：$G_{\theta }(u)(\text{y})=\sum _{k=1}^q\underset{Branch}{\underbrace{b_k\left(u(x_1),u(x_2),...,u(x_m)\right)}}\cdot \underset{Trunk}{\underbrace{t_k(\text{y})}}$.

7. 根据(5)，(6)，从理论上知道：$G_{\theta }(u)({\text{y}}_u)\approx G(u)({\text{y}}_u)$，$G(u^{(i)})y_j^{(i)}$中的$y_j^{(i)}$从初值和边界条件取点，计算损失：${\mathcal{L}}_{Operator}(\theta )=\frac{1}{NP}\sum _{i=1}^N\sum _{j=1}^P{\left|G_{\theta }(u^{(i)})y_j^{(i)}-G(u^{(i)})y_j^{(i)}\right|}^2$.

8. 在定义域内选择Q个点：${\text{y}}_r=y_{r1},y_{r2},...,y_{rQ}$.

9. 根据相关物理定律，计算损失：${\mathcal{L}}_{Physics}(\theta )=\frac{1}{NQ}\sum _{i=1}^N\sum _{j=1}^Q{\left|R_{\theta }^{(i)}(y_{r,j}^{(i)})-u^{(i)}(x_{r,j}^{(i)})\right|}^2$.

10. 计算整体损失：$\mathcal{L}(\theta )={\mathcal{L}}_{operator}(\theta )+{\mathcal{L}}_{Physics}(\theta )$.

11. 更新branch网络和trunk网络参数来最小化：$\mathcal{L}(\theta )$.

12. 选取$ϵ$作为终止条件，重复这个过程直到：$G_{\theta }(u)(x,t)\approx G(u)(x,t)$.

文末

DeepONet与PINN优缺点

特性	DeepONet	PINN
核心概念	学习输入输出之间的算子关系	将物理定律嵌入神经网络的损失函数中
数据需求	通常需要大量的数据	可以在物理约束下使用较少的数据
泛化能力	对不同输入函数有良好的泛化能力	通常针对特定问题，泛化能力有限
预测速度	训练后预测速度快	每次求解需要较多计算，预测速度相对慢
适用场景	需要快速多次预测、数据驱动的复杂系统	需要满足物理定律的精确求解、逆问题

附录

如果你还有关于deepxde的问题，欢迎去deepxde库查看！
也可以去github官方查看例子学习！