-
Notifications
You must be signed in to change notification settings - Fork 20
/
Copy pathutil.py
65 lines (46 loc) · 2.81 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
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 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