0914PINN.py

"""
@author: Maziar Raissi
"""

import sys
sys.path.insert(0, '../../Utilities/')

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import scipy.io
from scipy.interpolate import griddata
from pyDOE import lhs
from plotting import newfig, savefig
from mpl_toolkits.mplot3d import Axes3D
import time
import matplotlib.gridspec as gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable


np.random.seed(1234)
tf.set_random_seed(1234)


class PhysicsInformedNN:
    # Initialize the class
    def __init__(self, X_ac, V_ac, X_ex,V_ex, X_wall, X_space, X_louver, layers,rho,mu,g):
        # 引数として初期条件、境界条件を受けて、層を受けている
        self.rho = rho
        self.mu = mu
        self.g=g
               
        self.x_ac = X_ac[:,0:1]
        self.y_ac = X_ac[:,1:2]
        self.z_ac = X_ac[:,2:3]

        self.x_ex = X_ex[:,0:1]
        self.y_ex = X_ex[:,1:2]
        self.z_ex = X_ex[:,2:3]

        self.x_wall = X_wall[:,0:1]
        self.y_wall = X_wall[:,1:2]
        self.z_wall = X_wall[:,2:3]

        self.x_space = X_space[:,0:1]
        self.y_space = X_space[:,1:2]
        self.z_space = X_space[:,2:3]

        self.x_louver = X_louver[:,0:1]
        self.y_louver = X_louver[:,1:2]
        self.z_louver = X_louver[:,2:3]

        self.u_ac = V_ac[0]
        self.v_ac = V_ac[1]
        self.w_ac = V_ac[2]

        self.u_ex = V_ex[0]
        self.v_ex = V_ex[1]
        self.w_ex = V_ex[2]
        
        # Initialize NNs
        self.layers = layers
        self.weights, self.biases = self.initialize_NN(layers)
        
        # tf Placeholders        
        self.x_ac_tf = tf.placeholder(tf.float32, shape=[None, self.x_ac.shape[1]])
        self.y_ac_tf = tf.placeholder(tf.float32, shape=[None, self.y_ac.shape[1]])
        self.z_ac_tf = tf.placeholder(tf.float32, shape=[None, self.z_ac.shape[1]])

        self.x_ex_tf = tf.placeholder(tf.float32, shape=[None, self.x_ex.shape[1]])
        self.y_ex_tf = tf.placeholder(tf.float32, shape=[None, self.y_ex.shape[1]])
        self.z_ex_tf = tf.placeholder(tf.float32, shape=[None, self.z_ex.shape[1]])

        self.x_wall_tf = tf.placeholder(tf.float32, shape=[None, self.x_wall.shape[1]])
        self.y_wall_tf = tf.placeholder(tf.float32, shape=[None, self.y_wall.shape[1]])
        self.z_wall_tf = tf.placeholder(tf.float32, shape=[None, self.z_wall.shape[1]])

        self.x_space_tf = tf.placeholder(tf.float32, shape=[None, self.x_space.shape[1]])
        self.y_space_tf = tf.placeholder(tf.float32, shape=[None, self.y_space.shape[1]])
        self.z_space_tf = tf.placeholder(tf.float32, shape=[None, self.z_space.shape[1]])

        self.x_louver_tf = tf.placeholder(tf.float32, shape=[None, self.x_louver.shape[1]])
        self.y_louver_tf = tf.placeholder(tf.float32, shape=[None, self.y_louver.shape[1]])
        self.z_louver_tf = tf.placeholder(tf.float32, shape=[None, self.z_louver.shape[1]])
        
        self.u_ac_tf = tf.placeholder(tf.float32, shape=[None, self.u_ac.shape[1]])
        self.v_ac_tf = tf.placeholder(tf.float32, shape=[None, self.v_ac.shape[1]])
        self.z_ac_tf = tf.placeholder(tf.float32, shape=[None, self.z_ac.shape[1]])

        self.u_ex_tf = tf.placeholder(tf.float32, shape=[None, self.u_ex.shape[1]])
        self.v_ex_tf = tf.placeholder(tf.float32, shape=[None, self.v_ex.shape[1]])
        self.z_ex_tf = tf.placeholder(tf.float32, shape=[None, self.z_ac.shape[1]])
        

        # tf Graphs
        self.u_wall_pred, self.v_wall_pred, self.w_wall_pred =self.net_uvw(self.x_wall_tf, self.y_wall_tf, self.z_wall_tf)
        self.u_louver_pred, self.v_louver_pred, self.w_louver_pred =self.net_uvw(self.x_louver_tf, self.y_louver_tf, self.z_louver_tf)
        self.u_ac_pred, self.v_ac_pred, self.w_ac_pred =self.net_uvw(self.x_ac_tf, self.y_ac_tf, self.z_ac_tf)
        self.u_ex_pred, self.v_ex_pred, self.w_ex_pred =self.net_uvw(self.x_ex_tf, self.y_ex_tf, self.z_ex_tf)
        self.div_V = self.net_incomp(self.x_space_tf, self.y_space_tf, self.z_space_tf)
        self.ns_x, self.ns_y, self.ns_z = self.net_NS(self.x_space_tf, self.y_space_tf, self.z_space_tf)
        
        
        
        # Loss
        self.loss = tf.reduce_mean(tf.square(self.u_wall_pred)) + \
                    tf.reduce_mean(tf.square(self.v_wall_pred)) + \
                    tf.reduce_mean(tf.square(self.w_wall_pred)) + \
                    tf.reduce_mean(tf.square(self.u_louver_pred)) + \
                    tf.reduce_mean(tf.square(self.v_louver_pred)) + \
                    tf.reduce_mean(tf.square(self.w_louver_pred)) + \
                    tf.reduce_mean(tf.square(self.u_ac - self.u_ac_pred)) + \
                    tf.reduce_mean(tf.square(self.v_ac - self.v_ac_pred)) + \
                    tf.reduce_mean(tf.square(self.w_ac - self.w_ac_pred)) + \
                    tf.reduce_mean(tf.square(self.u_ex - self.u_ex_pred)) + \
                    tf.reduce_mean(tf.square(self.v_ex - self.v_ex_pred)) + \
                    tf.reduce_mean(tf.square(self.w_ex - self.w_ex_pred)) + \
                    tf.reduce_mean(tf.square(self.div_V)) + \
                    tf.reduce_mean(tf.square(self.ns_x)) + \
                    tf.reduce_mean(tf.square(self.ns_x)) + \
                    tf.reduce_mean(tf.square(self.ns_z))
        
        # Optimizers
        self.optimizer = tf.contrib.opt.ScipyOptimizerInterface(self.loss, 
                                                                method = 'L-BFGS-B', 
                                                                options = {'maxiter': 50000,
                                                                           'maxfun': 50000,
                                                                           'maxcor': 50,
                                                                           'maxls': 50,
                                                                           'ftol' : 1.0 * np.finfo(float).eps})
    
        self.optimizer_Adam = tf.train.AdamOptimizer() #自分で作った関数をもとにoptimizerを定義
        self.train_op_Adam = self.optimizer_Adam.minimize(self.loss)
                
        # tf session
        self.sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
                                                     log_device_placement=True))
        
        init = tf.global_variables_initializer()
        self.sess.run(init)
              
    def initialize_NN(self, layers):        
        weights = []
        biases = []
        num_layers = len(layers) 
        for l in range(0,num_layers-1):
            W = self.xavier_init(size=[layers[l], layers[l+1]])
            b = tf.Variable(tf.zeros([1,layers[l+1]], dtype=tf.float32), dtype=tf.float32)
            weights.append(W)
            biases.append(b)        
        return weights, biases
        
    def xavier_init(self, size):
        in_dim = size[0]
        out_dim = size[1]        
        xavier_stddev = np.sqrt(2/(in_dim + out_dim))
        return tf.Variable(tf.truncated_normal([in_dim, out_dim], stddev=xavier_stddev), dtype=tf.float32)
    
    # 入力からニューラルネットの出力を計算
    def neural_net(self, X, weights, biases):
        num_layers = len(weights) + 1
        
        H = 2.0*(X - self.lb)/(self.ub - self.lb) - 1.0
        for l in range(0,num_layers-2):
            W = weights[l]
            b = biases[l]
            H = tf.tanh(tf.add(tf.matmul(H, W), b))
        W = weights[-1]
        b = biases[-1]
        Y = tf.add(tf.matmul(H, W), b)
        return Y

    # ニューラルネットの出力から、uvwを計算
    def net_uvw(self, x, y,z):
        # 勾配を得る
        X = tf.concat([x,y,z],1)
        # neural_netを用いて予測値を得る（u,v）
        Y = self.neural_net(X, self.weights, self.biases)
        u = Y[:,0:1]
        v = Y[:,1:2]
        w = Y[:,2:3]

        return u, v, w

    # ニューラルネットワークの出力から、uvwの一階微分を計算
    def net_incomp(self, x, y,z):
        # 勾配を得る
        with tf.GradientTape() as tape1:
            X = tf.concat([x,y,z],1)
            # neural_netを用いて予測値を得る（u,v）
            Y = self.neural_net(X, self.weights, self.biases)
            u = Y[:,0:1]
            v = Y[:,1:2]
            w = Y[:,2:3]
        u_x=tape1.gradient(u,x)
        v_y=tape1.gradient(v,y)
        w_z=tape1.gradient(w,z)

        return u_x + v_y + w_z

    def net_NS(self, x, y, z):
        # 勾配を得る
        with tf.GradientTape() as tape1:
            with tf.GradientTape() as tape2:
                X = tf.concat([x,y,z],1)
                # neural_netを用いて予測値を得る（u,v）
                Y = self.neural_net(X, self.weights, self.biases)
                u = Y[:,0:1]
                v = Y[:,1:2]
                w = Y[:,2:3]
                p = Y[:,3:4]
            u_x=tape2.gradient(u,x)
            u_y=tape2.gradient(u,y)
            u_z=tape2.gradient(u,z)
            v_x=tape2.gradient(v,x)
            v_y=tape2.gradient(v,y)
            v_z=tape2.gradient(v,z)
            w_x=tape2.gradient(w,x)
            w_y=tape2.gradient(w,y)
            w_z=tape2.gradient(w,z)
            p_x=tape2.gradient(p,x)
            p_y=tape2.gradient(p,y)
            p_z=tape2.gradient(p,z)
        u_xx = tape1.gradient(u_x,x)
        v_yy = tape1.gradient(v_y,y)
        w_zz = tape1.gradient(w_z,z)

        # Navier-Stokes式を計算
        viscos_term = self.mu*(u_xx + v_yy + w_zz)
        ns_x = u * u_x +v * u_y + z * u_z + (p_x - viscos_term) / self.rho - self.g[0]
        ns_y = u * v_x +v * v_y + z * v_z + (p_y - viscos_term) / self.rho - self.g[1]
        ns_z = u * w_x +v * w_y + z * w_z + (p_z - viscos_term) / self.rho - self.g[2]

        return ns_x, ns_y, ns_z
    
    def callback(self, loss):
        print('Loss:', loss)
        
    def train(self, nIter):
        
        tf_dict = {self.x0_tf: self.x0, self.t0_tf: self.t0,
                   self.u0_tf: self.u0, self.v0_tf: self.v0,
                   self.x_lb_tf: self.x_lb, self.t_lb_tf: self.t_lb,
                   self.x_ub_tf: self.x_ub, self.t_ub_tf: self.t_ub,
                   self.x_f_tf: self.x_f, self.t_f_tf: self.t_f}
        
        start_time = time.time()
        for it in range(nIter):
            self.sess.run(self.train_op_Adam, tf_dict)
            
            # Print
            if it % 10 == 0:
                elapsed = time.time() - start_time
                loss_value = self.sess.run(self.loss, tf_dict)
                print('It: %d, Loss: %.3e, Time: %.2f' % 
                      (it, loss_value, elapsed))
                start_time = time.time()
                                                                                                                          
        self.optimizer.minimize(self.sess, 
                                feed_dict = tf_dict,         
                                fetches = [self.loss], 
                                loss_callback = self.callback)        
                                    
    
    def predict(self, X_star):
        
        tf_dict = {self.x0_tf: X_star[:,0:1], self.t0_tf: X_star[:,1:2]}
        
        u_star = self.sess.run(self.u0_pred, tf_dict)  
        v_star = self.sess.run(self.v0_pred, tf_dict)  
        
        
        tf_dict = {self.x_f_tf: X_star[:,0:1], self.t_f_tf: X_star[:,1:2]}
        
        f_u_star = self.sess.run(self.f_u_pred, tf_dict)
        f_v_star = self.sess.run(self.f_v_pred, tf_dict)
               
        return u_star, v_star, f_u_star, f_v_star
    
if __name__ == "__main__": 
     
    noise = 0.0        
    
    # Doman bounds
    V_ac=np.array([1.,1.,1.])
    V_ex=np.array([1.,1.,1.])
    g=np.array([1.,1.,1.])
    rho=1
    mu=1

    layers = [3, 100, 100, 100, 100, 4]
        
    # データの読み込み（要修正）
    data = scipy.io.loadmat('../Data/NLS.mat')
    
    t = data['tt'].flatten()[:,None]
    x = data['x'].flatten()[:,None]
    Exact = data['uu']
    Exact_u = np.real(Exact)
    Exact_v = np.imag(Exact)
    Exact_h = np.sqrt(Exact_u**2 + Exact_v**2)
    
    X, T = np.meshgrid(x,t)
    
    X_star = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))
    u_star = Exact_u.T.flatten()[:,None]
    v_star = Exact_v.T.flatten()[:,None]
    h_star = Exact_h.T.flatten()[:,None]
    
    ###########################
    
    idx_x = np.random.choice(x.shape[0], N0, replace=False)
    x0 = x[idx_x,:]
    u0 = Exact_u[idx_x,0:1]
    v0 = Exact_v[idx_x,0:1]
    
    idx_t = np.random.choice(t.shape[0], N_b, replace=False)
    tb = t[idx_t,:]
    
    X_f = lb + (ub-lb)*lhs(2, N_f)
            
    model = PhysicsInformedNN(X_ac, V_ac, X_ex,V_ex, X_wall, X_space, X_louver, layers,rho,mu,g)
             
    start_time = time.time()                
    model.train(50000)
    elapsed = time.time() - start_time                
    print('Training time: %.4f' % (elapsed))
    
        
    u_pred, v_pred, f_u_pred, f_v_pred = model.predict(X_star)
    h_pred = np.sqrt(u_pred**2 + v_pred**2)
            
    error_u = np.linalg.norm(u_star-u_pred,2)/np.linalg.norm(u_star,2)
    error_v = np.linalg.norm(v_star-v_pred,2)/np.linalg.norm(v_star,2)
    error_h = np.linalg.norm(h_star-h_pred,2)/np.linalg.norm(h_star,2)
    print('Error u: %e' % (error_u))
    print('Error v: %e' % (error_v))
    print('Error h: %e' % (error_h))

    
    U_pred = griddata(X_star, u_pred.flatten(), (X, T), method='cubic')
    V_pred = griddata(X_star, v_pred.flatten(), (X, T), method='cubic')
    H_pred = griddata(X_star, h_pred.flatten(), (X, T), method='cubic')

    FU_pred = griddata(X_star, f_u_pred.flatten(), (X, T), method='cubic')
    FV_pred = griddata(X_star, f_v_pred.flatten(), (X, T), method='cubic')     
    

    
    ######################################################################
    ############################# Plotting ###############################
    ######################################################################    
    
    X0 = np.concatenate((x0, 0*x0), 1) # (x0, 0)
    X_lb = np.concatenate((0*tb + lb[0], tb), 1) # (lb[0], tb)
    X_ub = np.concatenate((0*tb + ub[0], tb), 1) # (ub[0], tb)
    X_u_train = np.vstack([X0, X_lb, X_ub])

    fig, ax = newfig(1.0, 0.9)
    ax.axis('off')
    
    ####### Row 0: h(t,x) ##################    
    gs0 = gridspec.GridSpec(1, 2)
    gs0.update(top=1-0.06, bottom=1-1/3, left=0.15, right=0.85, wspace=0)
    ax = plt.subplot(gs0[:, :])
    
    h = ax.imshow(H_pred.T, interpolation='nearest', cmap='YlGnBu', 
                  extent=[lb[1], ub[1], lb[0], ub[0]], 
                  origin='lower', aspect='auto')
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    fig.colorbar(h, cax=cax)
    
    ax.plot(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (X_u_train.shape[0]), markersize = 4, clip_on = False)
    
    line = np.linspace(x.min(), x.max(), 2)[:,None]
    ax.plot(t[75]*np.ones((2,1)), line, 'k--', linewidth = 1)
    ax.plot(t[100]*np.ones((2,1)), line, 'k--', linewidth = 1)
    ax.plot(t[125]*np.ones((2,1)), line, 'k--', linewidth = 1)    
    
    ax.set_xlabel('$t$')
    ax.set_ylabel('$x$')
    leg = ax.legend(frameon=False, loc = 'best')
#    plt.setp(leg.get_texts(), color='w')
    ax.set_title('$|h(t,x)|$', fontsize = 10)
    
    ####### Row 1: h(t,x) slices ##################    
    gs1 = gridspec.GridSpec(1, 3)
    gs1.update(top=1-1/3, bottom=0, left=0.1, right=0.9, wspace=0.5)
    
    ax = plt.subplot(gs1[0, 0])
    ax.plot(x,Exact_h[:,75], 'b-', linewidth = 2, label = 'Exact')       
    ax.plot(x,H_pred[75,:], 'r--', linewidth = 2, label = 'Prediction')
    ax.set_xlabel('$x$')
    ax.set_ylabel('$|h(t,x)|$')    
    ax.set_title('$t = %.2f$' % (t[75]), fontsize = 10)
    ax.axis('square')
    ax.set_xlim([-5.1,5.1])
    ax.set_ylim([-0.1,5.1])
    
    ax = plt.subplot(gs1[0, 1])
    ax.plot(x,Exact_h[:,100], 'b-', linewidth = 2, label = 'Exact')       
    ax.plot(x,H_pred[100,:], 'r--', linewidth = 2, label = 'Prediction')
    ax.set_xlabel('$x$')
    ax.set_ylabel('$|h(t,x)|$')
    ax.axis('square')
    ax.set_xlim([-5.1,5.1])
    ax.set_ylim([-0.1,5.1])
    ax.set_title('$t = %.2f$' % (t[100]), fontsize = 10)
    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.8), ncol=5, frameon=False)
    
    ax = plt.subplot(gs1[0, 2])
    ax.plot(x,Exact_h[:,125], 'b-', linewidth = 2, label = 'Exact')       
    ax.plot(x,H_pred[125,:], 'r--', linewidth = 2, label = 'Prediction')
    ax.set_xlabel('$x$')
    ax.set_ylabel('$|h(t,x)|$')
    ax.axis('square')
    ax.set_xlim([-5.1,5.1])
    ax.set_ylim([-0.1,5.1])    
    ax.set_title('$t = %.2f$' % (t[125]), fontsize = 10)
    
    # savefig('./figures/NLS')