基于pytorch的PINN代码

import torchimport torch.autograd as autograd         # computation graphfrom torch import Tensor                  # tensor node in the computation graphimport torch.nn as nn                     # neural networksimport torch.optim as optim               # optimizers e.g. gradient descent, ADAM, etc.import matplotlib.pyplot as pltimport matplotlib.gridspec as gridspecfrom mpl_toolkits.axes_grid1 import make_axes_locatablefrom mpl_toolkits.mplot3D import Axes3Dimport matplotlib.tickerfrom sklearn.model_selection import train_test_splitimport numpy as npimport timefrom pyDOE import lhs         #Latin Hypercube Samplingimport scipy.io#Set default dtype to float32torch.set_default_dtype(torch.float)#PyTorch random number generatortorch.manual_seed(1234)# Random number generators in other librariesnp.random.seed(1234)# Device configurationdevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')print(device)if device == 'cuda':    print(torch.cuda.get_device_name())steps=10000lr=1e-1lr1 = 0.001layers = np.array([2,20,20,20,20,20,20,20,20,1]) #8 hidden layers#Nu: Number of training points # Nf: Number of collocation points (Evaluate PDE)N_u = 100 #Total number of data points for 'u'N_f = 10000 #Total number of collocation pointsnu = 0.01/np.pi #diffusion coeficientepoches = 1000def plot3D(x,t,y):  x_plot =x.squeeze(1)  t_plot =t.squeeze(1)  X,T= torch.meshgrid(x_plot,t_plot)  F_xt = y  fig,ax=plt.subplots(1,1)  cp = ax.contourf(T,X, F_xt,20,cmap="rainbow")  fig.colorbar(cp) # Add a colorbar to a plot  ax.set_title('F(x,t)')  ax.set_xlabel('t')  ax.set_ylabel('x')  plt.show()  ax = plt.axes(projection='3d')  ax.plot_surface(T.numpy(), X.numpy(), F_xt.numpy(),cmap="rainbow")  ax.set_xlabel('t')  ax.set_ylabel('x')  ax.set_zlabel('f(x,t)')  plt.show()def plot3D_Matrix(x, t, y):  X, T = x, t  F_xt = y  fig, ax = plt.subplots(1, 1)  cp = ax.contourf(T, X, F_xt, 20, cmap="rainbow")  fig.colorbar(cp)  # Add a colorbar to a plot  ax.set_title('F(x,t)')  ax.set_xlabel('t')  ax.set_ylabel('x')  plt.show()  ax = plt.axes(projection='3d')  ax.plot_surface(T.numpy(), X.numpy(), F_xt.numpy(), cmap="rainbow")  ax.set_xlabel('t')  ax.set_ylabel('x')  ax.set_zlabel('f(x,t)')  plt.show()def solutionplot(u_pred, X_u_train, u_train):    # https://GitHub.com/omniscientoctopus/Physics-InfORMed-Neural-Networks    fig, ax = plt.subplots()    ax.axis('off')    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(u_pred, interpolation='nearest', cmap='rainbow',                  extent=[T.min(), T.max(), X.min(), X.max()],                  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)' % (u_train.shape[0]), markersize=4,            clip_on=False)    line = np.linspace(x.min(), x.max(), 2)[:, None]    ax.plot(t[25] * np.ones((2, 1)), line, 'w-', linewidth=1)    ax.plot(t[50] * np.ones((2, 1)), line, 'w-', linewidth=1)    ax.plot(t[75] * np.ones((2, 1)), line, 'w-', linewidth=1)    ax.set_xlabel('$t$')    ax.set_ylabel('$x$')    ax.legend(frameon=False, loc='best')    ax.set_title('$u(x,t)$', fontsize=10)    '''     Slices of the solution at points t = 0.25, t = 0.50 and t = 0.75    '''    ####### Row 1: u(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, usol.T[25, :], 'b-', linewidth=2, label='Exact')    ax.plot(x, u_pred.T[25, :], 'r--', linewidth=2, label='Prediction')    ax.set_xlabel('$x$')    ax.set_ylabel('$u(x,t)$')    ax.set_title('$t = 0.25s$', fontsize=10)    ax.axis('square')    ax.set_xlim([-1.1, 1.1])    ax.set_ylim([-1.1, 1.1])    ax = plt.subplot(gs1[0, 1])    ax.plot(x, usol.T[50, :], 'b-', linewidth=2, label='Exact')    ax.plot(x, u_pred.T[50, :], 'r--', linewidth=2, label='Prediction')    ax.set_xlabel('$x$')    ax.set_ylabel('$u(x,t)$')    ax.axis('square')    ax.set_xlim([-1.1, 1.1])    ax.set_ylim([-1.1, 1.1])    ax.set_title('$t = 0.50s$', fontsize=10)    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.35), ncol=5, frameon=False)    ax = plt.subplot(gs1[0, 2])    ax.plot(x, usol.T[75, :], 'b-', linewidth=2, label='Exact')    ax.plot(x, u_pred.T[75, :], 'r--', linewidth=2, label='Prediction')    ax.set_xlabel('$x$')    ax.set_ylabel('$u(x,t)$')    ax.axis('square')    ax.set_xlim([-1.1, 1.1])    ax.set_ylim([-1.1, 1.1])    ax.set_title('$t = 0.75s$', fontsize=10)    plt.savefig('Burgers.png', dpi=500)class FCN(nn.Module):    def __init__(self, layers):        super().__init__()  # call __init__ from parent class        'activation function'        self.activation = nn.Tanh()        'loss function'        self.loss_function = nn.MSELoss(reduction='mean')        'Initialise neural network as a list using nn.Modulelist'        self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i + 1]) for i in range(len(layers) - 1)])        self.iter = 0        'Xavier Normal Initialization'        for i in range(len(layers) - 1):            nn.init.xavier_normal_(self.linears[i].weight.data, gain=1.0) #Xavier正态分布            wwww = self.linears[i].weight.data            # set biases to zero            nn.init.zeros_(self.linears[i].bias.data)    'foward pass'    def forward(self, x):        if torch.is_tensor(x) != True:            x = torch.from_numpy(x)        u_b = torch.from_numpy(ub).float().to(device)        l_b = torch.from_numpy(lb).float().to(device)        # preprocessing input        x = (x - l_b) / (u_b - l_b)  # feature scaling        # convert to float        a = x.float()        for i in range(len(layers) - 2):            z = self.linears[i](a)            a = self.activation(z)        a = self.linears[-1](a)        return a    def loss_BC(self, x, y):        loss_u = self.loss_function(self.forward(x), y)        return loss_u    def loss_PDE(self, X_train_Nf):        g = X_train_Nf.clone()        g.requires_grad = True        u = self.forward(g)        u_x_t = \        autograd.grad(u, g, torch.ones([X_train_Nf.shape[0], 1]).to(device), retain_graph=True, create_graph=True)[0]        u_xx_tt = autograd.grad(u_x_t, g, torch.ones(X_train_Nf.shape).to(device), create_graph=True)[0]        u_x = u_x_t[:, [0]]        u_t = u_x_t[:, [1]]        u_xx = u_xx_tt[:, [0]]        f = u_t + (self.forward(g)) * (u_x) - (nu) * u_xx        loss_f = self.loss_function(f, f_hat)        return loss_f    def loss(self, x, y, X_train_Nf):        loss_u = self.loss_BC(x, y)        loss_f = self.loss_PDE(X_train_Nf)        loss_val = loss_u + loss_f        return loss_val    'callable for optimizer'    def closure(self):        optimizer.zero_grad()        loss = self.loss(X_train_Nu, U_train_Nu, X_train_Nf)        print(loss)        loss.backward()        # self.iter += 1        #        # if self.iter % 100 == 0:        #     error_vec, _ = PINN.test()        #        #     print(loss, error_vec)        return loss    'test neural network'    def test(self):        u_pred = self.forward(X_test)        error_vec = torch.linalg.norm((u - u_pred), 2) / torch.linalg.norm(u, 2)  # Relative L2 Norm of the error (Vector)        u_pred = u_pred.cpu().detach().numpy()        u_pred = np.reshape(u_pred, (256, 100), order='F')        return error_vec, u_preddata = scipy.io.loadmat('./result/burgers_shock.mat')x = data['x']       # 256 points between -1 and 1 [256x1]t = data['t']       # 100 time points between 0 and 1 [100x1]usol = data['usol'] # solution of 256x100 grid pointsX, T = np.meshgrid(x,t)                         # makes 2 arrays X and T such that u(X[i],T[j])=usol[i][j] are a tupleplot3D(torch.from_numpy(x),torch.from_numpy(t),torch.from_numpy(usol)) #f_real was defined previously(function)print(x.shape,t.shape,usol.shape)print(X.shape,T.shape)X_test = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))# Domain boundslb = X_test[0]  # [-1. 0.]ub = X_test[-1] # [1.  0.99]u_true = usol.flatten('F')[:,None] #Fortran style (Column Major)print(lb,ub)'''Boundary Conditions'''#Initial Condition -1 =< x =<1 and t = 0left_X = np.hstack((X[0,:][:,None], T[0,:][:,None])) #L1left_U = usol[:,0][:,None]#Boundary Condition x = -1 and 0 =< t =<1bottom_X = np.hstack((X[:,0][:,None], T[:,0][:,None])) #L2bottom_U = usol[-1,:][:,None]#Boundary Condition x = 1 and 0 =< t =<1top_X = np.hstack((X[:,-1][:,None], T[:,0][:,None])) #L3top_U = usol[0,:][:,None]X_train = np.vstack([left_X, bottom_X, top_X])U_train = np.vstack([left_U, bottom_U, top_U])#choose random N_u points for trainingidx = np.random.choice(X_train.shape[0], N_u, replace=False)X_train_Nu = X_train[idx, :] #choose indices from  set 'idx' (x,t)U_train_Nu = U_train[idx,:]      #choose corresponding u'''Collocation Points'''# Latin Hypercube sampling for collocation points# N_f sets of tuples(x,t)X_train_Nf = lb + (ub-lb)*lhs(2,N_f)X_train_Nf = np.vstack((X_train_Nf, X_train_Nu)) # append training points to collocation pointsprint("Original shapes for X and U:",X.shape,usol.shape)print("Boundary shapes for the edges:",left_X.shape,bottom_X.shape,top_X.shape)print("Available training data:",X_train.shape,U_train.shape)print("Final training data:",X_train_Nu.shape,U_train_Nu.shape)print("Total collocation points:",X_train_Nf.shape)'Convert to tensor and send to GPU'X_train_Nf = torch.Tensor(X_train_Nf).float().to(device)X_train_Nu = torch.Tensor(X_train_Nu).float().to(device)U_train_Nu = torch.from_numpy(U_train_Nu).float().to(device)X_test = torch.from_numpy(X_test).float().to(device)u = torch.from_numpy(u_true).float().to(device)f_hat = torch.zeros(X_train_Nf.shape[0], 1).to(device)PINN = FCN(layers)PINN.to(device)'Neural Network Summary'print(PINN)params = list(PINN.parameters())'''Optimization'''optimizer1 = torch.optim.Adam(PINN.parameters(), lr1)'L-BFGS Optimizer'optimizer = torch.optim.LBFGS(PINN.parameters(), lr,  max_iter=steps,  max_eval=None,  tolerance_grad=1e-11,  tolerance_change=1e-11,  history_size=100,  line_search_fn='strong_wolfe')for epoch in range(epoches):    optimizer1.zero_grad()    loss = loss = PINN.loss(X_train_Nu, U_train_Nu, X_train_Nf)    print("Adam loss:",loss)    loss.backward()    optimizer1.step()start_time = time.time()optimizer.step(PINN.closure)elapsed = time.time() - start_timeprint('Training time: %.2f' % (elapsed))''' Model Accuracy '''error_vec, u_pred = PINN.test()print('Test Error: %.5f' % (error_vec))solutionplot(u_pred,X_train_Nu.cpu().detach().numpy(),U_train_Nu)x1=X_test[:,0]t1=X_test[:,1]arr_x1=x1.reshape(shape=X.shape).transpose(1,0).detach().cpu()arr_T1=t1.reshape(shape=X.shape).transpose(1,0).detach().cpu()arr_y1=u_predarr_y_test=usolplot3D(torch.from_numpy(x),torch.from_numpy(t),torch.from_numpy(usol))

来源地址：https://blog.csdn.net/m0_45861823/article/details/128919362