From 4a3f2812ace709cab9c9ffab80651cd55363eef2 Mon Sep 17 00:00:00 2001 From: Karthik Chandrashekara Date: Fri, 3 Feb 2023 11:10:13 +0100 Subject: [PATCH] Added plotting, animating functions inside Shot class and added functionality to plot only user-defined sections of sequence, corrected the plotting of sequence by multiplying switches for certain channels. --- readAndPlotSequence.py | 229 +++++++++++++++++++---------------------- 1 file changed, 108 insertions(+), 121 deletions(-) diff --git a/readAndPlotSequence.py b/readAndPlotSequence.py index 9d07c70..c3638f0 100644 --- a/readAndPlotSequence.py +++ b/readAndPlotSequence.py @@ -1,5 +1,4 @@ import numpy, imageio, os -from scipy import interpolate import matplotlib.pyplot as plt from qtutils import * from labscript_utils.connections import ConnectionTable @@ -7,122 +6,6 @@ from labscript_utils import device_registry from labscript_c_extensions.runviewer.resample import resample as _resample import h5py - -def generate_ylabel(channel_name): - if channel_name == 'AO_MOT_Grad_Coil_current': - label = '$ \\nabla B_{AHH}$' - elif channel_name == 'AO_MOT_CompZ_Coil_current': - label = '$B_{HH}$' - elif channel_name == 'AO_MOT_3D_freq': - label = '$\Delta \\nu$' - elif channel_name == 'AO_MOT_3D_amp': - label = '$P_{3D}$' - elif channel_name == 'AO_Red_Push_amp': - label = '$P_{Push}$' - elif channel_name == 'MOT_2D_Shutter': - label = '$P_{2D}$' - elif channel_name == 'AO_ODT1_Pow': - label = '$P_{cODT1}$' - elif channel_name == 'Imaging_RF_Switch': - label = '$P_{img}$' - elif channel_name == 'MOT_3D_Camera_Trigger': - label = '$Cam$' - return label - -def plotSequence(Channels, Switches, animate = False, idx = 0): - - axs = plt.subplots(len(Channels), figsize = (10, 6), constrained_layout=True, sharex=True)[1] - - for i, channel_name in enumerate(Channels): - - #markers_unscaled = sorted(list(shotObj.markers.keys())) - #scalehandler = ScaleHandler(markers_unscaled, markers_unscaled, shotObj.stop_time) - #unscaled_time = numpy.asarray(Traces[channel_name])[0] - #scaled_time = scalehandler.get_scaled_time(unscaled_time) - - channel_time = numpy.asarray(Traces[channel_name])[0] - channel_trace = numpy.asarray(Traces[channel_name])[1] - xmin = 0 - xmax = shotObj.stop_time - dx = 1e-9 - resampled_channel_trace = shotObj.resample(channel_time, channel_trace, xmin, xmax, shotObj.stop_time, dx)[1] - time = numpy.arange(xmin, xmax, (xmax-xmin)/len(resampled_channel_trace)) - - switch_time = numpy.asarray(Traces[Switches[i]])[0] - switch_trace = numpy.asarray(Traces[Switches[i]])[1] - xmin = 0 - xmax = shotObj.stop_time - dx = 1e-9 - resampled_switch_trace = shotObj.resample(switch_time, switch_trace, xmin, xmax, shotObj.stop_time, dx)[1] - - trace = numpy.multiply(resampled_channel_trace, resampled_switch_trace) - - if not animate: - axs[i].plot(time, trace, '-b') #'-ob' - axs[i].fill_between(time, trace, alpha=0.4) - else: - axs[i].plot(time[0:idx], trace[0:idx], '-b') #'-ob' - axs[i].fill_between(time[0:idx], trace[0:idx], alpha=0.4) - - axs[i].axvline(x=0, color = 'b', linestyle = '--') - axs[i].axvline(x=4, color = 'b', linestyle = '--') - axs[i].axvline(x=4.315, color = 'b', linestyle = '--') - axs[i].axvline(x=shotObj.stop_time, color = 'b', linestyle = '--') - axs[i].set_xlim(0, shotObj.stop_time) - axs[i].set_ylim(0, max(resampled_channel_trace) + 0.2) - if i == len(Channels)-1: - axs[i].set_xlabel('Time (s)', fontsize = 16) - axs[i].set_ylabel(generate_ylabel(channel_name), fontsize = 16) - - if not animate: - plt.savefig(f'seq.png', format='png', bbox_inches = "tight") - plt.show() - else: - plt.savefig(f'seq-{idx}.png') - plt.close() - -def animateSequence(Channels, Switches): - - SIZE = 6000 - STEP = 58 - - for i in range(2, SIZE, STEP): - plotSequence(Channels, Switches, animate = True, idx = i) - - with imageio.get_writer('seq_animated.gif', mode='i', fps = 24, loop = 1) as writer: - for i in range(2, SIZE, STEP): - image = imageio.imread(f'seq-{i}.png') - writer.append_data(image) - os.remove(f'seq-{i}.png') - -class ScaleHandler(): - - def __init__(self, input_times, target_positions, stop_time): - # input_times is a list (may be unsorted) of times which should be scaled evenly with target_length - # an input list of [1,2,4,6] and target_length of 1.0 will result in: - # get_scaled_time(1) -> 1 - # get_scaled_time(1.5) -> 1.5 - # get_scaled_time(3) -> 2.5 - # get_scaled_time(4) -> 3 - # get_scaled_time(5) -> 3.5 ... - self.org_stop_time = float(stop_time) - - if not all((x >= 0) and (x <= self.org_stop_time) for x in input_times): - raise Exception('shot contains at least one marker before t=0 and/or after the stop time. Non-linear time currently does not support this.') - - unscaled_times = sorted(input_times) - scaled_times = sorted(target_positions) - - - # append values for linear scaling before t=0 and after stop time - unscaled_times = [min(unscaled_times)-1e-9] + unscaled_times + [max(unscaled_times) + 1e-9] - scaled_times = [min(scaled_times)-1e-9] + scaled_times + [max(scaled_times) + 1e-9] - - self.get_scaled_time = interpolate.interp1d(unscaled_times, scaled_times, assume_sorted=True, bounds_error=False, fill_value='extrapolate') - self.get_unscaled_time = interpolate.interp1d(scaled_times, unscaled_times, assume_sorted=True, bounds_error=False, fill_value='extrapolate') - - self.scaled_stop_time = self.get_scaled_time(self.org_stop_time) - class Shot(object): def __init__(self, path): @@ -273,6 +156,101 @@ class Shot(object): return x_out, y_out + def find_nearest(self, array, value): + array = numpy.asarray(array) + idx = (numpy.abs(array - value)).argmin() + return idx, array[idx] + + def generate_ylabel(self, channel_name): + if channel_name == 'AO_MOT_Grad_Coil_current': + label = '$ \\nabla B_{AHH}$' + elif channel_name == 'AO_MOT_CompZ_Coil_current': + label = '$B_{HH}$' + elif channel_name == 'AO_MOT_3D_freq': + label = '$\Delta \\nu$' + elif channel_name == 'AO_MOT_3D_amp': + label = '$P_{3D}$' + elif channel_name == 'AO_Red_Push_amp': + label = '$P_{Push}$' + elif channel_name == 'MOT_2D_Shutter': + label = '$P_{2D}$' + elif channel_name == 'AO_ODT1_Pow': + label = '$P_{cODT1}$' + elif channel_name == 'Imaging_RF_Switch': + label = '$P_{img}$' + elif channel_name == 'MOT_3D_Camera_Trigger': + label = '$Cam$' + return label + + def plotSequence(self, Channels, Switches, PlotRange, animate = False, idx = 0): + + Traces = self.traces + + axs = plt.subplots(len(Channels), figsize = (10, 6), constrained_layout=True, sharex=True)[1] + + for i, channel_name in enumerate(Channels): + + channel_time = numpy.asarray(Traces[channel_name])[0] + channel_trace = numpy.asarray(Traces[channel_name])[1] + xmin = 0 + xmax = self.stop_time + dx = 1e-9 + resampled_channel_trace = self.resample(channel_time, channel_trace, xmin, xmax, self.stop_time, dx)[1] + time = numpy.arange(xmin, xmax, (xmax-xmin)/len(resampled_channel_trace)) + + switch_time = numpy.asarray(Traces[Switches[i]])[0] + switch_trace = numpy.asarray(Traces[Switches[i]])[1] + xmin = 0 + xmax = self.stop_time + dx = 1e-9 + resampled_switch_trace = self.resample(switch_time, switch_trace, xmin, xmax, self.stop_time, dx)[1] + + trace = numpy.multiply(resampled_channel_trace, resampled_switch_trace) + + TrimRange = [self.find_nearest(time, PlotRange[0])[0], self.find_nearest(time, PlotRange[1])[0]] + trimmed_time = time[TrimRange[0]:TrimRange[1]] + trimmed_trace = trace[TrimRange[0]:TrimRange[1]] + + if not animate: + axs[i].plot(trimmed_time, trimmed_trace, '-b') #'-ob' + axs[i].fill_between(trimmed_time, trimmed_trace, alpha=0.4) + else: + axs[i].plot(time[0:TrimRange[0] + idx], trace[0:TrimRange[0] + idx], '-b') #'-ob' + axs[i].fill_between(time[0:TrimRange[0] + idx], trace[0:TrimRange[0] + idx], alpha=0.4) + + axs[i].axvline(x=0, color = 'b', linestyle = '--') + axs[i].axvline(x=4, color = 'b', linestyle = '--') + axs[i].axvline(x=4.315, color = 'b', linestyle = '--') + axs[i].axvline(x=self.stop_time, color = 'b', linestyle = '--') + axs[i].set_xlim(0, self.stop_time) + axs[i].set_ylim(0, max(resampled_channel_trace) + 0.2) + if i == len(Channels)-1: + axs[i].set_xlabel('Time (s)', fontsize = 16) + axs[i].set_ylabel(self.generate_ylabel(channel_name), fontsize = 16) + + if not animate: + plt.savefig(f'seq.png', format='png', bbox_inches = "tight") + plt.show() + else: + plt.savefig(f'seq-{idx}.png') + plt.close() + + def animateSequence(self, Channels, Switches, PlotRange): + + SIZE = 6000 + STEP = 58 + + for i in range(2, SIZE, STEP): + self.plotSequence(Channels, Switches, PlotRange, animate = True, idx = i) + + with imageio.get_writer('seq_animated.gif', mode='i', fps = 24, loop = 1) as writer: + for i in range(2, SIZE, STEP): + image = imageio.imread(f'seq-{i}.png') + writer.append_data(image) + + for i in range(2, SIZE, STEP): + os.remove(f'seq-{i}.png') + @property def channels(self): if self._channels is None: @@ -306,7 +284,6 @@ if __name__ == "__main__": shotObj._load() Channels = list(shotObj.channels) - Traces = shotObj.traces """ 'prawn_clock_line_0', 'prawn_clock_line_1', 'Dummy_1', 'Imaging_RF_Switch', 'Imaging_Shutter', 'MOT_2D_Shutter', 'MOT_3D_RF_Switch', 'MOT_3D_Shutter', 'Push_Beam_Blue_Shutter', @@ -324,8 +301,18 @@ if __name__ == "__main__": # Channels = ['MOT_2D_Shutter', 'AO_Red_Push_amp', 'AO_MOT_3D_amp', 'AO_MOT_3D_freq', 'AO_MOT_Grad_Coil_current', 'AO_MOT_CompZ_Coil_current', 'AO_ODT1_Pow', 'Imaging_RF_Switch', 'MOT_3D_Camera_Trigger'] # Switches = ['MOT_2D_Shutter', 'Push_Beam_Red_Switch', 'MOT_3D_RF_Switch', 'MOT_3D_RF_Switch', 'MOT_Grad_Coil_Switch', 'MOT_CompZ_Coil_Switch', 'CDT1_Switch', 'Imaging_RF_Switch', 'MOT_3D_Camera_Trigger'] - plotSequence(Channels, Switches) - - # animateSequence(Channels, Switches) + """ Plot Full Sequence """ + # TimeRange = [0.0, shotObj.stop_time] + + """ Plot till loading of MOT """ + # TimeRange = [0.0, 4.0] + + """ Plot from loading of MOT till end of sequence""" + TimeRange = [4.0, shotObj.stop_time] + + """ Plot sequence """ + # shotObj.plotSequence(Channels, Switches, PlotRange = TimeRange) + """ Animate sequence """ + shotObj.animateSequence(Channels, Switches, PlotRange = TimeRange)