[PIDoN] Burger's Equation 예제 코드 분석 [1/2]

 예상치 못하게 학부 때 하던 인공지능이나 해석과 동 떨어진 일을 하게 되어, 근 1년 간 딥러닝에 손을 떼고 있었다. 아무리 안 쓰더라도 마냥 손 떼기도 그렇고, 언젠가는 내가 하는 일에도 공부했던 것들을 적용하기를 바라며, 다시 하나둘 공부해 봐야겠다는 생각이 들었다. 그에 대한 일환으로 일단 학부 때 했던 것들을 기억을 되살려 하나씩 기록하려고 한다. 가장 먼저 할 건 귀찮아서 미뤄뒀던 Physics-Informed DeepONet(PIDoN) 예제 코드에 대한 포스팅이다. 본 포스팅에 사용할 예제 코드는 "Learning the solution operator of parametric partial differential equations with physics-informed DeepONets" 논문의 Section "Burgers transport dynamics"의 예제 코드를 바탕으로 한다.

 

 PIDoN에 대한 간단한 설명은 아래 글을 통해 확인할 수 있다. 사실 개인적으로는 PINO에 비해 돌아가는 매커니즘 자체가 직관적이기도 하고, 연산자라는 개념 자체가 워낙 막강하다고 생각해, 언젠간 응용해 볼 수 있지 않을까 하는 기대감이 있다. (교수님 말씀으로는 연산자 학습에 대한 연구가 워낙 빠르게 되고 있어, PIDON, PINO 같은 것들도 이젠 구식이라곤 하지만...) 

 

2024.02.11 - [인공지능/Deep Learning] - [PIDoN] Physics-Informed DeepONet에 대하여

[PIDoN] Physics-Informed DeepONet에 대하여

 

Problem Setting

 

왜 물리 기반 인공지능들은 다들 Burger's Equation부터 시작하는지는 모르겠지만 (아마 가장 간단한 유체 방정식이라서 아닐까 싶다.), 이번 예제 코드도 Burger's Equation이며, 문제는 다음과 같다.

 

$$\frac{ds}{dt} + s\frac {ds}{dx} - ν\frac {d^2s}{dx^2},\quad (x, t) ∈ (0, 1) × (0, 1]$$

$$s(x, 0) = u(x), \quad x ∈ (0, 1)$$

 

$t∈(0,1)$ 이며, 초기 조건 $u(x)$는 $GRF \sim 𝒩(0,25^2(−Δ + 5^2I)^{−4})$ 내에서 생성되며, $v$는 점도 계수로 본 예제에서는 0.01의 수치가 사용된다. 경계 조건은 Periodic Boundary Condition으로 상세 경계 조건은 아래와 같다.

 

$$s(0, t) = s(1, t),$$

$$\frac {ds}{dt}(0, t) = \frac {ds}{dt}(1, t)$$

 

본 예제 코드는 초기 조건 $u(x)$를 통해 1D Burger's Equation의 전체 시공간 솔루션 $s(x,t)$에 매핑하는 솔루션 연산자를 학습한다.

 

Jax 라이브러리

 

 본 예제 코드는 jax를 기반으로 하고 있기에, 해당 코드에 사용되는 라이브러리에 대해 간단히 알아보고 가자.

 

random : 무작위 난수를 생성하는 모듈

grad : 함수를 미분하기 위한 모듈

vmap : 벡터화 매핑(배열 연산과 관련하여, 함수를 배열의 각 요소에 대해 자동으로 벡터화)을 제공하는 모듈

jit : Just-In-Time(JIT) 컴파일을 제공하는 모듈, 이로 인해 함수를 효율적으로 최적화된 컴파일 버전으로 변환하고, 더 빠르고 최적화된 수치 연산 실행 가능

optimizers : 최적화와 관련된 모듈

config : jax 라이브러리의 구성 설정에 관련된 기능을 제공하는 모듈

ravel_pytree : 데이터를 구성하는 복잡한 구조들을 1D 배열로 flatten 시켜주는 모듈

relu, elu : 활성화 함수를 제공하는 모듈

 

import numpy as onp
import scipy.io
from scipy.interpolate import griddata
import jax.numpy as np
from jax import random, grad, vmap, jit
from jax.example_libraries import optimizers
from jax.config import config
from jax.flatten_util import ravel_pytree
from jax.nn import relu, elu
import itertools
from functools import partial
from torch.utils import data
from tqdm import trange
import matplotlib.pyplot as plt

%matplotlib inline

 

Define Model

 

 PIDoN 모델을 학습시키기에 앞서 해당 모델에 사용되는 다층 퍼셉트론 모델(Multi-Layer Perceptron, MLP)과 Data Generator를 정의한다. MLP는 여러 개의 퍼셉트론 뉴런을 여러 층 쌓은 다층신경망 구조로, 입력층과 출력층 사이에 하나 이상의 은닉층을 가지고 있는 신경망이다. 흔히 역전파를 사용하여, 비선형 데이터를 해석하기 위해 사용되며, 본 예제 코드는 MLP 외에도 Modified MLP 모델을 넣어, 두 모델의 성능을 비교한다.

 

 Modified MLP는 기존의 Weight와 Bias을 단일로 이용한 MLP와 달리,  추가적인 난수로 생성 후, Input을 바탕으로 U와 V라는 변수를 만든다. 이를 MLP와 같은 원리로 생성된 output에 복합적으로 적용함으로써, 일반 MLP에 비해 더 큰 비선형성을 보여준다. Modified MLP를 MLP와 비교할 때 장점은 아래와 같다.

 

1. 더 강력한 비선형성

2. 복잡한 문제에 적합하다.

3. 모델의 유연성 향상 및 표현 능력 증가

 

 Modified MLP의 모델 구조는 아래와 같다.

 

$$ U = \phi(XW^1+b^1), V = \phi(XW^2+b^2) $$ $$ H^{(1)}=\phi(XW^{z,1}+b^{z,1}) $$ $$ Z^{(k)}=\phi(H^{(k)}W^{z, k}+b^{z, k}), k = 1,..., L $$ $$ H^{(k+1)}=(1-Z^{(k)})⊙U+Z^{(k)}⊙V, k=1,..., L $$ $$ f_{\theta}(x)=H^{(L+1)}W+b $$

 

 이러한 점이 적용되는 부분을 아래 코드에 표기해 두었다. 추가로 특이한 점은 난수를 생성할 때 random.PRNGkey라는 함수를 사용하였다는 점인데, 이는 numpy의 random 함수와 같은 역할을 한다. 해당 함수는 jax에서만 쓰이는 함수로 주요한 장점은 난수로 인한 특정 오류가 발생하였을 때, PRNGkey 안의 key 번호로 동일한 난수를 생성함으로써 오류를 재현할 수 있다는 점이다. 이는 numpy의 random 함수는 불가능하다. 

def MLP(layers, activation=relu):
  ''' Vanilla MLP'''
  def init(rng_key):
      def init_layer(key, d_in, d_out):
          k1, k2 = random.split(key)
          glorot_stddev = 1. / np.sqrt((d_in + d_out) / 2.)
          W = glorot_stddev * random.normal(k1, (d_in, d_out))
          b = np.zeros(d_out)
          return W, b
      key, *keys = random.split(rng_key, len(layers))
      params = list(map(init_layer, keys, layers[:-1], layers[1:]))
      return params
  def apply(params, inputs):
      for W, b in params[:-1]:
          outputs = np.dot(inputs, W) + b   ####################################
          inputs = activation(outputs)    ######################################
      W, b = params[-1]
      outputs = np.dot(inputs, W) + b
      return outputs
  return init, apply


# Define modified MLP
def modified_MLP(layers, activation=relu):
  def xavier_init(key, d_in, d_out):
      glorot_stddev = 1. / np.sqrt((d_in + d_out) / 2.)
      W = glorot_stddev * random.normal(key, (d_in, d_out))
      b = np.zeros(d_out)
      return W, b

  def init(rng_key):
      U1, b1 =  xavier_init(random.PRNGKey(12345), layers[0], layers[1])  #####
      U2, b2 =  xavier_init(random.PRNGKey(54321), layers[0], layers[1])  #####
      def init_layer(key, d_in, d_out):
          k1, k2 = random.split(key)
          W, b = xavier_init(k1, d_in, d_out)
          return W, b
      key, *keys = random.split(rng_key, len(layers))
      params = list(map(init_layer, keys, layers[:-1], layers[1:]))
      return (params, U1, b1, U2, b2)

  def apply(params, inputs):
      params, U1, b1, U2, b2 = params
      U = activation(np.dot(inputs, U1) + b1)   ################################
      V = activation(np.dot(inputs, U2) + b2)   ################################
      for W, b in params[:-1]:
          outputs = activation(np.dot(inputs, W) + b)   ########################
          inputs = np.multiply(outputs, U) + np.multiply(1 - outputs, V)  ######
      W, b = params[-1]
      outputs = np.dot(inputs, W) + b
      return outputs
  return init, apply

 

Physics-Informed DeepONet

 

  PIDoN을 사용하기 위해, 해당 클래스를 정의해 준다. PIDoN의 기초적인 개념이 되는 Branch Network와 Trunk Network를 앞서 정의한 Modified MLP로 정의해준다. optimizer는 최근 (1년 전까지는...) 많이 사용하는 adam를 적용해주는 모습을 볼 수 있다. 

 

# Define Physics-informed DeepONet model
class PI_DeepONet:
    def __init__(self, branch_layers, trunk_layers):
        # Network initialization and evaluation functions
        self.branch_init, self.branch_apply = modified_MLP(branch_layers, activation=np.tanh)
        self.trunk_init, self.trunk_apply = modified_MLP(trunk_layers, activation=np.tanh)

        # Initialize
        branch_params = self.branch_init(rng_key = random.PRNGKey(1234))
        trunk_params = self.trunk_init(rng_key = random.PRNGKey(4321))
        params = (branch_params, trunk_params)

        # Use optimizers to set optimizer initialization and update functions
        self.opt_init, \
        self.opt_update, \
        self.get_params = optimizers.adam(optimizers.exponential_decay(1e-3,
                                                                      decay_steps=2000,
                                                                      decay_rate=0.9))
        self.opt_state = self.opt_init(params)

        # Used to restore the trained model parameters
        _, self.unravel_params = ravel_pytree(params)

        # Logger
        self.itercount = itertools.count()
        self.loss_log = []
        self.loss_ics_log = []
        self.loss_bcs_log = []
        self.loss_res_log = []

 

Operator Net

 

 DeepONet 구조에서 가장 핵심적인 역할을 하는 Operator Net을 정의해준다. 해당 Network에서는 앞서 정의한 두 개의 서브 네트워크인 Branch Network와 Trunk Network로 구성된다. 입력 데이터를 Input으로 사용하는 branch Network와 입력 좌표의 잠재 표현을 추천하기 위해 $[t, x]$를 Input으로 사용하는 Trunk Network를 거쳐 $\mathbf {B}$와 $\mathbf{T}$를 생성한다. 이를 바탕으로 최종 출력을  $\mathbf{B}$와 $\mathbf {T}$의 요소별 곱을 합산하여 다음과 같이 계산한다.

$$Output = \sum_{i=1}^d B_i \cdot T_i$$

 

 이 과정은 Brunch Network와 Trunk Network가 곱해져 솔루션의 연산자를 생성하는 Operator의 주요한 개념이 적용된 부분이라고 볼 수 있다.

 

    # Define DeepONet architecture
    def operator_net(self, params, u, t, x):
        branch_params, trunk_params = params
        y = np.stack([t,x])
        B = self.branch_apply(branch_params, u)
        T = self.trunk_apply(trunk_params, y)
        outputs = np.sum(B * T)
        return   outputs

 

글을 마치며

 

 코드 자체가 길기 때문에, 어디서 끊어야 될지를 고민했는데 아무래도 개인적으로는 물리 기반 인공지능을 가장 많이 접하였기에, 가장 중요한 것을 손실 함수를 어떻게 정의했냐라고 생각한다. 그렇기 때문에 이번 포스팅에서는 전체적인 구조를 다루고, 다음 포스팅에 구체적인 손실 함수 설정을 적기로 하였다. 오랜만에 보는 코드라 가물가물하기도 하고, latex로 포스팅하는 것 자체도 낯설어서 틀린 부분이 있을 수 있는 점 감안해 주시길...

 

 논문과 해당 코드가 있는 Github 링크는 아래와 같다.

 

논문 링크 : https://arxiv.org/abs/2103.10974

Github 링크 : https://github.com/PredictiveIntelligenceLab/Physics-informed-DeepONets