-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain.py
125 lines (96 loc) · 4.61 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import torch
import torch.nn as nn
from torch.autograd import Variable
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.hidden_layer1 = nn.Linear(2,5)
self.hidden_layer2 = nn.Linear(5,5)
self.hidden_layer3 = nn.Linear(5,5)
self.hidden_layer4 = nn.Linear(5,5)
self.hidden_layer5 = nn.Linear(5,5)
self.output_layer = nn.Linear(5,1)
def forward(self, x,t):
inputs = torch.cat([x,t],axis=1) # combined two arrays of 1 columns each to one array of 2 columns
layer1_out = torch.sigmoid(self.hidden_layer1(inputs))
layer2_out = torch.sigmoid(self.hidden_layer2(layer1_out))
layer3_out = torch.sigmoid(self.hidden_layer3(layer2_out))
layer4_out = torch.sigmoid(self.hidden_layer4(layer3_out))
layer5_out = torch.sigmoid(self.hidden_layer5(layer4_out))
output = self.output_layer(layer5_out) ## For regression, no activation is used in output layer
return output
### (2) Model
net = Net()
net = net.to(device)
mse_cost_function = torch.nn.MSELoss() # Mean squared error
optimizer = torch.optim.Adam(net.parameters())
## PDE as loss function. Thus would use the network which we call as u_theta
def f(x,t, net):
u = net(x,t)
print(u)
# the dependent variable u is given by the network based on independent variables x,t
## Based on our f = du/dx - 2du/dt - u, we need du/dx and du/dt
u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0]
u_t = torch.autograd.grad(u.sum(), t, create_graph=True)[0]
pde = u_x - 2*u_t - u
return pde
## Data from Boundary Conditions
# u(x,0)=6e^(-3x)
## BC just gives us datapoints for training
# BC tells us that for any x in range[0,2] and time=0, the value of u is given by 6e^(-3x)
# Take say 500 random numbers of x
x_bc = np.random.uniform(low=0.0, high=2.0, size=(500,1))
t_bc = np.zeros((500,1))
# compute u based on BC
u_bc = 6*np.exp(-3*x_bc)
### (3) Training / Fitting
iterations = 1000
previous_validation_loss = 99999999.0
for epoch in range(iterations):
optimizer.zero_grad() # to make the gradients zero
# Loss based on boundary conditions
pt_x_bc = Variable(torch.from_numpy(x_bc).float(), requires_grad=False).to(device)
pt_t_bc = Variable(torch.from_numpy(t_bc).float(), requires_grad=False).to(device)
pt_u_bc = Variable(torch.from_numpy(u_bc).float(), requires_grad=False).to(device)
net_bc_out = net(pt_x_bc, pt_t_bc) # output of u(x,t)
mse_u = mse_cost_function(net_bc_out, pt_u_bc)
# Loss based on PDE
x_collocation = np.random.uniform(low=0.0, high=2.0, size=(500, 1))
t_collocation = np.random.uniform(low=0.0, high=1.0, size=(500, 1))
all_zeros = np.zeros((500, 1))
pt_x_collocation = Variable(torch.from_numpy(x_collocation).float(), requires_grad=True).to(device)
pt_t_collocation = Variable(torch.from_numpy(t_collocation).float(), requires_grad=True).to(device)
pt_all_zeros = Variable(torch.from_numpy(all_zeros).float(), requires_grad=False).to(device)
f_out = f(pt_x_collocation, pt_t_collocation, net) # output of f(x,t)
mse_f = mse_cost_function(f_out, pt_all_zeros)
# Combining the loss functions
loss = mse_u + mse_f
loss.backward() # This is for computing gradients using backward propagation
optimizer.step() # This is equivalent to : theta_new = theta_old - alpha * derivative of J w.r.t theta
with torch.autograd.no_grad():
print(epoch, "Traning Loss:", loss.data)
fig = plt.figure()
ax = fig.gca(projection='3d')
x = np.arange(0, 2, 0.02)
t = np.arange(0, 1, 0.02)
ms_x, ms_t = np.meshgrid(x, t)
## Just because meshgrid is used, we need to do the following adjustment
x = np.ravel(ms_x).reshape(-1, 1)
t = np.ravel(ms_t).reshape(-1, 1)
pt_x = Variable(torch.from_numpy(x).float(), requires_grad=True).to(device)
pt_t = Variable(torch.from_numpy(t).float(), requires_grad=True).to(device)
pt_u = net(pt_x, pt_t)
u = pt_u.data.cpu().numpy()
ms_u = u.reshape(ms_x.shape)
surf = ax.plot_surface(ms_x, ms_t, ms_u, cmap=cm.coolwarm, linewidth=0, antialiased=False)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
fig.colorbar(surf, shrink=0.5, aspect=5)
plt.show()
torch.save(net.state_dict(), "model_uxt.pt")