Skip to content

Commit

Permalink
restart capability
Browse files Browse the repository at this point in the history
  • Loading branch information
dreamer2368 committed Sep 12, 2024
1 parent f4fa8e1 commit 4468542
Show file tree
Hide file tree
Showing 8 changed files with 174 additions and 65 deletions.
32 changes: 16 additions & 16 deletions examples/burgers1d.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,7 @@
"from lasdi.latent_space import Autoencoder, initial_condition_latent\n",
"from lasdi.postprocess import compute_errors\n",
"\n",
"date = '08_28_2024_20_46'\n",
"bglasdi_results = np.load('results/bglasdi_' + date + '.npy', allow_pickle = True).item()"
"filename = 'lasdi_09_11_2024_20_20.npy'"
]
},
{
Expand All @@ -55,17 +54,16 @@
"outputs": [],
"source": [
"import yaml\n",
"from lasdi.workflow import initialize_physics, initialize_latent_space, ld_dict\n",
"from lasdi.workflow import initialize_trainer\n",
"from lasdi.param import ParameterSpace\n",
"\n",
"cfg_file = 'burgers1d.yml'\n",
"with open(cfg_file, 'r') as f:\n",
" config = yaml.safe_load(f)\n",
"\n",
"param_space = ParameterSpace(config)\n",
"physics = initialize_physics(param_space, config)\n",
"autoencoder = initialize_latent_space(physics, config)\n",
"sindy = ld_dict[config['latent_dynamics']['type']](autoencoder.n_z, physics.nt, config['latent_dynamics'])"
"restart_file = np.load(filename, allow_pickle=True).item()\n",
"\n",
"trainer, param_space, physics, autoencoder, sindy = initialize_trainer(config, restart_file)"
]
},
{
Expand All @@ -75,16 +73,15 @@
"metadata": {},
"outputs": [],
"source": [
"autoencoder_param = bglasdi_results['autoencoder_param']\n",
"from lasdi.gp import fit_gps\n",
"\n",
"autoencoder_param = restart_file['latent_space']['autoencoder_param']\n",
"\n",
"X_train = bglasdi_results['final_X_train']\n",
"coefs = bglasdi_results['coefs']\n",
"gp_dictionnary = bglasdi_results['gp_dictionnary']\n",
"fd_type = bglasdi_results['latent_dynamics']['fd_type']\n",
"X_train = restart_file['trainer']['X_train']\n",
"\n",
"paramspace_dict = bglasdi_results['parameters']\n",
"param_train = paramspace_dict['final_param_train']\n",
"param_grid = paramspace_dict['param_grid']\n",
"paramspace_dict = restart_file['parameters']\n",
"param_train = paramspace_dict['train_space']\n",
"param_grid = paramspace_dict['test_space']\n",
"test_meshgrid = paramspace_dict['test_meshgrid']\n",
"test_grid_sizes = paramspace_dict['test_grid_sizes']\n",
"\n",
Expand All @@ -93,7 +90,7 @@
"n_a_grid, n_w_grid = test_grid_sizes\n",
"a_grid, w_grid = test_meshgrid\n",
"\n",
"physics_dict = bglasdi_results['physics']\n",
"physics_dict = restart_file['physics']\n",
"t_grid = physics_dict['t_grid']\n",
"x_grid = physics_dict['x_grid']\n",
"t_mesh, x_mesh = np.meshgrid(t_grid, x_grid)\n",
Expand All @@ -116,6 +113,9 @@
"n_z = autoencoder_param['fc' + str(n_hidden + 1) + '_e.weight'].shape[0]\n",
"\n",
"autoencoder.load_state_dict(autoencoder_param)\n",
"Z = autoencoder.encoder(X_train)\n",
"coefs = sindy.calibrate(Z, physics.dt, compute_loss=False, numpy=True)\n",
"gp_dictionnary = fit_gps(param_space.train_space, coefs)\n",
"\n",
"n_coef = sindy.ncoefs\n",
"\n",
Expand Down
4 changes: 4 additions & 0 deletions examples/burgers1d.yml
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,10 @@ lasdi:
path_checkpoint: checkpoint
path_results: results

workflow:
use_restart: true
restart_file: restarts/burgers1d.restart.npy

parameter_space:
parameters:
- name: a
Expand Down
56 changes: 36 additions & 20 deletions src/lasdi/gplasdi.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,21 @@ def get_fom_max_std(autoencoder, Zis):

return m_index

# move optimizer parameters to device
def optimizer_to(optim, device):
for param in optim.state.values():
# Not sure there are any global tensors in the state dict
if isinstance(param, torch.Tensor):
param.data = param.data.to(device)
if param._grad is not None:
param._grad.data = param._grad.data.to(device)
elif isinstance(param, dict):
for subparam in param.values():
if isinstance(subparam, torch.Tensor):
subparam.data = subparam.data.to(device)
if subparam._grad is not None:
subparam._grad.data = subparam._grad.data.to(device)

class BayesianGLaSDI:
X_train = None

Expand Down Expand Up @@ -161,7 +176,8 @@ def train(self):
self.optimizer.step()

if loss.item() < self.best_loss:
torch.save(autoencoder_device.state_dict(), self.path_checkpoint + '/' + 'checkpoint.pt')
torch.save(autoencoder_device.cpu().state_dict(), self.path_checkpoint + '/' + 'checkpoint.pt')
autoencoder_device = self.autoencoder.to(device)
self.best_coefs = coefs
self.best_loss = loss.item()

Expand Down Expand Up @@ -202,26 +218,10 @@ def train(self):
self.timer.start("finalize")

if (self.best_coefs.shape[0] == ps.n_train):
state_dict = torch.load(self.path_checkpoint + '/' + 'checkpoint.pt')
self.autoencoder.load_state_dict(state_dict)
coefs = self.best_coefs

gp_dictionnary = fit_gps(ps.train_space, coefs)

bglasdi_results = {'autoencoder_param': self.autoencoder.state_dict(), 'final_X_train': self.X_train,
'coefs': coefs, 'gp_dictionnary': gp_dictionnary, 'lr': self.lr, 'n_iter': self.n_iter,
'n_greedy': self.n_greedy, 'sindy_weight': self.sindy_weight, 'coef_weight': self.coef_weight,
'n_samples' : self.n_samples,
}
bglasdi_results['physics'] = self.physics.export()
bglasdi_results['parameters'] = self.param_space.export()
# TODO(kevin): restart capability for timer.
bglasdi_results['timer'] = self.timer.export()
bglasdi_results['latent_dynamics'] = self.latent_dynamics.export()

date = time.localtime()
date_str = "{month:02d}_{day:02d}_{year:04d}_{hour:02d}_{minute:02d}"
date_str = date_str.format(month = date.tm_mon, day = date.tm_mday, year = date.tm_year, hour = date.tm_hour + 3, minute = date.tm_min)
np.save(self.path_results + '/' + 'bglasdi_' + date_str + '.npy', bglasdi_results)

self.timer.end("finalize")
self.timer.print()

Expand Down Expand Up @@ -267,4 +267,20 @@ def sample_fom(self):
self.X_train = torch.cat([self.X_train, new_X], dim = 0)
else:
# TODO(kevin): interface for offline FOM simulation
raise RuntimeError("Offline FOM simulation is not supported yet!")
raise RuntimeError("Offline FOM simulation is not supported yet!")

def export(self):
dict_ = {'X_train': self.X_train, 'lr': self.lr, 'n_iter': self.n_iter, 'n_samples' : self.n_samples,
'n_greedy': self.n_greedy, 'sindy_weight': self.sindy_weight, 'coef_weight': self.coef_weight,
'restart_iter': self.restart_iter, 'timer': self.timer.export(), 'optimizer': self.optimizer.state_dict()
}
return dict_

def load(self, dict_):
self.X_train = dict_['X_train']
self.restart_iter = dict_['restart_iter']
self.timer.load(dict_['timer'])
self.optimizer.load_state_dict(dict_['optimizer'])
if (self.device != 'cpu'):
optimizer_to(self.optimizer, self.device)
return
7 changes: 7 additions & 0 deletions src/lasdi/latent_dynamics/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,4 +56,11 @@ def sample(self, coefs_sample, z0_sample, t_grid):
def export(self):
param_dict = {'dim': self.dim, 'ncoefs': self.ncoefs}
return param_dict

# SINDy does not need to load parameters.
# Other latent dynamics might need to.
def load(self, dict_):
assert(self.dim == dict_['dim'])
assert(self.ncoefs == dict_['ncoefs'])
return

10 changes: 9 additions & 1 deletion src/lasdi/latent_space.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,4 +174,12 @@ def apply_attention(self, x, layer):
x, _ = a(x, x, x) # apply attention
x = x.squeeze(1) # Remove sequence dimension

return x
return x

def export(self):
dict_ = {'autoencoder_param': self.cpu().state_dict()}
return dict_

def load(self, dict_):
self.load_state_dict(dict_['autoencoder_param'])
return
20 changes: 16 additions & 4 deletions src/lasdi/param.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ def create_uniform_1dspace(config):
class ParameterSpace:
param_list = []
param_name = []
n_param = 0
train_space = None
test_space = None
n_test = 0
Expand All @@ -35,6 +36,7 @@ def __init__(self, config):
parser = InputParser(config['parameter_space'], name="param_space_input")

self.param_list = parser.getInput(['parameters'], datatype=list)
self.n_param = len(self.param_list)

self.param_name = []
for param in self.param_list:
Expand Down Expand Up @@ -131,13 +133,23 @@ def appendTrainSpace(self, param):
return

def export(self):
dict_ = {'final_param_train': self.train_space,
'param_grid': self.test_space,
dict_ = {'train_space': self.train_space,
'test_space': self.test_space,
'test_grid_sizes': self.test_grid_sizes,
'test_meshgrid': self.test_meshgrid,
'n_init': self.n_init}
return dict_

def loadTrainSpace(self):
raise RuntimeError("ParameterSpace.loadTrainSpace is not implemented yet!")
def load(self, dict_):
self.train_space = dict_['train_space']
self.test_space = dict_['test_space']
self.test_grid_sizes = dict_['test_grid_sizes']
self.test_meshgrid = dict_['test_meshgrid']

assert(self.n_init == dict_['n_init'])
assert(self.train_space.shape[1] == self.n_param)
assert(self.test_space.shape[1] == self.n_param)

self.n_train = self.train_space.shape[0]
self.n_test = self.test_space.shape[0]
return
16 changes: 15 additions & 1 deletion src/lasdi/timing.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,22 @@ def print(self):
return

def export(self):
for start in self.starts:
if (start is not None):
raise RuntimeError('Timer.export: cannot export while Timer is still ticking!')

param_dict = {}
param_dict["names"] = self.names
param_dict["calls"] = self.calls
param_dict["times"] = self.times
return param_dict
return param_dict

def load(self, dict_):
self.names = dict_['names']
self.calls = dict_['calls']
self.times = dict_['times']

assert(len(self.names) == len(self.calls))
assert(len(self.names) == len(self.times))
self.starts = [None] * len(self.names)
return
Loading

0 comments on commit 4468542

Please sign in to comment.