forked from imanchester/RobustRNN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_generalization.py
143 lines (116 loc) · 5.07 KB
/
test_generalization.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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import torch
import numpy as np
# import data.load_data as load_data
from scipy.optimize import minimize, NonlinearConstraint
import multiprocessing
import scipy.io as io
import argparse
import os
import models.rnn as rnn
import models.lstm as lstm
import models.RobustRnn as RobustRnn
import models.ciRNN as ciRNN
import models.dnb as dnb
import data.n_linked_msd as msd
torch.set_default_dtype(torch.float64)
torch.set_printoptions(precision=4)
multiprocessing.set_start_method('spawn', True)
def test_performance(model, period, sd, samples=100, sim_len=1000, wash_per=200):
model.eval()
sim = msd.msd_chain(N=4, T=sim_len, u_sd=sd, period=period, Ts=0.5, batchsize=samples)
loader = sim.sim_rand_ic(sim_len, samples, mini_batch_size=samples)
for idx, U, Y in loader:
with torch.no_grad():
Yest = model(U)
SE = (Y[:, :, wash_per:] - Yest[:, :, wash_per:]).norm(dim=(1, 2)).detach()
NSE = ((Y[:, :, wash_per:] - Yest[:, :, wash_per:]).norm(dim=(1, 2)) / Y[:, :, wash_per:].norm(dim=(1, 2))).detach()
res = {"NSE": NSE, "SE": SE}
return res
if __name__ == "__main__":
N = 4
train_seq_len = 1000
training_batches = 100
mini_batch_size = 1
nu = 1
ny = 1
width = 10
batches = training_batches
path = './results/msd/'
if not os.path.exists(path + 'lip/'):
os.mkdir(path + 'lip/')
os.mkdir(path + 'wcg/')
os.mkdir(path + 'wcp/')
print("Directory ", path, "lip/ Created")
def vary_amplitude(model):
print("\t Testing amplitude")
samples = 300
test_points = 21
period = 100
NSE_dist = np.zeros((test_points, samples))
SE_dist = np.zeros((test_points, samples))
amps = np.linspace(0.5, 10.5, test_points)
for idx in range(test_points):
res = test_performance(model, period, amps[idx], samples=samples)
NSE_dist[idx, :] = res["NSE"]
SE_dist[idx, :] = res["SE"]
return {"amps": amps, "NSE": NSE_dist, "SE": SE_dist, "period": period}
# Test generalization of Robust RNNs
print("Running tests on robust-RNN")
name = 'iqc-rnn_w10_gamma0.0_n4'
model = RobustRnn.RobustRnn(nu, width, ny, width, nBatches=batches, method='Neuron')
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
name = 'iqc-rnn_w10_gamma3.0_n4'
model = RobustRnn.RobustRnn(nu, width, ny, width, nBatches=batches, method='Neuron')
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
name = 'iqc-rnn_w10_gamma5.0_n4'
model = RobustRnn.RobustRnn(nu, width, ny, width, nBatches=batches, method='Neuron')
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
name = 'iqc-rnn_w10_gamma8.0_n4'
model = RobustRnn.RobustRnn(nu, width, ny, width, nBatches=batches, method='Neuron')
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
name = 'iqc-rnn_w10_gamma10.0_n4'
model = RobustRnn.RobustRnn(nu, width, ny, width, nBatches=batches, method='Neuron')
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
name = 'iqc-rnn_w10_gamma15.0_n4'
model = RobustRnn.RobustRnn(nu, width, ny, width, nBatches=batches, method='Neuron')
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
# # lstm
print("Running tests on LSTM")
name = 'lstm_w10_gamma0.0_n4'
model = lstm.lstm(nu, width, ny, layers=1, nBatches=batches)
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
# rnn
print("Running tests on RNN")
name = 'rnn_w10_gamma0.0_n4'
model = rnn.rnn(nu, width, ny, 1, nBatches=batches)
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
# cirnn
print("Running tests on cirnn")
name = 'cirnn_w10_gamma0.0_n4'
model = ciRNN.ciRNN(nu, width, ny, 1, nBatches=100)
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)
# srnn
print("Running tests on srnn")
name = 'dnb_w10_gamma0.0_n4'
model = dnb.dnbRNN(nu, width, ny, layers=1, nBatches=batches)
model.load_state_dict(torch.load(path + name + ".params"))
res = vary_amplitude(model)
io.savemat('./results/msd/generalization/amp_' + name + '.mat', res)