Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Restart capability and offline FOM interface #9

Merged
merged 11 commits into from
Oct 23, 2024
3 changes: 2 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ This python package requires updated prerequistes:
"scipy>=1.13.1",
"pyyaml>=6.0",
"matplotlib>=3.8.4",
"argparse>=1.4.0"
"argparse>=1.4.0",
"h5py"
```

### For LLNL LC Lassen users
Expand Down
192 changes: 139 additions & 53 deletions examples/burgers1d.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
"cells": [
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"id": "d67dad03-9d76-4891-82ff-7e19d1369a24",
"metadata": {},
"outputs": [],
Expand All @@ -16,116 +16,198 @@
},
{
"cell_type": "markdown",
"id": "cfffced1",
"id": "493f4313",
"metadata": {},
"source": [
"We assume all data generation/training is complete. If not done yet, please run on the terminal:\n",
"# Overall workflow and training\n",
"\n",
"Data generation/training can be performed by built-in executable `lasdi`. For this example of Burgers 1D equation, you can simply run on command-line terminal:\n",
"```\n",
"lasdi burgers1d.yml\n",
"```"
"```\n",
"\n",
"The workflow can be also manually constructed for those who prefer python scripts and for prototyping. Following code snippets show the high-level view of the workflow in the executable `lasdi`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b4966aa0",
"id": "1ead8f2d",
"metadata": {},
"outputs": [],
"source": [
"from lasdi.latent_space import Autoencoder, initial_condition_latent\n",
"from lasdi.postprocess import compute_errors\n",
"import yaml\n",
"from lasdi.workflow import initialize_trainer, run_samples, pick_samples\n",
"\n",
"cfg_file = 'burgers1d.yml'\n",
"with open(cfg_file, 'r') as f:\n",
" config = yaml.safe_load(f)\n",
"\n",
"date = '08_28_2024_20_46'\n",
"bglasdi_results = np.load('results/bglasdi_' + date + '.npy', allow_pickle = True).item()"
"trainer, param_space, physics, autoencoder, sindy = initialize_trainer(config)\n",
"\n",
"# generate initial training/test data\n",
"pick_samples(trainer, config)\n",
"run_samples(trainer, config)\n",
"# initial training given training data\n",
"trainer.train()\n",
"\n",
"while (trainer.restart_iter < trainer.max_iter):\n",
" if (trainer.restart_iter <= trainer.max_greedy_iter):\n",
" # perform greedy sampling to pick up new samples\n",
" pick_samples(trainer, config)\n",
" # update training data with newly picked samples\n",
" run_samples(trainer, config)\n",
"\n",
" # train over given training data\n",
" trainer.train()"
]
},
{
"cell_type": "markdown",
"id": "e1dfd2b6",
"id": "cfffced1",
"metadata": {},
"source": [
"Initialize physics solver, according to the config file `burgers1d.yml`"
"If you ran the command instead, a restart file is saved at the end of the training, which can be loaded for post-processing:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "68f93c32",
"id": "b4966aa0",
"metadata": {},
"outputs": [],
"source": [
"# Specify the restart file you have.\n",
"filename = 'lasdi_10_01_2024_17_09.npy'\n",
"\n",
"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)"
]
},
{
"cell_type": "markdown",
"id": "42758da9",
"metadata": {},
"source": [
"# Post-processing"
]
},
{
"cell_type": "markdown",
"id": "0cf48489",
"metadata": {},
"source": [
"Load data for post-processing:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ef6a1685",
"execution_count": 4,
"id": "dcdac0c2",
"metadata": {},
"outputs": [],
"source": [
"autoencoder_param = bglasdi_results['autoencoder_param']\n",
"coefs = trainer.best_coefs\n",
"X_train = trainer.X_train\n",
"X_test = trainer.X_test\n",
"\n",
"param_train = param_space.train_space\n",
"param_grid = param_space.test_space\n",
"test_meshgrid = param_space.test_meshgrid\n",
"test_grid_sizes = param_space.test_grid_sizes\n",
"n_init = param_space.n_init\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",
"n_a_grid, n_w_grid = test_grid_sizes\n",
"a_grid, w_grid = test_meshgrid\n",
"\n",
"t_grid = physics.t_grid\n",
"x_grid = physics.x_grid\n",
"t_mesh, x_mesh = np.meshgrid(t_grid, x_grid)\n",
"Dt, Dx = physics.dt, physics.dx\n",
"\n",
"time_dim, space_dim = t_grid.shape[0], x_grid.shape[0]\n",
"\n",
"n_coef = sindy.ncoefs"
]
},
{
"cell_type": "markdown",
"id": "c2b6e720",
"metadata": {},
"source": [
"They can be also loaded directly from restart file:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e796b0bc",
"metadata": {},
"outputs": [],
"source": [
"coefs = restart_file['trainer']['best_coefs']\n",
"X_train = restart_file['trainer']['X_train']\n",
"X_test = restart_file['trainer']['X_test']\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",
"n_init = paramspace_dict['n_init']\n",
"n_samples = 20\n",
"\n",
"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",
"Dt = physics_dict['dt']\n",
"Dx = physics_dict['dx']\n",
"\n",
"# total_time = bglasdi_results['total_time']\n",
"# start_train_phase = bglasdi_results['start_train_phase']\n",
"# start_fom_phase = bglasdi_results['start_fom_phase']\n",
"# end_train_phase = bglasdi_results['end_train_phase']\n",
"# end_fom_phase = bglasdi_results['end_fom_phase']\n",
"\n",
"data_test = np.load('data/data_test.npy', allow_pickle = True).item()\n",
"X_test = data_test['X_test']\n",
"\n",
"time_dim, space_dim = t_grid.shape[0], x_grid.shape[0]\n",
"n_coef = restart_file['latent_dynamics']['ncoefs']"
]
},
{
"cell_type": "markdown",
"id": "1262a0c3",
"metadata": {},
"source": [
"## Gaussian-process uncertainty evaluation\n",
"We evaluated the uncertainties of latent dynamics coefficients over 2d parameter space, with samples from GP prediction:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ef6a1685",
"metadata": {},
"outputs": [],
"source": [
"from lasdi.gp import fit_gps\n",
"from lasdi.gplasdi import sample_roms, average_rom\n",
"from lasdi.postprocess import compute_errors\n",
"from lasdi.gp import eval_gp\n",
"\n",
"n_hidden = len(autoencoder_param.keys()) // 4 - 1\n",
"hidden_units = [autoencoder_param['fc' + str(i + 1) + '_e.weight'].shape[0] for i in range(n_hidden)]\n",
"n_z = autoencoder_param['fc' + str(n_hidden + 1) + '_e.weight'].shape[0]\n",
"\n",
"autoencoder.load_state_dict(autoencoder_param)\n",
"n_samples = 20\n",
"autoencoder.cpu()\n",
"\n",
"n_coef = sindy.ncoefs\n",
"gp_dictionnary = fit_gps(param_space.train_space, coefs)\n",
"\n",
"from lasdi.gplasdi import sample_roms, average_rom\n",
"Zis_samples = sample_roms(autoencoder, physics, sindy, gp_dictionnary, param_grid, n_samples)\n",
"Zis_mean = average_rom(autoencoder, physics, sindy, gp_dictionnary, param_grid)\n",
"\n",
"print(Zis_mean.shape)\n",
"print(Zis_samples.shape)\n",
"\n",
"X_pred_mean = autoencoder.decoder(torch.Tensor(Zis_mean)).detach().numpy()\n",
"X_pred_samples = autoencoder.decoder(torch.Tensor(Zis_samples)).detach().numpy()\n",
"\n",
Expand All @@ -140,13 +222,17 @@
"avg_rel_error = avg_rel_error.reshape([n_w_grid, n_a_grid]).T\n",
"max_std = max_std.reshape([n_w_grid, n_a_grid]).T\n",
"\n",
"print(avg_rel_error.shape)\n",
"print(max_std.shape)\n",
"\n",
"from lasdi.gp import eval_gp\n",
"gp_pred_mean, gp_pred_std = eval_gp(gp_dictionnary, param_grid)"
]
},
{
"cell_type": "markdown",
"id": "2f31c964",
"metadata": {},
"source": [
"# Visualization"
]
},
{
"cell_type": "markdown",
"id": "389cdfb3",
Expand Down
48 changes: 48 additions & 0 deletions examples/burgers1d.offline.bash
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
#!/usr/bin/bash

check_result () {
# $1: Result output of the previous command ($?)
# $2: Name of the previous command
if [ $1 -eq 0 ]; then
echo "$2 succeeded"
else
echo "$2 failed"
exit -1
fi
}

# First stage will be PickSample, which will produce parameter point files in hdf5 format.
# Each lasdi command will save a restart file, which will be read on the next lasdi command.
# So all stages are run by the same command, directed differently by the restart file.
lasdi burgers1d.offline.yml
check_result $? initial-picksample

# Run/save FOM solution with offline FOM solver.
burgers1d burgers1d.offline.yml
check_result $? initial-runsample

# Collect FOM solution.
lasdi burgers1d.offline.yml
check_result $? initial-collect

# Train latent dynamics model.
lasdi burgers1d.offline.yml
check_result $? initial-train

for k in {0..8}
do
# Pick a new sample from greedy sampling
lasdi burgers1d.offline.yml
check_result $? pick-sample

# A cycle of offline FOM/CollectSamples
burgers1d burgers1d.offline.yml
check_result $? run-sample

lasdi burgers1d.offline.yml
check_result $? collect-sample

# Train latent dynamics model.
lasdi burgers1d.offline.yml
check_result $? train
done
61 changes: 61 additions & 0 deletions examples/burgers1d.offline.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
lasdi:
type: gplasdi
gplasdi:
# device: mps
n_samples: 20
lr: 0.001
max_iter: 2000
n_iter: 200
max_greedy_iter: 2000
ld_weight: 0.1
coef_weight: 1.e-6
path_checkpoint: checkpoint
path_results: results

workflow:
use_restart: true
restart_file: restarts/burgers1d.restart.npy
offline_greedy_sampling:
train_param_file: sampling/new_train.burgers1d.h5
test_param_file: sampling/new_test.burgers1d.h5
train_sol_file: sampling/new_Xtrain.burgers1d.h5
test_sol_file: sampling/new_Xtest.burgers1d.h5

parameter_space:
parameters:
- name: a
min: 0.7
max: 0.9
test_space_type: uniform
sample_size: 11
log_scale: false
- name: w
min: 0.9
max: 1.1
test_space_type: uniform
sample_size: 11
log_scale: false
test_space:
type: grid

latent_space:
type: ae
ae:
hidden_units: [100]
latent_dimension: 5

latent_dynamics:
type: sindy
sindy:
fd_type: sbp12
coef_norm_order: fro

physics:
type: burgers1d
burgers1d:
offline_driver: true
number_of_timesteps: 1001
simulation_time: 1.
grid_size: [1001]
xmin: -3.
xmax: 3.
Loading