def test_imagetool_bin(self, qtbot):
dat = self.make_regular_data()
it = ImageTool(dat)
# Set cursor index to (2, 1, 0), and bin in index by 2, so we average over three elements in x
it.info_bar.cursor_i[0].setValue(2)
it.info_bar.cursor_i[1].setValue(1)
it.info_bar.bin_i[0].setValue(2)
assert it.info_bar.bin_c[0].value() == pytest.approx(2*dat.delta[0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].yData, dat.values[:, 1, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].xData,
np.mean(dat.values[1:4, :, 0:1], axis=(0, 2)))
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].xData,
np.mean(dat.values[1:4, 1:2, :], axis=(0, 1)))
# Confirm that updating the bin width double spinbox property updates the bin width integer spinbox
it.info_bar.bin_c[0].setValue(4.1)
it.info_bar.bin_c[0].editingFinished.emit()
assert it.info_bar.bin_i[0].value() == 2
# Place the cursor at x = 3, and the index of x = 1.5 => 2
# Set the bin width index to 1. This guarantees no binning.
it.info_bar.cursor_c[0].setValue(3)
it.info_bar.cursor_c[0].editingFinished.emit()
it.info_bar.bin_i[0].setValue(1)
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].yData, dat.values[:, 1, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].xData,
np.mean(dat.values[2:3, :, 0:1], axis=(0, 2)))
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].xData,
np.mean(dat.values[2:3, 1:2, :], axis=(0, 1)))
# Cursor still at 3. Now the bin width is slightly larger than delta, and there will be averaging.
it.info_bar.bin_c[0].setValue(2.1)
it.info_bar.bin_c[0].editingFinished.emit()
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].yData, dat.values[:, 1, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].xData,
np.mean(dat.values[1:3, :, 0:1], axis=(0, 2)))
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].xData,
np.mean(dat.values[1:3, 1:2, :], axis=(0, 1)))
def test_imagetool_transpose(self, qtbot):
dat = self.make_regular_data()
it = ImageTool(dat)
qtbot.addWidget(it)
it.info_bar.cursor_i[0].setValue(2)
it.info_bar.cursor_i[1].setValue(1)
it.info_bar.bin_i[0].setValue(2)
dat = dat.transpose([1, 0, 2])
it.info_bar.transpose_request.emit([1, 0, 2])
# dat = dat.transpose({0: 1, 1: 0, 2: 2})
# it.info_bar.transpose_request.emit({0: 1, 1: 0, 2: 2})
assert it.info_bar.cursor_labels[0].text() == 'y2'
assert it.info_bar.cursor_labels[1].text() == 'x4'
assert it.info_bar.cursor_labels[2].text() == 'z2'
it.info_bar.bin_i[1].setValue(2)
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].xData,
dat.axes[0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].yData,
dat.axes[1])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].yData,
dat.axes[2])
np.testing.assert_almost_equal(
it.pg_win.lineplots_data['x'][0].yData,
np.mean(dat.values[:, 0:2, 0:1], axis=(1, 2)))
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].xData,
dat.values[0, :, 0])
np.testing.assert_almost_equal(
it.pg_win.lineplots_data['z'][0].xData,
np.mean(dat.values[0:1, 0:2, :], axis=(0, 1)))
def test_imagetool_numpy(self, qtbot):
mat = self.make_numpy_data()
it = ImageTool(mat)
assert it.pg_win.lineplots_data['x'][1] == 'h'
assert np.all(it.pg_win.lineplots_data['x'][0].xData == np.arange(4))
assert np.all(it.pg_win.lineplots_data['x'][0].yData == mat[:, 0, 0])
assert it.pg_win.lineplots_data['y'][1] == 'v'
assert np.all(it.pg_win.lineplots_data['y'][0].xData == mat[0, :, 0])
assert np.all(it.pg_win.lineplots_data['y'][0].yData == np.arange(2))
assert it.pg_win.lineplots_data['z'][1] == 'v'
assert np.all(it.pg_win.lineplots_data['z'][0].xData == mat[0, 0, :])
assert np.all(it.pg_win.lineplots_data['z'][0].yData == np.arange(2))
def test_imagetool_cursor(self, qtbot):
dat = self.make_regular_data()
it = ImageTool(dat)
it.info_bar.cursor_i[0].setValue(2)
it.info_bar.cursor_i[1].setValue(1)
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].yData, dat.values[:, 1, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].xData, dat.values[2, :, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].xData, dat.values[2, 1, :])
it.info_bar.cursor_i[0].setValue(1)
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].yData, dat.values[:, 1, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].xData, dat.values[1, :, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].xData, dat.values[1, 1, :])
def test_imagetool_regular(self, qtbot):
dat = self.make_regular_data()
it = ImageTool(dat)
assert it.pg_win.lineplots_data['x'][1] == 'h'
assert np.all(it.pg_win.lineplots_data['x'][0].xData == dat.axes[0])
assert np.all(it.pg_win.lineplots_data['x'][0].yData == dat.values[:, 0, 0])
assert it.pg_win.lineplots_data['y'][1] == 'v'
assert np.all(it.pg_win.lineplots_data['y'][0].xData == dat.values[0, :, 0])
assert np.all(it.pg_win.lineplots_data['y'][0].yData == dat.axes[1])
assert it.pg_win.lineplots_data['z'][1] == 'v'
assert np.all(it.pg_win.lineplots_data['z'][0].xData == dat.values[0, 0, :])
assert np.all(it.pg_win.lineplots_data['z'][0].yData == dat.axes[2])
def test_imagetool_transpose(self, qtbot):
dat = self.make_regular_data()
it = ImageTool(dat)
it.show()
qtbot.waitForWindowShown(it)
it.info_bar.cursor_i[0].setValue(2)
it.info_bar.cursor_i[1].setValue(1)
it.info_bar.bin_i[0].setValue(2)
dat = dat.transpose({0: 1, 1: 0, 2: 2})
it.transpose_data({0: 1, 1: 0, 2: 2})
assert it.info_bar.cursor_labels[0].text() == 'y2'
assert it.info_bar.cursor_labels[1].text() == 'x4'
assert it.info_bar.cursor_labels[2].text() == 'z2'
it.info_bar.bin_i[1].setValue(2)
qtbot.waitSignal(it.info_bar.bin_i[1].valueChanged)
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].xData, dat.axes[0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].yData, dat.axes[1])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].yData, dat.axes[2])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['x'][0].yData,
np.mean(dat.values[:, 0:2, 0:1], axis=(1, 2)))
np.testing.assert_almost_equal(it.pg_win.lineplots_data['y'][0].xData, dat.values[0, :, 0])
np.testing.assert_almost_equal(it.pg_win.lineplots_data['z'][0].xData,
np.mean(dat.values[0:1, 0:2, :], axis=(0, 1)))
class Arpes:
def __init__(self, xarray_obj):
self._obj = xarray_obj
self.is_kinetic = True
self.ef = None
self.it = None
@requires_ef
def map_isoenergy_k_irreg(self,
ke=None,
be=None,
binwidth=1,
phi0=0,
theta0=0,
azimuth=0):
if ke is None:
ke = be + self.ef
iso_e = self._obj.arpes.sel_kinetic(ke - binwidth,
ke + binwidth).sum('energy')
F, T = np.meshgrid(iso_e.arpes.slit, iso_e.arpes.perp, indexing='ij')
kx = 0.512 * np.sqrt(ke) * np.sin(np.pi / 180 * (F - phi0))
ky = 0.512 * np.sqrt(ke) * np.cos(np.pi / 180 * F) * np.sin(
np.pi / 180 * (T - theta0))
if azimuth != 0:
kxp = np.cos(azimuth * np.pi / 180) * kx + np.sin(
azimuth * np.pi / 180) * ky
kyp = np.cos(azimuth * np.pi / 180) * ky - np.sin(
azimuth * np.pi / 180) * kx
else:
kxp = kx
kyp = ky
iso_e = iso_e.assign_coords({
'kx': (('slit', 'perp'), kxp),
'ky': (('slit', 'perp'), kyp)
})
return iso_e
@requires_ef
def spectra_k_irreg(self, phi0):
KE, F = np.meshgrid(self._obj.arpes.energy,
self._obj.arpes.slit,
indexing='ij')
kx = 0.512 * np.sqrt(KE) * np.sin(np.pi / 180 * (F - phi0))
self._obj = self._obj.assign_coords({
'kx': (('energy', 'slit'), kx),
'binding': (('energy', 'slit'), KE - self.ef)
})
return self._obj
def normalize(self):
self._obj.values = (self._obj.values - np.min(self._obj.values)) / (
np.max(self._obj.values) - np.min(self._obj.values))
def set_gamma(self, slit=None, perp=None):
if slit is not None:
self._obj.coords['slit'].values -= slit
if perp is not None:
self._obj.coords['perp'].values -= perp
def guess_high_symmetry(self):
x = self._obj.values
x_r = np.flip(x)
cor = np.correlate(x, x_r, mode='full')
idx = (np.argmax(cor) + 1) / 2
return np.interp(idx, np.arange(self._obj.coords['slit'].size),
self._obj.coords['slit'].values)
def plot_spectra(self, kspace=False):
if kspace:
self._obj.plot(x='kx', y='binding')
else:
self._obj.plot(x='slit', y='energy')
def plot_map(self, kspace=False):
if self._obj.ndim == 2:
if 'slit' in self._obj.coords and 'perp' in self._obj.coords:
if kspace:
if 'kx' not in self._obj.coords or 'ky' not in self._obj.coords:
raise ValueError('You have not computed kspace yet!')
self._obj.plot(x='kx', y='ky')
else:
self._obj.plot(x='slit', y='perp')
else:
raise ValueError(
'Your data is more than two dimensions. Cannot plot a map.')
@requires_ef
def waterfall(self,
spacing=0.2,
linear=True,
kspace=False,
fig=None,
ax=None,
**lineparams):
if fig is None and ax is None:
fig, ax = plt.subplots(1)
edc_range = (
np.max(self._obj, axis=self._obj.dims.index('slit')).values -
np.min(self._obj, axis=self._obj.dims.index('slit')).values)
if linear:
spacing = spacing * np.max(edc_range)
offset = 0
for i in range(self._obj['energy'].size):
mdc = self._obj.isel({'energy': i}).values + offset
if kspace:
ax.plot(self._obj['kx'].isel({
'energy': i
}).values, mdc, **lineparams)
else:
ax.plot(self._obj.coords['slit'].values, mdc, **lineparams)
if linear:
offset += spacing
else:
offset += spacing * edc_range[i]
return fig, ax
@property
def kinetic(self):
return self._obj.coords['energy'].values
@requires_ef
@property
def binding(self):
return self._obj.coords['energy'].values - self.ef
@property
def energy(self):
return self._obj.coords['energy'].values
@property
def slit(self):
return self._obj.coords['slit'].values
@property
def perp(self):
return self._obj.coords['perp'].values
@requires_ef
def sel_binding(self, *args):
if len(args) == 2:
return self._obj.sel(
{'energy': slice(args[0] + self.ef, args[1] + self.ef)})
else:
raise ValueError('binding only accepts min and max')
def sel_kinetic(self, *args):
if len(args) == 2:
return self._obj.sel({'energy': slice(args[0], args[1])})
else:
return self._obj.sel({'energy': args[0]})
def sel_slit(self, *args):
if len(args) > 1:
return self._obj.sel({'slit': slice(args[0], args[1])})
else:
return self._obj.sel({'slit': args[0]})
def sel_perp(self, *args):
if len(args) > 1:
return self._obj.sel({'perp': slice(args[0], args[1])})
else:
return self._obj.sel({'perp': args[0]})
def guess_ef(self):
edc = self._obj.copy()
if edc.ndim > 1:
for dim_label in edc.dims:
if dim_label != 'energy':
edc = edc.sum(dim_label)
if edc.size > 150:
factor = int(np.ceil(edc.size / 100))
edc = edc.arpes.downsample({'energy': factor})
edc_y = edc.values
edc_x = edc.coords['energy'].values
return edc_x[np.argmin(np.diff(edc_y))]
@staticmethod
def guess_fermi_params(edc, bg_order=1):
"""
downsamples the data so that it contains at most 150 points
estimate Ef by minimizing the first derivative
fit the data below Ef to a polynomial of order bg_order
guess that the resolution is 4.4 meV (TODO: make this more robust by finding width of peak in derivative)
estimate y0 by suggesting the average of the final few values
"""
if edc.size > 150:
factor = int(np.ceil(edc.size / 100))
edc = edc.arpes.downsample({'energy': factor})
edc_y = edc.values
edc_x = edc.coords['energy'].values
i = np.argmin(np.diff(edc_y))
ef = edc_x[i]
p = np.polyfit(edc_x[:(i - 5)], edc_y[:(i - 5)], bg_order)
kBT = 0.001
y0 = np.mean(edc_y[-5:])
return (ef, kBT, y0) + tuple(p)
@staticmethod
def fermi_fit(edc, bg_order=1):
return curve_fit(fermi_fcn_bg,
edc.arpes.energy,
edc.values,
p0=Arpes.guess_fermi_params(edc, bg_order))
@staticmethod
def print_fit_params(popt):
formatter = 'Ef: {0} (eV)\n10-90: {1} (meV)\ny0: {2}\n'
for i in range(3, len(popt)):
formatter += 'bg x**{0}: '.format(len(popt) - i - 1)
formatter += '{' + str(i) + '}\n'
vals = list(popt)
vals[1] *= 4400
print(formatter.format(*vals))
@staticmethod
def dewarp_curve(p, x):
if len(p) == 1:
return p[0]
elif len(p) == 2:
return p[0] * x + p[1]
elif len(p) == 3:
return p[0] * x**2 + p[1] * x + p[2]
elif len(p) == 4:
return p[0] * x**3 + p[1] * x**2 + p[2] * x + p[3]
@staticmethod
def make_dewarp_curve(angles, ef):
p = np.polyfit(angles, ef, deg=2)
return partial(Arpes.dewarp_curve, p)
@staticmethod
def dewarp_spectra(spectra, dewarp):
ef_pos = dewarp(spectra.coords['slit'].values)
ef_min = np.min(ef_pos)
ef_max = np.max(ef_pos)
de = spectra.coords['energy'].values[1] - spectra.coords[
'energy'].values[0]
px_to_remove = int(round((ef_max - ef_min) / de))
dewarped = np.empty((spectra.coords['energy'].size - px_to_remove,
spectra.coords['slit'].size))
for i in range(spectra.coords['slit'].size):
rm_from_bottom = int(round((ef_pos[i] - ef_min) / de))
rm_from_top = spectra.coords['energy'].size - (px_to_remove -
rm_from_bottom)
dewarped[:, i] = spectra.values[rm_from_bottom:rm_from_top, i]
bottom_energy_offset = int(round((ef_max - ef_min) / de))
energy = spectra.coords['energy'].values[bottom_energy_offset:]
return xr.DataArray(dewarped,
coords={
'energy': energy,
'slit': spectra.coords['slit'].values
},
dims=['energy', 'slit'],
attrs=spectra.attrs)
@staticmethod
def create_dewarp(spectra):
N = spectra.coords['slit'].size
slit_angles = spectra.coords['slit'].values
ef_pos = np.empty(N)
for i in range(N):
edc = spectra.isel({'slit': i}).values
energy = spectra.coords['energy'].values
params = Arpes.guess_fermi_params(spectra.isel({'slit': i}), 2)
params, pcov = curve_fit(fermi_fcn_bg, energy, edc, p0=params)
ef_pos[i] = params[0]
dewarp = Arpes.make_dewarp_curve(slit_angles, ef_pos)
return dewarp
def downsample(self, dsf, operation='mean'):
"""
dsf is a dict of dims and the amount to downsample by (i.e. 2, 3, 4, etc)
operation can be sum or mean
downsampled coordinates are always averaged
"""
dat = self._obj.copy()
new_shape = list(dat.shape)
new_coords = {key: dat[key].values for key in dat.dims}
dat_mat = dat.values
for label, ds in dsf.items():
i = dat.dims.index(label)
extra = new_shape[i] % ds
if extra == 0:
extra = None
else:
extra *= -1
new_shape[i] = new_shape[i] // ds
new_coords[label] = new_coords[label][:extra].reshape(
(-1, ds)).mean(axis=1)
slicer = [slice(None, None)] * dat_mat.ndim
slicer[i] = slice(None, extra)
dat_mat = dat_mat[tuple(slicer)]
mat = bin_ndarray(dat_mat, new_shape, operation=operation)
return xr.DataArray(mat,
coords=new_coords,
dims=dat.dims,
attrs=dat.attrs)
def plot(self, layout=ImageTool.LayoutComplete):
self.it = ImageTool(self._obj, layout=layout)
self.it.show()
return self.it
def plot(self, layout=ImageTool.LayoutComplete):
self.it = ImageTool(self._obj, layout=layout)
self.it.show()
return self.it
delta = [
f['data'][axis].attrs['delta'] for axis in ['axis0', 'axis1', 'axis2']
]
coord_min = [
f['data'][axis].attrs['offset']
for axis in ['axis0', 'axis1', 'axis2']
]
f.close()
mydata = RegularDataArray(dat_mat, delta=delta, coord_min=coord_min)
return mydata
# Start Qt event loop unless running in interactive mode.
app = QtWidgets.QApplication([])
# Generate data
# my_data = gen_data()
my_data = from_file()
my_data = my_data.sel({
'x': slice(-6, 6.62),
'y': slice(-12.4, 17.33),
'z': slice(20.93, 21.59)
})
image_tool = ImageTool(my_data, layout=ImageTool.LayoutComplete)
image_tool.show()
if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
QtWidgets.QApplication.instance().exec_()
def test_imagetool_regular_show(self, qtbot):
dat = self.make_regular_data()
it = ImageTool(dat)
it.show()
def test_imagetool_numpy_show(self, qtbot):
mat = self.make_numpy_data()
it = ImageTool(mat)
it.show()
def test_imagetool_regular_show(self, qtbot):
dat = self.make_regular_data()
it = ImageTool(dat)
qtbot.addWidget(it)
assert len(it.info_bar.cursor_i) == 3
it.show()
def test_imagetool_numpy_show(self, qtbot):
mat = self.make_numpy_data()
it = ImageTool(mat)
qtbot.addWidget(it)
assert len(it.info_bar.cursor_i) == 3
it.show()