PINNは減衰を再現できるか

動機

僕は情報系の研究室で、損失関数に支配方程式からなる誤差項を追加することにより、支配方程式の解を学習する　Physics Informed Neural Networks(PINNs)という手法を研究しています。
ちょうど年末の帰省で暇になり、振動論について久しぶりに勉強したら、PINNはどれまで振動現象を再現できるのか試してみたいなあと思いました。具体的に何が気になったのかは数学のお話のところも見てください。

数学のお話

運動方程式として、$\frac{d^2x}{dt^2}+a\frac{dx}{dt}+bx=0$を考えます。ここで、$a>0$の時には、この運動方程式は速度に比例した抵抗力を受けながら運動する物体の運動を記述することになります。
この特性方程式は、$\lambda^2+a\lambda+b=0$であり、判別式は$D=a^2-4b$となるため、この判別式によって、解は以下のように分類されます。

D<0の場合

この時、特性方程式は二つの虚数解を持ちます。書き下すと、
$\lambda = \frac{-a \pm \sqrt{a^2-4b}}{2}=-\frac{a}{2} \pm i\frac{\sqrt{4b-a^2}}{2}$
よって、一般解は、二つの積分定数$C_1,C_2$を用いて、
$$x(t) = e^{-\frac{a}{2}t}(C_1\cos\frac{\sqrt{4b-a^2}}{2}t + C_2\sin\frac{\sqrt{4b-a^2}}{2}t)$$
この運動は、$\cos\frac{\sqrt{4b-a^2}}{2}t + \sin\frac{\sqrt{4b-a^2}}{2}t$の部分から分かるように振動していますが、$e^{-\frac{a}{2}t}$がかかっているので、時間が経過するにつれてこの振動は上から抑えられていきます。このような状態を、減衰振動と呼びます。

D=0の場合

この時、特性方程式は一つの重解を持ちます。よって、一般解は次のように書けます（この部分の詳しい話は、微分方程式の本などに詳しい記載が必ずあるので、そちらを参照してください。
$$x(t) = C_1e^{-\frac{a}{2}t} + C_2 t e^{-\frac{a}{2}t}$$
指数の方に直接$-\frac{a}{2}t$が乗っているので、この運動は急速に減衰します。
この運動は、臨界減衰と呼ばれます。

D>0の場合

この時、特性方程式は相異なる二つの実数解を持ちます。一般解は以下のように書けます。
$$ x(t) = C_1 e^{\frac{-a+\sqrt{a^2-4b}}{2}t} + C_2 e^{\frac{-a-\sqrt{a^2-4b}}{2}t} $$

と書けます。これは振動することなく急速に減衰するため、過減衰と呼ばれます。

改めて、再度モチベーションの説明

上記の説明から、支配方程式の抵抗力の項の係数aがある程度大きければ過減衰となり、そこから徐々に小さくしていくと、特性方程式が重解を持つ点で臨界減衰となり、それ以降は減衰振動となることがわかるかと思います。（直感通りですね！）
上のように解析的に解くことはできますが、ただ単に運動方程式の係数が変わっただけで大きく運動の様子が異なるこの減衰を、PINNは本当に学習ができるのか気になったので、実装してみました。

手法

$\hat{x}$をニューラルネットワークの出力として、損失を以下のように定義。
$$ L = L_{initial}+L_{boundary}+L_{PDE}$$
ここで、$L_{initial}$は、初期条件であり、$|\hat{x}(0)-x(0)|^2$で定義。
$L_{boundary}$は、境界条件であり$| \frac{d \hat{x}(0)}{d t} - \frac{ dx(0)}{dt}|^2$で定義。
$L_{PDE}$は、コロケーションポイントと呼ばれる定義域上の点で取られた点で、支配方程式からどれだけずれているのかを表現する項で、コロケーションポイントの数を$N$とすると、
$L_{PDE}=\frac{1}{N} \sum_{i=1}^N | \frac{d^2\hat{x}(t_i)}{dt^2}+a\frac{d\hat{x}(t_i)}{dt} +b\hat{x}(t_i) |^2$で定義。
定義域は、$t\in[0,10]$、コロケーションポイントの数は100として、SGDで学習が可能か試してみました。

実験

減衰振動について

$\ddot{x}+\frac{1}{2}\dot{x}+\frac{17}{16}x=0$を考えます。この一般解は、$x(t)=e^{-\frac{1}{4}t}(C_1 \cos t + C_2 \sin t)$とかけます。ここで、条件として$x(0)=0,v(0)=1$とすると、この特殊解として、
$$ x(t) = e^{-\frac{1}{4}t}\sin t$$
が得られます。　
さて、次のコードで実験してみました。

main.py

import torch
from torch import nn
import torch.optim as optim
import numpy as np

class NN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1,20),
            nn.Tanh(),
            nn.Linear(20,20),
            nn.Tanh(),
            nn.Linear(20,1)
        )
    def forward(self,t):
        x = self.net(t)
        return x

#定義域を決定
t_train = torch.linspace(0,20,1000).unsqueeze(1)
t_train.requires_grad=True

def loss(t,model):
    #initial condition
    t0 = torch.tensor([0.0],requires_grad=True)
    u_at_t0 = model(t0)
    ini_loss = (u_at_t0 )**2

    #v(0)=1の損失
    du_dt_at0 = torch.autograd.grad(u_at_t0,t0,grad_outputs=torch.ones_like(u_at_t0),create_graph=True,retain_graph=True)[0]
    bound_loss = (du_dt_at0-1)**2


    u = model(t)
    du_dt = torch.autograd.grad(u,t,grad_outputs=torch.ones_like(u),create_graph=True,retain_graph=True)[0]
    du_dtt = torch.autograd.grad(du_dt,t,grad_outputs=torch.ones_like(du_dt),create_graph=True,retain_graph=True)[0]
    #支配方程式の残差からなる損失
    f = du_dtt + du_dt/2 + 17*u/16
    PDE_loss = nn.MSELoss()(f,torch.zeros_like(f))

    loss = ini_loss + bound_loss + PDE_loss

    return loss

num_epochs = 100000
model = NN()
optimizer = optim.Adam(model.parameters(), lr = 1e-3)

for epoch in range(num_epochs):
    optimizer.zero_grad()
    losses = loss(t_train,model)
    losses.backward()
    optimizer.step()
    if epoch % 100 ==0:
        print(f"epoch:{epoch}, loss:{losses}")

結果

お、なかなか良さそう。

過減衰

支配方程式として
$\ddot{x}+\frac{1}{2}\dot{x}+\frac{1}{18}x=0$を考える。この一般解は、$x(t)=C_1 e^{-\frac{t}{3}} + C_2 e^{-\frac{t}{6}}$である。特解として、$x(t)=e^{-\frac{t}{3}}$を採用すると、$x(0)=1,v(0)=-\frac{1}{3}$である。損失の部分だけ、以下のように書き換えて実行。

うん。なかなかやるじゃん。