Skip to content

Commit 498087c

Browse files
author
Yuan Yin
committedApr 21, 2022
Add scripts for Darcy Flow.
1 parent a15cf3a commit 498087c

7 files changed

+680
-15
lines changed
 

‎GKN_darcy2d.py

+207
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
import torch.nn.functional as F
2+
3+
from torch_geometric.data import DataLoader
4+
from src.utils.utilities import *
5+
from torch_geometric.nn import NNConv
6+
7+
from timeit import default_timer
8+
9+
10+
class KernelNN3(torch.nn.Module):
11+
def __init__(self, width_node, width_kernel, depth, ker_in, in_width=1, out_width=1):
12+
super(KernelNN3, self).__init__()
13+
self.depth = depth
14+
15+
self.fc1 = torch.nn.Linear(in_width, width_node)
16+
17+
kernel = DenseNet([ker_in, width_kernel // 2, width_kernel, width_node**2], torch.nn.ReLU)
18+
self.conv1 = NNConv(width_node, width_node, kernel, aggr='mean')
19+
20+
self.fc2 = torch.nn.Linear(width_node, 1)
21+
22+
def forward(self, data):
23+
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
24+
x = self.fc1(x)
25+
for k in range(self.depth):
26+
x = self.conv1(x, edge_index, edge_attr)
27+
if k != self.depth - 1:
28+
x = F.relu(x)
29+
30+
x = self.fc2(x)
31+
return x
32+
33+
PATH = '/data/yiny/exp/mgko'
34+
TRAIN_PATH = f'{PATH}/data/piececonst_r241_N1024_smooth1.mat'
35+
TEST_PATH = f'{PATH}/data/piececonst_r241_N1024_smooth2.mat'
36+
37+
ms = [1600]
38+
for case in range(1):
39+
r = 1
40+
s = int(((241 - 1)/r) + 1)
41+
n = s**2
42+
m = ms[case]
43+
k = 1
44+
45+
radius_train = 0.5/8
46+
radius_test = 0.5/8
47+
print('resolution', s)
48+
49+
50+
ntrain = 100
51+
ntest = 100
52+
53+
54+
batch_size = 1
55+
batch_size2 = 1
56+
width = 64
57+
ker_width = 256
58+
depth = 4
59+
edge_features = 6
60+
node_features = 6
61+
62+
epochs = 200
63+
learning_rate = 0.001
64+
scheduler_step = 50
65+
scheduler_gamma = 0.5
66+
67+
68+
path = 'neurips1_GKN_s'+str(s)+'_ntrain'+str(ntrain)+'_kerwidth'+str(ker_width) + '_m0' + str(m)
69+
path_model = f'{PATH}/model/' + path
70+
path_train_err = f'{PATH}/results/' + path + 'train.txt'
71+
path_test_err = f'{PATH}/results/' + path + 'test.txt'
72+
path_runtime = f'{PATH}/results/' + path + 'time.txt'
73+
path_image = f'{PATH}/results/' + path
74+
75+
runtime = np.zeros(2, )
76+
t1 = default_timer()
77+
78+
79+
reader = MatReader(TRAIN_PATH)
80+
train_a = reader.read_field('coeff')[:ntrain,::r,::r].reshape(ntrain,-1)
81+
train_a_smooth = reader.read_field('Kcoeff')[:ntrain,::r,::r].reshape(ntrain,-1)
82+
train_a_gradx = reader.read_field('Kcoeff_x')[:ntrain,::r,::r].reshape(ntrain,-1)
83+
train_a_grady = reader.read_field('Kcoeff_y')[:ntrain,::r,::r].reshape(ntrain,-1)
84+
train_u = reader.read_field('sol')[:ntrain,::r,::r].reshape(ntrain,-1)
85+
86+
reader.load_file(TEST_PATH)
87+
test_a = reader.read_field('coeff')[:ntest,::r,::r].reshape(ntest,-1)
88+
test_a_smooth = reader.read_field('Kcoeff')[:ntest,::r,::r].reshape(ntest,-1)
89+
test_a_gradx = reader.read_field('Kcoeff_x')[:ntest,::r,::r].reshape(ntest,-1)
90+
test_a_grady = reader.read_field('Kcoeff_y')[:ntest,::r,::r].reshape(ntest,-1)
91+
test_u = reader.read_field('sol')[:ntest,::r,::r].reshape(ntest,-1)
92+
93+
94+
a_normalizer = GaussianNormalizer(train_a)
95+
train_a = a_normalizer.encode(train_a)
96+
test_a = a_normalizer.encode(test_a)
97+
as_normalizer = GaussianNormalizer(train_a_smooth)
98+
train_a_smooth = as_normalizer.encode(train_a_smooth)
99+
test_a_smooth = as_normalizer.encode(test_a_smooth)
100+
agx_normalizer = GaussianNormalizer(train_a_gradx)
101+
train_a_gradx = agx_normalizer.encode(train_a_gradx)
102+
test_a_gradx = agx_normalizer.encode(test_a_gradx)
103+
agy_normalizer = GaussianNormalizer(train_a_grady)
104+
train_a_grady = agy_normalizer.encode(train_a_grady)
105+
test_a_grady = agy_normalizer.encode(test_a_grady)
106+
107+
u_normalizer = UnitGaussianNormalizer(train_u)
108+
train_u = u_normalizer.encode(train_u)
109+
# test_u = y_normalizer.encode(test_u)
110+
111+
112+
113+
meshgenerator = RandomMeshGenerator([[0,1],[0,1]],[s,s], sample_size=m)
114+
data_train = []
115+
for j in range(ntrain):
116+
for i in range(k):
117+
idx = meshgenerator.sample()
118+
grid = meshgenerator.get_grid()
119+
edge_index = meshgenerator.ball_connectivity(radius_train)
120+
edge_attr = meshgenerator.attributes(theta=train_a[j,:])
121+
#data_train.append(Data(x=init_point.clone().view(-1,1), y=train_y[j,:], edge_index=edge_index, edge_attr=edge_attr))
122+
data_train.append(Data(x=torch.cat([grid, train_a[j, idx].reshape(-1, 1),
123+
train_a_smooth[j, idx].reshape(-1, 1), train_a_gradx[j, idx].reshape(-1, 1),
124+
train_a_grady[j, idx].reshape(-1, 1)
125+
], dim=1),
126+
y=train_u[j, idx], edge_index=edge_index, edge_attr=edge_attr, sample_idx=idx
127+
))
128+
129+
130+
meshgenerator = RandomMeshGenerator([[0,1],[0,1]],[s,s], sample_size=m)
131+
data_test = []
132+
for j in range(ntest):
133+
idx = meshgenerator.sample()
134+
grid = meshgenerator.get_grid()
135+
edge_index = meshgenerator.ball_connectivity(radius_test)
136+
edge_attr = meshgenerator.attributes(theta=test_a[j,:])
137+
data_test.append(Data(x=torch.cat([grid, test_a[j, idx].reshape(-1, 1),
138+
test_a_smooth[j, idx].reshape(-1, 1), test_a_gradx[j, idx].reshape(-1, 1),
139+
test_a_grady[j, idx].reshape(-1, 1)
140+
], dim=1),
141+
y=test_u[j, idx], edge_index=edge_index, edge_attr=edge_attr, sample_idx=idx
142+
))
143+
#
144+
train_loader = DataLoader(data_train, batch_size=batch_size, shuffle=True)
145+
test_loader = DataLoader(data_test, batch_size=batch_size2, shuffle=False)
146+
147+
t2 = default_timer()
148+
149+
print('preprocessing finished, time used:', t2-t1)
150+
device = torch.device('cuda:3')
151+
152+
model = KernelNN3(width, ker_width,depth,edge_features,in_width=node_features).to(device)
153+
print('model parameters:', get_n_params(model))
154+
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=5e-4)
155+
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma)
156+
157+
myloss = LpLoss(size_average=False)
158+
u_normalizer.to(device)
159+
ttrain = np.zeros((epochs, ))
160+
ttest = np.zeros((epochs,))
161+
model.train()
162+
for ep in range(epochs):
163+
t1 = default_timer()
164+
train_mse = 0.0
165+
train_l2 = 0.0
166+
for batch in train_loader:
167+
batch = batch.to(device)
168+
169+
optimizer.zero_grad()
170+
out = model(batch)
171+
mse = F.mse_loss(out.view(-1, 1), batch.y.view(-1,1))
172+
mse.backward()
173+
174+
l2 = myloss(
175+
u_normalizer.decode(out.view(batch_size, -1), sample_idx=batch.sample_idx.view(batch_size, -1)),
176+
u_normalizer.decode(batch.y.view(batch_size, -1), sample_idx=batch.sample_idx.view(batch_size, -1)))
177+
optimizer.step()
178+
train_mse += mse.item()
179+
train_l2 += l2.item()
180+
181+
scheduler.step()
182+
t2 = default_timer()
183+
184+
model.eval()
185+
test_l2 = 0.0
186+
with torch.no_grad():
187+
for batch in test_loader:
188+
batch = batch.to(device)
189+
out = model(batch)
190+
out = u_normalizer.decode(out.view(batch_size2,-1), sample_idx=batch.sample_idx.view(batch_size2,-1))
191+
test_l2 += myloss(out, batch.y.view(batch_size2, -1)).item()
192+
# test_l2 += myloss(out.view(batch_size2,-1), y_normalizer.encode(batch.y.view(batch_size2, -1))).item()
193+
194+
t3 = default_timer()
195+
ttrain[ep] = train_l2/(ntrain * k)
196+
ttest[ep] = test_l2/ntest
197+
198+
print(k, ntrain, ep, t2-t1, t3-t2, train_mse/len(train_loader), train_l2/(ntrain * k), test_l2/ntest)
199+
200+
runtime[0] = t2-t1
201+
runtime[1] = t3-t2
202+
np.savetxt(path_train_err, ttrain)
203+
np.savetxt(path_test_err, ttest)
204+
np.savetxt(path_runtime, runtime)
205+
torch.save(model, path_model)
206+
207+

‎GKN_orthogonal_burgers1d.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616
#
1717
########################################################################
1818

19-
PATH = ""
20-
TRAIN_PATH = "/data/migus/mgkn_revamped/burgers_data_R10.mat"
21-
TEST_PATH = "/data/migus/mgkn_revamped/burgers_data_R10.mat"
19+
PATH = "."
20+
TRAIN_PATH = f"{PATH}/data/burgers_data_R10.mat"
21+
TEST_PATH = f"{PATH}/data/burgers_data_R10.mat"
2222

2323
r = 16
2424
s = 2 ** 13 // r

‎MGKN_darcy2d.py

+199
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
import torch.nn.functional as F
2+
3+
from torch_geometric.data import DataLoader
4+
from src.utils.utilities import *
5+
6+
from timeit import default_timer
7+
from src.models.model_mgkn import KernelInduced
8+
from src.data.data_processing import convert_data_to_mesh
9+
10+
torch.manual_seed(0)
11+
np.random.seed(0)
12+
13+
14+
PATH = '.'
15+
TRAIN_PATH = f'{PATH}/data/piececonst_r241_N1024_smooth1.mat'
16+
TEST_PATH = f'{PATH}/data/piececonst_r241_N1024_smooth2.mat'
17+
18+
19+
r = 1
20+
s = int(((241 - 1)/r) + 1)
21+
n = s**2
22+
k = 1
23+
24+
# this is too large
25+
# m = [6400, 1600, 400, 100, 25]
26+
# radius_inner = [0.5/16, 0.5/8, 0.5/4, 0.5/2, 0.5]
27+
# radius_inter = [0.5/16 * 1.41, 0.5/8* 1.41, 0.5/4* 1.41, 0.5/2* 1.41]
28+
29+
for case in [1]:
30+
31+
print('!!!!!!!!!!!!!! case ', case, ' !!!!!!!!!!!!!!!!!!!!!!!!')
32+
33+
if case == 0:
34+
m = [1600, 400, 100, 25]
35+
radius_inner = [ 0.5/8, 0.5/4, 0.5/2, 0.5]
36+
radius_inter = [0.5/8* 1.41, 0.5/4* 1.41, 0.5/2* 1.41]
37+
38+
if case == 1:
39+
m = [1600, 400, 100]
40+
radius_inner = [0.5/8, 0.5/4, 0.5/2]
41+
radius_inter = [0.5/8* 1.41, 0.5/4* 1.41]
42+
43+
if case == 2:
44+
m = [1600, 400]
45+
radius_inner = [0.5/8, 0.5/4]
46+
radius_inter = [0.5/8* 1.41]
47+
48+
level = len(m)
49+
print('resolution', s)
50+
51+
ntrain = 100
52+
ntest = 100
53+
54+
# don't change this
55+
batch_size = 1
56+
batch_size2 = 1
57+
58+
width = 64
59+
ker_width = 256
60+
depth = 4
61+
edge_features = 6
62+
node_features = 6
63+
64+
epochs = 200
65+
learning_rate = 0.001
66+
scheduler_step = 10
67+
scheduler_gamma = 0.8
68+
69+
path = f'neurips1_multigraph_s'+str(s)+'_ntrain'+str(ntrain)+'_kerwidth'+str(ker_width) + '_m0' + str(m[0])
70+
path_model = f'{PATH}/model/' + path
71+
path_train_err = f'{PATH}/results/' + path + 'train.txt'
72+
path_test_err = f'{PATH}/results/' + path + 'test.txt'
73+
path_runtime = f'{PATH}/results/' + path + 'time.txt'
74+
path_image = f'{PATH}/results/' + path
75+
76+
runtime = np.zeros(2,)
77+
78+
t1 = default_timer()
79+
80+
81+
reader = MatReader(TRAIN_PATH)
82+
train_a = reader.read_field('coeff')[:ntrain,::r,::r].reshape(ntrain,-1)
83+
train_a_smooth = reader.read_field('Kcoeff')[:ntrain,::r,::r].reshape(ntrain,-1)
84+
train_a_gradx = reader.read_field('Kcoeff_x')[:ntrain,::r,::r].reshape(ntrain,-1)
85+
train_a_grady = reader.read_field('Kcoeff_y')[:ntrain,::r,::r].reshape(ntrain,-1)
86+
train_u = reader.read_field('sol')[:ntrain,::r,::r].reshape(ntrain,-1)
87+
88+
reader.load_file(TEST_PATH)
89+
test_a = reader.read_field('coeff')[:ntest,::r,::r].reshape(ntest,-1)
90+
test_a_smooth = reader.read_field('Kcoeff')[:ntest,::r,::r].reshape(ntest,-1)
91+
test_a_gradx = reader.read_field('Kcoeff_x')[:ntest,::r,::r].reshape(ntest,-1)
92+
test_a_grady = reader.read_field('Kcoeff_y')[:ntest,::r,::r].reshape(ntest,-1)
93+
test_u = reader.read_field('sol')[:ntest,::r,::r].reshape(ntest,-1)
94+
95+
96+
a_normalizer = GaussianNormalizer(train_a)
97+
train_a = a_normalizer.encode(train_a)
98+
test_a = a_normalizer.encode(test_a)
99+
as_normalizer = GaussianNormalizer(train_a_smooth)
100+
train_a_smooth = as_normalizer.encode(train_a_smooth)
101+
test_a_smooth = as_normalizer.encode(test_a_smooth)
102+
agx_normalizer = GaussianNormalizer(train_a_gradx)
103+
train_a_gradx = agx_normalizer.encode(train_a_gradx)
104+
test_a_gradx = agx_normalizer.encode(test_a_gradx)
105+
agy_normalizer = GaussianNormalizer(train_a_grady)
106+
train_a_grady = agy_normalizer.encode(train_a_grady)
107+
test_a_grady = agy_normalizer.encode(test_a_grady)
108+
109+
u_normalizer = UnitGaussianNormalizer(train_u)
110+
train_u = u_normalizer.encode(train_u)
111+
# test_u = y_normalizer.encode(test_u)
112+
113+
114+
data_train = convert_data_to_mesh([train_a, train_a_smooth, train_a_gradx, train_a_grady], train_u, ntrain, k,
115+
s, m, radius_inner, radius_inter)
116+
data_test = convert_data_to_mesh([test_a, test_a_smooth, test_a_gradx, test_a_grady], test_u, ntest, k, s,
117+
m, radius_inner, radius_inter)
118+
119+
train_loader = DataLoader(data_train, batch_size=batch_size, shuffle=True)
120+
test_loader = DataLoader(data_test, batch_size=batch_size2, shuffle=False)
121+
122+
t2 = default_timer()
123+
124+
print('preprocessing finished, time used:', t2-t1)
125+
device = torch.device('cuda:0')
126+
127+
cycle_type = "f"
128+
print('cycle type:', cycle_type)
129+
in_cycle_shared = True
130+
shared = False
131+
skip_connection = True
132+
print('in cycle shared:', in_cycle_shared, 'iteration shared:', shared)
133+
134+
135+
model = KernelInduced(width=width, ker_width=ker_width, depth=depth, ker_in=edge_features,
136+
points=m, level=level, in_width=node_features, out_width=1, cycle_type=cycle_type, shared=shared, skip_connection=skip_connection, in_cycle_shared=in_cycle_shared).to(device)
137+
init_weights(model, init_type='orthogonal', init_gain=1.414)
138+
139+
print('model parameters:', get_n_params(model))
140+
141+
142+
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=5e-4)
143+
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=scheduler_step, gamma=scheduler_gamma)
144+
145+
myloss = LpLoss(size_average=False)
146+
u_normalizer.to(device)
147+
ttrain = np.zeros((epochs,))
148+
ttest = np.zeros((epochs,))
149+
model.train()
150+
for ep in range(epochs):
151+
t1 = default_timer()
152+
train_mse = 0.0
153+
train_l2 = 0.0
154+
for batch in train_loader:
155+
batch = batch.to(device)
156+
157+
optimizer.zero_grad()
158+
out = model(batch)
159+
mse = F.mse_loss(out.view(-1, 1), batch.y.view(-1,1))
160+
# mse.backward()
161+
162+
l2 = myloss(
163+
u_normalizer.decode(out.view(batch_size, -1), sample_idx=batch.sample_idx.view(batch_size, -1)),
164+
u_normalizer.decode(batch.y.view(batch_size, -1), sample_idx=batch.sample_idx.view(batch_size, -1)))
165+
l2.backward()
166+
167+
optimizer.step()
168+
train_mse += mse.item()
169+
train_l2 += l2.item()
170+
171+
scheduler.step()
172+
t2 = default_timer()
173+
ttrain[ep] = train_l2 / (ntrain * k)
174+
175+
print(ep, t2 - t1, train_mse / len(train_loader), train_l2 / (ntrain * k))
176+
177+
runtime[0] = t2 - t1
178+
179+
t1 = default_timer()
180+
model.eval()
181+
test_l2 = 0.0
182+
with torch.no_grad():
183+
for batch in test_loader:
184+
batch = batch.to(device)
185+
out = model(batch)
186+
out = u_normalizer.decode(out.view(batch_size2, -1), sample_idx=batch.sample_idx.view(batch_size2, -1))
187+
test_l2 += myloss(out, batch.y.view(batch_size2, -1)).item()
188+
189+
ttest[ep] = test_l2 / ntest
190+
t2 = default_timer()
191+
print(ep, t2 - t1, test_l2 / ntest)
192+
193+
runtime[1] = t2 - t1
194+
195+
np.savetxt(path_train_err, ttrain)
196+
np.savetxt(path_test_err, ttest)
197+
np.savetxt(path_runtime, runtime)
198+
torch.save(model, path_model)
199+

‎MGKN_orthogonal_burgers1d.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616
#
1717
########################################################################
1818

19-
PATH = ""
20-
TRAIN_PATH = "/data/migus/mgkn_revamped/burgers_data_R10.mat"
21-
TEST_PATH = "/data/migus/mgkn_revamped/burgers_data_R10.mat"
19+
PATH = "."
20+
TRAIN_PATH = f"{PATH}/data/burgers_data_R10.mat"
21+
TEST_PATH = f"{PATH}/data/burgers_data_R10.mat"
2222

2323
r = 8
2424
s = 2 ** 13 // r

‎MLP_GCN_darcy2d.py

+248
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
1+
import torch.nn.functional as F
2+
3+
from torch_geometric.data import DataLoader
4+
from src.utils.utilities import *
5+
from torch_geometric.nn import GCNConv, MLP
6+
from torch_cluster import radius_graph
7+
8+
from timeit import default_timer
9+
10+
torch.manual_seed(0)
11+
np.random.seed(0)
12+
13+
print(torch.cuda.is_available())
14+
15+
16+
class GCN_Net(torch.nn.Module):
17+
def __init__(self, width, ker_width, depth, in_width=1, out_width=1, net_type='gcn'):
18+
super(GCN_Net, self).__init__()
19+
self.depth = depth
20+
self.width = width
21+
22+
self.fc_in = torch.nn.Linear(in_width, width)
23+
self.net_type = net_type
24+
25+
if net_type == 'gcn':
26+
self.conv1 = GCNConv(width, width)
27+
self.conv2 = GCNConv(width, width)
28+
self.conv3 = GCNConv(width, width)
29+
self.conv4 = GCNConv(width, width)
30+
elif net_type == 'mlp':
31+
self.conv1 = torch.nn.Linear(width, width)
32+
self.conv2 = torch.nn.Linear(width, width)
33+
self.conv3 = torch.nn.Linear(width, width)
34+
self.conv4 = torch.nn.Linear(width, width)
35+
36+
37+
self.fc_out1 = torch.nn.Linear(width, ker_width)
38+
self.fc_out2 = torch.nn.Linear(ker_width, 1)
39+
40+
41+
def forward(self, data):
42+
x, X, edge_index = data.x, data.pos, data.edge_index
43+
# if self.net_type == 'gcn':
44+
x = torch.cat([X, x], dim=1)
45+
# print(X.device)
46+
47+
x = self.fc_in(x)
48+
49+
for t in range(self.depth):
50+
if self.net_type == 'gcn':
51+
x = x + self.conv1(x, edge_index)
52+
x = F.relu(x)
53+
x = x + self.conv2(x, edge_index)
54+
x = F.relu(x)
55+
x = x + self.conv3(x, edge_index)
56+
x = F.relu(x)
57+
x = x + self.conv4(x, edge_index)
58+
x = F.relu(x)
59+
elif self.net_type == 'mlp':
60+
x = self.conv1(x)
61+
x = F.relu(x)
62+
x = self.conv2(x)
63+
x = F.relu(x)
64+
x = self.conv3(x)
65+
x = F.relu(x)
66+
x = self.conv4(x)
67+
x = F.relu(x)
68+
x = F.leaky_relu(self.fc_out1(x))
69+
x = self.fc_out2(x)
70+
return x
71+
72+
73+
PATH = '.'
74+
TRAIN_PATH = f'{PATH}/data/piececonst_r241_N1024_smooth1.mat'
75+
TEST_PATH = f'{PATH}/data/piececonst_r241_N1024_smooth2.mat'
76+
77+
78+
r = 1
79+
s = int(((241 - 1)/r) + 1)
80+
n = s**2
81+
k = 1
82+
83+
net_type = 'gcn'
84+
print(f'net_type: {net_type}')
85+
86+
print('resolution', s)
87+
88+
ntrain = 100
89+
ntest = 100
90+
91+
batch_size = 1
92+
batch_size2 = 1
93+
94+
gpu_id = 2
95+
96+
width = 64
97+
ker_width = 256
98+
depth = 4
99+
100+
node_features = 6
101+
102+
epochs = 200
103+
learning_rate = 1e-3
104+
scheduler_step = 10
105+
scheduler_gamma = 0.85
106+
107+
108+
109+
path = f'neurips4_GCN_s'+str(s)+'_ntrain'+str(ntrain)+'_kerwidth'+str(ker_width)
110+
path_model = f'{PATH}/model/' + path
111+
path_train_err = f'{PATH}/results/' + path + 'train.txt'
112+
path_test_err = f'{PATH}/results/' + path + 'test.txt'
113+
path_image = f'{PATH}/results/' + path
114+
115+
116+
t1 = default_timer()
117+
118+
119+
reader = MatReader(TRAIN_PATH)
120+
train_a = reader.read_field('coeff')[:ntrain,::r,::r].reshape(ntrain,-1)
121+
train_a_smooth = reader.read_field('Kcoeff')[:ntrain,::r,::r].reshape(ntrain,-1)
122+
train_a_gradx = reader.read_field('Kcoeff_x')[:ntrain,::r,::r].reshape(ntrain,-1)
123+
train_a_grady = reader.read_field('Kcoeff_y')[:ntrain,::r,::r].reshape(ntrain,-1)
124+
train_u = reader.read_field('sol')[:ntrain,::r,::r].reshape(ntrain,-1)
125+
126+
reader.load_file(TEST_PATH)
127+
test_a = reader.read_field('coeff')[:ntest,::r,::r].reshape(ntest,-1)
128+
test_a_smooth = reader.read_field('Kcoeff')[:ntest,::r,::r].reshape(ntest,-1)
129+
test_a_gradx = reader.read_field('Kcoeff_x')[:ntest,::r,::r].reshape(ntest,-1)
130+
test_a_grady = reader.read_field('Kcoeff_y')[:ntest,::r,::r].reshape(ntest,-1)
131+
test_u = reader.read_field('sol')[:ntest,::r,::r].reshape(ntest,-1)
132+
133+
134+
a_normalizer = GaussianNormalizer(train_a)
135+
train_a = a_normalizer.encode(train_a)
136+
test_a = a_normalizer.encode(test_a)
137+
as_normalizer = GaussianNormalizer(train_a_smooth)
138+
train_a_smooth = as_normalizer.encode(train_a_smooth)
139+
test_a_smooth = as_normalizer.encode(test_a_smooth)
140+
agx_normalizer = GaussianNormalizer(train_a_gradx)
141+
train_a_gradx = agx_normalizer.encode(train_a_gradx)
142+
test_a_gradx = agx_normalizer.encode(test_a_gradx)
143+
agy_normalizer = GaussianNormalizer(train_a_grady)
144+
train_a_grady = agy_normalizer.encode(train_a_grady)
145+
test_a_grady = agy_normalizer.encode(test_a_grady)
146+
147+
u_normalizer = UnitGaussianNormalizer(train_u)
148+
train_u = u_normalizer.encode(train_u)
149+
# test_u = y_normalizer.encode(test_u)
150+
151+
X, edge_index, _ = grid_edge(s, s)
152+
153+
data_train = []
154+
for j in range(ntrain):
155+
for i in range(k):
156+
x = torch.cat([train_a[j].reshape(-1, 1),
157+
train_a_smooth[j].reshape(-1, 1),
158+
train_a_gradx[j].reshape(-1, 1),
159+
train_a_grady[j].reshape(-1, 1)
160+
], dim=1)
161+
data_train.append(Data(x=x, pos=X, y=train_u[j], edge_index=radius_graph(X, r=0.5/8)))
162+
163+
print(x.shape)
164+
print(edge_index.shape)
165+
166+
167+
168+
data_test = []
169+
for j in range(ntest):
170+
x = torch.cat([test_a[j].reshape(-1, 1),
171+
test_a_smooth[j].reshape(-1, 1),
172+
test_a_gradx[j].reshape(-1, 1),
173+
test_a_grady[j].reshape(-1, 1)
174+
], dim=1)
175+
data_test.append(Data(x=x, pos=X, y=test_u[j], edge_index=radius_graph(X, r=0.5/8)))
176+
177+
print(x.shape)
178+
print(edge_index.shape)
179+
train_loader = DataLoader(data_train, batch_size=batch_size, shuffle=True)
180+
test_loader = DataLoader(data_test, batch_size=batch_size2, shuffle=False)
181+
t2 = default_timer()
182+
183+
print('preprocessing finished, time used:', t2-t1)
184+
device = torch.device(f'cuda:{gpu_id}')
185+
186+
model = GCN_Net(width=width, ker_width=ker_width, depth=depth, in_width=node_features, out_width=1, net_type=net_type).to(device)
187+
init_weights(model, init_type='normal', init_gain=0.1)
188+
189+
print('model parameters:', get_n_params(model))
190+
191+
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
192+
193+
myloss = LpLoss(size_average=False)
194+
ttrain = np.zeros((epochs, ))
195+
ttest = np.zeros((epochs,))
196+
197+
for ep in range(epochs):
198+
t1 = default_timer()
199+
train_mse = 0.0
200+
train_l2 = 0.0
201+
model.train()
202+
u_normalizer.to(device)
203+
for batch in train_loader:
204+
batch = batch.to(device)
205+
206+
optimizer.zero_grad()
207+
out = model(batch)
208+
mse = F.mse_loss(out.view(-1, 1), batch.y.view(-1,1))
209+
# mse.backward()
210+
211+
l2 = myloss(
212+
u_normalizer.decode(out.view(batch_size, -1)),
213+
u_normalizer.decode(batch.y.view(batch_size, -1)))
214+
l2.backward()
215+
216+
optimizer.step()
217+
train_mse += mse.item()
218+
train_l2 += l2.item()
219+
220+
# scheduler.step()
221+
t2 = default_timer()
222+
ttrain[ep] = train_l2 / (ntrain * k)
223+
224+
print(ep, t2 - t1, train_mse / len(train_loader), train_l2 / (ntrain * k))
225+
226+
if ep % 10 == 0:
227+
model.eval()
228+
test_l2 = 0.0
229+
with torch.no_grad():
230+
t1 = default_timer()
231+
for batch in test_loader:
232+
batch = batch.to(device)
233+
out = model(batch)
234+
l2 = myloss(u_normalizer.decode(out.view(batch_size2, -1)), batch.y.view(batch_size2, -1))
235+
test_l2 += l2.item()
236+
237+
238+
ttest[ep] = test_l2 / ntest
239+
t2 = default_timer()
240+
print(ep, t2 - t1, test_l2 / (ntest*k) )
241+
torch.save(model, path_model+str(ep))
242+
243+
244+
245+
np.savetxt(path_train_err, ttrain)
246+
np.savetxt(path_test_err, ttest)
247+
248+

‎MLP_GCN_orthogonal_burgers1d.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
#
1616
########################################################################
1717

18-
PATH = ""
19-
TRAIN_PATH = "/data/migus/mgkn_revamped/burgers_data_R10.mat"
20-
TEST_PATH = "/data/migus/mgkn_revamped/burgers_data_R10.mat"
18+
PATH = "."
19+
TRAIN_PATH = f"{PATH}/data/burgers_data_R10.mat"
20+
TEST_PATH = f"{PATH}/data/burgers_data_R10.mat"
2121

2222
r = 8
2323
s = 2 ** 13 // r

‎README.md

+17-6
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,37 @@
11
# Multi-scale Physical Representations for Approximating PDE Solutions with Graph Neural Operators
22

33
This repository contains the code for the paper:
4-
- [Multi-scale Physical Representations for Approximating PDE Solutions with Graph Neural Operators](https://openreview.net/forum?id=rx9TVZJax5)
4+
5+
- [Multi-scale Physical Representations for Approximating PDE Solutions with Graph Neural Operators](https://openreview.net/forum?id=rx9TVZJax5)
56

67
In this work, we build on Fourier Neural Operator (FNO) and numerical schemes to create a novel architecture to tackle multi-scale physical representations.
78

89
## Requirements
10+
911
- [PyTorch](https://pytorch.org/)
10-
- [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/)
12+
- [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/)
1113

1214
## Files
15+
1316
- MGKN_orthogonal_burgers1d.py is the Multipole Graph Neural Operator (MGKN) file for the 1D Burgers equation.
1417
- GKN_orthogonal_burgers1d.py is the Graph Kernel Network (GKN) file for the 1D Burgers equation.
1518
- MLP_GCN_orthogonal_burgers1d.py is the multilayer perceptron (MLP) and graph convolutional network (GCN) file for the 1D Burgers equation.
16-
- src contains the code for the graph data creation, the different models and some utilities files from the original FNO code.
19+
- src contains the code for the graph data creation, the different models and some utilities files from the original FNO code.
20+
1721
## Datasets
18-
We used the Burgers equation and Darcy flow equation datasets from the FNO paper. The data generation configuration can be found in their paper.
22+
23+
We used the Burgers equation and Darcy flow equation datasets from the GNO paper. The data generation configuration can be found in their paper.
24+
1925
- [PDE datasets](https://drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-ZaXrKUZhp7pjt-)
2026

21-
Their [github](https://github.com/zongyi-li/fourier_neural_operator) provides addtional information.
27+
See [github](https://github.com/zongyi-li/fourier_neural_operator) for addtional information.
2228

2329
### Usage
24-
```
30+
31+
```bash
2532
python3 MGKN_orthogonal_burgers1d.py
2633
```
34+
35+
### Acknowledgement
36+
37+
This code is adapted from [this repository](https://github.com/zongyi-li/graph-pde) by Zongyi Li.

0 commit comments

Comments
 (0)
Please sign in to comment.