神经网络解热传导方程

1
2
3
4
5
6
7
8
9
10
11

# import matplotlib.pyplot as plt
# import numpy as np
# from 

# import torch.nn.functional as F
# import torch.optim as optim
# import matplotlib.pyplot as plt
# from mpl_toolkits.mplot3d import Axes3D
# from scipy.stats import norm
# from matplotlib import cm
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import torch 
import torch.nn as nn

class Net(nn.Module):
    # NL: the number of hidden layers
    # NN: the number of vertices in each layer
    def __init__(self, NL, NN):
        super(Net, self).__init__()
        self.input_layer = nn.Linear(3, NN)
        self.hidden_layers = nn.ModuleList([nn.Linear(NN, NN) for _ in range(NL)])
        self.output_layer = nn.Linear(NN, 1)
        self.act = nn.Sigmoid()
        
    def forward(self, x):
        o = self.act(self.input_layer(x))
        for _, hl in enumerate(self.hidden_layers):
            o = self.act(hl(o))
        out = self.output_layer(o)
        return out

\begin{cases}u_t-u_{yy}-u_{yy}=(\pi-2\pi^2)(\cos(\pi t)+\sin\pi t))\sin(\pi x)\sin(\pi y),\text{on}[0,4]\times[0,1]\times[0,1] \\ u(0,x,y)=0,\text{on}[0,1]\times[0,1]\\ u(t,0,y)=u(t,1,y)=0,\text{on}[0,4]\times[0,1],\\ u(t,x,0)=u(t,1,x,1)=0,\text{on}[0,4]\times[0,1],\\ u(t,x,0)=u(t,x,1)=0,\text{on}[0,4]\times[0,1],\end{cases}

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
class Heat():
    def __init__(self, net, te, xe, ye):
        self.net = net
        self.te = te
        self.xe = xe
        self.ye = ye

    def sample(self, size=2**8):
        te = self.te
        xe = self.xe
        ye = self.ye

        x = torch.cat((torch.rand([size, 1]) * te, 
                       torch.rand([size, 1]) * xe, 
                       torch.rand([size, 1]) * ye), dim=1)
        
        x_initial = torch.cat((torch.zeros(size, 1), 
                               torch.rand([size, 1]) * xe, 
                               torch.rand([size, 1]) * ye), dim=1)
        
        x_boundary_left = torch.cat((torch.rand([size, 1]) * te, 
                                     torch.zeros([size, 1]), 
                                     torch.rand(size, 1) * ye), dim=1)
        
        x_boundary_right = torch.cat((torch.rand([size, 1]) * te, 
                                      torch.ones([size, 1]), 
                                      torch.rand(size, 1) * ye), dim=1)
        
        x_boundary_up = torch.cat((torch.rand([size, 1]) * te, 
                                   torch.rand([size, 1]) * xe, 
                                   torch.ones(size, 1)), dim=1)
        
        x_boundary_down = torch.cat((torch.rand([size, 1]) * te, 
                                     torch.rand([size, 1]) * xe, 
                                     torch.zeros(size, 1)), dim=1)
        
        return x, x_initial, x_boundary_left, x_boundary_right, x_boundary_up, x_boundary_down


    def loss_func(self, size=2**8):

        x, x_initial, x_boundary_left, x_boundary_right, x_boundary_up, x_boundary_down = self.sample(size=size)
        x = torch.autograd.Variable(x, requires_grad=True)

        d = torch.autograd.grad(self.net(x), x, grad_outputs=torch.ones_like(self.net(x)), create_graph=True)
        dt = d[0][:, 0].reshape(-1, 1)
        dx = d[0][:, 1].reshape(-1, 1)
        dy = d[0][:, 2].reshape(-1, 1)
        dxx = torch.autograd.grad(dx, x, grad_outputs=torch.ones_like(dx), create_graph=True)[0][:, 1].reshape(-1, 1)
        dyy = torch.autograd.grad(dy, x, grad_outputs=torch.ones_like(dy), create_graph=True)[0][:, 2].reshape(-1, 1)

        f = torch.pi * (torch.cos(torch.pi*x[:, 0])) * (torch.sin(torch.pi*x[:, 1])) * (torch.sin(torch.pi*x[:, 2]))\
             + 2 * torch.pi * torch.pi * (torch.sin(torch.pi*x[:, 0])) * (torch.sin(torch.pi*x[:, 1])) * (torch.sin(torch.pi*x[:, 2]))

        diff_error = (dt - dxx - dyy - f.reshape(-1, 1))**2
        # initial condition
        init_error = (self.net(x_initial)) ** 2
        # boundary condition
        bd_left_error = (self.net(x_boundary_left)) ** 2
        bd_right_error = (self.net(x_boundary_right)) ** 2
        bd_up_error = (self.net(x_boundary_up)) ** 2
        bd_down_error = (self.net(x_boundary_down)) ** 2

        return torch.mean(diff_error + init_error + bd_left_error + bd_right_error + bd_up_error + bd_down_error)
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
class Train():
    def __init__(self, net, heateq, BATCH_SIZE):
        self.errors = []
        self.BATCH_SIZE = BATCH_SIZE
        self.net = net
        self.model = heateq

    def train(self, epoch, lr):
        optimizer = torch.optim.Adam(self.net.parameters(), lr)
        avg_loss = 0
        for e in range(epoch):
            optimizer.zero_grad()
            loss = self.model.loss_func(self.BATCH_SIZE)
            avg_loss = avg_loss + float(loss.item())
            loss.backward()
            optimizer.step()
            if e % 50 == 49:
                loss = avg_loss/50
                print("Epoch {} - lr {} -  loss: {}".format(e, lr, loss))
                avg_loss = 0

                error = self.model.loss_func(2**8)
                self.errors.append(error.detach())

    def get_errors(self):
        return self.errors
1
2
3
4
5
6
7
8
9
10
11
net = Net(NL=2, NN=20)
te = 4
xe = 1
ye = 1
heatequation = Heat(net, te, xe, ye)
train = Train(net, heatequation, BATCH_SIZE=2**8)
train.train(epoch=10**5, lr=0.0001)
torch.save(net, 'net_model.pkl')
errors = train.get_errors()


Epoch 49 - lr 0.0001 -  loss: 49.82045875549316
Epoch 99 - lr 0.0001 -  loss: 50.029335556030276
Epoch 149 - lr 0.0001 -  loss: 50.067277603149414
...
Epoch 99949 - lr 0.0001 -  loss: 4.609520344734192
Epoch 99999 - lr 0.0001 -  loss: 4.6677580118179325

1
2
3
4
x_test = torch.cat((torch.full([size, 1], 0.4), torch.linspace(0,2,2000).unsqueeze(-1)), dim=1)
y_test = net(x_test)
from matplotlib import pyplot as plt
plt.plot(torch.linspace(0,2,2000).numpy(), y_test[:,0].detach().numpy())