karthik
/
Plotting
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

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.

master
Karthik 2 years ago
parent
159f46345c
commit
4a3f2812ac
1 changed files with 108 additions and 121 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 227
      readAndPlotSequence.py

227
readAndPlotSequence.py
View File

@ -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)
""" 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]
# animateSequence(Channels, Switches)
""" Plot sequence """
# shotObj.plotSequence(Channels, Switches, PlotRange = TimeRange)
""" Animate sequence """
shotObj.animateSequence(Channels, Switches, PlotRange = TimeRange)
Write Preview
Loading…
Cancel
Save