Skip to content

Commit

Permalink
Restart capability and offline FOM interface (#9)
Browse files Browse the repository at this point in the history 
* restart capability

* run_samples and collect_samples workflow.

* pick_samples workflow and factorizing train/greedy sampling.

* burgers1d.ipynb has the entire workflow now.

* update README with prerequisites.

* offline fom wrapper

* offline bash script example

* separate param_space from physics.

* offline fom export

* template for prototyping.

* refactored between picksample and runsample.
  • Loading branch information
@dreamer2368
dreamer2368 authored Oct 23, 2024
1 parent abda363 commit d484837
Show file tree
Hide file tree
Showing 16 changed files with 1,244 additions and 220 deletions.
3 changes: 2 additions & 1 deletion README.md
View file
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
View file
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
View file
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
View file
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

0 comments on commit d484837

Please sign in to comment.