diff --git a/acquire_automatic.py b/acquire_automatic.py index d0bc862..4dec919 100644 --- a/acquire_automatic.py +++ b/acquire_automatic.py @@ -6,15 +6,22 @@ import matplotlib.pyplot as plt import redpitaya_scpi as scpi import numpy as np -from scipy.fftpack import fft +from numpy.fft import fft, fftfreq +# from scipy.fft import fft # use numpy import math import util +from timeit import timeit +from datetime import datetime +from scipy.signal import butter, filtfilt IP = 'rp-f0c04a.local' rp_s = scpi.scpi(IP) print('Connected to ' + IP) +# plt.rcParams.update({ +# "text.usetex": True +# }) -def run_one_shot(use_freq, max_freq, start_freq=1, end_freq=1000, ampl=0.1, decimation=8192, +def run_one_shot(start_freq=1, end_freq=1000, ampl=0.1, gen_dec=8192, acq_dec=256, num_acq=1, store_data=False, plot_data=False, filename='data.h5py'): """Runs one shot of driving the speaker with a waveform and collecting the relevant data. @@ -30,15 +37,16 @@ def run_one_shot(use_freq, max_freq, start_freq=1, end_freq=1000, ampl=0.1, deci N = 16384 # Number of samples in buffer SMPL_RATE_DEC1 = 125e6 # sample rate for decimation=1 in Samples/s (Hz) - smpl_rate = SMPL_RATE_DEC1//decimation - burst_time = N / smpl_rate + gen_smpl_rate = SMPL_RATE_DEC1//gen_dec + acq_smpl_rate = SMPL_RATE_DEC1//acq_dec + burst_time = N / gen_smpl_rate wave_form = 'ARBITRARY' freq = 1 / burst_time - t, y = util.bounded_frequency_waveform(start_freq, end_freq, length=N, sample_rate=smpl_rate, - use_freq=use_freq, max_freq=max_freq) + t, y = util.bounded_frequency_waveform(start_freq, end_freq, length=N, sample_rate=gen_smpl_rate, invert=True) y = util.linear_convert(y) # convert range of waveform to [-1, 1] to properly set ampl + if plot_data: plt.plot(t, y) plt.show() @@ -49,94 +57,122 @@ def run_one_shot(use_freq, max_freq, start_freq=1, end_freq=1000, ampl=0.1, deci ##### Generation ##### # Function for configuring Source - rp_s.sour_set(1, wave_form, ampl, freq, data=y) + rp_s.sour_set(1, wave_form, 1, freq, data=ampl*y) # Enable output rp_s.tx_txt('OUTPUT1:STATE ON') rp_s.tx_txt('SOUR1:TRig:INT') + # print("output enabled", datetime.now().time()) ##### Acqusition ##### + pd_data = [] + speaker_data = [] + vel_data = [] # Function for configuring Acquisition - rp_s.acq_set(dec=decimation, trig_delay=0) - rp_s.tx_txt('ACQ:START') - time.sleep(1) - rp_s.tx_txt('ACQ:TRig CH2_PE') + rp_s.tx_txt('ACQ:RST') + rp_s.acq_set(dec=acq_dec, trig_delay=8192) + rp_s.tx_txt('ACQ:START') + # ! OS 2.00 or higher only ! ## + time.sleep(0.4) + rp_s.tx_txt('ACQ:TRig NOW') # CH2_PE + time.sleep(0.4) # Wait for trigger while 1: rp_s.tx_txt('ACQ:TRig:STAT?') # Get Trigger Status if rp_s.rx_txt() == 'TD': # Triggered? + # print("td", datetime.now().time()) break - ## ! OS 2.00 or higher only ! ## while 1: rp_s.tx_txt('ACQ:TRig:FILL?') if rp_s.rx_txt() == '1': + # print("filled", datetime.now().time()) break - - ##### Analaysis ##### + ##### Analysis ##### # Read data and plot function for Data Acquisition - pd_data = np.array(rp_s.acq_data(chan=1, convert=True)) # Volts - speaker_data = np.array(rp_s.acq_data(chan=2, convert=True)) # Volts - velocity_data, converted_signal, freq = util.velocity_waveform(speaker_data, smpl_rate, use_freq, max_freq) - displ_data, _, _ = util.displacement_waveform(speaker_data, smpl_rate, use_freq, max_freq) - y_vel, y_converted, _ = util.velocity_waveform(ampl*y, smpl_rate, use_freq, max_freq) - time_data = np.linspace(0, N-1, num=N) / smpl_rate - + pds = np.array(rp_s.acq_data(chan=1, convert=True)) + speaker = np.array(rp_s.acq_data(chan=2, convert=True)) + acq_time_data = np.linspace(0, N-1, num=N) / acq_smpl_rate + if plot_data: - fig, ax = plt.subplots(nrows=3) - - ax[0].plot(time_data, pd_data, color='blue', label='Observed PD') - ax[0].plot(time_data, speaker_data, color='black', label='Observed Drive') - ax[0].plot(time_data, ampl*y, label='Drive Output', alpha=0.5) - ax[0].legend() - ax[0].set_ylabel('Amplitude (V)') - ax[0].set_xlabel('Time (s)') - - ax[1].plot(freq, np.abs(fft(speaker_data)), color='black', label='Observed Drive', marker='.') - ax[1].plot(freq, np.abs(converted_signal), color='green', label='Expected Observed Vel', marker='.') - ax[1].plot(freq, np.abs(fft(ampl*y)), color='blue', label='Expected Drive', marker='.') - ax[1].plot(freq, np.abs(y_converted), color='orange', label='Expected Ideal Vel', marker='.') - ax[1].loglog() - ax[1].set_xlabel('Frequency (Hz)') - ax[1].set_ylabel('$|\^{V}|$') - ax[1].legend() - - ax[2].plot(time_data, velocity_data, color='black', label='Observed Drive') - ax[2].plot(time_data, y_vel, label='Drive Output') - ax[2].set_ylabel('Expected Vel (Microns/s)') - ax[2].set_xlabel('Time (s)') - ax[2].legend() - plt.tight_layout() - plt.show() - + pd_data = pds + speaker_data = speaker + # pd_data.append(pds) # Volts + # speaker_data.append(speaker) # Volts + # vel_data.append(vels) + if store_data: # Store data in h5py file path = "/Users/angelajia/Code/College/SMI/data/" # filename = "training_data.h5py" file_path = os.path.join(path, filename) - entries = { - 'Time (s)': time_data, - 'Speaker (V)': speaker_data, - 'Speaker (Microns/s)': velocity_data, - 'PD (V)': pd_data, - 'Speaker (Microns)': displ_data + 'Speaker (V)': speaker, + 'Speaker (Microns/s)': vels, + 'PD (V)': pds } - util.write_data(file_path, entries) + + if plot_data: + gen_time_data = np.linspace(0, N-1, num=N) / gen_smpl_rate + # pd_data = np.concatenate(pd_data) + # speaker_data = np.concatenate(speaker_data) + # vel_data = np.concatenate(vel_data) + y_vel, y_converted, _ = util.velocity_waveform(ampl*y, gen_smpl_rate) + avg_speaker = np.mean(np.reshape(speaker, (-1, 32)), axis=1) + print(avg_speaker.shape) + vel_data, vel_converted, freqs= util.velocity_waveform(avg_speaker, acq_smpl_rate) + + fig, ax = plt.subplots(nrows=4) + + ax[0].plot(pd_data, color='blue', label='Observed PD')# , marker='.') + ax[0].plot(speaker_data, color='black', label='Observed Drive') #, marker='.') + # ax[0].plot(time_data, ampl*y, label='Drive Output', alpha=0.5) + ax[0].legend() + ax[0].set_ylabel('Amplitude (V)') + ax[0].set_xlabel('Samples') + + #ax[1].set_title(r'$\tilde{F}^{-1}(VelTransferFunc(f) * \tilde{F}(Voltage Time Series)(f) )$') + ax[1].set_title("Vel for acq_dec=256") + ax[1].plot(vel_data) + ax[1].set_xlabel('Samples') + ax[1].set_ylabel('Expected Vel (Microns/s)') + + # ax[3].set_title("Vel for gen_dec=8192") + # ax[3].plot(gen_time_data, y_vel, label='Drive Output', alpha=0.7) + # ax[3].set_ylabel('Expected Vel (Microns/s)') + # ax[3].set_xlabel('Time (s)') + # ax[1].legend() + + # ax[2].set_title("Generated Speaker Voltage (gen_dec=8192)") + # ax[2].plot(gen_time_data, ampl*y) + # ax[2].set_xlabel('Time (s)') + # ax[2].set_ylabel('Amplitude (V)') + + ax[2].plot(fftfreq(speaker.shape[0], d=1/acq_smpl_rate), np.abs(fft(speaker, norm='ortho')), marker='.') + ax[3].plot(freqs, np.abs(vel_converted), marker='.') + plt.tight_layout() + plt.show() ##### Reset when closing program ##### rp_s.tx_txt('GEN:RST') rp_s.tx_txt('ACQ:RST') -num_shots = 2000 + +filenames = ['test_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_2kx1shots_randampl.h5py', + 'val_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_2kx1shots_randampl.h5py'] +shots = [1, 0] amplitude = 0 -end_freq = 1000 -max_freq = 10*end_freq -for i in range(num_shots): - amplitude = np.random.uniform(0.1, 0.6) - if i % 400 == 0: - print(f"{i}: ampl = {amplitude}") - run_one_shot(True, max_freq, 30, end_freq, ampl=amplitude, decimation=256, store_data=True, plot_data=False, - filename='test_max10kHz_30to1kHz_2kshots_dec=256_randampl.h5py') - # print(i) \ No newline at end of file +acq_num = 1 +print("start", datetime.now().time()) +for (file, num_shots) in zip(filenames, shots): + print(file) + print(num_shots) + for i in range(num_shots): + amplitude = 0.1 # np.random.uniform(0.1, 0.6) + if i % 500 == 0: + print("\t", datetime.now().time()) + print(f"\t{i*acq_num}: ampl = {amplitude}") + run_one_shot(100, 101, ampl=amplitude, gen_dec=8192, acq_dec=256, num_acq=acq_num, store_data=False, plot_data=True, + filename=file) + print("end", datetime.now().time()) \ No newline at end of file diff --git a/acquire_continuous.py b/acquire_continuous.py index 6292be4..9542bcc 100644 --- a/acquire_continuous.py +++ b/acquire_continuous.py @@ -12,12 +12,12 @@ print('Connected to ' + IP) wave_form = "SINE" -freq = 1000 +freq = 120 ampl = 0.1 N = 16384 # Number of samples in buffer SMPL_RATE_DEC1 = 125e6 # sample rate for decimation=1 in Samples/s (Hz) -decimation = 32 +decimation = 256# 32 smpl_rate = SMPL_RATE_DEC1//decimation # Reset Generation and Acquisition diff --git a/decoder.py b/decoder.py index 14254ba..f7425a6 100644 --- a/decoder.py +++ b/decoder.py @@ -60,8 +60,8 @@ def configure_optimizers(self): # We will reduce the learning rate by factor of gamma after 100 and 150 # epochs scheduler = optim.lr_scheduler.MultiStepLR(optimizer, - milestones=[100, 150, 200], - gamma=1) + milestones=[50, 100, 150], #changed from [100, 150, 200] + gamma=1) return [optimizer], [scheduler] def get_loss_function(self, loss_hparams): diff --git a/model_analysis.py b/model_analysis.py deleted file mode 100644 index 2d2a29b..0000000 --- a/model_analysis.py +++ /dev/null @@ -1,99 +0,0 @@ -import os, sys, glob -import h5py -import numpy as np -import matplotlib.pyplot as plt -import matplotlib as mpl -import torch -import lightning as L -from torch.utils.data import Dataset, DataLoader - -import train -import decoder - -N = 16384 -SMPL_RATE_DEC1 = 125e6 -decimation = 256 -smpl_rate = SMPL_RATE_DEC1/decimation -time_data = np.linspace(0, N-1, num=N) / smpl_rate - -valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\valid_30to1kHz_2kshots_dec=256_randampl.h5py' -train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\training_30to1kHz_10kshots_dec=256_randampl.h5py' -test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_30to1kHz_2kshots_dec=256_randampl.h5py' - -model_tag = "fu7qnzar" -step = 128 -batch_size = 128 - -runner = train.TrainingRunner(train_file, valid_file, test_file, step=step) -model, result = runner.load_model(model_tag) - -train_mode = True - -if train_mode: - iter_loader = iter(runner.train_loader) -else: - iter_loader = iter(runner.test_loader) -# for i in range(9): -# next(iter_loader) -batch_val = next(iter_loader) -batches = [batch_val] - -for batch in batches: - ## Plot results - inputs = batch[0] - targets = batch[1] - - print(inputs.shape) - print(targets.shape) - outputs_val, _ = model.predict_step(batch, 1, test_mode=True) - - outputs_val = torch.squeeze(outputs_val).cpu().detach().numpy() - inputs_val = torch.squeeze(inputs).cpu().detach().numpy() - targets_val = torch.squeeze(targets).cpu().detach().numpy() - - targets_squeezed_val = np.squeeze(targets_val) - print("targets shape", targets_squeezed_val.shape) - outputs_squeezed_val = np.squeeze(outputs_val) - print("preds shape", outputs_squeezed_val.shape) - - fig, ax = plt.subplots(2) - fig.set_size_inches(8, 6) - - if train_mode: - inputs_flattened = inputs_val.flatten() - ax[0].plot(inputs_flattened) - ax[0].set_title('Training PD Trace (scrambled buffer)', fontsize=15) - ax[0].set_ylabel('V', fontsize=15) - ax[0].set_xticks([]) - ax[0].tick_params(axis='y', which='major', labelsize=13) - # ax[0].set_xlabel('Time (s)') - ax[1].plot(targets_squeezed_val, marker='.', label='Target') - ax[1].set_title('Velocity', fontsize=15) - ax[1].set_ylabel('um/s', fontsize=15) - ax[1].set_xlabel('Batch Idx', fontsize=15) - - ax[1].plot(outputs_squeezed_val, marker='.', label='Pred') - ax[1].legend(prop={'size': 12}) - ax[1].tick_params(axis='both', which='major', labelsize=13) - else: - idx = 0 - ax[0].plot(time_data, inputs_val[idx]) - ax[0].set_title('Test PD Trace (contiguous buffer)', fontsize=15) - ax[0].set_ylabel('V', fontsize=15) - ax[0].set_xticks([]) - ax[0].tick_params(axis='y', which='major', labelsize=13) - # ax[0].set_xlabel('Time (s)') - num_groups = outputs_squeezed_val.shape[1] # 127 - start_idxs = torch.arange(num_groups) * step - - ax[1].plot(time_data[start_idxs], targets_squeezed_val[idx], marker='.', label='Target') - ax[1].set_title('Velocity', fontsize=15) - ax[1].set_ylabel('um/s', fontsize=15) - ax[1].set_xlabel('Time (s)', fontsize=15) - - ax[1].plot(time_data[start_idxs], outputs_squeezed_val[idx], marker='.', label='Pred') - ax[1].legend(prop={'size': 12}) - ax[1].tick_params(axis='both', which='major', labelsize=13) - - fig.tight_layout() - plt.show() \ No newline at end of file diff --git a/models.py b/models.py index 0b17e7b..f4a2c17 100644 --- a/models.py +++ b/models.py @@ -37,7 +37,7 @@ def __init__(self, input_size, output_size, in_channels=1, activation='LeakyReLU self.fc_layers = nn.Sequential( # length of input = 64 filters * length of 10 left nn.Linear(640, 16), - nn.ReLU(), + act_fn_by_name[activation], nn.Linear(16, output_size) ) diff --git a/plot_2dhists.py b/plot_2dhists.py new file mode 100644 index 0000000..680c93a --- /dev/null +++ b/plot_2dhists.py @@ -0,0 +1,121 @@ +import os, sys, glob +import h5py +import numpy as np +import matplotlib.pyplot as plt +import matplotlib as mpl +import torch +from torch import optim, nn +import lightning as L +from torch.utils.data import Dataset, DataLoader +from matplotlib.colors import LogNorm + +import train +import decoder + +if __name__ == '__main__': + if len(sys.argv) == 1: + N = 16384 + SMPL_RATE_DEC1 = 125e6 + decimation = 256 + smpl_rate = SMPL_RATE_DEC1 / decimation + time_data = np.linspace(0, N - 1, num=N) / smpl_rate + + # valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\valid_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_2kx1shots_randampl.h5py' + # test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_2kx1shots_randampl.h5py' + # train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\train_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_10kx1shots_randampl.h5py' + valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\val_1to1kHz_invertspectra_trigdelay8192_sleep100ms_2kx1shots_dec=256_8192_randampl.h5py' + test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_1to1kHz_invertspectra_trigdelay8192_sleep100ms_2kx1shots_dec=256_8192_randampl.h5py' + train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\train_1to1kHz_invertspectra_trigdelay8192_sleep100ms_10kx1shots_dec=256_8192_randampl.h5py' + + model_tag = "z4dscrhq" #"qn4eo8zu"#"z4dscrhq" + step = 64 + group_size = 256 + + runner = train.TrainingRunner(train_file, valid_file, test_file, step=step) + model, result = runner.load_model(model_tag) + + modes = ['train', 'valid', 'test'] + + # only works for datasets using test_mode=True + for mode in ['test']: + if mode == 'valid': + iter_loader = iter(runner.valid_loader_testmode) + elif mode == 'test': # test mode + iter_loader = iter(runner.test_loader) + else: + raise ValueError('Invalid mode') + + preds = torch.empty(0) + truths = torch.empty(0) + input_segments = torch.empty(0) + + num_batches = 10 # ceil[2k/128] + for i in range(num_batches): + print(i) + batch = next(iter_loader) + inputs = torch.squeeze(batch[0]) + targets = torch.squeeze(batch[1]) + outputs = torch.squeeze(model.predict_step(batch, 1, test_mode=True)[0]) # store only y_hat + # print("Input shape", inputs.shape) # [batch_size, 16384] + # print("Target shape", targets.shape) # [batch_size, num_groups] + # print("Pred shape", outputs.shape) # [batch_size, num_groups] + + preds = torch.cat((preds, torch.flatten(outputs)), dim=0) + truths = torch.cat((truths, torch.flatten(targets)), dim=0) + + preds = preds.cpu().detach().numpy() + truths = truths.cpu().detach().numpy() + + num_samples = preds.shape[0] + print("samples", num_samples) + losses = (preds - truths)**2 + print("losses", losses.shape) + + fig, axs = plt.subplots(1, 3, figsize=(18, 6), sharex=True, sharey=True, tight_layout=True) #, gridspec_kw={'height_ratios': [1, 1, 1]}) + num_bins = 51 # odd to see if 0-velocity over/under-predicted + + range_min = min(np.min(truths), np.min(preds)) + range_max = max(np.max(truths), np.max(preds)) + range_max = max(np.abs(range_min), range_max) # find max magnitude in all data + hist_range = [[-range_max, range_max]] * 2 + # 2D Histogram for Counts + hist_counts, xedges_counts, yedges_counts = np.histogram2d(truths, preds, bins=num_bins, range=hist_range) + masked_hist_counts = np.ma.masked_where(hist_counts == 0, hist_counts) + im_counts_log = axs[0].imshow(masked_hist_counts.T, extent=[xedges_counts[0], xedges_counts[-1], yedges_counts[0], yedges_counts[-1]], + origin='lower', aspect='equal', cmap='coolwarm', norm=LogNorm()) + axs[0].set_xlabel('Expected Velocity (um/s)') + axs[0].set_ylabel('Predicted Velocity (um/s)') + axs[0].set_title('Counts for Predicted vs Expected Velocity') + cbar_0 = fig.colorbar(im_counts_log, ax=axs[0], orientation='horizontal', pad=0.1) + cbar_0.set_label('Counts Per Bin') + # avg_mse = hist_mse_sum / hist_counts # avg mse per bin + + # 2D Histogram for Average MSELoss + hist_mse_sum, xedges, yedges = np.histogram2d(truths, preds, bins=num_bins, weights=losses, range=hist_range) + # masked_hist = np.ma.masked_where(hist_mse_sum == 0, hist_mse_sum) + cax = axs[1].imshow((hist_mse_sum/masked_hist_counts).T, extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]], + origin='lower', aspect='equal', cmap='coolwarm', norm=LogNorm()) + # with np.errstate(divide='ignore', invalid='ignore'): + # cax = axs[1].imshow((hist_mse_sum / hist_counts).T, extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]], origin='lower', aspect='auto') + cbar_1 = fig.colorbar(cax, ax=axs[1], orientation='horizontal', pad=0.1) + cbar_1.set_label('Average MSE Loss Per Bin') + axs[1].set_xlabel('Expected Velocity (um/s)') + # axs[1].set_ylabel('Predicted Velocity') + axs[1].set_title('Avg MSELoss for Predicted vs Expected Velocity') + + # 2D Histogram for Summed MSELoss + sum_ax = axs[2].imshow(hist_mse_sum.T, extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]], + origin='lower', aspect='equal', cmap='coolwarm', norm=LogNorm()) + cbar_sum = fig.colorbar(sum_ax, ax=axs[2], orientation='horizontal', pad=0.1) + cbar_sum.set_label('Summed MSE Loss Per Bin') + axs[2].set_xlabel('Expected Velocity (um/s)') + # axs[2].set_ylabel('Predicted Velocity') + axs[2].set_title('Summed MSELoss for Predicted vs Expected Velocity') + + for ax in axs: + ax.axline((0, 0), slope=1, color='black') + ax.axline((0, 0), slope=-1, color='black') + ax.set_aspect(np.diff(ax.get_xlim()) / np.diff(ax.get_ylim())) + plt.show() + else: + print("Error: Unsupported number of command-line arguments") diff --git a/plot_predictions.py b/plot_predictions.py new file mode 100644 index 0000000..be1bf08 --- /dev/null +++ b/plot_predictions.py @@ -0,0 +1,109 @@ +import os, sys, glob +import h5py +import numpy as np +import matplotlib.pyplot as plt +import matplotlib as mpl +import torch +import lightning as L +from torch.utils.data import Dataset, DataLoader + +import train +import decoder + +N = 16384 +SMPL_RATE_DEC1 = 125e6 +decimation = 256 +smpl_rate = SMPL_RATE_DEC1 / decimation +time_data = np.linspace(0, N - 1, num=N) / smpl_rate + +valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\val_1to1kHz_invertspectra_trigdelay8192_sleep100ms_2kx1shots_dec=256_8192_randampl.h5py' +test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_1to1kHz_invertspectra_trigdelay8192_sleep100ms_2kx1shots_dec=256_8192_randampl.h5py' +train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\train_1to1kHz_invertspectra_trigdelay8192_sleep100ms_10kx1shots_dec=256_8192_randampl.h5py' + +model_tag = "z4dscrhq" +step = 64 +batch_size = 128 + +runner = train.TrainingRunner(train_file, valid_file, test_file, step=step) +model, result = runner.load_model(model_tag) + +modes = ['train', 'valid', 'test'] + +for mode in modes: + if mode == 'train': + iter_loader = iter(runner.train_loader) + elif mode == 'valid': + iter_loader = iter(runner.valid_loader_testmode) + else: # test mode + iter_loader = iter(runner.test_loader) + for i in range(2): + next(iter_loader) + batch_val = next(iter_loader) # batch 3 + batches = [batch_val] + next(iter_loader) + batches.append(next(iter_loader)) # batch 5 + next(iter_loader) + batches.append(next(iter_loader)) # batch 7 + + for batch in batches: + ## Plot results + inputs = batch[0] + targets = batch[1] + + print(inputs.shape) + print(targets.shape) + outputs_val, _ = model.predict_step(batch, 1, test_mode=True) + + outputs_val = torch.squeeze(outputs_val).cpu().detach().numpy() + inputs_val = torch.squeeze(inputs).cpu().detach().numpy() + targets_val = torch.squeeze(targets).cpu().detach().numpy() + + targets_squeezed_val = np.squeeze(targets_val) + print("targets shape", targets_squeezed_val.shape) + outputs_squeezed_val = np.squeeze(outputs_val) + print("preds shape", outputs_squeezed_val.shape) + + fig, ax = plt.subplots(2) + fig.set_size_inches(8, 6) + + if mode == 'train': + inputs_flattened = inputs_val.flatten() + ax[0].plot(inputs_flattened) + ax[0].set_title('Training PD Trace (scrambled buffer)', fontsize=15) + ax[0].set_ylabel('V', fontsize=15) + ax[0].set_xticks([]) + ax[0].tick_params(axis='y', which='major', labelsize=13) + # ax[0].set_xlabel('Time (s)') + ax[1].plot(targets_squeezed_val, marker='.', label='Target') + ax[1].set_title('Velocity', fontsize=15) + ax[1].set_ylabel('um/s', fontsize=15) + ax[1].set_xlabel('Batch Idx', fontsize=15) + + ax[1].plot(outputs_squeezed_val, marker='.', label='Pred') + ax[1].legend(prop={'size': 12}) + ax[1].tick_params(axis='both', which='major', labelsize=13) + else: + idx = 0 # use first in batch of 128 shots + ax[0].plot(time_data, inputs_val[idx]) + if mode == 'valid': + ax[0].set_title('Validation PD Trace (contiguous buffer)', fontsize=15) + elif mode == 'test': + ax[0].set_title('Test PD Trace (contiguous buffer)', fontsize=15) + ax[0].set_ylabel('V', fontsize=15) + ax[0].set_xticks([]) + ax[0].tick_params(axis='y', which='major', labelsize=13) + # ax[0].set_xlabel('Time (s)') + num_groups = outputs_squeezed_val.shape[1] # 127 + start_idxs = torch.arange(num_groups) * step + + ax[1].plot(time_data[start_idxs], targets_squeezed_val[idx], marker='.', label='Target') + ax[1].set_title('Velocity', fontsize=15) + ax[1].set_ylabel('um/s', fontsize=15) + ax[1].set_xlabel('Time (s)', fontsize=15) + + ax[1].plot(time_data[start_idxs], outputs_squeezed_val[idx], marker='.', label='Pred') + ax[1].legend(prop={'size': 12}) + ax[1].tick_params(axis='both', which='major', labelsize=13) + + fig.tight_layout() + plt.show() \ No newline at end of file diff --git a/plot_quartiles.py b/plot_quartiles.py new file mode 100644 index 0000000..507ef94 --- /dev/null +++ b/plot_quartiles.py @@ -0,0 +1,147 @@ +import os, sys, glob +import h5py +import numpy as np +import matplotlib.pyplot as plt +import matplotlib as mpl +import torch +from torch import optim, nn +import lightning as L +from torch.utils.data import Dataset, DataLoader +from matplotlib.colors import LogNorm + +import train +import decoder + +if __name__ == '__main__': + if len(sys.argv) == 1: + N = 16384 + SMPL_RATE_DEC1 = 125e6 + decimation = 256 + smpl_rate = SMPL_RATE_DEC1 / decimation + time_data = np.linspace(0, N - 1, num=N) / smpl_rate + + # valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\valid_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_2kx1shots_randampl.h5py' + # test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_2kx1shots_randampl.h5py' + # train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\train_1to1kHz_misaligned_invertspectra_trigdelay8192_sleep100ms_10kx1shots_randampl.h5py' + valid_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\val_1to1kHz_invertspectra_trigdelay8192_sleep100ms_2kx1shots_dec=256_8192_randampl.h5py' + test_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\test_1to1kHz_invertspectra_trigdelay8192_sleep100ms_2kx1shots_dec=256_8192_randampl.h5py' + train_file = 'C:\\Users\\aj14\\Desktop\\SMI\\data\\train_1to1kHz_invertspectra_trigdelay8192_sleep100ms_10kx1shots_dec=256_8192_randampl.h5py' + + model_tag = "z4dscrhq" #"qn4eo8zu"#"z4dscrhq" + step = 64 + group_size = 256 + + runner = train.TrainingRunner(train_file, valid_file, test_file, step=step) + model, result = runner.load_model(model_tag) + + modes = ['train', 'valid', 'test'] + + # only works for datasets using test_mode=True + iter_loader = iter(runner.test_loader) + + preds = torch.empty(0) + truths = torch.empty(0) + inputs = torch.empty(0) + + num_batches = 10 # ceil[2k/128] + for i in range(num_batches): + print(i) + batch = next(iter_loader) + input_buffers = torch.squeeze(batch[0]) + targets = torch.squeeze(batch[1]) + outputs = torch.squeeze(model.predict_step(batch, 1, test_mode=True)[0]) # store only y_hat + # print("Input shape", inputs.shape) # [batch_size, 16384] + # print("Target shape", targets.shape) # [batch_size, num_groups] + # print("Pred shape", outputs.shape) # [batch_size, num_groups] + + preds = torch.cat((preds, outputs), dim=0) + truths = torch.cat((truths, targets), dim=0) + inputs = torch.cat((inputs, input_buffers), dim=0) + + preds = preds.cpu().detach().numpy() # [num_batches * batch_size, num_groups] + truths = truths.cpu().detach().numpy() # [num_batches * batch_size, num_groups] + inputs = inputs.cpu().detach().numpy() # [num_batches * batch_size, 16384] + + print("preds shape", preds.shape) + print("truths shape", truths.shape) + print("inputs shape", inputs.shape) + + losses = np.mean((preds - truths)**2, axis=1) # avg_mse_loss per buffer, [num_batches * batch_size] + print("losses shape", losses.shape) + + sorted_idxs = np.argsort(losses) + # all sorted arrays should have same shape as their unsorted originals + sorted_losses = losses[sorted_idxs] + print("sorted losses shape", sorted_losses.shape) + sorted_truths = truths[sorted_idxs] + print("sorted truths shape", sorted_truths.shape) + sorted_preds = preds[sorted_idxs] + print("sorted preds shape", sorted_preds.shape) + sorted_inputs = inputs[sorted_idxs] + print("sorted inputs shape", sorted_inputs.shape) + + num_buffers = sorted_losses.shape[0] + q2_start_idx = num_buffers // 4 # End of first quartile + q3_start_idx = num_buffers // 2 # End of second quartile + q4_start_idx = 3 * num_buffers // 4 # End of third quartile + + first_quartile = { + "losses": sorted_losses[:q2_start_idx], + "truths": sorted_truths[:q2_start_idx], + "preds": sorted_preds[:q2_start_idx], + "inputs": sorted_inputs[:q2_start_idx] + } + + fourth_quartile = { + "losses": sorted_losses[q4_start_idx:], + "truths": sorted_truths[q4_start_idx:], + "preds": sorted_preds[q4_start_idx:], + "inputs": sorted_inputs[q4_start_idx:] + } + + print("First quartile buffer count:", first_quartile["inputs"].shape) + print("Fourth quartile buffer count:", fourth_quartile["inputs"].shape) + + def plot_quartile_examples(quartiles, num_exs=4): + fig, axes = plt.subplots(num_exs * 2, 2, figsize=(15, num_exs * 5), + gridspec_kw={'height_ratios': [1, 2] * num_exs}) + + time_data = np.linspace(0, N - 1, num=N) / smpl_rate + num_groups = truths.shape[1] + start_idxs = torch.arange(num_groups) * step + for i in range(num_exs): + for j, quartile in enumerate(quartiles): + # Plot PD trace + axes[i*2, j].plot(time_data, quartile['inputs'][i], color='blue') + axes[i*2, j].set_ylabel('V') + axes[i*2, j].set_xticks([]) + + # Plot velocities + axes[i*2 + 1, j].plot(time_data[start_idxs], quartile['truths'][i], + label='Target', color='blue', + marker='o', markersize=2) + axes[i*2 + 1, j].plot(time_data[start_idxs], quartile['preds'][i], + label=f'Pred (Avg MSE: {quartile["losses"][i]:.2f})', color='orange', marker='o', + markersize=2) + axes[i*2 + 1, j].set_ylabel('µm/s') + axes[i*2 + 1, j].legend() + + if i == num_exs - 1: # a bottom plot + axes[i*2 + 1, j].set_xlabel('Time (s)') + else: + axes[i*2 + 1, j].set_xticks([]) + + fig.text(0.25, 0.98, '1st Quartile (Best Performance)', ha='center', fontsize=15) + fig.text(0.75, 0.98, '4th Quartile (Worst Performance)', ha='center', fontsize=15) + plt.tight_layout(rect=[0, 0, 1, 0.96]) # Make space for suptitles + plt.show() + + # Call the function with the first and fourth quartile data + plot_quartile_examples( + quartiles=[ + first_quartile, # Best quartile data + fourth_quartile # Worst quartile data + ] + ) + else: + print("Error: Unsupported number of command-line arguments") \ No newline at end of file diff --git a/train.py b/train.py index 80e03c8..97e8c65 100644 --- a/train.py +++ b/train.py @@ -187,6 +187,8 @@ def set_dataloaders(self, batch_size=128): 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_testmode = self.get_custom_dataloader( + True, 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): diff --git a/util.py b/util.py index 40b1413..63639b8 100644 --- a/util.py +++ b/util.py @@ -1,11 +1,9 @@ import numpy as np -from scipy.fftpack import fft, ifft, fftfreq +from numpy.fft import fft, ifft, fftfreq import h5py - 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 @@ -18,20 +16,20 @@ def bounded_frequency_waveform(start_frequency, end_frequency, length, sample_ra # 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 - 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]))) # 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)) + phi = 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)]) - y = np.real(ifft(spectrum)) + spectrum = spectrum * c * np.exp(1j*phi) + if invert: + spectrum = np.divide(spectrum, A(freq) * np.exp(1j*phase(freq))) + spectrum = np.hstack([spectrum, np.zeros_like(spectrum)]) + y = np.real(ifft(spectrum, norm="ortho")) y = np.fft.fftshift(y) return t, y @@ -74,7 +72,7 @@ def write_data(file_path, entries): col_data.shape[1]), chunks=True) - + # Constants from calibration_rp using RPRPData.csv f0 = 257.20857316296724 Q = 15.804110908084784 @@ -121,23 +119,19 @@ def displacement_waveform(speaker_data, sample_rate, use_freq, max_freq): converted displacement waveform in frequency domain, frequency array (Hz) """ - speaker_spectrum = fft(speaker_data) + speaker_spectrum = fft(speaker_data, norm="ortho") n = speaker_data.size - if use_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))) - y = np.real(ifft(converted_signal)) + 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, norm="ortho")) return y, converted_signal, freq - -def velocity_waveform(speaker_data, sample_rate, use_freq, max_freq): +def velocity_waveform(speaker_data, sample_rate): """Calculates the corresponding velocity waveform based on the given voltage waveform using calibration. @@ -150,17 +144,14 @@ def velocity_waveform(speaker_data, sample_rate, use_freq, max_freq): converted velocity waveform in frequency domain, frequency array (Hz) """ - speaker_spectrum = fft(speaker_data) + speaker_spectrum = fft(speaker_data, norm="ortho") n = speaker_data.size - if use_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))) - v = np.real(ifft(converted_signal)) + A(freq) * np.where(freq < 0, np.exp(-1j*phase(-freq)), np.exp(1j*phase(freq))) + v = np.real(ifft(converted_signal, norm="ortho")) return v, converted_signal, freq