-
Notifications
You must be signed in to change notification settings - Fork 7
/
util.py
89 lines (67 loc) · 3.59 KB
/
util.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
import numpy as np
import torch.nn as nn
import copy
def get_data(x_range, y_range, x_num, y_num):
x = np.linspace(x_range[0], x_range[1], x_num)
t = np.linspace(y_range[0], y_range[1], y_num)
x_mesh, t_mesh = np.meshgrid(x, t)
data = np.concatenate((np.expand_dims(x_mesh, -1), np.expand_dims(t_mesh, -1)), axis=-1)
b_left = data[0, :, :]
b_right = data[-1, :, :]
b_upper = data[:, -1, :]
b_lower = data[:, 0, :]
res = data.reshape(-1, 2)
return res, b_left, b_right, b_upper, b_lower
def get_n_params(model):
pp = 0
for p in list(model.parameters()):
nn = 1
for s in list(p.size()):
nn = nn * s
pp += nn
return pp
def make_time_sequence(src, num_step=5, step=1e-4):
dim = num_step
src = np.repeat(np.expand_dims(src, axis=1), dim, axis=1) # (N, L, 2)
for i in range(num_step):
src[:, i, -1] += step * i
return src
def make_space_time_sequence(src, space_num_step=5, space_step=1e-4, time_num_step=5, time_step=1e-4):
dim = space_num_step * time_num_step
src = np.repeat(np.expand_dims(src, axis=1), dim, axis=1) # (N, L, 2)
for i in range(time_num_step):
for j in range(space_num_step):
src[:, i * space_num_step + j, -1] += time_step * i
src[:, i * space_num_step + j, 0] += space_step * (j - space_num_step // 2)
return src
def get_clones(module, N):
return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
def get_data_3d(x_range, y_range, t_range, x_num, y_num, t_num):
step_x = (x_range[1] - x_range[0]) / float(x_num - 1)
step_y = (y_range[1] - y_range[0]) / float(y_num - 1)
step_t = (t_range[1] - t_range[0]) / float(t_num - 1)
x_mesh, y_mesh, t_mesh = np.mgrid[x_range[0]:x_range[1] + step_x:step_x, y_range[0]:y_range[1] + step_y:step_y,
t_range[0]:t_range[1] + step_t:step_t]
data = np.concatenate((np.expand_dims(x_mesh, -1), np.expand_dims(y_mesh, -1), np.expand_dims(t_mesh, -1)), axis=-1)
res = data.reshape(-1, 3)
x_mesh, y_mesh, t_mesh = np.mgrid[x_range[0]:x_range[0] + step_x:step_x, y_range[0]:y_range[1] + step_y:step_y,
t_range[0]:t_range[1] + step_t:step_t]
b_left = np.squeeze(
np.concatenate((np.expand_dims(x_mesh, -1), np.expand_dims(y_mesh, -1), np.expand_dims(t_mesh, -1)), axis=-1))[
1:-1].reshape(-1, 3)
x_mesh, y_mesh, t_mesh = np.mgrid[x_range[1]:x_range[1] + step_x:step_x, y_range[0]:y_range[1] + step_y:step_y,
t_range[0]:t_range[1] + step_t:step_t]
b_right = np.squeeze(
np.concatenate((np.expand_dims(x_mesh, -1), np.expand_dims(y_mesh, -1), np.expand_dims(t_mesh, -1)), axis=-1))[
1:-1].reshape(-1, 3)
x_mesh, y_mesh, t_mesh = np.mgrid[x_range[0]:x_range[1] + step_x:step_x, y_range[0]:y_range[0] + step_y:step_y,
t_range[0]:t_range[1] + step_t:step_t]
b_lower = np.squeeze(
np.concatenate((np.expand_dims(x_mesh, -1), np.expand_dims(y_mesh, -1), np.expand_dims(t_mesh, -1)), axis=-1))[
1:-1].reshape(-1, 3)
x_mesh, y_mesh, t_mesh = np.mgrid[x_range[0]:x_range[1] + step_x:step_x, y_range[1]:y_range[1] + step_y:step_y,
t_range[0]:t_range[1] + step_t:step_t]
b_upper = np.squeeze(
np.concatenate((np.expand_dims(x_mesh, -1), np.expand_dims(y_mesh, -1), np.expand_dims(t_mesh, -1)), axis=-1))[
1:-1].reshape(-1, 3)
return res, b_left, b_right, b_upper, b_lower