diff --git a/decoder.py b/decoder.py index 91d39da..982a9c5 100644 --- a/decoder.py +++ b/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,18 +90,18 @@ 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) @@ -108,18 +109,18 @@ def test_step(self, batch, batch_idx): 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 diff --git a/interferometers.py b/interferometers.py new file mode 100644 index 0000000..a5aad26 --- /dev/null +++ b/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() + \ No newline at end of file diff --git a/models.py b/models.py index a64a044..1bfcc4b 100644 --- a/models.py +++ b/models.py @@ -5,23 +5,24 @@ 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( @@ -29,7 +30,7 @@ def __init__(self, input_size, output_size, ch_in=1, activation='LeakyReLU'): 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 - \ No newline at end of file diff --git a/train.py b/train.py index 79d904b..52230a1 100644 --- a/train.py +++ b/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 \ No newline at end of file + return model, result diff --git a/util.py b/util.py index fdd0613..8fcbed3 100644 --- a/util.py +++ b/util.py @@ -2,46 +2,48 @@ from scipy.fftpack import fft, ifft, fftfreq import h5py -def bounded_frequency_waveform(start_frequency, end_frequency, length, sample_rate, use_freq, max_freq): - """Generates a random waveform within the given frequency range of a given length. - + +def bounded_frequency_waveform(start_frequency, end_frequency, length, sample_rate): + """Generates a random waveform within the given frequency range of a given length. + Args: start_frequency (float): the lower bound of the valid frequency range end_frequency (float): the upper bound of the valid frequency range - length (int): the number of values to generate + length (int): the number of values to generate sample_rate (float): the rate at which to sample values - + Returns: [1darr, 1darr]: the array of time points and amplitude points in time domain """ # Create an evenly spaced time array t = np.linspace(0, 1.0, length, False) # 1 second # Generate a random frequency spectrum between the start and end frequencies - if use_freq: - freq = np.linspace(0, max_freq, length//2, False) # replaced sample_rate/2 with end_frequency - else: - freq = np.linspace(0, sample_rate/2, length//2, False) + # if use_freq: + # freq = np.linspace(0, max_freq, length // 2, False) # replaced sample_rate/2 with end_frequency + + freq = np.linspace(0, sample_rate / 2, length // 2, False) spectrum = np.random.uniform(0, 1, len(freq)) spectrum = np.where((freq >= start_frequency) & (freq <= end_frequency), spectrum, 0) - c = np.random.rayleigh(np.sqrt(spectrum*(freq[1]-freq[0]))) + c = np.random.rayleigh(np.sqrt(spectrum * (freq[1] - freq[0]))) # See Phys. Rev. A 107, 042611 (2023) ref 28 for why we use the Rayleigh distribution here # Unless we use this distribution, the random noise will not be Gaussian distributed phase = np.random.uniform(-np.pi, np.pi, len(freq)) # Use the inverse Fourier transform to convert the frequency domain signal back to the time domain # Also include a zero phase component - spectrum = np.hstack([c*spectrum*np.exp(1j*phase), np.zeros_like(spectrum)]) + spectrum = np.hstack([c * spectrum * np.exp(1j * phase), np.zeros_like(spectrum)]) y = np.real(ifft(spectrum)) y = np.fft.fftshift(y) return t, y + def linear_convert(data, new_min=-1, new_max=1): - """Linearly scales data to a new range. Default is [-1, 1]. - + """Linearly scales data to a new range. Default is [-1, 1]. + Args: data (1darr): data to scale - new_min (float, optional): new minimum value for data. Defaults to -1. - new_max (float, optional): new maximum value for data. Defaults to 1. - + new_min (float, optional): new minimum value for data. Defaults to -1. + new_max (float, optional): new maximum value for data. Defaults to 1. + Returns: 1darr: the newly scaled data """ @@ -51,10 +53,11 @@ def linear_convert(data, new_min=-1, new_max=1): new_range = new_max - new_min return new_min + new_range * (data - old_min) / old_range + def write_data(file_path, entries): - """Add data to a given dataset in 'file'. Creates dataset if it doesn't exist; - otherwise, appends. - Args: + """Add data to a given dataset in 'file'. Creates dataset if it doesn't exist; + otherwise, appends. + Args: file_path (string): the name of the output HDF5 file to which to append data entries (dict): dictionary of column name & corresponding data """ @@ -66,9 +69,10 @@ def write_data(file_path, entries): f[col_name][-1:] = new_data else: f.create_dataset(col_name, - data=np.expand_dims(col_data, axis=0), - maxshape=(None, col_data.shape[0]), - chunks=True) + data=np.expand_dims(col_data, axis=0), + maxshape=(None, col_data.shape[0]), + chunks=True) + # Constants from calibration_rp using RPRPData.csv f0 = 257.20857316296724 @@ -76,9 +80,10 @@ def write_data(file_path, entries): k = 33.42493417407945 c = -3.208233068626455 + def A(f): - """Calculates the expected displacement of the speaker at an inputted drive amplitude 'ampl' for a given frequency 'f', - based on the calibration fit at 0.2Vpp. + """Calculates the expected displacement of the speaker at an inputted drive amplitude 'ampl' for a given frequency 'f', + based on the calibration fit at 0.2Vpp. Args: f (1darr): frequencies at which to calculate expected displacement @@ -86,7 +91,8 @@ def A(f): Returns: 1darr: expected displacement/V_ampl in microns/V """ - return (k * f0**2) / np.sqrt((f0**2 - f**2)**2 + f0**2*f**2/Q**2) + return (k * f0**2) / np.sqrt((f0**2 - f**2)**2 + f0**2 * f**2 / Q**2) + def phase(f): """Calculates the phase delay between the speaker voltage waveform and the photodiode response @@ -98,11 +104,12 @@ def phase(f): Returns: 1darr: phase in radians """ - return np.arctan2(f0/Q*f, f**2 - f0**2) + c + return np.arctan2(f0 / Q * f, f**2 - f0**2) + c + def displacement_waveform(speaker_data, sample_rate, use_freq, max_freq): """Calculates the corresponding displacement waveform based on the given voltage waveform - using calibration. + using calibration. Args: speaker_data (1darr): voltage waveform for speaker @@ -116,25 +123,27 @@ def displacement_waveform(speaker_data, sample_rate, use_freq, max_freq): speaker_spectrum = fft(speaker_data) n = speaker_data.size if use_freq: - sample_spacing = 1/(2*max_freq) + sample_spacing = 1 / (2 * max_freq) else: - sample_spacing = 1/sample_rate - freq = fftfreq(n, d=sample_spacing) # units: cycles/s = Hz - + sample_spacing = 1 / sample_rate + freq = fftfreq(n, d=sample_spacing) # units: cycles/s = Hz + # Multiply signal by transfer func in freq domain, then return to time domain - converted_signal = speaker_spectrum * A(freq) * np.where(freq < 0, np.exp(-1j*phase(-freq)), np.exp(1j*phase(freq))) + converted_signal = speaker_spectrum * A(freq) * np.where(freq < 0, + np.exp(-1j * phase(-freq)), np.exp(1j * phase(freq))) y = np.real(ifft(converted_signal)) return y, converted_signal, freq + def velocity_waveform(speaker_data, sample_rate, use_freq, max_freq): """Calculates the corresponding velocity waveform based on the given voltage waveform - using calibration. + using calibration. Args: speaker_data (1darr): voltage waveform for speaker sample_rate (float): sample rate used to generate voltage waveform - + Returns: [1darr, 1darr, 1darr]: converted velocity waveform (microns/s) in time domain, converted velocity waveform in frequency domain, @@ -143,13 +152,14 @@ def velocity_waveform(speaker_data, sample_rate, use_freq, max_freq): speaker_spectrum = fft(speaker_data) n = speaker_data.size if use_freq: - sample_spacing = 1/(2*max_freq) + sample_spacing = 1 / (2 * max_freq) else: - sample_spacing = 1/sample_rate - freq = fftfreq(n, d=sample_spacing) # units: cycles/s = Hz - + sample_spacing = 1 / sample_rate + freq = fftfreq(n, d=sample_spacing) # units: cycles/s = Hz + # Multiply signal by transfer func in freq domain, then return to time domain - converted_signal = 1j*freq * speaker_spectrum * A(freq) * np.where(freq < 0, np.exp(-1j*phase(-freq)), np.exp(1j*phase(freq))) + converted_signal = 1j * freq * speaker_spectrum * \ + A(freq) * np.where(freq < 0, np.exp(-1j * phase(-freq)), np.exp(1j * phase(freq))) v = np.real(ifft(converted_signal)) - return v, converted_signal, freq \ No newline at end of file + return v, converted_signal, freq