Main.py

#!/usr/bin/env python3
import numpy as np
import timeit
import os
import sys
import tensorflow as tf
from tensorflow import keras
from tensorflow.python.ops.numpy_ops import np_config
from matplotlib import pyplot as plt
import matplotlib
import scipy

sys.path.append('/Users/nwade/PycharmProjects/StochasticPINN')
np_config.enable_numpy_behavior()

class SC_Heat_2D():
    def __init__(self, input_dict):
        self.Network = input_dict['Network']
        self.dataonly = input_dict['dataonly']
        self.numcp_x = input_dict['Num_cp_x']
        self.numcp_y = input_dict['Num_cp_y']
        self.numcp = self.numcp_y * self.numcp_x
        self.nelem = Input_dict['nelem']
        self.NData = input_dict['Num_Data']
        self.numgp = input_dict['ngp']
        self.Nbatch = input_dict['batchSize']
        self.num_bc_points = input_dict['num_bc_points']
        self.totalpoints = self.Nbatch*self.numcp*self.numgp**2
        self.coords_dom = input_dict['coordinate_domain']
        self.Element_size = Input_dict['Element_size']
        self.mesh = np.mgrid[self.coords_dom[0][0]:self.coords_dom[0][1] + self.Element_size:self.Element_size,
                    self.coords_dom[1][0]:self.coords_dom[1][1] + self.Element_size:self.Element_size]

        self.irsize = input_dict['IR_size']
        self.gauss_legendre = np.polynomial.legendre.leggauss(self.numgp)

        if self.numgp == 1:
            self.gp_weights = tf.transpose(tf.tile([tf.cast(self.gauss_legendre[1], dtype=tf.float32)],
                                                   [1, self.Nbatch * self.numcp]))
        else:
            self.gp_weights = tf.transpose(tf.tile([tf.cast(tf.repeat(self.gauss_legendre[1], 2, axis=0),
                                           dtype=tf.float32)], [1, self.Nbatch * self.numcp]))

        self.computeIntegrationRegionVectors()
        self.Npt = Input_dict['Npt']  # total number of collocation points
        self.Num_epochs = Input_dict['Num_epochs']  # number of epochs to trains
        self.yhat = Input_dict['Y_hat']
        self.ndim = Input_dict['ndim']
        self.data_name = Input_dict['model_data_name']
        self.model_name = Input_dict['model_name']
        self.loadold = Input_dict['loadold']
        self.model_load = Input_dict['2DHeat_model_min_error.h5']
        self.nDofs = (self.nelem[0]+1)*(self.nelem[1]+1)
        self.niter = Input_dict['niter']
        self.layer_u = Input_dict['layer_u'][1:]
        self.input_u = layer_u[0]
        self.error_write = Input_dict['error_write']
        self.bc_weight = Input_dict['bc_weight']
        self.Weighted = Input_dict['inputs_weighted']
        self.inp_update = Input_dict['inputs_updated']
        self.NWbins = Input_dict['NWbins']
        self.Path = Input_dict['Path']
        self.Num_solves = Input_dict['Num_solves']
        self.Max_weight = Input_dict['Max_weight']
        self.Learning_rate = Input_dict['Learning_rate']
        self.Learning_rate_decay = Input_dict['Learning_rate_decay']
        self.Learning_rate_step = Input_dict['Learning_rate_step']
        self.Start = Input_dict['Start']
        self.Data_load = Input_dict['Data_load']
        self.weight_method = Input_dict['weight_method']
        self.Ramping_percent = Input_dict['Ramping_percent']
        self.alpha = Input_dict['alpha']
        self.mode = Input_dict['mode']

        data_loc_x = np.reshape(np.linspace(self.coords_dom[0][0], self.coords_dom[0][1], self.nelem[0]+1, dtype=np.float32),[self.nelem[0]+1,1])
        data_loc_y = np.reshape(np.linspace(self.coords_dom[1][0], self.coords_dom[1][1], self.nelem[1]+1, dtype=np.float32),[self.nelem[1]+1,1])
        self.FEA_Mesh = np.zeros([(self.nelem[0] + 1) * (self.nelem[1] + 1), 2], dtype=np.float32)
        self.FEA_Mesh_size = (self.nelem[0] + 1) * (self.nelem[1] + 1)
        index = 0
        for nx in range(self.nelem[0] + 1):
            for ny in range(self.nelem[1] + 1):
                self.FEA_Mesh[index, 0] = data_loc_x[nx, 0]
                self.FEA_Mesh[index, 1] = data_loc_y[ny, 0]
                index = index + 1


        # Fixed Model Hyper Parameters
        self.starter_learning_rate = self.Learning_rate
        self.lr_schedule = keras.optimizers.schedules.ExponentialDecay(
            initial_learning_rate=self.starter_learning_rate,
            decay_steps=self.Learning_rate_step, decay_rate=self.Learning_rate_decay, staircase=True)
        self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.lr_schedule)
        self.lossFunc = tf.keras.losses.MeanSquaredError()
        self.bc = 700
        self.buildModel()
        (self.xdata, self.R_data) = self.establish_training_dataset()
        self.yhat = self.R_data[:,0:ndim]

        xr = tf.cast(tf.reshape(tf.linspace(self.coords_dom[0][0], self.coords_dom[0][1], self.num_bc_points), [self.num_bc_points, 1]),
                     dtype=tf.float32)
        yr = tf.cast(tf.reshape(tf.linspace(self.coords_dom[1][0], self.coords_dom[1][1], self.num_bc_points), [self.num_bc_points, 1]),
                     dtype=tf.float32)
        xmin = tf.cast(tf.ones([self.num_bc_points, 1]) * self.coords_dom[0][0], dtype=tf.float32)
        ymin = tf.cast(tf.ones([self.num_bc_points, 1]) * self.coords_dom[1][0], dtype=tf.float32)
        xmax = tf.cast(tf.ones([self.num_bc_points, 1]) * self.coords_dom[0][1], dtype=tf.float32)
        ymax = tf.cast(tf.ones([self.num_bc_points, 1]) * self.coords_dom[1][1], dtype=tf.float32)
        self.xi_0 = tf.tile(tf.concat([xr, xmin, xr, xmax], axis=0), [self.Nbatch, 1])
        self.yi_0 = tf.tile(tf.concat([ymin, yr, ymax, yr], axis=0), [self.Nbatch, 1])

    def computeIntegrationRegionVectors(self):
        p_loc_x = np.linspace(self.coords_dom[0][0], self.coords_dom[0][1], nelem[0]+1)[1:-2] + 0.02 / 2
        p_loc_y = np.linspace(self.coords_dom[1][0], self.coords_dom[1][1], nelem[1]+1)[1:-2] + 0.02 / 2

        self.c_points_x = p_loc_x[np.round(np.linspace(0, p_loc_x.size - 1, self.numcp_x)).astype(int)]
        self.c_points_y = p_loc_y[np.round(np.linspace(0, p_loc_y.size - 1, self.numcp_x)).astype(int)]

        if self.numgp == 2:
            self.delta_x = Element_size / 2 / self.gauss_legendre[0][1]
            self.delta_y = Element_size / 2 / self.gauss_legendre[0][1]
        if self.numgp == 1:
            self.delta_x = ((self.c_points_x[2] - self.c_points_x[1]) / 2)
            self.delta_y = ((self.c_points_x[2] - self.c_points_x[1]) / 2)

        xa = self.c_points_x[0] - self.delta_x
        xb = self.c_points_x[0] + self.delta_x
        ya = self.c_points_y[0] - self.delta_y
        yb = self.c_points_y[0] + self.delta_y
        cpx = xa + self.delta_x
        cpy = ya + self.delta_y

        self.ir_loc = np.zeros([2, self.numcp_x, self.numcp_y])
        for n in range(self.numcp_x):
            for nn in range(self.numcp_x):
                self.ir_loc[0,n,nn] = self.c_points_x[n]
                self.ir_loc[1,n,nn] = self.c_points_x[nn]

        shift_x = np.tile(np.tile(np.repeat(self.gauss_legendre[0], self.numgp), [self.numcp, 1]), [self.Nbatch, 1]) * self.delta_x
        shift_y = np.tile(np.tile(np.tile(self.gauss_legendre[0], [1, self.numgp]), [self.numcp, 1]), [self.Nbatch, 1]) * self.delta_y
        x_loc_new = np.tile(np.tile(np.reshape(np.repeat(self.c_points_x, self.numgp**2), [self.numcp_x, self.numgp**2]), [self.numcp_y, 1]), [self.Nbatch, 1])
        y_loc_new = np.tile(np.reshape(np.repeat(np.repeat(self.c_points_y, self.numgp**2), self.numcp_x), [self.numcp_y * self.numcp_x, self.numgp**2]), [self.Nbatch, 1])

        self.x_loc = np.reshape(x_loc_new + shift_x, [self.Nbatch * self.numgp**2 * self.numcp, 1])
        self.y_loc = np.reshape(y_loc_new + shift_y, [self.Nbatch * self.numgp**2 * self.numcp, 1])

        self.x_loc_batch = np.reshape(np.repeat(np.tile(self.c_points_x, self.numcp_y), self.Nbatch),
                                      [self.numcp * self.Nbatch, 1])
        self.y_loc_batch = np.reshape(np.repeat(np.repeat(self.c_points_y, self.numcp_x), self.Nbatch),
                                      [self.numcp * self.Nbatch, 1])

        # Build the test function vectors
        gauss_loc_2d = []
        gauss_weights_2d = []
        for n in range(self.numgp):
            for nn in range(self.numgp):
                gauss_loc_2d.append([self.gauss_legendre[0][n], self.gauss_legendre[0][nn]])
                gauss_weights_2d.append([self.gauss_legendre[1][n], self.gauss_legendre[1][nn]])

        w = np.zeros(self.numgp ** 2)
        dwdx = np.zeros(self.numgp ** 2)
        dwdy = np.zeros(self.numgp ** 2)
        for n in range(self.numgp ** 2):
            x = cpx + gauss_loc_2d[n][0] * self.delta_x
            y = cpy + gauss_loc_2d[n][1] * self.delta_y

            if self.numgp == 2:
                w[n] = (x - xa)**2*(x - xb)**2*(y - ya)**2*(y - yb)**2/((-xa/2 + xb/2)**4*(-ya/2 + yb/2)**4)
                dwdx[n] = (x - xa)**2*(2*x - 2*xb)*(y - ya)**2*(y - yb)**2/((-xa/2 + xb/2)**4*(-ya/2 + yb/2)**4) + \
                          (x - xb)**2*(2*x - 2*xa)*(y - ya)**2*(y - yb)**2/((-xa/2 + xb/2)**4*(-ya/2 + yb/2)**4)
                dwdy[n] = (x - xa)**2*(x - xb)**2*(y - ya)**2*(2*y - 2*yb)/((-xa/2 + xb/2)**4*(-ya/2 + yb/2)**4) + \
                          (x - xa)**2*(x - xb)**2*(y - yb)**2*(2*y - 2*ya)/((-xa/2 + xb/2)**4*(-ya/2 + yb/2)**4)

        self.w = tf.reshape(tf.cast(np.tile(np.tile(w, self.numcp), self.Nbatch), dtype=tf.float32),
                            [self.totalpoints, 1])
        self.dvdx = tf.reshape(tf.cast(np.tile(np.tile(dwdx, self.numcp), self.Nbatch), dtype=tf.float32),
                               [self.totalpoints, 1])
        self.dvdy = tf.reshape(tf.cast(np.tile(np.tile(dwdy, self.numcp), self.Nbatch), dtype=tf.float32),
                               [self.totalpoints, 1])

        return

    def buildModel(self):
        if self.Network == 'Deep':

            deep_branch_input = tf.keras.Input(shape=(self.ndim,))
            deep_branch = deep_branch_input
            for j in range(len(self.layer_u)):
                deep_branch = tf.keras.layers.Dense(self.layer_u[j], activation='tanh',
                                                    kernel_initializer='glorot_normal')(deep_branch)

            # Deep Trunk
            deep_trunk_input = tf.keras.Input(shape=(2,))
            deep_trunk = deep_trunk_input
            for j in range(len(self.layer_u)):
                deep_trunk = tf.keras.layers.Dense(self.layer_u[j], activation='tanh',
                                                   kernel_initializer='glorot_normal')(deep_trunk)

            # Combine Branch and Trunk Outputs
            combined = tf.keras.layers.Multiply()([deep_branch, deep_trunk])

            # Output Layer
            output = tf.keras.layers.Dense(1, activation='linear')(combined)
            final_output = tf.keras.layers.Dense(1,activation='linear', kernel_initializer='zeros', bias_initializer=tf.keras.initializers.Constant(self.bc))(output)

            # Final Model
            model = tf.keras.Model(inputs=[deep_branch_input, deep_trunk_input], outputs=final_output)

            # Compile Model
            self.optimizer_instance = self.optimizer  # Store optimizer explicitly
            model.compile(optimizer=self.optimizer_instance, loss=self.lossFunc)
        else:
            model = tf.keras.Sequential()

            # Input layer
            model.add(tf.keras.Input(shape=(self.input_u,)))

            # Hidden layers
            for j in range(len(self.layer_u)):
                model.add(tf.keras.layers.Dense(self.layer_u[j], activation='tanh', kernel_initializer='glorot_normal'))

            # Output layers
            model.add(tf.keras.layers.Dense(1, activation='linear', kernel_initializer='glorot_normal', use_bias=False))
            model.add(tf.keras.layers.Dense(1, activation='linear', trainable=False,
                                            kernel_initializer=tf.keras.initializers.Constant(1),
                                            bias_initializer=tf.keras.initializers.Constant(self.bc)))

            # Compile the model
            self.optimizer_instance = self.optimizer  # Store optimizer explicitly
            model.compile(optimizer=self.optimizer_instance, loss=self.lossFunc)

        self.model = model


    def establish_training_dataset(self):
        combined = np.genfromtxt(self.Data_load, delimiter=',')[self.Start - 1:self.Start - 1 + self.Npt, :]
        np.savetxt(self.data_name, combined, delimiter=',')

        # read from file
        def parse_fnc(line):
            string_vals = tf.strings.split([line], sep=',').values
            return tf.strings.to_number(string_vals, tf.float32)

        xd = tf.data.TextLineDataset([self.data_name]).repeat()
        xd = xd.map(map_func=parse_fnc)
        xd = xd.batch(self.Nbatch)
        return xd, combined

    def load_training_dataset(self):
        def parse_fnc(line):
            string_vals = tf.strings.split([line], sep=',').values
            return tf.strings.to_number(string_vals, tf.float32)

        xd = tf.data.TextLineDataset([self.data_name]).repeat()
        xd = xd.map(map_func=parse_fnc)
        xd = xd.batch(self.Nbatch)
        return xd

    def load_weights_dataset(self, mode):
        def parse_fnc(line):
            string_vals = tf.strings.split([line], sep=',').values
            return tf.strings.to_number(string_vals, tf.float32)

        if mode==1:
            xd = tf.data.TextLineDataset(['Npt_weights.txt']).repeat()
            xd = xd.map(map_func=parse_fnc)
            xd = xd.batch(self.Nbatch)
        else:
            xd = tf.data.TextLineDataset(['Npt_bc_weights.txt']).repeat()
            xd = xd.map(map_func=parse_fnc)
            xd = xd.batch(self.Nbatch)
        return xd
    def buildDataVector(self, xdata):
        data = iter(xdata.enumerate(tf.cast(1, tf.int64)))
        (_, solution) = data.next()
        if self.NData > self.Nbatch:
            data_s = tf.reshape(solution[:, self.ndim:self.ndim + self.nDofs], [self.Nbatch * self.nDofs, 1])
            s_input = tf.reshape(tf.tile(solution[:, 0:self.ndim], [1, self.nDofs]),
                                 [self.Nbatch * self.nDofs, self.ndim])
            count = self.NData - self.Nbatch

            while count > self.Nbatch:
                (_, solution) = data.next()
                data_s_2 = tf.reshape(solution[:, self.ndim:self.ndim + self.nDofs], [self.Nbatch * self.nDofs, 1])
                s_input_2 = tf.reshape(tf.tile(solution[:, 0:self.ndim], [1, self.nDofs]),
                                       [self.Nbatch * self.nDofs, self.ndim])
                data_s = tf.concat([data_s, data_s_2], 0)
                s_input = tf.concat([s_input, s_input_2], 0)
                count = count - self.Nbatch

            data_s_2 = tf.reshape(solution[0:count, self.ndim:self.ndim + self.nDofs], [(count) * self.nDofs, 1])
            s_input_2 = tf.reshape(tf.tile(solution[0:count, 0:self.ndim], [1, self.nDofs]), [(count) * self.nDofs, self.ndim])
            self.data_s = tf.concat([data_s, data_s_2], 0)
            s_input = tf.concat([s_input, s_input_2], 0)
        else:
            self.data_s = tf.reshape(solution[0:self.NData, self.ndim:self.ndim + self.nDofs],
                                     [self.NData * self.nDofs, 1])
            s_input = tf.reshape(tf.tile(solution[0:self.NData, 0:self.ndim], [1, self.nDofs]),
                                 [self.nDofs * self.NData, self.ndim])
        x_data_loc = tf.cast(
            tf.tile(tf.reshape(self.FEA_Mesh[:, 0], [(self.nelem[0] + 1) * (self.nelem[1] + 1), 1]), [self.NData, 1]),
            dtype=tf.float32)
        y_data_loc = tf.cast(
            tf.tile(tf.reshape(self.FEA_Mesh[:, 1], [(self.nelem[0] + 1) * (self.nelem[1] + 1), 1]), [self.NData, 1]),
            dtype=tf.float32)
        self.data_input = tf.concat([tf.concat([s_input, x_data_loc], 1), y_data_loc], 1)
        return

    @tf.function
    def predict_u(self, input):
        if self.Network == 'Deep':
            u = self.model([input[:, 0:self.ndim], input[:, self.ndim:self.ndim+2]])
        else:
            u = self.model(input)
        return u

    @tf.function
    def net_predict_r(self, xi):
        with tf.GradientTape() as gp1:
            gp1.watch(xi)
            u = self.predict_u(xi)

        grad_u = gp1.gradient(u, xi)
        dUnn_dx = tf.gather(grad_u, [self.ndim], axis=1)
        dUnn_dy = tf.gather(grad_u, [self.ndim+1], axis=1)

        s = tf.reshape(self.fkls, [self.Nbatch * self.numcp * self.numgp ** 2, 1])
        k = tf.ones([tf.size(dUnn_dx), 1])*100
        residual = 0.5*((k*self.dvdx * dUnn_dx) + (k*self.dvdy * dUnn_dy)) - (s * self.w)
        return residual

    @tf.function
    def lossMinimize(self, xit, W_r, W_b, scale_step):
        if self.dataonly:
            tvars = self.model.trainable_variables
            with tf.GradientTape() as g:
                l1 = self.model.compute_loss(xit, self.data_s, self.predict_u(self.data_input))
                l2 = 0
                l3 = 0
                lss = l1 + l2 + l3
            grads = g.gradient(lss, tvars)  # `tvars` are the trainable variables
            self.model.optimizer.apply_gradients(zip(grads, tvars))
            return lss, l1, l2, l3, tf.zeros([self.Npt,]), tf.zeros([self.Npt,])
        else:
            tvars = self.model.trainable_variables
            w1 = self.Nbatch / self.Npt
            w2 = self.bc_weight
            w3 = 1
            l1 = tf.cast(0, dtype=tf.float32)
            xi = tf.cast(tf.reshape(tf.repeat(xit, repeats=self.numcp * self.numgp ** 2, axis=0),
                                    [self.Nbatch * self.numcp * self.numgp ** 2, self.ndim]), dtype=tf.float32)
            xi = tf.concat([tf.concat([xi, self.x_loc], 1), self.y_loc], 1)
            xi_bc = tf.cast(tf.repeat(xit, self.num_bc_points * 4, axis=0), dtype=tf.float32)
            xi_bc = tf.concat([tf.concat([xi_bc, self.xi_0], 1), self.yi_0], 1)
            with tf.GradientTape() as g:
                bc = self.predict_u(xi_bc)
                bc_vector = tf.reshape(bc, [self.Nbatch, self.num_bc_points * 4])
                bc_lss = tf.reduce_sum(tf.math.abs(bc_vector - self.bc), axis=1) / (self.num_bc_points * 4)

                l2 = tf.math.reduce_mean(tf.reshape(bc_lss, [bc_lss.size, 1]) ** 2 * W_b)

                residual = self.net_predict_r(xi)
                cp_vector = tf.reshape(tf.reduce_sum(tf.reshape(residual, [self.Nbatch * self.numcp, self.numgp ** 2]), axis=1),
                    [self.Nbatch, self.numcp])
                res_lss = tf.reduce_sum(tf.math.abs(cp_vector), axis=1) / self.numcp
                l3 = tf.math.reduce_mean(tf.reshape(res_lss, [res_lss.size, 1]) ** 2 * W_r)
                lss = w1 * l1 + w2 * l2 + w3 * l3
            grads = g.gradient(lss, tvars)  # `tvars` are the trainable variables
            self.model.optimizer.apply_gradients(zip(grads, tvars))
            return lss, l1, l2, l3, bc_lss, res_lss


    def compute_weights(self, data):
        mean_data = np.mean(data)
        std_data = np.std(data)
        coeff_variation = np.array(std_data / mean_data)

        # Normalize the data
        min_data = tf.reduce_min(data)
        max_data = tf.reduce_max(data)
        data = (data - min_data) / (max_data - min_data)
        Source_kde = scipy.stats.gaussian_kde(data)
        P_s = Source_kde(data)
        w = 1 - coeff_variation * P_s
        for n in range(Npt):
            if w[n] < 0.01:
                w[n] = 0.01
        final_weights = w / (sum(w) / Npt)
        return final_weights

    def train(self):
        start_time = timeit.default_timer()
        self.buildDataVector(self.xdata)

        fname = self.model_name[0:-9] + '_output.txt'
        fname_err = self.model_name[0:-9] + '_error_history.txt'
        fid = open(fname, 'w')
        fid.close()
        fid = open(fname_err, 'w')
        fid.write('Iteration #, Mean relative error between SCPINN and FE at SC points\n')

        fid.close()
        i1 = 0
        iloop = tf.constant(self.error_write)
        xdatait = iter(self.xdata.enumerate(0))
        min_error = 10
        Data = np.loadtxt(self.data_name, delimiter=',')
        if Npt == 1:
            Data = np.reshape(Data, [1, Data.size])
            self.FE_sol = np.reshape(Data[0, ndim:self.FEA_Mesh_size + ndim], [1, self.FEA_Mesh_size])
            FEA_data =  np.reshape(Data[0, 0:self.FEA_Mesh_size + ndim], [1, self.FEA_Mesh_size + ndim])
        else:
            FEA_data = Data[:, 0:self.FEA_Mesh_size + ndim]
            self.FE_sol = Data[:, ndim:self.FEA_Mesh_size + ndim]

        index = 0
        self.input_id = int((index % (self.Npt / self.Nbatch)) * self.Nbatch)

        W_r = np.ones([self.Npt])
        with open('Npt_weights.txt', 'w') as fid:
            for n in range(0, self.Npt):
                fid.write('%0.8f \n' % W_r[n])

        W_b = np.ones([self.Npt])
        with open('Npt_bc_weights.txt', 'w') as fid:
            for n in range(0, self.Npt):
                fid.write('%0.8f \n' % W_b[n])

        while i1 < self.niter:
            for _ in range(iloop):
                loop_residuals = np.zeros((int(self.Npt / self.Nbatch), self.Nbatch))
                loop_bc = np.zeros((int(self.Npt / self.Nbatch), self.Nbatch))
                for j in range(int(self.Npt / self.Nbatch)):  # Use `j` for indexing batches
                    (_, inpdata) = xdatait.next()

                    # Extract inputs
                    xi = inpdata[:, 0:self.ndim]
                    self.fkls = inpdata[:, self.ndim + self.nDofs:]
                    scale_step = min(i1+1, self.Num_epochs * self.Ramping_percent)
                    # Perform loss minimization
                    (lss, l1, l2, l3, boundary_lss, res_lss) = self.lossMinimize(xi, W_r, W_b, scale_step)

                    # Store the residual into the preallocated array
                    loop_residuals[j, :] = res_lss
                    loop_bc[j, :] = boundary_lss
                    index += 1
                    self.input_id = int((index % (self.Npt / self.Nbatch)) * self.Nbatch)

            i1 += iloop
            if self.dataonly:
                xi_temp = tf.cast(tf.reshape(tf.repeat(xi, repeats=self.numcp * self.numgp ** 2, axis=0),
                                        [self.Nbatch * self.numcp * self.numgp ** 2, self.ndim]), dtype=tf.float32)
                xi_temp = tf.concat([tf.concat([xi_temp, self.x_loc], 1), self.y_loc], 1)
                xi_bc = tf.cast(tf.repeat(xi, self.num_bc_points * 4, axis=0), dtype=tf.float32)
                xi_bc = tf.concat([tf.concat([xi_bc, self.xi_0], 1), self.yi_0], 1)
                bc = self.predict_u(xi_bc)
                bc_vector = tf.reshape(bc, [self.Nbatch, self.num_bc_points * 4])
                bc_lss = tf.reduce_sum(tf.math.abs(bc_vector - self.bc), axis=1) / (self.num_bc_points * 4)
                residual = self.net_predict_r(xi_temp)
                cp_vector = tf.reshape(
                    tf.reduce_sum(tf.reshape(residual, [self.Nbatch * self.numcp, self.numgp ** 2]), axis=1),
                    [self.Nbatch, self.numcp])
                res_lss = tf.reduce_sum(tf.math.abs(cp_vector), axis=1) / self.numcp
                loop_residuals[j, :] = res_lss
                loop_bc[j, :] = boundary_lss
            loop_residuals = tf.squeeze(loop_residuals, axis=0)
            loop_bc = tf.squeeze(loop_bc, axis=0)

            if self.Weighted:
                W_r = self.compute_weights(loop_residuals)
                W_b = self.compute_weights(loop_bc)
            elapsed_time = timeit.default_timer() - start_time
            print('It: %d,Epoch#: %d, Loss_u: %.3e,Time: %.2f' % (
                i1, i1 // (self.Npt / self.Nbatch), lss, elapsed_time))

            res_npt = np.reshape(loop_residuals, [self.Npt, ])
            bc_npt = np.reshape(loop_bc, [self.Npt, ])
            epoch = int(i1.numpy() // (self.Npt / self.Nbatch))

            u_pred = np.zeros([Npt, self.nDofs])
            # Computes the error at each interval
            for n in range(0, self.nDofs):
                temp_input = tf.convert_to_tensor(np.append(FEA_data[:, 0:ndim], np.ones([Npt, 1]) * self.FEA_Mesh[n, :], 1))
                u_pred[:, n] = np.array(model.predict_u(temp_input))[:, 0]
            L2_error = tf.reduce_sum(np.linalg.norm(self.FE_sol - u_pred, ord=2, axis=1) / np.linalg.norm(self.FE_sol, ord=2, axis=1)) / self.Npt
            Max_error = tf.reduce_max(tf.abs(self.FE_sol-u_pred))
            Mean_error = tf.reduce_mean(tf.abs(self.FE_sol-u_pred))

            Max_errors = [np.max(np.abs(self.FE_sol[n, :] - u_pred[n, :])) for n in range(self.Npt)]
            Total_loss = [res_npt[n] ** 2 + self.bc_weight * bc_npt[n] ** 2 for n in range(self.Npt)]
            Weighted_total_loss = [res_npt[n] ** 2 * W_r[n] + self.bc_weight * bc_npt[n] ** 2 * W_b[n] for n in range(self.Npt)]

            # Save the individual inputs measurements
            res_history = 'Residual_history/Res_history_epoch_' + str(i1.numpy() // (self.Npt / self.Nbatch)) + '.txt'
            with open(res_history, 'a') as fid:
                for n in range(0, self.Npt):
                    fid.write(
                        'Npt: %d, L2_Error: %.2e, Max_Error: %.2e, Loss: %.2e, Residual: %.2e, Boundary: %.2e, Weight: %.4e, B_Weight: %.4e\n' %
                        (n, np.linalg.norm(self.FE_sol[n, :] - u_pred[n, :], ord=2, axis=0) / np.linalg.norm(
                            self.FE_sol[n, :], ord=2, axis=0),
                         Max_errors[n], Total_loss[n], res_npt[n]**2, self.bc_weight * bc_npt[n] ** 2, W_r[n], W_b[n]))


            if makeplots:
                colors = np.arange(self.Npt)
                Scale_W_r = max(W_r) * .015
                Scale_W_b = max(W_b) * .015
                Scale_Max_errors = max(Max_errors) * .0025
                Scale_Total_loss = max(Total_loss) * .025

                # Scatter plot for max error vs W_r[n]
                plt.figure(figsize=(10, 5))
                plt.scatter(W_r[:self.Npt], Max_errors, c=colors, cmap='viridis')
                plt.colorbar(label='Index')
                plt.xlabel('W_r')
                plt.ylabel('Max Absolute Error')

                # Add index labels to each point
                for n in range(self.Npt):
                    plt.text(W_r[n] + Scale_W_r, Max_errors[n] - Scale_Max_errors, str(n), fontsize=8, ha='center',
                             va='center')
                plt.savefig('Residual_plots/Scatter_Wr_epoch_' + str(i1.numpy() // (self.Npt / self.Nbatch)) + '.png')
                plt.close()

                # Scatter plot for max error vs W_b[n]
                plt.figure(figsize=(10, 5))
                plt.scatter(W_b[:self.Npt], Max_errors, c=colors, cmap='viridis')
                plt.colorbar(label='Index')
                plt.xlabel('W_b')
                plt.ylabel('Max Absolute Error')

                # Add index labels to each point
                for n in range(self.Npt):
                     plt.text(W_b[n] + Scale_W_b, Max_errors[n] - Scale_Max_errors, str(n), fontsize=8, ha='center',
                             va='center')
                plt.savefig('Residual_plots/Scatter_Wb_epoch_' + str(i1.numpy() // (self.Npt / self.Nbatch)) + '.png')
                plt.close()

                # Scatter plot for Max error vs Total Loss
                plt.figure(figsize=(10, 5))
                plt.scatter(Total_loss, Max_errors, c=colors, cmap='viridis')
                plt.colorbar(label='Index')
                plt.xlabel('Total Loss')
                plt.ylabel('Max Absolute Error')
                # Add index labels to each point
                for n in range(self.Npt):
                    plt.text(Total_loss[n] + Scale_Total_loss, Max_errors[n] - Scale_Max_errors, str(n), fontsize=8,
                             ha='center', va='center')

                # Save the plot
                plot_path_wb = 'Residual_plots/Scatter_Total_loss_epoch_' + str(
                    i1.numpy() // (self.Npt / self.Nbatch)) + '.png'
                plt.savefig(plot_path_wb)
                plt.close()

                # Scatter plot for Total Loss vs W_b[n]
                plt.figure(figsize=(10, 5))
                plt.scatter(W_b, Total_loss, c=colors, cmap='viridis')
                plt.colorbar(label='Index')
                plt.xlabel('W_b')
                plt.ylabel('Total Loss')

                # Add index labels to each point
                for n in range(self.Npt):
                    plt.text(W_b[n] + Scale_W_b, Total_loss[n] - Scale_Total_loss / 5, str(n), fontsize=8, ha='center',
                              va='center')
                plt.savefig('Residual_plots/Scatter_Wb_Total_loss_epoch_' + str(
                     i1.numpy() // (self.Npt / self.Nbatch)) + '.png')
                plt.close()


                p_loss = np.array(Weighted_total_loss) / np.sum(Weighted_total_loss)
                sorted_indices = np.argsort(Max_errors)
                sorted_max_error = np.array(Max_errors)[sorted_indices]
                sorted_pro_loss = p_loss[sorted_indices]

                # Compute the cumulative sum of pro_loss for the sorted max_error values
                cumulative_sum = np.cumsum(sorted_pro_loss)
                plt.figure(figsize=(10, 5))
                plt.plot(sorted_max_error, cumulative_sum, marker='.', linestyle='-')
                plt.xlim(0, 50)

                # Customize the plot
                plt.title('CDF of Total Weighted Loss')
                plt.xlabel('Max Error')
                plt.ylabel('Percent of Total Weighted Loss')
                plt.grid(True)
                plt.savefig('Residual_plots/Cumulative_loss_epoch_' + str(
                    i1.numpy() // (self.Npt / self.Nbatch)) + '.png')
                plt.close()

                sorted_indices = np.argsort(res_npt)
                sorted_res = np.array(res_npt)[sorted_indices]
                sorted_pro_loss = p_loss[sorted_indices]

                # # Scatter plot for Total Loss vs W_r[n]
                plt.figure(figsize=(10, 5))
                plt.scatter(W_r, Total_loss, c=colors, cmap='viridis')
                plt.colorbar(label='Index')
                plt.xlabel('W_r')
                plt.ylabel('Total Loss')

                # Add index labels to each point
                for n in range(self.Npt):
                    plt.text(W_r[n] + Scale_W_r, Total_loss[n] - Scale_Total_loss / 5, str(n), fontsize=8, ha='center',
                             va='center')
                plt.savefig('Residual_plots/Scatter_Wr_Total_loss_epoch_' + str(
                    i1.numpy() // (self.Npt / self.Nbatch)) + '.png')
                plt.close()

            #Saves images of the solution
            Solution_image = np.where(np.max(self.FE_sol)==self.FE_sol)[0][0]
            X = np.outer(np.linspace(0, 1, int(self.nelem[0]) + 1), np.ones(int(self.nelem[1]) + 1))
            Y = np.outer(np.linspace(0, 1, int(self.nelem[1]) + 1), np.ones(int(self.nelem[0]) + 1)).T
            solution = np.reshape(u_pred[Solution_image, :], [int(self.nelem[0]) + 1, int(self.nelem[0]) + 1])
            solution_FEA = np.reshape(self.FE_sol[Solution_image, :], [int(self.nelem[0]) + 1, int(self.nelem[0]) + 1])

            fig = plt.figure(figsize=plt.figaspect(.333))
            ax = fig.add_subplot(1, 3, 1, projection='3d')
            ax.plot_surface(X, Y, solution, cmap=matplotlib.cm.coolwarm)
            plt.title('Model Solution')
            ax = fig.add_subplot(1, 3, 2, projection='3d')
            ax.plot_surface(X, Y, solution_FEA, cmap=matplotlib.cm.coolwarm)
            plt.title('FEA Solution')
            ax = fig.add_subplot(1, 3, 3, projection='3d')
            ax.plot_surface(X, Y, solution_FEA - solution, cmap=matplotlib.cm.coolwarm)
            plt.title('Absolute Error')
            label = ('Input #' + str(Solution_image))
            fig.text(0.5, 0.95, label, ha='center', va='center', fontsize=14)
            #label = 'Input weight was: ' + str(np.round(W_r[Solution_image].numpy(),2))
            label = 'W_r: ' + str(np.round(W_r[Solution_image],2)) + '    W_b: ' + str(np.round(W_b[Solution_image],2))
            fig.text(0.5, 0.05, label, ha='center', va='center', fontsize=12)
            plt.savefig('Solution_images/Npt_' + str(int(Solution_image)) + '_epoch_' + str(epoch) + '.png', bbox_inches='tight')
            plt.close(fig)

            with open(fname_err, 'a') as fid:
                fid.write('It: %d, Epoch#: %d, L2_Error: %.3e, Max_Error: %.3e, Mean_Error: %.3e, Time: %.2f\n' % \
                          (i1, i1 // (self.Npt / self.Nbatch), L2_error, Max_error, Mean_error, elapsed_time))
            if Max_error < min_error:
                min_error = Max_error
                self.model.save('2DHeat_model_min_error.h5')

            self.model.save('Model_history/2DHeat_model.h5')
def loadBaselineInputs(path, pts, Start,Data_load, ndim):
    yhat = np.genfromtxt(Data_load, delimiter=',')[Start - 1:Start - 1 + pts, 0:ndim]
    vhat = np.loadtxt(path + '/V_hat.txt', delimiter=',')
    wlam = np.loadtxt(path + '/W_lam.txt', delimiter=',')

    return yhat, wlam, vhat


if __name__ == '__main__':
    # User Inputs
    # Stochastic Distribution Veriables
    load_baseline = True  # In this version of code alway set to true. Generate soucre before running Main.py
    Dist_form = 1
    Level = 'Low'
    ndim = 4  # number of stochastic dimensions
    Npt = 400  # Total number of training samples
    Start = 1  # Select which input from the data is the start index for training

    # Network Configuration Veriables
    dataonly = False # Select True for data driven training
    makeplots = True
    loadold = False # Set to true if restarting a training
    Network_method = 'CNN' # Choice between 'CNN' [Fully Connected Network] and 'Deep' [DeepONet]
    Learning_rate = 2e-3
    Learning_rate_decay = .985
    Learning_rate_step = 500
    Num_hidden_layers = 4
    Hidden_layer_size = 100
    Num_cp_x = 24
    Num_cp_y = 24
    num_bc_points = 74
    Num_epochs = 75000
    batchSize = np.min([1000, Npt])

    # Loss Funciton Veriables
    bc_weight = 1e7
    inputs_weighted = False #Select True for weighted loss training
    weight_method = 'KDE' # Method to distribute weights
    Ramping_percent = 3.5 # Ramping rate if selected
    alpha = 0.4 # Weighting coeeficent
    mode = 1

    # File Locations
    p = '/Users/nwade/PycharmProjects/StochasticPINN/' # change to your file directory
    Study_name = 'Data_1' # Folder Name where files will be stored
    Trial_num = 1
    error_write = 1000  # Select how often error values should be exported
    # End of Inputs

    # Main Code Set up
    tf.config.run_functions_eagerly(False)

    if not os.path.exists(p + Study_name):
        os.makedirs(p + Study_name)
    File_path = '/Users/nwade/PycharmProjects/StochasticPINN/' + Study_name + '/Trial_' + str(Trial_num)
    Data_load = p + 'FE_data/Solution_data_' + str(int(Dist_form)) + '/' + Level + '/Ndim_' + str(ndim) + '_CP_' + str(int(Num_cp_x)) + '_data.txt'
    if not os.path.exists(File_path):
        os.makedirs(File_path)
        os.makedirs(File_path + str('/Solution_images'))
        os.makedirs(File_path + str('/Residual_plots'))
        os.makedirs(File_path + str('/Residual_history'))
        os.makedirs(File_path + str('/Model_history'))
        os.chdir(File_path)
    else:
        os.chdir(File_path)

    print('Trial#: ' + str(Trial_num))
    # PDE problem specific parameters
    x_l = float(0)
    x_u = float(1)
    y_l = float(0)
    y_u = float(1)
    Element_size = 0.02

    # Core model parameters
    Num_Data = Npt
    layer_u = [ndim + 2]
    for n in range(0, Num_hidden_layers):
        layer_u.append(Hidden_layer_size)
    niter = Num_epochs * (Npt // batchSize)
    NW_bins = 20
    Max_weight = 50
    Num_solves = 1
    Num_per_solve = Npt / Num_solves

    # Integration parameters
    IR_size = 4 * Element_size
    ngp = 2

    # Output parameters
    model_name = '2DHeat_model.h5'
    model_load = '2DHeat_model_min_error.h5'
    model_FEA_name = '2DHeat_FEA_data.h5'
    model_data_name = '2DHeat_data.txt'

    # Initialize variables
    y_length = y_u - y_l
    x_length = x_u - x_l
    coordinate_domain = [[int(x_l), int(x_u)], [int(y_l), int(y_u)]]
    nelem = [int(x_length / Element_size), int(y_length / Element_size)]

    Dist_loc = p + 'FE_Data/Solution_data_' + str(int(Dist_form)) + '/' + Level
    (Y_hat, W_lam, V_hat) = loadBaselineInputs(Dist_loc, Npt, Start, Data_load, ndim)

    with open('Trial_parameters.txt', 'w') as fid:
        fid.write('Network Structure: ' + Network_method + '\n')
        fid.write('Data Only: ' + str(dataonly) + '\n')
        fid.write('Weighted: ' + str(inputs_weighted) + '\n')
        fid.write('Npt: %d\n' % Npt)
        fid.write('Distribution: ' + str(Dist_form) + '\n')
        fid.write('Level: ' + str(Level) + '\n')
        fid.write('Method: ' + weight_method + '\n')
        fid.write('Alpha: %0.2e\n' % alpha)
        fid.write('Ramping Percent: %0.2e\n' % Ramping_percent)
        fid.write('Mode: %d\n' % mode)
        fid.write('Learning Rate: %.2e\n' % Learning_rate)
        fid.write('Learning Decay: %.2e\n' % Learning_rate_decay)
        fid.write('Learning Step: %0.2e\n' % Learning_rate_step)
        fid.write('Boundary Weight: %.2e\n' % bc_weight)
        fid.write('Boundary Points: %d\n' % num_bc_points)
        fid.write('Start Number: %d\n' % Start)
        fid.write('Network depth: %d\n' % Num_hidden_layers)
        fid.write('Layer Size: %d\n' % Hidden_layer_size)
        fid.write('CP_x: %d\n' % Num_cp_x)
        fid.write('CP_y: %d\n' % Num_cp_y)

    Input_dict = {}
    Input_dict['Network'] = Network_method
    Input_dict['dataonly'] = dataonly
    Input_dict['loadold'] = loadold
    Input_dict['Trial_num'] = Trial_num
    Input_dict['Num_cp_x'] = Num_cp_x
    Input_dict['Num_cp_y'] = Num_cp_y
    Input_dict['Num_Data'] = Num_Data
    Input_dict['Num_epochs'] = Num_epochs
    Input_dict['Npt'] = Npt
    Input_dict['ndim'] = ndim
    Input_dict['layer_u'] = layer_u
    Input_dict['batchSize'] = batchSize
    Input_dict['num_bc_points'] = num_bc_points
    Input_dict['niter'] = niter
    Input_dict['Element_size'] = Element_size
    Input_dict['nelem'] = nelem
    Input_dict['IR_size'] = IR_size
    Input_dict['ngp'] = ngp
    Input_dict['model_name'] = model_name
    Input_dict['model_FEA_name'] = model_FEA_name
    Input_dict['model_data_name'] = model_data_name
    Input_dict['2DHeat_model_min_error.h5'] = model_load
    Input_dict['error_write'] = error_write
    Input_dict['coordinate_domain'] = coordinate_domain
    Input_dict['nelem'] = nelem
    Input_dict['Y_hat'] = Y_hat
    Input_dict['W_lam'] = W_lam
    Input_dict['V_hat'] = V_hat
    Input_dict['bc_weight'] = bc_weight
    Input_dict['inputs_weighted'] = inputs_weighted
    Input_dict['NWbins'] = NW_bins
    Input_dict['Path'] = p
    Input_dict['Num_solves'] = Num_solves
    Input_dict['Max_weight'] = Max_weight
    Input_dict['Learning_rate'] = Learning_rate
    Input_dict['Learning_rate_decay'] = Learning_rate_decay
    Input_dict['Learning_rate_step'] = Learning_rate_step
    Input_dict['Start'] = Start
    Input_dict['Data_load'] = Data_load
    Input_dict['weight_method'] = weight_method
    Input_dict['Ramping_percent'] = Ramping_percent
    Input_dict['alpha'] = alpha
    Input_dict['mode'] = mode

    # Trains the model
    model = SC_Heat_2D(Input_dict)
    model.train()