Modularizing encoder/decoder in Autoencoder class

README.md

@@ -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

examples/burgers1d.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "d67dad03-9d76-4891-82ff-7e19d1369a24",
    "metadata": {},
    "outputs": [],
@@ -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",
@@ -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",

examples/burgers1d.offline.bash

@@ -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

examples/burgers1d.offline.yml

@@ -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.