├── .gitignore ├── configs ├── __pycache__ │ ├── easyAttn.cpython-38.pyc │ ├── lstm.cpython-38.pyc │ ├── nomenclature.cpython-38.pyc │ ├── selfAttn.cpython-38.pyc │ └── vae.cpython-38.pyc ├── easyAttn.py ├── lstm.py ├── nomenclature.py ├── selfAttn.py └── vae.py ├── lib ├── POD.py ├── __pycache__ │ ├── datas.cpython-38.pyc │ ├── init.cpython-38.pyc │ ├── model.cpython-38.pyc │ ├── pp_space.cpython-38.pyc │ ├── pp_time.cpython-38.pyc │ ├── runner_lantent_predictor.cpython-38.pyc │ ├── runners.cpython-38.pyc │ └── train.cpython-38.pyc ├── datas.py ├── init.py ├── model.py ├── pp_space.py ├── pp_time.py ├── runners.py └── train.py ├── license.md ├── main.py ├── models ├── VAE_easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt ├── checkpoints │ ├── Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000_bestVal.pth.tar │ └── Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000_final.pth.tar └── pretrained │ ├── LSTM_64in_64dmodel_1next_2dim_Noneemb_128hideen_4nlayer_Noneoutact_100Epoch_135000N_TrueES_50P.pt │ ├── Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000.pth.tar │ ├── easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt │ └── selfAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt ├── nns ├── RNNs.py ├── __pycache__ │ ├── AutoEncoder.cpython-38.pyc │ ├── EmbedTransformerEncoder.cpython-38.pyc │ ├── Embedding.cpython-38.pyc │ ├── FullTransformer.cpython-38.pyc │ ├── RNNs.cpython-38.pyc │ ├── RevTransformer.cpython-38.pyc │ ├── Transformer.cpython-38.pyc │ ├── attns.cpython-38.pyc │ ├── beta_vae.cpython-38.pyc │ ├── easyAttns.cpython-38.pyc │ ├── embedding.cpython-38.pyc │ ├── layers.cpython-38.pyc │ ├── transformer.cpython-38.pyc │ └── vae.cpython-38.pyc ├── attns.py ├── beta_vae.py ├── embedding.py ├── layers.py └── transformer.py ├── readme.md ├── requirements.txt ├── res ├── easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.npz └── modes_Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000.hdf5 └── utils ├── __.py ├── __pycache__ ├── chaotic.cpython-38.pyc ├── config.cpython-38.pyc ├── datas.cpython-38.pyc ├── figs_time.cpython-38.pyc ├── model.cpython-38.pyc ├── plot.cpython-38.pyc ├── plot_time.cpython-38.pyc ├── plt_rc_setup.cpython-38.pyc ├── pp.cpython-38.pyc ├── reverse_config.cpython-38.pyc └── train.cpython-38.pyc ├── figs.py ├── figs_time.py ├── io.py └── plt_rc_setup.py /.gitignore: -------------------------------------------------------------------------------- 1 | data/ 2 | -------------------------------------------------------------------------------- /configs/__pycache__/easyAttn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/configs/__pycache__/easyAttn.cpython-38.pyc -------------------------------------------------------------------------------- /configs/__pycache__/lstm.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/configs/__pycache__/lstm.cpython-38.pyc -------------------------------------------------------------------------------- /configs/__pycache__/nomenclature.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/configs/__pycache__/nomenclature.cpython-38.pyc -------------------------------------------------------------------------------- /configs/__pycache__/selfAttn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/configs/__pycache__/selfAttn.cpython-38.pyc -------------------------------------------------------------------------------- /configs/__pycache__/vae.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/configs/__pycache__/vae.cpython-38.pyc -------------------------------------------------------------------------------- /configs/easyAttn.py: -------------------------------------------------------------------------------- 1 | class easyAttn_config: 2 | """ 3 | A class of configuration of Transformer predictor 4 | """ 5 | from configs.vae import VAE_config 6 | 7 | 8 | in_dim = 64 9 | out_dim = 64 # The output sequence length 10 | d_model = 64 11 | 12 | time_proj = 64 # The projection on time, which is used for new embedding stragtegy 13 | next_step = 1 14 | nmode = VAE_config.latent_dim # Should be consistent as the modes 15 | 16 | 17 | num_head = 4 18 | attn_type = "easy" 19 | embed = "time" # sin / cos/ time 20 | num_block = 4 # Number of layer 21 | 22 | is_res_attn = True 23 | is_res_proj = True 24 | proj_dim = 128 25 | 26 | act_proj = "relu" 27 | is_output = True 28 | out_act = None 29 | 30 | Epoch = 100 31 | Batch_size = 256 32 | lr = 1e-3 33 | 34 | train_split = 0.8 35 | num_train = 135000 36 | 37 | early_stop = True 38 | 39 | if early_stop == True: 40 | patience = 50 # 30 or 50 41 | else: 42 | patience = 0 43 | 44 | -------------------------------------------------------------------------------- /configs/lstm.py: -------------------------------------------------------------------------------- 1 | class lstm_config: 2 | """ 3 | A class of config for LSTM Predictor 4 | """ 5 | from configs.vae import VAE_config 6 | 7 | in_dim = 64 8 | d_model = 64 9 | next_step = 1 10 | nmode = VAE_config.latent_dim 11 | 12 | num_layer = 4 13 | embed = None 14 | 15 | hidden_size = 128 16 | 17 | is_output = True 18 | out_act = None 19 | 20 | Epoch = 100 21 | Batch_size = 256 22 | lr = 1e-3 23 | 24 | train_split = 0.8 25 | val_split = 0.2 26 | num_train = 135000 27 | 28 | early_stop = True 29 | 30 | if early_stop == True: 31 | patience = 50 32 | else: 33 | patience = 0 34 | 35 | 36 | 37 | 38 | -------------------------------------------------------------------------------- /configs/nomenclature.py: -------------------------------------------------------------------------------- 1 | """ 2 | Functions for naming the models in the present study 3 | @yuningw 4 | """ 5 | 6 | 7 | def Name_VAE(cfg): 8 | 9 | """ 10 | A function to name the VAE checkpoint 11 | 12 | Args: 13 | cfg: A class of VAE model configuration 14 | 15 | Returns: 16 | name: A string for VAE name 17 | """ 18 | 19 | name = f"Re{cfg.Re}_" +\ 20 | f'smallerCNN_'+\ 21 | f'beta{cfg.beta}_' + \ 22 | f'wDecay{cfg.decWdecay}_'+\ 23 | f'dim{cfg.latent_dim}_'+\ 24 | f'lr{cfg.lr}OneCycleLR{cfg.lr_end}_'+\ 25 | f'bs{cfg.batch_size}_'+\ 26 | f'epochs{cfg.epochs}' 27 | 28 | return name 29 | 30 | 31 | 32 | def Make_Transformer_Name(cfg): 33 | """ 34 | A function to name the VAE checkpoint 35 | 36 | Args: 37 | cfg: A class of Transorfmer model configuration 38 | 39 | Returns: 40 | name: A string for Transformer model 41 | """ 42 | 43 | case_name = f"{cfg.attn_type}Attn_"+\ 44 | f"{cfg.in_dim}in_" +\ 45 | f'{cfg.d_model}dmodel_' +\ 46 | f'{cfg.next_step}next_'+\ 47 | f'{cfg.nmode}dim_'+\ 48 | f"{cfg.embed}emb_"+\ 49 | f"{cfg.num_head}h_"+\ 50 | f"{cfg.num_block}nb_"+\ 51 | f"{cfg.proj_dim}ff_"+\ 52 | f"{cfg.act_proj}act_"+\ 53 | f"{cfg.out_act}outact_"+\ 54 | f"{cfg.Epoch}Epoch_"+\ 55 | f"{cfg.num_train}N_"+\ 56 | f"{cfg.early_stop}ES_"+\ 57 | f"{cfg.patience}P" 58 | return case_name 59 | 60 | 61 | 62 | def Make_LSTM_Name(cfg): 63 | """ 64 | A function to name the LSTM checkpoint 65 | 66 | Args: 67 | cfg: A class of LSTM model configuration 68 | 69 | Returns: 70 | name: A string for LSTM model 71 | """ 72 | 73 | case_name = f"LSTM_"+\ 74 | f"{cfg.in_dim}in_"+\ 75 | f"{cfg.d_model}dmodel_"+\ 76 | f"{cfg.next_step}next_"+\ 77 | f"{cfg.nmode}dim_"+\ 78 | f"{cfg.embed}emb_"+\ 79 | f"{cfg.hidden_size}hideen_"+\ 80 | f"{cfg.num_layer}nlayer_"+\ 81 | f"{cfg.out_act}outact_"+\ 82 | f"{cfg.Epoch}Epoch_"+\ 83 | f"{cfg.num_train}N_"+\ 84 | f"{cfg.early_stop}ES_"+\ 85 | f"{cfg.patience}P" 86 | 87 | return case_name 88 | 89 | 90 | 91 | -------------------------------------------------------------------------------- /configs/selfAttn.py: -------------------------------------------------------------------------------- 1 | class selfAttn_config: 2 | """ 3 | A class of configuration of Transformer predictor 4 | """ 5 | from configs.vae import VAE_config 6 | 7 | 8 | in_dim = 64 9 | out_dim = 64 # The output sequence length 10 | d_model = 64 11 | 12 | time_proj = 64 # The projection on time, which is used for new embedding stragtegy 13 | next_step = 1 14 | 15 | nmode = VAE_config.latent_dim # Should be consistent as the modes 16 | 17 | 18 | num_head = 4 19 | attn_type = "self" # self or selfconv or easy 20 | 21 | embed = "time" # sin / cos/ time 22 | num_block = 4 # Number of layer 23 | 24 | is_res_attn = True 25 | is_res_proj = True 26 | proj_dim = 128 27 | 28 | act_proj = "relu" 29 | is_output = True 30 | out_act = None 31 | 32 | Epoch = 100 33 | Batch_size = 256 34 | lr = 1e-3 35 | 36 | train_split = 0.8 37 | num_train = 135000 38 | 39 | early_stop = True 40 | 41 | if early_stop == True: 42 | patience = 50 # 30 or 50 43 | else: 44 | patience = 0 45 | 46 | -------------------------------------------------------------------------------- /configs/vae.py: -------------------------------------------------------------------------------- 1 | class VAE_config: 2 | """ 3 | A class of configuration of vae model 4 | """ 5 | 6 | Re = 40 # Reynolds number for the case 7 | 8 | latent_dim = 2 9 | 10 | delta_t = 1 11 | 12 | batch_size = 256 13 | 14 | lr = 2e-4 15 | 16 | lr_end = 1e-5 17 | 18 | epochs = 1000 19 | 20 | beta = 0.005 21 | 22 | beta_init = 0.001 23 | 24 | downsample = 1 # We set = 1, change if you need 25 | 26 | beta_warmup = 20 27 | 28 | n_test = 200 29 | 30 | encWdecay = 0 31 | 32 | decWdecay = 0 33 | 34 | DATA_TO_GPU = False 35 | 36 | -------------------------------------------------------------------------------- /lib/POD.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import time 3 | import matplotlib.pyplot as plt 4 | import numpy as np 5 | from sklearn.utils.extmath import randomized_svd 6 | 7 | 8 | class POD: 9 | def __init__(self, datafile, n_test, re, path, n_modes=100, delta_t=1) -> None: 10 | """ 11 | A runner for POD 12 | 13 | Args: 14 | 15 | datafile : (Str) Path of training data 16 | n_test : (int) Number of test timesteps (not used for POD) 17 | re : (int) Reynolds number 18 | path : (str) Path to save POD results 19 | n_modes : (int) Number of POD modes to calculate 20 | delta_t : (int) Steps between snapshots to use 21 | """ 22 | 23 | self.datafile = datafile 24 | self.n_modes = n_modes 25 | self.delta_t = delta_t 26 | self.n_test = n_test 27 | self.re = re 28 | self.casename = f'POD_Re{self.re}_dt{self.delta_t}_ntst{self.n_test}_nm{self.n_modes}.npz' 29 | self.filename = path + self.casename 30 | 31 | print(f"POD file name:\n {self.filename}") 32 | 33 | def load_data(self): 34 | # load data 35 | with h5py.File(self.datafile, 'r') as f: 36 | u_scaled = f['UV'][:] 37 | mean = f['mean'][:] 38 | std = f['std'][()] 39 | 40 | u_scaled = np.moveaxis(u_scaled, -1, 1) 41 | self.mean = np.moveaxis(mean, -1, 1) 42 | self.std = np.moveaxis(std, -1, 1) 43 | 44 | u_scaled = u_scaled[::self.delta_t] 45 | 46 | n_total = u_scaled.shape[0] 47 | self.n_train = n_total - self.n_test 48 | print(f"N train: {self.n_train:d}, N test: {self.n_test:d}, N total {n_total:d}") 49 | print(f'u_scaled {u_scaled.shape}') 50 | 51 | self.u_train = u_scaled[:self.n_train] 52 | self.u_test = u_scaled[self.n_train:] 53 | 54 | def get_POD(self): 55 | 56 | try: 57 | self.load_POD() 58 | print('POD loaded from file') 59 | except: 60 | print('Calculating POD') 61 | self.calc_POD() 62 | 63 | def load_POD(self): 64 | d = np.load(self.filename) 65 | self.temporal_modes = d['tm'] 66 | self.spatial_modes = d['sm'] 67 | self.eigVal = d['eig'] 68 | 69 | def calc_POD(self): 70 | 71 | u_train_flat = self.u_train.reshape(self.u_train.shape[0], -1) 72 | u_test_flat = self.u_test.reshape(self.u_test.shape[0], -1) 73 | 74 | print(f'POD u_train: {u_train_flat.shape}') 75 | print(f'POD u_test: {u_test_flat.shape}') 76 | 77 | print(f'U shape: {u_train_flat.shape}') 78 | # C matrix 79 | print('Calc C matrix', end="") 80 | start = time.time() 81 | C = u_train_flat.T.dot(u_train_flat) 82 | C = C / (self.n_train - 1) 83 | print(f': {(time.time() - start):.1f}s') 84 | print(f'C shape: {C.shape}') 85 | 86 | # SVD 87 | print('Calc SVD', end="") 88 | start = time.time() 89 | self.spatial_modes, self.eigVal, _ = randomized_svd(C, n_components=self.n_modes, random_state=0) 90 | print(f': {(time.time() - start):.1f}s') 91 | print(f'spatial_modes {self.spatial_modes.shape}') 92 | 93 | self.temporal_modes = u_train_flat.dot(self.spatial_modes) 94 | 95 | print(f'temporal_modes shape: {self.temporal_modes.shape}') 96 | print(f'spatial_modes shape: {self.spatial_modes.shape}') 97 | 98 | np.savez_compressed( 99 | file=self.filename, 100 | tm=self.temporal_modes, 101 | sm=self.spatial_modes, 102 | eig=self.eigVal 103 | ) 104 | 105 | def eval_POD(self): 106 | from lib.pp_space import get_Ek 107 | 108 | self.Ek_nm = np.zeros(self.n_modes) 109 | 110 | for nm in range(1, self.n_modes + 1): 111 | u_train_rec = self.temporal_modes[:, :nm].dot(self.spatial_modes[:, :nm].T).reshape(self.u_train.shape) 112 | 113 | self.Ek_nm[nm - 1] = get_Ek(self.u_train, u_train_rec) 114 | # print(f'POD train E = {self.Ek_nm[nm-1]:.4f}, {nm} modes') 115 | 116 | print(f'E POD: {self.Ek_nm}') 117 | 118 | 119 | # Test code 120 | if __name__ == "__main__": 121 | # data 122 | datafile = "../data/Data2PlatesGap1Re40_Alpha-00_downsampled_v6.hdf5" 123 | 124 | re = 40 125 | delta_t = 1 126 | n_modes = 10 127 | n_test = 200 // delta_t 128 | 129 | POD = POD(datafile, n_test, re, '../res/', n_modes, delta_t) 130 | POD.load_data() 131 | POD.get_POD() 132 | POD.eval_POD() 133 | -------------------------------------------------------------------------------- /lib/__pycache__/datas.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/datas.cpython-38.pyc -------------------------------------------------------------------------------- /lib/__pycache__/init.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/init.cpython-38.pyc -------------------------------------------------------------------------------- /lib/__pycache__/model.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/model.cpython-38.pyc -------------------------------------------------------------------------------- /lib/__pycache__/pp_space.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/pp_space.cpython-38.pyc -------------------------------------------------------------------------------- /lib/__pycache__/pp_time.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/pp_time.cpython-38.pyc -------------------------------------------------------------------------------- /lib/__pycache__/runner_lantent_predictor.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/runner_lantent_predictor.cpython-38.pyc -------------------------------------------------------------------------------- /lib/__pycache__/runners.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/runners.cpython-38.pyc -------------------------------------------------------------------------------- /lib/__pycache__/train.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/lib/__pycache__/train.cpython-38.pyc -------------------------------------------------------------------------------- /lib/datas.py: -------------------------------------------------------------------------------- 1 | """ 2 | Functions for loading data 3 | 4 | author: @yuningw 5 | """ 6 | import h5py 7 | import numpy as np 8 | 9 | 10 | ############################### 11 | ## beta-VAE 12 | ############################### 13 | 14 | #--------------------------------------------------------------------- 15 | def loadData(file, printer=False): 16 | """ 17 | Read flow field dataset 18 | 19 | Args: 20 | file : (str) Path of database 21 | 22 | printer : (bool) print the summary of datasets 23 | 24 | Returns: 25 | 26 | u_scaled: (NumpyArray) The scaled data 27 | 28 | mean : (float) mean of data 29 | 30 | std : (float) std of data 31 | 32 | """ 33 | 34 | with h5py.File(file, 'r') as f: 35 | u_scaled = f['UV'][:] 36 | mean = f['mean'][:] 37 | std = f['std'][()] 38 | 39 | u_scaled = np.moveaxis(u_scaled, -1, 1) 40 | mean = np.moveaxis(mean, -1, 1) 41 | std = np.moveaxis(std, -1, 1) 42 | 43 | if printer: 44 | print('u_scaled: ', u_scaled.shape) 45 | print('mean: ', mean.shape) 46 | print('std: ', std) 47 | 48 | return u_scaled, mean, std 49 | 50 | 51 | #--------------------------------------------------------------------- 52 | def get_vae_DataLoader(d_train, n_train, device, batch_size): 53 | """ 54 | make tensor data loader for training 55 | 56 | Args: 57 | d_train: (NumpyArray) Train DataSet 58 | 59 | n_train : (int) Training samples 60 | 61 | device : (str) Device 62 | 63 | batch_size: (int) Batch size 64 | 65 | 66 | Return: 67 | train_dl, val_dl: The train and validation DataLoader 68 | """ 69 | import torch 70 | from torch.utils.data import DataLoader, TensorDataset 71 | 72 | if ('cuda' in device): 73 | train_dl = torch.utils.data.DataLoader(dataset=torch.from_numpy(d_train[:n_train]).to(device), 74 | batch_size=batch_size, 75 | shuffle=True, num_workers=0) 76 | val_dl = torch.utils.data.DataLoader(dataset=torch.from_numpy(d_train[n_train:]).to(device), 77 | batch_size=batch_size, 78 | shuffle=False, num_workers=0) 79 | else: 80 | train_dl = torch.utils.data.DataLoader(dataset=torch.from_numpy(d_train[:n_train]), batch_size=batch_size, 81 | shuffle=True, pin_memory=True, num_workers=4, 82 | persistent_workers=True) 83 | val_dl = torch.utils.data.DataLoader(dataset=torch.from_numpy(d_train[n_train:]), batch_size=batch_size, 84 | shuffle=False, pin_memory=True, num_workers=4, 85 | persistent_workers=True) 86 | 87 | return train_dl, val_dl 88 | 89 | 90 | 91 | ############################### 92 | ## Temporal Prediction 93 | ############################### 94 | 95 | #--------------------------------------------------------------------- 96 | def make_Sequence(cfg,data): 97 | """ 98 | Generate time-delay sequence data 99 | 100 | Args: 101 | cfg: A class contain the configuration of data 102 | data: A numpy array follows [Ntime, Nmode] shape 103 | 104 | Returns: 105 | X: Numpy array for Input 106 | Y: Numpy array for Output 107 | """ 108 | 109 | from tqdm import tqdm 110 | import numpy as np 111 | 112 | if len(data.shape) <=2: 113 | data = np.expand_dims(data,0) 114 | seqLen = cfg.in_dim 115 | nSamples = (data.shape[1]-seqLen) 116 | X = np.empty([nSamples, seqLen, data.shape[-1]]) 117 | Y = np.empty([nSamples, cfg.next_step,data.shape[-1]]) 118 | # Fill the input and output arrays with data 119 | k = 0 120 | for i in tqdm(np.arange(data.shape[0])): 121 | for j in np.arange(data.shape[1]-seqLen- cfg.next_step): 122 | X[k] = data[i, j :j+seqLen] 123 | Y[k] = data[i, j+seqLen :j+seqLen+cfg.next_step] 124 | k = k + 1 125 | print(f"The training data has been generated, has shape of {X.shape, Y.shape}") 126 | 127 | return X, Y 128 | 129 | 130 | #--------------------------------------------------------------------- 131 | def make_DataLoader(X,y,batch_size, 132 | drop_last=False,train_split = 0.8): 133 | """ 134 | make tensor data loader for training 135 | 136 | Args: 137 | X: Tensor of features 138 | y: Tensor of target 139 | batch_size: Batch size 140 | drop_last: If drop the last batch which does not have same number of mini batch 141 | train_split: A ratio of train and validation split 142 | 143 | Return: 144 | train_dl, val_dl: The train and validation DataLoader 145 | """ 146 | 147 | from torch.utils.data import DataLoader, TensorDataset,random_split 148 | try: 149 | dataset = TensorDataset(X,y) 150 | except: 151 | print("The data is not torch.tenor!") 152 | 153 | len_d = len(dataset) 154 | train_size = int(train_split * len_d) 155 | valid_size = len_d - train_size 156 | 157 | train_d , val_d = random_split(dataset,[train_size, valid_size]) 158 | 159 | train_dl = DataLoader(train_d,batch_size=batch_size,drop_last=drop_last,shuffle=True) 160 | val_dl = DataLoader(val_d,batch_size=batch_size,drop_last=drop_last,shuffle=True) 161 | 162 | return train_dl, val_dl -------------------------------------------------------------------------------- /lib/init.py: -------------------------------------------------------------------------------- 1 | """ 2 | Initialisation and setup before running 3 | @yuningw 4 | """ 5 | 6 | class pathsBib: 7 | 8 | data_path = 'data/' 9 | model_path = 'models/'; 10 | res_path = 'res/' 11 | fig_path = 'figs/' 12 | log_path = 'logs/' 13 | chekp_path = model_path + 'checkpoints/' 14 | pretrain_path = model_path + "pretrained/" 15 | 16 | #------------------------------------------------- 17 | def init_env(Re=40): 18 | """ 19 | A function for initialise the path and data 20 | 21 | Args: 22 | Re : The corresponding Reynolds number of the case 23 | 24 | Returns: 25 | 26 | data_file : (str) File name 27 | 28 | """ 29 | from configs.vae import VAE_config 30 | assert(Re == VAE_config.Re), print(f"ERROR: Please match the config of vae {VAE_config.Re}!") 31 | 32 | is_init_path = init_path() 33 | 34 | datafile = None 35 | 36 | if is_init_path: 37 | is_acquired, datafile = acquire_data(Re) 38 | else: 39 | print(f"ERROR: Init Path failed!") 40 | 41 | return datafile 42 | 43 | 44 | #------------------------------------------------- 45 | def init_path(): 46 | """ 47 | Initialisation of all the paths 48 | 49 | Returns: 50 | is_init_path : (bool) if initialise success 51 | """ 52 | import os 53 | from pathlib import Path 54 | 55 | is_init_path = False 56 | try: 57 | print("#"*30) 58 | print(f"Start initialisation of paths") 59 | path_list =[i for _,i in pathsBib.__dict__.items() if type(i)==str and "/" in i] 60 | print(path_list) 61 | for pth in path_list: 62 | # if "/" in pth: 63 | Path(pth).mkdir(exist_ok=True) 64 | print(f"INIT:\t{pth}\tDONE") 65 | print("#"*30) 66 | is_init_path = True 67 | except: 68 | print(f"Error: FAILD to inital path, Please check setup for your path!") 69 | 70 | return is_init_path 71 | 72 | 73 | #------------------------------------------------- 74 | def acquire_data(Re=40): 75 | 76 | """ 77 | Acquisition of dataset from zendo 78 | 79 | Args: 80 | Re : The corresponding Reynolds number of the case 81 | 82 | Returns: 83 | is_acquired : (bool) A flag for results 84 | data_file : (str) File name 85 | """ 86 | 87 | import urllib.request 88 | import os 89 | import time 90 | is_acuqired = False 91 | datfile = None 92 | if Re == 40: 93 | datfile = pathsBib.data_path + "Data2PlatesGap1Re40_Alpha-00_downsampled_v6.hdf5" 94 | else: 95 | print(f"Error: Data might be too large to download, please get it manually") 96 | 97 | try: 98 | if not os.path.exists(datfile): 99 | try: 100 | print(f"{datfile}") 101 | print(f"INFP: Not found, trying to download example dataset") 102 | st = time.time() 103 | urllib.request.urlretrieve('https://zenodo.org/records/10501216/files/Data2PlatesGap1Re40_Alpha-00_downsampled_v6.hdf5?download=1', 104 | datfile) 105 | et = time.time() 106 | ct = et - st 107 | print(f"File downloaded successfully to {datfile}, Time Cost: {ct:.2f}s") 108 | is_acuqired = True 109 | 110 | except Exception as e: 111 | print(f"Failed to download sample dataset. Error: {e}") 112 | else: 113 | print(f"INFO: Data exists in local!") 114 | 115 | is_acuqired = True 116 | 117 | except: 118 | print(f"Error: Failed loading data") 119 | 120 | print("#"*30) 121 | 122 | return is_acuqired, datfile 123 | 124 | -------------------------------------------------------------------------------- /lib/model.py: -------------------------------------------------------------------------------- 1 | """ 2 | Function for obtain the models from nns library 3 | 4 | @yuningw 5 | """ 6 | 7 | import torch 8 | from torch import nn 9 | 10 | def get_vae(latent_dim): 11 | """ 12 | A function to obtain the beta-VAE model 13 | 14 | Args: 15 | latent_dim : (int) The latent-variable dimension 16 | 17 | Returns: 18 | model : (torch.nn.Module)The beta-VAE model 19 | """ 20 | 21 | from nns.beta_vae import VAE 22 | 23 | return VAE(latent_dim) 24 | 25 | 26 | 27 | def get_predictors(name): 28 | """ 29 | A function to obatin the model for temporal dynamics predictor in latent space 30 | 31 | Args: 32 | name : (str) Name between; attn, self and lstm 33 | 34 | Returns: 35 | model : (torch.nn.Module) The model 36 | 37 | filename : (str) The filename to save the model 38 | 39 | config : (class) The configuration of the model 40 | """ 41 | 42 | assert(name=="self" or name == "easy" or name == "lstm"), print("ERROR: Given Name is not Valid!") 43 | 44 | if name == "easy": 45 | from configs.easyAttn import easyAttn_config as cfg 46 | from configs.nomenclature import Make_Transformer_Name 47 | from nns.transformer import easyTransformerEncoder 48 | try: 49 | model = easyTransformerEncoder( d_input = cfg.in_dim, 50 | d_output = cfg.next_step, 51 | seqLen = cfg.nmode, 52 | d_proj = cfg.time_proj, 53 | d_model = cfg.d_model, 54 | d_ff = cfg.proj_dim, 55 | num_head = cfg.num_head, 56 | num_layer = cfg.num_block,) 57 | except: 58 | print("ERROR: Parameter NOT MATCHED!") 59 | exit() 60 | 61 | filename = Make_Transformer_Name(cfg) 62 | print(f"Easy-Attention-based Transformer has been generated") 63 | print(f"FileName: {filename}") 64 | return model, filename, cfg 65 | 66 | elif name == "self": 67 | from configs.selfAttn import selfAttn_config as cfg 68 | from configs.nomenclature import Make_Transformer_Name 69 | from nns.transformer import EmbedTransformerEncoder 70 | try: 71 | model = EmbedTransformerEncoder(d_input = cfg.in_dim, 72 | d_output= cfg.next_step, 73 | n_mode = cfg.nmode, 74 | d_proj = cfg.time_proj, 75 | d_model = cfg.d_model, 76 | d_ff = cfg.proj_dim, 77 | num_head = cfg.num_head, 78 | num_layer = cfg.num_block) 79 | except: 80 | print("ERROR: Parameter NOT MATCHED!") 81 | exit() 82 | 83 | filename = Make_Transformer_Name(cfg) 84 | print(f"Self-Attention-based Transformer has been generated") 85 | print(f"FileName: {filename}") 86 | return model, filename, cfg 87 | 88 | 89 | elif name =="lstm": 90 | from configs.lstm import lstm_config as cfg 91 | from configs.nomenclature import Make_LSTM_Name 92 | from nns.RNNs import LSTMs 93 | try: 94 | model = LSTMs( 95 | d_input = cfg.in_dim, 96 | d_model = cfg.d_model, 97 | nmode = cfg.nmode, 98 | embed = cfg.embed, 99 | hidden_size = cfg.hidden_size, 100 | num_layer = cfg.num_layer, 101 | is_output = cfg.is_output, out_dim= cfg.next_step, out_act= cfg.out_act 102 | ) 103 | except: 104 | print("ERROR: Parameter NOT MATCHED!") 105 | exit() 106 | 107 | filename = Make_LSTM_Name(cfg) 108 | print(f"LSTM has been generated") 109 | print(f"FileName: {filename}") 110 | return model, filename, cfg 111 | 112 | else: 113 | 114 | print(f"Error: There is no options!") 115 | exit() 116 | 117 | 118 | 119 | -------------------------------------------------------------------------------- /lib/pp_space.py: -------------------------------------------------------------------------------- 1 | """ 2 | Post-processing and analysis algorithm for beta-VAE in latent space and physic space 3 | 4 | author : @alsora 5 | editting: @yuningw 6 | 7 | """ 8 | 9 | import torch 10 | import numpy as np 11 | import h5py 12 | from lib.init import pathsBib 13 | 14 | ################################ 15 | ### Main programme for spatial analysis 16 | ############################### 17 | def spatial_Mode( fname, 18 | model, 19 | latent_dim, 20 | train_data, 21 | test_data, 22 | dataset_train, 23 | dataset_test, 24 | mean, std, 25 | device, 26 | if_order = True, 27 | if_nlmode = True, 28 | if_Ecumt = True, 29 | if_Ek_t = True, 30 | ): 31 | """ 32 | The main function for spatial mode analysis and generate the dataset 33 | 34 | Args: 35 | 36 | fname : (str) The file name 37 | 38 | latent_dim : (int) The latent dimension 39 | 40 | train_data : (NumpyArray) Dataset for training 41 | 42 | test_data : (NumpyArray) Dataset for test 43 | 44 | dataset_train : (dataloader) DataLoader for training data 45 | 46 | dataset_test : (dataloader) DataLoader for test data 47 | 48 | mean : (NumpyArray) The mean of flow database 49 | 50 | std : (NumpyArray) The std of flow database 51 | 52 | device : (str) The device for the computation 53 | 54 | order : (NumpyArray) A array which contains the ranking results 55 | 56 | Ecum : (NumpyArray) accumlative Energy obtained for each mode 57 | 58 | Ecum_test : (NumpyArray) accumlative Energy obtained for each mode 59 | 60 | NLvalues : (NumpyArray) The used range of value 61 | 62 | NLmodes : (NumpyArray) The non-linear spatial mode 63 | 64 | Ek_t : (List) A list of enery of each snapshot in dataset 65 | 66 | if_order : (bool) IF ranking the mode 67 | 68 | if_nlmode : (bool) IF generate non-linear mode 69 | 70 | if_Ecumt : (bool) IF compute accumulative energy 71 | 72 | if_Ek_t : (bool) IF compute evolution of energy 73 | 74 | Returns: 75 | 76 | if_save : (bool) If successfully save file 77 | 78 | """ 79 | 80 | print(f"INFO: Start spatial mode generating") 81 | if if_order: 82 | order, Ecum = get_order(model, latent_dim, 83 | train_data, 84 | dataset_train, 85 | std, device) 86 | print(f"INFO: RANKING DONE") 87 | else: 88 | order = None 89 | Ecum = None 90 | 91 | 92 | if if_nlmode: 93 | NLvalues, NLmodes = getNLmodes(model, order[0], latent_dim, device) 94 | print("INFO: Non-linear mode generated") 95 | else: 96 | NLmodes = None 97 | NLvalues = None 98 | 99 | if if_Ecumt: 100 | 101 | Ecum_test = get_EcumTest(model, latent_dim, test_data, dataset_test, std, device, order) 102 | print('INFO: Test E_cum generated') 103 | else: 104 | Ecum_test = None 105 | 106 | if if_Ek_t: 107 | Ek_t = get_Ek_t(model=model, data=test_data, device=device) 108 | 109 | else: 110 | Ek_t = None 111 | 112 | is_save = createModesFile(fname, model, latent_dim, 113 | dataset_train, dataset_test, 114 | mean, std, device, 115 | order, Ecum, Ecum_test, 116 | NLvalues, NLmodes,Ek_t) 117 | 118 | if is_save: print("INFO: Successfuly DONE!") 119 | 120 | return is_save 121 | 122 | ################################ 123 | ### Basic function for using VAE 124 | ############################### 125 | 126 | #-------------------------------------------------------- 127 | def encode(model, data, device): 128 | """ 129 | Use encoder to compress flow field into latent space 130 | Args: 131 | model : (nn.Module) Pytorch module for beta-VA 132 | 133 | data : (DataLoader) DataLoader of data to be encoded 134 | 135 | device : (str) The device for the computation 136 | 137 | Returns: 138 | 139 | means : (NumpyArray) The mu obtained in latent space 140 | 141 | logvars : (NumpyArray) The sigma obtained in latent space 142 | """ 143 | mean_list = [] 144 | logvar_list = [] 145 | with torch.no_grad(): 146 | for batch in data: 147 | batch = batch.to(device, non_blocking=True) 148 | mean, logvariance = torch.chunk(model.encoder(batch), 2, dim=1) 149 | 150 | mean_list.append(mean.cpu().numpy()) 151 | logvar_list.append(logvariance.cpu().numpy()) 152 | 153 | means = np.concatenate(mean_list, axis=0) 154 | logvars = np.concatenate(logvar_list, axis=0) 155 | 156 | return means, logvars 157 | 158 | 159 | #-------------------------------------------------------- 160 | def decode(model, data, device): 161 | """ 162 | Use decoder to reconstruct flow field back to physical space 163 | 164 | Args: 165 | model : (nn.Module) Pytorch module for beta-VA 166 | 167 | data : (NumpyArray) The latent vectors required to be reconstructed 168 | 169 | device : (str) The device for the computation 170 | 171 | Returns: 172 | 173 | rec : (NumpyArray) The reconstruction of the flow fields. 174 | 175 | """ 176 | dataset = torch.utils.data.DataLoader(dataset=torch.from_numpy(data), batch_size=512, 177 | shuffle=False, num_workers=2) 178 | rec_list = [] 179 | with torch.no_grad(): 180 | for batch in dataset: 181 | batch = batch.to(device) 182 | rec_list.append(model.decoder(batch).cpu().numpy()) 183 | 184 | return np.concatenate(rec_list, axis=0) 185 | 186 | 187 | 188 | def get_samples(model, dataset_train, dataset_test, device): 189 | """ 190 | 191 | A function for quickly obtain a restructed flow field for the propose of testing or visualisation 192 | 193 | We obtain snapshot through training and test data, respectively 194 | 195 | Args: 196 | model : (nn.Module) Pytorch module for beta-VA 197 | 198 | dataset_train : (DataLoader) DataLoader of training data 199 | 200 | dataset_test : (DataLoader) DataLoader of test data 201 | 202 | device : (str) The device for the computation 203 | 204 | Returns: 205 | 206 | rec_train : (NumpyArray) The reconstruction from training dataset. 207 | 208 | rec_test : (NumpyArray) The reconstruction from test dataset. 209 | 210 | true_train : (NumpyArray) The corresponding ground truth from training dataset. 211 | 212 | true_test : (NumpyArray) The corresponding ground truth from test dataset. 213 | """ 214 | with torch.no_grad(): 215 | if dataset_train != None: 216 | for batch_train in dataset_train: 217 | batch_train = batch_train.to(device, non_blocking=True) 218 | rec_train, _, _ = model(batch_train) 219 | 220 | rec_train = rec_train.cpu().numpy()[-1] 221 | true_train = batch_train.cpu().numpy()[-1] 222 | 223 | break 224 | else: 225 | rec_train = None 226 | true_train = None 227 | 228 | if dataset_test != None: 229 | for batch_test in dataset_test: 230 | batch_test = batch_test.to(device, non_blocking=True) 231 | rec_test, _, _ = model(batch_test) 232 | 233 | rec_test = rec_test.cpu().numpy()[-1] 234 | true_test = batch_test.cpu().numpy()[-1] 235 | 236 | break 237 | else: 238 | rec_test = None 239 | true_test = None 240 | 241 | return rec_train, rec_test, true_train, true_test 242 | 243 | 244 | ################################ 245 | ### Spatial-mode generate and analysis 246 | ############################### 247 | 248 | #-------------------------------------------------------- 249 | def calcmode(model, latent_dim, mode, device): 250 | """ 251 | 252 | Generate the non-linear mode with unit vector 253 | 254 | Args: 255 | model : (nn.Module) Pytorch module for beta-VA 256 | 257 | latent_dim : (int) Latent-dimension adpot for beta-VA 258 | 259 | mode : (int) The indice of the mode zi 260 | 261 | Returns: 262 | 263 | mode : (NumpyArray) The spatial mode for zi 264 | """ 265 | 266 | z_sample = np.zeros((1, latent_dim), dtype=np.float32) 267 | 268 | z_sample[:, mode] = 1 269 | 270 | with torch.no_grad(): 271 | mode = model.decoder(torch.from_numpy(z_sample).to(device)).cpu().numpy() 272 | 273 | return mode 274 | 275 | 276 | #-------------------------------------------------------- 277 | def get_spatial_modes(model, latent_dim, device): 278 | 279 | """ 280 | Algorithm for optain the spatial mode from beta-VAE Decoder. 281 | For latent variable i, We use the unit vector v (where vi = 1) as the input to obtain the spatial mode for each latent variables 282 | Also, we compute the spatial mode use zeros vectors as input 283 | 284 | Args: 285 | 286 | model : (torch.nn.Module) The beta-VAE model 287 | 288 | latent_dim : (int) The latent dimension we employed 289 | 290 | device : (str) The device for the computation 291 | 292 | Returns: 293 | 294 | modes : The 295 | 296 | """ 297 | 298 | 299 | with torch.no_grad(): 300 | zero_output = model.decoder(torch.from_numpy(np.zeros((1, latent_dim), dtype=np.float32)).to(device)).cpu().numpy() 301 | 302 | modes = np.zeros((latent_dim, zero_output.shape[1], zero_output.shape[2], zero_output.shape[3])) 303 | 304 | for mode in range(latent_dim): 305 | modes[mode, :, :, :] = calcmode(model, latent_dim, mode, device) 306 | 307 | return zero_output, modes 308 | 309 | 310 | #-------------------------------------------------------- 311 | def getNLmodes(model, mode, latent_dim, device): 312 | """ 313 | Algorithm for optain single spatial mode from beta-VAE Decoder. 314 | 315 | For latent variable i, We use the vector v (where vi = 1) with a value within a range 316 | as the input to obtain the spatial mode for each latent variables 317 | 318 | Args: 319 | 320 | model : (torch.nn.Module) The beta-VAE model 321 | 322 | mode : (int) The indice of the mode zi 323 | 324 | latent_dim : (int) The latent dimension we employed 325 | 326 | device : (str) The device for the computation 327 | 328 | Returns: 329 | 330 | NLvalues : (NumpyArray) The used range of value 331 | 332 | NLmodes : (NumpyArray) The non-linear spatial mode 333 | 334 | """ 335 | 336 | zero_output = decode(model, np.zeros((1, latent_dim), dtype=np.float32), device) 337 | NLvalues = np.arange(-2, 2.1, .1) 338 | NLmodes = np.zeros((NLvalues.shape[0], zero_output.shape[1], zero_output.shape[2], zero_output.shape[3]), 339 | dtype=np.float32) 340 | 341 | for idx, value in enumerate(NLvalues): 342 | latent = np.zeros((1, latent_dim), dtype=np.float32) 343 | latent[0, mode] = value 344 | NLmodes[idx,:,:,:] = decode(model, latent, device) 345 | 346 | return NLvalues, NLmodes 347 | 348 | 349 | #-------------------------------------------------------- 350 | def get_order( model, latent_dim, 351 | data, dataset, 352 | std, device): 353 | """ 354 | Algorithm for ranking the obtained spatial modes according to the yield accumlated energy level 355 | For more detail please check the paper 356 | 357 | Args: 358 | 359 | model : (torch.nn.Module) The beta-VAE model 360 | 361 | 362 | latent_dim : (int) The latent dimension we employed 363 | 364 | data : (NumpyArray) The flow database 365 | 366 | dataset : (torch.Dataloader) The dataloader of the flow data 367 | 368 | std : (NumpyArray) The std of flow database 369 | 370 | device : (str) The device for the computation 371 | 372 | Returns: 373 | 374 | m : (NumpyArray) The ranking result (order) of each mode 375 | 376 | Ecum : (NumpyArray) accumlative Energy obtained for each mode 377 | 378 | """ 379 | import time 380 | import numpy as np 381 | print('#'*30) 382 | print('Ordering modes') 383 | 384 | modes, _ = encode(model, dataset, device) 385 | 386 | print(modes.shape) 387 | u = data[:, :, :, :] * std[:, :, :, :] 388 | m = np.zeros(latent_dim, dtype=int) 389 | n = np.arange(latent_dim) 390 | 391 | Ecum = [] 392 | partialModes = np.zeros_like(modes, dtype=np.float32) 393 | 394 | for i in range(latent_dim): 395 | Eks = [] 396 | for j in n: # for mode in remaining modes 397 | start = time.time() 398 | print(m[:i], j, end="") 399 | partialModes *= 0 400 | partialModes[:, m[:i]] = modes[:, m[:i]] 401 | partialModes[:, j] = modes[:, j] 402 | u_pred = decode(model, partialModes, device) * std 403 | Eks.append(get_Ek(u, u_pred)) 404 | elapsed = time.time() - start 405 | print(f' : Ek={Eks[-1]:.4f}, elapsed: {elapsed:.2f}s') 406 | Eks = np.array(Eks).squeeze() 407 | ind = n[np.argmax(Eks)] 408 | m[i] = ind 409 | n = np.delete(n, np.argmax(Eks)) 410 | Ecum.append(np.max(Eks)) 411 | print('Adding: ', ind, ', Ek: ', np.max(Eks)) 412 | print('#'*30) 413 | Ecum = np.array(Ecum) 414 | print(f"Rank finished, the rank is {m}") 415 | print(f"Cumulative Ek is {Ecum}") 416 | 417 | return np.array(m), Ecum 418 | 419 | 420 | 421 | ################################ 422 | ### Assessment on Energy 423 | ############################### 424 | #-------------------------------------------------------- 425 | def get_Ek(original, rec): 426 | 427 | """ 428 | Calculate energy percentage reconstructed 429 | 430 | Args: 431 | original : (NumpyArray) The ground truth 432 | 433 | rec : (NumpyArray) The reconstruction from decoder 434 | 435 | Returns: 436 | 437 | The energy percentage for construction. Note that it is the Ek/100 !! 438 | """ 439 | 440 | import numpy as np 441 | 442 | TKE_real = original[:, 0, :, :] ** 2 + original[:, 1, :, :] ** 2 443 | 444 | u_rec = rec[:, 0, :, :] 445 | v_rec = rec[:, 1, :, :] 446 | 447 | return 1 - np.sum((original[:, 0, :, :] - u_rec) ** 2 + (original[:, 1, :, :] - v_rec) ** 2) / np.sum(TKE_real) 448 | 449 | 450 | 451 | #-------------------------------------------------------- 452 | def get_Ek_t(model, data, device): 453 | """ 454 | 455 | Get the Reconstructed energy for snapshots 456 | 457 | Args: 458 | 459 | model : (torch.nn.Module) The beta-VAE model 460 | 461 | data : (NumpyArray) The flow database 462 | 463 | device : (str) The device for the computation 464 | 465 | Returns: 466 | 467 | Ek_t : (List) A list of enery of each snapshot in dataset 468 | 469 | 470 | """ 471 | 472 | dataloader = torch.utils.data.DataLoader(dataset=torch.from_numpy(data), 473 | batch_size=1, 474 | shuffle=False, 475 | pin_memory=True, 476 | num_workers=2) 477 | 478 | rec_list = [] 479 | with torch.no_grad(): 480 | for batch in dataloader: 481 | batch = batch.to(device) 482 | rec, _, _ = model(batch) 483 | rec_list.append(rec.cpu().numpy()) 484 | 485 | rec = np.concatenate(rec_list, axis=0) 486 | 487 | print(rec.shape) 488 | print(data.shape) 489 | 490 | Ek_t = np. zeros((rec.shape[0])) 491 | for i in range(rec.shape[0]): 492 | Ek_t[i] = get_Ek(data[np.newaxis, i], rec[np.newaxis, i]) 493 | 494 | return Ek_t 495 | 496 | 497 | 498 | 499 | 500 | #-------------------------------------------------------- 501 | def get_EcumTest(model, latent_dim, data, dataset, std, device, order): 502 | 503 | """ 504 | Get the accumlative energy of test database 505 | 506 | Args: 507 | 508 | model : (torch.nn.Module) The beta-VAE model 509 | 510 | 511 | latent_dim : (int) The latent dimension we employed 512 | 513 | data : (NumpyArray) The flow database 514 | 515 | dataset : (torch.Dataloader) The dataloader of the flow data 516 | 517 | std : (NumpyArray) The std of flow database 518 | 519 | device : (str) The device for the computation 520 | 521 | order : (NumpyArray) A array which contains the ranking results 522 | 523 | Returns: 524 | 525 | Ecum : (NumpyArray) accumlative Energy obtained for each mode 526 | 527 | """ 528 | 529 | modes, _ = encode(model, dataset, device) 530 | print(modes.shape) 531 | u = data[:, :, :, :] * std[:, :, :, :] 532 | 533 | Ecum = [] 534 | for i in range(latent_dim): 535 | partialModes = np.zeros_like(modes, dtype=np.float32) 536 | partialModes[:, order[:i+1]] = modes[:, order[:i+1]] 537 | u_pred = decode(model, partialModes, device) 538 | u_pred *= std[:, :, :, :] 539 | Ecum.append(get_Ek(u, u_pred)) 540 | print(order[:i+1], Ecum[-1]) 541 | 542 | return np.array(Ecum) 543 | 544 | 545 | 546 | 547 | 548 | ################################ 549 | ### I/O 550 | ############################### 551 | #-------------------------------------------------------- 552 | def createModesFile(fname, 553 | model, 554 | latent_dim, 555 | dataset_train, dataset_test, 556 | mean, std, 557 | device, 558 | order, 559 | Ecum, Ecum_test, 560 | NLvalues, NLmodes, 561 | Ek_t): 562 | """ 563 | 564 | Function for integrating all the obtained results and save it as fname 565 | 566 | Args: 567 | 568 | fname : (str) The file name 569 | 570 | latent_dim : (int) The latent dimension 571 | 572 | dataset_train : (dataloader) DataLoader for training data 573 | 574 | dataset_test : (dataloader) DataLoader for test data 575 | 576 | mean : (NumpyArray) The mean of flow database 577 | 578 | std : (NumpyArray) The std of flow database 579 | 580 | device : (str) The device for the computation 581 | 582 | order : (NumpyArray) A array which contains the ranking results 583 | 584 | Ecum : (NumpyArray) accumlative Energy obtained for each mode 585 | 586 | Ecum_test : (NumpyArray) accumlative Energy obtained for each mode 587 | 588 | NLvalues : (NumpyArray) The used range of value 589 | 590 | NLmodes : (NumpyArray) The non-linear spatial mode 591 | 592 | Ek_t : (List) A list of enery of each snapshot in dataset 593 | 594 | Returns: 595 | 596 | is_save : (bool) 597 | 598 | """ 599 | 600 | if_save = False 601 | 602 | print(f"Start post-processing") 603 | 604 | 605 | means_train, stds_train = encode(model, dataset_train, device) 606 | print('INFO: Latent Variable Train Generated') 607 | means_test, stds_test = encode(model, dataset_test, device) 608 | print('INFO: Latent Variable Test Generated') 609 | 610 | zero_output, modes = get_spatial_modes(model, latent_dim, device) 611 | print('INFO: Spatial mode generated') 612 | 613 | if order is None: 614 | order = np.arange(latent_dim) 615 | 616 | 617 | with h5py.File(fname + ".hdf5", 'w') as f: 618 | 619 | # Mean velocity of flow data 620 | f.create_dataset('mean', data=mean) 621 | # std of flow data 622 | f.create_dataset('std', data=std) 623 | 624 | # mu 625 | f.create_dataset('vector', data=means_train) 626 | f.create_dataset('vector_test', data=means_test) 627 | 628 | # Sigma 629 | f.create_dataset('stds_vector', data=stds_train) 630 | f.create_dataset('stds_vector_test', data=stds_test) 631 | 632 | # Spatial modes: Unit vector input 633 | f.create_dataset('modes', data=modes) 634 | # Sptail modes: zeros vector input 635 | f.create_dataset('zero_output', data=zero_output) 636 | # Spatial modes: Non-linear value adopted 637 | f.create_dataset('NLvalues', data=NLvalues) 638 | # Spatial modes: Non-linear modes 639 | f.create_dataset('NLmodes', data=NLmodes) 640 | 641 | # Ranking results for spatial modes 642 | f.create_dataset('order', data=order) 643 | 644 | #Assessments: Accumlative energy 645 | f.create_dataset('Ecum', data=Ecum) 646 | #Assessments: Accumlative energy on test dataset 647 | f.create_dataset('Ecum_test', data=Ecum_test) 648 | #Assessments: E_k for each single snapshots 649 | f.create_dataset('Ek_t', data=Ek_t) 650 | 651 | f.close() 652 | 653 | # Extra file for time-series prediction. 654 | with h5py.File(pathsBib.data_path + 'latent_data' + ".h5py", 'w') as f: 655 | # mu 656 | f.create_dataset('vector', data=means_train) 657 | f.create_dataset('vector_test', data=means_test) 658 | f.close() 659 | 660 | 661 | if_save = True 662 | 663 | 664 | print(f"INFO: Post-processing results has been saved as dataset: {fname}") 665 | 666 | return if_save -------------------------------------------------------------------------------- /lib/pp_time.py: -------------------------------------------------------------------------------- 1 | """ 2 | Post-processing and evaluation for time series prediction 3 | """ 4 | 5 | ##################################### 6 | # Assessment of temporal dynamic prediction 7 | #################################### 8 | 9 | #-------------------------------------------------------- 10 | def make_Prediction(test_data, model, 11 | device, in_dim, next_step): 12 | """ 13 | Function for generat the prediction data 14 | 15 | Args: 16 | test_data : A numpy array of test data, with shape of [Ntime, Nmode] 17 | model : A torch.nn.Module object as model 18 | device : String of name of device 19 | in_dim : Integar of TimeDelay size 20 | next_step : Future time step to predict 21 | 22 | Returns: 23 | preds : A numpy array of the prediction 24 | """ 25 | from copy import deepcopy 26 | import torch 27 | from tqdm import tqdm 28 | 29 | 30 | 31 | model.eval() 32 | model.to(device) 33 | Preds = deepcopy(test_data) 34 | seq_len = max([Preds.shape[0],Preds.shape[1]]) 35 | print(f"The sequence length = {seq_len}") 36 | 37 | for i in tqdm(range(in_dim,seq_len-next_step)): 38 | 39 | feature = Preds[None,i-in_dim:i,:] 40 | 41 | x = torch.from_numpy(feature) 42 | x = x.float().to(device) 43 | pred = model(x) 44 | 45 | pred = pred.cpu().detach().numpy() 46 | 47 | Preds[i:i+next_step,:] = pred[0,:,:] 48 | 49 | return Preds 50 | 51 | #-------------------------------------------------------- 52 | def Sliding_Window_Error(test_data, 53 | model, device, 54 | in_dim,window = 5): 55 | """ 56 | Compute the sliding window error on test dataset 57 | Args: 58 | test_data : A numpy array of test data [Ntime, Nmode] 59 | model : A torch.nn.Module as model 60 | device : String of name of device 61 | in_dim : Integar of input dimension 62 | window : The size of window for evaluation, default = 100 63 | 64 | Returns: 65 | error_l2 : A numpy arrary of sliding window error, shape = [window,] 66 | 67 | """ 68 | import torch 69 | import copy 70 | from tqdm import tqdm 71 | import numpy as np 72 | 73 | def l2norm(predictions, targets): 74 | return np.sqrt( np.sum( (predictions - targets) ** 2, axis=1 ) ) 75 | 76 | model.to(device) 77 | model.eval() 78 | 79 | SeqLen = test_data.shape[0] 80 | error = None 81 | for init in tqdm(range(in_dim,SeqLen-window,1)): 82 | temporalModes_pred = copy.deepcopy(test_data) 83 | 84 | for timestep in range(init, init+window): 85 | 86 | data = temporalModes_pred[None, (timestep-in_dim):timestep, :] 87 | feature = torch.from_numpy(data).float().to(device) 88 | pred = model(feature) 89 | temporalModes_pred[timestep,:] = pred[0].detach().cpu().numpy() 90 | 91 | if error is None: 92 | error = l2norm(temporalModes_pred[init:init+window,:], test_data[init:init+window,:]) 93 | n = 1 94 | else: 95 | error = error + l2norm(temporalModes_pred[init:init+window,:], test_data[init:init+window,:]) 96 | n += 1 97 | 98 | print(n) 99 | error_l2 = error / n 100 | print(error.shape) 101 | 102 | return error_l2 103 | 104 | #-------------------------------------------------------- 105 | def l2Norm_Error(Pred, Truth): 106 | """ 107 | Compute the l2-norm proportional error for the prediction of temporal evolution 108 | Args: 109 | Pred : Numpy array has shape of [Ntime, Nmode] 110 | Truth : Numpy array has shape of [Ntime, Nmode] 111 | 112 | Returns: 113 | error : A vector of error for each mode 114 | """ 115 | 116 | import sys 117 | import numpy as np 118 | from numpy.linalg import norm 119 | 120 | 121 | if Pred.shape != Truth.shape: 122 | print("The size of array does not match!") 123 | sys.exit() 124 | 125 | error = norm((Pred-Truth),axis=1)\ 126 | /(norm(Truth,axis=1)) 127 | 128 | return error.mean()*100 129 | 130 | 131 | #-------------------------------------------------------- 132 | def make_physical_prediction(vae, 133 | pred_latent, 134 | true_latent, 135 | device): 136 | """ 137 | Reconstruct the latent-space prediction to the physical space 138 | 139 | Args: 140 | 141 | vae : (lib.runners.vaeRunner) The module for employed VAE 142 | 143 | pred_latent : (Numpyrray) The latent-space prediction 144 | 145 | true_latent : (NumpyArray) The latent-space ground truth from VAE 146 | 147 | device : (str) The device for computation 148 | 149 | Returns: 150 | 151 | rec_vae : (NumpyArray) Reconstruction from reference latent variable 152 | 153 | rec_pred : (NumpyArray) Reconstruction from the predicted latent variable 154 | """ 155 | import numpy as np 156 | from lib.pp_space import decode 157 | 158 | rec_vae = decode(vae.model,true_latent,device) 159 | 160 | rec_pred = decode(vae.model,pred_latent,device) 161 | 162 | return rec_vae, rec_pred 163 | 164 | 165 | 166 | 167 | ##################################### 168 | # Poincare Map for chaotic nature analysis 169 | #################################### 170 | 171 | """ 172 | Implement Poincare Maps on test data 173 | 174 | @yuningw 175 | """ 176 | 177 | #-------------------------------------------------------- 178 | def Zero_Cross(data,postive_dir = True): 179 | """ 180 | Function to find the cross section betwee positive and negative of a 1D Vector 181 | Args: 182 | data : 1D numpy arrary, object data 183 | postive_dir : bool, Find the positive direction or negative 184 | 185 | Returns: 186 | cross_idx : Indices of the cross-section points 187 | x_values : Estimation of the position of the zero crossing 188 | within the interval between crossings_idx 189 | and crossings_idx_next 190 | 191 | """ 192 | import numpy as np 193 | zero_crossings = np.where( np.diff(np.signbit(data)) ) 194 | if (postive_dir): 195 | wherePos = np.where( np.diff(data) > 0 ) 196 | else: 197 | wherePos = np.where( np.diff(data) < 0 ) 198 | 199 | cross_idx = np.intersect1d(zero_crossings, wherePos) 200 | cross_idx_next = cross_idx +1 201 | 202 | x_values = cross_idx - data[list(cross_idx)]/\ 203 | (data[list(cross_idx_next)] - data[list(cross_idx)]) 204 | 205 | 206 | return cross_idx, x_values 207 | 208 | 209 | 210 | #-------------------------------------------------------- 211 | def Intersection(data,planeNo = 0,postive_dir = True): 212 | """ 213 | Compute the intersections of time-series data w.r.t each temporal mode 214 | 215 | Args: 216 | data : A 2D numpy array has shape of [Ntime, Nmodes] 217 | planeNo : Integar, the No. of plane to compute the intersections 218 | postive_dir : bool, choose which direction 219 | 220 | Returns: 221 | InterSec : The intersection data in numpy array 222 | """ 223 | import numpy as np 224 | import sys 225 | if len(data.shape) !=2: 226 | print("The data should have 2 dimensions") 227 | sys.exit() 228 | 229 | SeqLen, Nmode = data.shape[0],data.shape[-1] 230 | zero_cross, x_value = Zero_Cross(data = data[:,planeNo], 231 | postive_dir = postive_dir) 232 | 233 | # Create InterSec to store the results 234 | InterSec = np.zeros((zero_cross.shape[0],Nmode)) 235 | for mode in range(0,Nmode): 236 | InterSec[:,mode] = np.interp(x_value, np.arange(SeqLen), data[:,mode]) 237 | 238 | return InterSec 239 | 240 | 241 | 242 | #-------------------------------------------------------- 243 | def PDF(InterSecX,InterSecY, 244 | xmin = -1,xmax = 1,x_grid = 50, 245 | ymin = -1,ymax = 1,y_grid = 50): 246 | 247 | """ 248 | Compute the joint PDF of X and Y 249 | Args: 250 | InterSecX : numpy array of data 1 251 | InterSecY : numpy array of data 2 252 | 253 | xmin, xmax, x_grid : The limitation of InterSecX and number of grid to be plot for contour 254 | ymin, ymax, y_grid : The limitation of InterSecY and number of grid to be plot for contour 255 | 256 | Returns: 257 | xx, yy: The meshgrid of InterSecX and InterSecY according to the limitation and number of grids 258 | pdf : The joint pdf of InterSecX and InterSecY 259 | """ 260 | import numpy as np 261 | import scipy.stats as st 262 | # Create meshgrid acorrding 263 | xx, yy = np.mgrid[xmin:xmax:1j*x_grid, ymin:ymax:1j*y_grid] 264 | 265 | 266 | positions = np.vstack([xx.ravel(), yy.ravel()]) 267 | values = np.vstack([InterSecX, InterSecY]) 268 | kernel = st.gaussian_kde(values) 269 | pdf = np.reshape(kernel(positions).T, xx.shape) 270 | 271 | return xx,yy,pdf -------------------------------------------------------------------------------- /lib/runners.py: -------------------------------------------------------------------------------- 1 | """ 2 | Runners for the VAE and temporal-dynamic prediction in latent space 3 | @yuningw 4 | """ 5 | 6 | import os 7 | import time 8 | from pathlib import Path 9 | import h5py 10 | import numpy as np 11 | import torch 12 | from torch import nn 13 | 14 | from lib.init import pathsBib 15 | from lib.train import * 16 | from lib.model import * 17 | from lib.pp_time import * 18 | from lib.pp_space import spatial_Mode 19 | from lib.datas import * 20 | 21 | os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" 22 | 23 | 24 | 25 | #################################################### 26 | ### RUNNER for beta-VAE 27 | #################################################### 28 | 29 | class vaeRunner(nn.Module): 30 | def __init__(self, device, datafile) -> None: 31 | """ 32 | A runner for beta-VAE 33 | 34 | Args: 35 | 36 | device : (Str) The device going to use 37 | 38 | datafile : (Str) Path of training data 39 | """ 40 | 41 | from configs.vae import VAE_config as cfg 42 | from configs.nomenclature import Name_VAE 43 | 44 | super(vaeRunner,self).__init__() 45 | print("#"*30) 46 | 47 | self.config = cfg 48 | self.filename = Name_VAE(self.config) 49 | 50 | self.datafile = datafile 51 | 52 | self.device = device 53 | 54 | self.model = get_vae(self.config.latent_dim) 55 | 56 | self.model.to(device) 57 | 58 | self.fmat = '.pth.tar' 59 | 60 | print(f"INIT betaVAE, device: {device}") 61 | print(f"Case Name:\n {self.filename}") 62 | #------------------------------------------------- 63 | def run(self): 64 | self.train() 65 | self.infer(model_type='final') 66 | 67 | #------------------------------------------------- 68 | def train(self): 69 | print("#"*30) 70 | print("INFO: Start Training ") 71 | self.get_data() 72 | self.compile() 73 | self.fit() 74 | self.train_dl = None 75 | self.val_dl = None 76 | print(f"INFO: Training finished, cleaned the data loader") 77 | print("#"*30) 78 | 79 | #------------------------------------------------- 80 | def infer(self, model_type): 81 | print("#"*30) 82 | self.load_pretrain_model(model_type=model_type) 83 | print("INFO: Model has been loaded!") 84 | 85 | self.get_test_data() 86 | print("INFO: test data has been loaded!") 87 | self.post_process() 88 | 89 | print(f"INFO: Inference ended!") 90 | print("#"*30) 91 | 92 | 93 | #------------------------------------------------- 94 | 95 | def get_data(self): 96 | """ 97 | 98 | Generate the DataLoader for training 99 | 100 | """ 101 | 102 | # datafile = 103 | 104 | u_scaled, self.mean, self.std = loadData(self.datafile) 105 | ## Down-Sample the data with frequency 106 | ## since we already down-sampled it for database, we skip it here 107 | 108 | u_scaled = u_scaled[::self.config.downsample] 109 | n_total = u_scaled.shape[0] 110 | self.n_train = n_total - self.config.n_test 111 | print(f"INFO: Data Summary: N train: {self.n_train:d}," + \ 112 | f"N test: {self.config.n_test:d},"+\ 113 | f"N total {n_total:d}") 114 | 115 | self.train_dl, self.val_dl = get_vae_DataLoader( d_train=u_scaled, 116 | n_train=self.n_train, 117 | device= self.device, 118 | batch_size= self.config.batch_size) 119 | print( f"INFO: Dataloader generated, Num train batch = {len(self.train_dl)} \n" +\ 120 | f"Num val batch = {len(self.val_dl)}") 121 | 122 | 123 | 124 | 125 | #------------------------------------------------- 126 | def compile(self): 127 | """ 128 | 129 | Compile the optimiser, schedulers and loss function for training 130 | 131 | 132 | """ 133 | 134 | from torch.optim import lr_scheduler 135 | 136 | print("#"*30) 137 | print(f"INFO: Start Compiling") 138 | 139 | encoder_params = list(self.model.encoder.parameters()) 140 | decoder_params = list(self.model.decoder.parameters()) 141 | 142 | # get optimizer 143 | self.opt = torch.optim.Adam( 144 | [ {'params': encoder_params, 'weight_decay': self.config.encWdecay}, 145 | {'params': decoder_params, 'weight_decay': self.config.decWdecay}], 146 | lr=self.config.lr, weight_decay=0) 147 | 148 | self.opt_sch = lr_scheduler.OneCycleLR(self.opt, 149 | max_lr=self.config.lr, 150 | epochs=self.config.epochs, # total_steps = epochs * steps_per_epoch 151 | steps_per_epoch=len(self.train_dl), 152 | div_factor=2, 153 | final_div_factor=self.config.lr/self.config.lr_end, 154 | pct_start=0.2) 155 | 156 | self.beta_sch = betaScheduler( startvalue =self.config.beta_init, 157 | endvalue =self.config.beta, 158 | warmup =self.config.beta_warmup) 159 | 160 | print(f"INFO: Compiling Finished!") 161 | 162 | 163 | #------------------------------------------------- 164 | 165 | def fit(self): 166 | """ 167 | 168 | Training beta-VAE 169 | 170 | """ 171 | from torch.utils.tensorboard import SummaryWriter 172 | from utils.io import save_checkpoint 173 | 174 | print(f"Training {self.filename}") 175 | logger = SummaryWriter(log_dir=pathsBib.log_path + self.filename) 176 | 177 | bestloss = 1e6 178 | loss = 1e6 179 | 180 | for epoch in range(1, self.config.epochs + 1): 181 | self.model.train() 182 | beta = self.beta_sch.getBeta(epoch, prints=False) 183 | loss, MSE, KLD, elapsed, collapsed = train_epoch(model=self.model, 184 | data=self.train_dl, 185 | optimizer=self.opt, 186 | beta=beta, 187 | device=self.device) 188 | self.model.eval() 189 | loss_test, MSE_test, KLD_test, elapsed_test = test_epoch(model=self.model, 190 | data=self.val_dl, 191 | beta=beta, 192 | device=self.device) 193 | 194 | self.opt_sch.step() 195 | 196 | printProgress(epoch=epoch, 197 | epochs=self.config.epochs, 198 | loss=loss, 199 | loss_test=loss_test, 200 | MSE=MSE, 201 | KLD=KLD, 202 | elapsed=elapsed, 203 | elapsed_test=elapsed_test, 204 | collapsed=collapsed) 205 | 206 | logger.add_scalar('General loss/Total', loss, epoch) 207 | logger.add_scalar('General loss/MSE', MSE, epoch) 208 | logger.add_scalar('General loss/KLD', KLD, epoch) 209 | logger.add_scalar('General loss/Total_test', loss_test, epoch) 210 | logger.add_scalar('General loss/MSE_test', MSE_test, epoch) 211 | logger.add_scalar('General loss/KLD_test', KLD_test, epoch) 212 | logger.add_scalar('Optimizer/LR', self.opt_sch.get_last_lr()[0], epoch) 213 | 214 | if (loss_test < bestloss and epoch > 100): 215 | bestloss = loss_test 216 | checkpoint = {'state_dict': self.model.state_dict(), 'optimizer_dict': self.opt.state_dict()} 217 | ckp_file = f'{pathsBib.chekp_path}/{self.filename}_bestVal' + self.fmat 218 | save_checkpoint(state=checkpoint, path_name=ckp_file) 219 | print(f'## Checkpoint. Epoch: {epoch}, test loss: {loss_test}, saving checkpoint {ckp_file}') 220 | 221 | checkpoint = {'state_dict': self.model.state_dict(), 'optimizer_dict': self.opt.state_dict()} 222 | ckp_file = f'{pathsBib.chekp_path}/{self.filename}_final' + self.fmat 223 | save_checkpoint(state=checkpoint, path_name=ckp_file) 224 | print(f'Checkpoint. Final epoch, loss: {loss}, test loss: {loss_test}, saving checkpoint {ckp_file}') 225 | 226 | 227 | #------------------------------------------------- 228 | 229 | def load_pretrain_model(self,model_type='pre'): 230 | """ 231 | 232 | Load the pretrained model for beta VAE 233 | 234 | Args: 235 | 236 | model_type : ['pre', 'val','final'] (str) Choose from pre-trained, best valuation and final model 237 | 238 | """ 239 | 240 | model_type_all = ['pre','val','final'] 241 | assert(model_type in model_type_all), print('ERROR: No type of the model matched') 242 | 243 | if model_type == 'pre': model_path = pathsBib.pretrain_path + self.filename + self.fmat 244 | elif model_type == 'val' : model_path = pathsBib.chekp_path + self.filename + '_bestVal' + self.fmat 245 | elif model_type == 'final' : model_path = pathsBib.chekp_path + self.filename + '_final' + self.fmat 246 | 247 | 248 | try: 249 | ckpoint = torch.load(model_path, map_location= self.device) 250 | 251 | except: 252 | print("ERROR: Model NOT found!") 253 | exit() 254 | stat_dict = ckpoint['state_dict'] 255 | self.model.load_state_dict(stat_dict) 256 | print(f'INFO: the state dict has been loaded!') 257 | 258 | #------------------------------------------------- 259 | 260 | def get_test_data(self): 261 | """ 262 | 263 | Generate the DataLoder for test 264 | 265 | """ 266 | from torch.utils.data import DataLoader 267 | 268 | u_scaled, self.mean, self.std = loadData(self.datafile) 269 | 270 | u_scaled = u_scaled[::self.config.downsample] 271 | n_total = u_scaled.shape[0] 272 | self.n_train = n_total - self.config.n_test 273 | 274 | print(f"INFO: Data Summary: N train: {self.n_train:d}," + \ 275 | f"N test: {self.config.n_test:d},"+\ 276 | f"N total {n_total:d}") 277 | 278 | self.train_d, self.test_d = u_scaled[:self.n_train] ,u_scaled[self.n_train:] 279 | 280 | self.train_dl = DataLoader(torch.from_numpy(self.train_d), 281 | batch_size=1, 282 | shuffle=False, 283 | pin_memory=True, 284 | num_workers=2) 285 | self.test_dl = DataLoader(torch.from_numpy(self.test_d), 286 | batch_size=1, 287 | shuffle=False, 288 | pin_memory=True, 289 | num_workers=2) 290 | 291 | print(f"INFO: Dataloader generated, Num Test batch = {len(self.test_dl)}") 292 | 293 | 294 | 295 | #------------------------------------------------- 296 | def post_process(self): 297 | 298 | """ 299 | 300 | Post-processing for Beta-VAE 301 | 302 | """ 303 | 304 | assert (self.test_dl != None), print("ERROR: NOT able to do post-processing without test data!") 305 | 306 | fname = pathsBib.res_path + "modes_" + self.filename 307 | 308 | if_save_spatial = spatial_Mode(fname, 309 | model=self.model,latent_dim=self.config.latent_dim, 310 | train_data=self.train_d,test_data=self.test_d, 311 | dataset_train=self.train_dl,dataset_test=self.test_dl, 312 | mean = self.mean, std = self.std, 313 | device= self.device, 314 | if_order= True, 315 | if_nlmode= True, 316 | if_Ecumt= True, 317 | if_Ek_t= True 318 | ) 319 | if if_save_spatial: 320 | print(f"INFO: Spatial Modes finished!") 321 | else: 322 | print(f'ERROR: Spatial modes has not saved!') 323 | 324 | 325 | #################################################### 326 | ### RUNNER for Temporal-dynamics Prediction 327 | #################################################### 328 | 329 | 330 | class latentRunner(nn.Module): 331 | def __init__(self,name,device): 332 | """ 333 | A runner for latent space temporal-dynmaics prediction 334 | 335 | Args: 336 | 337 | name : (str) The model choosed for temporal-dynamics prediction 338 | 339 | device : (Str) The device going to use 340 | 341 | """ 342 | 343 | super(latentRunner,self).__init__() 344 | print("#"*30) 345 | print(f"INIT temporal predictor: {name}, device: {device}") 346 | self.device = device 347 | self.model,self.filename, self.config = get_predictors(name) 348 | 349 | self.NumPara = sum(p.numel() for p in self.model.parameters() if p.requires_grad) 350 | 351 | self.fmat = '.pt' 352 | print(f"INFO: The model has been generated, num of parameter is {self.NumPara}") 353 | print(f"Case Name:\n {self.filename}") 354 | 355 | 356 | #------------------------------------------------- 357 | 358 | def train(self): 359 | print("#"*30) 360 | print("INFO: Start Training ") 361 | self.get_data() 362 | self.compile() 363 | self.fit() 364 | self.train_dl = None 365 | self.val_dl = None 366 | print(f"INFO: Training finished, cleaned the data loader") 367 | print("#"*30) 368 | 369 | #------------------------------------------------- 370 | def infer(self, model_type = 'pre', 371 | if_window=True, 372 | if_pmap=True): 373 | """ 374 | 375 | Inference and evaluation of the model 376 | 377 | Args: 378 | 379 | model_type: (str) The type of model to load 380 | 381 | if_window : (str) If compute the sliding-widnow error 382 | 383 | if_pmap : (str) If compute the Poincare Map 384 | 385 | """ 386 | 387 | print("#"*30) 388 | print("INFO: Start post-processing") 389 | # self.com 390 | self.load_pretrain_model(model_type=model_type) 391 | 392 | self.post_process(if_window,if_pmap) 393 | print(f"INFO: Inference ended!") 394 | print("#"*30) 395 | 396 | #------------------------------------------------- 397 | 398 | 399 | def get_data(self): 400 | """ 401 | Get the latent space variable data for training and validation 402 | """ 403 | try: 404 | hdf5 = h5py.File(pathsBib.data_path + "latent_data.h5py") 405 | data = np.array(hdf5['vector']) 406 | except: 407 | print(f"Error: DataBase not found, please check path or keys") 408 | 409 | X,Y = make_Sequence(self.config,data=data) 410 | self.train_dl, self.val_dl =make_DataLoader(torch.from_numpy(X),torch.from_numpy(Y), 411 | batch_size=self.config.Batch_size, 412 | drop_last=False, 413 | train_split=self.config.train_split) 414 | print(f"INFO: DataLoader Generated!") 415 | del data, X, Y 416 | 417 | #------------------------------------------------- 418 | 419 | def compile(self): 420 | """ 421 | Compile the model with optimizer, scheduler and loss function 422 | """ 423 | self.loss_fn = torch.nn.MSELoss() 424 | self.opt = torch.optim.Adam(self.model.parameters(),lr = self.config.lr, eps=1e-7) 425 | self.opt_sch = [ 426 | torch.optim.lr_scheduler.ExponentialLR(optimizer=self.opt, gamma= (1 - 0.01)) 427 | ] 428 | 429 | #------------------------------------------------- 430 | 431 | def fit(self): 432 | """ 433 | Training Model, we use the fit() function 434 | """ 435 | 436 | s_t = time.time() 437 | history = fitting( self.device, 438 | self.model, 439 | self.train_dl, 440 | self.loss_fn, 441 | self.config.Epoch, 442 | self.opt, 443 | self.val_dl, 444 | scheduler=self.opt_sch, 445 | if_early_stop=self.config.early_stop, 446 | patience=self.config.patience) 447 | e_t = time.time() 448 | cost_time = e_t - s_t 449 | 450 | print(f"INFO: Training FINISH, Cost Time: {cost_time:.2f}s") 451 | 452 | check_point = { "model":self.model.state_dict(), 453 | "history":history, 454 | "time":cost_time} 455 | 456 | torch.save(check_point,pathsBib.model_path + self.filename + self.fmat) 457 | print(f"INFO: The checkpoints has been saved!") 458 | 459 | #------------------------------------------------- 460 | 461 | 462 | def load_pretrain_model(self,model_type='pre'): 463 | """ 464 | 465 | Load the pretrained model for beta VAE 466 | 467 | Args: 468 | 469 | model_type : ['pre', 'val','final'] (str) Choose from pre-trained, best valuation and final model 470 | 471 | """ 472 | 473 | model_type_all = ['pre','val','final'] 474 | assert(model_type in model_type_all), print('ERROR: No type of the model matched') 475 | 476 | if model_type == 'pre': model_path = pathsBib.pretrain_path + self.filename + self.fmat 477 | elif model_type == 'val' : model_path = pathsBib.chekp_path + self.filename + '_bestVal' + self.fmat 478 | elif model_type == 'final' : model_path = pathsBib.chekp_path + self.filename + '_final' + self.fmat 479 | try: 480 | ckpoint = torch.load(model_path, map_location= self.device) 481 | 482 | except: 483 | print("ERROR: Model NOT found!") 484 | exit() 485 | stat_dict = ckpoint['model'] 486 | 487 | self.model.load_state_dict(stat_dict) 488 | self.history = ckpoint['history'] 489 | 490 | 491 | print(f'INFO: the state dict has been loaded!') 492 | print(self.model.eval) 493 | 494 | 495 | #------------------------------------------------- 496 | 497 | def post_process(self,if_window=True,if_pmap=True): 498 | """ 499 | Post Processing of the temporal-dynamics predcition 500 | Args: 501 | 502 | if_window : (bool) If compute the sliding-window error 503 | 504 | if_pmap : (bool) If compute the Poincare Map 505 | """ 506 | 507 | try: 508 | hdf5 = h5py.File(pathsBib.data_path + "latent_data.h5py") 509 | test_data = np.array(hdf5['vector_test']) 510 | except: 511 | print(f"Error: DataBase not found, please check path or keys or try run the vae first") 512 | 513 | print(f"INFO: Test data loaded, SIZE = {test_data.shape}") 514 | Preds = make_Prediction(test_data = test_data, 515 | model = self.model, 516 | device = self.device, 517 | in_dim = self.config.in_dim, 518 | next_step = self.config.next_step) 519 | 520 | if if_window: 521 | print(f"Begin to compute the sliding window error") 522 | window_error = Sliding_Window_Error(test_data, 523 | self.model, 524 | self.device, 525 | self.config.in_dim) 526 | else: 527 | window_error = np.nan 528 | 529 | 530 | if if_pmap: 531 | planeNo = 0 532 | postive_dir = True 533 | lim_val = 2.5 # Limitation of x and y bound when compute joint pdf 534 | grid_val = 50 535 | InterSec_pred = Intersection(Preds, planeNo=planeNo,postive_dir=postive_dir) 536 | InterSec_test = Intersection(test_data, planeNo=planeNo,postive_dir=postive_dir) 537 | else: 538 | InterSec_pred = np.nan 539 | InterSec_test = np.nan 540 | 541 | 542 | np.savez_compressed( 543 | pathsBib.res_path + self.filename + ".npz", 544 | p = Preds, 545 | g = test_data, 546 | e = window_error, 547 | pmap_g = InterSec_test, 548 | pmap_p = InterSec_pred 549 | ) -------------------------------------------------------------------------------- /lib/train.py: -------------------------------------------------------------------------------- 1 | """ 2 | functions for training 3 | 4 | @yuningw 5 | """ 6 | 7 | import torch 8 | import torch.nn as nn 9 | import time 10 | import numpy as np 11 | 12 | 13 | ############################################## 14 | # Functions used for training beta-VAE 15 | ################################################ 16 | 17 | class betaScheduler: 18 | """Schedule beta, linear growth to max value""" 19 | 20 | def __init__(self, startvalue, endvalue, warmup): 21 | self.startvalue = startvalue 22 | self.endvalue = endvalue 23 | self.warmup = warmup 24 | 25 | def getBeta(self, epoch, prints=False): 26 | 27 | if epoch < self.warmup: 28 | beta = self.startvalue + (self.endvalue - self.startvalue) * epoch / self.warmup 29 | if prints: 30 | print(beta) 31 | return beta 32 | else: 33 | return self.endvalue 34 | 35 | 36 | def loss_function(reconstruction, data, mean, logvariance, beta): 37 | MSELoss = nn.MSELoss(reduction='mean').cuda() 38 | MSE = MSELoss(reconstruction, data) 39 | 40 | KLD = -0.5 * torch.mean(1 + logvariance - mean.pow(2) - logvariance.exp()) 41 | 42 | loss = MSE + KLD * beta 43 | 44 | return loss, MSE, KLD 45 | 46 | def train_epoch(model, data, optimizer, beta, device): 47 | start_epoch_time = time.time() 48 | 49 | loss_batch = [] # Store batch loss 50 | MSE_batch = [] # Store batch MSE 51 | KLD_batch = [] # Store batch KLD 52 | logVar_batch = [] # Store batch logVar to count collapsed modes 53 | 54 | for batch in data: 55 | #batch_noise = gaussian_noise(batch, 0.2, device) 56 | if not batch.is_cuda: 57 | batch = batch.to(device, non_blocking=True) 58 | #batch_noise = batch_noise.to(device, non_blocking=True) 59 | 60 | rec, mean, logvariance = model(batch)# + batch_noise) 61 | loss, MSE, KLD = loss_function(rec, batch, mean, logvariance, beta) 62 | 63 | loss_batch.append(loss.item()) 64 | MSE_batch.append(MSE.item()) 65 | KLD_batch.append(KLD.item()) 66 | logVar_batch.append(np.exp(0.5* np.mean(logvariance.detach().cpu().numpy(), 0))) 67 | 68 | optimizer.zero_grad() 69 | loss.backward() 70 | optimizer.step() 71 | 72 | 73 | return sum(loss_batch) / len(loss_batch),\ 74 | sum(MSE_batch) / len(MSE_batch),\ 75 | sum(KLD_batch) / len(KLD_batch),\ 76 | time.time() - start_epoch_time, \ 77 | (np.mean(np.stack(logVar_batch, axis=0), 0) < 0.1).sum() # count collapsed modes 78 | 79 | def test_epoch(model, data, beta, device): 80 | start_epoch_time = time.time() 81 | 82 | with torch.no_grad(): 83 | loss_batch = [] # Store batch loss 84 | MSE_batch = [] # Store batch MSE 85 | KLD_batch = [] # Store batch KLD 86 | 87 | for batch in data: 88 | if not batch.is_cuda: 89 | batch = batch.to(device, non_blocking=True) 90 | 91 | rec, mean, logvariance = model(batch) 92 | loss, MSE, KLD = loss_function(rec, batch, mean, logvariance, beta) 93 | 94 | loss_batch.append(loss.item()) 95 | MSE_batch.append(MSE.item()) 96 | KLD_batch.append(KLD.item()) 97 | 98 | return sum(loss_batch) / len(loss_batch),\ 99 | sum(MSE_batch) / len(MSE_batch),\ 100 | sum(KLD_batch) / len(KLD_batch),\ 101 | time.time() - start_epoch_time 102 | 103 | def gaussian_noise(x, var, device): 104 | 105 | """ 106 | 107 | Add gaussian noise on std term for reprarameterization 108 | 109 | 110 | 111 | """ 112 | return torch.normal(0.0, var, size=x.shape) 113 | 114 | 115 | 116 | def printProgress(epoch, epochs, loss, loss_test, MSE, KLD, elapsed, elapsed_test, collapsed): 117 | print(f"Epoch: {epoch:3d}/{epochs:d}, Loss: {loss:2.4f}, Loss_test: {loss_test:2.4f}, MSE: {MSE:2.4f}, KLD: {KLD:2.4f}, collapsed: {collapsed:2d}, time train: {elapsed:2.3f}, time test: {elapsed_test:2.3f}") 118 | 119 | 120 | 121 | ############################################## 122 | # Functions used for training temporal-dynamics predictor 123 | ################################################ 124 | 125 | 126 | class EarlyStopper: 127 | def __init__(self, patience=1, min_delta=0): 128 | self.patience = patience 129 | self.min_delta = min_delta 130 | self.counter = 0 131 | self.min_validation_loss = np.inf 132 | 133 | def early_stop(self, validation_loss): 134 | if validation_loss < self.min_validation_loss: 135 | self.min_validation_loss = validation_loss 136 | self.counter = 0 137 | elif validation_loss > (self.min_validation_loss + self.min_delta): 138 | self.counter += 1 139 | if self.counter >= self.patience: 140 | return True 141 | return False 142 | 143 | 144 | 145 | 146 | 147 | def fitting(device, 148 | model, 149 | dl, 150 | loss_fn, 151 | Epoch, 152 | optimizer:torch.optim.Optimizer, 153 | val_dl = None, 154 | scheduler:list= None, 155 | if_early_stop = True,patience = 10, 156 | ): 157 | 158 | """ 159 | A function for training loop 160 | 161 | Args: 162 | device : the device for training, which should match the model 163 | 164 | model : The model to be trained 165 | 166 | dl : A dataloader for training 167 | 168 | loss_fn : Loss function 169 | 170 | Epochs : Number of epochs 171 | 172 | optimizer : The optimizer object 173 | 174 | val_dl : The data for validation 175 | 176 | scheduler : A list of traning scheduler 177 | 178 | 179 | Returns: 180 | history: A dict contains training loss and validation loss (if have) 181 | 182 | """ 183 | 184 | from tqdm import tqdm 185 | 186 | history = {} 187 | history["train_loss"] = [] 188 | 189 | if val_dl: 190 | history["val_loss"] = [] 191 | 192 | model.to(device) 193 | print(f"INFO: The model is assigned to device: {device} ") 194 | 195 | if scheduler is not None: 196 | print(f"INFO: The following schedulers are going to be used:") 197 | for sch in scheduler: 198 | print(f"{sch.__class__}") 199 | 200 | print(f"INFO: Training start") 201 | 202 | if if_early_stop: 203 | early_stopper = EarlyStopper(patience=patience,min_delta=0) 204 | print("INFO: Early-Stopper prepared") 205 | 206 | for epoch in range(Epoch): 207 | ##### 208 | #Training step 209 | ##### 210 | model.train() 211 | loss_val = 0; num_batch = 0 212 | for batch in tqdm(dl): 213 | x, y = batch 214 | x = x.to(device).float(); y =y.to(device).float() 215 | optimizer.zero_grad() 216 | 217 | pred = model(x) 218 | loss = loss_fn(pred,y) 219 | loss.backward() 220 | optimizer.step() 221 | 222 | 223 | 224 | loss_val += loss.item()/x.shape[0] 225 | num_batch += 1 226 | 227 | history["train_loss"].append(loss_val/num_batch) 228 | 229 | if scheduler is not None: 230 | lr_now = 0 231 | for sch in scheduler: 232 | sch.step() 233 | lr_now = sch.get_last_lr() 234 | print(f"INFO: Scheduler updated, LR = {lr_now} ") 235 | 236 | if val_dl: 237 | ##### 238 | #Valdation step 239 | ##### 240 | loss_val = 0 ; num_batch = 0 241 | model.eval() 242 | for batch in (val_dl): 243 | x, y = batch 244 | x = x.to(device).float(); y =y.to(device).float() 245 | pred = model(x) 246 | loss = loss_fn(pred,y) 247 | 248 | loss_val += loss.item()/x.shape[0] 249 | num_batch += 1 250 | 251 | history["val_loss"].append(loss_val/num_batch) 252 | 253 | train_loss = history["train_loss"][-1] 254 | val_loss = history["val_loss"][-1] 255 | print( 256 | f"At Epoch = {epoch},\n"+\ 257 | f"Train_loss = {train_loss}\n"+\ 258 | f"Val_loss = {val_loss}\n" 259 | ) 260 | 261 | if if_early_stop: 262 | if early_stopper.early_stop(loss_val/num_batch): 263 | print("Early-stopp Triggered, Going to stop the training") 264 | break 265 | return history 266 | 267 | 268 | 269 | 270 | 271 | -------------------------------------------------------------------------------- /license.md: -------------------------------------------------------------------------------- 1 | # Creative Commons Attribution 4.0 International 2 | 3 | Creative Commons Corporation (“Creative Commons”) is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an “as-is” basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible. 4 | 5 | ### Using Creative Commons Public Licenses 6 | 7 | Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses. 8 | 9 | * __Considerations for licensors:__ Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. [More considerations for licensors](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensors). 10 | 11 | * __Considerations for the public:__ By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor’s permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. [More considerations for the public](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensees). 12 | 13 | ## Creative Commons Attribution 4.0 International Public License 14 | 15 | By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions. 16 | 17 | ### Section 1 – Definitions. 18 | 19 | a. __Adapted Material__ means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image. 20 | 21 | b. __Adapter's License__ means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License. 22 | 23 | c. __Copyright and Similar Rights__ means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights. 24 | 25 | d. __Effective Technological Measures__ means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements. 26 | 27 | e. __Exceptions and Limitations__ means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material. 28 | 29 | f. __Licensed Material__ means the artistic or literary work, database, or other material to which the Licensor applied this Public License. 30 | 31 | g. __Licensed Rights__ means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license. 32 | 33 | h. __Licensor__ means the individual(s) or entity(ies) granting rights under this Public License. 34 | 35 | i. __Share__ means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them. 36 | 37 | j. __Sui Generis Database Rights__ means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world. 38 | 39 | k. __You__ means the individual or entity exercising the Licensed Rights under this Public License. __Your__ has a corresponding meaning. 40 | 41 | ### Section 2 – Scope. 42 | 43 | a. ___License grant.___ 44 | 45 | 1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to: 46 | 47 | A. reproduce and Share the Licensed Material, in whole or in part; and 48 | 49 | B. produce, reproduce, and Share Adapted Material. 50 | 51 | 2. __Exceptions and Limitations.__ For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions. 52 | 53 | 3. __Term.__ The term of this Public License is specified in Section 6(a). 54 | 55 | 4. __Media and formats; technical modifications allowed.__ The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material. 56 | 57 | 5. __Downstream recipients.__ 58 | 59 | A. __Offer from the Licensor – Licensed Material.__ Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License. 60 | 61 | B. __No downstream restrictions.__ You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material. 62 | 63 | 6. __No endorsement.__ Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i). 64 | 65 | b. ___Other rights.___ 66 | 67 | 1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise. 68 | 69 | 2. Patent and trademark rights are not licensed under this Public License. 70 | 71 | 3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties. 72 | 73 | ### Section 3 – License Conditions. 74 | 75 | Your exercise of the Licensed Rights is expressly made subject to the following conditions. 76 | 77 | a. ___Attribution.___ 78 | 79 | 1. If You Share the Licensed Material (including in modified form), You must: 80 | 81 | A. retain the following if it is supplied by the Licensor with the Licensed Material: 82 | 83 | i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated); 84 | 85 | ii. a copyright notice; 86 | 87 | iii. a notice that refers to this Public License; 88 | 89 | iv. a notice that refers to the disclaimer of warranties; 90 | 91 | v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable; 92 | 93 | B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and 94 | 95 | C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License. 96 | 97 | 2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information. 98 | 99 | 3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable. 100 | 101 | 4. If You Share Adapted Material You produce, the Adapter's License You apply must not prevent recipients of the Adapted Material from complying with this Public License. 102 | 103 | ### Section 4 – Sui Generis Database Rights. 104 | 105 | Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material: 106 | 107 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database; 108 | 109 | b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material; and 110 | 111 | c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database. 112 | 113 | For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights. 114 | 115 | ### Section 5 – Disclaimer of Warranties and Limitation of Liability. 116 | 117 | a. __Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.__ 118 | 119 | b. __To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.__ 120 | 121 | c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability. 122 | 123 | ### Section 6 – Term and Termination. 124 | 125 | a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically. 126 | 127 | b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates: 128 | 129 | 1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or 130 | 131 | 2. upon express reinstatement by the Licensor. 132 | 133 | For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License. 134 | 135 | c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License. 136 | 137 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License. 138 | 139 | ### Section 7 – Other Terms and Conditions. 140 | 141 | a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed. 142 | 143 | b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License. 144 | 145 | ### Section 8 – Interpretation. 146 | 147 | a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License. 148 | 149 | b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions. 150 | 151 | c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor. 152 | 153 | d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority. 154 | 155 | > Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the “Licensor.” Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at [creativecommons.org/policies](http://creativecommons.org/policies), Creative Commons does not authorize the use of the trademark “Creative Commons” or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses. 156 | > 157 | > Creative Commons may be contacted at creativecommons.org -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | """ 2 | Main program 3 | 4 | NOTE: The "run" in running mode here means we do both train and infer 5 | 6 | @yuningw 7 | """ 8 | 9 | import torch 10 | import argparse 11 | from lib import init, POD 12 | from lib.runners import vaeRunner, latentRunner 13 | from utils.figs_time import vis_temporal_Prediction 14 | from utils.figs import vis_bvae, vis_pod 15 | 16 | 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('-nn',default="easy", type=str, help="Choose the model for time-series prediction: easy, self OR lstm") 19 | parser.add_argument('-re',default=40, type=int, help="40 OR 100, Choose corresponding Reynolds number for the case") 20 | parser.add_argument('-m', default="test", type=str, help='Switch the mode between train, infer and run') 21 | parser.add_argument('-t', default="pre", type=str, help='The type of saved model: pre/val/final') 22 | parser.add_argument('-pod',default=True, type=bool, help='Compute POD') 23 | args = parser.parse_args() 24 | 25 | device = ('cuda' if torch.cuda.is_available() else "cpu") 26 | 27 | 28 | if __name__ == "__main__": 29 | 30 | ## Env INIT 31 | datafile = init.init_env(args.re) 32 | 33 | 34 | ## Beta-VAE 35 | bvae = vaeRunner(device,datafile) 36 | if args.m == 'train': 37 | bvae.train() 38 | elif args.m == 'test': 39 | bvae.infer(args.t) 40 | elif args.m == 'run': 41 | bvae.run() 42 | 43 | 44 | ## POD 45 | if args.pod: 46 | POD = POD.POD(datafile, n_test=bvae.config.n_test, re=args.re, 47 | path='res/', n_modes=10, delta_t=bvae.config.delta_t) 48 | POD.load_data() 49 | POD.get_POD() 50 | POD.eval_POD() 51 | 52 | # Time-series prediction runner 53 | lruner = latentRunner(args.nn,device) 54 | if args.m == 'train': 55 | lruner.train() 56 | elif args.m == 'test': 57 | lruner.infer(args.t) 58 | elif args.m == 'run': 59 | lruner.train() 60 | lruner.infer(args.t) 61 | 62 | vis_bvae(init.pathsBib.res_path + "modes_" + bvae.filename + ".hdf5", 63 | init.pathsBib.log_path + bvae.filename) 64 | vis_pod(POD) 65 | vis_temporal_Prediction(model_type=args.nn, predictor=lruner, vae=bvae) 66 | -------------------------------------------------------------------------------- /models/VAE_easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/models/VAE_easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt -------------------------------------------------------------------------------- /models/checkpoints/Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000_bestVal.pth.tar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/models/checkpoints/Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000_bestVal.pth.tar -------------------------------------------------------------------------------- /models/checkpoints/Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000_final.pth.tar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/models/checkpoints/Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000_final.pth.tar -------------------------------------------------------------------------------- /models/pretrained/LSTM_64in_64dmodel_1next_2dim_Noneemb_128hideen_4nlayer_Noneoutact_100Epoch_135000N_TrueES_50P.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/models/pretrained/LSTM_64in_64dmodel_1next_2dim_Noneemb_128hideen_4nlayer_Noneoutact_100Epoch_135000N_TrueES_50P.pt -------------------------------------------------------------------------------- /models/pretrained/Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000.pth.tar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/models/pretrained/Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000.pth.tar -------------------------------------------------------------------------------- /models/pretrained/easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/models/pretrained/easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt -------------------------------------------------------------------------------- /models/pretrained/selfAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/models/pretrained/selfAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.pt -------------------------------------------------------------------------------- /nns/RNNs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Recurrent network model in torch.nn.Module 3 | 4 | @author Yuning Wang 5 | """ 6 | import torch 7 | from torch.nn import Module, ModuleList, Linear, Dropout, LayerNorm 8 | from nns.embedding import SineActivation, CosineActivation, PosEncoding 9 | import torch.nn as nn 10 | 11 | class LSTMs(nn.Module): 12 | def __init__(self, d_input, 13 | d_model, 14 | nmode, 15 | embed = None, 16 | hidden_size = 64, 17 | num_layer = 1, 18 | is_output = True, 19 | out_dim = 1, 20 | out_act = "tanh", 21 | is_output_bias = True) -> None: 22 | """ 23 | Module for Long Short-term Memory architecture 24 | 25 | Note that by default, we do not use embedding layer for RNN model since it is able to recognize the sequence data 26 | 27 | Args: 28 | d_input : (Int) The input dimension 29 | 30 | d_model : (Int) The projection dimension 31 | 32 | nmode : (Int) Number of mode 33 | 34 | embed : (Str) Embed layer, if not use embedding then args is None 35 | 36 | hidden_size : (Str) Number of hidden cells in RNN 37 | 38 | num_layer : (Str) Number of layers used in RNN 39 | 40 | is_output : (Bool) If use output layer 41 | 42 | out_act : (Str) The activation function used for output 43 | 44 | is_out_put_bias : (Bool) If use bias for output layer 45 | 46 | 47 | 48 | 49 | """ 50 | 51 | super(LSTMs,self).__init__() 52 | 53 | self.hidden_size = hidden_size 54 | self.num_layer = num_layer 55 | 56 | self.is_output = is_output 57 | self.out_dim = out_dim 58 | if embed == "sin": 59 | self.embed = SineActivation(nmode,nmode,d_model) 60 | if embed == "cos": 61 | self.embed = CosineActivation(nmode,nmode,d_model) 62 | if embed == "posenc": 63 | self.embed = PosEncoding(d_input,nmode,d_model) 64 | 65 | if embed == None: 66 | try: d_model == d_input 67 | except: print("INFO: NO Embedding! \nThe input dimension should be same as the d_model") 68 | self.embed = None 69 | 70 | 71 | self.lstms = nn.LSTM(nmode,hidden_size,num_layer,batch_first=True ) 72 | for name, param in self.lstms.named_parameters(): 73 | if "bias" in name: 74 | nn.init.zeros_(param) 75 | print(f"INFO: {name} has been initilised") 76 | if 'weight' in name: 77 | nn.init.xavier_uniform_(param) 78 | print(f"INFO: {name} has been initilised") 79 | 80 | 81 | if is_output: 82 | self.out = nn.Linear(hidden_size,nmode, 83 | bias=is_output_bias) 84 | nn.init.xavier_uniform_(self.out.weight) 85 | nn.init.zeros_(self.out.bias) 86 | if out_act =="tanh": 87 | self.out_act = nn.Tanh() 88 | elif out_act =="elu": 89 | self.out_act = nn.ELU() 90 | else: 91 | self.out_act = None 92 | 93 | def forward(self,x): 94 | 95 | if self.embed != None: 96 | x = self.embed(x) 97 | 98 | hidden, cell = self.init_hidden(x.shape[0],device=x.device) 99 | 100 | x, (hn,cn) = self.lstms(x,(hidden.detach(),cell.detach())) 101 | 102 | 103 | if self.is_output: 104 | if self.out_act is not None: 105 | return self.out_act(self.out(x[:,-self.out_dim:,:])) 106 | else: 107 | return self.out(x[:,-self.out_dim:,:]) 108 | def init_hidden(self,batch_size,device): 109 | hidden = torch.zeros(self.num_layer, 110 | batch_size, 111 | self.hidden_size).to(device) 112 | 113 | cell = torch.zeros(self.num_layer, 114 | batch_size, 115 | self.hidden_size).to(device) 116 | 117 | 118 | return hidden, cell -------------------------------------------------------------------------------- /nns/__pycache__/AutoEncoder.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/AutoEncoder.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/EmbedTransformerEncoder.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/EmbedTransformerEncoder.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/Embedding.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/Embedding.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/FullTransformer.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/FullTransformer.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/RNNs.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/RNNs.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/RevTransformer.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/RevTransformer.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/Transformer.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/Transformer.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/attns.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/attns.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/beta_vae.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/beta_vae.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/easyAttns.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/easyAttns.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/embedding.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/embedding.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/layers.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/layers.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/transformer.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/transformer.cpython-38.pyc -------------------------------------------------------------------------------- /nns/__pycache__/vae.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/nns/__pycache__/vae.cpython-38.pyc -------------------------------------------------------------------------------- /nns/attns.py: -------------------------------------------------------------------------------- 1 | """ 2 | Attention module used in Transformer 3 | 4 | Here we use the easy-attention mechanism for time-series prediction in latent space 5 | We use a stack of encoders only for the prediction, leveraging time-delay embedding 6 | """ 7 | import torch 8 | import torch.nn as nn 9 | import math 10 | 11 | class MultiHeadAttention(nn.Module): 12 | def __init__(self, d_model, num_heads): 13 | """ 14 | Multi-head attention with mask 15 | 16 | Args: 17 | 18 | d_model : The dimension of the embedding 19 | 20 | num_heads : Number of heads to be split 21 | 22 | Funcs: 23 | __init__ 24 | 25 | scaled_dot_product_attention(self, Q, K, V, mask): Implement the scaled-dot attention for each head 26 | 27 | split_heads() : Split the sequnce by the number of heads 28 | 29 | combine_heads() : Combine the sequnce by the number of heads 30 | 31 | forward() : Forward prop for the module 32 | 33 | 34 | """ 35 | super(MultiHeadAttention, self).__init__() 36 | assert d_model % num_heads == 0, "d_model must be divisible by num_heads" 37 | 38 | # Basic data dimension 39 | self.d_model = d_model 40 | self.num_heads = num_heads 41 | self.d_k = d_model // num_heads 42 | 43 | self.W_q = nn.Linear(d_model, d_model) 44 | self.W_k = nn.Linear(d_model, d_model) 45 | self.W_v = nn.Linear(d_model, d_model) 46 | self.W_o = nn.Linear(d_model, d_model) 47 | 48 | def scaled_dot_product_attention(self, Q, K, V, mask=None): 49 | """ 50 | Implement the scaled dot attention 51 | 52 | Args: 53 | Q : Query projected by WQ 54 | K : Key projected by WK 55 | V : Value projected by WV 56 | mask : Mask that implmented on attention score 57 | 58 | Returns: 59 | output : The sequence that be encoded by attention mechanism 60 | """ 61 | attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) 62 | 63 | if mask is not None: 64 | 65 | attn_scores = attn_scores.masked_fill(mask == 0, -1e9) 66 | 67 | attn_probs = torch.softmax(attn_scores, dim=-1) 68 | output = torch.matmul(attn_probs, V) 69 | return output 70 | 71 | def split_heads(self, x): 72 | """ 73 | Split the sequence into multi-heads 74 | 75 | Args: 76 | x : Input sequence shape = [B, S, N] 77 | 78 | Returns: 79 | x : sequence with shape = [B, H, S, N//H] 80 | """ 81 | batch_size, seq_length, d_model = x.size() 82 | return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2) 83 | 84 | def combine_heads(self, x): 85 | """ 86 | Combine the sequence into multi-heads 87 | 88 | Args: 89 | x : Input sequence shape = [B, H, S, N//H] 90 | 91 | Returns: 92 | x : sequence with shape = [B, S, N] 93 | """ 94 | 95 | 96 | batch_size, _, seq_length, d_k = x.size() 97 | return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model) 98 | 99 | def forward(self, Q, K, V, mask=None): 100 | """ 101 | Forward prop of the attention module 102 | 103 | Args: 104 | Q, K, V : The input tensor used as Query, Key and Value 105 | mask : The mask used on attention score 106 | 107 | Returns: 108 | output : The encoded sequence by attention module 109 | 110 | """ 111 | Q = self.split_heads(self.W_q(Q)) 112 | K = self.split_heads(self.W_k(K)) 113 | V = self.split_heads(self.W_v(V)) 114 | 115 | attn_output = self.scaled_dot_product_attention(Q, K, V, mask) 116 | output = self.W_o(self.combine_heads(attn_output)) 117 | return output 118 | 119 | 120 | 121 | 122 | class DenseEasyAttn(nn.Module): 123 | 124 | def __init__(self, d_model, 125 | seqLen, 126 | num_head,) -> None: 127 | """ 128 | Dense Easy attention mechansim used in transformer model for the time-series prediction and reconstruction 129 | 130 | Args: 131 | 132 | d_model : The embedding dimension for the input tensor 133 | 134 | seqLen : The length of the sequence 135 | 136 | num_head : The number of head to be used for multi-head attention 137 | 138 | """ 139 | super(DenseEasyAttn,self).__init__() 140 | 141 | assert d_model % num_head == 0, "dmodel must be divible by number of heads" 142 | 143 | self.d_model = d_model 144 | self.d_k = d_model // num_head 145 | self.num_heads = num_head 146 | # Create the tensors 147 | self.Alpha = nn.Parameter(torch.randn(size=(num_head,seqLen,seqLen),dtype=torch.float),requires_grad=True) 148 | self.WV = nn.Parameter(torch.randn(size=(d_model,d_model) ,dtype=torch.float),requires_grad=True) 149 | # Initialisation 150 | nn.init.xavier_uniform_(self.Alpha) 151 | nn.init.xavier_uniform_(self.WV) 152 | 153 | def split_heads(self, x): 154 | """ 155 | Split the sequence into multi-heads 156 | 157 | Args: 158 | x : Input sequence shape = [B, S, N] 159 | 160 | Returns: 161 | x : sequence with shape = [B, H, S, N//H] 162 | """ 163 | batch_size, seq_length, d_model = x.size() 164 | return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2) 165 | 166 | def combine_heads(self, x): 167 | """ 168 | Combine the sequence into multi-heads 169 | 170 | Args: 171 | x : Input sequence shape = [B, H, S, N//H] 172 | 173 | Returns: 174 | x : sequence with shape = [B, S, N] 175 | """ 176 | batch_size, _, seq_length, d_k = x.size() 177 | return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model) 178 | 179 | 180 | def forward(self,x:torch.Tensor): 181 | """ 182 | Forward prop for the easyattention module 183 | 184 | Following the expression: x_hat = Alpha @ Wv @ x 185 | 186 | Args: 187 | 188 | self : The self objects 189 | 190 | x : A tensor of Input data 191 | 192 | Returns: 193 | 194 | x : The tensor be encoded by the moudle 195 | 196 | """ 197 | # Obtain the value of batch size 198 | B,_,_ = x.shape 199 | # We repeat the tensor into same number of batch size that input has 200 | Wv = self.WV.repeat(B,1,1) 201 | # Implement matmul along the batch 202 | V = torch.bmm(Wv,x) 203 | # Split heads for value 204 | V_h = self.split_heads(V) 205 | # Implement the learnable attention tensor 206 | Alpha = self.Alpha.repeat(B,1,1,1) 207 | # Gain Attention by matrix multiplication 208 | x = Alpha @ V_h 209 | # Combine the output back to the original size 210 | x = self.combine_heads(x) 211 | 212 | return x 213 | 214 | 215 | 216 | 217 | 218 | -------------------------------------------------------------------------------- /nns/beta_vae.py: -------------------------------------------------------------------------------- 1 | """ 2 | The archiecture for beta-VAE 3 | @alsolra 4 | """ 5 | 6 | import torch 7 | from torch import nn 8 | 9 | 10 | class VAE(nn.Module): 11 | def __init__(self, latent_dim): 12 | super().__init__() 13 | 14 | self.latent_dim = latent_dim 15 | self.encoder = self.buildEncoder(latent_dim) 16 | self.decoder = self.buildDecoder(latent_dim) 17 | 18 | def buildEncoder(self, latent_dim): 19 | encoder = nn.Sequential( 20 | 21 | nn.Conv2d(in_channels=2, out_channels=8, kernel_size=3, stride=2, padding=1), 22 | nn.ELU(), 23 | 24 | nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=2, padding=1), 25 | nn.ELU(), 26 | 27 | nn.ConstantPad3d((0, 1, 0, 0), 0), 28 | nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1), 29 | nn.ELU(), 30 | 31 | nn.ConstantPad3d((0, 0, 0, 1), 0), 32 | nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1), 33 | nn.ELU(), 34 | 35 | nn.ConstantPad3d((0, 1, 0, 0), 0), 36 | nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1), 37 | nn.ELU(), 38 | 39 | nn.ConstantPad3d((0, 0, 0, 1), 0), 40 | nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1), 41 | nn.ELU(), 42 | 43 | nn.Flatten(start_dim=1, end_dim=-1), 44 | 45 | nn.Linear(2560, 256), 46 | nn.ELU(), 47 | 48 | nn.Linear(256, latent_dim * 2), 49 | ) 50 | return encoder 51 | 52 | 53 | def buildDecoder(self, latent_dim): 54 | decoder = nn.Sequential( 55 | 56 | nn.Linear(latent_dim, 256), 57 | nn.ELU(), 58 | 59 | nn.Linear(256, 256 * 5 * 2), 60 | nn.ELU(), 61 | 62 | nn.Unflatten(dim=1, unflattened_size=(256, 2, 5)), 63 | 64 | nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=3, stride=2, padding=1, output_padding=1), 65 | nn.ConstantPad2d((0, 0, 0, -1), 0), 66 | nn.ELU(), 67 | 68 | nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=3, stride=2, padding=1, output_padding=1), 69 | nn.ConstantPad2d((0, -1, 0, 0), 0), 70 | nn.ELU(), 71 | 72 | nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=3, stride=2, padding=1, output_padding=1), 73 | nn.ConstantPad2d((0, 0, 0, -1), 0), 74 | nn.ELU(), 75 | 76 | nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=3, stride=2, padding=1, output_padding=1), 77 | nn.ConstantPad2d((0, -1, 0, 0), 0), 78 | nn.ELU(), 79 | 80 | nn.ConvTranspose2d(in_channels=16, out_channels=8, kernel_size=3, stride=2, padding=1, output_padding=1), 81 | nn.ConstantPad2d((0, 0, 0, 0), 0), 82 | nn.ELU(), 83 | 84 | nn.ConvTranspose2d(in_channels=8, out_channels=2, kernel_size=3, stride=2, padding=1, output_padding=1), 85 | 86 | ) 87 | return decoder 88 | 89 | def sample(self, mean, logvariance): 90 | """ 91 | Implementing reparameterlisation trick 92 | """ 93 | 94 | std = torch.exp(0.5 * logvariance) 95 | epsilon = torch.rand_like(std) 96 | 97 | return mean + epsilon*std 98 | 99 | def forward(self, data): 100 | 101 | mean_logvariance = self.encoder(data) 102 | 103 | mean, logvariance = torch.chunk(mean_logvariance, 2, dim=1) 104 | 105 | z = self.sample(mean, logvariance) 106 | 107 | reconstruction = self.decoder(z) 108 | 109 | return reconstruction, mean, logvariance 110 | -------------------------------------------------------------------------------- /nns/embedding.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import Module 4 | 5 | ###################################################### 6 | #### Time2Vec architecture 7 | ###################################################### 8 | 9 | def t2v(tau, f, w, b, w0, b0, arg=None): 10 | if arg: 11 | v1 = f(torch.matmul(tau, w) + b, arg) 12 | else: 13 | v1 = f(torch.matmul(tau, w) + b) 14 | 15 | 16 | v2 = torch.matmul(tau, w0) + b0 17 | 18 | return torch.cat([v1, v2], -1) 19 | 20 | class SineActivation(Module): 21 | def __init__(self, in_features,nmodes, out_features): 22 | super(SineActivation, self).__init__() 23 | 24 | self.w0 = nn.parameter.Parameter(torch.randn(in_features,nmodes)) 25 | self.b0 = nn.parameter.Parameter(torch.randn(nmodes)) 26 | 27 | self.w = nn.parameter.Parameter(torch.randn(in_features,out_features-nmodes)) 28 | self.b = nn.parameter.Parameter(torch.randn(out_features-nmodes)) 29 | 30 | self.f = torch.sin 31 | 32 | nn.init.xavier_uniform_(self.w0) 33 | nn.init.xavier_uniform_(self.w) 34 | 35 | 36 | def forward(self, tau): 37 | return t2v(tau, self.f, self.w, self.b, self.w0, self.b0) 38 | 39 | class CosineActivation(Module): 40 | def __init__(self, in_features,nmodes, out_features): 41 | super(CosineActivation, self).__init__() 42 | self.out_features = out_features 43 | self.w0 = nn.parameter.Parameter(torch.randn(in_features,nmodes)) 44 | self.b0 = nn.parameter.Parameter(torch.randn(nmodes)) 45 | 46 | self.w = nn.parameter.Parameter(torch.randn(in_features, out_features-nmodes)) 47 | self.b = nn.parameter.Parameter(torch.randn(out_features-nmodes)) 48 | 49 | self.f = torch.cos 50 | 51 | nn.init.xavier_uniform_(self.w0) 52 | nn.init.xavier_uniform_(self.w) 53 | 54 | def forward(self, tau): 55 | return t2v(tau, self.f, self.w, self.b, self.w0, self.b0) 56 | 57 | 58 | 59 | 60 | 61 | import math 62 | 63 | 64 | class PosEncoding(Module): 65 | def __init__(self,in_features,nmodes,out_features ) -> None: 66 | super(PosEncoding,self).__init__() 67 | 68 | self.proj = nn.Linear(nmodes,out_features) 69 | 70 | pe = torch.zeros((in_features, out_features)) 71 | 72 | position = torch.arange(0, in_features).unsqueeze(1) 73 | div_term = torch.exp((torch.arange(0, out_features, 2, dtype=torch.float) * 74 | -(math.log(10000.0) / out_features))) 75 | 76 | pe[:, 0::2] = torch.sin(position.float() * div_term) 77 | pe[:, 1::2] = torch.cos(position.float() * div_term) 78 | self.pe = pe 79 | 80 | nn.init.xavier_normal_(self.proj.weight) 81 | nn.init.zeros_(self.proj.bias) 82 | 83 | def forward(self,x): 84 | x_emb = self.proj(x) 85 | 86 | return x_emb + self.pe[:, :x_emb.size(-1)].to(x.device) 87 | 88 | 89 | def Positional_Encoding(Batch,d_model, nmode): 90 | """ 91 | :param d_model: dimension of the model 92 | :param nmode: nmode of positions 93 | :return: nmode*d_model position matrix 94 | """ 95 | if d_model % 2 != 0: 96 | raise ValueError("Cannot use sin/cos positional encoding with " 97 | "odd dim (got dim={:d})".format(d_model)) 98 | pe = torch.zeros(Batch,nmode, d_model) 99 | position = torch.arange(0, nmode).unsqueeze(1) 100 | div_term = torch.exp((torch.arange(0, d_model, 2, dtype=torch.float) * 101 | -(math.log(10000.0) / d_model))) 102 | pe[:,:, 0::2] = torch.sin(position.float() * div_term) 103 | pe[:,:, 1::2] = torch.cos(position.float() * div_term) 104 | 105 | return pe 106 | 107 | 108 | 109 | 110 | """ 111 | 112 | Create a new embedding strategy for time and space embedding 113 | 114 | @ yuningw 115 | 116 | """ 117 | 118 | import torch 119 | 120 | from torch import nn 121 | 122 | import numpy as np 123 | 124 | 125 | class TimeSpaceEmbedding(nn.Module): 126 | 127 | """" 128 | 129 | A embedding module based on both time and space 130 | Args: 131 | 132 | d_input : The input size of timedelay embedding 133 | 134 | n_mode : The number of modes/dynamics in the time series 135 | 136 | d_expand : The projection along the time 137 | 138 | d_model : The projection along the space 139 | 140 | """ 141 | 142 | def __init__(self, d_input, n_mode, 143 | d_expand,d_model ): 144 | 145 | super(TimeSpaceEmbedding, self).__init__() 146 | 147 | self.spac_proj = nn.Linear(n_mode,d_model) 148 | 149 | self.time_proj = nn.Conv1d(d_input, d_expand,1) 150 | 151 | self.time_avgpool = nn.AvgPool1d(2) 152 | self.time_maxpool = nn.MaxPool1d(2) 153 | self.time_compress = nn.Linear(d_model, d_model) 154 | self.act = nn.Identity() 155 | 156 | 157 | nn.init.xavier_uniform_(self.spac_proj.weight) 158 | nn.init.xavier_uniform_(self.time_proj.weight) 159 | nn.init.xavier_uniform_(self.time_compress.weight) 160 | 161 | def forward(self, x): 162 | 163 | # Along space projection 164 | x = self.spac_proj(x) 165 | 166 | # Along the time embedding 167 | x = self.time_proj(x) 168 | timeavg = self.time_avgpool(x) 169 | timemax = self.time_maxpool(x) 170 | tau = torch.cat([timeavg, timemax],-1) 171 | out = self.act(self.time_compress(tau)) 172 | return out 173 | -------------------------------------------------------------------------------- /nns/layers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Script for the encoder and decoder layer for full transformer 3 | """ 4 | 5 | from nns.attns import MultiHeadAttention, DenseEasyAttn 6 | from torch import nn 7 | 8 | class PositionWiseFeedForward(nn.Module): 9 | def __init__(self, 10 | d_model, 11 | d_ff, 12 | activation="relu"): 13 | """ 14 | The nn.Module for Feed-Forward network in transformer encoder/decoder layer 15 | 16 | Args: 17 | d_model : (Int) The dimension of embedding 18 | 19 | d_ff : (Int) The projection dimension in the FFD 20 | 21 | activation : (Str) Activation function used in network 22 | 23 | """ 24 | 25 | super(PositionWiseFeedForward, self).__init__() 26 | self.fc1 = nn.Linear(d_model, d_ff) 27 | self.fc2 = nn.Linear(d_ff, d_model) 28 | 29 | if activation == "relu": 30 | self.act = nn.ReLU() 31 | if activation == "gelu": 32 | self.act = nn.GELU() 33 | if activation == "elu": 34 | self.act = nn.ELU() 35 | 36 | 37 | def forward(self, x): 38 | return self.fc2(self.act(self.fc1(x))) 39 | 40 | 41 | 42 | class EncoderLayer(nn.Module): 43 | def __init__(self, d_model, num_heads, d_ff, dropout, act_proj): 44 | """ 45 | nn.Module for transformer Encoder layer 46 | 47 | Args: 48 | d_model : (Int) The embedding dimension 49 | 50 | num_heads : (Int) The number of heads used in attention module 51 | 52 | d_ff : (Int) Projection dimension used in Feed-Forward network 53 | 54 | dropout : (Float) The dropout value to prevent from pverfitting 55 | 56 | act_proj : (Str) The activation function used in the FFD 57 | """ 58 | super(EncoderLayer, self).__init__() 59 | self.self_attn = MultiHeadAttention(d_model, num_heads) 60 | self.feed_forward = PositionWiseFeedForward(d_model, d_ff,act_proj) 61 | self.norm1 = nn.LayerNorm(d_model) 62 | self.norm2 = nn.LayerNorm(d_model) 63 | self.dropout = nn.Dropout(dropout) 64 | 65 | def forward(self, x, mask): 66 | """ 67 | The forward prop for the module 68 | Args: 69 | x : Input sequence 70 | 71 | mask : the mask used for attention, usually be the src_mask 72 | 73 | Returns: 74 | x : The encoded sequence in latent space 75 | """ 76 | attn_output = self.self_attn(x, x, x, mask) 77 | 78 | x = self.norm1(x + self.dropout(attn_output)) 79 | ff_output = self.feed_forward(x) 80 | x = self.norm2(x + self.dropout(ff_output)) 81 | return x 82 | 83 | 84 | 85 | class easyEncoderLayer(nn.Module): 86 | def __init__(self, d_model, seqLen, num_heads, d_ff, dropout, act_proj): 87 | """ 88 | nn.Module for transformer Encoder layer 89 | 90 | Args: 91 | d_model : (Int) The embedding dimension 92 | 93 | seqLen : (Int) The length of the input sequence 94 | 95 | num_heads : (Int) The number of heads used in attention module 96 | 97 | 98 | d_ff : (Int) Projection dimension used in Feed-Forward network 99 | 100 | dropout : (Float) The dropout value to prevent from pverfitting 101 | 102 | act_proj : (Str) The activation function used in the FFD 103 | """ 104 | super(easyEncoderLayer, self).__init__() 105 | self.attn = DenseEasyAttn(d_model, seqLen, num_heads) 106 | self.feed_forward = PositionWiseFeedForward(d_model, d_ff,act_proj) 107 | self.norm1 = nn.LayerNorm(d_model) 108 | self.norm2 = nn.LayerNorm(d_model) 109 | self.dropout = nn.Dropout(dropout) 110 | 111 | def forward(self, x): 112 | """ 113 | The forward prop for the module 114 | Args: 115 | x : Input sequence 116 | 117 | Returns: 118 | x : The encoded sequence in latent space 119 | """ 120 | attn_output = self.attn(x) 121 | 122 | x = self.norm1(x + self.dropout(attn_output)) 123 | ff_output = self.feed_forward(x) 124 | x = self.norm2(x + self.dropout(ff_output)) 125 | return x 126 | 127 | -------------------------------------------------------------------------------- /nns/transformer.py: -------------------------------------------------------------------------------- 1 | """ 2 | The transfomer encoders using new embeding layer 3 | 4 | """ 5 | 6 | from nns.layers import EncoderLayer, easyEncoderLayer 7 | from nns.embedding import TimeSpaceEmbedding 8 | from torch import nn 9 | 10 | class EmbedTransformerEncoder(nn.Module): 11 | 12 | 13 | def __init__(self, d_input, d_output, n_mode, 14 | d_proj, d_model, d_ff, 15 | num_head,num_layer, 16 | act_proj= "relu", 17 | dropout= 1e-5) -> None: 18 | """ 19 | A transformer-based architecture using temporal-spatial embedding and a stack of encoder 20 | here we use self attention as the attention mechanism 21 | """ 22 | super(EmbedTransformerEncoder,self).__init__() 23 | 24 | self.embed = TimeSpaceEmbedding(d_input, n_mode, d_proj, d_model) 25 | 26 | self.encoders = nn.ModuleList([ EncoderLayer(d_model= d_model, num_heads=num_head ,d_ff = d_ff,act_proj=act_proj,dropout=dropout ) for _ in range(num_layer)]) 27 | 28 | self.cf = nn.Conv1d(d_proj, d_output,1) 29 | 30 | self.of = nn.Linear(d_model,n_mode) 31 | 32 | nn.init.xavier_uniform_(self.cf.weight) 33 | nn.init.xavier_uniform_(self.of.weight) 34 | nn.init.zeros_(self.cf.bias) 35 | nn.init.zeros_(self.of.bias) 36 | 37 | def forward(self, src): 38 | enc_input = self.embed(src) 39 | # Leave the residual for forward porp 40 | enc_res = 0 41 | for enc_layer in self.encoders: 42 | enc_input = enc_layer(enc_input+enc_res,None) 43 | enc_res = enc_input 44 | 45 | x = self.cf(enc_input) 46 | x = self.of(x) 47 | 48 | return x 49 | 50 | 51 | 52 | class easyTransformerEncoder(nn.Module): 53 | """ 54 | A transformer-based architecture using temporal-spatial embedding and a stack of encoder 55 | here we use self attention as the attention mechanism 56 | """ 57 | 58 | def __init__(self, d_input, d_output, seqLen, 59 | d_proj, d_model, d_ff, 60 | num_head,num_layer, 61 | act_proj= "relu", 62 | dropout= 1e-5) -> None: 63 | super(easyTransformerEncoder,self).__init__() 64 | 65 | self.embed = TimeSpaceEmbedding(d_input, seqLen, d_proj, d_model) 66 | 67 | self.encoders = nn.ModuleList([ easyEncoderLayer(d_model = d_model, 68 | seqLen = d_proj, 69 | num_heads = num_head , 70 | d_ff = d_ff, 71 | act_proj = act_proj, 72 | dropout = dropout ) for _ in range(num_layer)]) 73 | 74 | self.cf = nn.Conv1d(d_proj, d_output,1) 75 | 76 | self.of = nn.Linear(d_model,seqLen) 77 | 78 | nn.init.xavier_uniform_(self.cf.weight) 79 | nn.init.xavier_uniform_(self.of.weight) 80 | nn.init.zeros_(self.cf.bias) 81 | nn.init.zeros_(self.of.bias) 82 | 83 | def forward(self, src): 84 | enc_input = self.embed(src) 85 | # print(enc_input.shape) 86 | # Leave the residual for forward porp 87 | 88 | for enc_layer in self.encoders: 89 | enc_input = enc_layer(enc_input) 90 | 91 | x = self.cf(enc_input) 92 | # print(x.shape) 93 | x = self.of(x) 94 | 95 | return x 96 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # $\beta$-Variational autoencoders and transformers for reduced-order modelling of fluid flow 2 | 3 | ## Introduction 4 | The code in this repository features a Python implementation of reduced-order model (ROM) of turbulent flow using $\beta$-variational autoencoders and transformer neural network. More details about the implementation and results from the training are available in ["$\beta$-Variational autoencoders and transformers for reduced-order modelling of fluid flow",Alberto Solera-Rico, Carlos Sanmiguel Vila, M. A. Gómez, Yuning Wang, Abdulrahman Almashjary, Scott T. M. Dawson, Ricardo Vinuesa](https://doi.org/10.1038/s41467-024-45578-4) 5 | 6 | ## Data availabilty 7 | 1. We share the down-sampled data in [zenodo](https://zenodo.org/records/10501216). 8 | 9 | 2. We share the pre-trained models of $\beta$-VAE, transformers and LSTM with this repository. 10 | 11 | ## Training and inference 12 | 13 | + To train and inference the easy-attention-based transformer, please run: 14 | 15 | python main.py -re 40 -m run -nn easy 16 | 17 | 18 | ## Structure 19 | 20 | + *data*: Dataset used for the present study 21 | 22 | + *lib*: The main code used in the present study 23 | 24 | + *utils*: Support functions for visualisation, etc. 25 | 26 | + *configs*: Configurations of hyper parameters for models 27 | 28 | + *nns*: The architecture of neural networks 29 | 30 | + *res*: Storage of prediction results 31 | 32 | + *figs*: Storaging figures output 33 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | h5py~=3.9.0 2 | numpy~=1.24.4 3 | torch~=2.0.1 4 | scipy~=1.10.1 5 | matplotlib~=3.7.1 6 | scikit-learn~=1.3.1 7 | tqdm~=4.65.0 8 | tbparse~=0.0.7 9 | seaborn~=0.13.0 10 | pandas~=2.0.3 -------------------------------------------------------------------------------- /res/easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/res/easyAttn_64in_64dmodel_1next_2dim_timeemb_4h_4nb_128ff_reluact_Noneoutact_100Epoch_135000N_TrueES_50P.npz -------------------------------------------------------------------------------- /res/modes_Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000.hdf5: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/res/modes_Re40_smallerCNN_beta0.005_wDecay0_dim2_lr0.0002OneCycleLR1e-05_bs256_epochs1000.hdf5 -------------------------------------------------------------------------------- /utils/__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Chaotic Nature Analysis Using Poinare Map 3 | """ 4 | 5 | import scipy.stats as st 6 | import matplotlib.pyplot as plt 7 | import numpy as np 8 | # https://towardsdatascience.com/simple-example-of-2d-density-plots-in-python-83b83b934f67 9 | 10 | 11 | 12 | def t_zero_cross(data,direction_positive = True): 13 | zero_crossings = np.where( np.diff(np.signbit(data)) ) 14 | if (direction_positive): 15 | positive = np.where( np.diff(data) > 0 ) 16 | else: 17 | positive = np.where( np.diff(data) < 0 ) 18 | 19 | crossings_idx = np.intersect1d(zero_crossings, positive) 20 | crossings_idx_next = crossings_idx +1 21 | 22 | x_values = crossings_idx - data[list(crossings_idx)]/(data[list(crossings_idx_next)] - data[list(crossings_idx)]) 23 | 24 | 25 | return crossings_idx, x_values 26 | 27 | 28 | def PDF(ax,x,y,xmin,xmax,ymin,ymax,color): 29 | 30 | # Create meshgrid 31 | xx, yy = np.mgrid[xmin:xmax:50j, ymin:ymax:50j] 32 | 33 | positions = np.vstack([xx.ravel(), yy.ravel()]) 34 | values = np.vstack([x, y]) 35 | kernel = st.gaussian_kde(values) 36 | f = np.reshape(kernel(positions).T, xx.shape) 37 | 38 | cset = ax.contour(xx, yy, f, colors=color) 39 | ax.set_aspect('equal', 'box') 40 | ax.set_xlim(xmin, xmax) 41 | ax.grid(visible=True,markevery=1,color='gainsboro', zorder=1) 42 | 43 | def completePDFmatrixFigure(series, positive_direction, plane, fig=None, axs=None, color='k'): 44 | lims=1 45 | nmodes = series.shape[1] 46 | 47 | if fig is None: 48 | fig, axs = plt.subplots(nmodes, nmodes, figsize=(25, 25)) 49 | 50 | zeros, x_cross = t_zero_cross(series[:,plane],positive_direction) 51 | 52 | intersections = np.zeros((zeros.shape[0],nmodes)) 53 | for axis in range(0,nmodes): 54 | intersections[:,axis] = np.interp(x_cross, np.arange(series.shape[0]), series[:,axis]) 55 | 56 | 57 | for i in range(0,nmodes): 58 | for j in range(0,nmodes): 59 | if(i==j or j==plane or i==plane or j>i): 60 | axs[i,j].set_visible(False) 61 | continue 62 | PDF(axs[i,j], intersections[:,i], intersections[:,j],-lims,lims,-lims,lims,color) 63 | #axs[i,j].set_title('Section r{}=0'.format(ax+1)) 64 | axs[i,j].set_xlabel('$r_{}$'.format(i), fontsize="large") 65 | axs[i,j].set_ylabel('$r_{}$'.format(j), fontsize="large") 66 | 67 | fig.tight_layout() 68 | return fig, axs 69 | 70 | 71 | def detailPDFfigure(series, positive_direction, plane, fig=None, axs=None, color='k', i=1, j=1, model='',dpi= 100): 72 | lims=1 73 | nmodes = series.shape[1] 74 | 75 | if fig is None: 76 | fig, axs = plt.subplots(1, 1, figsize=(3, 3), dpi=dpi) 77 | 78 | zeros, x_cross = t_zero_cross(series[:,plane],positive_direction) 79 | 80 | intersections = np.zeros((zeros.shape[0],nmodes)) 81 | for axis in range(0,nmodes): 82 | intersections[:,axis] = np.interp(x_cross, np.arange(series.shape[0]), series[:,axis]) 83 | 84 | 85 | PDF(axs, intersections[:,i], intersections[:,j],-lims,lims,-lims,lims,color) 86 | axs.text(0.08, 0.08, '$r_{}=0$'.format(plane+1), fontsize=10, 87 | transform=axs.transAxes, bbox=dict(facecolor='white', alpha=0.4)) 88 | 89 | axs.set_xlabel('$r_{10}$', fontsize="large") #OJO 90 | axs.set_ylabel('$r_{}$'.format(j+1), fontsize="large") 91 | 92 | axs.spines['top'].set_visible(False) 93 | axs.spines['right'].set_visible(False) 94 | 95 | fig.tight_layout() 96 | return fig, axs -------------------------------------------------------------------------------- /utils/__pycache__/chaotic.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/chaotic.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/config.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/config.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/datas.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/datas.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/figs_time.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/figs_time.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/model.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/model.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plot.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/plot.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plot_time.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/plot_time.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/plt_rc_setup.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/plt_rc_setup.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/pp.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/pp.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/reverse_config.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/reverse_config.cpython-38.pyc -------------------------------------------------------------------------------- /utils/__pycache__/train.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KTH-FlowAI/beta-Variational-autoencoders-and-transformers-for-reduced-order-modelling-of-fluid-flows/c20c065d9149cb8d53307e0bf3261401233c901f/utils/__pycache__/train.cpython-38.pyc -------------------------------------------------------------------------------- /utils/figs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Function supports visualisation 3 | 4 | Author: @yuningw & @alsolera 5 | 6 | Editting: @yuningw 7 | """ 8 | 9 | import matplotlib.pyplot as plt 10 | import numpy as np 11 | 12 | 13 | # -------------------------------------------------------- 14 | def plot_training(logFile, path): 15 | """Plot from tensorboard file""" 16 | from tbparse import SummaryReader 17 | 18 | def loadLog(logFile): 19 | df = SummaryReader(logFile, pivot=True).scalars 20 | df = df[['step', 'General loss/KLD', 'General loss/KLD_test', 'General loss/MSE', 21 | 'General loss/MSE_test', 'General loss/Total', 22 | 'General loss/Total_test']] 23 | list(df.columns) 24 | return df['step'].to_numpy(), df['General loss/KLD'].to_numpy(), df['General loss/MSE'].to_numpy(), df[ 25 | 'General loss/MSE_test'].to_numpy(), df['General loss/KLD_test'].to_numpy() 26 | 27 | step, KL, MSE, MSE_t, KL_t = loadLog(logFile) 28 | 29 | fig, ax1 = plt.subplots(figsize=(5, 2)) 30 | ax2 = ax1.twinx() # instantiate a second axes that shares the same x-axis 31 | 32 | ax1.plot(step, KL, '-', label='KL', color='lightseagreen') 33 | ax2.plot(step, MSE, '-', label='MSE', color='saddlebrown') 34 | ax1.plot(step, KL_t, ':', label='KL test', color='lightseagreen') 35 | ax2.plot(step, MSE_t, ':', label='MSE test', color='saddlebrown') 36 | 37 | ax1.set_xlabel('Train step') 38 | # ax1.axis(xmin=0, xmax=1000) 39 | # ax1.axis(ymin=1) 40 | ax1.set_ylabel('KL loss', color='lightseagreen') 41 | ax1.tick_params(axis='y', labelcolor='lightseagreen') 42 | 43 | ax2.set_ylabel('Reconstruction loss', color='saddlebrown') 44 | ax2.tick_params(axis='y', labelcolor='saddlebrown') 45 | # ax2.axis(ymax=0.5) 46 | # plt.ylabel('Average error (L2)') 47 | # plt.legend() 48 | 49 | ax1.spines['top'].set_visible(False) 50 | ax2.spines['top'].set_visible(False) 51 | ax1.grid(color='whitesmoke', zorder=1) 52 | 53 | plt.legend() 54 | 55 | plt.savefig(path + 'training.png', format='png', bbox_inches="tight") 56 | 57 | 58 | # -------------------------------------------------------- 59 | 60 | def annot_max(x, y, ax=None): 61 | import numpy as np 62 | 63 | xmax = x[np.argmax(y)] 64 | ymax = y.max() 65 | text = "$f/f_c={:.3f}$".format(xmax) 66 | if not ax: 67 | ax = plt.gca() 68 | bbox_props = dict(boxstyle="square,pad=0.3", fc="w", ec="k", lw=0.72) 69 | arrowprops = dict(arrowstyle="->", connectionstyle="angle,angleA=0,angleB=60") 70 | kw = dict(xycoords='data', textcoords="axes fraction", 71 | arrowprops=arrowprops, bbox=bbox_props, ha="right", va="top") 72 | ax.annotate(text, xy=(xmax, ymax), xytext=(0.94, 0.96), **kw) 73 | 74 | 75 | # -------------------------------------------------------- 76 | 77 | def plotCompleteModes(modes, temporalModes, numberModes, fs, order, path): 78 | """ 79 | Plot the obtained spatial modes and temporal evolution of the latent variables in frequency domain, using welch method 80 | 81 | Args: 82 | 83 | modes : (NumpyArray) Spatial modes 84 | 85 | temporal_modes : (NumpyArray) Latent varibles from VAE 86 | 87 | numberModes : (int) Number of modes to be ploted 88 | 89 | fs : (int) Sampling frequency of welch metod 90 | 91 | order : (list/NumpyArray) The ranked results of modes accroding to energy level 92 | 93 | path : (str) Path to Save figure 94 | """ 95 | from scipy import signal 96 | import numpy as np 97 | 98 | fig, ax = plt.subplots(numberModes, 3, figsize=(16, numberModes * 1.6), sharex='col') 99 | 100 | # for mode in range(numberModes): 101 | 102 | for i, mode in np.ndenumerate(order): 103 | i = i[0] 104 | Uplot = modes[mode, 0, :, :] 105 | Vplot = modes[mode, 1, :, :] 106 | 107 | Ulim = max(abs(Uplot.flatten())) 108 | Vlim = max(abs(Vplot.flatten())) 109 | 110 | im = ax[i, 0].imshow(Uplot, cmap="RdBu_r", vmin=-Ulim, vmax=Ulim, 111 | extent=[-9, 87, -14, 14]) 112 | # ax[mode,0].set_title('Mode ' + str(mode) + ', u') 113 | ax[i, 0].set_ylabel('y/c') 114 | fig.colorbar(im, ax=ax[i, 0], shrink=0.8, aspect=10) 115 | 116 | im = ax[i, 1].imshow(Vplot, cmap="RdBu_r", vmin=-Vlim, vmax=Vlim, 117 | extent=[-9, 87, -14, 14]) 118 | ax[i, 1].set_title('Mode ' + str(i + 1)) 119 | # ax[mode,1].set_ylabel('y/c') 120 | fig.colorbar(im, ax=ax[i, 1], shrink=0.8, aspect=10) 121 | 122 | f, Pxx_den = signal.welch(temporalModes[:, mode], fs, nperseg=512 * 4) 123 | ax[i, 2].plot(f, Pxx_den, color='lightseagreen') 124 | ax[i, 2].axis(xmin=0, xmax=.15) 125 | annot_max(f, Pxx_den, ax=ax[i, 2]) 126 | ax[i, 2].grid(color='whitesmoke', zorder=1) 127 | ax[i, 2].spines['top'].set_visible(False) 128 | ax[i, 2].spines['right'].set_visible(False) 129 | if mode == (numberModes - 1): 130 | ax[i, 0].set_xlabel('$x/c$') 131 | ax[i, 1].set_xlabel('$x/c$') 132 | ax[i, 2].set_xlabel('$f/f_c$') 133 | 134 | plt.savefig(path + 'modes.png', format='png', bbox_inches="tight") 135 | 136 | 137 | # -------------------------------------------------------- 138 | 139 | def correlationMatrix(temporalModes, order, path): 140 | """ 141 | Visualisation of the correlation matrix to demonstrate the orthogonality 142 | 143 | Args: 144 | 145 | temporalModes : (NumpyArray) Latent variables encoded by VAE 146 | 147 | order : (list/NumpyArray) The ranked results of modes accroding to energy level 148 | 149 | path : (str) Path to Save figure 150 | 151 | """ 152 | import pandas as pd 153 | import seaborn as sns 154 | import numpy as np 155 | 156 | df = pd.DataFrame(temporalModes[:, order]) 157 | 158 | # Create the correlation matrix 159 | corr = df.corr() 160 | 161 | # Generate a mask for the upper triangle; True = do NOT show 162 | mask = np.zeros_like(corr, dtype=bool) 163 | mask[np.triu_indices_from(mask)] = True 164 | 165 | # Set up the matplotlib figure 166 | f, ax = plt.subplots(figsize=(4, 3)) 167 | 168 | # Generate a colormap 169 | cmap = sns.color_palette("icefire", as_cmap=True) 170 | 171 | n_modes = temporalModes.shape[1] 172 | 173 | axis_labels = np.arange(1, n_modes + 1).tolist() 174 | axis_labels = list(map(str, axis_labels)) 175 | if n_modes >= 20: 176 | for i in range(0, n_modes, 2): 177 | axis_labels[i] = '' 178 | 179 | # Draw the heatmap with the mask and correct aspect ratio 180 | # More details at https://seaborn.pydata.org/generated/seaborn.heatmap.html 181 | sns.heatmap( 182 | np.abs(corr), # The data to plot 183 | # mask=mask, # Mask some cells 184 | cmap=cmap, # What colors to plot the heatmap as 185 | annot=False, # Should the values be plotted in the cells? 186 | fmt=".2f", 187 | vmax=1, # The maximum value of the legend. All higher vals will be same color 188 | vmin=0, # The minimum value of the legend. All lower vals will be same color 189 | # center=0.75, # The center value of the legend. With divergent cmap, where white is 190 | square=True, # Force cells to be square 191 | linewidths=0.5, # Width of lines that divide cells 192 | # cbar_kws={"shrink": .5} # Extra kwargs for the legend; in this case, shrink by 50% 193 | # norm=LogNorm() 194 | xticklabels=axis_labels, 195 | yticklabels=axis_labels, 196 | ) 197 | plt.xticks(rotation=0) 198 | plt.yticks(rotation=0) 199 | 200 | plt.savefig(path + 'matrix.png', format='png', bbox_inches="tight") 201 | 202 | 203 | # -------------------------------------------------------- 204 | 205 | def plotTemporalSeries(modes, path): 206 | """ 207 | 208 | Visualize the temproal evolution of the latent variables 209 | 210 | modes : (NumpyArray) The latent variables from VAE 211 | 212 | path : (str) Path to Save Figure 213 | 214 | """ 215 | latent_dim = modes.shape[1] 216 | fig, ax = plt.subplots(latent_dim, 1, figsize=(16, latent_dim * 1.6), sharex='col') 217 | 218 | for i in range(modes.shape[1]): 219 | ax[i].plot(modes[:1000, i], label='Mode ' + str(i)) 220 | ax[i].axis(ymin=-2.5, ymax=2.5) 221 | 222 | ax[i].legend() 223 | ax[i].grid() 224 | 225 | plt.savefig(path + 'series.png', format='png', bbox_inches="tight") 226 | 227 | 228 | # -------------------------------------------------------- 229 | def plotEcum(Ecum, path): 230 | """ 231 | Show the accumlative energy level 232 | 233 | Ecum : (NumpyArray) Obtained energy level 234 | 235 | path : (str) Path to Save Figure 236 | 237 | """ 238 | import numpy as np 239 | 240 | fig = plt.figure() 241 | plt.plot(np.arange(1, Ecum.shape[0] + 1), Ecum) 242 | plt.xlabel('Number of modes') 243 | plt.ylabel('Cumulative Ek') 244 | plt.grid() 245 | 246 | plt.savefig(path + 'Ecum.png', format='png', bbox_inches="tight") 247 | 248 | 249 | # -------------------------------------------------------- 250 | def plotNLmodeField(modes, values, path): 251 | """ 252 | Visualize the non-linear mode 253 | 254 | Args: 255 | 256 | modes : (NumpyArray) The spatial modes using decoder. 257 | 258 | values : (float) A non-zero value as the element in latent vector 259 | 260 | path : (str) Path to Save figure 261 | 262 | """ 263 | import numpy as np 264 | 265 | valuesToPlot = np.array([-2., -1., 0., 1., 2.]) 266 | 267 | # Find the indices of values_array elements in main_array 268 | indices = [] 269 | for value in valuesToPlot: 270 | matching_indices = np.where(np.abs(values - value) < 1e-3)[0] 271 | indices.extend(matching_indices) 272 | 273 | indices_array = np.array(indices) 274 | values = values[indices_array] 275 | 276 | Ulim = max(abs(modes[indices_array, 0, :, :].flatten())) 277 | Vlim = max(abs(modes[indices_array, 1, :, :].flatten())) 278 | 279 | fig, ax = plt.subplots(len(indices_array), 2, figsize=(10, 8), sharex='col', sharey='row') 280 | 281 | for idx, value in enumerate(values): 282 | print(idx, value) 283 | Uplot = modes[indices_array[idx], 0, :, :] 284 | Vplot = modes[indices_array[idx], 1, :, :] 285 | 286 | im = ax[idx, 0].imshow(Uplot, cmap="RdBu_r", vmin=-Ulim, vmax=Ulim, extent=[-9, 87, -14, 14]) 287 | # ax[idx, 0].set_title('Value {}'.format(round(value, 1))) 288 | ax[idx, 0].set_ylabel('y/c') 289 | fig.colorbar(im, ax=ax[idx, 0], shrink=0.7, aspect=10) 290 | 291 | im = ax[idx, 1].imshow(Vplot, cmap="RdBu_r", vmin=-Vlim, vmax=Vlim, extent=[-9, 87, -14, 14]) 292 | # ax[idx,1].set_title('Mode {}, value {}, v'.format(mode, round(NLvalues[value],2))) 293 | fig.colorbar(im, ax=ax[idx, 1], shrink=0.7, aspect=10) 294 | 295 | fig.text(0.48, 0.965 - idx * .184, '$s_1 = {}$'.format(round(value, 1)), fontsize=11, ha="center") 296 | 297 | ax[idx, 0].set_xlabel('x/c') 298 | ax[idx, 1].set_xlabel('x/c') 299 | 300 | fig.set_tight_layout(True) 301 | 302 | plt.savefig(path + 'NLfield.png', format='png', bbox_inches="tight") 303 | 304 | 305 | def vis_bvae(modes_file, training_file): 306 | """ 307 | Visualisation of the beta-VAE results 308 | 309 | Args: 310 | modes_file : The file saves the post-processing results of VAE 311 | 312 | training_file : The history and log of training the model 313 | 314 | """ 315 | import h5py 316 | from lib.init import pathsBib 317 | from pathlib import Path 318 | 319 | path = pathsBib.fig_path + 'bVAE/' 320 | Path(path).mkdir(exist_ok=True) 321 | 322 | try: 323 | plot_training(training_file) 324 | except: 325 | print(f'Warning: {training_file} could not be loaded') 326 | 327 | with h5py.File(modes_file, 'r') as f: 328 | print("Keys: %s" % f.keys()) 329 | temporalModes = f['vector'][:, :] 330 | temporalModes_test = f['vector_test'][:, :] 331 | modes = f['modes'][:, :] 332 | order = f['order'][:] 333 | Ecum = f['Ecum'][:] 334 | NLvalues = f['NLvalues'][:] 335 | NLmodes = f['NLmodes'][:] 336 | 337 | # Re40 case is sampled at 1tc, Re100 case is sampled at tc/5 338 | if 'Re40' in modes_file: 339 | fs = 1 340 | else: 341 | fs = 5 342 | 343 | plotNLmodeField(NLmodes, NLvalues, path) 344 | plotCompleteModes(modes, temporalModes, modes.shape[0], fs, order, path) 345 | plotTemporalSeries(temporalModes, path) 346 | correlationMatrix(temporalModes, order, path) 347 | plotEcum(Ecum, path) 348 | 349 | def vis_pod(POD): 350 | """ 351 | Visualisaton of POD results 352 | 353 | Args: 354 | 355 | POD : (lib.POD.POD) The running object for implmenting POD 356 | 357 | """ 358 | 359 | import h5py 360 | from lib.init import pathsBib 361 | from pathlib import Path 362 | 363 | path = pathsBib.fig_path + 'POD/' 364 | Path(path).mkdir(exist_ok=True) 365 | 366 | # Re40 case is sampled at 1tc, Re100 case is sampled at tc/5 367 | if POD.re==40: 368 | fs = 1 369 | else: 370 | fs = 5 371 | 372 | shape = [POD.spatial_modes.shape[1], 373 | POD.u_train.shape[1], 374 | POD.u_train.shape[2], 375 | POD.u_train.shape[3]] 376 | sm = np.swapaxes(POD.spatial_modes, 0, 1).reshape(shape) 377 | 378 | plotCompleteModes(sm, 379 | POD.temporal_modes, 380 | POD.n_modes, 381 | fs, 382 | range(POD.n_modes), 383 | path) 384 | 385 | plotTemporalSeries(POD.temporal_modes, path) 386 | plotEcum(POD.Ek_nm, path) 387 | -------------------------------------------------------------------------------- /utils/figs_time.py: -------------------------------------------------------------------------------- 1 | 2 | """ 3 | The Visualisation of time-series prediction results 4 | 5 | @yuningw 6 | """ 7 | 8 | from lib.runners import latentRunner, vaeRunner 9 | 10 | 11 | #----------------------------------------------------------- 12 | class colorplate: 13 | """ 14 | Color in HTML format for plotting 15 | """ 16 | 17 | red = "r" 18 | blue = "b" 19 | yellow = "y" 20 | cyan = "c" 21 | black = "k" 22 | 23 | #----------------------------------------------------------- 24 | def vis_temporal_Prediction( 25 | vae :vaeRunner, 26 | predictor :latentRunner, 27 | model_type, 28 | if_loss = True, 29 | if_evo = True, 30 | if_window = True, 31 | if_pmap_s = False, 32 | if_pmap_all = False, 33 | if_snapshot = True 34 | ): 35 | """ 36 | Visualisation of the temporal-dynamics prediction results 37 | 38 | Args: 39 | 40 | vae : (lib.runners.vaeRunner) The module for employed VAE 41 | 42 | predictor : (lib.runners.latentRunner) The module for latent-space temporal-dynamics predictions 43 | 44 | model_type : (str) The type of model used (easy/self/lstm) 45 | 46 | if_loss : (bool) If plot the loss evolution 47 | 48 | if_evo : (bool) If plot the temporal evolution of latent mode 49 | 50 | if_window : (bool) If plot the l2-norm error horizon 51 | 52 | if_pmap_s : (bool) If plot the single Poincare Map 53 | 54 | if_pmap_all : (bool) If plot all the Poincare Maps 55 | 56 | if_snapshot : (bool) If plot the flow field reconsturction 57 | 58 | """ 59 | 60 | from pathlib import Path 61 | from lib.init import pathsBib 62 | import numpy as np 63 | from lib.pp_time import make_physical_prediction 64 | 65 | figPath = pathsBib.fig_path + model_type + '/' 66 | case_name = predictor.filename 67 | datPath = pathsBib.res_path + case_name + '.npz' 68 | 69 | print("#"*30) 70 | print(f"Start visualisation:\nSave Fig to:{figPath}\nLoad data from:{datPath}") 71 | 72 | try: 73 | d = np.load(datPath) 74 | g = d['g'] 75 | p = d['p'] 76 | e = d['e'] 77 | pmap_g = d['pmap_g'] 78 | pmap_p = d['pmap_p'] 79 | 80 | except: 81 | print(f"ERROR: FAILD loading data") 82 | 83 | Path(figPath).mkdir(exist_ok=True) 84 | 85 | if if_loss: 86 | plot_loss(predictor.history, save_file= figPath + "loss_" + case_name + '.jpg' ) 87 | print(f'INFO: Loss Evolution Saved!') 88 | 89 | if if_evo and (g.any() !=None) and (p.any() != None): 90 | plot_signal(g,p, save_file=figPath + "signal_" + case_name + '.jpg' ) 91 | print(f"INFO: Prediction Evolution Saved!") 92 | 93 | if if_window and (e.any() != None): 94 | plot_pred_horizon_error(e, colorplate.blue, save_file=figPath + "horizon_" + case_name + '.jpg' ) 95 | print(f"INFO: l2-norm error Horizion Saved!") 96 | 97 | ## Poincare Map 98 | planeNo = 0 99 | postive_dir = True 100 | lim_val = 2.5 # Limitation of x and y bound when compute joint pdf 101 | grid_val = 50 102 | i = 1; j =1 103 | 104 | if if_pmap_s and (pmap_g.any() != None) and (pmap_p.any() != None): 105 | plotSinglePoincare( planeNo, i, j, 106 | pmap_p,pmap_g, 107 | lim_val, grid_val, 108 | save_file = figPath + f'Pmap_{i}_{j}_' + case_name + '.jpg') 109 | print(f"INFO: Single Poincare Map of {i}, {j} Saved!") 110 | 111 | if if_pmap_all and (pmap_g.any() != None) and (pmap_p.any() != None): 112 | plotCompletePoincare(predictor.config.latent_dim,planeNo, 113 | pmap_p, pmap_g, 114 | lim_val, grid_val, 115 | save_file = None, 116 | dpi = 200) 117 | print(f"INFO: Complete Poincare Map Saved!") 118 | 119 | if if_snapshot and (p.any() != None) and (g.any() != None): 120 | VAErec, pred = make_physical_prediction(vae=vae,pred_latent=p,true_latent=g,device=vae.device) 121 | 122 | stepPlot = int(predictor.config.in_dim + 1) # Here we test the prediction purely based on the predicted variables 123 | 124 | predFieldFigure(vae.test_d,VAErec,pred, 125 | vae.std,vae.mean, 126 | stepPlot = stepPlot, 127 | model_name= model_type, 128 | save_file = figPath + "recSnapShot_" + case_name + '.jpg') 129 | print(f"INFO: Reconstruted Snapshot at {stepPlot} Saved!") 130 | return 131 | 132 | 133 | #----------------------------------------------------------- 134 | def plot_loss(history, save_file, dpi = 200): 135 | """ 136 | Plot the loss evolution during training 137 | Args: 138 | history : A dictionary contains loss evolution in list 139 | save_file : Path to save the figure, if None, then just show the plot 140 | dpi : The dpi for save the file 141 | 142 | Returns: 143 | A fig for loss 144 | 145 | """ 146 | import matplotlib.pyplot as plt 147 | import utils.plt_rc_setup 148 | from utils.figs_time import colorplate as cc 149 | 150 | fig, axs = plt.subplots(1,1,figsize=(12,5)) 151 | 152 | axs.semilogy(history["train_loss"],lw = 2, c = cc.blue) 153 | if len(history["val_loss"]) != 0 : 154 | axs.semilogy(history["val_loss"],lw = 2, c = cc.red) 155 | axs.set_xlabel("Epoch") 156 | axs.set_ylabel("MSE Loss") 157 | 158 | if save_file != None: 159 | plt.savefig(save_file, bbox_inches='tight', dpi= dpi) 160 | 161 | 162 | #----------------------------------------------------------- 163 | def plot_signal(test_data, Preds, save_file:None, dpi = 200): 164 | """ 165 | Plot the temproal evolution of prediction and test data 166 | Args: 167 | test_data : A numpy array of test data 168 | Preds : A numpy array of prediction data 169 | 170 | save_file : Path to save the figure, if None, then just show the plot 171 | dpi : The dpi for save the file 172 | 173 | Returns: 174 | A fig for temporal dynamic of ground truth and predictions on test data 175 | 176 | """ 177 | import sys 178 | 179 | try: 180 | test_data.shape == Preds.shape 181 | except: 182 | print("The prediction and test data must have same shape!") 183 | sys.exit() 184 | 185 | import matplotlib.pyplot as plt 186 | import utils.plt_rc_setup 187 | from utils.figs_time import colorplate as cc 188 | 189 | Nmode = min(test_data.shape[0],test_data.shape[-1]) 190 | 191 | fig, axs = plt.subplots(Nmode,1,figsize=(16,2.5* Nmode),sharex=True) 192 | 193 | for i, ax in enumerate(axs): 194 | ax.plot(test_data[:,i],c = cc.black,lw = 1.5) 195 | ax.plot(Preds[:,i],c = cc.blue,lw = 1.5) 196 | ax.set_ylabel(f"M{i+1}") 197 | axs[-1].set_xlabel("t",fontsize=20) 198 | 199 | ax.legend(["Ground truth",'Prediction']) 200 | 201 | if save_file != None: 202 | plt.savefig(save_file, bbox_inches='tight', dpi= dpi) 203 | 204 | 205 | #----------------------------------------------------------- 206 | 207 | def plot_pred_horizon_error(window_err, Color, save_file, dpi=300): 208 | """ 209 | Viusalize the latent-space prediction horizon error 210 | 211 | Args: 212 | 213 | window_err : (NumpyArray) The horizon of l2-norm error of prediction 214 | 215 | Color : (str) The color for the line 216 | 217 | save_file : Path to save the figure, if None, then just show the plot 218 | 219 | dpi : The dpi for save the file 220 | 221 | """ 222 | import matplotlib.pyplot as plt 223 | 224 | fig, axs = plt.subplots(1,1, figsize = (8,4)) 225 | 226 | axs.plot(window_err, lw = 3, c = Color) 227 | axs.set_xlabel("Prediction steps",fontsize = 20) 228 | axs.set_ylabel(r"$\epsilon$",fontsize = 20) 229 | 230 | if save_file != None: 231 | plt.savefig(save_file, bbox_inches='tight', dpi= dpi) 232 | 233 | return 234 | 235 | 236 | #----------------------------------------------------------- 237 | def plotSinglePoincare( planeNo, i, j, 238 | InterSec_pred,InterSec_test, 239 | lim_val, grid_val, 240 | save_file:None, 241 | dpi = 200): 242 | """ 243 | 244 | Visualisation of a single Poincare Map for test data and prediction 245 | 246 | Args: 247 | 248 | planeNo : (int) The plane no to compute the intersection 249 | 250 | i : (int) The Number of the mode on x-Axis 251 | 252 | j : (int) The Number of the mode on y-Axis 253 | 254 | lim_val : (float) Limitation of region on the map 255 | 256 | grid_val : (int) Number of the mesh grid 257 | 258 | save_file : (str) Path to save the file 259 | 260 | dpi : (int) The dpi for the image 261 | 262 | """ 263 | import matplotlib.pyplot as plt 264 | import utils.plt_rc_setup 265 | from utils.figs_time import colorplate as cc 266 | from lib.pp_time import PDF 267 | 268 | fig, axs = plt.subplots(1, 1, figsize=(5, 5)) 269 | 270 | _, _, pdf_test = PDF(InterSecX=InterSec_test[:, i], 271 | InterSecY=InterSec_test[:, j], 272 | xmin=-lim_val, xmax=lim_val, 273 | ymin=-lim_val, ymax=lim_val, 274 | x_grid=grid_val, y_grid=grid_val, 275 | ) 276 | 277 | xx, yy, pdf_pred = PDF(InterSecX=InterSec_pred[:, i], 278 | InterSecY=InterSec_pred[:, j], 279 | xmin=-lim_val, xmax=lim_val, 280 | ymin=-lim_val, ymax=lim_val, 281 | x_grid=grid_val, y_grid=grid_val, 282 | ) 283 | 284 | axs.contour(xx, yy, pdf_test, colors=cc.black) 285 | axs.contour(xx, yy, pdf_pred, colors='lightseagreen') 286 | axs.set_xlim(-lim_val, lim_val) 287 | axs.text(0.80, 0.08, '$r_{}=0$'.format(planeNo+1), fontsize=16, 288 | transform=axs.transAxes, bbox=dict(facecolor='white', alpha=0.4)) 289 | axs.set_xlabel(f"$r_{i + 1}$", fontsize='large') 290 | axs.set_ylabel(f"$r_{j + 1}$", fontsize='large') 291 | axs.set_aspect('equal', "box") 292 | axs.spines['top'].set_visible(False) 293 | axs.spines['right'].set_visible(False) 294 | axs.grid(visible=True, markevery=1, color='gainsboro', zorder=1) 295 | if save_file != None: 296 | plt.savefig(save_file, bbox_inches = 'tight', dpi = dpi) 297 | 298 | 299 | return 300 | 301 | 302 | #----------------------------------------------------------- 303 | def plotCompletePoincare(Nmodes,planeNo, 304 | InterSec_pred, InterSec_test, 305 | lim_val, grid_val, 306 | save_file = None, 307 | dpi = 200): 308 | """ 309 | 310 | Visualisation of whole Poincare Maps for test data and prediction 311 | 312 | 313 | Args: 314 | 315 | planeNo : (int) The plane no to compute the intersection 316 | 317 | lim_val : (float) Limitation of region on the map 318 | 319 | grid_val : (int) Number of the mesh grid 320 | 321 | save_file : (str) Path to save the file 322 | 323 | dpi : (int) The dpi for the image 324 | 325 | 326 | """ 327 | import matplotlib.pyplot as plt 328 | import utils.plt_rc_setup 329 | from utils.figs_time import colorplate as cc 330 | from lib.pp_time import PDF 331 | 332 | 333 | fig, axs = plt.subplots(Nmodes, Nmodes, 334 | figsize=(5 * Nmodes, 5 * Nmodes), 335 | sharex=True, sharey=True) 336 | 337 | for i in range(0, Nmodes): 338 | for j in range(0, Nmodes): 339 | if i == j or j == planeNo or i == planeNo or j > i: 340 | axs[i, j].set_visible(False) 341 | continue 342 | 343 | _, _, pdf_test = PDF(InterSecX=InterSec_test[:, i], 344 | InterSecY=InterSec_test[:, j], 345 | xmin=-lim_val, xmax=lim_val, 346 | ymin=-lim_val, ymax=lim_val, 347 | x_grid=grid_val, y_grid=grid_val, 348 | ) 349 | 350 | xx, yy, pdf_pred = PDF(InterSecX=InterSec_pred[:, i], 351 | InterSecY=InterSec_pred[:, j], 352 | xmin=-lim_val, xmax=lim_val, 353 | ymin=-lim_val, ymax=lim_val, 354 | x_grid=grid_val, y_grid=grid_val, 355 | ) 356 | 357 | axs[i, j].contour(xx, yy, pdf_test, colors=cc.black) 358 | axs[i, j].contour(xx, yy, pdf_pred, colors='lightseagreen') 359 | axs[i, j].set_xlim(-lim_val, lim_val) 360 | axs[i, j].set_xlabel(f"$r_{i + 1}$", fontsize='large') 361 | axs[i, j].set_ylabel(f"$r_{j + 1}$", fontsize='large') 362 | axs[i, j].set_aspect('equal', "box") 363 | axs[i, j].spines['top'].set_visible(False) 364 | axs[i, j].spines['right'].set_visible(False) 365 | axs[i, j].grid(visible=True, markevery=1, color='gainsboro', zorder=1) 366 | 367 | if save_file != None: 368 | plt.savefig(save_file, bbox_inches = 'tight', dpi = dpi) 369 | 370 | return 371 | 372 | 373 | 374 | #----------------------------------------------------------- 375 | 376 | def predFieldFigure(true, VAErec, pred, std_data, mean_data, stepPlot, model_name, save_file, dpi=200): 377 | 378 | """ 379 | 380 | Visualise the flow fields reconstructed by the latent-space prediction from the transformer/lstm 381 | 382 | true : (NumpyArray) The ground truth 383 | 384 | VAErec : (NumpyArray) The reconstruction from VAE ONLY 385 | 386 | pred : (NumpyArray) The reconstruction from the prediction of transformer 387 | 388 | std_data : (NumpyArray) Std of flow fields 389 | 390 | mean_data : (NumpyArray) Mean of flow fields 391 | 392 | model_name : (str) The name of the predictor model: easy/self/lstm 393 | 394 | save_file : (str) Path to save the file 395 | 396 | dpi : (int) The dpi for the image 397 | 398 | """ 399 | 400 | import matplotlib.pyplot as plt 401 | fig, ax = plt.subplots(2, 3, figsize=(11, 4), sharex='col', sharey='row') 402 | 403 | Umax = 1.5 404 | Umin = 0 405 | Vlim = 1 406 | 407 | # From dataset 408 | true_u = (true[stepPlot, 0, :, :] * std_data[0, 0, :, :] + mean_data[0, 0, :, :]).squeeze() 409 | true_v = (true[stepPlot, 1, :, :] * std_data[0, 1, :, :] + mean_data[0, 1, :, :]).squeeze() 410 | 411 | vae_u = (VAErec[stepPlot, 0, :, :] * std_data[0, 0, :, :] + mean_data[0, 0, :, :]).squeeze() 412 | vae_v = (VAErec[stepPlot, 1, :, :] * std_data[0, 1, :, :] + mean_data[0, 1, :, :]).squeeze() 413 | 414 | pred_u = (pred[stepPlot, 0, :, :] * std_data[0, 0, :, :] + mean_data[0, 0, :, :]).squeeze() 415 | pred_v = (pred[stepPlot, 1, :, :] * std_data[0, 1, :, :] + mean_data[0, 1, :, :]).squeeze() 416 | 417 | im = ax[0, 0].imshow(true_u, 418 | cmap="RdBu_r", vmin=Umin, vmax=Umax, extent=[-9, 87, -14, 14]) 419 | ax[0, 0].set_title('True u\n($t+$' + (str(stepPlot) if stepPlot > 1 else "") + '$t_c$)') 420 | fig.colorbar(im, ax=ax[0, 0], shrink=0.7, ticks=([0, 0.5, 1, 1.5])) 421 | 422 | im = ax[1, 0].imshow(true_v, 423 | cmap="RdBu_r", vmin=-Vlim, vmax=Vlim, extent=[-9, 87, -14, 14]) 424 | ax[1, 0].set_title('True v\n($t+$' + (str(stepPlot) if stepPlot > 1 else "") + '$t_c$)') 425 | fig.colorbar(im, ax=ax[1, 0], shrink=0.7) 426 | 427 | # Encoded and decoded 428 | im = ax[0, 1].imshow(vae_u, 429 | cmap="RdBu_r", vmin=Umin, vmax=Umax, extent=[-9, 87, -14, 14]) 430 | ax[0, 1].set_title(r'$\beta$-VAE' + ' u\n($t+$' + (str(stepPlot) if stepPlot > 1 else "") + '$t_c$)') 431 | fig.colorbar(im, ax=ax[0, 1], shrink=0.7, ticks=([0, 0.5, 1, 1.5])) 432 | 433 | im = ax[1, 1].imshow(vae_v, 434 | cmap="RdBu_r", vmin=-Vlim, vmax=Vlim, extent=[-9, 87, -14, 14]) 435 | ax[1, 1].set_title(r'$\beta$-VAE' + ' v\n($t+$' + (str(stepPlot) if stepPlot > 1 else "") + '$t_c$)') 436 | fig.colorbar(im, ax=ax[1, 1], shrink=0.7) 437 | 438 | # Encoded, predicted and decoded 439 | im = ax[0, 2].imshow(pred_u, 440 | cmap="RdBu_r", vmin=Umin, vmax=Umax, extent=[-9, 87, -14, 14]) 441 | ax[0, 2].set_title(r'$\beta$-VAE + ' + model_name + ' u\n($t+$' + (str(stepPlot) if stepPlot > 1 else "") + '$t_c$)') 442 | fig.colorbar(im, ax=ax[0, 2], shrink=0.7, ticks=([0, 0.5, 1, 1.5])) 443 | 444 | im = ax[1, 2].imshow(pred_v, 445 | cmap="RdBu_r", vmin=-Vlim, vmax=Vlim, extent=[-9, 87, -14, 14]) 446 | ax[1, 2].set_title(r'$\beta$-VAE + ' + model_name + ' v\n($t+$' + (str(stepPlot) if stepPlot > 1 else "") + '$t_c$)') 447 | fig.colorbar(im, ax=ax[1, 2], shrink=0.7) 448 | 449 | ax[1, 0].set_xlabel('x/c') 450 | ax[1, 1].set_xlabel('x/c') 451 | ax[1, 2].set_xlabel('x/c') 452 | ax[0, 0].set_ylabel('y/c') 453 | ax[1, 0].set_ylabel('y/c') 454 | 455 | # fig.set_tight_layout(True) 456 | 457 | if save_file != None: 458 | plt.savefig(save_file, bbox_inches = 'tight', dpi = dpi) 459 | 460 | return fig, ax 461 | -------------------------------------------------------------------------------- /utils/io.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utility for saving and loading 3 | 4 | @yuningw 5 | """ 6 | 7 | import torch 8 | 9 | def save_checkpoint(state, path_name): 10 | 11 | torch.save(state, path_name) 12 | print('Saved checkpoint') 13 | 14 | 15 | def load_checkpoint(model, path_name, optimizer=None): 16 | 17 | print('Loading checkpoint') 18 | print(path_name) 19 | 20 | checkpoint = torch.load(path_name) 21 | model.load_state_dict(checkpoint['state_dict']) 22 | if optimizer is not None: 23 | optimizer.load_state_dict(checkpoint['optimizer_dict']) 24 | 25 | print('Loaded checkpoint') 26 | 27 | -------------------------------------------------------------------------------- /utils/plt_rc_setup.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | plt.rc("font",family = "serif") 3 | plt.rc("font",size = 20) 4 | plt.rc("axes",labelsize = 16, linewidth = 2) 5 | plt.rc("legend",fontsize= 12, handletextpad = 0.3) 6 | plt.rc("xtick",labelsize = 16) 7 | plt.rc("ytick",labelsize = 16) --------------------------------------------------------------------------------