Script to plot an experiment sequence executed for a shot as extracted from the shot file.

DyLevelStructure/visualizeDyLevelStructureWithTransitions.py

@@ -71,7 +71,9 @@ def parse_NIST_data(path, min_J, max_J, max_wavenumber):
 def plot_level_structure_with_red_and_blue_transitions(*args, **kwargs):
     
     #start plotting
-    f, ax = plt.subplots(figsize=(4, 8))
+    f, ax = plt.subplots(figsize=(4.5, 7))
+    #plt.subplots_adjust(top=0.973, bottom=0.121, left=0.244, right=0.959, hspace=0.2, wspace=0.2)
+
     named_colors = ['r', 'm']
     Red_Blue_colors = ['#ab162a', '#cf5246', '#eb9172', '#fac8af', '#faeae1', '#e6eff4', '#bbdaea', '#7bb6d6', '#3c8abe', '#1e61a5']
 
@@ -79,53 +81,55 @@ def plot_level_structure_with_red_and_blue_transitions(*args, **kwargs):
     plot_handle = sns.scatterplot(x='J', y='Wavenumber', data = dataframe, s=2000, hue = 'Parity', palette = sns.color_palette(named_colors), marker = '_', linewidth=1.5, legend=False)
     
     #write electronic configuration for GS
-    ax.text(gs_J + 0.15, gs_wavenumber + 400, '$6s^2$') 
+    #ax.text(gs_J + 0.15, gs_wavenumber + 400, '$6s^2$') 
     #draw guide line for GS
     #plt.axhline(y=gs_wavenumber, color='m', linestyle='--', linewidth=1, alpha=0.5)
     
     #write wavelength of red transition
-    ax.text(red_J - 0.4, red_wavenumber * 0.5, '$626.082 ~ \mathrm{nm}$', color = '#db2929')
+    ax.text(red_J - 0.4, red_wavenumber * 0.5, '626.082 nm', color = '#db2929', fontsize = 16, fontweight = 'bold')
-    ax.text(red_J - 0.4, red_wavenumber * 0.46, '$(\\Gamma = 2\\pi\\times 136 ~ \mathrm{kHz})$', fontsize = 8, color = '#db2929')
+    #ax.text(red_J - 0.49, red_wavenumber * 0.42, '$(\\Gamma = 2\\pi\\times 136 ~ \mathrm{kHz})$', fontsize = 12, color = '#db2929')
     #draw red transition arrow
     ax.annotate('', 
             xy=(red_J, red_wavenumber),
             xytext=(gs_J, gs_wavenumber), 
-            arrowprops=dict(color='#db2929', alpha=0.8, width=1.5),
+            arrowprops=dict(color='#db2929', width=1.5),
             horizontalalignment='right',
             verticalalignment='top')
 
     #write electronic configuration for triplet excited state
-    ax.text(red_J + 0.35, red_wavenumber + 200, '$6s6p(^3P_1)$', fontsize = 10) 
+    #ax.text(red_J + 0.35, red_wavenumber + 200, '$6s6p(^3P_1)$', fontsize = 10) 
     #draw guide line for triplet excited state
     plt.axhline(y=red_wavenumber, color='m', linestyle='--', linewidth=1, alpha=0.5)
 
     #write wavelength of red transition
-    ax.text(blue_J - 1.5, blue_wavenumber * 0.55, '$421.291~ \mathrm{nm}$', color = '#2630ea')
+    ax.text(blue_J - 1.8, blue_wavenumber * 0.4, '421.291 nm', color = '#2630ea', fontsize = 16, fontweight = 'bold')
-    ax.text(blue_J - 1.55, blue_wavenumber * 0.52, '$(\\Gamma = 2\\pi\\times 32.2 ~ \mathrm{MHz})$', fontsize = 8, color = '#2630ea')
+    #ax.text(blue_J - 1.85, blue_wavenumber * 0.4, '$(\\Gamma = 2\\pi\\times$', fontsize = 12, color = '#2630ea') #$(\\Gamma = 2\\pi\\times 32.2 ~ \mathrm{MHz})$
+    #ax.text(blue_J - 1.5, blue_wavenumber * 0.35, '$32.2 ~ \mathrm{MHz})$', fontsize = 12, color = '#2630ea')
     #draw blue transition arrow
     ax.annotate('', 
             xy=(blue_J, blue_wavenumber),
             xytext=(gs_J, gs_wavenumber), 
-            arrowprops=dict(color='#2630ea', alpha=0.8, width=3.5),
+            arrowprops=dict(color='#2630ea', width=3.5),
             horizontalalignment='right',
             verticalalignment='top')
 
     #write electronic configuration for singlet excited state
-    ax.text(blue_J + 0.35, blue_wavenumber + 200, '$6s6p(^1P_1)$', fontsize = 10) 
+    #ax.text(blue_J + 0.35, blue_wavenumber + 200, '$6s6p(^1P_1)$', fontsize = 10) 
     #draw guide line for singlet excited state
     plt.axhline(y=blue_wavenumber, color='m', linestyle='--', linewidth=1, alpha=0.5)
 
     #figure options
     f.canvas.draw()
-    plt.xlabel('J', fontsize=16)
+    plt.xlabel('J', fontsize=20, fontweight = 'bold')
     #plt.ylabel('$\\tilde{v}~(cm^{-1})$', fontsize=16)
-    plt.ylabel('Wavelength (nm)', fontsize=16)
+    plt.ylabel('Wavelength (nm)', fontsize=20, fontweight = 'bold')
     
     plot_handle.set_xticks(range(min_J-1, max_J+2))
     ax.get_xticklabels()[0].set_visible(False)
     ax.get_xticklabels()[-1].set_visible(False)
     ax.get_xticklines()[0].set_visible(False)
     ax.get_xticklines()[-2].set_visible(False)
+    ax.set_xticklabels(ax.get_xticks(), weight='bold')
 
     yticklabels = [item.get_text() for item in ax.get_yticklabels()]
     yticklabels = ['' if item.startswith('−') or item.startswith('0') else item for item in yticklabels]
@@ -133,14 +137,15 @@ def plot_level_structure_with_red_and_blue_transitions(*args, **kwargs):
     new_yticks = np.arange(min(yticks), max_wavenumber, 4000)
     plot_handle.set_yticks(new_yticks)
     new_yticklabels = [round(1e7/item) if item != 0 else item for item in new_yticks]
-    ax.set_yticklabels(new_yticklabels)
+    ax.set_yticklabels(new_yticklabels, weight='bold')
     ax.get_yticklabels()[0].set_visible(False)
     #ax.get_yticklabels()[-1].set_visible(False)
     ax.get_yticklines()[0].set_visible(False)
     #ax.get_yticklines()[-2].set_visible(False)
 
-    plt.tick_params(axis='both', which='major', labelsize=14)
+    plt.tick_params(axis='both', which='major', labelsize=16)
 
+    #f.tight_layout()
     #plt.show()
     
     f.savefig(Path(home_path + os.sep + 'result.pdf'), format='pdf', bbox_inches = "tight")
@@ -149,7 +154,8 @@ if __name__ == '__main__':
     
     min_J = 8
     max_J = 9
-    max_wavenumber = 24500.0
+    #max_wavenumber = 24500.0
+    max_wavenumber = 28500.0
 
     #Ground State: [Xe]4f^{10} 6s^2(^5I_8)
     gs_J = 8

readAndPlotSequence.py

@@ -0,0 +1,331 @@
+import numpy, imageio, os
+from scipy import interpolate
+import matplotlib.pyplot as plt
+from qtutils import *
+from labscript_utils.connections import ConnectionTable
+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):
+        self.path = path
+
+        # Store list of traces
+        self._traces = None
+        # store list of channels
+        self._channels = None
+        # store list of markers
+        self._markers = None
+        
+        # store list of shutter changes and callibrations
+        self._shutter_times = None
+        self._shutter_calibrations = {}
+
+        # Load connection table
+        self.connection_table = ConnectionTable(path)
+
+        # open h5 file
+        with h5py.File(path, 'r') as file:
+            # Get master pseudoclock
+            self.master_pseudoclock_name = file['connection table'].attrs['master_pseudoclock']
+            if isinstance(self.master_pseudoclock_name, bytes):
+                self.master_pseudoclock_name = self.master_pseudoclock_name.decode('utf8')
+            else:
+                self.master_pseudoclock_name = str(self.master_pseudoclock_name)
+
+            # get stop time
+            self.stop_time = file['devices'][self.master_pseudoclock_name].attrs['stop_time']
+
+            self.device_names = list(file['devices'].keys())
+
+            # Get Shutter Calibrations
+            if 'calibrations' in file and 'Shutter' in file['calibrations']:
+                for name, open_delay, close_delay in numpy.array(file['calibrations']['Shutter']):
+                    name = name.decode('utf8') if isinstance(name, bytes) else str(name)
+                    self._shutter_calibrations[name] = [open_delay, close_delay]
+
+    def _load(self):
+        if self._channels is None:
+            self._channels = {}
+        if self._traces is None:
+            self._traces = {}
+        if self._markers is None:
+            self._markers = {}
+        if self._shutter_times is None:
+            self._shutter_times = {}
+
+        self._load_markers()
+        # Let's walk the connection table, starting with the master pseudoclock
+        master_pseudoclock_device = self.connection_table.find_by_name(self.master_pseudoclock_name)
+
+        self._load_device(master_pseudoclock_device)
+
+    def _load_markers(self):
+        with h5py.File(self.path, 'r') as file:
+            if "time_markers" in file:
+                for row in file["time_markers"]:
+                    self._markers[row['time']] = {'color': row['color'].tolist()[0], 'label': row['label']}
+            elif "runviewer" in file:
+                for time, val in file["runviewer"]["markers"].attrs.items():
+                    props = val.strip('{}}').rsplit(",", 1)
+                    color = list(map(int, props[0].split(":")[1].strip(" ()").split(",")))
+                    label = props[1].split(":")[1]
+                    self._markers[float(time)] = {'color': color, 'label': label}
+            if 0 not in self._markers:
+                self._markers[0] = {'color': [0,0,0], 'label': 'Start'}
+            if self.stop_time not in self._markers:
+                self._markers[self.stop_time] = {'color': [0,0,0], 'label' : 'End'}
+
+    def add_trace(self, name, trace, parent_device_name, connection):
+        name = str(name)
+        self._channels[name] = {'device_name': parent_device_name, 'port': connection}
+        self._traces[name] = trace
+
+        # add shutter times
+        con = self.connection_table.find_by_name(name)
+        if con.device_class == "Shutter" and 'open_state' in con.properties:
+            self.add_shutter_times([(name, con.properties['open_state'])])
+
+    # Temporary solution to physical shutter times
+    def add_shutter_times(self, shutters):
+        for name, open_state in shutters:
+            x_values, y_values = self._traces[name]
+            if len(x_values) > 0:
+                change_indices = numpy.where(y_values[:-1] != y_values[1:])[0]
+                change_indices += 1 # use the index of the value that is changed to
+                change_values = list(zip(x_values[change_indices], y_values[change_indices]))
+                change_values.insert(0, (x_values[0], y_values[0])) # insert first value
+                self._shutter_times[name] = {x_value + (self._shutter_calibrations[name][0] if y_value == open_state else self._shutter_calibrations[name][1]): 1 if y_value == open_state else 0 for x_value, y_value in change_values}
+
+    def _load_device(self, device, clock=None):
+        try:
+            module = device.device_class
+            device_class = device_registry.get_runviewer_parser(module)
+            device_instance = device_class(self.path, device)
+            
+            clocklines_and_triggers = device_instance.get_traces(self.add_trace, clock)
+            
+            
+            for name, trace in clocklines_and_triggers.items():
+                child_device = self.connection_table.find_by_name(name)
+                for grandchild_device_name, grandchild_device in child_device.child_list.items():
+                    self._load_device(grandchild_device, trace)
+
+        except Exception:
+            pass
+
+    def resample(self, data_x, data_y, xmin, xmax, stop_time, num_pixels):
+        """This is a function for downsampling the data before plotting
+        it. Unlike using nearest neighbour interpolation, this method
+        preserves the features of the plot. It chooses what value to
+        use based on what values within a region are most different
+        from the values it's already chosen. This way, spikes of a short
+        duration won't just be skipped over as they would with any sort
+        of interpolation."""
+        # TODO: Only finely sample the currently visible region. Coarsely sample the rest
+        # x_out = numpy.float32(numpy.linspace(data_x[0], data_x[-1], 4000*(data_x[-1]-data_x[0])/(xmax-xmin)))
+        x_out = numpy.float64(numpy.linspace(xmin, xmax, 3 * 2000 + 2))
+        y_out = numpy.empty(len(x_out) - 1, dtype=numpy.float64)
+        data_x = numpy.float64(data_x)
+        data_y = numpy.float64(data_y)
+
+        # TODO: investigate only resampling when necessary.
+        #       Currently pyqtgraph sometimes has trouble rendering things
+        #       if you don't resample. If a point is far off the graph,
+        #       and this point is the first that should be drawn for stepMode,
+        #       because there is a long gap before the next point (which is
+        #       visible) then there is a problem.
+        #       Also need to explicitly handle cases where none of the data
+        #       is visible (which resampling does by setting NaNs)
+        #
+        # x_data_slice = data_x[(data_x>=xmin)&(data_x<=xmax)]
+        # print len(data_x)
+        # if len(x_data_slice) < 3*2000+2:
+        #    x_out = x_data_slice
+        #    y_out = data_y[(data_x>=xmin)&(data_x<=xmax)][:-1]
+        #    logger.info('skipping resampling')
+        # else:
+        resampling = True
+
+        if resampling:
+            _resample(data_x, data_y, x_out, y_out, numpy.float64(stop_time))
+            # self.__resample4(data_x, data_y, x_out, y_out, numpy.float32(stop_time))
+        else:
+            x_out, y_out = data_x, data_y
+
+        return x_out, y_out
+
+    @property
+    def channels(self):
+        if self._channels is None:
+            self._load()
+
+        return self._channels.keys()
+
+    @property
+    def markers(self):
+        if self._markers is None:
+            self._load()
+        return self._markers
+
+    @property
+    def traces(self):
+        if self._traces is None:
+            self._load()
+        return self._traces
+    
+    @property
+    def shutter_times(self):
+        if self._shutter_times is None:
+            self._load()
+        return self._shutter_times
+
+if __name__ == "__main__":
+
+    filepath = 'C:/Users/Karthik/Desktop/PreTalk/Plot Sequence/2023-01-25_0069_ODT_Imaging_9.h5'
+    
+    shotObj = Shot(filepath)
+    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', 
+    'Push_Beam_Blue_Switch', 'Push_Beam_Red_Shutter', 'Push_Beam_Red_Switch', 'CDT1_Switch', 'CDT2_Switch', 'MOT_3D_Camera_Trigger', 'MOT_3D_Camera_trigger', 'MOT_CompX_Coil_Switch', 
+    'MOT_CompY_Coil_Switch', 'MOT_CompZ_Coil_Switch', 'MOT_Grad_Coil_Switch', 'ODT_Axis_Camera_Trigger', 'ODT_Axis_Camera_trigger', 'AO_Blue_Push_amp', 'AO_Blue_Push_freq', 'AO_Imaging_amp', 
+    'AO_Imaging_freq', 'AO_MOT_3D_amp', 'AO_MOT_3D_freq', 'AO_MOT_CompX_Coil_current', 'AO_MOT_CompX_Coil_voltage', 'AO_MOT_CompY_Coil_current', 'AO_MOT_CompY_Coil_voltage', 'AO_MOT_CompZ_Coil_current', 
+    'AO_MOT_CompZ_Coil_voltage', 'AO_MOT_Grad_Coil_current', 'AO_MOT_Grad_Coil_voltage', 'AO_Red_Push_amp', 'AO_Red_Push_freq', 'AO_Dummy', 'AO_ODT1_Mod', 'AO_ODT1_Pow', 'AO_ODT2_Pow'
+    """
+
+    """ Without Imaging """
+    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']
+    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']
+
+    """ With Imaging """
+    # 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)
+    
+