def plot_q_(images, r_support, Eq, qs_center, rs_center):
kxu, kyu, Equ = asbp.unroll_q(r_support, Eq, qs_center)
Equ = Equ * mathx.expj((rs_center[0] * kxu + rs_center[1] * kyu))
images[0].setImage(abs(Equ))
images[0].setRect(pg.axes_to_rect(kxu / 1e3, kyu / 1e3))
images[1].setImage(np.angle(Equ) / (2 * np.pi))
images[1].setRect(pg.axes_to_rect(kxu / 1e3, kyu / 1e3))
def set_log(self, x, y, time, h_log, v_log, **kwargs):
x = x.squeeze()
y = y.squeeze()
self.vlog.setImage(image=v_log, autoLevels=True, **kwargs)
self.hlog.setImage(image=h_log, autoLevels=True, **kwargs)
if len(time) > 1:
self.hlog.setRect(pg.axes_to_rect(x, time))
self.vlog.setRect(pg.axes_to_rect(time, y))
def plot_r_q_polar(self, flat=False, tilt=False, show=True):
"""Plot approximate plane profile and surface z relative to z_mean."""
app = self.app
Er = app.Er_flat if flat else app.Er
Eq = math.fft2(Er)
if not tilt:
Er = Er * mathx.expj(-(app.qs_center[0] * app.x +
app.qs_center[1] * app.y))
Eq = Eq * mathx.expj(app.rs_center[0] * app.kx +
app.rs_center[1] * app.ky)
glw = pg.GraphicsLayoutWidget()
glw.addLabel(self.title_str)
glw.nextRow()
gl = glw.addLayout()
plot = gl.addAlignedPlot(labels={'left': 'y (mm)', 'bottom': 'x (mm)'})
x, y, zu = sa.unroll_r(self.rs_support, self.z, self.rs_center)
image = plot.image((zu - self.mean_z) * 1e3,
rect=pg.axes_to_rect(x * 1e3, y * 1e3),
lut=pg.get_colormap_lut('bipolar'))
gl.addHorizontalSpacer(10)
gl.addColorBar(image=image, rel_row=2, label='Relative z (mm)')
glw.nextRow()
glw.addLabel('Approximate planar profile')
glw.nextRow()
gl = glw.addLayout()
plots = plotting.plot_r_q_polar(app.rs_support, Er, app.rs_center,
app.qs_center, gl, Eq)
glw.resize(830, 675)
if show:
glw.show()
return glw, plots
def set(self,
x,
y,
image,
horz_proj=None,
vert_proj=None,
pen=pg.mkPen(),
**kwargs):
"""Set image and projection data
TODO: more flexibility and more consistent semantics with options on
pyqtgraph plotting commands
Args:
x (vector of length m): the x axis
y (vector of length n): the y axis
image (mxn array): the image
horz_proj (vector of length m): horizontal projection. If None, use
sum over y axis of image
vert_proj (vector of length n): vertical projection. If None, use
sum over x axis of image
pen: passed on to the projection setData methods
lut: passed on to image setImage method
"""
if any([x is None, y is None, image is None]):
return
x = x.squeeze()
y = y.squeeze()
if horz_proj is None:
horz_proj = image.sum(1)
if vert_proj is None:
vert_proj = image.sum(0)
self.image.setImage(image, autoLevels=False, **kwargs)
self.image.setRect(pg.axes_to_rect(x, y))
self.horz_proj.setData(x, horz_proj, pen=pen)
self.vert_proj.setData(vert_proj, y, pen=pen)
def set_Eq_image_item(item, r_support, Eq, qs_center=(0, 0), quantity='waves'):
kx, ky, Equ = sa.unroll_q(r_support, Eq, qs_center)
if quantity == 'amplitude':
data = abs(Equ)
levels = 0, data.max()
elif quantity == 'waves':
data = np.angle(Equ)/(2*np.pi)
levels = -0.5, 0.5
else:
raise ValueError('Unknown quantity %s.', quantity)
item.setImage(data)
item.setRect(pg.axes_to_rect(kx/1e3, ky/1e3))
item.setLevels(levels)
return item
def make_Er_image_item(r_support, Er, rs_center=(0, 0), quantity='waves'):
x, y, Eru = sa.unroll_r(r_support, Er, rs_center)
if quantity == 'amplitude':
data = abs(Eru)
lut = pg.get_colormap_lut()
levels = 0, data.max()
elif quantity == 'waves':
data = np.angle(Eru)/(2*np.pi)
lut = pg.get_colormap_lut('bipolar')
levels = -0.5, 0.5
item = pg.ImageItem(data, lut=lut)
item.setRect(pg.axes_to_rect(x*1e3, y*1e3))
item.setLevels(levels)
return item
def plot_r(images, beam, scale=1):
profile = beam.profile
xu, yu, Eru = asbp.unroll_r(profile.rs_support, profile.Er * scale,
profile.rs_center)
tilt_mode = self.tilt_mode_widget.currentIndex()
if not isinstance(beam.profile, asbp.PlaneProfile):
q_profile = profile.app
else:
q_profile = profile
if tilt_mode == 0:
q0s = q_profile.centroid_qs_flat
else:
q0s = q_profile.peak_qs
q0s += np.asarray(
(self.tilt_nudge_x_widget.value(),
self.tilt_nudge_y_widget.value())) / 1e3 * profile.k
Eru = Eru * mathx.expj(-(q0s[0] *
(xu - profile.peak_rs[0]) + q0s[1] *
(yu - profile.peak_rs[1]) +
np.angle(profile.peak_Er)))
images[0].setImage(abs(Eru))
images[0].setRect(pg.axes_to_rect(xu * 1e3, yu * 1e3))
images[1].setImage(np.angle(Eru) / (2 * np.pi))
images[1].setRect(pg.axes_to_rect(xu * 1e3, yu * 1e3))
def set_pulse_2D(self, **kwargs):
axes, intensity, phase = self.process_pulse_2D_kwargs(**kwargs)
if intensity is not None:
self.images.intensity.setImage(intensity.T, autoLevels=False)
if phase is not None:
kwargs = {}
if self.symmetric_phase_colormap:
mn = phase.min()
mx = phase.max()
rng = max(mn, mx)
kwargs['levels'] = -rng, rng
self.images.phase.setImage(phase.T, autoLevels=False, **kwargs)
if axes is not None:
rect = pg.axes_to_rect(*axes)
self.images.intensity.setRect(rect)
self.images.phase.setRect(rect)
def show(self):
rms = (self.ix.var(axis=(-1, -2)) + self.iy.var(axis=(-1, -2)))**0.5
ixmu = self.ix.mean(axis=(-1, -2))
iymu = self.iy.mean(axis=(-1, -2))
glw = pg.GraphicsLayoutWidget()
# title = 'Input & output aperture side length %.1f mm. %dx%d plane waves, each with %dx%d rays.'%(
# self.field_side_length*1e3, num_spots, num_spots, num_rays, num_rays)
# glw.addLabel(title, colspan=2)
# glw.nextRow()
gl = glw.addLayout()
spot_plot = gl.addAlignedPlot(labels={
'left': 'field y (mm)',
'bottom': 'field x (mm)'
},
title='Spots')
index = 0
for ix_, iy_ in zip(self.ix, self.iy):
for ix, iy in zip(ix_, iy_):
item = pg.ScatterPlotItem(ix.ravel() * 1e3,
iy.ravel() * 1e3,
pen=None,
brush=pg.tableau20[index % 20],
size=2)
spot_plot.addItem(item)
index += 1
gl = glw.addLayout()
rms_plot = gl.addAlignedPlot(labels={
'left': 'theta y (mrad)',
'bottom': 'theta x (mrad)'
},
title='RMS spot size')
rms_image = rms_plot.image(rms * 1e6,
rect=pg.axes_to_rect(
self.thetax * 1e3, self.thetay * 1e3),
lut=pg.get_colormap_lut())
gl.addHorizontalSpacer()
gl.addColorBar(image=rms_image, rel_row=2, label='Size (micron)')
glw.resize(920, 400)
glw.show()
return glw
def set_image(self,
x,
y,
image,
pen=pg.mkPen('k'),
datetimestr=None,
**kwargs):
"""
Args:
kwargs: passed on to image.setImage
"""
if any([x is None, y is None, image is None]):
return
horz_proj = image.sum(1)
vert_proj = image.sum(0)
x = x.squeeze()
y = y.squeeze()
self.image.setImage(image, autoLevels=False, **kwargs)
self.image.setRect(pg.axes_to_rect(x, y))
self.hproj.setData(x, horz_proj, pen=pen)
self.vproj.setData(vert_proj, y, pen=pen)
if datetimestr is not None:
self.image_cornertext.setText(datetimestr, color='00FF00')
"""
Check out how ImageItem behaves with nan values and how to make a mask e.g. to
grey out undefined phase values.
More investigation than testing.
"""
import numpy as np
import pyqtgraph_extended as pg
##
x=np.arange(100)
y=np.arange(120)[:,None]
f=np.exp(-((x-50)**2+(y-60)**2)/200)
fp=f.copy()
fp[0,0]=float('nan')
fp[f<0.2]=float('nan')
plt=pg.plot()
im=pg.ImageItem(fp)
im.setRect(pg.axes_to_rect(x,y))
im.setLookupTable(pg.get_colormap_lut())
im.setLevels((0,1))
plt.addItem(im)
##
f3=pg.makeARGB(f,pg.get_colormap_lut(),levels=(0,1),useRGBA=True)[0]#[:,:,[1,2,3,0]]
f3[f<0.1]=[128,128,128,128]
plt=pg.plot()
im=pg.ImageItem(f3)
im.setRect(pg.axes_to_rect(x,y))
#im.setLookupTable(pg.get_colormap_lut())
plt.addItem(im)
def set_E(self, **kwargs):
"""Set signal, calculate phase space distribution and update plots.
Args:
Ex: signal in x domain
Ek: signal in k domain
lut: Pyqtgraph lookup table
S_transform (str): None or 'log10'
log10_range (int): decades for log color scale
bound_cond (str): passed on to :func:`fourier.spectrogram', one of
'pad','odd-cyclic','cyclic'
"""
try:
self.log10_range = kwargs['log10_range']
except KeyError:
pass
try:
self.S_transform = kwargs['S_transform']
except KeyError:
pass
assert self.S_transform in (None, 'log10')
try:
self.ftd = kwargs['ftd']
except KeyError:
pass
try:
self.Ek = kwargs['Ek']
if self.Ek is not None:
self.Ex = None
except KeyError:
pass
try:
self.Ex = kwargs['Ex']
if self.Ex is not None:
self.Ek = None
except KeyError:
pass
try:
self.bound_cond = kwargs['bound_cond']
except KeyError:
pass
if self.Ex is None and self.Ek is None:
return
if self.Ex is None:
self.Ex = self.ftd.inv_trans(self.Ek)
else:
self.Ek = self.ftd.trans(self.Ex)
Ix = abs(self.Ex.squeeze())**2
if self.scale['Ix'] == 'max':
Ix /= Ix.max()
else:
Ix /= self.scale['Ix']
Ik = abs(self.Ek.squeeze())**2
if self.scale['Ik'] == 'max':
Ik /= Ik.max()
else:
Ik /= self.scale['Ik']
self.lines['Ix'].setData(self.ftd.x.squeeze() / self.scale['x'], Ix)
self.lines['Ik'].setData(Ik, self.ftd.k.squeeze() / self.scale['k'])
xs, ks, S = fourier.spectrogram(self.ftd.x.squeeze(), 1,
self.Ex.squeeze(), 32,
self.ftd.k.squeeze(), self.bound_cond)
if self.scale['S'] == 'max':
S /= S.max()
else:
S /= self.scale['S']
if self.S_transform == 'log10':
S[S == 0] = S[S != 0].min()
S = np.log10(S)
self.image.setImage(S.T)
self.image.setRect(
pg.axes_to_rect(xs / self.scale['x'], ks / self.scale['k']))
if self.S_transform == 'log10':
mx = S.max()
self.image.setLevels((mx - self.log10_range, mx))
def _update(self):
p = self.profile
if p is not None:
# Set relative z plot.
x, y, z = sa.unroll_r(p.rs_support, p.z, p.rs_center)
self.dz_image.setImage((z - p.mean_z) * 1e3)
self.dz_image.setRect(pg.axes_to_rect(x, y, 1e3))
field_index = self.field_combo_box.currentIndex()
if field_index == 0:
Er = p.Er
Eq = p.Eq
elif field_index == 1:
Er = p.app.Er
Eq = p.app.Eq
elif field_index == 2:
Er = p.app.Er_flat
Eq = p.app.Eq_flat
else:
raise ValueError('Unknown mode %d.', field_index)
if self.remove_tilt_check_box.isChecked():
Er = Er * mathx.expj(-(p.qs_center[0] *
(p.x - p.rs_center[0]) +
p.qs_center[1] *
(p.y - p.rs_center[1])))
Eq = Eq * mathx.expj(p.rs_center[0] *
(p.app.kx - p.app.qs_center[0]) +
p.app.rs_center[1] *
(p.app.ky - p.app.qs_center[1]))
if self.remove_constant_check_box.isChecked():
Er = Er * mathx.expj(-np.angle(Er[p.r_center_indices]))
Eq = Eq * mathx.expj(-np.angle(Eq[p.app.q_center_indices]))
plotting.set_Er_image_item(self.images.r.abs, p.rs_support, Er,
p.rs_center, 'amplitude')
plotting.set_Er_image_item(self.images.r.phase, p.rs_support, Er,
p.rs_center, 'waves')
plotting.set_Eq_image_item(self.images.q.abs, p.rs_support, Eq,
p.qs_center, 'amplitude')
plotting.set_Eq_image_item(self.images.q.phase, p.rs_support, Eq,
p.qs_center, 'waves')
title = p.title_str
idx = title.find('qs_center')
self.label.setText('<br>'.join((title[:idx], title[idx:])))
else:
for images_ in self.images:
for image in images_:
image.setImage(None)
self.label.setText('')
if self.line is None or p is None:
x = np.asarray([])
y = np.asarray([])
kx = np.asarray([])
ky = np.asarray([])
else:
point = self.line.origin
vector = self.line.vector
# Scatter plot item wants 1D arrays.
x, y = rt1.to_xy(point.reshape((-1, 4)))
kx, ky = rt1.to_xy(vector.reshape((-1, 4))) * p.k
if field_index in (1, 2):
# Remove spherical wavefront from transverse wavenumbers.
kx -= (x - p.rs_center[0]) * p.k / p.rocs[0]
ky -= (y - p.rs_center[1]) * p.k / p.rocs[1]
self.scatters.r.abs.setData(x * 1e3, y * 1e3)
self.scatters.r.phase.setData(x * 1e3, y * 1e3)
self.scatters.q.abs.setData(kx / 1e3, ky / 1e3)
self.scatters.q.phase.setData(kx / 1e3, ky / 1e3)
def plot_polar_image(r,
phi,
data,
layout=None,
r_label_phi=math.pi / 4,
grid_r=None,
levels=None,
lut=None,
cbar_label=None,
phi_0=0,
phi_dir=1,
r_label_fmt='%g',
aspect_locked=True,
mask_array=None,
theta_repeats=None,
grid_color='w',
osamp=1):
"""
Args:
r (uniformly sampled vector of shape (n,1) or (n,)): radii
phi (uniformly sampled vector of shape (m,)): azimuthal angles in radians
data (array of shape (n,m)): data to plot
layout (pyqtgraph GraphicsLayout): if None, a GraphicsLayoutWidget is created.
An AlignedPlotItem is added to the layout.
mask_array (array of same shape as data): mask for data - where it is negative
data will be greyed out. The pyqtgraph ImageItem will be in RGB rather
than scalar format, and the color bar will not be adjustable.
"""
assert usv.is_usv(r)
assert usv.is_usv(phi)
if layout is None:
layout_ = pg.GraphicsLayoutWidget()
else:
layout_ = layout
plt = layout_.addAlignedPlot()
plt.hideAxis('left')
plt.hideAxis('bottom')
if lut is None:
lut = pg.get_colormap_lut()
if isinstance(lut, str):
lut = pg.get_colormap_lut(lut)
phip = phi_0 + phi_dir * phi
x, y, data = mathx.polar_reg_grid_to_rect(r.squeeze(), phip, data,
theta_repeats, osamp)
if mask_array is not None:
_, _, mask_array = mathx.polar_reg_grid_to_rect(
r.squeeze(), phip, mask_array, theta_repeats, osamp)
if levels is None:
data_masked = data[mask_array >= 0]
levels = data_masked.min(), data_masked.max()
image_data = pg.makeARGB(data.T, lut, levels=levels, useRGBA=True)[0]
image_data[mask_array.T < 0] = [128, 128, 128, 128]
image = plt.image(image_data,
rect=pg.axes_to_rect(x, y),
levels=(0, 255))
else:
image = plt.image(data.T, rect=pg.axes_to_rect(x, y), lut=lut)
if levels is not None:
image.setLevels(levels)
if grid_r is None:
grid_r = np.arange(1, 5) / 5 * r[-1]
if len(grid_r) > 0:
plot_radial_grid_lines(plt,
grid_r,
grid_color,
r_label_fmt,
label_angle=phi_0 +
phi_dir * np.asarray(r_label_phi))
phi_lines = (np.round(phi[-1] /
(math.pi / 2)) - np.arange(4)) * math.pi / 2
plot_azimuthal_grid_lines(plt,
phi_lines,
grid_r[-1],
grid_color,
'%gπ',
unit=math.pi,
phi_0=phi_0,
phi_dir=phi_dir)
plt.setAspectLocked(aspect_locked)
plt.setXRange(x[0], x[-1], padding=0)
plt.setYRange(y[0], y[-1], padding=0)
if cbar_label is not None:
layout_.addHorizontalSpacer(5)
# Add color bar
cbar = layout_.addColorBar(label=cbar_label, rel_row=2)
if mask_array is not None:
cbar.setManual(lut, levels)
else:
cbar.setImage(image)
if layout is None:
layout_.show()
return layout_
else:
return plt
