PDE反问题描述
地球电势$u$由Poisson方程控制
$$ \nabla \cdot (\beta \nabla u) = q $$
地壳中的电导率$\beta$常常是不知道的，所以实测几个点的电势$u$，然后通过求解PDE优化问题反推出电导率$\beta$。
$$ \begin{align} & \text{min}_{\beta} J = \sum_{i}^{M} (u_{i} - u^{obs}_{i})^2 \\ & \text{subject to } \nabla \cdot (\beta \nabla u) = q \\ \end{align} $$
正向求解
假定一个电导率分布$\beta$，求解方程得到电势$u$，对比$u$与观测值$u^{obs}_{i}$。 此时存在误差，因此需要修改$\beta$，重复上述正向求解过程，最终得到准确的$\beta$使得$u$与观测值$u^{obs}_{i}$相同。
伴随优化基本理论
为了便于叙述伴随优化，将PDE细节略去，上述PDE简化为残差形式
$$ \vec{R}(\vec{u}, \vec{\beta}) = 0 $$
因为$\vec{R} = 0$，所以对设计变量$\beta$求导如下
$$ \begin{align} & \frac{d\vec{R}}{d\vec{\beta}} = 0 = \frac{\partial \vec{R}}{\partial \vec{\beta}} + \frac{\partial \vec{R}}{\partial \vec{u}} \frac{\partial \vec{u}}{\partial \vec{\beta}} \\ \rightarrow & \frac{\partial \vec{R}}{\partial \vec{u}} \frac{\partial \vec{u}}{\partial \vec{\beta}} = -\frac{\partial \vec{R}}{\partial \vec{\beta}} \\ \end{align} $$
目标函数$J$对某个设计变量$\beta_{i}$的敏感度定义如下
$$ \frac{d J}{d \beta_{i}} = \frac{\partial J}{\partial \beta_{i}} + \frac{\partial J}{\partial \vec{u}} \frac{\partial \vec{u}}{\partial \beta_{i}} $$
$\frac{\partial J}{\partial \vec{u}}$很好求，直接从$J$的定义中求导得到。 我们现在不想求$\frac{\partial \vec{u}}{\partial \beta_{i}}$。 所以我们引入新的变量$\vec{\psi} = \frac{\partial J}{\partial \vec{R}}$，满足如下关系
$$ \vec{\psi} \frac{\partial \vec{R}}{\partial \vec{u}} = \frac{\partial J}{\partial \vec{u}} $$
写成$Ax=b$的形式如下
$$ \left(\frac{\partial \vec{R}}{\partial \vec{u}} \right)^T \vec{\psi} = \frac{\partial J}{\partial \vec{u}} $$
那么就可以利用$\vec{\psi}$重写目标函数$J$对某个设计变量$\beta_{i}$的敏感度
$$ \begin{align} \frac{\partial J}{\partial \vec{u}} \frac{\partial \vec{u}}{\partial \beta_{i}} &= \left(\frac{\partial \vec{R}}{\partial \vec{u}} \right)^T \vec{\psi} \frac{\partial \vec{u}}{\partial \beta_{i}} \\ &= \vec{\psi}^T \frac{\partial \vec{R}}{\partial \vec{u}} \frac{\partial \vec{u}}{\partial \beta_{i}} \\ &= \vec{\psi}^T \left( -\frac{\partial \vec{R}}{\partial \beta_{i}} \right) \\ \end{align} $$
从而得到
$$ \begin{align} \frac{d J}{d \beta_{i}} &= \frac{\partial J}{\partial \beta_{i}} + \frac{\partial J}{\partial \vec{u}} \frac{\partial \vec{u}}{\partial \beta_{i}} \\ &= \frac{\partial J}{\partial \beta_{i}} + \vec{\psi}^T \left( -\frac{\partial \vec{R}}{\partial \beta_{i}} \right) \\ \end{align} $$
$\frac{\partial J}{\partial \beta_{i}}$常常是容易计算的。 $\frac{\partial \vec{R}}{\partial \beta_{i}}$和$\frac{\partial \vec{R}}{\partial \vec{u}}$都可以利用自动微分之类的工具在正向由$\vec{u}, \vec{\beta}$计算$\vec{R}$时计算得到。 $\frac{\partial J}{\partial \beta_{i}}$是用于更新$\beta$的梯度。
从上述过程中可以看出，伴随方法是帮助计算梯度的一种方式， 优化方法还是基于梯度的优化方法。
伴随优化求解PDE反问题的例子
$$ \begin{align} & \text{min}_{\beta} J = \sum_{i}^{M} (u_{i} - u^{obs}_{i})^2 \\ & \text{subject to } \nabla \cdot (\beta \nabla u) = q \\ & \text{with } u = u_{bnd} \text{ on } \partial \Omega \\ & \text{with } \beta = \beta_{1} / u + \beta_{2} = 13221/u + 27 \\ \end{align} $$
问题定义： $\beta_{1}=13221, \beta_{2}=27$对应一个$\vec{\beta}$分布，进而有解$\vec{u}^{ref}$。 我们现在只有真解$\vec{u}^{ref}$，想知道$\beta_{1}=13221, \beta_{2}=27$对应的$\vec{\beta}$分布。
首先实现一个正向的计算，求解PDE方程，得到当前$\vec{\beta}$下解$\vec{u}$。 我们能够指定$\beta_{1}=13221, \beta_{2}=27$，得到真解$\vec{u}^{ref}$。 正向计算是外层Newton迭代包住内层线性求解迭代， 每步Newton迭代调用线性代数求解器收敛，外层直至收敛。 然后伴随优化是Newton迭代之上再包一层迭代。 每步伴随迭代计算$\frac{\partial J}{\partial \vec{\beta}}$，从而更新$\vec{\beta}$。
完整的代码库见jingchangshi/DemoAdjOpt， 基于Jax编写，从而可以直接利用Jax的自动微分和JIT加速技术。