Browse Source
Script to plot an experiment sequence executed for a shot as extracted from the shot file.
master
Script to plot an experiment sequence executed for a shot as extracted from the shot file.
master
Karthik
2 years ago
2 changed files with 352 additions and 15 deletions
@ -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) |
|||
|
|||
|
Write
Preview
Loading…
Cancel
Save
Reference in new issue