[Python] Implement Physics informed Neural Network using pytorch

I found a very interesting paper, Physics Informed Deep Learning (Part I): Data-driven Solutions of Nonlinear Partial Differential Equations and want to give it a trial. For this, I create a dummy problem and implement what I understand from the paper.

Problem Statement


Suppose, I want to solve the ODE dy/dx = cos(x) with initial conditions y(0)=y(2*pi)=0. Actually, we can easily guess the analytic solution y(x)=sin(x). But I want to see how the model predict the solution using PINN.

# import libraries
import torch
import torch.autograd as autograd # computation graph
import torch.nn as nn # neural networks
import torch.optim as optim # optimizers e.g. gradient descent, ADAM, etc.
import matplotlib.pyplot as plt
import numpy as np

#Set default dtype to float32
torch.set_default_dtype(torch.float)
#PyTorch random number generator
torch.manual_seed(1234)
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

Model Architecture


## Model Architecture
class FCN(nn.Module):
##Neural Network
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, 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.weight.data, gain=1.0)
# set biases to zero
nn.init.zeros_(self.linears.bias.data)

# foward pass
def forward(self,x):
if torch.is_tensor(x) != True:
x = torch.from_numpy(x)
a = x.float()
for i in range(len(layers)-2):
z = self.linears(a)
a = self.activation(z)
a = self.linears[-1](a)
return a
# Loss Functions
#Loss PDE
def lossPDE(self,x_PDE):
g=x_PDE.clone()
g.requires_grad=True #Enable differentiation
f=self.forward(g)
f_x=autograd.grad(f,g,torch.ones([x_PDE.shape[0],1]).to(device),\
retain_graph=True, create_graph=True)[0]
loss_PDE=self.loss_function(f_x,PDE(g))
return loss_PDE

Generate data


# generate training and evaluation points
x = torch.linspace(min,max,total_points).view(-1,1)
y = torch.sin(x)
print(x.shape, y.shape)

# Set Boundary conditions:
# Actually for this problem
# we don't need extra boundary constraint
# as it was concided with x_PDE point & value
# BC_1=x[0,:]
# BC_2=x[-1,:]
# print(BC_1,BC_2)
# x_BC=torch.vstack([BC_1,BC_2])
# print(x_BC)
x_PDE = x[1:-1,:]
print(x_PDE.shape)

x_PDE=x_PDE.float().to(device)
# x_BC=x_BC.to(device)
#Create Model
layers = np.array([1,50,50,50,50,1])
model = FCN(layers)
print(model)
model.to(device)
params = list(model.parameters())
optimizer = torch.optim.Adam(model.parameters(),lr=lr,amsgrad=False)

Train Neural Network


for i in range(500):
yh = model(x_PDE)
loss = model.loss_PDE(x_PDE) # use mean squared error
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i%(500/10)==0:
print(loss)

predict the solution using PINN


# predict the solution beyond training set
x = torch.linspace(0,max+max,total_points).view(-1,1)
yh=model(x.to(device))
y=torch.sin(x)
#Error
print(model.lossBC(x.to(device)))

y_plot=y.detach().numpy()
yh_plot=yh.detach().cpu().numpy()
fig, ax1 = plt.subplots()
ax1.plot(x,y_plot,color='blue',label='Real')
ax1.plot(x,yh_plot,color='red',label='Predicted')
ax1.set_xlabel('x',color='black')
ax1.set_ylabel('f(x)',color='black')
ax1.tick_params(axis='y', color='black')
ax1.legend(loc = 'upper left')




But the end result was so disappointing. The model was unable to learn the simple ODE. I was wondering the model architecture of mine may have some issue which I couldn't figure out myself. Can anyone suggest any improvement?

Continue reading...