def set_tracewidths(self, mkplot=True, npix=5):
# Start by doing a simple box car extraction
tlist = self.get_tracelist()
ntrace = len(tlist)
arr = self.get_flat2d()
out, outivar = boxcar_cutout(tlist, arr, npix=npix)
# Do a crude normalization
norm = out.sum(axis=1)
out1 = out/(norm[:,np.newaxis,:]+1.e-10)
# Genetaye profiles
csec, qaz = weighted_average_clip(out1, outivar)
widths = []
if mkplot :
plotfn = '%s_trace_width_qa.pdf'%self.name
plotcc = PdfPages(plotfn)
for ii in range(ntrace):
# Do the actual fits
fitcc = GaussFit(np.arange(2*npix+1), csec[ii,:])
fitcc.optimize(frel=1.e-4, fabs=1.e-6)
widths.append(fitcc.xopt[0])
if mkplot :
plt.clf()
plt.plot(csec[ii, :])
plt.plot(fitcc.model(fitcc.xopt), 'ro')
plotcc.savefig()
# Clean up
plotcc.close()
self.traces['widths'] = widths
def gaussfit(self):
if self._gf is not None and (self._xx.min(), self._xx.max(), self._xx.sum()) == self._gf_xx and (self._yy.min(), self._yy.max(), self._yy.sum()) == self._gf_yy:
return self._gf
self._gf = GaussFit(self._xx, self._yy, fit_const=False)
self._gf_xx = (self._xx.min(), self._xx.max(), self._xx.sum())
self._gf_yy = (self._yy.min(), self._yy.max(), self._yy.sum())
return self._gf
def gauss_slice(self, debug=False):
mask = self._get_mask()
mean_list = np.zeros_like(self.averaged_x_axis, dtype=float)
std_list = mean_list.copy()
#for i in range(np.shape(self.image_averaged)[1]):
for i in np.argwhere(mask)[:, 0]:
gf = GaussFit(self.y_axis, self.image_averaged[:, i])
if self.y_axis.min() < gf.mean < self.y_axis.max():
mean_list[i] = gf.mean
std_list[i] = abs(gf.sigma)
if debug:
import matplotlib.pyplot as plt
plt.figure()
sp = plt.subplot(1, 1, 1)
gf.plot_data_and_fit(sp)
plt.show()
import pdb
pdb.set_trace()
return mean_list, std_list
def profile_from_blmeas(file_, tt_halfrange, charge, energy_eV, subtract_min=False, zero_crossing=1):
bl_meas = data_loader.load_blmeas(file_)
time_meas0 = bl_meas['time_profile%i' % zero_crossing]
current_meas0 = bl_meas['current%i' % zero_crossing]
if subtract_min:
current_meas0 = current_meas0 - current_meas0.min()
if tt_halfrange is None:
tt, cc = time_meas0, current_meas0
else:
current_meas0 *= charge/current_meas0.sum()
gf_blmeas = GaussFit(time_meas0, current_meas0)
time_meas0 -= gf_blmeas.mean
time_meas1 = np.arange(-tt_halfrange, time_meas0.min(), np.diff(time_meas0).mean())
time_meas2 = np.arange(time_meas0.max(), tt_halfrange, np.diff(time_meas0).mean())
time_meas = np.concatenate([time_meas1, time_meas0, time_meas2])
current_meas = np.concatenate([np.zeros_like(time_meas1), current_meas0, np.zeros_like(time_meas2)])
tt, cc = time_meas, current_meas
return BeamProfile(tt, cc, energy_eV, charge)
### Forward track with elegant for measured beam profile
ms.figure(
'Forward and backward tracking - Ignore natural beamsize (200 nm emittance)'
)
subplot = ms.subplot_factory(2, 3)
sp_ctr = 1
bl_meas = data_loader.load_blmeas(bl_meas_file)
simulator = elegant_matrix.get_simulator(magnet_file)
time_meas0 = bl_meas['time_profile1']
current_meas0 = bl_meas['current1']
current_meas0 *= 200e-12 / current_meas0.sum()
gf_blmeas = GaussFit(time_meas0, current_meas0)
time_meas0 -= gf_blmeas.mean
time_meas1 = np.arange(-tt_halfrange, time_meas0.min(),
np.diff(time_meas0).mean())
time_meas2 = np.arange(time_meas0.max(), tt_halfrange,
np.diff(time_meas0).mean())
time_meas = np.concatenate([time_meas1, time_meas0, time_meas2])
current_meas = np.concatenate(
[np.zeros_like(time_meas1), current_meas0,
np.zeros_like(time_meas2)])
sim1, mat_dict, wf_dicts, disp_dict = simulator.simulate_streaker(
time_meas,
current_meas,
timestamp, (gap, 10e-3), (beam_offset, 0),
def plot_img_and_proj(self, sp, x_factor=None, y_factor=None, plot_proj=True, log=False, revert_x=False, plot_gauss=True, slice_dict=None, xlim=None, ylim=None, cmapname='hot'):
def unit_to_factor(unit):
if unit == 'm':
factor = 1e3
elif unit == 's':
factor = 1e15
elif unit == 'eV':
factor = 1e-6
else:
factor = 1
return factor
if x_factor is None:
x_factor = unit_to_factor(self.x_unit)
if y_factor is None:
y_factor = unit_to_factor(self.y_unit)
x_axis, y_axis, image = self.x_axis, self.y_axis, self.image
if xlim is None:
index_x_min, index_x_max = 0, len(x_axis)
else:
index_x_min, index_x_max = sorted([np.argmin((x_axis-xlim[1])**2), np.argmin((x_axis-xlim[0])**2)])
if ylim is None:
index_y_min, index_y_max = 0, len(y_axis)
else:
index_y_min, index_y_max = sorted([np.argmin((y_axis-ylim[1])**2), np.argmin((y_axis-ylim[0])**2)])
x_axis = x_axis[index_x_min:index_x_max]
y_axis = y_axis[index_y_min:index_y_max]
image = image[index_y_min:index_y_max,index_x_min:index_x_max]
extent = [x_axis[0]*x_factor, x_axis[-1]*x_factor, y_axis[0]*y_factor, y_axis[-1]*y_factor]
if log:
image_ = np.clip(image, 1, None)
log = np.log(image_)
else:
log = image
sp.imshow(log, aspect='auto', extent=extent, origin='lower', cmap=plt.get_cmap(cmapname))
if slice_dict is not None:
old_lim = sp.get_xlim(), sp.get_ylim()
for key_mean, key_sigma, color in [
#('slice_mean_rms', 'slice_rms_sq', 'blue'),
('slice_cut_mean', 'slice_cut_rms_sq', 'orange'),
#('slice_mean', 'slice_sigma_sq', 'red'),
]:
xx = slice_dict['slice_x']*x_factor
yy = slice_dict[key_mean]*y_factor
yy_err = np.sqrt(slice_dict[key_sigma])*y_factor
sp.errorbar(xx, yy, yerr=yy_err, color=color, ls='None', marker='None', lw=.75)
sp.set_xlim(*old_lim[0])
sp.set_ylim(*old_lim[1])
if plot_proj:
proj = image.sum(axis=-2)
proj_plot = (y_axis.min() +(y_axis.max()-y_axis.min()) * proj/proj.max()*0.3)*y_factor
sp.plot(x_axis*x_factor, proj_plot, color='red')
if plot_gauss:
gf = GaussFit(x_axis, proj_plot-proj_plot.min(), fit_const=False)
sp.plot(x_axis*x_factor, gf.reconstruction+proj_plot.min(), color='orange')
#import matplotlib.pyplot as plt
#plt.figure()
#sp = plt.subplot(1,1,1)
#gf.plot_data_and_fit(sp)
#plt.show()
#import pdb; pdb.set_trace()
proj = image.sum(axis=-1)
proj_plot = (x_axis.min() +(x_axis.max()-x_axis.min()) * proj/proj.max()*0.3)*x_factor
sp.plot(proj_plot, y_axis*y_factor, color='red')
if plot_gauss:
gf = GaussFit(y_axis, proj_plot-proj_plot.min(), fit_const=False)
sp.plot(gf.reconstruction+proj_plot.min(), y_axis*y_factor, color='orange')
if revert_x:
xlim = sp.get_xlim()
sp.set_xlim(*xlim[::-1])
def y_to_eV(self, dispersion, energy_eV, ref_y=None):
if ref_y is None:
ref_y = GaussFit(self.y_axis, np.sum(self.image, axis=-1)).mean
#print('y_to_eV', ref_y*1e6, ' [um]')
E_axis = (self.y_axis-ref_y) * dispersion * energy_eV
return self.child(self.image, self.x_axis, E_axis, y_unit='eV', ylabel='$\Delta$ E (MeV)'), ref_y
def fit_slice(self, smoothen_first=True, smoothen=100e-6, intensity_cutoff=None, charge=1, rms_sigma=5, noise_cut=0.1):
y_axis = self.y_axis
n_slices = len(self.x_axis)
pixelsize = abs(y_axis[1] - y_axis[0])
smoothen = smoothen/pixelsize
slice_mean = []
slice_sigma = []
slice_gf = []
slice_rms = []
slice_mean_rms = []
slice_cut_rms = []
slice_cut_mean = []
for n_slice in range(n_slices):
intensity = self.image[:,n_slice]
if smoothen_first:
yy_conv = gaussian_filter1d(intensity, smoothen)
gf0 = GaussFit(y_axis, yy_conv, fit_const=True)
p0 = gf0.popt
else:
p0 = None
try:
gf = GaussFit(y_axis, intensity, fit_const=True, p0=p0, raise_=True)
except RuntimeError:
slice_mean.append(np.nan)
slice_sigma.append(np.nan)
slice_gf.append(None)
slice_mean_rms.append(np.nan)
slice_rms.append(np.nan)
slice_cut_rms.append(np.nan)
slice_cut_mean.append(np.nan)
else:
where_max = y_axis[np.argmax(intensity)]
mask_rms = np.logical_and(
y_axis > where_max - abs(gf.sigma)*rms_sigma,
y_axis < where_max + abs(gf.sigma)*rms_sigma)
y_rms = y_axis[mask_rms]
data_rms = intensity[mask_rms]
if np.sum(data_rms) == 0:
mean_rms, rms = 0, 0
else:
mean_rms = np.sum(y_rms*data_rms)/np.sum(data_rms)
rms = np.sum((y_rms-mean_rms)**2*data_rms)/np.sum(data_rms)
slice_gf.append(gf)
slice_mean.append(gf.mean)
slice_sigma.append(gf.sigma**2)
slice_rms.append(rms)
slice_mean_rms.append(mean_rms)
intensity = intensity.copy()
intensity[np.logical_or(y_axis<mean_rms-1.5*rms, y_axis>mean_rms+1.5*rms)]=0
prof_y = intensity-intensity.min()
if np.all(prof_y == 0):
slice_cut_rms.append(0)
slice_cut_mean.append(0)
else:
profile = AnyProfile(y_axis, prof_y)
profile.cutoff2(noise_cut)
profile.crop()
slice_cut_rms.append(profile.rms()**2)
slice_cut_mean.append(profile.mean())
# Debug bad gaussfits
#if 101e-15 < self.x_axis[n_slice] < 104e-15:
#if abs(gf.sigma) < 1e5:
# import matplotlib.pyplot as plt
# num = plt.gcf().number
# plt.figure()
# plt.suptitle('Debug 38 fs')
# sp = plt.subplot(1,1,1)
# gf.plot_data_and_fit(sp)
# sp.legend()
# plt.figure(num)
# plt.show()
# import pdb; pdb.set_trace()
proj = np.sum(self.image, axis=-2)
proj = proj / np.sum(proj) * charge
current = proj / (self.x_axis[1] - self.x_axis[0])
slice_dict = {
'slice_x': self.x_axis,
'slice_mean': np.array(slice_mean),
'slice_sigma_sq': np.array(slice_sigma),
'slice_rms_sq': np.array(slice_rms),
'slice_mean_rms': np.array(slice_mean_rms),
'slice_gf': slice_gf,
'slice_intensity': proj,
'slice_current': current,
'slice_cut_rms_sq': np.array(slice_cut_rms),
'slice_cut_mean': np.array(slice_cut_mean),
}
if intensity_cutoff:
mask = proj > proj.max()*intensity_cutoff
for key, value in slice_dict.items():
if hasattr(value, 'shape') and value.shape == mask.shape:
slice_dict[key] = value[mask]
return slice_dict