PINN for Phonon BTE
PINN for Phonon BTE
> Li, R., Lee, E., & Luo, T. (2021). Physics-informed neural networks for solving multiscale mode-resolved phonon Boltzmann transport equation. Materials Today Physics, 19, 100429.
https://github.com/RuiyangLi6/PINN-pBTE/tree/main/BTE_2D_Square
作者开源了他的PINN程序，不过实在是有些不太好读，在读的时候总想起哈尔的移动城堡..
这篇文章一方面是翻译一下这些code，另一方面也是介绍一下采用确定性方法是怎么求解声子BTE的，以及用NN解声子BTE的Loss是如何构建的. 采用随机性方法（蒙特卡洛)求解BTE可以参考之前写过的蒙特卡洛中的频率抽样和动力学蒙特卡洛. 这篇文章里定义的函数完全采用了作者开源的程序，并没有进行修改.
求解体系
作者算了一些不同维度的case，这里拿上面这个2维体系作为例子，把两侧边界调整成周期性边界，上方热沉边界取窄再附加上漫反射边界条件凑成热流边界就变成了热扩展问题的典型结构.
控制方程
稳态SMRT近似下energy-deviation-based phonon BTE具有以下的形式， \[ v_g \boldsymbol{s} \cdot \nabla e=\frac{e^{e q}-e}{\tau} \] 对立体角和频率积分，可以得到能量守恒方程， \[ \nabla \cdot \boldsymbol{q}=\sum_{p} \int_{0}^{\omega_{\max , p}} \int_{4 \pi} \frac{e^{e q}-e}{4\pi \tau} d \Omega d \omega=0 \] 其中热流的表达式为（这里作者的文章里有个typo，没有除以\(4\pi\)，不过不影响结果） \[ \boldsymbol{q}=\sum_{p} \int_{0}^{\omega_{\max , p}} \int_{4 \pi} \frac{\boldsymbol{v} e}{4\pi} d \Omega d \omega \] 联立BTE和能量守恒方程并给出合适的边界条件，我们就可以求解出分布函数的分布，并最终得到体系的温度分布. 在具体介绍PINN如何离散并求解它们之前，需要先弄清楚分布函数是哪些量的函数，从而才能理解我们的input向量究竟是什么形式的.
\(e(\boldsymbol{x}, \boldsymbol{s}, k, p) = \hbar\omega D(\omega, p)\left[ f(\boldsymbol{x}, \boldsymbol{s}, k, p) - f^{BE}(T_{ref})\right]\)是声子能量偏差函数，\(e^{eq}(k, p, T) = \hbar\omega D(\omega, p)\left[f^{BE}(T) - f^{BE}(T_{ref})\right]\)是平衡声子能量偏差函数. 在选定了\(T_{ref}\)以后，\(f^{BE}(T_{ref})\)就是确定的，求解\(e\)就等价于求解\(f\). 之所以定义这样一个能量偏差函数，是因为这样做之后平衡分布\(e^{eq}\)的可以做出一定近似，从而简化了求解. \[ f^{BE}(T) = \frac{1}{\exp[\hbar\omega/ (k_BT)]}, \quad D(\omega, p)=\frac{k^2}{2\pi^2 v_g} \] \(D(\omega, p)\)为声子态密度，色散关系描述了\(\omega\)和\(k\)之间的依赖关系，同时给出了\(\omega_m\)频率上限，根据色散关系我们可以求出群速度\(v_g=\partial \omega / \partial k\). 在求解phonon BTE之前，弛豫时间和色散关系需要我们预先确定好，所以当色散关系给定以后，\(D(\omega, p)\)这一项我们是完全已知的. 声子支的含义是不同支的色散关系和弛豫时间是不一样的，比如下面这张图分别是Si的横支声子（TA）色散和纵支声子（LA）色散. 由于光学支群速度很小，对传热贡献不大，所以在声子BTE的求解过程中，我们一般只考虑三个声学，即一个LA和两个TA. 假设我们一个声子支离散成10段，那么三个声子支就需要离散成30段，考虑更多的声子支的处理方式都是一样的. \(D(\omega, p)\)里面的\(p\)，就来自于不同声子支的色散关系的区别.
在稳态下声子分布函数\(f\)是空间坐标\(\boldsymbol{x}\)，方向向量\(\boldsymbol{s}(\cos\theta, \sin\theta\cos\phi, \sin\theta\sin\phi)\)，波数\(k\)（等价于角频率\(\omega = \omega(k, p)\)，他们之间的关系由给定的声子色散已经确定了），声子支\(p\)的函数. 也就是说如果要采用确定性方法来求解\(f\)的话，我们需要建立一个多维数组来描述\(f\)的分布，最外层是空间分布，空间上的每一个点上都储存着一个三维数组，相当于空间中的每个点都储存着一个球，球半径方向的距离代表波矢大小\(k\)，球面的方向为\(\boldsymbol{s}\). 如果我们在空间上离散了\(Nx\)个点，\(\theta\)方向离散了\(N_1\)个点，\(\phi\)方向离散了\(N_2\)个点（也就是把整个球面立体角离散了\(N_s = N_1 \times N_2\)个点），声子波矢离散\(N_k\)个点（\(N_k\)已经包含了各个声子支的离散），那么最终描述\(f\)的数组长度就是\(N_x\times N_s \times N_k\).
小温差近似
在小温差近似下，\(e^{eq}\)可以近似为 \[ e^{eq}(k, p, T) = \hbar\omega D(\omega, p)\left[f^{BE}(T) - f^{BE}(T_{ref})\right] \approx C(\omega, p)\left(T-T_{r e f}\right) \] 其中 \[ C(\omega, p) = \hbar\omega D(\omega, p) \frac{\partial f^{BE}}{\partial T} \] 这个近似在蒙特卡洛模拟中计算温度时也广泛用到，这个简化的核心是实现了频率和温度的解耦，我们是基于能量守恒方程来获得体系的温度分布的. 从上面的能量守恒方程中可以看到，我们需要对整个声子频谱做积分. 如果频率和温度是耦合的，不同温度下这个积分的数值都需要重新计算，想要从积分值反推回温度也是一个很复杂的非线性过程. 而一旦频率和温度解耦了，我们就相当于把温度从积分中提出来了，不同温度下这个积分值都是一样的，从而简化了求解.
把这个近似代入到能量守恒方程中，可以得到 \[ (T - T_{ref}) \sum_p \int_0^{\omega_m} \frac{4\pi C(\omega, p) }{\tau}\mathrm{d}\omega = \sum_p \int_0^{\omega_m} \int_{4\pi} \frac{e }{\tau} \mathrm{d}\Omega \mathrm{d}\omega \] 于是温度就可以表示成， \[ T - T_{ref} = \frac{1}{4\pi} \left(\sum_p \int_0^{\omega_m} \int_{4\pi} \frac{e}{\tau} v_g \mathrm{d}\Omega \mathrm{d}k \right)\times \left(\sum_p \int_0^{\omega_m} \frac{C(\omega, p) }{\tau}v_g\mathrm{d}k\right)^{-1} \] 多出来的群速度是把频率积分转换成波矢积分时产生的.
PINN的网格离散
下面的程序源自于作者提供的mesh_2d.py文件，空间网格采用Sobol sequence离散450个点，文章这里做了一个非均匀的空间离散，靠近上面热源的地方多撒了些点. 立体角采用高斯-勒让德插值离散\(12\times 12\)个网格，离散对象为\(\cos\theta\in [-1, 1]\)和\(\phi \in [0, \pi]\)，下面这个TwoD_train_mesh(Nx, Nt, N1, N2, Nk)这几个输入参数分别是Nx空间点的总数，Nt大概数决定了上面的点的密度，N1是\(\cos\theta\)离散的网格数量，N2是\(\phi\)离散的网格数量，Nk是一个声子支离散的网格数量，作者对每一支离散了10个网格. param(k, p, T)是采用文章中指定的色散模型计算对应\(k, p\)处的v（群速度）,tau（弛豫时间）,C（比热），p=0是TA支，p=1是LA支. TwoD_vt(k,Tr)输入离散后的波矢\(k\)数组和参考温度，输出TA支和LA支的声子自由程. TwoD_test_mesh(Nx,N1,N2,Nk)这个和TwoD_train_mesh差不多，只不过是生成测试集的离散网格，这时候不需要对上面做加密.
1 | # File: mesh_2d.py |
看一下离散出来的点的分布，
1 | # 用于观察空间离散点的形状 |
传统数值方法求解PDE需要依赖网格间差分来近似梯度项，但是NN直接建立了input和output的关系，直接对NN求梯度就可以了，网格点之间也没有直接的相互作用，因此也不需要建立规则的网格. 但是NN预测的效果很依赖于在体系温度梯度大撒上更多的点，而且每换一个体系、换一个边界条件或者换一个材料就需要重新train一个NN，想用它来处理某些实际问题我不太觉得是一个很好的选择.. 在反问题的求解上也是一个麻烦的方案.
NN搭建
这篇文章train了两个NN，Net1用来预测平衡部分\(e^{eq}\)，Net0用来预测非平衡部分\(e - e^{eq}\)，从小温差近似可以看到 \[
e^{e q}(k, p, T) / C(\omega, p) \approx \left(T-T_{r e f}\right)
\] \(e_{eq}\)除以对应模态的比热得到的就是温度，而温度只和空间分布有关，\(C(\omega, p)\)又是我们已知的，因此用来预测平衡部分的NN的输入没有必要是\((\boldsymbol{x}, \boldsymbol{s}, k, p)\)这么高维度的向量，只是空间坐标的函数就可以了. 在得到了预测值后，我们可以再乘上对应模态比热还原回去. 我们可以再用体系的特征温差\(\mathrm{d}T\)把输出调整到无量纲的形式（\(T_{h} - T_{ref} = T_{ref} - T_{c} = \mathrm{d}T\)）. 此时Net1的输入和输出就是 \[
\operatorname{Net1}(x, y, L) = \frac{e^{e q}(k, p, T)}{C(\omega, p)\mathrm{d}T} \approx \frac{T-T_{r e f}}{\mathrm{d}T}
\] L是体系的特征尺寸，我们希望训练好的NN能够预测同一体系不同特征尺寸下的结果，比如希望网络能够预测10nm ~ 100um的结果，那么输入的L就可以选为[0, 1, 2, 3, 4]，每一个数值对应一个对数后的特征尺度. 我们让Net0和Net1的输出量纲是一样的，这时比较好对他们统一操作，于是Net0的输入和输出就是 \[
\operatorname{Net0}(x, y, \cos\theta, \phi, k, MFP, p, L) = \frac{e^{neq}(x, y, \cos\theta, \phi, k, p)}{C(\omega, p)\mathrm{d}T} = \frac{e(x, y, \cos\theta, \phi, k, p) - e^{eq}(k, p, T)}{C(\omega, p)\mathrm{d}T}
\] 这个MFP就是对应\(k\)处的自由程，\(MFP = v_g \times \tau\). （我不知道为什么这里要输入MFP，或许是作者尝试后把它也当作input训练效果比较好吧）在有了Net0和Net1以后，我们就可以得到\(e=e^{neq} + e^{eq}\)了，通过\(e\)就可以定义内部节点和边界节点处的Loss. 由于声子的自由程跨度很大且体系的特征尺寸也是跨量级变化的，因此输入到网络中的\(MFP\)实际上是\(\log(MFP\times 10^{11})\). 为了让输出也保持在相同量级上，不至于让各层权重变化太大，作者的实际网络输出的值在上述的原始表达式上又除以了$ $.
1 | # File: model.py |
Loss构建
Loss有三个部分.
一是各频率的声子需要满足玻尔兹曼方程， \[
\begin{aligned}
& \Vert \mathbf{v_g} \cdot \nabla e - \frac{e^{neq}}{\tau} \Vert^2 = 0 \\
\Rightarrow & \Vert v_g \times (\cos\theta, \sin\theta\cos\phi) \cdot (\frac{\partial e}{\partial (Lx^*)}, \frac{\partial e}{\partial (Ly^*)}) - \frac{e^{neq}}{\tau} \Vert^2 = 0 \\
\Rightarrow & \Vert \cos\theta \frac{\partial e}{\partial x^*} + \sin\theta\cos\phi \frac{\partial e}{\partial y^*} - \frac{e^{neq}}{\text{MFP}/L} \Vert^2 = 0
\end{aligned}
\] 下面程序里的e0 = net0(e0_in)*(10**(vt0-L))*dT和e1 = net0(e1_in)*(10**(vt1-L))*dT分别是TA支和LA支的非平衡部分，即\(e^{neq} / C(\omega, p)\)，eEq = net1(torch.cat((x,L),1))*dT是平衡部分，即\(e^{eq} / C(\omega, p)\). 后面的 1
2
3
4e0_x = torch.autograd.grad(e0+eEq,x,grad_outputs=torch.ones_like(x).to(device),create_graph=True)[0]
e1_x = torch.autograd.grad(e1+eEq,x,grad_outputs=torch.ones_like(x).to(device),create_graph=True)[0]
e0_y = torch.autograd.grad(e0+eEq,y,grad_outputs=torch.ones_like(y).to(device),create_graph=True)[0]
e1_y = torch.autograd.grad(e1+eEq,y,grad_outputs=torch.ones_like(y).to(device),create_graph=True)[0]1
2loss_1 = ((mu*e0_x+eta*e0_y) + e0/(10**vt0/10**L))/dT
loss_2 = ((mu*e1_x+eta*e1_y) + e1/(10**vt1/10**L))/dTmu和eta分别是离散的\(\cos\theta\)和\(\sin\theta\sin\phi\)数组.
二是声子能量守恒方程， \[
T - T_{ref} = \frac{1}{4\pi} \left(\sum_p \int_0^{\omega_m} \int_{4\pi} \frac{e}{\tau} v_g \mathrm{d}\Omega \mathrm{d}k \right)\times \left(\sum_p \int_0^{\omega_m} \frac{C(\omega, p) }{\tau}v_g\mathrm{d}k\right)^{-1}
\] 下面程序里的 1
sum_e = torch.matmul(e.reshape(-1,Ns), w).reshape(-1,1).to(device)
delta_T就是按照上述公式计算的得到的体系温差
1 | deltaT = torch.matmul(sum_e.reshape(-1,Nk*Np),C*wk/tau*v/(4*np.pi)).reshape(-1,1)\ |
而Net1的输出是 \[
\operatorname{Net1}(x, y, L) = \frac{e^{e q}(k, p, T)}{C(\omega, p)\mathrm{d}T} \approx \frac{T-T_{r e f}}{\mathrm{d}T}
\] 所以我们让eEq = net1(torch.cat((x,y,L),1))*dT，就可以定义这部分的Loss了， 1
loss_3 = (deltaT - eEq)/dT
\[
\nabla \cdot \boldsymbol{q}=\sum_{p} \int_{0}^{\omega_{\max , p}} \int_{4 \pi} \frac{e^{e q}-e}{4\pi \tau} d \Omega d \omega=0
\] 其中热流的表达式为 \[
\boldsymbol{q}=\sum_{p} \int_{0}^{\omega_{\max , p}} \int_{4 \pi} \frac{\boldsymbol{v_g} e}{4\pi } d \Omega d \omega
\] 我们可以直接代入热流，求散度，也可以得到一个能量守恒的Loss， \[
\nabla \cdot \boldsymbol{q}=\sum_{p} \int_{0}^{\omega_{\max , p}} \int_{4 \pi} \left( \frac{\partial e}{\partial x}v_g \cos\theta + \frac{\partial e}{\partial y}v_g\sin\theta\sin\phi \right)\frac{v_g}{4\pi} d \Omega d k
\] 1
2
3sum_ex = torch.matmul(e_x.reshape(-1,Ns), w*mu[0:Ns].reshape(-1,1)).reshape(-1,1)
sum_ey = torch.matmul(e_y.reshape(-1,Ns), w*eta[0:Ns].reshape(-1,1)).reshape(-1,1)
dq = torch.matmul((sum_ex+sum_ey).reshape(-1,Nk*Np),C*wk*v**2/(4*np.pi)).reshape(-1,1)C(omega, p)是因为作者定义的\(e\)变量实际上是\(e/C(\omega, p)\). 这个TC相当于对不同尺度的数值做一下调整. 1
loss_4 = (dq/TC)
第三部分是边界条件，等温边界吸收所有射向它的声子，并以边界温度的平衡分布向体系内部发射声子， \[ e(\mathbf{x_b}, \vec{s}, k, p) = e^{eq}(k, p, T_b), \vec{s} \cdot \vec{n_b} > 0 \] 等价于 \[ e(\mathbf{x_b}, \vec{s}, k, p) / C(\omega, p) = T_b - T_{ref}, \vec{s} \cdot \vec{n_b} > 0 \] 程序下面的ec0和ec1代表冷边界分布，eh0和eh1代表热边界分布，ec1和eh1是在冷热边界交界处附近额外增加了一些点，可以看Main Loop中的程序. 文章里假设\(T_h - T_{ref} = T_{ref} - T_b = \mathrm{d}T = 0.5\). 根据前面对NN的input和output的讨论可以看到，程序里的ec和eh实际代表着 \[ \frac{e}{C(\omega, p)\mathrm{d}T} \] 于是等式就变成了 \[ \frac{e}{C(\omega, p)\mathrm{d}T} = \frac{T_b - T_{ref}}{\mathrm{d}T} = \pm 1 \] 要求\(\vec{s} \cdot \vec{n_b} > 0\)，只选择一半的立体角数组代入边界条件就可以了. 于是就可以定义边界条件的Loss了，
1 | loss_5 = (ec1 + 1) |
散射边界的表达式为 \[ e\left(\boldsymbol{x}_{b}, \mathbf{s}, k, p\right)=\frac{1}{\pi} \int_{\mathbf{s}^{\prime} \cdot \boldsymbol{n}_{b}<0} e\left(\boldsymbol{x}_{b}, \mathbf{s}^{\prime}, k, p\right)\left|\mathbf{s}^{\prime} \cdot \boldsymbol{n}_{b}\right| d \Omega, s \cdot \boldsymbol{n}_{b}>0 \] 周期性边界的表达式为 \[ e\left(\boldsymbol{x}_{b_{1}}, \boldsymbol{s}, k, p\right)-e^{e q}\left(k, p, T_{b_{1}}\right)=e\left(\boldsymbol{x}_{b_{2}}, \boldsymbol{s}, k, p\right)-e^{e q}\left(k, p, T_{b_{2}}\right) \] 都可以用类似的方式来处理，至此我们得到了用NN训练Phonon BTE所有的Loss.
1 | # File: bte_train.py |
主程序
1 | python main.py |
可以开始训练了. 模拟的体系变了之后，需要重新train所有的net.
1 | # main.py |