【PINNs】Helmholtz方程式を解く

①はじめに

昨今、データドリブンアプローチによる学習が盛んであるが、データドリブンアプローチでは内挿域における精度は良好であるが、外挿域では精度が悪化することが一般に知られている。その対策として、事前学習済みのモデルを活かして別タスクに適応させる転移学習、Fine-Tuningが一般である。
　一方、データドリブンアプローチとは異なるアプローチの手段として、方程式に則ったLoss関数を定義し学習させることで、方程式ベースの学習モデルで内挿域から外挿域までの予測精度を高めようとするPhysics Informed Neural Networks(通称、PINN もしくは PINNs）の開発が盛んになりつつある。
本記事では、自分自身のPINNに対する理解を深めるため、helmholtz方程式に則ったLoss関数を定義し学習させることとした。helmholtz方程式を選定したのは、24年4月時点で日本国内の記事では少ないため（一方、Burgers方程式の記事は多くある）

②PINNの作り方

1.方程式を定義する
2.初期条件を定義する
3.境界条件を定義する
4.Loss関数を定義する
基本的にこのような手順が必要である。
ではHelmholtz方程式(時間依存のない2次元波動方程式)で見ていこう

方程式（左辺が方程式、右辺が厳密解）
$$ \frac{\partial^2 u}{\partial x^2} +\frac{\partial^2 u}{\partial y^2} +(4\pi)^2u = (4π)^2sin(4πx)sin(4πy) $$

$$ x\in[0,1], y\in[0,1] $$

境界条件
$$ u(x=0,y)=u(x,y=0)=0, u(x=1,y)=u(x,y=1)=0 $$

初期条件は時間項がないため無し
ここまでで登場人物が出そろった。

では少し手作業でLoss関数を仕上げていく。
まずは方程式の式変形
$$ Loss_{PINN}=\frac{\partial^2 u}{\partial x^2} +\frac{\partial^2 u}{\partial y^2} +(4\pi)^2u - (4π)^2sin(4πx)sin(4πy) $$
簡単に書くと
$$ Loss_{PINN}=\text{方程式} - \text{厳密解} $$
となる。では、厳密解がない場合どうしたらよいか？
CAEなどの数値シミュレーションの結果を厳密解の代理にすればよいので、
$$ Loss_{PINN}=\text{方程式} - \text{CAE値} $$
となる。

これらの$Loss_{PINN}$を0になるように学習すれば方程式に則る。
同様に、境界条件も任意の$u$が境界条件上では0になるよう$Loss_{bound}$ を仕上げる
$$ Loss_{bound} = \sum|| u(x_{i},y_{i}) - u_{bc}(x_{i},y_{i}) ||^2 $$

そして、$ Loss_{PINN}$と$Loss_{bound}$の合計を$Loss$関数として0になるように学習させる
$$ Loss = Loss_{PINN} + Loss_{bound} $$

PINNSの山場は$Loss$関数の設定なのでこれが出来たらあとは学習するだけ。
それでは実装コードを次の章で説明

③実装してみる

0.準備
↓
Ⅰ.Loss関数の設定
↓
Ⅱ.空間のサンプリング
↓
Ⅲ.DNNの設定
↓
Ⅳ.学習
↓
Ⅴ.検証
この流れで説明します

③-0.準備

pytorchで構築しますので、必要なライブラリをインポートします

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import scipy.io as sio
from scipy.io import loadmat

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import grad
np.random.seed(42)
torch.manual_seed(42)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

③-Ⅰ.Loss関数の設定

②章で言及したLossをコーディングするだけ
ポイントはtorch.autograd.gradを使って偏微分方程式を間違えないように記述するのみ。

criteria = nn.MSELoss()
# np.piをtorch.tensorに変換する(GPU使用)
pi_tensor = torch.tensor(np.pi,device=device)

def physics_informed_loss(x, y, net):
  u = net(x, y)

  u_x = torch.autograd.grad(
    u, x,
    grad_outputs=torch.ones_like(u),
    retain_graph=True,
    create_graph=True,
    allow_unused=True
  )[0]
  u_xx = torch.autograd.grad(
    u_x, x,
    grad_outputs=torch.ones_like(u_x),
    retain_graph=True,
    create_graph=True,
    allow_unused=True
  )[0]

  u_y = torch.autograd.grad(
    u, y,
    grad_outputs=torch.ones_like(u),
    retain_graph=True,
    create_graph=True,
    allow_unused=True
  )[0]

  u_yy = torch.autograd.grad(
    u_y, y,
    grad_outputs=torch.ones_like(u_y),
    retain_graph=True,
    create_graph=True,
    allow_unused=True
  )[0]

  # Loss_PINNの定義
  pinn_loss = (u_xx)+( u_yy )+( (4*pi_tensor)*(4*pi_tensor)*u ) - (4*pi_tensor)*(4*pi_tensor)*torch.sin(4*pi_tensor*x)*torch.sin(4*pi_tensor*y)
  # Loss_PINNをゼロにするように設定(MSE)
  zeros_t = torch.zeros(pinn_loss.size()).to(device) 
  pinn_loss_ = criteria(pinn_loss, zeros_t)
  return pinn_loss_

def boundary_condition_loss(x, y, net, u_bc):
  u = net(x, y)
  bc_condition_loss = criteria(u, u_bc)
  return bc_condition_loss

#
x_bc = np.array([0.0, 1.0])
y_bc = np.array([0.0, 1.0])
bc_sample_size = 500
t_bc = np.linspace(0, 1.0, bc_sample_size)

# x = 0.0
X_bc1 = np.zeros([bc_sample_size, 3])
X_bc1[:,0] = 0.0
X_bc1[:,1] = t_bc
X_bc1[:,2] = y_bc[0]

# x = 1.0
X_bc2 = np.zeros([bc_sample_size, 3])
X_bc2[:,0] = 1.0
X_bc2[:,1] = t_bc
X_bc2[:,2] = y_bc[1]

# y = 0.0
Y_bc1 = np.zeros([bc_sample_size, 3])
Y_bc1[:,0] = t_bc
Y_bc1[:,1] = 0.0
Y_bc1[:,2] = x_bc[0]

# y = 1.0
Y_bc2 = np.zeros([bc_sample_size, 3])
Y_bc2[:,0] = t_bc
Y_bc2[:,1] = 1.0
Y_bc2[:,2] = x_bc[1]

# stack
X_bc_stack = np.vstack([X_bc1, X_bc2, Y_bc1, Y_bc2])
u_bc_stack = np.zeros(X_bc_stack.shape[0])

X_bc_t = torch.tensor(X_bc_stack, requires_grad=True).float().to(device)
u_bc_t = torch.tensor(u_bc_stack, requires_grad=True).float().to(device).unsqueeze(dim=1)

③-Ⅱ.空間のサンプリング

$x,y$を5000点サンプリングしておく

# sampling point
x_ = np.linspace(0,1,100)
y_ = np.linspace(0,1,100)
X, Y = np.meshgrid(x_, y_, indexing='ij')
x_flat = X.flatten()
y_flat = Y.flatten()

# sampling size
sampling_size = 5000
random_idx = np.random.choice(np.arange(x_flat.shape[0]),size=sampling_size,replace=False)

#
# sampling
#
x_sampled = x_flat[random_idx]
y_sampled = y_flat[random_idx]
X_sampled = np.zeros([sampling_size,2])
X_sampled[:,0] = x_sampled
X_sampled[:,1] = y_sampled
X_sample_t = torch.tensor(X_sampled,requires_grad=True).float().to(device)

③-Ⅲ.DNNの設定

入力は$x,y$の2因子、隠れ層は8層、最終層はLinearでつないで$u$を求める

class PINN(torch.nn.Module):
  def __init__(self,activation='relu'):
    super().__init__()
    self.regressor = nn.Linear(5, 1)
    self.activation = activation
    self.linear1 = self.linear(2,5,activation = self.activation)
    self.linear2 = self.linear(5,20,activation = self.activation)
    self.linear3 = self.linear(20,80,activation = self.activation)
    self.linear4 = self.linear(80,160,activation = self.activation)
    self.linear5 = self.linear(160,160,activation = self.activation)
    self.linear6 = self.linear(160,80,activation = self.activation)
    self.linear7 = self.linear(80,40,activation = self.activation)
    self.linear8 = self.linear(40,5,activation = self.activation)

  def linear(self, in_features, out_features, activation='relu'):
    layers = [nn.Linear(in_features, out_features)]
    if activation == 'relu':
      layers.append(nn.ReLU(inplace=True))
    elif activation == 'tanh':
      layers.append(nn.Tanh())
    else:
      layers.append(nn.Sigmoid())
    net = nn.Sequential(*layers)
    return net
  def forward(self,x,t):
    inputs = torch.cat([x,t],axis=1)
    out = self.linear1(inputs)
    out = self.linear2(out)
    out = self.linear3(out)
    out = self.linear4(out)
    out = self.linear5(out)
    out = self.linear6(out)
    out = self.linear7(out)
    out = self.linear8(out)
    out = self.regressor(out)
    return out

③-Ⅳ.学習

活性化関数はReLU,学習率は0.001,エポック数は10000で学習させる

net = PINN(activation='relu').to(device)
optimizer = optim.Adam(net.parameters(), lr=0.001)
num_epoch = 10000

losses=[]
pinn_losses=[]
bnd_losses=[]

for epoch in range(num_epoch):
  optimizer.zero_grad()

  # PINN loss
  x_sampled = X_sample_t[:,0].unsqueeze(dim=-1).to(device)
  y_sampled = X_sample_t[:,1].unsqueeze(dim=-1).to(device)
  pinn_loss = physics_informed_loss(x_sampled, y_sampled, net)

  # bc loss
  x_bnd = X_bc_t[:,0].unsqueeze(dim=-1).to(device)
  y_bnd = X_bc_t[:,1].unsqueeze(dim=-1).to(device)
  bnd_loss = boundary_condition_loss(x_bnd,y_bnd,net,u_bc_t)
  # Loss_total
  loss = pinn_loss  + bnd_loss
  loss.backward()
  optimizer.step()
  losses.append(loss.item())
  pinn_losses.append(pinn_loss.item())
  bnd_losses.append(bnd_loss.item())

  if epoch % 500 == 0:
    loss_ = loss.item()
    pinn_loss_ = pinn_loss.item()
    #ini_loss_ = ini_loss.item()
    bnd_loss_ = bnd_loss.item()
    print(f'epoch:{epoch:.3e},loss:{loss_:.3e},pinn:{pinn_loss_:.3e},bnd:{bnd_loss:.3e}')

学習の結果、$Loss$の推移から正常に学習ができている様子が見て取れる

$Loss_{PINN}$も正常に学習できている様子が見て取れる

$Loss_{boundary}$もオーダーは小さいが
最初オーバーシュート気味だった。
後半にかけて正常に学習できている様子が見て取れる

③-V.検証

作ったPINNで推論してみる

X_test = np.zeros([x_flat.shape[0],2])
X_test[:,0] = x_flat
X_test[:,1] = y_flat

X_test_t = torch.tensor(X_test).float().to(device)
x_test = X_test_t[:,0].unsqueeze(dim=-1)
y_test = X_test_t[:,1].unsqueeze(dim=-1)
u_pred = net(x_test , y_test)

xtest=x_test.cpu().numpy()
ytest=y_test.cpu().numpy()
upred=u_pred.cpu().detach().numpy()

厳密解とPINNの予測結果を比較する

# x, y の範囲と刻み幅の設定
x_range = np.arange(0, 1.001, 0.001)
y_range = np.arange(0, 1.001, 0.001)

# exact_dataを格納する辞書の初期化
exact_data = {'x': [], 'y': [], 'z': []}

# x, y のループを回して exact を計算し exact_data に格納
for x in x_range:
    for y in y_range:
        #exact = (4*np.pi)*(4*np.pi)*np.sin(4*np.pi*x)*np.sin(4*np.pi*y)
        exact = np.sin(4*np.pi*x)*np.sin(4*np.pi*y)
        exact_data['x'].append(x)
        exact_data['y'].append(y)
        exact_data['z'].append(exact)

# exact_dataをnumpy配列に変換
exact_data['x'] = np.array(exact_data['x'])
exact_data['y'] = np.array(exact_data['y'])
exact_data['z'] = np.array(exact_data['z'])
# exact_dataからx, y, zを取得
x = exact_data['x']
y = exact_data['y']
z = exact_data['z']

# PINNの結果のプロット
plt.figure(figsize=(16, 6))

plt.subplot(1, 2, 1)  # 1行2列の1番目のサブプロット
plt.scatter(xtest, ytest, c=upred, cmap='viridis')
plt.colorbar(label='prediction')
plt.xlabel('x')
plt.ylabel('y')
plt.title('PINNS')

# 厳密解のプロット
plt.subplot(1, 2, 2)  # 1行2列の2番目のサブプロット
plt.contourf(x.reshape(1001, 1001), y.reshape(1001, 1001), z.reshape(1001, 1001), cmap='viridis')
plt.colorbar(label='exact solution')
plt.xlabel('x')
plt.ylabel('y')
plt.title('exact solution')

plt.tight_layout()  # レイアウトの調整
plt.show()

これらを実行すると下図の結果が得られる
左がPINNS　右が厳密解

しっかりと方程式通りの結果が得られていることを確認することが出来た

まとめ

今回はPINNSの実装を通して、より理解することができた
また、CAEのサロゲートモデルを構築する際に有効な手段となる可能性があるなと感じることもできた。一方、multi-physicsの問題ではどの方程式で学ばせるか？など難儀なところもあるなと感じた。今回実装したコードがGithubにて公開する.