maziarraissi / PINNs

Physics Informed Deep Learning: Data-driven Solutions and Discovery of Nonlinear Partial Differential Equations
https://maziarraissi.github.io/PINNs
MIT License
3.45k stars 1.21k forks source link

some improvement need to run on tensorflow v2 #48

Open M-R-FARHADI opened 1 year ago

M-R-FARHADI commented 1 year ago

class PhysicsInformedNN: def init(self, x, y, t, u, v, layers):

    X = np.concatenate([x, y, t], 1)

    self.lb = X.min(0)
    self.ub = X.max(0)

    self.X = X

    self.x = X[:,0:1]
    self.y = X[:,1:2]
    self.t = X[:,2:3]

    self.u = u
    self.v = v

    self.layers = layers

    # Initialize NN
    self.weights, self.biases = self.initialize_NN(layers)        

    # Initialize parameters
    self.lambda_1 = tf.Variable([0.0], dtype=tf.float32)
    self.lambda_2 = tf.Variable([0.0], dtype=tf.float32)

    # tf placeholders and graph
    self.sess = tf.compat.v1.Session()

    self.x_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.x.shape[1]])
    self.y_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.y.shape[1]])
    self.t_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.t.shape[1]])

    self.u_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.u.shape[1]])
    self.v_tf = tf.compat.v1.placeholder(tf.float32, shape=[None, self.v.shape[1]])

    self.u_pred, self.v_pred, self.p_pred, self.f_u_pred, self.f_v_pred = self.net_NS(self.x_tf, self.y_tf, self.t_tf)

    self.loss = tf.reduce_sum(tf.square(self.u_tf - self.u_pred)) + \
                tf.reduce_sum(tf.square(self.v_tf - self.v_pred)) + \
                tf.reduce_sum(tf.square(self.f_u_pred)) + \
                tf.reduce_sum(tf.square(self.f_v_pred))

    self.adam = tf.compat.v1.train.AdamOptimizer(learning_rate=1e-3)
    self.train_op_adam = self.adam.minimize(self.loss)

    init = tf.compat.v1.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.random.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

def net_NS(self, x, y, t):
    lambda_1 = self.lambda_1
    lambda_2 = self.lambda_2

    psi_and_p = self.neural_net(tf.concat([x,y,t], 1), self.weights, self.biases)
    psi = psi_and_p[:,0:1]
    p = psi_and_p[:,1:2]

    u = tf.gradients(psi, y)[0]
    v = -tf.gradients(psi, x)[0]  

    u_t = tf.gradients(u, t)[0]
    u_x = tf.gradients(u, x)[0]
    u_y = tf.gradients(u, y)[0]
    u_xx = tf.gradients(u_x, x)[0]
    u_yy = tf.gradients(u_y, y)[0]

    v_t = tf.gradients(v, t)[0]
    v_x = tf.gradients(v, x)[0]
    v_y = tf.gradients(v, y)[0]
    v_xx = tf.gradients(v_x, x)[0]
    v_yy = tf.gradients(v_y, y)[0]

    p_x = tf.gradients(p, x)[0]
    p_y = tf.gradients(p, y)[0]

    f_u = u_t + lambda_1*(u*u_x + v*u_y) + p_x - lambda_2*(u_xx + u_yy) 
    f_v = v_t + lambda_1*(u*v_x + v*v_y) + p_y - lambda_2*(v_xx + v_yy)

    return u, v, p, f_u, f_v

def callback(self, loss, lambda_1, lambda_2):
    print('Loss: %.3e, l1: %.3f, l2: %.5f' % (loss, lambda_1, lambda_2))

def train(self, nIter): 

    tf_dict = {self.x_tf: self.x, self.y_tf: self.y, self.t_tf: self.t,
               self.u_tf: self.u, self.v_tf: self.v}

    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)
            lambda_1_value = self.sess.run(self.lambda_1)
            lambda_2_value = self.sess.run(self.lambda_2)
            print('It: %d, Loss: %.3e, l1: %.3f, l2: %.5f, Time: %.2f' % 
                  (it, loss_value, lambda_1_value, lambda_2_value, elapsed))
            start_time = time.time()

def predict(self, x_star, y_star, t_star):

    tf_dict = {self.x_tf: x_star, self.y_tf: y_star, self.t_tf: t_star}

    u_star = self.sess.run(self.u_pred, tf_dict)
    v_star = self.sess.run(self.v_pred, tf_dict)
    p_star = self.sess.run(self.p_pred, tf_dict)

    return u_star, v_star, p_star
juliosdutra commented 1 year ago

Thank you for this update!

tirtho109 commented 9 months ago

Thank you for your code!