-
Notifications
You must be signed in to change notification settings - Fork 8
/
ns_uno2d_main.py
65 lines (57 loc) · 2.79 KB
/
ns_uno2d_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
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from data_load_navier_stocks import *
import matplotlib.pyplot as plt
from navier_stokes_uno2d import UNO, UNO_P
import operator
import random
from functools import reduce
from functools import partial
from ns_train_2d import train_model
from timeit import default_timer
from utilities3 import *
from Adam import Adam
from torchsummary import summary
import gc
import math
#hyper Parameters
# Learning rate = 0.001
# Weight deacy = 1e-5
S = 64 # resolution SxS
T_in = 10 # input time interval (0 - T_in)
T_f = 40 # output time interval (T_in - = T_in+T_f)
ntrain = 4000 # number of training instances
ntest = 500 # number of test instances
nval = 500 # number of validation instances
batch_size = 16
width = 32 # Uplifting dimesion
inwidth = 14 # dimension of UNO input ( 10 time step + position embedding )
epochs = 500
# Following code load data from two separate files containing Navier-Stokes equation simulation
train_a_1, train_u_1, test_a_1, test_u_1 = load_NS_("path to navier stokes simulation with viscosity 1e-3 with 1200 instances"\
,2000,500,Sample_num = 2500,T_in=T_in, T = T_f, size = S)
train_a_2, train_u_2, test_a_2, test_u_2 = load_NS_("path to navier stokes simulation with viscosity 1e-3 with 1200 instances"\
,2000 ,500,Sample_num = 2500,T_in=T_in, T = T_f, size=S)
a = torch.cat([train_a_1,train_a_2,test_a_1,test_a_2], dim = 0)
u = torch.cat([train_u_1,train_u_2,test_u_1,test_u_2],dim = 0)
indexs = [i for i in range(a.shape[0])]
random.shuffle(indexs)
train_a,val_a,test_a = a[indexs[:ntrain]], a[indexs[ntrain:ntrain+nval]],a[indexs[ntrain+nval:]]
train_u,val_u,test_u = u[indexs[:ntrain]],u[indexs[ntrain:ntrain+nval]],u[indexs[ntrain+nval:]]
train_a = train_a.reshape(ntrain,S,S,T_in)
val_a = val_a.reshape(nval,S,S,T_in)
test_a = test_a.reshape(ntest,S,S,T_in)
gc.collect()
train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(train_a, train_u),\
batch_size=batch_size, shuffle=True)
val_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(val_a, val_u),\
batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(test_a, test_u), \
batch_size=batch_size, shuffle=False)
model = UNO(inwidth,width)
summary(model, (64, 64,10))
train_model(model,train_loader,val_loader,test_loader, ntrain,nval,ntest,\
weight_path = 'UNO-10e3.pt',T_f=T_f,batch_size=batch_size,epochs=epochs,learning_rate= 0.001,scheduler_step= 100,scheduler_gamma= 0.5,weight_decay = 1e-5)