merged discrepancy dropout in new discrepancy branch

.gitignore

@@ -5,4 +5,4 @@ results/
 *.npz
 *.sur
 *.csv
-.linfa/
+.linfa/

To_do.txt

@@ -0,0 +1,3 @@
+- Plot lf model plus discrepancy at arbitrary values of the variables and calibrarion model parameters. 
+- Add dropouts as an option to the mlp component, and add optional scheduler for the dropout probability. 
+- Generalize the code so it works on multiple outputs and also works for the NoFAS surrogate

linfa/discrepancy.py

@@ -19,38 +19,63 @@ def __init__(self, model_name,
  var_grid_out,
  dnn_arch=None, 
  dnn_activation='relu', 
+ dnn_dropout=None,
  model_folder='./', 
  surrogate=None, 
  device='cpu'):
-
- self.device = device
- self.input_size = input_size
- self.output_size = output_size
+
  self.model_name = model_name
  self.model_folder = model_folder
- self.is_trained = False
 
- # Assign LF model
- self.lf_model = lf_model
+ if(var_grid_in is None):
 
- # Store variable grid locations
- self.var_grid_in=var_grid_in
- # Output variables - multiple noisy observations 
- # are available for each combination of variables
- self.var_grid_out=var_grid_out
+ self.device = None
+ self.input_size = None
+ self.output_size = None
+ self.dnn_arch = None
+ self.dnn_activation = None
+ self.dnn_dropout = None
+ self.is_trained = None
+ self.lf_model = None
+ self.var_grid_in = None
+ self.var_grid_out = None
+ self.var_in_avg = None
+ self.var_in_std = None
+ self.var_out_avg = None
+ self.var_out_std = None
+ self.surrogate = None
 
- # Input/output statistics
- self.var_in_avg = torch.mean(var_grid_in,dim=0)
- if(len(self.var_grid_in) == 1):
- self.var_in_std = torch.zeros_like(self.var_in_avg)
  else:
- self.var_in_std = torch.std(var_grid_in,dim=0)
- # If there are multiple outputs, we will define one file for each output
- self.var_out_avg = torch.mean(var_grid_out)
- self.var_out_std = torch.std(var_grid_out)
 
- # Create surrogate
- self.surrogate = FNN(input_size, output_size, arch=dnn_arch, device=self.device, init_zero=True) if surrogate is None else surrogate
+ self.device = device
+ self.input_size = input_size
+ self.output_size = output_size
+ self.dnn_arch = dnn_arch
+ self.dnn_activation = dnn_activation
+ self.dnn_dropout = dnn_dropout
+ self.is_trained = False
+
+ # Assign LF model
+ self.lf_model = lf_model
+
+ # Store variable grid locations
+ self.var_grid_in=var_grid_in
+ # Output variables - multiple noisy observations 
+ # are available for each combination of variables
+ self.var_grid_out=var_grid_out
+
+ # Input/output statistics
+ self.var_in_avg = torch.mean(var_grid_in,dim=0)
+ if(len(self.var_grid_in) == 1):
+ self.var_in_std = torch.zeros_like(self.var_in_avg)
+ else:
+ self.var_in_std = torch.std(var_grid_in,dim=0)
+ # If there are multiple outputs, we will define one file for each output
+ self.var_out_avg = torch.mean(var_grid_out)
+ self.var_out_std = torch.std(var_grid_out)
+
+ # Create surrogate
+ self.surrogate = FNN(input_size, output_size, arch=self.dnn_arch, device=self.device, init_zero=True, dropout=dnn_dropout) if surrogate is None else surrogate
 
  def surrogate_save(self):
  """Save surrogate model to [self.name].sur and [self.name].npz
@@ -60,7 +85,19 @@ def surrogate_save(self):
 
  """
  # Save model state dictionary
- torch.save(self.surrogate.state_dict(), self.model_folder + self.model_name + '.sur')
+ dict_to_save = {}
+ dict_to_save['weights'] = self.surrogate.state_dict()
+ dict_to_save['grid_in'] = self.var_grid_in
+ dict_to_save['grid_stats_in'] = [self.var_in_avg,self.var_in_std]
+ dict_to_save['grid_out'] = self.var_grid_out
+ dict_to_save['grid_stats_out'] = [self.var_out_avg,self.var_out_std]
+ dict_to_save['trained'] = self.is_trained
+ dict_to_save['input_size'] = self.input_size
+ dict_to_save['output_size'] = self.output_size
+ dict_to_save['dnn_arch'] = self.dnn_arch
+ dict_to_save['device'] = self.device
+ # Save entire dictionary
+ torch.save(dict_to_save, self.model_folder +'/'+ self.model_name + '.sur')
 
  def surrogate_load(self):
  """Load surrogate model from [self.name].sur and [self.name].npz
@@ -69,7 +106,18 @@ def surrogate_load(self):
  None
  """
  # Read back the state dictionary from file
- self.surrogate.load_state_dict(torch.load(self.model_folder + self.model_name + '.sur'))
+ load_dict = torch.load(self.model_folder +'/'+ self.model_name + '.sur')
+ self.var_grid_in = load_dict['grid_in']
+ self.var_in_avg,self.var_in_std = load_dict['grid_stats_in']
+ self.var_grid_out = load_dict['grid_out']
+ self.var_out_avg,self.var_out_std = load_dict['grid_stats_out']
+ self.is_trained = load_dict['trained']
+ self.input_size = load_dict['input_size']
+ self.output_size = load_dict['output_size']
+ self.dnn_arch = load_dict['dnn_arch']
+ self.device = load_dict['device']
+ self.surrogate = FNN(self.input_size, self.output_size, arch=self.dnn_arch, device=self.device, init_zero=True)
+ self.surrogate.load_state_dict(load_dict['weights'])
 
  def update(self, batch_x, max_iters=10000, lr=0.01, lr_exp=0.999, record_interval=50, store=True, reg=False, reg_penalty=0.0001):
  """Train surrogate model with pre-grid.
@@ -79,6 +127,7 @@ def update(self, batch_x, max_iters=10000, lr=0.01, lr_exp=0.999, record_interva
  print('--- Training model discrepancy')
  print('')
 
+ # Set it as trained
  self.is_trained = True
 
  # LF model output at the current batch
@@ -104,6 +153,7 @@ def update(self, batch_x, max_iters=10000, lr=0.01, lr_exp=0.999, record_interva
  for i in range(max_iters):
  # Set surrogate in training mode
  self.surrogate.train() 
+
  # Surrogate returns a table with rows as batches and columns as variables considered 
  disc = self.surrogate(var_grid)
 
@@ -133,8 +183,14 @@ def update(self, batch_x, max_iters=10000, lr=0.01, lr_exp=0.999, record_interva
  print('')
  print('--- Surrogate model pre-train complete')
  print('') 
+ # Save if needed
  if store:
  self.surrogate_save()
+ # Put it in eval model if no dropouts are present
+ if(self.dnn_dropout is not None):
+ self.surrogate.train()
+ else:
+ self.surrogate.eval()
 
  def forward(self, var):
  """Function to evaluate the surrogate
@@ -157,7 +213,7 @@ def forward(self, var):
  else:
  return res
 
-def test_surrogate():
+def test_discrepancy():
 
  import matplotlib.pyplot as plt
  from linfa.models.discrepancy_models import PhysChem
@@ -170,7 +226,7 @@ def test_surrogate():
  model = PhysChem(var_grid)
 
  # Generate true data
- model.genDataFile(dataFileNamePrefix='observations', use_true_model=True, store=True, num_observations=10) 
+ model.genDataFile(dataFileNamePrefix='observations', use_true_model=True, store=True, num_observations=3) 
 
  # Get data from true model at the same TP conditions 
  var_data = np.loadtxt('observations.csv',skiprows=1,delimiter=',')
@@ -179,22 +235,29 @@ def test_surrogate():
 
  # Define emulator and pre-train on global grid
  discrepancy = Discrepancy(model_name='discrepancy_test',
- lf_model=model.solve_lf,
+ lf_model=model.solve_t,
  input_size=2,
  output_size=1,
  var_grid_in=var_data_in,
- var_grid_out=var_data_out)
-
+ var_grid_out=var_data_out,
+ dnn_arch=[64,64],
+ dnn_activation='relu',
+ dnn_dropout=[0.2,0.5], 
+ activation='silu')
+
  # Create a batch of samples for the calibration parameters
  batch_x = model.defParams
 
  # Update the discrepancy model
- discrepancy.update(batch_x, max_iters=1000, lr=0.001, lr_exp=0.9999, record_interval=100)
+ discrepancy.update(batch_x, max_iters=10000, lr=0.001, lr_exp=0.9999, record_interval=100)
 
  # Plot discrepancy
  fig = plt.figure()
  ax = fig.add_subplot(projection='3d')
- ax.scatter(model.var_in[:,0].detach().numpy(), model.var_in[:,1].detach().numpy(), discrepancy.forward(model.var_in).detach().numpy(), marker='o')
+ for loopA in range(50):
+ ax.scatter(model.var_in[:,0].detach().numpy(), model.var_in[:,1].detach().numpy(), model.solve_t(batch_x)+discrepancy.forward(model.var_in).detach().numpy(), color='blue', marker='o')
+ for loopA in range(var_data_out.size(1)):
+ ax.scatter(model.var_in[:,0].detach().numpy(), model.var_in[:,1].detach().numpy(), var_data_out[:,loopA].detach().numpy(), color='red', marker='D', s=5)
  ax.set_xlabel('Temperature')
  ax.set_ylabel('Pressure')
  ax.set_zlabel('Coverage')
@@ -203,5 +266,5 @@ def test_surrogate():
 # TEST SURROGATE
 if __name__ == '__main__':
 
- test_surrogate()
+ test_discrepancy()
 

linfa/eval_model_from_chkpt.py

@@ -0,0 +1,84 @@
+import os,torch
+from linfa.discrepancy import Discrepancy
+from linfa.maf import MAF, RealNVP
+from run_experiment import load_exp_from_file
+
+def eval_discrepancy(file_path,test_data):
+
+ # Read in data
+ exp_name = os.path.basename(file_path)
+ dir_name = os.path.dirname(file_path)
+
+ # Create new discrepancy
+ dicr = Discrepancy(model_name = exp_name, 
+ model_folder = dir_name,
+ lf_model = None,
+ input_size = None,
+ output_size = None,
+ var_grid_in = None,
+ var_grid_out = None)
+ dicr.surrogate_load()
+
+ # Evaluate discrepancy over test grid
+ return dicr.forward(test_data) 
+
+def eval_model(exp_chkpt_file,nf_chkpt_file,discr_chkpt_file,num_calib_samples,test_data):
+
+ # Load experiment from file
+ exp = load_exp_from_file(exp_chkpt_file)
+
+ # Create NF model from experiment
+ if exp.flow_type == 'maf':
+ nf = MAF(exp.n_blocks, exp.input_size, exp.hidden_size, exp.n_hidden, None,
+ exp.activation_fn, exp.input_order, batch_norm=exp.batch_norm_order)
+ elif exp.flow_type == 'realnvp': # Under construction
+ nf = RealNVP(exp.n_blocks, exp.input_size, exp.hidden_size, exp.n_hidden, None,
+ batch_norm=exp.batch_norm_order)
+
+ # Read state dictionary 
+ nf.state_dict(torch.load(nf_chkpt_file)) 
+
+ # Sample calibration parameter realizations
+ x00 = nf.base_dist.sample([num_calib_samples])
+ xkk, _ = nf(x00)
+
+ # Evaluate discrepancy at tp data "test_data"
+ res_discr = eval_discrepancy(discr_chkpt_file,test_data)
+
+ # Solve models 
+ # Need to change this to be evaluated at arbitraty temperatures and pressures.
+ if(exp.transform is None):
+ res_lf = exp.model.solve_t(xkk)
+ else:
+ res_lf = exp.model.solve_t(exp.transform.forward(xkk))
+
+ # CURRENTLY NO NOISE IS ADDED, NEED TO BE IMPLEMENTED IF APPROPRIATE!!!
+
+ # return 
+ return res_lf + res_discr
+
+# MAIN CODE
+if __name__ == "__main__":
+
+ # Assign files
+ exp_chkpt_file = './tests/results/test_lf_with_disc_hf_data_TP1/experiment.pt'
+ nf_chkpt_file = './tests/results/test_lf_with_disc_hf_data_TP1/test_lf_with_disc_hf_data_TP1_3000.nf'
+ discr_chkpt_file = './tests/results/test_lf_with_disc_hf_data_TP1/test_lf_with_disc_hf_data_TP1'
+ # 
+ num_calib_samples = 100
+
+ # Set the grid for 
+ min_dim_1 = 400.0
+ max_dim_1 = 500.0
+ min_dim_2 = 2.0
+ max_dim_2 = 3.0
+ num_1d_grid_points = 5
+ #
+ test_grid_1 = torch.linspace(min_dim_1, max_dim_1, num_1d_grid_points)
+ test_grid_2 = torch.linspace(min_dim_2, max_dim_2, num_1d_grid_points)
+ grid_t, grid_p = torch.meshgrid(test_grid_1, test_grid_2, indexing='ij')
+ test_data = torch.cat((grid_t.reshape(-1,1), grid_p.reshape(-1,1)),1)
+
+ res = eval_model(exp_chkpt_file,nf_chkpt_file,discr_chkpt_file,num_calib_samples,test_data)
+
+ print(res.size())

linfa/mlp.py

@@ -5,7 +5,7 @@
 class FNN(nn.Module):
  """Fully Connected Neural Network"""
 
- def __init__(self, input_size, output_size, arch=None, activation='relu', device='cpu',init_zero=False):
+ def __init__(self, input_size, output_size, arch=None, activation='relu', device='cpu', init_zero=False, dropout=None):
  """
  Args:
  input_size (int): Input size for FNN
@@ -32,6 +32,12 @@ def __init__(self, input_size, output_size, arch=None, activation='relu', device
  self.fc.append(nn.Linear(self.neuron_size[layer-1], self.neuron_size[layer]).to(device))
  # Assign output layer
  self.fc.append(nn.Linear(self.neuron_size[len(self.neuron_size)-1], output_size).to(device))
+ # Add dropout layers
+ self.dropout = dropout
+ if(self.dropout is not None):
+ self.dropout_visible = nn.Dropout(p=self.dropout[0])
+ self.dropout_hidden = nn.Dropout(p=self.dropout[1])
+
 
  def forward(self, x):
  """
@@ -41,17 +47,37 @@ def forward(self, x):
  Returns:
  torch.Tensor. Assumed to be a batch.
  """
+ if(self.dropout is not None):
+
+ x = self.dropout_visible(x)
+
  for loopA in range(len(self.fc)-1):
- if(self.activation == 'relu'):
- x = F.relu(self.fc[loopA](x))
+
+ if(self.activation == 'relu'): 
+
+ if(self.dropout is not None):
+ x = self.dropout_hidden(F.relu(self.fc[loopA](x)))
+ else:
+ x = F.relu(self.fc[loopA](x))
+
  elif(self.activation == 'silu'):
- x = F.silu(self.fc[loopA](x))
+
+ if(self.dropout is not None):
+ x = self.dropout_hidden(F.silu(self.fc[loopA](x)))
+ else:
+ x = F.silu(self.fc[loopA](x))
+
  elif(self.activation == 'tanh'):
- x = F.tanh(self.fc[loopA](x))
+
+ if(self.dropout is not None):
+ x = self.dropout_hidden(F.tanh(self.fc[loopA](x)))
+ else:
+ x = F.tanh(self.fc[loopA](x))
+
  else:
  print('Invalid activation string.')
  exit(-1)
 
  # Last layer with linear activation
- x = self.fc[len(self.fc)-1](x)
+ x = self.fc[-1](x)
  return x