How to code PINN for a simple PDE

username: Ali_geo
cooked: <p>Dear community,<br>
I am trying to solve a simlpe PDE using PINN. My PDE is a 1D heat transfer:</p>
<blockquote>
<p>dT/dt = lamda/(rho*cp) * d2T/dt2<br>
I tried some simple finite difference siólutions and it worked out for me:<br>
import torch<br>
import torch.nn as nn<br>
import matplotlib.pyplot as plt<br>
import numpy as np<br>
from pyDOE import lhs<br>
######### Finite difference solution</p>
<h1>geometry:</h1>
<p>L = 0.1 # length of the rod</p>
<h1>mesh:</h1>
<p>dx = 0.001<br>
nx = int(L/dx) + 1<br>
x = np.linspace(0, L, nx)</p>
<h1>temporal grid:</h1>
<p>t_sim = 60<br>
dt = 0.01<br>
nt = int (t_sim/dt)</p>
<h1>parametrization</h1>
<p>rho, lamda, cp = 1, 1, 1</p>
<h1>IC</h1>
<p>t_ic = 50</p>
<h1>BC</h1>
<p>t_left = 70 # left side with 60 °C temperature<br>
t_right = 40 # right side with 40 °C temperature</p>
<h1>Results</h1>
<p>T = np.ones(nx) * t_ic<br>
all_T = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
for i in range (0, nt):<br>
Tn = T.copy()<br>
T[1:-1] = Tn[1:-1] + dt/(dx++2) * (lamda/cp) * (Tn[2:] - 2*Tn[1:-1] + Tn[0:-2])<br>
T[0] = t_left<br>
T[-1] = t_right<br>
<span class="hashtag-raw">#T</span>[-1] = T[-2] # Neumann<br>
all_T.append(Tn)<br>
x = torch.linspace(0, L, nx, dtype=torch.float32)<br>
t = torch.linspace(0, t_sim, nt, dtype=torch.float32)<br>
T, X = torch.meshgrid(t,x)<br>
Temps = np.concatenate (all_T).reshape(nt,nx)<br>
The following code visualizes the results:<br>
fig,ax=plt.subplots(1,1)<br>
cp = ax.contourf(X, T, Temps, cmap=“rainbow”)<br>
fig.colorbar(cp) # Add a colorbar to a plot<br>
ax.set_xlabel(‘x’)<br>
ax.set_ylabel(‘time’)<br>
plt.show()</p>
</blockquote>
<p>Which is somthing like:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/c/8c624992ea1a8bc4dff330f5060396b7e817c4e8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/8c624992ea1a8bc4dff330f5060396b7e817c4e8" title="fig_FD"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/c/8c624992ea1a8bc4dff330f5060396b7e817c4e8_2_682x500.jpeg" alt="fig_FD" data-base62-sha1="k1Tqoanitalyw80B3MG6r4eygU8" width="682" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/c/8c624992ea1a8bc4dff330f5060396b7e817c4e8_2_682x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/c/8c624992ea1a8bc4dff330f5060396b7e817c4e8_2_1023x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/c/8c624992ea1a8bc4dff330f5060396b7e817c4e8_2_1364x1000.jpeg 2x" data-dominant-color="6AD9E2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">fig_FD</span><span class="informations">2520×1846 115 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Then, I tried to solve the same PDE using PINN but results are not acceptable at all. This is my code:</p>
<blockquote>
<p>x_test = torch.hstack((X.transpose(1,0).flatten()[:,None], T.transpose(1,0).flatten()[:,None]))<br>
y_test = torch.from_numpy(Temps) # I suppose it is the ground truth<br>
lb = x_test[0] # lower boundary<br>
ub = x_test[-1] # upper boundary<br>
left_x = torch.hstack((X[:,0][:,None], T[:,0][:,None])) # x and t of left boundary<br>
left_y = torch.ones(left_x.shape[0], 1) * t_left # Temperature of left boundary<br>
right_x = torch.hstack((X[:,-1][:,None], T[:,0][:,None]))  # x and t of right boundary<br>
right_y = torch.ones(right_x.shape[0], 1) * t_right # Temperature of right boundary<br>
bottom_x = torch.hstack((X[0,1:-1][:,None], T[0,1:-1][:,None])) # x and t of IC<br>
bottom_y = torch.ones(bottom_x.shape[0], 1) * t_ic # Temperature of IC<br>
X_train = torch.vstack([left_x, right_x, bottom_x]) # all of the BCs and IC<br>
Y_train = torch.vstack([left_y, right_y, bottom_y])<br>
No = 100<br>
idx = np.random.choice(left_x.shape[0], No, replace=False)<br>
X_train_No = X_train[idx,:] # random points of BC and IC for training<br>
Y_train_No = Y_train[idx,:]<br>
N_f = 750<br>
X_train_Nf = lb + (ub-lb)<em>lhs(2,N_f)<br>
X_train_Nf = torch.vstack((X_train_Nf, X_train_No)) # random points of domain for training<br>
f_hat = torch.zeros(X_train_Nf.shape[0],1, dtype=torch.float32) # zero array for loss of PDE<br>
class FCN(nn.Module):<br>
#<span class="hashtag-raw">#Neural</span> Network<br>
def <strong>init</strong>(self,layers):<br>
super().<strong>init</strong>() <span class="hashtag-raw">#call</span> <strong>init</strong> from parent class<br>
self.activation = nn.Tanh()<br>
self.loss_function = nn.MSELoss(reduction =‘mean’)<br>
‘Initialise neural network as a list using nn.Modulelist’<br>
self.linears = nn.ModuleList([nn.Linear(layers[i], layers[i+1]) for i in range(len(layers)-1)])<br>
self.iter = 0<br>
‘Xavier Normal Initialization’<br>
for i in range(len(layers)-1):<br>
nn.init.xavier_normal_(self.linears[i].weight.data, gain=1.0)<br>
nn.init.zeros_(self.linears[i].bias.data)<br>
‘foward pass’<br>
def forward(self,x):<br>
if torch.is_tensor(x) != True:<br>
x = torch.from_numpy(x)<br>
a = x.float()<br>
for i in range(len(layers)-2):<br>
z = self.linears<a>i</a><br>
a = self.activation(z)<br>
a = self.linears<a>-1</a><br>
return a<br>
‘Loss Functions’<br>
<span class="hashtag-raw">#Loss</span> BC<br>
def lossBC(self,x_BC,y_BC):<br>
loss_BC=self.loss_function(self.forward(x_BC),y_BC)<br>
return loss_BC<br>
<span class="hashtag-raw">#Loss</span> PDE<br>
def lossPDE(self,x_PDE):<br>
g=x_PDE.clone()<br>
g.requires_grad=True # Enable differentiation<br>
f=self.forward(g)<br>
f_x_t=torch.autograd.grad(f,g,torch.ones([g.shape[0],1]),retain_graph=True, create_graph=True)[0] <span class="hashtag-raw">#first</span> derivative<br>
f_xx_tt=torch.autograd.grad(f_x_t,g,torch.ones(g.shape), create_graph=True)[0]<span class="hashtag-raw">#second</span> derivative<br>
f_t = f_x_t[:,[1]]<br>
f_xx = f_xx_tt[:,[0]]<br>
f = f_t - (1/(1</em>1))*f_xx<br>
return self.loss_function(f,f_hat)</p>
<pre><code>def loss(self,x_BC,y_BC,x_PDE):
    loss_bc=self.lossBC(x_BC.float(),y_BC.float())
    loss_pde=self.lossPDE(x_PDE.float())
    return loss_bc.float()+loss_pde.float()
</code></pre>
<p>layers=np.array([2, 50, 50, 50, 50, 1])<br>
PINN=FCN(layers)<br>
optimizer = torch.optim.Adam(PINN.parameters(), lr=0.001)<br>
def closure():<br>
optimizer.zero_grad()<br>
loss = PINN.loss(X_train_No, Y_train_No, X_train_Nf)<br>
loss.backward()<br>
return loss<br>
for i in range(10000):<br>
loss = optimizer.step(closure)<br>
if (i+1)%2500==0:<br>
with torch.no_grad():<br>
test_loss=PINN.lossBC(x_test,y_test.flatten().view(-1,1))<br>
print(‘Training:’, loss.cpu().detach().numpy(), ‘/ Testing’, test_loss.cpu().detach().numpy())<br>
y1 = PINN(x_test)<br>
arr_y1 = np.array (y1.reshape(shape=[nx,nt]).transpose(1,0).detach())<br>
fig,ax=plt.subplots(1,1)<br>
cp = ax.contourf(X, T, arr_y1, cmap=“rainbow”)<br>
fig.colorbar(cp) # Add a colorbar to a plot<br>
ax.set_xlabel(‘x’)<br>
ax.set_ylabel(‘time’)<br>
plt.show()</p>
</blockquote>
<p>Sorry for raising suchlong problem, I very much appreciate any help.</p>