ML multicolor simulation

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):
@@ -24,14 +25,16 @@ def __init__(self, model_name, model_hparams, optimizer_name,
         self.save_hyperparameters()
         # Create model
         self.model = self.create_model(model_name, model_hparams)
-        # print(summary(self.model, input_size=(misc_hparams['batch_size'], 1, 256)))
+        print(summary(self.model, input_size=(misc_hparams['batch_size'],
+                                              model_hparams['in_channels'],
+                                              256)))
         # Create loss module
         self.loss_function = nn.MSELoss()
-        
+
         self.step = misc_hparams["step"]
 
-        torch.set_float32_matmul_precision('medium')
-        
+        torch.set_float32_matmul_precision('high')
+
     def create_model(self, model_name, model_hparams):
         if model_name in model_dict:
             return model_dict[model_name](**model_hparams)
@@ -40,7 +43,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":
@@ -54,26 +57,27 @@ def configure_optimizers(self):
         else:
             assert False, f'Unknown optimizer: "{self.hparams.optimizer_name}"'
 
-        # We will reduce the learning rate by 0.1 after 100 and 150 epochs
+        # We will reduce the learning rate by factor of gamma after 100 and 150
+        # epochs
         scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
                                                    milestones=[50, 100, 150], #changed from [100, 150, 200]
-                                                   gamma=0.1)
+                                                   gamma=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
         x, y = batch
         preds = self.model(x)
         loss = self.loss_function(preds, y)
-        acc = (preds == y).float().mean()
-        self.log("train_acc", acc, on_step=False, on_epoch=True)
+        # acc = (preds == y).float().mean()
+        # self.log("train_acc", acc, on_step=False, on_epoch=True)
         self.log("train_loss", loss, on_epoch=True, prog_bar=True)
         return loss
 
@@ -85,41 +89,44 @@ def validation_step(self, batch, batch_idx):
         # print(f"y size {y.size()}")
         preds = self.model(x)
         loss = self.loss_function(preds, y)
-        acc = (preds == y).float().mean()
-        self.log("val_acc", acc, on_step=False, on_epoch=True)
-        self.log("val_loss", loss, prog_bar=False)
+        # acc = (preds == y).float().mean()
+        # self.log("val_acc", acc, on_step=False, on_epoch=True)
+        self.log("val_loss", loss, on_epoch=True, prog_bar=True)
         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)
+        # acc = (preds == y_tot).float().mean()
+        # self.log("test_acc", acc, on_step=False, on_epoch=True)
+        self.log("test_loss", loss, on_epoch=True, prog_bar=True)
         return loss
 
-    def predict_step(self, batch, batch_idx, test_mode=False, dataloader_idx=0):
-        x_tot, y_tot = batch 
+    def predict_step(self, batch, batch_idx, test_mode=True, dataloader_idx=0):
+        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
+            print(x_tot.shape, y_tot.shape)
+            num_groups = y_tot.shape[1]
+            group_size = x_tot.shape[1] - (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, :]
+                x = x.reshape(x.shape[0], x.shape[2], x.shape[1])
+                # [batch_size, group_size, num_channels]
                 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
+        return y_hat, y_tot, x_tot

interferometers.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+
+import numpy as np
+import matplotlib.pyplot as plt
+import util
+from tqdm import tqdm
+import h5py
+
+
+class MichelsonInterferometer:
+    def __init__(self, wavelength, displacement_amplitude, phase):
+        self.wavelength = wavelength # in microns
+        self.displacement_amplitude = displacement_amplitude # in microns
+        self.phase = phase # in radians, stands for random position offset
+        self.displacement = None # in microns
+        self.velocity = None # in microns/s
+        self.time = None # in seconds
+
+    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 set_displacement(self, displacement, time):
+        # TODO: add checks and tests here, that the displacement is in right
+        #  range, sampling, etc.
+        self.displacement = displacement
+        self.time = time
+
+    def get_interferometer_output(self, start_frequency, end_frequency,
+                                  measurement_noise_level, length,
+                                  sample_rate, displacement, time):
+        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]
+
+        if displacement is None:
+            self.time, self.displacement = self.get_displacement(start_frequency,
+                                                 end_frequency, length, sample_rate)
+        else:
+            self.set_displacement(displacement, time)
+
+        interference = np.cos(2 * np.pi / self.wavelength * self.displacement
+                              + self.phase)
+
+        signal = E0**2 + ER**2 + 2 * E0 * ER * interference
+
+        return signal
+
+    def get_buffer(self, start_frequency=0, end_frequency=1e3,
+                   displacement=None, time=None):
+        self.signal = self.get_interferometer_output(start_frequency,
+                                                     end_frequency, 0.3,
+                                                     8192, 1e6, displacement,
+                                                     time)
+
+        self.velocity = np.diff(self.displacement)
+        self.velocity = np.insert(self.velocity, 0, self.velocity[0])
+        self.velocity /= (self.time[1] - self.time[0])
+
+        # Remove DC offset
+        self.signal = self.signal - np.mean(self.signal)
+
+        return self.time, self.signal, self.displacement, self.velocity
+
+    def plot_buffer(self):
+        time, signal, displacement, velocity = self.get_buffer(0, 1e3)
+
+        # time = time[0:256]
+        # signal = signal[0:256]
+        # displacement = displacement[0:256]
+        # velocity = velocity[0:256]
+
+        fig, ax1 = plt.subplots(figsize=(18, 6))
+        ax1.plot(time, signal, color='b')
+        ax1.set_xlabel('Time (s)')
+        ax1.set_ylabel('Signal (V)', color='b')
+        ax1.tick_params('y', colors='b')
+
+        ax2 = ax1.twinx()
+        ax2.plot(time, displacement, color='r')
+        ax2.plot(time, velocity, color='g')
+        ax2.set_ylabel('Displacement (microns)', color='r')
+        ax2.tick_params('y', colors='r')
+
+        plt.tight_layout()
+        plt.show()
+
+
+def write_pretraining_data(num_shots, num_channels, file_path):
+    if num_channels == 1:
+        interferometer = MichelsonInterferometer(0.5, 5, np.pi / 4)
+        for _ in tqdm(range(num_shots)):
+            _, signal, _, velocity = interferometer.get_buffer()
+            signal = np.expand_dims(signal, axis=-1)
+            velocity = np.expand_dims(velocity, axis=-1)
+            # Want to end up with these shapes in h5 file:
+            # signal: (num_shots, buffer_size, 1)
+            # velocity: (num_shots, buffer_size, 1)
+            entries = {"signal": signal, "velocity": velocity}
+            util.write_data(file_path, entries)
+    elif num_channels == 2:
+        interferometer1 = MichelsonInterferometer(1.064, 5, np.pi / 4)
+        interferometer2 = MichelsonInterferometer(0.532, 5, 3 * np.pi / 4)
+        for _ in tqdm(range(num_shots)):
+            time, signal1, displacement, velocity = interferometer1.get_buffer()
+            _, signal2, _, _ = interferometer2.get_buffer(displacement=displacement, time=time)
+            signal = np.stack((signal1, signal2), axis=-1)
+            velocity = np.expand_dims(velocity, axis=-1)
+            # Want to end up with these shapes in h5 file:
+            # signal: (num_shots, buffer_size, num_channels)
+            # velocity: (num_shots, buffer_size, 1)
+            entries = {"signal": signal, "velocity": velocity}
+            # print(signal.shape)
+            # print(velocity.shape)
+            util.write_data(file_path, entries)
+
+def plot_pretraining_data(file_path):
+    with h5py.File(file_path, 'r') as f:
+        signal = np.array(f['signal'])
+        velocity = np.array(f['velocity'])
+        print(signal.shape)
+        print(velocity.shape)
+        for shot_idx in range(signal.shape[0]):
+            fig, ax1 = plt.subplots(figsize=(18, 6))
+            ax1.plot(signal[shot_idx, :, 0], color='b', label='Signal 1')
+            ax1.plot(signal[shot_idx, :, 1], color='g', label='Signal 2')
+            ax1.set_xlabel('Time (s)')
+            ax1.set_ylabel('Signal (V)', color='b')
+            ax1.tick_params('y', colors='b')
+
+            ax2 = ax1.twinx()
+            ax2.plot(velocity[shot_idx, :, 0], color='r')
+            ax2.set_ylabel('Velocity (microns/second)', color='r')
+
+            plt.tight_layout()
+            plt.show()
+
+
+if __name__ == '__main__':
+    np.random.seed(0x5EED + 4)
+    write_pretraining_data(10, 2,
+                           "/Users/nolanpeard/Desktop/SMI_sim/valid_double.h5")
+    # plot_pretraining_data("/Users/nolanpeard/Desktop/test.h5")
+

main.py

@@ -13,15 +13,25 @@
     if len(sys.argv) == 1:
         """Run functions in this scratch area. 
         """
-        valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\valid_max10kHz_30to1kHz_2kshots_dec=256_randampl.h5py'
-        train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\training_max10kHz_30to1kHz_10kshots_dec=256_randampl.h5py'
-        test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_max10kHz_30to1kHz_2kshots_dec=256_randampl.h5py'
+        # valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\valid_max10kHz_30to1kHz_2kshots_dec=256_randampl.h5py'
+        # train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\training_max10kHz_30to1kHz_10kshots_dec=256_randampl.h5py'
+        # test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_max10kHz_30to1kHz_2kshots_dec=256_randampl.h5py'
+
+        train_file = "/Users/nolanpeard/Desktop/SMI_sim/train_double.h5"
+        valid_file = "/Users/nolanpeard/Desktop/SMI_sim/valid_double.h5"
+        test_file = "/Users/nolanpeard/Desktop/SMI_sim/test_double.h5"
 
         print('begin main', datetime.datetime.now())
-        step_list = [256, 128, 64, 32] # step sizes for rolling input
-        for step in step_list:
-            runner = train.TrainingRunner(train_file, valid_file, test_file, step)
-            runner.scan_hyperparams()
+        # step_list = [256]#, 128, 64, 32] # step sizes for rolling input
+        # for step in step_list:
+        #     runner = train.TrainingRunner(train_file, valid_file, test_file,
+        #                                   step)
+        #     runner.scan_hyperparams()
+
+        runner = train.TrainingRunner(train_file, valid_file, test_file,
+                                      step=256)
+        #runner.plot_predictions(model_name="CNN", model_id="tdwhpu2l")
+        runner.plot_predictions(model_name="CNN", model_id="e8vpuie1")
 
     else:
         print("Error: Unsupported number of command-line arguments")

models.py

@@ -5,37 +5,49 @@
 
 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'):
+    def __init__(self, input_size, output_size, in_channels=1, activation='LeakyReLU'):
         super(CNN, self).__init__()
-        self.ch_in = ch_in
+        self.in_channels = in_channels
+        print("self.in_channels = ", self.in_channels)
         self.conv_layers = nn.Sequential(
-            nn.Conv1d(ch_in, 16, kernel_size=7),  # Lout = 250, given L = 256
-            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], 
-            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], 
-            nn.MaxPool1d(2),  # Lout = 26, given L = 53
-            nn.Dropout(0.1), 
-            nn.Conv1d(64, 64, kernel_size=7),  # Lout = 20, given L = 26
-            act_fn_by_name[activation], 
-            nn.MaxPool1d(2)  # Lout = 10, given L = 20
+            nn.Conv1d(in_channels, 16, kernel_size=7),
+            # Lout = 250, given L = 256
+            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],
+            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],
+            nn.MaxPool1d(2),
+            # Lout = 26, given L = 53
+            nn.Dropout(0.1),
+            nn.Conv1d(64, 64, kernel_size=7),
+            # Lout = 20, given L = 26
+            act_fn_by_name[activation],
+            nn.MaxPool1d(2)
+            # Lout = 10, given L = 20
         )
         self.fc_layers = nn.Sequential(
+            # length of input = 64 filters * length of 10 left
             nn.Linear(640, 16),
             act_fn_by_name[activation],
             nn.Linear(16, output_size)
         )
-        
+
     def forward(self, x):
+        # print("x.shape", x.shape)
         out = self.conv_layers(x)
         # print(f"post conv out size: {out.size()}")  # [128, 64, 10]
-        out = out.view(out.size(0), self.ch_in, -1)
-        # print(f"post conv out reshaped size: {out.size()}")  # confirmed [128, 1, 640]
+        out = out.view(out.size(0), 1, -1)
+        # print(f"post conv out reshaped size: {out.size()}")  # confirmed [
+        # 128, 1, 640]
         out = self.fc_layers(out)  # expect out [128, 1, 1]
         # print(f"post fc out size: {out.size()}") # confirmed: [128, 1, 1]
         return out
-