def build_quadratic_fermi_edge_correction(arr: xr.DataArray,
fit_limit=0.001,
eV_slice=None,
plot=False) -> lf.model.ModelResult:
# TODO improve robustness here by allowing passing in the location of the fermi edge guess
# We could also do this automatically by using the same method we use for step detection to find the edge of the
# spectrometer image
if eV_slice is None:
approximate_fermi_level = arr.S.find_spectrum_energy_edges().max()
eV_slice = slice(approximate_fermi_level - 0.4,
approximate_fermi_level + 0.4)
else:
approximate_fermi_level = 0
sum_axes = exclude_hemisphere_axes(arr.dims)
edge_fit = broadcast_model(
GStepBModel,
arr.sum(sum_axes).sel(eV=eV_slice),
'phi',
params={'center': {
'value': approximate_fermi_level
}})
size_phi = len(arr.coords['phi'])
not_nanny = (np.logical_not(np.isnan(arr)) * 1).sum('eV') > size_phi * 0.30
condition = np.logical_and(edge_fit.F.s('center') < fit_limit, not_nanny)
quadratic_corr = QuadraticModel().guess_fit(edge_fit.F.p('center'),
weights=condition * 1)
if plot:
edge_fit.F.p('center').plot()
plt.plot(arr.coords['phi'], quadratic_corr.best_fit)
return quadratic_corr
def fit_for_effective_mass(data: DataType, fit_kwargs=None):
"""
Performs an effective mass fit by first fitting for Lorentzian lineshapes and then fitting a quadratic
model to the result. This is an alternative to global effective mass fitting.
In the case that data is provided in anglespace, the Lorentzian fits are performed in anglespace
before being converted to momentum where the effective mass is extracted.
We should probably include uncertainties here.
:param data:
:param fit_kwargs: Passthrough for arguments to `broadcast_model`,
used internally to obtain the Lorentzian peak locations
:return:
"""
if fit_kwargs is None:
fit_kwargs = {}
data = normalize_to_spectrum(data)
mom_dim = [d for d in ['kp', 'kx', 'ky', 'kz', 'phi', 'beta', 'theta'] if d in data.dims][0]
results = broadcast_model([LorentzianModel, AffineBackgroundModel], data, mom_dim, **fit_kwargs)
if mom_dim in {'phi', 'beta', 'theta'}:
forward = convert_coordinates_to_kspace_forward(data)
final_mom = [d for d in ['kx', 'ky', 'kp', 'kz'] if d in forward][0]
eVs = results.F.p('a_center').values
kps = [forward[final_mom].sel(eV=eV, **dict([[mom_dim, ang]]), method='nearest') for
eV, ang in zip(eVs, data.coords[mom_dim].values)]
quad_fit = QuadraticModel().fit(eVs, x=np.array(kps))
return HBAR_SQ_EV_PER_ELECTRON_MASS_ANGSTROM_SQ / (2 * quad_fit.params['a'].value)
quad_fit = QuadraticModel().guess_fit(results.F.p('a_center'))
return HBAR_SQ_EV_PER_ELECTRON_MASS_ANGSTROM_SQ / (2 * quad_fit.params['a'].value)
def build_photon_energy_fermi_edge_correction(arr: xr.DataArray,
plot=False,
energy_window=0.2):
edge_fit = broadcast_model(
GStepBModel,
arr.sum(exclude_hv_axes(
arr.dims)).sel(eV=slice(-energy_window, energy_window)), 'hv')
return edge_fit
def build_cycle_fermi_edge_correction(data: DataType, energy_range=None):
arr = normalize_to_spectrum(data)
if 'pixels' in arr.dims or 'phi' in arr.dims:
arr = arr.S.region_sel('wide_angular')
if energy_range is None:
energy_range = slice(-0.1, 0.1)
arr = arr.S.sum_other(['eV', 'cycle'])
return broadcast_model(GStepBModel, arr.sel(eV=energy_range), 'cycle')
def fs_gap(data: DataType, shape=None, energy_range=None):
data = normalize_to_spectrum(data)
if energy_range is None:
energy_range = slice(-0.1, None)
data.sel(eV=energy_range)
reduction = None
if shape is None:
# Just rebin the data along 'phi'
reduction = {'phi': 16}
data = rebin(data, reduction=reduction, shape=shape)
return broadcast_model(GStepBModel, data, ['phi', 'beta'])
def on_add_new_peak(selection):
amplitude = data.sel(**selection).mean().item()
selection = selection[data.dims[0]]
center = (selection.start + selection.stop) / 2
sigma = (selection.stop - selection.start)
model_settings.append({
'center': {
'value': center,
'min': center - sigma,
'max': center + sigma
},
'sigma': {
'value': sigma
},
'amplitude': {
'min': 0,
'value': amplitude
}
})
model_defs.append(LorentzianModel)
if model_defs:
results = broadcast_model(model_defs,
for_fit,
'fit_dim',
params=compute_parameters())
result = results.results[0].item()
if result is not None:
# residual
for_residual = data.copy(deep=True)
for_residual.values = result.residual
residual_view.data = for_residual
# fit_result
for_best_fit = data.copy(deep=True)
for_best_fit.values = result.best_fit
fitted_view.data = for_best_fit
# initial_fit_result
for_initial_fit = data.copy(deep=True)
for_initial_fit.values = result.init_fit
initial_fit_view.data = for_initial_fit
ax_fitted.set_ylim(ax_initial.get_ylim())
def fit_fermi_edge(data, energy_range=None):
"""
Fits a Fermi edge. Not much easier than doing it manually, but this can be
useful sometimes inside procedures where you don't want to reimplement this logic.
:param data:
:param energy_range:
:return:
"""
if energy_range is None:
energy_range = slice(-0.1, 0.1)
broadcast_directions = list(data.dims)
broadcast_directions.remove('eV')
assert len(broadcast_directions) == 1 # for now we don't support more
edge_fit = broadcast_model(GStepBModel, data.sel(eV=energy_range), broadcast_directions[0])
return edge_fit
def build_direct_fermi_edge_correction(arr: xr.DataArray,
fit_limit=0.001,
energy_range=None,
plot=False,
along='phi'):
"""
Builds a direct fermi edge correction stencil.
This means that fits are performed at each value of the 'phi' coordinate
to get a list of fits. Bad fits are thrown out to form a stencil.
This can be used to shift coordinates by the nearest value in the stencil.
:param copper_ref:
:param args:
:param kwargs:
:return:
"""
if energy_range is None:
energy_range = slice(-0.1, 0.1)
exclude_axes = ['eV', along]
others = [d for d in arr.dims if d not in exclude_axes]
edge_fit = broadcast_model(GStepBModel,
arr.sum(others).sel(eV=energy_range),
along).results
def sieve(c, v):
return v.item().params['center'].stderr < 0.001
corrections = edge_fit.T.filter_coord(
along, sieve).T.map(lambda x: x.params['center'].value)
if plot:
corrections.plot()
return corrections
def fermi_edge_reference(data, title=None, ax=None, out=None, norm=None, **kwargs):
warnings.warn('Not automatically correcting for slit shape distortions to the Fermi edge')
sum_dimensions = {'cycle', 'phi', 'kp', 'kx'}
sum_dimensions.intersection_update(set(data.dims))
summed_data = data.sum(*list(sum_dimensions))
broadcast_dimensions = summed_data.dims
broadcast_dimensions.remove('eV')
if len(broadcast_dimensions) == 1:
edge_fit = broadcast_model(GStepBModel, summed_data.sel(eV=slice(-0.1, 0.1)), broadcast_dimensions[0])
else:
warnings.warn('Could not product fermi edge reference. Too many dimensions: {}'.format(broadcast_dimensions))
return
centers = apply_dataarray(edge_fit, np.vectorize(lambda x: x.params['center'].value, otypes=[np.float]))
widths = apply_dataarray(edge_fit, np.vectorize(lambda x: x.params['width'].value, otypes=[np.float]))
if ax is None:
_, ax = plt.subplots(figsize=(8, 5))
if title is None:
title = data.S.label.replace('_', ' ')
plot_centers = centers.plot(norm=norm, ax=ax)
plot_widths = widths.plot(norm=norm, ax=ax)
ax.set_xlabel(label_for_dim(data, ax.get_xlabel()))
ax.set_ylabel(label_for_dim(data, ax.get_ylabel()))
ax.set_title(title, font_size=14)
if out is not None:
plt.savefig(path_for_plot(out), dpi=400)
return path_for_plot(out)
plt.show()