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.
This commit is contained in:
parent
159f46345c
commit
4a3f2812ac
@ -1,5 +1,4 @@
|
|||||||
import numpy, imageio, os
|
import numpy, imageio, os
|
||||||
from scipy import interpolate
|
|
||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
from qtutils import *
|
from qtutils import *
|
||||||
from labscript_utils.connections import ConnectionTable
|
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
|
from labscript_c_extensions.runviewer.resample import resample as _resample
|
||||||
import h5py
|
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):
|
class Shot(object):
|
||||||
|
|
||||||
def __init__(self, path):
|
def __init__(self, path):
|
||||||
@ -273,6 +156,101 @@ class Shot(object):
|
|||||||
|
|
||||||
return x_out, y_out
|
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
|
@property
|
||||||
def channels(self):
|
def channels(self):
|
||||||
if self._channels is None:
|
if self._channels is None:
|
||||||
@ -306,7 +284,6 @@ if __name__ == "__main__":
|
|||||||
shotObj._load()
|
shotObj._load()
|
||||||
|
|
||||||
Channels = list(shotObj.channels)
|
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',
|
'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']
|
# 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']
|
# 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)
|
""" Plot Full Sequence """
|
||||||
|
# TimeRange = [0.0, shotObj.stop_time]
|
||||||
|
|
||||||
# animateSequence(Channels, Switches)
|
""" 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)
|
||||||
|
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user