forked from PlasmaControl/PlasmaEvolution
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathprediction_helpers.py
221 lines (210 loc) · 12.2 KB
/
prediction_helpers.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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
import configparser
import torch
#torch.manual_seed(0)
import dataSettings
import numpy as np
import os
from train_helpers import make_bucket
from torch.nn.utils.rnn import pad_sequence, unpad_sequence
import re
import glob
from customModels import IanRNN, IanMLP, HiroLRAN
from dataSettings import get_denormalized_dic,normalizations
from customDatasetMakers import state_to_dic, dic_to_state
import time
models={'IanRNN': IanRNN, 'IanMLP': IanMLP, 'HiroLRAN': HiroLRAN}
MAX_NUMBER_OF_TIMES=300
### recall x_test is the in_samples from customDatasetMakers.ian_dataset, y_test is the out_samples.
### x_test contains normalized rofiles at t and actuators at t and t+1, y_test contains normalized profiles at t+1
# from y_test, get the real profiles at t+1 to t+prediction_length with warmup removed. This is used to compare with ML output. Also gets the real parameters
def get_ml_truth(y_test,
profiles, parameters,
recorded_profiles=['zipfit_etempfit_rho','zipfit_itempfit_rho','zipfit_trotfit_rho'],
prediction_length=-1, nwarmup=0, use_fancy_normalization=False):
num_samples=len(y_test)
num_profiles=len(profiles)
# just make this bigger than you think it needs to be
y=np.ones((num_samples,num_profiles,MAX_NUMBER_OF_TIMES,dataSettings.nx))*np.nan
y_params = np.ones((num_samples, len(parameters), MAX_NUMBER_OF_TIMES))*np.nan
for sample_ind in range(num_samples):
output_dic=state_to_dic(y_test[sample_ind], profiles, parameters)
#### get input stuff (profile warmup and actuator trajectories
# only needed for
'''input_dic=state_to_dic(x_test[sample_ind], profiles, parameters, calculations, actuators)
if use_fancy_normalization:
for actuator in actuators:
output_dic[actuator]=input_dic[actuator][:,0]'''
####
denormed_dic=get_denormalized_dic(output_dic, use_fancy_normalization=use_fancy_normalization)
for profile_ind,profile in enumerate(profiles):
num_times=len(denormed_dic[profile][nwarmup:])
y[sample_ind,profile_ind,:num_times]=denormed_dic[profile][nwarmup:]
for param_ind,param in enumerate(parameters):
num_times=len(denormed_dic[param][nwarmup:])
y_params[sample_ind,param_ind,:num_times]=denormed_dic[param][nwarmup:]
return y[:,:,:prediction_length], y_params[:,:,:prediction_length]
# get the profiles for warmup times from x_test
def get_ml_profile_warmup(x_test,
profiles, parameters, calculations, actuators,
recorded_profiles=['zipfit_etempfit_rho','zipfit_itempfit_rho','zipfit_trotfit_rho'],
recorded_parameters=[],
nwarmup=3, use_fancy_normalization=False):
num_samples=len(x_test)
profile_warmup=np.ones((num_samples,len(recorded_profiles),nwarmup+1,dataSettings.nx))*np.nan
parameter_warmup = np.ones((num_samples, len(recorded_parameters), nwarmup+1))*np.nan
for sample_ind in range(num_samples):
output_dic=state_to_dic(x_test[sample_ind], profiles, parameters, calculations, actuators)
denormed_dic=get_denormalized_dic(output_dic, use_fancy_normalization=use_fancy_normalization)
for profile_ind,profile in enumerate(recorded_profiles):
profile_warmup[sample_ind,profile_ind]=denormed_dic[profile][:nwarmup+1]
for param_ind,param in enumerate(recorded_parameters):
parameter_warmup[sample_ind,param_ind]=denormed_dic[param][:nwarmup+1]
return profile_warmup, parameter_warmup
# get the actuator trajectory during warmup and prediction times
def get_ml_actuator_trajectory(x_test,
profiles, parameters, calculations, actuators,
prediction_length=20,
nwarmup=0,
use_fancy_normalization=False):
num_samples=len(x_test)
actuator_trajectory=np.ones((num_samples,len(actuators),MAX_NUMBER_OF_TIMES))*np.nan
for sample_ind in range(num_samples):
output_dic=state_to_dic(x_test[sample_ind], profiles, parameters, calculations, actuators)
denormed_dic=get_denormalized_dic(output_dic, use_fancy_normalization=use_fancy_normalization)
for actuator_ind,actuator in enumerate(actuators):
num_times=len(denormed_dic[actuator])
actuator_trajectory[sample_ind,actuator_ind,:num_times]=denormed_dic[actuator][:,0]
actuator_trajectory[sample_ind,actuator_ind,num_times]=denormed_dic[actuator][-1,1]
return actuator_trajectory[:,:,:prediction_length+nwarmup+1]
# given an x_test, y_test, and a model, outputs the predictions for prediction_length, excluding nwarmup
def get_ml_predictions(x_test, y_test,
profiles, parameters, calculations, actuators,
considered_models,
recorded_profiles=['zipfit_etempfit_rho','zipfit_itempfit_rho','zipfit_trotfit_rho'],
recorded_actuators=['pinj'],
recorded_parameters=[],
prediction_length=20,nwarmup=0,
use_fancy_normalization=False,
bucket_size=10000):
test_x_buckets = make_bucket(x_test, bucket_size)
test_y_buckets = make_bucket(y_test, bucket_size)
test_length_buckets = [[len(arr) for arr in bucket] for bucket in test_x_buckets]
num_keys=len(x_test)
num_profiles=len(recorded_profiles)
num_parameters=len(recorded_parameters)
yhat=np.ones((num_keys,num_profiles,MAX_NUMBER_OF_TIMES,dataSettings.nx))*np.nan
yhat_parameters = np.ones((num_keys, num_parameters, MAX_NUMBER_OF_TIMES))*np.nan
begin_time=time.time()
prev_time=begin_time
evaluation_begin_time=time.time()
prev_time=evaluation_begin_time
with torch.no_grad():
sample_ind=0
for which_bucket in range(len(test_x_buckets)):
x_bucket=test_x_buckets[which_bucket]
y_bucket=test_y_buckets[which_bucket]
length_bucket=torch.tensor(test_length_buckets[which_bucket])
padded_x=pad_sequence(x_bucket, batch_first=True)
padded_y=pad_sequence(y_bucket, batch_first=True)
#padded_x=padded_x.to(device)
#padded_y=padded_y.to(device)
# only save simulations after warmup is over
# see note above, taking out ability to ensemble models
# since the ethos should be considering different ML and sim
# models on equal footing
model_output=torch.zeros_like(padded_y)
for model in considered_models:
#model=considered_models[0]
model_output+=model(padded_x, reset_probability=0, nwarmup=nwarmup)
model_output/=len(considered_models)
unpadded_output=unpad_sequence(model_output, length_bucket, batch_first=True)
for which_output,output in enumerate(unpadded_output):
output_dic=state_to_dic(output, profiles, parameters)
#### get input stuff (profile warmup and actuator trajectories
# only needed for
input_dic=state_to_dic(x_test[sample_ind], profiles, parameters, calculations, actuators)
if use_fancy_normalization:
for actuator in actuators:
output_dic[actuator]=input_dic[actuator][:,-1]
####
denormed_dic=get_denormalized_dic(output_dic, use_fancy_normalization=use_fancy_normalization)
for profile_ind,profile in enumerate(recorded_profiles):
num_times=len(denormed_dic[profile][nwarmup:])
yhat[sample_ind,profile_ind,:num_times]=denormed_dic[profile][nwarmup:]
for parameter_ind, parameter in enumerate(recorded_parameters):
num_times=len(denormed_dic[parameter][nwarmup:])
yhat_parameters[sample_ind,parameter_ind,:num_times]=denormed_dic[parameter][nwarmup:]
sample_ind+=1
print(f'Bucket {which_bucket+1}/{len(test_x_buckets)} took {time.time()-prev_time:0.0f}s')
prev_time=time.time()
print(f'Took {time.time()-begin_time:.2f} s')
return yhat[:,:,:prediction_length], yhat_parameters[:,:,:prediction_length]
def get_ml_profiles_with_warmup(profiles, warmup_ups):
return np.concatenate((warmup_ups, profiles), axis=2)
# get the predicted profiles at t+1 given the full history, outputs are normalized
def get_fast_profile_prediction(x_test_sample, model):
with torch.no_grad():
model_output = model(x_test_sample, reset_probability=1)
return model_output[:,-1:, :]
def get_considered_models(config_filename, ensemble=True, epoch=None):
config=configparser.ConfigParser()
config.read(config_filename)
output_filename_base=config['model']['output_filename_base']
output_dir=config['model']['output_dir']
model_type=config['model']['model_type']
profiles=config['inputs']['profiles'].split()
actuators=config['inputs']['actuators'].split()
parameters=config['inputs']['parameters'].split()
calculations=config['inputs']['calculations'].split()
state_length=len(profiles)*dataSettings.nx+len(parameters)
actuator_length=len(actuators)
calculation_length=len(calculations)*dataSettings.nx
considered_models=[]
epoch_specification=''
if epoch is not None:
epoch_specification=f'EPOCH{epoch}'
if ensemble:
regex=f'{output_filename_base}[0-9]*{epoch_specification}.tar'
all_model_files=glob.glob(os.path.join(output_dir, regex))
# glob is not as powerful, the star is for everything - so whittle down further so it only
# accepts repeats of numbers
all_model_files=[model_file for model_file in all_model_files if re.match(regex,os.path.basename(model_file))]
# exclude models under the median loss
#losses = []
#for model_file in all_model_files:
# saved_state=torch.load(model_file, map_location=torch.device('cpu'))
# losses.append(np.min([saved_state['val_losses'][-i] for i in range(10)]))
#max_loss = np.median(losses)
for model_file in all_model_files:
saved_state=torch.load(model_file, map_location=torch.device('cpu'))
if True: #np.min([saved_state['val_losses'][-i] for i in range(10)])<max_loss:
model=models[model_type](input_dim=state_length+calculation_length+2*actuator_length, output_dim=state_length,
**saved_state['model_hyperparams'])
model.load_state_dict(saved_state['model_state_dict'])
considered_models.append(model)
print(f'{len(considered_models)} models used')
#print(f'{len(considered_models)}/{len(all_model_files)} models used (i.e. only loss<{max_loss:0.2e})')
else:
model_file=os.path.join(output_dir, f'{output_filename_base}{epoch_specification}.tar')
saved_state=torch.load(model_file, map_location=torch.device('cpu'))
model=models[model_type](input_dim=state_length+calculation_length+2*actuator_length, output_dim=state_length,
**saved_state['model_hyperparams'])
model.load_state_dict(saved_state['model_state_dict'])
considered_models=[model]
print(f'Using {model_file}')
return considered_models
def get_fake_actuator_state(normalized_true_state, profiles, parameters, actuators):
fake_actuator_state=normalized_true_state.clone()
fake_actuator_dic=state_to_dic(fake_actuator_state, profiles=profiles, parameters=parameters, actuators=actuators)
for actuator in actuators:
arr = fake_actuator_dic[actuator][0]
freeze_index = len(arr)//2 - 40
perturb_index = len(arr)//2 + 30
perturb_length = 20
arr[freeze_index:-1] = torch.tensor([arr[freeze_index]]*len(arr[freeze_index:-1]))
if (actuator=='D_tot'):
perturb_index = len(arr)//2 + 30
perturb_length = 30
arr[perturb_index:perturb_index + perturb_length] = torch.tensor([((np.sin(np.pi*i/perturb_length))*arr[perturb_index] + arr[perturb_index]) for i in range(perturb_length)])
fake_actuator_state = dic_to_state(fake_actuator_dic, profiles, parameters, actuators=actuators)
return fake_actuator_state