wrote class for simulating classic interferometers to perform pre-tra…

…ining, testing of new concepts

decoder.py

@@ -8,6 +8,7 @@
 
 model_dict = {"CNN": CNN}
 
+
 class VelocityDecoder(L.LightningModule):
     def __init__(self, model_name, model_hparams, optimizer_name,
                  optimizer_hparams, misc_hparams):
@@ -27,11 +28,11 @@ def __init__(self, model_name, model_hparams, optimizer_name,
         print(summary(self.model, input_size=(misc_hparams['batch_size'], 1, 256)))
         # Create loss module
         self.loss_function = nn.MSELoss()
-        
+
         self.step = misc_hparams["step"]
 
         torch.set_float32_matmul_precision('medium')
-        
+
     def create_model(self, model_name, model_hparams):
         if model_name in model_dict:
             return model_dict[model_name](**model_hparams)
@@ -40,7 +41,7 @@ def create_model(self, model_name, model_hparams):
 
     def forward(self, x):
         return self.model(x)
-    
+
     def configure_optimizers(self):
         # We will support Adam or SGD as optimizers.
         if self.hparams.optimizer_name == "Adam":
@@ -59,13 +60,13 @@ def configure_optimizers(self):
                                                    milestones=[100, 150, 200],
                                                    gamma=0.1)
         return [optimizer], [scheduler]
-    
+
     def get_loss_function(self, loss_hparams):
         if loss_hparams["loss_name"] == "mse":
             self.loss_function = nn.MSELoss()
         else:
             assert False, f'Unknown loss: "{loss_hparams["loss_name"]}"'
-            
+
     def training_step(self, batch, batch_idx):
         # training_step defines the train loop.
         # it is independent of forward
@@ -89,37 +90,37 @@ def validation_step(self, batch, batch_idx):
         self.log("val_acc", acc, on_step=False, on_epoch=True)
         self.log("val_loss", loss, prog_bar=False)
         return loss
-    
+
     def test_step(self, batch, batch_idx):
         x_tot, y_tot = batch  # [batch_size, 1, buffer_size], [batch_size, 1, num_groups]
         num_groups = y_tot.shape[2]
         group_size = x_tot.shape[2] - (num_groups - 1) * self.step
         preds = []
         for i in range(num_groups):
-            start_idx = i*self.step
-            x = x_tot[:, :, start_idx:start_idx+group_size]
+            start_idx = i * self.step
+            x = x_tot[:, :, start_idx:start_idx + group_size]
             preds.append(self.model(x).flatten())
         preds = torch.unsqueeze(torch.transpose(torch.stack(preds), dim0=0, dim1=1),
-                              dim=1)  # [batch_size, 1, num_groups]
+                                dim=1)  # [batch_size, 1, num_groups]
         loss = self.loss_function(preds, y_tot)
         acc = (preds == y_tot).float().mean()
         self.log("test_acc", acc, on_step=False, on_epoch=True)
         self.log("test_loss", loss, prog_bar=True)
         return loss
 
     def predict_step(self, batch, batch_idx, test_mode=False, dataloader_idx=0):
-        x_tot, y_tot = batch 
+        x_tot, y_tot = batch
         if test_mode:
             # x_tot, y_tot: [batch_size, 1, buffer_size], [batch_size, 1, num_groups]
             num_groups = y_tot.shape[2]
             group_size = x_tot.shape[2] - (num_groups - 1) * self.step
             y_hat = []
             for i in range(num_groups):
-                start_idx = i*self.step
-                x = x_tot[:, :, start_idx:start_idx+group_size]  # [batch_size, 1, group_size]
+                start_idx = i * self.step
+                x = x_tot[:, :, start_idx:start_idx + group_size]  # [batch_size, 1, group_size]
                 y_hat.append(self.model(x).flatten())
             y_hat = torch.unsqueeze(torch.transpose(torch.stack(y_hat), dim0=0, dim1=1),
-                                  dim=1)  # [batch_size, 1, num_groups]
+                                    dim=1)  # [batch_size, 1, num_groups]
         else:
             y_hat = self.model(x_tot)
         return y_hat, y_tot

interferometers.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+
+import numpy as np
+import matplotlib.pyplot as plt
+import util
+
+class MichelsonInterferometer:
+    def __init__(self, wavelength, displacement_amplitude, phase):
+        self.wavelength = wavelength # in microns
+        self.displacement_amplitude = displacement_amplitude # in microns
+        self.phase = phase
+        self.displacement = None
+        self.velocity = None
+
+    def get_displacement(self, start_frequency, end_frequency, length, sample_rate):
+        # Get a random displacement in time, resets each time it is called
+        time, displacement = util.bounded_frequency_waveform(start_frequency,
+                                                end_frequency, length, sample_rate)
+
+        # scale max displacement to the desired amplitude
+        displacement = displacement/np.max(np.abs(displacement)) * self.displacement_amplitude
+
+        return time, displacement
+
+    def interferometer_output(self, start_frequency, end_frequency,
+                              measurement_noise_level, length, sample_rate):
+        E0 = 1 + measurement_noise_level * util.bounded_frequency_waveform(1e3,
+                                                                          1e6,
+                                                                  length, sample_rate)[1]
+        ER = 0.1 + measurement_noise_level * util.bounded_frequency_waveform(1e3,
+                                                                          1e6,
+                                                                    length, sample_rate)[1]
+
+        self.time, self.displacement = self.get_displacement(start_frequency,
+                                                 end_frequency, length, sample_rate)
+
+        interference = np.cos(2 * np.pi / self.wavelength * self.displacement)
+
+        signal = E0**2 + ER**2 + 2 * E0 * ER * interference
+
+        return signal
+
+    def get_pretraining_data(self, start_frequency, end_frequency):
+        self.signal = self.interferometer_output(start_frequency,
+                                                 end_frequency, 0.1, 8192, 1e6)
+
+        self.velocity = np.diff(self.displacement)
+        self.velocity = np.insert(self.velocity, 0, self.velocity[0])
+        self.velocity /= (self.time[1] - self.time[0])
+
+        return self.time, self.signal, self.displacement, self.velocity
+
+    def plot_pretraining_data(self):
+        time, signal, displacement, velocity = self.get_pretraining_data(0, 1e3)
+
+        plt.plot(time, signal)
+        plt.plot(time, displacement)
+        plt.plot(time, velocity)
+        plt.tight_layout()
+        plt.show()
+
+if __name__ == '__main__':
+    interferometer = MichelsonInterferometer(0.5, 5, 0)
+    interferometer.plot_pretraining_data()
+

models.py

@@ -5,31 +5,32 @@
 
 act_fn_by_name = {'LeakyReLU': nn.LeakyReLU(), 'ReLU': nn.ReLU()}
 
+
 class CNN(nn.Module):
     def __init__(self, input_size, output_size, ch_in=1, activation='LeakyReLU'):
         super(CNN, self).__init__()
         self.ch_in = ch_in
         self.conv_layers = nn.Sequential(
             nn.Conv1d(ch_in, 16, kernel_size=7),  # Lout = 250, given L = 256
-            act_fn_by_name[activation], 
+            act_fn_by_name[activation],
             nn.MaxPool1d(2),  # Lout = 125, given L = 250
             nn.Conv1d(16, 32, kernel_size=7),  # Lout = 119, given L = 125
-            act_fn_by_name[activation], 
+            act_fn_by_name[activation],
             nn.MaxPool1d(2),  # Lout = 59, given L = 119
             nn.Conv1d(32, 64, kernel_size=7),  # Lout = 53, given L = 59
-            act_fn_by_name[activation], 
+            act_fn_by_name[activation],
             nn.MaxPool1d(2),  # Lout = 26, given L = 53
-            nn.Dropout(0.1), 
+            nn.Dropout(0.1),
             nn.Conv1d(64, 64, kernel_size=7),  # Lout = 20, given L = 26
-            act_fn_by_name[activation], 
+            act_fn_by_name[activation],
             nn.MaxPool1d(2)  # Lout = 10, given L = 20
         )
         self.fc_layers = nn.Sequential(
             nn.Linear(640, 16),
             nn.ReLU(),
             nn.Linear(16, output_size)
         )
-        
+
     def forward(self, x):
         out = self.conv_layers(x)
         # print(f"post conv out size: {out.size()}")  # [128, 64, 10]
@@ -38,4 +39,3 @@ def forward(self, x):
         out = self.fc_layers(out)  # expect out [128, 1, 1]
         # print(f"post fc out size: {out.size()}") # confirmed: [128, 1, 1]
         return out
-

train.py

@@ -1,6 +1,8 @@
 #!/usr/bin/env python3
 
-import os, sys, glob
+import os
+import sys
+import glob
 import torch
 from torch.utils.data import Dataset, DataLoader
 import h5py
@@ -17,6 +19,8 @@
 import time
 
 # Define a custom Dataset class
+
+
 class VelocityDataset(Dataset):
     def __init__(self, test_mode, h5_file, step, group_size=256):
         self.h5_file = h5_file
@@ -32,14 +36,14 @@ def __init__(self, test_mode, h5_file, step, group_size=256):
         self.opened_flag = False
         self.test_mode = test_mode
 
-    def open_hdf5(self, rolling=True, step=256, group_size=256, ch_in = 1):
+    def open_hdf5(self, rolling=True, step=256, group_size=256, ch_in=1):
         """Set up inputs and targets. For each shot, buffer is split into groups of sequences.
-        Inputs include grouped photodiode trace of 'group_size', spaced interval 'step' apart for each buffer. 
+        Inputs include grouped photodiode trace of 'group_size', spaced interval 'step' apart for each buffer.
         Targets include average velocity of each group.
         Input shape is [num_shots * num_groups, ch_in, group_size]. Target shape is [num_shots * num_groups, ch_in, 1],
         where num_groups = (buffer_len - group_size)/step + 1, given that buffer_len - group_size is a multiple of step.
         Input and target shape made to fit (N, C_in, L_in) in PyTorch Conv1d doc.
-        If the given 'group_size' and 'step' do not satisfy the above requirement, 
+        If the given 'group_size' and 'step' do not satisfy the above requirement,
         the data will not be cleanly grouped.
 
         Args:
@@ -53,9 +57,9 @@ def open_hdf5(self, rolling=True, step=256, group_size=256, ch_in = 1):
         self.file = h5py.File(self.h5_file, 'r')
         pds = torch.Tensor(np.array(self.file['PD (V)']))  # [num_shots, buffer_size]
         vels = torch.Tensor(np.array(self.file['Speaker (Microns/s)']))  # [num_shots, buffer_size]
-        
+
         if rolling:
-            # ROLLING INPUT INDICES 
+            # ROLLING INPUT INDICES
             num_groups = (pds.shape[1] - group_size) // step + 1
             start_idxs = torch.arange(num_groups) * step  # starting indices for each group
             idxs = torch.arange(group_size)[:, None] + start_idxs
@@ -73,10 +77,11 @@ def open_hdf5(self, rolling=True, step=256, group_size=256, ch_in = 1):
             if self.test_mode:
                 assert False, 'test_mode not implemented for step input. use rolling step=256'
             else:
-                self.inputs = torch.cat(torch.split(pds, group_size, dim=1), dim=0)  # [num_shots * num_groups, group_size]
+                # [num_shots * num_groups, group_size]
+                self.inputs = torch.cat(torch.split(pds, group_size, dim=1), dim=0)
                 grouped_vels = torch.cat(torch.split(vels, group_size, dim=1), dim=0)
                 self.targets = torch.unsqueeze(torch.mean(grouped_vels, dim=1), dim=1)  # [num_shots * num_groups, 1]
-                
+
         if ch_in == 1:
             self.inputs = torch.unsqueeze(self.inputs, dim=1)
             self.targets = torch.unsqueeze(self.targets, dim=1)
@@ -93,7 +98,7 @@ def __len__(self):
 
     def __getitem__(self, idx):
         # print("getitem entered")
-        if not self.opened_flag: #not hasattr(self, 'h5_file'):
+        if not self.opened_flag:  # not hasattr(self, 'h5_file'):
             self.open_hdf5(step=self.step)
             self.opened_flag = True
             # print("open_hdf5 finished")
@@ -106,6 +111,7 @@ def __getitem__(self, idx):
     #         print("open_hdf5 in getitems")
     #     return FloatTensor(self.inputs[indices]), FloatTensor(self.targets[indices])
 
+
 class TrainingRunner:
     def __init__(self, training_h5, validation_h5, testing_h5, step=256,
                  velocity_only=True):
@@ -132,7 +138,7 @@ def __init__(self, training_h5, validation_h5, testing_h5, step=256,
         # directories
         self.checkpoint_dir = "./checkpoints"
         print('TrainingRunner initialized', datetime.datetime.now())
-        
+
     def get_custom_dataloader(self, test_mode, h5_file, batch_size=128, shuffle=True,
                               velocity_only=True):
         # if velocity_only:
@@ -144,11 +150,12 @@ def get_custom_dataloader(self, test_mode, h5_file, batch_size=128, shuffle=True
                                 pin_memory=True)
         print("dataloader initialized")
         return dataloader
-    
+
     def set_dataloaders(self, batch_size=128):
         self.batch_size = batch_size
         self.train_loader = self.get_custom_dataloader(False, self.training_h5, batch_size=self.batch_size)
-        self.valid_loader = self.get_custom_dataloader(False, self.validation_h5, batch_size=self.batch_size, shuffle=False)
+        self.valid_loader = self.get_custom_dataloader(
+            False, self.validation_h5, batch_size=self.batch_size, shuffle=False)
         self.test_loader = self.get_custom_dataloader(True, self.testing_h5, batch_size=self.batch_size, shuffle=False)
 
     def train_model(self, model_name, save_name=None, **kwargs):
@@ -208,30 +215,29 @@ def train_model(self, model_name, save_name=None, **kwargs):
         logger.experiment.finish()
 
         return model, result
-    
+
     def scan_hyperparams(self):
-        lr_list = [1e-3, 1e-4] # [1e-3, 1e-4, 1e-5]
-        act_list = ['LeakyReLU'] #, 'ReLU']
-        optim_list = ['Adam'] #, 'SGD']
-        for lr, activation, optim in product(lr_list, act_list, optim_list): #, 1e-2, 3e-2]:
+        lr_list = [1e-3, 1e-4]  # [1e-3, 1e-4, 1e-5]
+        act_list = ['LeakyReLU']  # , 'ReLU']
+        optim_list = ['Adam']  # , 'SGD']
+        for lr, activation, optim in product(lr_list, act_list, optim_list):  # , 1e-2, 3e-2]:
             model_config = {"input_size": self.input_size,
                             "output_size": self.output_size,
                             "activation": activation}
             optimizer_config = {"lr": lr}
-                                #"momentum": 0.9,}
+            # "momentum": 0.9,}
             misc_config = {"batch_size": self.batch_size, "step": self.step}
 
             self.train_model(model_name="CNN",
-                            model_hparams=model_config,
-                            optimizer_name=optim,
-                            optimizer_hparams=optimizer_config,
-                            misc_hparams=misc_config)
-
+                             model_hparams=model_config,
+                             optimizer_name=optim,
+                             optimizer_hparams=optimizer_config,
+                             misc_hparams=misc_config)
 
     def load_model(self, model_tag, model_name='CNN'):
         # Check whether pretrained model exists. If yes, load it and skip training
         pretrained_filename = os.path.join(self.checkpoint_dir, model_name, "SMI", model_tag,
-                                            "checkpoints", "*" + ".ckpt")
+                                           "checkpoints", "*" + ".ckpt")
         print(pretrained_filename)
         if os.path.isfile(glob.glob(pretrained_filename)[0]):
             pretrained_filename = glob.glob(pretrained_filename)[0]
@@ -248,10 +254,10 @@ def load_model(self, model_tag, model_name='CNN'):
 
             # Test best model on validation and test set
             val_result = trainer.test(model, dataloaders=self.valid_loader,
-                                        verbose=False)
+                                      verbose=False)
             test_result = trainer.test(model, dataloaders=self.test_loader,
-                                        verbose=False)
+                                       verbose=False)
             result = {"test": test_result[0]["test_acc"],
-                        "val": val_result[0]["test_acc"]}
+                      "val": val_result[0]["test_acc"]}
 
-            return model, result
+            return model, result