Packaging core routines

.gitignore

@@ -0,0 +1,8 @@
+__pycache__
+.DS_Store
+*.egg-info
+dist
+examples/data
+examples/results
+examples/checkpoint
+*.npy

examples/burgers1d.yml

@@ -0,0 +1,53 @@
+lasdi:
+  type: gplasdi
+  gplasdi:
+    # device: mps
+    n_samples: 20
+    lr: 0.001
+    n_iter: 28000
+    n_greedy: 2000
+    max_greedy_iter: 28000
+    sindy_weight: 0.1
+    coef_weight: 1.e-6
+    path_checkpoint: checkpoint/
+    path_results: results/
+    # This probably should be specified on physics part.
+    time_dim: 1001
+    space_dim: 1001
+    xmin: -3.
+    xmax: 3.
+    tmax: 1.
+    # This probably should be specified as parameters.
+    a_min: 0.7
+    a_max: 0.9
+    w_min: 0.9
+    w_max: 1.1
+    n_a_grid: 21
+    n_w_grid: 21
+
+latent_space:
+  type: ae
+  ae:
+  # we should not specify fom_dimension here. temporarily.
+    fom_dimension: 1001
+    hidden_units: [100]
+    latent_dimension: 5
+
+physics:
+  type: burgers1d
+  burgers1d:
+    time_dim: 1001
+    space_dim: 1001
+    xmin: -3.
+    xmax: 3.
+    tmax: 1.
+    initial_train:
+      train_data: data/data_train.npy
+      test_data: data/data_test.npy
+      # This probably should be specified as parameters.
+      a_min: 0.7
+      a_max: 0.9
+      w_min: 0.9
+      w_max: 1.1
+      n_a_grid: 21
+      n_w_grid: 21

pyproject.toml

@@ -0,0 +1,36 @@
+[project]
+name = "lasdi"
+version = "2.0.0-dev"
+authors = [
+  { name="Christophe Bonneville", email="[email protected]" },
+  { name="Kevin (Seung Whan) Chung", email="[email protected]" },
+  { name="Youngsoo Choi", email="[email protected]" }
+]
+description = "LaSDI: Parametric latent space dynamics identification"
+readme = "README.md"
+requires-python = ">=3.12"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+  "torch>=2.3.0",
+  "numpy>=1.26.4",
+  "scikit-learn>=1.4.2",
+  "scipy>=1.13.1",
+  "pyyaml>=6.0",
+  "matplotlib>=3.8.4",
+  "argparse>=1.4.0"
+]
+
+[project.urls]
+Homepage = "https://github.com/LLNL/GPLaSDI"
+Issues = "https://github.com/LLNL/GPLaSDI/issues"
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project.scripts]
+lasdi = "lasdi.workflow:main [config_file]"

src/lasdi/enums.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+class NextStep(Enum):
+    Initial = 1
+    Train = 2
+    RunSample = 3
+    CollectSample = 4
+
+class Result(Enum):
+    Unexecuted = 1
+    Success = 2
+    Fail = 3
+    Complete = 4

src/lasdi/gplasdi.py

@@ -0,0 +1,255 @@
+from .sindy import *
+from .interp import *
+from .latent_space import *
+from .enums import *
+import torch
+import time
+import numpy as np
+
+def find_sindy_coef(Z, Dt, n_train, time_dim, loss_function):
+
+    '''
+
+    Computes the SINDy loss, reconstruction loss, and sindy coefficients
+
+    '''
+
+    loss_sindy = 0
+    loss_coef = 0
+
+    dZdt, Z = compute_sindy_data(Z, Dt)
+    sindy_coef = []
+
+    for i in range(n_train):
+
+        dZdt_i = dZdt[i, :, :]
+        Z_i = torch.cat([torch.ones(time_dim - 1, 1), Z[i, :, :]], dim = 1)
+        # coef_matrix_i = Z_i.pinverse() @ dZdt_i
+        coef_matrix_i = torch.linalg.lstsq(Z_i, dZdt_i).solution
+
+        loss_sindy += loss_function(dZdt_i, Z_i @ coef_matrix_i)
+        loss_coef += torch.norm(coef_matrix_i)
+
+        sindy_coef.append(coef_matrix_i.detach().numpy())
+
+    return loss_sindy, loss_coef, sindy_coef
+
+class BayesianGLaSDI:
+    def __init__(self, autoencoder, model_parameters):
+
+        '''
+
+        This class runs a full GPLaSDI training. It takes into input the autoencoder defined as a PyTorch object and the
+        dictionnary containing all the training parameters.
+        The "train" method with run the active learning training loop, compute the reconstruction and SINDy loss, train the GPs,
+        and sample a new FOM data point.
+
+        '''
+
+        self.autoencoder = autoencoder
+
+        # TODO(kevin): we need to simply possess physics object.
+        self.time_dim = model_parameters['time_dim']
+        self.space_dim = model_parameters['space_dim']
+        # self.n_z = model_parameters['n_z']
+
+        # TODO(kevin): we need to simply possess physics object.
+        xmin = model_parameters['xmin']
+        xmax = model_parameters['xmax']
+        self.Dx = (xmax - xmin) / (self.space_dim - 1)
+        assert(self.Dx > 0.)
+
+        tmax = model_parameters['tmax']
+        self.Dt = tmax / (self.time_dim - 1)
+
+        self.x_grid = np.linspace(xmin, xmax, self.space_dim)
+        self.t_grid = np.linspace(0, tmax, self.time_dim)
+
+        # TODO(kevin): generalize physics
+        # self.initial_condition = model_parameters['initial_condition']
+        from .physics.burgers1d import initial_condition
+        self.initial_condition = initial_condition
+
+        # TODO(kevin): generalize parameters
+        self.a_min = model_parameters['a_min']
+        self.a_max = model_parameters['a_max']
+        self.w_min = model_parameters['w_min']
+        self.w_max = model_parameters['w_max']
+
+        self.n_a_grid = model_parameters['n_a_grid']
+        self.n_w_grid = model_parameters['n_w_grid']
+        a_grid = np.linspace(self.a_min, self.a_max, self.n_a_grid)
+        w_grid = np.linspace(self.w_min, self.w_max, self.n_w_grid)
+        self.a_grid, self.w_grid = np.meshgrid(a_grid, w_grid)
+
+        self.param_grid = np.hstack((self.a_grid.reshape(-1, 1), self.w_grid.reshape(-1, 1)))
+        self.n_test = self.param_grid.shape[0]
+
+        self.n_samples = model_parameters['n_samples']
+        self.lr = model_parameters['lr']
+        self.n_iter = model_parameters['n_iter']
+        self.n_greedy = model_parameters['n_greedy']
+        self.max_greedy_iter = model_parameters['max_greedy_iter']
+        self.sindy_weight = model_parameters['sindy_weight']
+        self.coef_weight = model_parameters['coef_weight']
+
+        self.optimizer = torch.optim.Adam(autoencoder.parameters(), lr = self.lr)
+        self.MSE = torch.nn.MSELoss()
+
+        self.path_checkpoint = model_parameters['path_checkpoint']
+        self.path_results = model_parameters['path_results']
+
+        from os.path import dirname
+        from pathlib import Path
+        Path(dirname(self.path_checkpoint)).mkdir(parents=True, exist_ok=True)
+        Path(dirname(self.path_results)).mkdir(parents=True, exist_ok=True)
+
+        device = model_parameters['device'] if 'device' in model_parameters else 'cpu'
+        if (device == 'cuda'):
+            assert(torch.cuda.is_available())
+            self.device = device
+        elif (device == 'mps'):
+            assert(torch.backends.mps.is_available())
+            self.device = device
+        else:
+            self.device = 'cpu'
+
+        self.best_loss = np.Inf
+        self.restart_iter = 0
+
+        return
+
+    def train(self):
+
+        device = self.device
+        autoencoder_device = self.autoencoder.to(device)
+        X_train_device = self.X_train.to(device)
+
+        tic_start = time.time()
+        start_train_phase = []
+        start_fom_phase = []
+        end_train_phase = []
+        end_fom_phase = []
+
+
+        start_train_phase.append(tic_start)
+
+        for iter in range(self.restart_iter, self.n_iter):
+
+            self.optimizer.zero_grad()
+            Z = autoencoder_device.encoder(X_train_device)
+            X_pred = autoencoder_device.decoder(Z)
+            Z = Z.cpu()
+
+            loss_ae = self.MSE(X_train_device, X_pred)
+            loss_sindy, loss_coef, sindy_coef = find_sindy_coef(Z, self.Dt, self.n_train, self.time_dim, self.MSE)
+
+            max_coef = np.abs(np.array(sindy_coef)).max()
+
+            loss = loss_ae + self.sindy_weight * loss_sindy / self.n_train + self.coef_weight * loss_coef / self.n_train
+
+            loss.backward()
+            self.optimizer.step()
+
+            if loss.item() < self.best_loss:
+                torch.save(autoencoder_device.state_dict(), self.path_checkpoint + 'checkpoint.pt')
+                self.best_sindy_coef = sindy_coef
+                self.best_loss = loss.item()
+
+            print("Iter: %05d/%d, Loss: %3.10f, Loss AE: %3.10f, Loss SI: %3.10f, Loss COEF: %3.10f, max|c|: %04.1f, "
+                  % (iter + 1, self.n_iter, loss.item(), loss_ae.item(), loss_sindy.item(), loss_coef.item(), max_coef),
+                  end = '')
+
+            if self.n_train < 6:
+                print('Param: ' + str(np.round(self.param_train[0, :], 4)), end = '')
+
+                for i in range(1, self.n_train - 1):
+                    print(', ' + str(np.round(self.param_train[i, :], 4)), end = '')
+                print(', ' + str(np.round(self.param_train[-1, :], 4)))
+
+            else:
+                print('Param: ...', end = '')
+                for i in range(5):
+                    print(', ' + str(np.round(self.param_train[-6 + i, :], 4)), end = '')
+                print(', ' + str(np.round(self.param_train[-1, :], 4)))
+
+
+            if ((iter > self.restart_iter) and (iter < self.max_greedy_iter) and (iter % self.n_greedy == 0)):
+
+                end_train_phase.append(time.time())
+
+                print('\n~~~~~~~ Finding New Point ~~~~~~~')
+                # TODO(kevin): need to re-write this part.
+
+                start_fom_phase.append(time.time())
+                # X_train = X_train_device.cpu()
+                autoencoder = autoencoder_device.cpu()
+                autoencoder.load_state_dict(torch.load(self.path_checkpoint + 'checkpoint.pt'))
+
+                if len(self.best_sindy_coef) == self.n_train:
+                    sindy_coef = self.best_sindy_coef
+
+                Z0 = initial_condition_latent(self.param_grid, self.initial_condition, self.x_grid, autoencoder)
+
+                interpolation_data = build_interpolation_data(sindy_coef, self.param_train)
+                gp_dictionnary = fit_gps(interpolation_data)
+                n_coef = interpolation_data['n_coef']
+
+                coef_samples = [interpolate_coef_matrix(gp_dictionnary, self.param_grid[i, :], self.n_samples, n_coef, sindy_coef) for i in range(self.n_test)]
+                Zis = [simulate_uncertain_sindy(gp_dictionnary, self.param_grid[i, 0], self.n_samples, Z0[i], self.t_grid, sindy_coef, n_coef, coef_samples[i]) for i in range(self.n_test)]
+
+                a_index, w_index, m_index = get_max_std(autoencoder, Zis, self.n_a_grid, self.n_w_grid)
+
+                # TODO(kevin): implement save/load the new parameter
+                self.new_param = self.param_grid[m_index, :]
+                self.restart_iter = iter
+                next_step, result = NextStep.RunSample, Result.Success
+                end_fom_phase.append(time.time())
+                print('New param: ' + str(np.round(self.new_param, 4)) + '\n') 
+                return result, next_step
+
+        if len(self.best_sindy_coef) == self.n_train:
+            sindy_coef = self.best_sindy_coef
+        interpolation_data = build_interpolation_data(sindy_coef, self.param_train)
+        gp_dictionnary = fit_gps(interpolation_data)
+
+        tic_end = time.time()
+        total_time = tic_end - tic_start
+
+        bglasdi_results = {'autoencoder_param': self.autoencoder.state_dict(), 'final_param_train': self.param_train,
+                           'final_X_train': self.X_train, 'param_grid': self.param_grid,
+                           'sindy_coef': sindy_coef, '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_init': self.n_init, 'n_samples' : self.n_samples, 'n_a_grid' : self.n_a_grid, 'n_w_grid' : self.n_w_grid,
+                           'a_grid' : self.a_grid, 'w_grid' : self.w_grid,
+                           't_grid' : self.t_grid, 'x_grid' : self.x_grid, 'Dt' : self.Dt, 'Dx' : self.Dx,
+                           # TODO(kevin): need to fix timer.
+                           'total_time' : total_time, 'start_train_phase' : start_train_phase,
+                           'start_fom_phase' : start_fom_phase, 'end_train_phase' : end_train_phase, 'end_fom_phase' : end_fom_phase}
+
+        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)
+
+        next_step, result = None, Result.Complete
+        return result, next_step
+
+    def sample_fom(self):
+
+        # TODO(kevin): generalize this physics part.
+        from .physics.burgers1d import solver
+
+        new_a, new_b = self.new_param[0], self.new_param[1]
+        # TODO(kevin): generalize this physics part.
+        u0 = self.initial_condition(new_a, new_b, self.x_grid)
+
+        maxk = 10
+        convergence_threshold = 1e-8
+        new_X = solver(u0, maxk, convergence_threshold, self.time_dim - 1, self.space_dim, self.Dt, self.Dx)
+        new_X = new_X.reshape(1, self.time_dim, self.space_dim)
+        new_X = torch.Tensor(new_X)
+
+        self.X_train = torch.cat([self.X_train, new_X], dim = 0)
+        self.param_train = np.vstack((self.param_train, np.array([[new_a, new_b]])))
+        self.n_train += 1